Distilling the Knowledge in a Neural Network (2015)

#知識蒸留

August 24, 2019

概要

蒸留(Distilling)は、既存のモデルの予測性能をできるだけ落とさずに、より小さいモデルを作るための手法である。複数のモデルからなるモデルや正則化された大きなモデルのように予測性能は高いが、同時に計算コストも高くつく。蒸留の目的は本番運用に耐えられるデプロイ可能なモデル作ることにである。 Distilling the Knowledge in a Neural Networkは、出力層の活性化関数に温度つきソフトマックスを使った多クラス分類のモデルを蒸留する手法を提案、評価した。

温度$T$の温度つきソフトマックスは次の式で表せる。$T=1$であれば通常のソフトマックスである。

$$ q\_i = \frac{\exp(z\_i/T)}{\sum\_{j}\exp(z\_j/T)} $$

蒸留に使う訓練データセットは、ラベルなしのデータか、蒸留するモデルと同じ訓練データのどちらかである。ラベルなしのデータセットであれば、高温度のソフトマックス関数と、もとのモデルを使い、データセットに確信度を割りあてる。次に、同じ高温度で小さいモデルを学習する。損失関数を交差エントロピー、ラベルである確信度を$q_i$とすると交差エントロピーの勾配は

$$ \frac{\sigma C}{\sigma z\_i} = \frac{1}{T}(q\_i - p\_i)=\frac{1}{T}\left(\frac{e^{z\_i/T}}{\sum\_j e^{z\_j/T}}-\frac{e^{v\_i/T}}{\sum\_j e^{v\_j/T}}\right) $$

となる。$T$が大きいとき、マクローリン展開より$e^x \approx 1+x$であるから、

$$ \frac{\sigma C}{\sigma z\_i}\approx \frac{1}{T}\left(\frac{1+z\_i/T}{N+\sum\_jz\_j/T} - \frac{1+v\_i/T}{N+\sum\_jv\_j/T}\right) $$

であり、さらに、$\sum_j z_j = \sum_j v_j =0$であれば、勾配は

$$ \frac{\sigma C}{\sigma z\_i} \approx \frac{1}{NT^2}(z\_i - v\_i) $$

のように単純化できる。

蒸留された小さいモデルで学習するときは$T=1$で確信度を求める。

なお、蒸留するモデルと同じ訓練データで小さいモデルを学習するときは、蒸留するモデルの出力に対する交差エントロピーと訓練データのラベルの分布に対する交差エントロピーの重みつき平均を損失関数に設定する。ただし、上の式より、モデルの出力に対応する交差エントロピーの勾配は$\frac{1}{T^2}$だけ訓練データのラベルの方とスケールが異なるので、勾配に対して$T^2$をかけて訓練データのラベルに対応する勾配とスケールを等しくする必要がある。

感想

論文ではもう一つの貢献として、上でまとめた蒸留に加え、特定のクラスの分類に特化した複数のモデルからなるモデルの学習方法の導入がある。クラス数が多い分類のタスクを解くときの参考にしたい。