蒸留レシピ図鑑 — logit・feature・attention・self をどう使い分けるか

蒸留のレシピは「教師のどこを学生に合わせるか」で決まります。出力(logit)・中間層(FitNet)・注意行列・自己蒸留の4種を、1枚の表と4本の式で並べ、教師がAPI越しなのか同族アーキなのかといった条件から選び方を決める地図にします。

Distilling the Knowledge in a Neural Network

「真似る」場所は、出口だけではない

料理の師匠から学ぶとき、真似る対象はいくつもあります。完成した皿の味だけを真似ることもできるし、下ごしらえで野菜がどう切られているかを真似ることもできる。師匠が鍋のどこを見ているか、その視線ごと真似る道もある。あるいは昨日の自分の出来映えを今日の手本にする、という独学の道もあります。

知識蒸留のレシピの違いは、突き詰めるとこれだけです。教師モデルのどこを学生モデルに合わせるか。合わせる場所が変われば、教師に要求する情報も、損失関数の形も、うまくいく条件も変わります。

蒸留の土台 — なぜ正解ラベルより教師の出力のほうが情報量が多いのか — は知識蒸留を1からで扱いました。この記事はその次の地図、「で、実際どれを使うのか」です。

4つのレシピを1枚の表で

この表の各行を、式・体感・コードの順に開いていきます。

レシピ1: 出力蒸留 — いちばん少ない材料で作れる

もっとも古典的で、いまも主力なのが出力蒸留です。教師と学生の最終出力の確率分布を近づけるだけで、教師の内部には一切触れません。

言い換えると「正解ラベルに合わせる分と、温度を上げてぼかした教師の分布に合わせる分を、$\alpha$ の割合で混ぜる」だけです。$y$ は正解ラベル、$\mathrm{CE}$ は交差エントロピー、$p_S$ は学生の確率、$p_T^{(T)}$・$p_S^{(T)}$ は温度 $T$ でならした教師・学生の分布、$\mathrm{KL}$ は2つの分布のズレを測る量（KL情報量を1から）です。

温度 $T$ は softmax に入る値を $T$ で割る操作で、$T$ を上げるほど分布は平らになります。平らになると「1位ではないが2位の候補」の情報が浮かび上がり、そこに教師の知識が宿っている、というのが蒸留の出発点でした。式(1)の頭に付いている $T^2$ は補正項です。温度で割った分だけ勾配が $1/T^2$ に縮むので、掛け戻さないと温度を変えるたびに実効的な学習率が変わってしまうのです。

このレシピの強みは要求の少なさです。教師の重みが手に入らなくても、構造がまったく違っていても、語彙さえ共有していれば実行できます。だから「まずこれを試す」が既定の順序になります。

弱みは強みと対称で、出口の1点からしか情報が流れないこと。学生が深いと下のほうの層まで信号が届きにくく、教師が持っている「途中でどんな特徴に分解しているか」は、確率分布に潰れてしまってからでは復元できません。ここから先の3つのレシピは、いずれもこの一点への対処だと思って読めます。

レシピ2: 中間層蒸留 — 途中の下ごしらえを合わせる

出口だけでは足りないなら、途中も合わせよう、というのがFitNetの発想です。教師のある層の出力（hint）と、学生のある層の出力（guided layer）を直接近づけます。

この式が言っているのは要するに、比べる場所を出口から途中へ移した、というだけのことです。教師と学生の棚を1段ずつ選び、そこに載っているものがどれだけ違うかを罰する。出口の確率分布ではなく、答えを作る途中の中間生成物どうしを突き合わせています。