AI・機械学習
上級

知識蒸留(Knowledge Distillation)(チシキジョウリュウ)

知識蒸留とは、Geoffrey Hintonらが2015年に提案した手法で、教師モデルの出力確率分布(soft label)を温度パラメータで制御しながら生徒モデルに転移する技術であり、LLMの軽量化の理論的基盤となっている。

0 回閲覧
0 いいね

知識蒸留(Knowledge Distillation)は、2015年にGeoffrey Hinton、Oriol Vinyals、Jeff Deanが論文「Distilling the Knowledge in a Neural Network」で提案した技術である。教師モデルのsoftmax出力を温度パラメータTで平滑化し、通常のhard label(one-hot)では得られない「クラス間の類似性情報」を生徒モデルに転移することが核心的なアイデアである。

理論的基盤

知識蒸留の損失関数は以下の2項の加重和で構成される:

  • 蒸留損失(Distillation Loss): 教師と生徒のsoft label間のKL-divergence。温度T=1〜20で制御。Tが大きいほど確率分布が平滑化され、低確率クラスの情報も転移される
  • タスク損失(Task Loss): 生徒の出力とground truthのcross-entropy。蒸留損失とのバランスはα(0.1〜0.9)で調整

最終損失 = α × 蒸留損失 × T² + (1-α) × タスク損失

T²の項は、温度Tで平滑化するとgradientの大きさが1/T²倍になるための補正である。実験的にはT=4〜8、α=0.5〜0.7が多くのタスクで最適とされる。

Soft LabelとHard Labelの違い

特性Hard LabelSoft Label (T=1)Soft Label (T=4)
表現形式one-hot [0,0,1,0][0.01,0.04,0.90,0.05][0.15,0.22,0.38,0.25]
情報量最小(正解のみ)中(確率分布)最大(クラス間関係)
学習効率低い中程度高い
ノイズ耐性高い中程度低い(平滑化しすぎると劣化)

LLMにおける知識蒸留の拡張

LLMでは元論文の分類タスク向け手法を以下のように拡張している:

  • シーケンスレベル蒸留: 各トークン位置での教師の出力分布を模倣。Kim & Rush (2016)が提案。計算コストが高いが精度は最良
  • Word-Level KD: 各位置のtop-K logitのみを蒸留に使用。メモリ効率が良く、70B→7Bの蒸留でも実用的
  • Attention Transfer: 教師のAttention Mapを生徒が模倣。Zagoruykoら(2017)が提案。レイヤー数が異なるモデル間でも適用可能
  • Contrastive Distillation: 正例(教師の出力)と負例(ランダム出力)のコントラスト学習。2024年にSnowflakeが提案し、Embedding蒸留で効果を実証

温度パラメータの影響

温度Tは蒸留品質に決定的な影響を与える:

  • T=1: 通常のsoftmax。教師の自信度をそのまま反映。dominant classに偏る
  • T=2〜4: 軽度の平滑化。top-5クラスの関係が明確になる。小規模タスクに適する
  • T=4〜8: 中度の平滑化。LLMのトークン蒸留で最も汎用的な範囲
  • T=8〜20: 強度の平滑化。ほぼ一様分布に近づく。語彙サイズ32,000〜128,000のLLMでは情報が薄まりすぎるリスク

代表的な知識蒸留フレームワーク

フレームワーク開発元言語LLM対応特徴
TextBrewerHFLPythonGPT-2/BERT系教育用途に最適、日本語対応
DistilBERTHugging FacePythonBERT系事前学習済み蒸留モデル提供
TinyLlamaZhang et al.PythonLlama系1.1Bパラメータ、3T tokenで学習
MiniLLMMicrosoftPythonGPT/OPT系逆KL-divergenceで安定学習
GKDGooglePythonGemma/T5系Generalized KD、on-policy蒸留

蒸留の理論的限界

知識蒸留にはいくつかの理論的限界が知られている:

  • Capacity Gap: 教師と生徒のパラメータ数の差が大きすぎると(100倍以上)、蒸留効率が急落する。Teacher Assistant(中間サイズモデル)を挟む段階的蒸留で緩和可能
  • Dark Knowledge の消失: 語彙サイズが巨大(128K以上)な場合、温度を上げてもlong-tail tokenの情報は事実上ゼロに近い
  • 分布シフト: 蒸留時の入力分布と推論時の分布が異なると性能劣化。On-policy蒸留(生徒自身の出力を入力に使用)で緩和

よくある質問(FAQ)

Q1: 知識蒸留と転移学習の違いは何ですか? A: 転移学習は事前学習済みモデルの重みを新タスクに再利用する技術で、モデル構造は変わらない。知識蒸留は教師の「振る舞い」(出力分布)を異なるアーキテクチャの生徒に転移する技術で、モデルサイズの大幅な削減が可能な点が根本的に異なる。

Q2: 温度パラメータはどう設定すべきですか? A: LLMの蒸留ではT=4〜8が汎用的な出発点。タスク特化ならT=2〜4で十分な場合が多い。語彙サイズが大きい(64K以上)モデルではT=4程度に抑えないとlong-tail tokenの情報が消失する。ハイパーパラメータ探索で最適値を見つけることを推奨する。

Q3: オープンソースモデルの蒸留にライセンス上の制約はありますか? A: Llama 3系はMeta Llama Licenseで商用利用可能だが、月間7億MAU超の場合は別途ライセンスが必要。Mistralの蒸留はApache 2.0で制約なし。GPT-4等のAPIモデルからの蒸留は、OpenAIの利用規約で競合モデルの学習への使用が禁止されているため注意が必要。

まとめ

  • 知識蒸留はHinton(2015)が提案した教師→生徒の知識転移技術
  • 温度パラメータTでsoft labelの平滑度を制御し、クラス間の関係を転移
  • LLMではシーケンスレベル蒸留・Attention Transfer等に拡張
  • Capacity Gapを避けるため、教師と生徒のサイズ比は10倍以内が推奨