2026-08-29·無料·論文·11分で読めます

蒸留の数学 — なぜ「柔らかい答え」から学べるのか

蒸留の損失はなぜ KL(教師||生徒) なのか、温度Tは何をしているのか、そして実装で必ず出てくる T² 倍は何の補正なのか。ソフトラベルが正解ラベルより多くを語る理由を、式の導出と動く図で1から追います。

#distillation #kl-divergence #softmax #temperature #training #math

対象textタスクtraining

Distilling the Knowledge in a Neural Network

一次資料 — この記事の根拠

この解説の公開 2026-08-29

Distilling the Knowledge in a Neural NetworkarXiv:1503.02531論文ページ·PDF

正解だけを写すか、採点コメントごと写すか

同じ問題集を解いた先輩のノートを借りるとします。借り方は2通りあります。1つは正解の選択肢だけを写すこと。「問3はイ」。もう1つは、先輩が余白に書いた迷いの跡ごと写すことです。「問3はイ。ただしウとはかなり紛らわしい。アとエは論外」。

どちらのノートからよく学べるかは明らかです。正解だけのノートは「イである」としか言っていませんが、迷いの跡のあるノートは「イとウは似ている」「アとエは全く別物」という選択肢どうしの関係まで教えてくれます。同じ1問から受け取る情報量が違うのです。

機械学習でこの2つはそれぞれハードラベルとソフトラベルと呼ばれます。訓練データに付いている正解は前者、訓練済みの大きなモデル（教師）が出力する確率分布は後者です。蒸留（distillation）とは、教師のソフトラベルを小さなモデル（生徒）に写す訓練手続きのことで、この記事はその数学的な中身だけを扱います。手法の全体像や系列レベル蒸留、実際のモデル事例は知識蒸留を1からにあります。

ハードラベルは何を捨てているのか

クラスが $C$ 個ある分類を考えます。ハードラベルは「 $C$ 個のうちどれか1つ」を指すだけなので、1件のデータが運ぶ情報は最大でも $\log_2 C$ ビットです。10クラスなら3.3ビット強、それ以上は増えません。

一方、教師の出力は $C$ 個の確率が並んだベクトルです。合計が1という制約があるので自由度は $C-1$ 個。つまり1件のデータが、1つの答えではなく $C-1$ 個の実数を生徒に指定します。「犬0.7、狼0.25、猫0.03、車0.02」というベクトルは、正解が犬であることに加えて「狼は惜しい」「車は論外」という順位と距離まで伝えている。この余剰分をHintonらは dark knowledge（暗黙知）と呼びました。

大事なのは、これが教師が間違えた分の情報だという点です。教師が完璧に正解だけを1.0で出す（ハードラベルと同じになる）なら、蒸留する意味はありません。教師の確信のなさ、迷いのパターンそのものが教材になっています。

温度つきsoftmax

ところが実際の教師モデルは、訓練が進むほど出力が尖ります。「犬0.999、狼0.0008、猫0.0001…」のようになると、せっかくの順位情報が数値としてほぼ消えてしまう。そこで温度 $T$ を入れて分布をわざと平らにします。

p_i(T) = \frac{\exp(z_i / T)}{\sum_{j=1}^{C} \exp(z_j / T)}

(1)

$z_i$ はクラス $i$ のロジット（softmaxに通す前の生の出力値）、 $T$ は正の実数です。式(1)は「ロジットを $T$ で割ってから、いつものsoftmaxに通す」と言っているだけです。 $T=1$ なら普通のsoftmax。 $T$ を大きくするとロジットの差が縮むので分布は平らになり、 $T$ を0に近づけると差が拡大して、最大値のクラスだけが1になる（＝ハードラベルに戻る）。

つまり温度は「正解だけを見る」と「全クラスの関係を見る」の間の連続的なつまみです。実際に動かしてみてください。

FIG 1温度スライダーを右へ動かすと分布が平らになり、2位以下の順位差が読めるようになる。左端まで戻すと1本だけが立ち、ハードラベルと変わらなくなる

蒸留の損失: なぜ KL なのか

