Knowledge Distillation

知識蒸留はモデル圧縮技術の一つで、小型の「生徒」ニューラルネットワークを学習させ、より大きく高精度な「教師」モデルの動作を再現します。生徒は正解ラベルに加えて、またはその代わりに、教師の出力確率や内部特徴から学習します。この情報量の多い教師信号により、ラベルだけで学習した同じアーキテクチャよりも高い精度を達成できます。その結果、教師の品質を多く維持しながら、スマートフォン、組み込みボード、その他のエッジ AIハードウェアで使える小型かつ高速なモデルが得られます。Buciluǎ、Caruana、Niculescu-Mizil は 2006 年にアンサンブルの知識を 1 つのモデルに圧縮できることを示しました。2015 年にはHinton、Vinyals、Deanがこの手法を知識蒸留と名付け、多くの手法で現在も使われている、温度で平滑化した教師信号を導入しました。

知識蒸留の仕組み#

ラベルから分かるのは、画像にオオカミが写っているということだけです。学習済み教師のsoftmax出力からは、画像がハスキーにやや似ていて、猫にも少し似ているものの、車とは似ていないことまで分かります。暗黙知と呼ばれることもある、誤ったクラス間の相対確率がソフトターゲットとして生徒に伝わります。softmax の前にロジットを 1 より大きい温度 T で割ると、分布が平坦になり、小さな確率からも十分に学習できるようになります。蒸留損失は、教師と生徒の温度で平滑化した分布間の Kullback-Leibler ダイバージェンスに T² を掛けたもので、通常のラベル損失と勾配の大きさを同程度に保ちます。

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

T=1 では教師はオオカミクラスに 0.875 を割り当て、猫には 0.006、車には 0.0 しか割り当てません。T=4 では同じロジットが 0.493、0.299、0.141、0.067 となり、オオカミは猫よりハスキーに似ており、猫は車よりオオカミに似ていることが生徒に伝わります。生徒はここでは 1.244 となるこの損失を最小化します。通常は、正解ラベルに対する通常のタスク損失との重み付き和として最小化します。教師は学習中ずっと固定され、順伝播のみを実行します。

蒸留の種類#

• 応答ベース: ソフトターゲット方式のように、生徒が教師の最終出力に一致するよう学習します。教師から必要なのはロジットだけです。DistilBERTは事前学習中にこの手法を言語モデリング損失と隠れ状態損失に組み合わせ、言語理解能力の 97% を維持しながら BERT を 40% 小型化し、60% 高速化しました。
• 特徴ベース: 生徒が教師の中間特徴マップに一致するよう学習し、学習可能な小型プロジェクターでチャネル数の差を補います。FitNetsはこの「ヒント」を導入し、教師よりパラメーター数が約 10.4 分の 1 の生徒を学習させ、CIFAR-10 で教師を上回りました。
• シーケンスレベル: 言語モデルの教師がテキストを生成し、生徒はその出力でファインチューニングされます。DeepSeek は、Qwen2.5 と Llama 3 をベースにしたR1 蒸留モデルを、DeepSeek-R1 で選別した 80 万件のサンプルでファインチューニングしました。

特徴ベース蒸留は、1 つのクラス分布ではなく数千個のボックスとスコアを出力する検出器に適しています。Ultralytics YOLO は、学習時にdistill_model教師を指定すると、この方法を使います。固定された大規模 YOLO26 モデルが、検出ヘッドに入力される 3 つのネック特徴マップを提供します。生徒は教師の信頼度で重み付けされ、dis（デフォルトは 6.0）でスケーリングされた L2 損失のもと、1×1 畳み込みプロジェクターを介して特徴マップを一致させます。COCO val2017では、蒸留したYOLO26 nano モデルは 41.5 mAP_50-95を達成します。これは、教師なしで学習した同じモデルの 40.9 を上回り、推論コストは同じです。

from ultralytics import YOLO

# 生徒: デプロイするコンパクトな YOLO26 nano モデル
model = YOLO("yolo26n.pt")

# 教師: 学習をガイドし、その後破棄する大規模な YOLO26 チェックポイント
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

蒸留と関連用語の比較#

• 転移学習: モデルの重みを新しいタスクの初期値として再利用します。蒸留では、別の、通常はより小さいアーキテクチャを学習させ、同じタスクにおけるモデルの動作を模倣させます。
• モデルプルーニング: 学習済みネットワークから重み、チャネル、またはレイヤーを削除します。精度を回復するため、プルーニングしたモデルを元のモデルからの蒸留でファインチューニングすることがよくあります。
• モデル量子化: INT8 などの低精度で重みと活性化を保存します。同じアーキテクチャを小型化する一方、蒸留はより小さいアーキテクチャを選択する手法であり、両者は併用できます。
• モデルアンサンブル: 推論時に複数モデルの出力を平均します。アンサンブルを 1 つの生徒モデルに蒸留することが、この手法が生まれた当初の動機でした。
• データセット蒸留: 名前は似ていますが、モデルではなくデータセットを圧縮します。Wangらは、MNISTの学習画像6万枚をクラスごとに1枚ずつ、合計10枚の合成画像に圧縮しました。これらの画像を使うと、初期化を固定したネットワークの精度を、数回の勾配ステップで元の水準に近づけることができます。