モデルの蒸留（Distillation）——大型モデルから小型モデルを生成する技術

GPT-4のような巨大なモデルは優れた性能を発揮するが、推論コストが高く、スマートフォンや組み込みデバイスでの実行は現実的ではない。この問題を解決するアプローチのひとつが「知識蒸留（Knowledge Distillation）」だ。大型の教師モデル（Teacher Model）の「知識」を小型の生徒モデル（Student Model）に転移させ、軽量ながら高性能なモデルを作り出す技術だ。

知識蒸留の基本アイデア

通常の機械学習では、モデルは正解ラベル（ハードラベル）から直接学習する。たとえば画像分類であれば、「この画像は猫（確率1.0）、犬（確率0.0）、鳥（確率0.0）」というワンホット形式の教師信号を使う。

しかし大型モデルが出力する確率分布（ソフトラベル）はより豊かな情報を含んでいる。「この画像は猫（確率0.80）、豹（確率0.12）、虎（確率0.05）、犬（確率0.02）、…」というように、大型モデルは正解以外のカテゴリへの確率分布によって「猫は豹に似ている」「犬とは違う」という暗黙の知識を表現している。

知識蒸留のアイデアは、このソフトラベルを生徒モデルの学習信号として使うことだ。生徒モデルは単に「正解を当てる」だけでなく、「教師モデルが見ているのと同じように世界を見る」ことを学ぶ。

温度パラメータの役割

知識蒸留において「温度（Temperature）」パラメータは核心的な役割を果たす。

ニューラルネットワークの最終層のSoftmax関数に温度 T を導入する：

softmax(z_i / T) = exp(z_i / T) / Σ_j exp(z_j / T)

T=1（標準的なSoftmax）では確率分布は比較的シャープで、最も確率の高いクラスに確率が集中する傾向がある。

T > 1（高温）にすると確率分布が「なだらか」になる。あまり確率の高くないクラスにも有意な確率が割り当てられ、教師モデルが持つ「クラス間の類似性に関する情報」がより明示的に表れる。

知識蒸留では高い温度（T=2〜5程度）で教師モデルのソフトラベルを生成し、同じ温度で生徒モデルも学習させる。最終的に生成・推論時は T=1 に戻す。

知識蒸留の損失関数

典型的な知識蒸留の損失関数は2つの項から構成される：

L = α × L_hard + (1-α) × L_soft

L_hard（ハード損失）： 正解ラベルを使った通常のクロスエントロピー損失。生徒モデルが基本的な分類タスクを正しく学ぶための項。

L_soft（ソフト損失）： 教師モデルのソフトラベルと生徒モデルの出力のKLダイバージェンス。教師の「知識」を転移させるための項。

α はバランスを制御するハイパーパラメータだ。タスクによって最適な値は異なるが、α=0.5程度から試すことが多い。

層間の知識転移：フィーチャーマッチング

ソフトラベルを使う手法は「応答ベースの知識蒸留」と呼ばれる。より高度なアプローチとして、中間層の特徴量（フィーチャーマップ）自体を転移させる「フィーチャーベースの知識蒸留」がある。

FitNets（2015年）はこのアイデアを実装した初期の手法だ。教師モデルの中間層の表現を「ヒント」として、生徒モデルが同様の中間表現を学ぶよう誘導する。

LLMの分野では、教師モデルの注意スコア（Attention Weights）を生徒モデルに転移させる手法（TinyBERT等）も効果的とされている。注意機構が「どの単語に注目すべきか」を学ぶことで、より意味的に適切な中間表現が得られる。

自己蒸留（Self-Distillation）

通常の知識蒸留は異なるサイズの教師・生徒モデルを必要とするが、同じモデル内または同じモデルの異なるエポックを使う「自己蒸留」という手法もある。

Born-Again Networksでは、同じアーキテクチャのモデルを繰り返し蒸留し、世代を重ねるごとに性能が向上することを示した。第1世代のモデルを教師として第2世代を学習し、第2世代を教師として第3世代を学習する——このプロセスを繰り返すことで、アーキテクチャを変えずに性能が改善できる。

なぜ同じサイズの蒸留で性能が上がるのか？ソフトラベルによる正則化効果と、「より良い初期化された教師からの学習」というアンサンブル効果が組み合わさっているためだと考えられている。

SpeculativeDecoding：蒸留の応用

知識蒸留の概念の応用として注目されているのが、LLMの推論高速化技術「SpeculativeDecoding（投機的デコーディング）」だ。

通常、LLMのデコーディングは1トークンずつ逐次的に行われる。SpeculativeDecodingでは：

1. 小型の「ドラフトモデル」（生徒に相当）が複数のトークンを高速に生成する
2. 大型の「検証モデル」（教師に相当）が生成されたトークン列をまとめて検証する
3. 検証モデルが承認したトークンはそのまま採用、不承認の場合は検証モデルの出力で置き換える

検証は並列処理でき、Transformerの計算効率上、複数トークンをまとめて処理するのとほぼ同コストで済む。ドラフトモデルの精度が高いほど、検証モデルとの差し替えが少なくなり、高速化効果が大きくなる。

実装によっては2〜3倍の推論高速化が達成されており、大型モデルの実用的な高速化手法として注目を集めている。

LLMの蒸留：実践上の課題

大規模言語モデルの知識蒸留には、画像分類と異なる特有の課題がある。

出力空間の広さ： 画像分類は1,000クラスだが、LLMの語彙は数万〜10万以上だ。ソフトラベル全体を正確に計算・保存するのは計算・ストレージ的に困難。

タスクの多様性： LLMは翻訳・要約・コーディング・質問応答など、多種多様なタスクを扱う。特定タスクに特化した蒸留を行うか、汎用的な能力を保ちながら蒸留するかのバランスが難しい。

能力のギャップ： 教師と生徒の能力差が大きすぎると、生徒が教師の出力を模倣しきれない「キャパシティーギャップ」問題が生じる。

知識の選択： 教師モデルのどの「知識」を転移させるかの選択が重要で、タスクに無関係な知識を転移しても効果は薄い。

エッジAIへの貢献

知識蒸留は、スマートフォン・IoTデバイス・車載システムなどのエッジデバイスでAIを動かす「エッジAI」の実現に大きく貢献している。

クラウド上の大型モデルを蒸留して小型化することで、インターネット接続なしにデバイス上でリアルタイム処理ができるようになる。プライバシーが重要な音声認識・顔認識などの用途で、データをクラウドに送らずにエッジで処理できるメリットは大きい。

量子化（INT8/INT4）とプルーニング（不要なパラメータの削除）を知識蒸留と組み合わせることで、さらなる小型化が可能だ。

まとめ

知識蒸留は「ソフトラベルの活用」という単純なアイデアから始まり、大型モデルが持つ暗黙の知識を小型モデルに転移する強力な技術に発展した。温度パラメータによるソフトラベルの調整、フィーチャーマッチングによる中間表現の転移、SpeculativeDecodingへの応用など、多様な発展形が生まれている。

LLMの民主化という観点からも、知識蒸留は重要な役割を担う。巨大な計算資源なしに、高性能なLLMの能力を小型モデルに詰め込む技術は、AIの恩恵をより多くのユースケースとデバイスに届けるための鍵となっている。