書記が数学やるだけ#896 knowledge distillationとDistilBERT

知識蒸留（knowledge distillation）について。

問題

説明

知識蒸留（knowledge distillation）は、教師モデルに蓄積された学習済み知識を生徒モデルへ転移し、精度をほぼ維持しながら計算コストやモデルサイズを削減する手法である。教師モデルは大規模で高性能なニューラルネットワークであり、生徒モデルは小規模で計算効率に優れたモデルである。

これをBERTに適用したのがDistilBERTであり、レイヤー数をBERTの半分にすることでパラメータ数を約40％削減しつつ、知識蒸留によって精度低下を最小限に抑えている。

解答

knowledge distillationでは、教師モデルと生徒モデルそれぞれのロジットに温度付きソフトマックスを適用し、その確率分布に基づくクロスエントロピー損失を用いる。

この損失の勾配は、高温極限においてロジットの差に比例し、温度の二乗およびサンプル数に反比例することが示される。

import numpy as np
import matplotlib.pyplot as plt

# 教師モデルのロジット例
logits = np.array([3.0, 1.0, 0.0])

# 温度のリスト
temperatures = [0.5, 1.0, 2.0, 5.0]

def softmax(z, T=1.0):
    """温度付きソフトマックス"""
    z = z / T
    exp_z = np.exp(z - np.max(z))  # 安定化のため最大値を引く
    return exp_z / exp_z.sum()

# グラフ描画
plt.figure(figsize=(8,5))

for T in temperatures:
    probs = softmax(logits, T)
    plt.plot(range(len(logits)), probs, marker='o', label=f'T={T}')

plt.xticks(range(len(logits)), [f'Class {i}' for i in range(len(logits))])
plt.ylim(0,1)
plt.xlabel('Class')
plt.ylabel('Probability')
plt.title('Softmax Output with Different Temperatures')
plt.legend()
plt.grid(True)
plt.show()

・温度 T を上げるとソフトマックス出力は平滑化され，クラス間の確率差が小さくなる。

勾配は(q −p)/T に比例するので，確率差 q−p が小さくなり、さらに 1/T の係数があるため，勾配の大きさは温度を上げるほど小さくなる。

・高温極限で見ると勾配はロジット差 (zi −vi ) に比例するが係数が 1/(NT^2) となるため，温度が大きいほどロジット差に対する感度は T^(−2)で小さくなる。

・高温にすると全体が平滑化されるが，相対的に非常に負のロジット由来の差分はさらに影響が小さくなることから，ノイズが多い・確信度の低いクラス（大きく負のログit）からの不安定な信号を抑制できる

・T を上げると蒸留信号は滑らかになり安定するが，学習信号自体は弱くなるため，蒸留ではこのトレードオフを温度と損失重みで調整する必要がある。

生徒モデルの3つの損失について。

(1) Distillation loss：教師モデルと生徒モデルの出力分布を温度付きソフトマックスで比較し、クロスエントロピーとして計算する。教師モデルが持つ知識（クラス間の相対的な確率分布）を生徒に伝達することで、正解ラベルだけでなく他のクラスに対する教師の「暗黙の判断」（ソフトターゲット）も学習できる。

(2) Masked language model (MLM) loss：入力の一部を [MASK] トークンに置き換え、その位置の単語を正解ラベルと照合して予測するクロスエントロピー損失。これは教師がいなくても独立に学習可能な自己教師ありタスクであり、生徒が言語モデリング能力（単語予測の基礎能力）を維持することで、教師への依存を補完し、教師の誤差や限界に引きずられないようにする。

(3) Cosine embedding loss：教師モデルと生徒モデルの隠れ状態ベクトル（中間表現）の方向を一致させるために、コサイン類似度を最大化するように設計された損失。これにより出力だけでなく内部表現（特徴空間の構造）も教師に近づけることができ、生徒が浅いモデルであっても教師の意味空間を模倣できる。単に最終出力を合わせるだけでなく、中間的な言語理解の表現をも継承することが可能になる。

本記事のもくじはこちら：