Coda

メモ Distilling the Knowledge in a Neural Network

August 24, 2019

概要

表題にあるニューラルネットワークの蒸留についての論文を紹介する。 蒸留は、既存のモデルを使い、できるだけ予測性能を落とさずに、より小さいモデルを作るための学習手法である。 既存のモデルとして想定されているのは、複数のモデルからなるモデルや正則化された大きなモデルのように予測性能は高いが計算コストが高いものであり、 蒸留の目的は本番の運用に耐えられるデプロイ可能なモデル作ることにある。 本論文は、出力層の活性化関数に温度つきソフトマックスを使った多クラス分類のモデルを蒸留する手法を提案し、実験により手法を評価している。

蒸留に必要な訓練データには、ラベルなしのデータか蒸留するモデルと同じ訓練データのどちらかを使う。 次の説明では、ラベルなしのデータを使うことを前提に説明する。

はじめに、ラベルなしのデータに対して蒸留するモデルを活性化関数を温度つきソフトマックスとして確信度を割り当て、教師データを用意する。 温度\(T\)の温度つきソフトマックスは次の式で与えられる。

$$ 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\)をかけて訓練データのラベルに対応する勾配とスケールを等しくする必要がある。

感想

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


論文はこちらからダウンロードできます。