【実装】蒸留を自作する — 100行で小さなモデルを育てる

蒸留の損失は20行で書けます。にもかかわらず自作するとほぼ全員が同じ場所で転ぶ — KLの向き、reductionの選び方、T²の付け忘れ。教師の凍結から損失、訓練ループ、教師なしベースラインとの比較、温度スイープ、そして「実装が壊れていないこと」を確かめる4つの検算までを、100行のコードで通します。

Distilling the Knowledge in a Neural Network

ライブラリを呼ばずに書くと、何が見えるか

蒸留の損失関数は短いものです。長くても20行、実質は3行です。だからこそ既存のトレーナに1行足すだけで済ませてしまい、うまくいかなかったときに何も分からなくなります。温度を上げても結果が動かない。生徒の精度がベースラインを下回る。原因が損失の式なのか、教師の扱いなのか、そもそも比較の仕方なのか、切り分ける手がかりがありません。

理屈の側は蒸留の数学で片が付いています。この記事はその続きで、動く実験装置を自分で組む側に回ります。全部で100行ほど。書いてみると分かるのは、蒸留の難しさが損失の数行ではなく、その周りにあることです。

比喩: つまみのある装置は、自分で配線したときだけ信用できる

温度計のついた装置を渡されて「このつまみを回すと味が変わります」と言われたとします。回しても何も変わらない。このとき考えられることが2つあります。つまみが効かない位置にあるのか、それともつまみが中の機構につながっていないのか。

外から見ている限り、この2つは区別できません。配線を自分でやった人だけが、「つまみは確かにつながっている、だから効かないのは設定のせいだ」と言えます。蒸留における温度 $T$ はまさにこのつまみで、後で見るように配線を1か所忘れると、回しても何も起きない装置が簡単に出来上がります。しかも学習自体は普通に進むので、壊れていることに気づけません。

直感: 実装の本体は損失ではなく「3本の走行」

自作の目的は損失関数を持つことではなく、比べられる状態を作ることです。最低3本走らせます。

• A 教師: 大きいモデルを普通に訓練する。これが性能の実質的な上限
• B 素の生徒: 小さいモデルを教師なしで、正解ラベルだけで訓練する。これが下限
• C 蒸留した生徒: 同じ小さいモデルを、教師のソフトラベル込みで訓練する

意味のある数字はただ一つ、CがBとAの間のどこにいるかです。よくある誤りは、Cの精度だけを見て「蒸留が効いた」と言ってしまうことです。生徒アーキテクチャがそもそも十分強くて、教師など無くても同じ精度が出ていたかもしれない。Bを走らせていなければ、その可能性を否定できません。逆にCがAに肉薄していても、AとBの差が最初から小さければ、蒸留は何もしていないのと同じです。

だから実装で最も大事な規律は、BとCで損失以外のすべてを揃えることです。モデル構造、初期化のシード、エポック数、学習率スケジュール、バッチ順序、データ拡張。1つでもずれると、その差が蒸留の効果として計上されます。

仕組み: 損失を1本の式に固定する

まず温度つきsoftmaxで、教師と生徒の分布を作ります。

$z_i$ はクラス $i$ のロジット（softmaxを通す前の生の出力値）、$T$ は正の実数です。式(1)は「ロジットを $T$ で割ってから、いつものsoftmaxに通す」と言っているだけで、$T$ を大きくするほど分布は平らになり、2位以下の順位が読めるようになります。

そのうえで、実装する損失はこれです。

$p$ が教師の分布、$q$ が生徒の分布、$y$ が本物の正解ラベル、$H$ は交差エントロピー、$\alpha$ は0から1の配合比です。式(2)は「教師の分布に合わせる項と、正解に当てる項を $\alpha$ で混ぜる」と読みます。ハード項だけ $q(1)$、つまり温度を1に戻してあるのは、そこが「本当の正解に当てる」担当だからです。頭の $T^2$ は、温度を上げるとソフト項の勾配が $1/T^2$ で縮むぶんを打ち消す補正で、導出は蒸留の数学にあります。

コードにする前に、つまみを実際に回しておいてください。教師の出力がどう平らになるかを見ておくと、後で温度スイープの結果を読むときの目が変わります。

実装1: 教師を用意して、凍らせる

教師は訓練済みモデルをそのまま使います。ここでやることは3つだけです。

teacher.eval()                      # Dropout / BatchNorm を推論モードに固定
for p in teacher.parameters():
    p.requires_grad_(False)         # 勾配を流さない

with torch.no_grad():
    teacher_logits = teacher(x)     # 毎ステップ、生徒と同じ入力で

eval() を忘れると、Dropoutが有効なまま教師が出力を返します。同じ画像を2回見せると違う分布が返ってくる状態で、ターゲットがただのノイズになります。学習は進むので気づきにくく、「蒸留すると精度が下がる」の典型的な原因です。requires_grad_(False) と no_grad は、教師に無駄な計算とメモリを使わないためのものです。

教師の推論が重いなら、ロジットを事前に1回だけ計算してディスクに置く手もあります。毎エポック教師を走らせなくて済むぶん速くなりますが、条件が付きます。データ拡張がランダムなら、キャッシュしたロジットは「別の画像に対する答え」になってしまうので使えません。拡張を固定するか、拡張後の入力ごとにキャッシュするかの二択です。LLMの蒸留では語彙数ぶんのロジットを全トークンについて保存することになるので、容量が現実的でなくなり、上位k個だけを残す運用になります。