知識蒸留（knowledge distillation）は、モデル蒸留とも呼ばれ、Machine learningにおける技術であり、大規模で複雑なモデル（教師モデル）から、より小さく単純なモデル（生徒モデル）へ知識を転移させる手法である。その目的は、大規模モデルの予測精度を維持しつつ、計算効率の高い評価が可能なコンパクトなモデルを作り出すことにある。これは、大規模モデルの実行コストが高くつくモバイルデバイスなどのリソース制約のあるハードウェアにモデルを展開する際に特に有用である。このプロセスは、大規模モデルの出力、特にそのソフトな確率が、最終的なクラスラベルだけよりも豊かな情報を含んでおり、この情報を用いて小さなモデルを効果的に訓練できるという考え方を利用している。
知識蒸留は、既存のモデル自体のサイズを縮小する手法（例えば、パラメータあたりのビット数を減らすなど）を指すモデル圧縮とは区別される。モデル圧縮は通常、アーキテクチャとパラメータ数を維持するが、蒸留は新しい小さなモデルを訓練することを伴う。この技術は、物体検出、音響モデル、自然言語処理、そして最近では非グリッドデータを扱うグラフニューラルネットワークなど、さまざまな分野で成功裏に応用されている。
歴史的背景
知識蒸留の概念は、Deep learningやNeural network研究の広範な分野にルーツを持つ。モデル間で知識を転移するという考え方はさまざまな形で探求されてきたが、現代的な定式化は、2015年にジェフリー・ヒントン、オリオル・ヴィニャルス、ジェフ・ディーンによって発表された論文「Distilling the Knowledge in a Neural Network」に帰せられることが多い。この研究は、ソフトマックス関数における温度の使用を形式化し、教師モデルの出力を軟化させることで、より効果的な知識転移を可能にした。それ以来、知識蒸留は機械学習ツールボックスにおける標準的な技術となり、多くの変種と応用が生まれている。
基本原理
知識蒸留の背後にある基本原理は、非常に深いNeural networkやモデルのアンサンブルなどの大規模モデルは、小さなモデルよりも高い知識容量を持つというものである。しかし、この容量は完全には活用されない可能性があり、また、大規模モデルの評価は、その容量のどれだけが使用されるかに関係なく計算コストが高い。蒸留は、大規模モデルに符号化された知識を、有効性を失うことなく小さなモデルに転移することを目的とする。小さなモデルは評価コストが低いため、性能の低いハードウェアに展開できる。
重要な洞察は、訓練済みモデルのソフトな出力、すなわちさまざまなクラスに確率を割り当てる出力が、モデルの内部表現に関する情報を含んでいるという点である。例えば、モデルが猫の画像を正しく分類するとき、'猫'に高い確率を割り当てるかもしれないが、'犬'や'狐'にも小さな非ゼロの確率を割り当てるかもしれない。これらの小さな確率は、クラス間の微妙な類似性や関係を符号化しており、ハードラベル（最も可能性の高いクラス）だけを使用すると失われてしまう。知識蒸留は、このソフトな情報を活用して、生徒モデルにより豊かな表現を教える。
数学的定式化
典型的な分類タスクでは、ニューラルネットワークの最終層はソフトマックス関数を使用してロジット（生のスコア）を確率に変換する。標準的なソフトマックスは次のように与えられる：
\( y_i(x) = \frac{e^{z_i(x)}}{\sum_j e^{z_j(x)}} \)
ここで、\( z_i(x) \) は入力 \( x \) に対するクラス \( i \) のロジットである。知識蒸留では、温度パラメータ \( t \) が導入され、ソフトマックスは次のように修正される：
\( y_i(x|t) = \frac{e^{z_i(x)/t}}{\sum_j e^{z_j(x)/t}} \)
より高い温度 \( t \) はよりソフトな確率分布を生成し、確率がクラス間でより均等に広がることを意味する。このソフトな分布は、クラス間の関係についてのより多くの情報を明らかにする。なぜなら、二次的なクラスに対するモデルの信頼度がより顕著になるからである。
蒸留中、生徒モデルは転移セット（元の訓練データまたは新しい、場合によってはラベルなしのデータ）で訓練される。損失関数は通常、生徒の出力と教師の出力との間の交差エントロピーであり、両方とも高温で計算される。単一の入力 \( x \) に対する損失は次のようになる：
\( E(x|t) = -\sum_i \hat{y}_i(x|t) \log y_i(x|t) \)
ここで、\( \hat{y}_i(x|t) \) は教師の出力、\( y_i(x|t) \) は生徒の出力である。高温を使用すると出力のエントロピーが増加し、生徒が学習するためのより多くの情報を提供し、異なるレコード間の勾配の分散を減らすため、より高い学習率が可能になる。
訓練プロセス
知識蒸留の訓練プロセスは、通常、いくつかのステップを含む。まず、大規模な教師モデルが標準的な手法を用いて元のデータセットで訓練される。このモデルは、単一の大規模ネットワークまたはモデルのアンサンブルであり得る。教師が訓練されたら、そのソフトな出力が転移セットに対して生成され、多くの場合高温で行われる。次に、生徒モデルは、上記の交差エントロピー損失を使用して、これらのソフトな出力を模倣するように訓練される。
場合によっては、損失は、利用可能であれば、生徒の出力とグラウンドトゥルースラベルとの間の交差エントロピーで補強される。この組み合わせにより、生徒は教師の知識と実際のターゲットラベルの両方を学習できる。温度は通常、訓練中にアニーリングされ、高温から始まり徐々に1に減少し、ソフトなターゲットの学習からハードラベルによる微調整へと移行する。
応用
知識蒸留は、Artificial intelligenceの多くの分野で成功裏に応用されている。コンピュータビジョンでは、物体検出や画像分類に使用され、エッジデバイスでのモデル実行を可能にする。音声認識では、音響モデルが蒸留の恩恵を受け、遅延とメモリ使用量を削減する。Natural language processingでは、蒸留はTransformer (architecture)ベースの大規模モデル、例えばLarge language modelを、モバイルデバイスやリアルタイムアプリケーションに展開できる小さなバージョンに圧縮するために使用される。例えば、Generative AIシステムのような大規模モデルは、特定のタスク用に小さなモデルに蒸留され、元の性能の多くを維持できる。
最近では、知識蒸留は、ソーシャルネットワークや分子構造などの非グリッドデータを扱うグラフニューラルネットワークにも拡張されている。この拡張は、さまざまなデータタイプとモデルアーキテクチャにわたるこの技術の汎用性を示している。
変種と拡張
特定の課題に対処するために、知識蒸留のいくつかの変種が開発されている。注目すべき変種の1つは逆知識蒸留であり、知識が小さなモデルから大きなモデルへ転移される。これはあまり一般的ではないが、小さなモデルが特定のデータで訓練されており、そのデータに対する大きなモデルの性能を向上させることが目標であるシナリオで有用であり得る。
他の拡張には、注意ベースの蒸留（生徒が教師の注意マップを模倣することを学習する）や、特徴ベースの蒸留（中間表現がターゲットとして使用される）が含まれる。これらの方法は、教師からより詳細な情報を捕捉することにより、知識転移の忠実度を向上させることを目的とする。
利点と限界
知識蒸留の主な利点は、精度を大きく失うことなく、リソース制約のあるデバイスに高性能モデルを展開できることである。これは、モバイルアプリ、組み込みシステム、リアルタイム推論などのアプリケーションにとって重要である。蒸留はまた、より高速な推論時間とより低いエネルギー消費につながる可能性がある。
しかし、限界もある。訓練プロセスには、十分に訓練された教師が必要であり、これは計算コストが高い場合がある。さらに、生徒モデルは、特に容量ギャップが大きすぎる場合、教師と同じレベルの性能を常に達成できるとは限らない。温度と転移セットの選択も慎重な調整を必要とする。
他の技術との関係
知識蒸留は、他のモデル最適化技術と組み合わせて使用されることが多い。例えば、Model Pruningは、蒸留後に生徒モデルに適用して、そのサイズをさらに縮小できる。Data Augmentationは、転移セットを拡張して生徒の一般化を向上させるために使用できる。蒸留はまた、ソフトなターゲットが漸進的学習の一形態と見なせるため、Curriculum Learningとも関連している。
Deep learningの文脈では、蒸留はTransfer learningの一形態であり、あるモデルから別のモデルへ知識が転移される。また、Reinforcement Learning from AI Feedback (RLAIF)（AIフィードバックからの強化学習）とも密接に関連しているが、後者はモデルを人間の好みに合わせることに焦点を当てている。
将来の方向性
モデルが成長を続けるにつれて、特にLarge language modelの分野では、知識蒸留はさらに重要になる可能性が高い。研究者たちは、OpenAI、Anthropic、Google DeepMindによって開発されたような巨大なモデルから、より小さく効率的なモデルへ知識を蒸留する方法を探求している。これは、高度なAI機能へのアクセスを民主化するために重要であり、小規模な組織や個人開発者が、広範な計算リソースを必要とせずに強力なモデルを使用できるようにする。
もう1つの研究分野はオンライン蒸留であり、教師と生徒が同時に訓練され、また、複数のモデルが互いに学習する協調蒸留である。これらのアプローチは、蒸留プロセスの効率と有効性を向上させることを目的とする。
要約すると、知識蒸留は、モデル圧縮と知識転移のための強力な技術である。大規模モデルのソフトな出力を活用することにより、正確かつ効率的なコンパクトモデルの作成を可能にし、現代の機械学習展開戦略の基盤となっている。