生徒に「教師と同じ分布を出せ」と言いたい。2つの確率分布のズレを測る標準的な量がKLダイバージェンスです。教師の温度つき分布を $p$ 、生徒の温度つき分布を $q$ として、

\mathcal{L}_{\text{soft}} = D_{\mathrm{KL}}(p \,\|\, q) = \sum_{i=1}^{C} p_i \log \frac{p_i}{q_i}

(2)

これは「教師の分布を基準にしたとき、生徒の分布がどれだけ余計か」を測る量です（詳しくはKL情報量を1から）。

ここで一度、対数を分けてみます。

$D_{\mathrm{KL}}(p\|q) = \underbrace{-\sum_i p_i \log q_i}_{\text{交差エントロピー } H(p,q)} \;-\; \underbrace{\left(-\sum_i p_i \log p_i\right)}_{\text{教師のエントロピー } H(p)}$

右辺の第2項は教師だけで決まる定数です。生徒のパラメータをどう動かしても変わりません。したがって $D_{\mathrm{KL}}(p\|q)$ を最小化することと、交差エントロピー $H(p,q)$ を最小化することは、勾配としては完全に同じです。実装がKLで書かれていても、ソフトターゲットつき交差エントロピーで書かれていても、学習結果は変わりません。この一致はエントロピーと交差エントロピーの延長線上にあります。

向きにも意味があります。 $D_{\mathrm{KL}}(p\|q)$ は教師 $p$ が大きいところで $q$ が小さいと激しく罰する形です。つまり「教師が可能性を認めた選択肢を、生徒が0にしてはいけない」という圧力になる。逆向きの $D_{\mathrm{KL}}(q\|p)$ にすると「生徒はどこか1つのモードに逃げ込めば済む」ようになり、まさに捨てたかったハードラベル的な振る舞いに近づきます。蒸留が $D_{\mathrm{KL}}(p\|q)$ を使うのは、この非対称性を利用しているからです。

$T^2$ はどこから来るのか

実装を見ると、ソフト側の損失に必ず $T^2$ が掛かっています。これは経験則ではなく、勾配の大きさから出てきます。

式(2)を生徒のロジット $z_i$ で微分すると、

\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_i} = \frac{1}{T}\left(q_i - p_i\right)

(3)

「生徒の確率と教師の確率の差を、 $T$ で割ったもの」です。ロジットを $T$ で割ってから通したので、連鎖律で $1/T$ が1つ出てきます。

ところが $1/T$ はこれで終わりではありません。 $T$ が大きいとき $q_i - p_i$ 自体も小さくなるからです。 $T$ が大きければ $z_i/T$ は0に近いので $\exp(z_i/T) \approx 1 + z_i/T$ と近似でき、ロジットの平均を0とすれば $q_i \approx \frac{1}{C} + \frac{z_i}{CT}$ 、教師側も同様に $p_i \approx \frac{1}{C} + \frac{v_i}{CT}$ （ $v_i$ は教師のロジット）となります。差を取ると、

$\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_i} \approx \frac{1}{T} \cdot \frac{z_i - v_i}{CT} = \frac{z_i - v_i}{C\,T^2}$

勾配は $1/T^2$ で小さくなる。だから $T^2$ を掛けて打ち消す。これで温度を4から8に変えてもソフト側の効き具合が変わらず、後述の $\alpha$ だけで配合を制御できます。 $T^2$ を掛け忘れると、温度を上げた瞬間にソフト項が実質的に無効化されます。

この近似式にはおまけの意味もあります。分子が $z_i - v_i$ ということは、高温の蒸留は「生徒のロジットを教師のロジットに合わせる二乗誤差」に近づくということです。温度を上げるほど、確率の一致から生ロジットの一致へと目標が滑らかに移ります。

全体の損失と実装

現場ではソフトラベルだけでなく、本物の正解（ハードラベル）も併用します。

\mathcal{L} = \alpha\, T^2\, D_{\mathrm{KL}}\!\left(p(T) \,\|\, q(T)\right) \;+\; (1-\alpha)\, H\!\left(y,\, q(1)\right)

(4)

