知識蒸留：教師から生徒へ有用な予測を伝える

定義

知識蒸留では教師モデルが生徒モデルを導く。生徒は小さいことが多いが、この手法を定義するのはモデルサイズではなく追加の学習信号である。古典的な分類の蒸留では、教師は一つの正解候補だけでなくクラス全体の確率分布を与え、代替候補の関係を示す。[1]

仕組み

softmax の前に教師と生徒のロジットを温度 T で割る。高い温度は分布を滑らかにする。代表的な目的関数は L = (1-alpha)*CE(y, p_s) + alpha*T^2*KL(p_t(T) || p_s(T)) である。ラベル項の p_s は温度1を使い、蒸留項は同じ T で両モデルを比較する。Alpha が相対的な重みを決める。T の二乗を掛けると温度による勾配のスケーリングをおおむね補正できるが、全温度で同じ挙動になる保証ではない。[1][2]

実践手順

学習、検証、テストのデータを分離する。教師から生徒の学習へテスト情報を持ち込まない。通常のオフライン蒸留では教師を固定して評価モードにし、教師の勾配を保存しない。対応する入力を両モデルに与え、クラスの並びをそろえる。[2]

ラベルのみで学習する生徒の基準を作り、温度と損失の重みを検証データで選ぶ。KL はライブラリが求める確率または対数確率の規約に従い、クラスとサンプルに対する集約方法を記録する。選択が終わるまで最終テストを使わない。教師のバージョン、前処理、転移用データ、生徒の構造を記録する。

説明用の例

問い合わせ分類の説明用の教師が、配送0.70、返金0.25、アカウント0.05を出したとする。配送というハードラベルだけでは代替候補の違いは分からないが、分布は返金がより近い競合解釈であることを示す。生徒は正解ラベルとともにこの構造を学べる。数値は説明用であり、性能結果や特定の温度変換を表すものではない。

限界と検証

教師は誤りや偏りも伝え得る。非常に小さな生徒は有用な違いを再現する容量が不足し、ドメインの変化は自信の高い教師出力さえ誤解を招くものにし得る。分布の一致だけではタスク性能の改善を証明できない。[1][2] ラベル空間やトークン語彙が異なる場合には明示的な対応付けが必要であり、単純なクラス別目的関数はそのまま使えない。

蒸留した生徒、ラベルのみの生徒、教師を独立した指標、重要な部分集団、遅延、メモリで比較する。学習損失だけでなく、配備する前処理や出力挙動を評価する。本記事のためのローカル蒸留実験や圧縮ベンチマークは実施していない。

参考資料

[1] Hinton, Vinyals and Dean: Distilling the Knowledge in a Neural Network [2] PyTorch: Knowledge Distillation Tutorial