知識蒸留 (knowledge distillation)
定義
知識蒸留knowledge distillationとは、大きく重いモデルが学習した知識を小さく軽いモデルへ移し、小さいモデルが大きいモデルに匹敵する性能を出せるようにする方法論である。このとき、大きく重いモデルを教師teacher、小さく軽いモデルを生徒studentという。モデル圧縮model compressionの一分野である。
説明1 2
性能の良いニューラルネットワークは概して大きい。パラメータが数億個に達するモデルや、複数のモデルの予測を平均するアンサンブルensembleは優れた性能を示すが、計算量と記憶容量を多く要求するため、携帯電話や組み込み機器のように資源が制限された環境に配備するのが難しい。学習に使うモデルと実際の配備に使うモデルの要求条件が互いに異なるのである。知識蒸留はこの隔たりを埋めるため、知識を失わずにモデルを配備に適した形へと変える。
重いモデルの知識を小さいモデルへ移すというとき、まず思い浮かぶ方法は、教師モデルが当てた正解、すなわち正解ラベルをそのまま生徒モデルの学習に使うことである。教師モデルを$\operatorname{Teacher}(x)$、生徒モデルを$\operatorname{Student}(x; \theta)$とするとき、入力$x$に対する二つのモデルの出力の差をデータ全体にわたって最小化するのである。
$$ \mathcal{L}(\theta) = \mathbb{E}_{x}\left[ \bigl\lVert \operatorname{Teacher}(x) - \operatorname{Student}(x; \theta) \bigr\rVert^{2} \right] $$
しかし、教師の出力をこのように一つに定まった予測としてのみ見るならば、結局のところ正解ラベルだけで小さいモデルを直接学習させることと大差ない。知識蒸留の核心は、学習されたモデルが持つ知識をパラメータの具体的な値と同一視しない点にある。知識をもう少し抽象的に眺めると、それは入力ベクトルを出力ベクトルへ移す写像mappingだと言える。複数のクラスを分類するモデルは正解クラスの確率を最大化するように学習されるが、その副産物として不正解クラスにもそれぞれ異なる確率を割り当てる。
たとえばBMWの自動車の写真を分類するとき、それをゴミ収集車と誤認する確率は非常に小さいとしても、ニンジンと誤認する確率よりははるかに大きい。このように不正解どうしの相対的な確率は、教師モデルがデータをどのように一般化generalizeするかに関する豊かな情報を含んでいる。知識蒸留はまさにこの情報を生徒モデルへ伝える。正解一つだけを教える代わりに、教師モデルが各クラスに割り当てた確率分布全体を学習目標とするのである。このようにクラス全体にわたって滑らかに広がった確率分布をソフトターゲットsoft targetといい、正解だけが$1$で残りは$0$であるワンホット形式の正解ラベルをハードターゲットhard targetという。ソフトターゲットはハードターゲットより学習事例一つあたりに多くの情報を含んでおり、事例間の勾配の分散も小さい。したがって、生徒モデルはもとの教師モデルを学習させたときよりはるかに少ないデータと高い学習率でも学習されうる。
温度付きソフトマックスと損失関数
分類のためのニューラルネットワークは、通常最後にソフトマックス層を置き、各クラスに対して計算された値であるロジットlogit $z_{i}$を確率$q_{i}$へ変換する。
$$ q_{i} = \dfrac{\exp(z_{i})}{\sum_{j} \exp(z_{j})} $$
知識蒸留では、ここに温度temperature $T$を導入して、次のように変形されたソフトマックス関数を使う。
$$ q_{i} = \dfrac{\exp(z_{i} / T)}{\sum_{j} \exp(z_{j} / T)} $$
$T = 1$ならば通常のソフトマックスと同じで、$T$を大きくするほどクラスに対する確率分布はより滑らかになる。
上の図は、同じロジットに温度だけを変えて適用した結果である。$T = 1$では確率が正解クラス一つにほとんど集中しているが、$T$を大きくするほど分布が次第に均等に広がり、不正解クラスに付けられた小さな確率どうしの相対的な大きさが現れてくる。
教師モデルのロジットを$v_{i}$、生徒モデルのロジットを$z_{i}$とし、温度$T$でそれぞれが作るソフトな確率を$p_{i}$、$q_{i}$としよう。生徒モデルは次の二つの損失関数の加重平均で学習される。
$$ \mathcal{L} = \alpha T^{2}\mathcal{L}_{\text{soft}} + (1 - \alpha)\mathcal{L}_{\text{hard}} $$
- $\mathcal{L}_{\text{soft}}$は、温度$T$における教師のソフトターゲット$p_{i}$と生徒のソフトな予測$q_{i}$の間のクロスエントロピーである。
- $\mathcal{L}_{\text{hard}}$は、温度$T = 1$における正解ラベルと生徒の予測の間のクロスエントロピーである。
- $\alpha$は二つの項の比重を調節する重みである。
温度$T$はソフトターゲットがどれほど滑らかであるかを調節する。MNISTのように教師モデルがほとんど常に正解を非常に高い確信で当てる課題では、学習された関数に関する情報の大部分がごく小さい確率どうしの比率の中に隠れている。（上の図を参照）たとえばある数字$2$は、$3$である確率が$10^{-6}$、$7$である確率が$10^{-9}$と付けられうるが、これらの値はどの$2$が$3$に似ていて、どの$2$が$7$に似ているかを教えてくれる貴重な情報である。しかし確率が$0$にあまりに近いため、通常のクロスエントロピーにはほとんど影響を与えない。温度を上げるとこれらの小さな確率が際立ち、その中に含まれる類似度の構造が学習に反映される。上の図で$T$を$1$から$3$へ上げると、$3$である確率が$10^{-6}$から$0.1$へ、$7$である確率が$10^{-9}$から$0.01$へと大きくなり、二つのクラスの間の相対的な大きさがはるかによく現れる。しかし$T = 10$のように温度を高く設定しすぎると、正解クラスと不正解クラスの区別がぼやけてしまう。
一方、ソフトターゲットが作る勾配の大きさは$1/T^{2}$に比例して小さくなる。そこでハードターゲットと併用するときは、ソフト項に$T^{2}$を掛けて、二つの項の相対的な寄与が温度に関わらずおおむね維持されるようにする。これが損失関数において$\mathcal{L}_{\text{soft}}$の前に$T^{2}$が付く理由である。
一般化
知識蒸留は実は、生徒モデルのロジットを教師モデルのロジットに合わせることの一般化と見なせる。生徒モデルが予測した確率を$q_{i} = e^{z_{i}/T} / \sum_{j} e^{z_{j}/T}$、教師モデルが予測した確率を$p_{i} = e^{v_{i}/T} / \sum_{j} e^{v_{j}/T}$としよう。すると、$i$番目のクラスに対する生徒のロジット$z_{i}$に寄与するクロスエントロピーの勾配は次のとおりである。
$$ \begin{align*} \dfrac{\partial C}{\partial z_{i}} &= \dfrac{\partial}{\partial z_{i}} \left( -\sum_{k} p_{k} \log q_{k} \right) \\ &= \dfrac{\partial}{\partial z_{i}} \left[ -\sum_{k} p_{k} \left( \dfrac{z_{k}}{T} - \log \sum_{j} e^{z_{j}/T} \right) \right] \\ &= \dfrac{\partial}{\partial z_{i}} \left( -\dfrac{1}{T}\sum_{k} p_{k} z_{k} + \log \sum_{j} e^{z_{j}/T} \right) \\ &= -\dfrac{p_{i}}{T} + \dfrac{1}{T} \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} \\ &= \dfrac{1}{T}\left( \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} - p_{i} \right) \\ &= \dfrac{1}{T}\left( \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} - \dfrac{e^{v_{i}/T}}{\sum_{j} e^{v_{j}/T}} \right) \end{align*} $$
温度$T$がロジットの大きさに比べて十分に大きければ、指数関数のマクローリン級数展開により、$e^{x/T} \approx 1 + x/T$と近似できるので、次を得る。
$$ \dfrac{\partial C}{\partial z_{i}} \approx \dfrac{1}{T}\left( \dfrac{1 + z_{i}/T}{N + \sum_{j} z_{j}/T} - \dfrac{1 + v_{i}/T}{N + \sum_{j} v_{j}/T} \right) $$
ここで$N$はクラスの個数である。すべてのクラスに対するロジットの平均を$0$と仮定すると$(\sum_{j} z_{j} = \sum_{j} v_{j} = 0)$、上の式は次のように簡単になる。
$$ \dfrac{\partial C}{\partial z_{i}} \approx \dfrac{1}{N T^{2}}(z_{i} - v_{i}) $$
すなわち、温度が高い極限において知識蒸留は$\frac{1}{2}(z_{i} - v_{i})^{2}$を最小化すること、つまり生徒と教師のロジットの差を縮めることと同じになる。逆に温度が低ければ、平均よりはるかに小さい（非常に負である）ロジットにはほとんど注意を払わなくなる。このようなロジットは教師モデルの学習過程でほとんど制約を受けておらず、雑音でありうるため、これを無視するほうが役に立つこともある。生徒モデルが教師の知識をすべて収めるには小さすぎるときは、中程度の温度が最もよく機能することが知られている。