$\alpha$ は0から1の配合比、 $y$ は正解のワンホットラベル、 $q(1)$ は温度を1に戻した生徒の分布です。ハード項まで温度をかけないのがポイントで、ここは「本当の正解に当てる」担当だからです。

PyTorchならこれだけです。

import torch.nn.functional as F

def kd_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),   # input は必ず log 確率
        F.log_softmax(teacher_logits / T, dim=-1),   # target 側も log で渡す
        reduction="batchmean",                       # 'mean' はクラス数でも割ってしまう
        log_target=True,
    ) * (T * T)                                      # 勾配の 1/T^2 を打ち消す
    hard = F.cross_entropy(student_logits, labels)   # こちらは T=1 のまま
    return alpha * soft + (1 - alpha) * hard

F.kl_div の引数は (input, target) の順で $D_{\mathrm{KL}}(\text{target} \| \text{input})$ を計算します。生徒が input、教師が target。前節で見た向きと逆に書くと、学習は進むのに生徒がモード落ちする（教師が認めた候補を潰す）挙動になります。

現場ではこう使う

誰がいつ触るか: 推論コストを下げたい機械学習エンジニアが、大きな社内モデルを本番用の小型モデルに置き換えるときの標準手順です。分類器の軽量化、大型LLMからの小型LLM作成、自己蒸留（同じ大きさの生徒に自分の出力を教える正則化）など、目的は違っても損失の形はここで書いたものが基本形になります。

触るパラメータ: 実質2つ、温度 $T$ と配合比 $\alpha$ です。 $T$ は教師の出力がどれだけ尖っているかで決めます。学習曲線ではなく、まず教師の出力を数件ダンプして2位以下の確率を目で見るのが速い。上位1つが0.999で埋まっているなら $T$ を上げる余地があり、すでに0.6/0.3/0.1のように分かれているなら $T$ は小さめで足ります。 $\alpha$ は「教師をどれだけ信じるか」で、教師が生徒より明確に強いときはソフト側を重く、教師の精度が怪しいドメインではハード側を重くします。

事故になる落とし穴:

$T^2$ の付け忘れ: 温度を上げるほどソフト項が勝手に弱まり、「 $T$ を上げると効かない」という誤った結論に至ります。ハイパラ探索の結果ごと歪むので厄介です。
reduction の選択: PyTorchの kl_div で 'mean' を使うとクラス数でも割られ、損失がクラス数分の1になります。数学的な定義と一致するのは 'batchmean' です。
推論時に温度を戻し忘れる: $T$ は訓練時の道具です。デプロイした生徒を $T=4$ のまま動かすと、確信度がすべて平らになり、閾値でフィルタする下流の処理が全部壊れます。
教師を eval() と no_grad にしない: DropoutやBatchNormが有効なままだと教師の出力がバッチごとに揺れ、ターゲットがノイズになります。教師へ勾配が流れれば計算とメモリも無駄になります。
語彙・クラス定義の不一致: LLMの蒸留で教師と生徒のトークナイザが違うと、ロジットの $i$ 番目が指すものが両者で別物になります。この場合はロジット同士のKLが定義できないので、教師の生成文を訓練データにする系列レベルの蒸留に切り替える必要があります。

面接で問われる形: 「ソフトラベルが効くのはなぜか」に対しては、1件あたりに与える制約が1個から $C-1$ 個に増え、その余剰がクラス間の類似度構造を伝えるから、と答えれば筋が通ります。続けて「では $T\to 0$ ではどうなるか」と聞かれたら、ハードラベルでの学習に一致し、蒸留の利点が消えるです。

まとめ

蒸留の損失は $D_{\mathrm{KL}}(p\|q)$ 。教師のエントロピーは定数なので、勾配としてはソフトターゲットつき交差エントロピーと同一
温度 $T$ は「正解だけを見る」と「クラス間の関係を見る」を結ぶ連続的なつまみ。 $T\to 0$ でハードラベルに戻る
ソフト側の勾配は $1/T^2$ で縮むので、 $T^2$ を掛けて配合比 $\alpha$ の意味を温度から独立させる
向きは $D_{\mathrm{KL}}(\text{教師} \| \text{生徒})$ 。逆向きにすると生徒がモードに逃げる