自己蒸留は、単一のニューラルネットワークが教師と生徒の両方の役割を担う知識蒸留の一形態である。従来の蒸留が、より大規模または複雑なモデルがより小規模なモデルを導くのとは異なり、自己蒸留は両方の役割に同じアーキテクチャを使用する。モデルはトレーニングデータに対してソフト予測（確率分布）を生成し、これらの予測は真のラベルに加えて追加のトレーニングターゲットとして使用される。このプロセスは反復的、または多世代方式で適用でき、各世代のモデルは前の世代の出力から学習する。この手法は、外部の教師モデルや追加のラベル付きデータを必要とせずに、モデルの精度、キャリブレーション、堅牢性を向上させることが示されている。
この概念は、2010年代半ばに普及した知識蒸留というより広い分野から生まれた。標準的な蒸留では、高容量の教師ネットワークがソフトラベルを生成し、より小規模な生徒ネットワークがそれを模倣することで、クラス間の類似性や不確実性に関する知識を伝達する。自己蒸留は、モデルが自身の知識を蒸留することで、別個の教師の必要性を排除する。これは、より大規模な教師のトレーニングにかかる計算コストを回避し、教師と生徒のセットアップに伴うアーキテクチャ上の制約を排除するため、特に魅力的である。研究者らは、単一のモデルであっても、自身のソフトターゲットでトレーニングすると性能が向上することを発見しており、この現象は理論的および実証的な調査を大いに促している。
メカニズムとトレーニング手順
典型的な自己蒸留のセットアップでは、トレーニングプロセスは段階的に進行する。まず、ニューラルネットワークがハードラベル（ワンホットエンコードされた正解）とクロスエントロピーなどの損失関数を使用して、標準的な分類タスクでトレーニングされる。この初期トレーニングの後、モデルのソフトマックス出力が各トレーニング例について記録される。これらのソフト出力は、確率分布をシャープにしたり平坦にしたりするために温度スケーリングされることが多く、教師の知識として機能する。次の段階では、同じモデルアーキテクチャが再初期化され（または既存の重みが微調整され）、元のハードラベルと前の段階のソフトターゲットの両方を含む複合損失を使用してトレーニングされる。温度パラメータはソフトターゲットの滑らかさを制御し、温度が高いほど、クラス間の関係を明らかにするソフトな分布が生成される。
この手順は複数の世代にわたって繰り返すことができる。例えば、世代0でトレーニングされたモデルは世代1用のソフトラベルを生成し、世代1は世代2用のソフトラベルを生成する、といった具合である。興味深いことに、性能は最初の数世代で向上し、その後横ばいまたはわずかに低下することが多い。この向上は、ソフトターゲットがハードラベルでは提供されないクラス間の類似性に関する情報をエンコードするため、モデルがより正則化された決定境界を学習することに起因する。トレーニング目標は通常、クロスエントロピー損失と蒸留損失を組み合わせ、教師のソフトターゲットの影響のバランスをとるハイパーパラメータによって重み付けされる。
理論的説明
自己蒸留が機能する理由を説明するために、いくつかの理論が提案されている。顕著な説明の1つは、ソフトターゲットの暗黙の正則化効果に関するものである。ハードラベルはモデルにすべての確率質量を正しいクラスに割り当てることを強制し、過信な予測や過学習につながる可能性がある。対照的に、ソフトターゲットはモデルが類似したクラス間で確率を分配することを促し、ラベル平滑化の一種として機能する。これにより、ノイズに対するモデルの感度が低下し、一般化が向上する。
別の研究系統では、自己蒸留を「ダークナレッジ」の概念、つまり正しくないクラス間の相対確率が有用な情報を運ぶという考え方と結び付けている。モデルが自身のソフトターゲットでトレーニングされると、このダークナレッジが効果的に増幅され、より堅牢な特徴表現が得られる。さらに、一部の研究者は自己蒸留をエントロピー正則化の一形態と見なしており、モデルが確信しすぎることを罰し、より滑らかな決定境界を促進する。
最適化の観点からは、自己蒸留はブートストラップの一種と見なすことができる。モデルは自身の予測を反復的に洗練し、各世代が次の世代にわずかに優れたターゲットを提供する。このプロセスは期待値最大化に似ており、モデルはターゲットの生成とそれへの適合を交互に行う。理論的分析により、特定の条件下では、自己蒸留は適切にキャリブレーションされたモデルに対応する不動点に収束し、精度が向上することが示されている。
ディープラーニングへの応用
自己蒸留は、Deep learningのさまざまな領域で応用されている。コンピュータビジョンでは、特にラベル付きデータが限られているシナリオで、画像分類モデルの改善に使用されている。例えば、CIFAR-10やImageNetなどのデータセットでトレーニングされたモデルは、自己蒸留された場合に一貫した精度向上を示している。この手法は半教師あり学習でも効果的であり、モデルがラベルなしデータ用のソフトラベルを生成して、実効的なトレーニングセットを拡張する。
Natural language processingでは、自己蒸留はTransformer (architecture)ベースのモデルやLarge language modelに適用されている。例えば、言語モデルは、自身が生成した応答でトレーニングして、一貫性と事実の正確性を向上させることができる。このアプローチは、モデルがラベルなしコーパスで自身の予測を反復的にラベル付けする自己トレーニングに関連している。自己蒸留は、機械翻訳などのSequence-to-Sequence (Seq2Seq)タスクでも使用されており、モデル自身の翻訳が追加のトレーニングデータとして機能する。
分類と生成以外にも、自己蒸留はModel Pruningと圧縮に適応されている。モデルを縮小された幅または深さでそれ自体に蒸留することにより、研究者は元の性能の大部分を保持するコンパクトなモデルを実現できる。これは、QualcommやArm Holdingsなどの限られた計算リソースを持つエッジデバイスにモデルを展開する場合に特に有用である。
他の手法との関係
自己蒸留は、他のいくつかのトレーニング方法と概念的な類似点を共有している。これは、トレーニング例が意味のある順序で提示されるCurriculum Learningと密接に関連している。自己蒸留では、ソフトターゲットはモデルの理解を徐々に洗練させるため、カリキュラムの一種と見なすことができる。また、ソフトターゲットが効果的に新しいトレーニング信号を生成するため、Data Augmentationとも重複する。一部の研究者は、正則化をさらに強化するために、自己蒸留をDropoutやBatch Normalizationと組み合わせている。
この手法は、モデルが自身または別のモデルのフィードバックを使用して改善するReinforcement Learning from AI Feedback (RLAIF)（AIフィードバックからの強化学習）にも関連している。自己蒸留では、フィードバックはモデル自身の確率分布であり、ソフトな報酬信号として機能する。さらに、自己蒸留は、教師と生徒が同一のアーキテクチャを共有する知識蒸留の特殊なケースと見なすことができ、典型的な異種混合のセットアップとは対照的である。
実証的知見とベンチマーク
実証研究により、自己蒸留は複数のアーキテクチャとデータセットにわたってテスト精度を一貫して向上させることが実証されている。例えば、CIFAR-100では、自己蒸留されたResidual Network (ResNet)は、ハードラベルのみでトレーニングされたベースラインと比較して、1〜2%の精度向上を達成すると報告されている。ImageNetでは、向上は小さいものの依然として有意であり、多くの場合0.5〜1%程度である。この向上は、モデルが過剰パラメータ化されている場合やデータセットが小さい場合により顕著であり、自己蒸留が正則化器として機能することを示唆している。
キャリブレーションも自己蒸留が優れている分野である。自己蒸留でトレーニングされたモデルは、予測確率が真の精度をよりよく反映することを意味する、より低い期待キャリブレーション誤差を持つ傾向がある。これは、医療診断や自動運転など、決定の信頼性が重要なアプリケーションにとって重要である。TeslaやWaymoシステムでは、適切にキャリブレーションされたモデルは過信なエラーのリスクを軽減する。
自己蒸留は、敵対的攻撃に対する堅牢性も向上させることが示されている。ソフトターゲットによって誘発される滑らかな決定境界は、小さな摂動が予測を反転させることを困難にする。これはビジョンモデルと言語モデルの両方で観察されているが、その効果は専用の敵対的トレーニング手法と比較すると控えめである。
バリエーションと拡張
特定の課題に対処するために、自己蒸留のいくつかのバリエーションが提案されている。1つの変種は「ボーンアゲイン・ネットワーク」と呼ばれ、教師とアーキテクチャが同一であるがランダム初期化が異なる生徒モデルのトレーニングを伴う。生徒は教師のソフトターゲットでトレーニングされ、このプロセスは複数の世代にわたって繰り返される。このアプローチは一貫した改善をもたらすことが示されており、後の世代が初期の世代よりも優れた性能を発揮することもある。
もう1つの拡張は「オンライン自己蒸留」であり、モデルは別々の段階ではなく、同じトレーニング実行中に自身の予測を蒸留する。これは、モデルの出力の移動平均を維持するか、共有分類器を使用することで実現される。オンライン手法は複数のトレーニングパスを回避するためより効率的であるが、安定性に欠ける可能性がある。
自己蒸留は、トランスフォーマーのMulti-Head Attentionメカニズムとも組み合わされている。例えば、モデルは自身のアテンションヘッドから単一のヘッドに知識を蒸留し、解釈可能性と性能を向上させることができる。これは、機械翻訳や要約で使用されるEncoder-Decoder Architectureアーキテクチャに特に関連する。
課題と限界
その利点にもかかわらず、自己蒸留は普遍的に効果的というわけではない。場合によっては、特に非常に小規模なモデルやノイズの多いデータセットでは、向上は無視できるか、あるいはマイナスになることもある。この手法はまた、温度や蒸留の重みなどの追加のハイパーパラメータを導入し、調整が必要となる。不適切な設定は、過小適合や過度の平滑化につながり、モデルが保守的になりすぎて識別力を失う可能性がある。
もう1つの限界は、計算オーバーヘッドである。自己蒸留は別個の教師のトレーニングを回避するが、それでも複数のトレーニングパスやソフトターゲットの保存が必要であり、大規模なデータセットではメモリを大量に消費する可能性がある。数十億のパラメータを持つLarge language modelの場合、すべてのトレーニング例に対してソフトターゲットを生成して保存することは非現実的である。研究者は、データのサブセットの使用やソフトターゲットの圧縮などの近似を提案しているが、これらは追加の複雑さを導入する。
自己蒸留の理論的理解は未完成のままである。実証結果は有望であるが、観察されたすべての現象を説明する統一的な枠組みは存在しない。いくつかの研究では、自己蒸留が初期モデルに存在するバイアスを増幅し、最適以下の解につながる可能性があることが示されている。これは、公平性が重要視されるアプリケーションでは特に懸念される。
今後の方向性
進行中の研究では、自己蒸留をより効率的で堅牢にする方法が模索されている。1つの方向性は、トレーニング中のモデルの信頼度に基づいて調整される適応温度スケジュールの開発である。もう1つは、焦点損失やラベル平滑化の変種など、キャリブレーション用に特別に設計されたLoss Functionsとの自己蒸留の統合である。
Generative AIの文脈では、自己蒸留は生成テキストの事実的一貫性を改善する方法として調査されている。モデルが自身の応答を蒸留することにより、幻覚を回避することを学習できる。これは、高リスク環境で展開されるOpenAIやAnthropicのモデルに特に関連する。
自己蒸留はまた、より多様なトレーニング信号を生成するためにData Augmentation技術と組み合わされている。例えば、モデルは同じ入力の拡張バージョンに対してソフトターゲットを生成し、一貫性を強制して不変性を向上させることができる。このアプローチは、半教師ありおよび自己教師あり学習パラダイムで有望を示している。
最後に、自己蒸留の理論的基礎を理解することへの関心が高まっている。研究者は情報理論と最適化のツールを使用して、自己蒸留がいつ、なぜ役立つかを特徴付けている。これは、この手法をいつ適用するか、またそのハイパーパラメータをどのように設定するかについての原則的なガイドラインにつながる可能性がある。
結論
自己蒸留は、モデルが自分自身から学習できるという、シンプルでありながら強力なアイデアを表している。自身の予測をソフトターゲットとして使用することにより、外部の監督なしで、より良い一般化、キャリブレーション、堅牢性を達成する。この手法は、小規模な畳み込みネットワークから大規模なトランスフォーマーまで、さまざまなタスクとアーキテクチャで検証されている。計算コストと理論的理解の点で課題は残っているが、自己蒸留は機械学習ツールボックスにおいて定番であり続ける可能性が高い。モデルがサイズと複雑さを増し続けるにつれて、自己蒸留は追加のデータやアーキテクチャの変更なしに性能を向上させるためのスケーラブルな方法を提供する。