AI論文を手軽にキャッチアップできるpodcastのニュースレターを試験的に配信しています。
こんな感じのpodcastでスキマ時間に聞くだけで論文の概要をキャッチアップできます!
サンプル:
https://open.spotify.com/episode/6OoOKIRmRAOrzk81cHp9hJ?si=V-xY1lp5SqavsXjSaK5jeg
元論文
この記事を読んで興味ある方はこちらのフォームから登録してみてください。
登録はこちらから → https://forms.gle/KbcYAz2AmDeKrGYT7
「蒸留」ってなに?
皆さん「蒸留」という言葉からは、お酒を作る場面やアロマオイルを抽出する場面を思い浮かべるかもしれません。実はAIの世界にも「知識蒸留」というテクニックがあります。これは、簡単に言うと「大きくて賢いモデル(教師モデル)が持っている知識を、小さくて軽量なモデル(生徒モデル)に教え込む」方法です。
先生が自分の知恵を要点だけノートにまとめて、生徒に渡すようなイメージですね。大規模なAIモデルは高性能ですが、サイズが大きく計算コストも莫大です。知識蒸留を使えば、その高性能な教師モデルの知識の“エッセンス”を抽出して生徒モデルに移すことで、性能をなるべく維持しつつモデルを小型化できます。
例えば、ウイスキーの蒸留では大量のもろみから芳醇な原酒を得ますが、知識蒸留では巨大モデルからコンパクトなモデルを得るわけです。
結果として、小さなモデルでもスマホや組み込み機器など限られた計算資源の環境で、高度なAI機能を実現できるようになるのです。
図1: 知識蒸留のイメージ。
左の大きな教師ネットワークが持つ知識を抽出(distillation)し、右の小さな生徒ネットワークに転送します。これにより、生徒モデルは教師モデルに近い精度を保ちつつ、計算資源を大幅に節約できます。
今回の記事では、この「知識蒸留」について解説します。
知識蒸留の基本概念とメカニズム
教師モデルと生徒モデルの関係
知識蒸留の中心にいるのが「教師モデル(Teacher Model)」と「生徒モデル(Student Model)」です。教師モデルは大規模で高性能なモデルで、莫大なデータで訓練されて高い精度を持っています。一方、生徒モデルはそれよりもずっと小型で軽量なモデルです。
知識蒸留では、教師モデルに一度大量のデータを学習させて知識を詰め込み、その後で教師モデルの出力を使って生徒モデルを訓練します。要は、 「大きな先生がテストの解答集を作り、小さな生徒はそれを見て勉強する」 ような構図です。
生徒モデルの訓練時、教師モデルの役割は「目標(ターゲット)を提供すること」です。ただし普通の教師あり学習のように正解ラベルだけを使うのではありません。教師モデルが入力に対して出力する確率分布(もしくはロジット(logits)値そのもの)を、生徒モデルの学習ターゲットとして利用します。
教師モデルの出力する確率分布には、正解ラベル以外のクラスに対する「暗黙の情報(dark knowledge)」が含まれています。
例えば画像分類で教師モデルが「この画像は犬:90%、猫:8%、その他:2%」と予測した場合、単に「犬(正解)」という一言よりも、「猫にもちょっと似ているんだな」といったニュアンスが伝わりますよね。このソフトなターゲットこそが、生徒モデルにとって貴重な学習信号になります。
ソフトターゲットと温度パラメータ
教師モデルの出力する確率分布をそのまま使うのではなく、温度パラメータという仕掛けを使って分布を調整するのが知識蒸留の肝となる工夫です。
ソフトターゲットとは、教師モデルの出力ロジットに対してソフトマックス関数を適用して得られる確率分布のことです。
このとき温度$T$を導入したソフトマックスは次式で定義されます:
$$
q_i^{(T)} = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} \tag{1}
$$
ここで$z_i$は教師モデルの各クラス$i$のロジット値、$T$は温度です。通常の確率分布は$T=1$の場合ですが、$T$を大きくすると出力分布がよりなだらか(ソフト)になります。
直感的には、温度を上げることで教師モデルの「自信満々な予測」を少し和らげ、間違いのクラスにもそれなりの確率を割り振るようにするのです。こうすることで、生徒モデルは「教師が各選択肢をどれくらいの割合で考慮しているか」という微妙なニュアンスまで学習できます。
温度を高くしたソフトターゲットからは、教師モデルが持つ暗黙の知識(例えば「犬と猫は少し似ているから出力確率も猫にそれなりに割り振られている」など)を引き出せます。一方で温度を下げて$T=1$に近づけると分布は尖った形になり、従来のハードな正解ラベルに近づきます。
訓練時には、生徒モデルの出力も同じ温度$T$でソフトマックスを計算し、教師モデルのソフトターゲットとの間でKLダイバージェンス(またはクロスエントロピー)を最小化するよう損失を定義します。また同時に、生徒モデルの出力と正解ラベルとの通常のクロスエントロピー損失も組み合わせて学習させることが多いです。Hintonらの提案した知識蒸留の損失関数は例えば以下のような形になります。
$$
L = (1-\alpha) \cdot H(y_{\text{true}}, p_S) + \alpha \cdot T^2 \cdot KL(p_T^{(T)} \parallel p_S^{(T)})
$$
ここで$H$は交差エントロピー、$KL$はKL距離、$p_T^{(T)}$は教師モデルの温度$T$でのソフトターゲット分布、$p_S^{(T)}$は生徒モデルの温度$T$でのソフトマックス出力です。
【※】$T^2$でスケーリングするのは、温度を上げると損失値が小さくなる効果を打ち消すためのテクニックです(興味があればHintonらの論文で詳しく述べられています)。
パラメータ$\alpha$はハードターゲット(正解ラベル)の損失とソフトターゲットの損失のバランスを調整するハイパーパラメータです。例えば$\alpha=0.1$とすればソフトターゲット損失を全体の10%程度に抑えるイメージです。こうした損失関数を使って生徒モデルを訓練することで、教師モデルの振る舞いをできるだけ真似させようというわけです。
知識蒸留の3つの手法
一口に「知識蒸留」と言っても、どの情報を蒸留(転移)するかによっていくつか種類があります。主に以下の3つの手法が代表的です。それぞれイメージとしては「教師の何を生徒に学ばせるか」の違いです。
ロジット蒸留(応答ベース蒸留)
教師モデルの最終出力(ロジットや確率分布)を生徒モデルに真似させる手法です。知識蒸留の中で最も基本でポピュラーな方法で、Hintonらが提案したオリジナルの蒸留手法もこれに当たります。
教師モデルの予測する確率分布(ソフトターゲット)そのものが“知識”と見なされ、それをそのまま生徒に与えます。実装も比較的シンプルで、タスクを問わず広く使われています。デメリットとしては、教師モデルの中間層に含まれる情報までは伝えられない点があります。
特徴蒸留(特徴マップ蒸留)
教師モデルの中間層の出力、すなわち特徴マップや埋め込み表現を生徒モデルに学習させる手法です。例えばCNN(畳み込みニューラルネット)の場合は畳み込み層が出力する特徴マップ、Transformerの場合は各層の隠れ状態ベクトルなどが該当します。教師の持つ高レベルな特徴表現を生徒にも獲得させることで、最終出力だけでは得られない深い知識まで伝えようという狙いです。
具体的には、教師と生徒の対応する中間層同士の出力の距離(例えばL2ノルム誤差)を損失に加えたり、教師の特徴を圧縮・変換したものを生徒に予測させたりする方法があります。
代表的な研究として、教師の中間層の一部を「ヒント層」としてその出力を生徒が再現するように学習させたFitNet (2015)などが知られています。
また近年のモデルでは単に特徴を真似るだけでなく、Attentionマップ(どの入力部分に注目しているかの分布)を伝達したり、あるいは全層のアクティベーション分布の統計量を一致させる、といった工夫をする手法も提案されています。
特徴蒸留は中間表現を活用できる分、よりリッチな知識を伝えられますが、教師と生徒でネットワーク構造(層の数や次元)が大きく異なると「どの層の特徴を対応付けるか?」という問題が生じ、調整が難しくなることがあります。
アテンション蒸留
こちらは特にTransformer系モデルで注目される蒸留手法で、教師モデルの持つ注意機構(Attention)の情報を生徒モデルに移すものです。ある意味、上の「特徴蒸留」の一種とも言えますが、重要なので分けて説明します。
Transformerでは各層で「どの単語(トークン)に注目して次の出力を計算するか」というAttention行列があります。アテンション蒸留では、教師モデルのAttention行列(もしくはそれから派生するアテンションマップ)を、生徒モデルの対応する層にできるだけ再現させるようにします。
教師が入力のどの部分を重要だと考えているか、その思考パターンをそのまま弟子に叩き込むイメージです。BERTの圧縮では、この注意機構の蒸留が精度向上に寄与することが報告されており、例えばTinyBERTの手法では各層のAttention出力と隠れ状態の両方を蒸留する二段階学習を行うことで、4層のTinyBERTでもBERT-baseの約97%の性能を発揮したとされています。
Attention蒸留は特に自然言語処理系で威力を発揮する手法で、教師と生徒のアーキテクチャが類似している場合には効果が高いです。
以上のように、出力(応答)そのものを真似る・中間特徴を真似る・注意の向け方を真似るという3つの観点で蒸留が行われます。それぞれ単独で使われることもあれば、組み合わせて使われることもあります。実際の研究ではこれらに加えて、複数の教師モデルの知識をまとめて生徒に教える「アンサンブル蒸留」や、生徒自身も他の生徒と教え合う「相互学習」など、バリエーションも豊富です。しかし根底にあるアイデアは共通で、「小さなモデルに大きなモデルの知恵を継承させる」ことに他なりません。
PyTorchで知識蒸留をやってみよう
では、知識蒸留を実際にコードで実践するイメージを掴むために、PyTorchによる簡単な実装例を見てみましょう。ここでは、既に訓練済みの教師モデルとランダム初期化の生徒モデルがあり、教師から生徒への蒸留で分類タスクの精度を上げる場面を考えます。
import torch
import torch.nn.functional as F
# 前提: teacher_model と student_model は定義済み、teacher_modelはすでに学習済み
teacher_model.eval() # 教師モデルは学習済みなので推論モードに
alpha = 0.1 # ソフトターゲット損失の重み
T = 2.0 # 温度パラメータ
for images, labels in train_loader: # 学習データでループ
# 1. 教師モデルと生徒モデルの予測を取得
with torch.no_grad():
teacher_logits = teacher_model(images) # 教師モデルのロジット
student_logits = student_model(images) # 生徒モデルのロジット
# 2. 教師モデルのロジットからソフトターゲット確率分布を計算
teacher_probs = F.softmax(teacher_logits / T, dim=1)
# 3. 生徒モデルのロジットにも温度Tを適用し、ソフトマックス -> distillation損失計算
student_log_probs = F.log_softmax(student_logits / T, dim=1)
distill_loss = F.kl_div(student_log_probs, teacher_probs, reduction='batchmean') * (T * T)
# 4. 生徒モデルの通常のラベルに対するクロスエントロピー損失(ハードターゲット損失)
hard_loss = F.cross_entropy(student_logits, labels)
# 5. 二つの損失を合成
loss = alpha * distill_loss + (1 - alpha) * hard_loss
# 6. 生徒モデルのパラメータを更新
optimizer.zero_grad()
loss.backward()
optimizer.step()
上記のコードでは、教師モデルの出力をソフトターゲット(teacher_probs)として計算し、生徒モデルの出力との間のKLダイバージェンスをdistill_lossとして求めています。また、生徒モデルの通常の分類誤差hard_lossも計算し、両者を重み$\alpha$で足し合わせたものを最終的な損失としています。
ポイントは、教師モデルの推論結果teacher_logitsをそのまま使うのではなく、温度$T$で割ってからソフトマックスすることでソフトターゲットを得ている点です。こうすることで冒頭で説明したように教師の出力分布がなだらかになり、生徒にとって学習しやすい目標になります。
最後にその損失に基づいて生徒モデルのパラメータを勾配降下法で更新していけば、知識蒸留による訓練が進んでいきます。 上の例ではコードを簡潔にするため細部(精度評価や学習率調整など)は省いていますが、基本的な流れはこのようになります。PyTorchの公式Tutorialにも知識蒸留の解説とコード例があり、より詳細な実験手法が紹介されています。興味がある方はぜひ参考にしてください。
知識蒸留はどこで使われている?
知識蒸留は研究分野だけでなく、実際のプロダクト開発や様々な企業のAIモデル最適化で活用されています。ここでは具体的な活用事例をいくつか紹介しましょう。
DeepSeek社(中国)
つい先日、AI業界を驚かせた中国のスタートアップDeepSeekが発表した大規模言語モデル**「DeepSeek R1」は、OpenAIのモデル(o1)に匹敵する性能を示しつつ、開発コスト推定600万ドルという低さでした。この驚異的な効率性を実現した要因の一つがモデルの知識蒸留だと分析されています。
DeepSeek R1では、高性能な大規模モデル(教師)から小型モデル(生徒)への知識蒸留を組み合わせ、チェインオブソート(思考の連鎖)や強化学習といった手法と併用することで、低コストでも高精度なLLMを作り上げたとのことです。
Google(BERTの圧縮)
Googleが生んだ言語モデルBERTは高精度ですが巨大で扱いにくいモデルでした。そこで研究者たちはBERTを小型化するために知識蒸留を活用しています。例えばHugging Face社のチームはBERTを約40%のサイズに圧縮したDistilBERTを公開し、元の97%以上の言語理解能力を保持しつつ推論速度を60%高速化することに成功しました。
またGoogle研究チーム自身もMobileBERTという小型モデルを開発しています。MobileBERTは特殊なボトルネック構造を導入しつつ教師モデルからの知識蒸留で学習することで、BERT_baseより4.3倍小さく、5.5倍高速でありながら、GLUEベンチマークでのスコアがBERT_baseに僅か0.6ポイント差まで迫る性能を達成しました。
つまり、知識蒸留を駆使すれば「あのBERTをスマホ上でサクサク動かす」なんてことも夢ではなくなるわけです。
OpenAI(GPTシリーズでの知識転送)
OpenAIもまた、大規模モデルの知識を活用してより小さなモデルを高性能化する取り組みを行っています。2024年10月にはOpenAIのAPIでモデル蒸留(Model Distillation)機能が発表されました
。これは、たとえばGPT-4等の高性能なフロンティアモデルの出力を大量に集めてデータセット化し、それを使ってより軽量なモデル(例: GPT-4o miniといった小型版)をファインチューニングすることで、高価なモデルに近い性能を安価に実現しようというものです。
OpenAIの発表によれば、このワークフローにより開発者は煩雑だった蒸留プロセス(大規模モデルでのデータ生成からチューニングまで)を一括して行えるようになり、特定タスクにおいて大規模モデル並みの性能を小規模モデルで達成できるといいます。
実際、OpenAIの社内ではGPT-4から小型モデルへの知識蒸留によって推論コストを抑えたモデル開発が進んでいるとも噂されています(例えばGPT-3.5相当の性能を持つ軽量モデルなど)。GPTシリーズの知識転送は、我々ユーザーが安価に高性能AIを使える未来につながるかもしれません。
Meta(LLaMAとオープンソースLLMの軽量化)
Meta(旧Facebook)が公開したLLaMAシリーズは、研究目的で提供された大規模言語モデルで、パラメータ数は小さいながら質の高いデータで訓練することで高性能を実現したモデルです。LLaMA自体は最初から小型高性能を目指したものですが、その登場に刺激を受けたオープンソースコミュニティはさらに知識蒸留を活用したモデル軽量化に挑戦しています。
Stanford大学のAlpacaプロジェクトでは、テキスト指示に応答できるChatGPTの挙動を再現するよう、ChatGPT(GPT-3.5)に大量の質問応答ペアを作らせてLLaMA 7Bモデルに学習させました。これはChatGPTの知識をLLaMAに「蒸留」した一種の実験で、Alpacaは驚くほど流暢に指示に応答できるようになりました。さらに進化したのがVicunaというモデルです。Vicuna-13BはLLaMA 13Bをベースに、ユーザーが共有したChatGPTとの対話データ7万件でファインチューニングされたチャットボットモデルで、GPT-4による評価ではChatGPTやGoogle Bardの90%に匹敵する性能を示したと報告されています。
驚くべきことに、このモデルの学習コストはたった300ドル相当とされています。
つまり、大規模モデル(ChatGPT)の持つ会話スキルをデータという形で蒸留し、小さめのオープンモデルに移植したわけです。Metaが提供したLLaMAという土台と、オープンコミュニティによる知識蒸留の組み合わせが、安価で高性能な対話AIを生み出した好例と言えるでしょう。 以上、世界の最前線で知識蒸留がどのように使われているかを見てきました。巨大モデルをそのまま使うのではなく、小さなモデルに知恵を継承させて効率化する発想が各所で活躍しているのが分かりますね。
最後に
最後に、「で、知識蒸留を学ぶと何が嬉しいの?」という点についてお話しします。 現在のAI業界では、モデルを巨大化すれば性能が上がる一方で、計算コスト・メモリ使用量・エネルギー消費も跳ね上がります。そこでモデル圧縮や最適化のニーズが非常に高まっています。
知識蒸留はそうしたモデル最適化手法の中でも強力かつ汎用性の高い技術です。実際、NLPをはじめとする分野ではモデル圧縮・高速化は活発に研究・産業応用されており、ユーザーに低遅延なサービスを提供するために不可欠となっています。
また、環境負荷やコストの観点からも、効率の良いAIが求められます。知識蒸留を活用すれば性能をなるべく落とさずにモデルを小さく圧縮できるため、こうした要求に応える技術として今後さらに重要になるはずです。大規模モデルありきの時代だからこそ、そのエッセンスを抽出して使い倒す知識蒸留が光るわけです。
まとめると、知識蒸留は「AIモデル界のエコ技術」であり、明日のAIエンジニアにとってほぼ必須とも言えるスキルになりつつあります。巨大モデルの時代に、小さなモデルで勝負する知恵を持つ――それがこれからのスマートなAI活用法と言えるでしょう。ぜひこの機会に知識蒸留の考え方と技術を押さえておいてください。あなたの作るモデルが明日、誰かのスマホやデバイスの中で軽快に動く未来が来るかもしれません。
追記
最後まで読んでくれてありがとうございます!
AI論文を手軽にキャッチアップできるpodcastのニュースレターを試験的に配信しています。
こんな感じのpodcastでスキマ時間に聞くだけで論文の概要をキャッチアップできます!
サンプル:
https://open.spotify.com/episode/6OoOKIRmRAOrzk81cHp9hJ?si=V-xY1lp5SqavsXjSaK5jeg
元論文
この記事を読んで興味ある方はこちらのフォームから登録してみてください。
登録はこちらから → https://forms.gle/KbcYAz2AmDeKrGYT7