知識蒸留を1から — 大きなモデルを小さく写す

正解ラベルには「犬と猫は似ている」と書かれていません。でも訓練済みモデルの出力にはそれが入っている — この差が知識蒸留の全部です。ソフトラベルと温度つきsoftmax、T²の由来、系列レベル蒸留、合成データ蒸留、そしてDeepSeek-R1が推論の手続きごと小さいモデルへ写した話まで、前提知識ゼロから。

Distilling the Knowledge in a Neural Network

「正解」だけでは、先生の知識は伝わらない

犬・猫・トラック・飛行機の4択で画像を分類するモデルを考えます。ある犬の写真に付いている正解ラベルは「犬」、数字にすれば $(1, 0, 0, 0)$ です。データセットが持っている情報はこれで全部です。

ところが、十分に訓練された大きなモデルに同じ写真を見せると、出力はたとえば $(0.90,\ 0.09,\ 0.005,\ 0.005)$ になります。「犬だ。ただし猫の可能性も少しある。トラックや飛行機ではまずない」。ここに書かれている 犬:猫 = 10:1、犬:トラック = 180:1 という比率は、正解ラベルには一切含まれていません。しかしこの比率こそ、そのモデルが大量のデータから学び取った「世界の似ている／似ていない構造」です。

Hinton らはこれを dark knowledge（暗黙の知識） と呼びました。知識蒸留（Knowledge Distillation）とは、大きなモデル（教師）が持つこの暗黙の知識を、小さなモデル（生徒）へ写し取る技術です。

比喩: 答案を写す生徒と、迷い方まで写す生徒

同じ問題集を解く2人の生徒を想像してください。

1人目は解答欄の丸だけを見ます。「問3の答えはウ」。それ以上のことは分かりません。 2人目は先生の採点コメントまで見ます。「答えはウ。ただしイと迷った人が多い。アとエは論外」。

2人目のほうが、同じ問題数からはるかに多くを学べます。「イとウは紛らわしい」という情報が、正解そのものとは別に手に入るからです。これが知識蒸留の直感で、実務上の意味は2つあります。

1つめ: 1サンプルあたりの情報量が増える。 one-hotラベルは「正解のクラス番号」しか伝えませんが、教師の出力は全クラス分の数字を伝えます。だから同じデータ量でも学習が速く進み、ラベルの付いていないデータすら学習に使えます（教師に予測させればソフトラベルが作れるので）。

2つめ: 正則化として効く。 「絶対に犬、他は0」と教え込むより「たぶん犬、猫も少し」と教えるほうが、生徒は極端な自信を持ちにくくなります。ラベルスムージングを、固定値ではなく教師の判断で行っているようなものです。

仕組み: 温度で分布を「ぼかす」

ただし、素直に教師の出力を使うと困ることがあります。よく訓練されたモデルの出力は $(0.99,\ 0.009,\ 0.0009,\ \dots)$ のように尖りがちで、肝心の「猫とトラックのどちらがより犬っぽいか」の差が小さすぎて、勾配にほとんど乗らないのです。

そこで 温度つきsoftmax を使います。ロジット（softmaxに入れる前の生のスコア）を $T$ で割ってから softmax にかけます。

$z_i$ はクラス $i$ のロジット、$T$ は温度と呼ばれる正の数、$p_i(T)$ は温度 $T$ のもとでのクラス $i$ の確率です。$T = 1$ なら普通のsoftmaxで、$T$ を大きくするほど分母の差が縮まり、分布は平らに近づきます。逆に $T$ を0に近づけると、最大値だけが1になる one-hot に潰れます。

つまり温度とは、「教師にどれくらい自信のなさを表明させるか」のつまみです。$T = 4$ くらいに上げると $(0.99, 0.009, \dots)$ が $(0.6, 0.2, 0.13, 0.07)$ 程度にほぐれ、埋もれていた順位関係が学習信号として使える大きさになります。

損失関数: 2つの目標を混ぜる

生徒には2つのことを同時に教えます。「正解ラベルに合わせる」ことと「教師の分布に合わせる」ことです。

第1項は普通の教師あり学習です。$y$ が正解ラベル、$p^S(1)$ が温度1（＝そのまま）での生徒の出力で、$\mathcal{L}_{\mathrm{CE}}$ はクロスエントロピー。第2項が蒸留の項で、$p^T(T)$ と $p^S(T)$ は温度 $T$ での教師と生徒の分布、$\mathrm{KL}$ は2つの分布のズレを測る量です（KLダイバージェンスを1から理解するで詳しく扱っています）。$\alpha$ は両者の配合比で、0.9のように蒸留側を重くすることが多いです。

言葉にすれば、「答え合わせは普通にやりつつ、教師のぼかした分布との距離も同時に縮めろ」という指示です。