知識蒸留とは

知識蒸留の目標は、より小さなモデルをトレーニングして、より大規模かつ複雑なモデルを模倣するようにすることです。従来のディープラーニングの目的は、人工ニューラル・ネットワークをトレーニングして、その予測をトレーニング・データ・セットで提供されるアウトプット例に近づけることですが、知識蒸留における主な目的は、生徒ネットワークをトレーニングして、教師ネットワークによる予測と一致させることです。

知識蒸留（KD）は、多くの層と学習可能なパラメータを持つ大規模なディープ・ニューラル・ネットワークに最もよく使用されています。このプロセスは、数十億個のパラメーターを持つ大規模な生成AIモデルの継続的な増加に特に関連しています。

この概念は、2006年に発表された「Model Compression（邦訳：モデル圧縮）」という論文に端を発しています。Caruanaらは、当時最先端の分類モデルであった、数百個の基本レベル分類ツールで構成される巨大なアンサンブル・モデルを使用して大規模なデータ・セットにラベル付けをし、その後、従来の教師あり学習を通じて、その新しくラベル付けされたデータ・セットで単一のニューラル・ネットワークをトレーニングしました。このコンパクトなモデルは、「1000分の1の小ささで、さらに高速」で、大規模なアンサンブル・モデルの結果と一致しました。¹

知識蒸留技術は、それ以来、自然言語処理（NLP）、音声認識、画像認識、物体検出など、さまざまな分野で効果的に活用されてきました。近年、知識蒸留の研究は、大規模言語モデル（LLM）にとって特に重要となっています。LLMでは、KDは、よく使用されている有料モデルから、より小規模でアクセスしやすいオープンソース・モデルに高度機能を移行するのに効果的な手段とみなされています。

多くの現実世界の状況では、人工知能モデルの精度と容量だけでは、モデルを有用なものにするのに十分ではありません。利用可能な時間、メモリー、資金、計算技術も予算内に収める必要があります。

特定のタスクに対して最高のパフォーマンスを発揮するモデルは、多くの場合、実際の使用ケースでは大きすぎたり、遅すぎたり、高価すぎたりします。一方、これらのモデルは多くの場合、その規模と膨大な量のトレーニング・データで事前トレーニングされていることにより、他にはない品質を備えています。これらの新たな能力は、GPTやLlamaなどの自己回帰言語モデルで特に顕著であり、シーケンス内の次の単語を単純に予測するという明示的なトレーニング目標を超えた能力を発揮します。逆に、小さなモデルは高速で計算負荷も少ないですが、はるかに多くのパラメーターを持つ大きなモデルほどの精度、洗練度、知識容量がありません。

2015年に発表された画期的な論文「Distilling the Knowledge in a Neural Network（邦訳：ニューラル・ネットワークにおける知識の蒸留）」で、Hintonらは、トレーニングを目的の異なる2つの段階に分割することでこれらの制限を回避できることを発見しました。著者らは、次のような類推を示しました。多くの昆虫には、環境からエネルギーと栄養素を抽出するために最適化された幼虫の形態と、移動と繁殖のために最適化されたまったく異なる成虫の形態がありますが、従来のディープラーニングでは、要件が異なるにもかかわらず、トレーニング段階と展開段階の両方で同じモデルが使用されます。

Hintonらは、自然とCaruanaらの研究の両方からインスピレーションを得て、大規模で扱いにくいモデルのトレーニングがデータから構造を抽出する最善の方法である場合は、こうしたトレーニングに価値があると示唆する一方で、その知識をリアルタイム・デプロイメントに適した小さなモデルに転送するための別の種類のトレーニングである蒸留を提案しました。²

知識蒸留技術は、教師モデルのアウトプットを複製するだけでなく、その「思考プロセス」をエミュレートすることを目的としています。LLMが活躍する昨今、KDにより、スタイル、推論能力、人間の好みや価値観への適合などの抽象的な性質を伝達できるようになりました。³

さらに、小さいモデルは基本的に説明しやすいこともメリットと言えます。なぜなら、数千億個ものパラメーターを持つモデルでは、ニューラル・ネットワークのさまざまな部分への寄与を解釈することが困難だからです。大規模な「ブラックボックス」モデルによって学習された表現をより単純なモデルに転送することは、医療診断や分子発見などの分野で変革的な洞察を明らかにするのに役立ちます。⁴

知識蒸留（KD）は、特定のニューラル・ネットワーク・アーキテクチャーに依存せず、教師ネットワークと生徒ネットワークが同じアーキテクチャーを持つ必要もありません。つまり、あらゆるディープラーニング・モデルに適用できます。

KDは、人工ニューラル・ネットワークが「普遍的な近似器」であるという事実を利用しています。つまり、十分なトレーニング・データと十分に大きな隠れ層があれば、ニューラル・ネットワークは任意の関数を任意の精度で近似できるということです。⁵

従来の機械学習では、トレーニングされたモデルの「知識」は、学習されたパラメーター、つまりニューラル・ネットワーク全体で発生するさまざまな数学的演算に適用される可変の重み（およびバイアス）によって識別され、ネットワークのアウトプットの特定の部分が他の部分に与える影響を増幅または低減します。このように知識を見ると、あるモデルがどのようにして規模や構造が異なる別のモデルの知識を吸収できるのかがわかりにくくなります。

そこで、Hintonらは、知識を単に「インプット・ベクトルからアウトプット・ベクトルへの学習されたマッピング」と扱うことにより、より抽象的かつ柔軟な見方を可能にしました。つまりこの方法では、KDはモデルの知識を、トレーニングで学習した厳密な数学的パラメーターとして解釈するのではなく、トレーニング後にモデルが新しいデータにどのように一般化されるかとして解釈します。

知識の蒸留法は、知識のこの代替的な理解を通じて、特定のインプットに対する教師モデルの最終的なアウトプットを模倣するだけでなく、教師モデルがその最終的なアウトプットに到達するために実行する推論手順も模倣するように生徒モデルをトレーニングすることを目指します。概念的には、これは思考の連鎖（CoT）プロンプトによる命令チューニングと同様に機能し、ステップ・バイ・ステップの論理的根拠を明確に表現することを教えることで、LLM応答の品質を向上させます。

従来の教師あり学習または自己教師あり学習では、損失関数は、さまざまなインプットにわたるモデルのアウトプットと「正しい」アウトプット（または真実）間の相違（または損失）を表すベクトルを生成します。勾配降下法などの最適化アルゴリズムを使用して、このベクトルの傾き（または勾配）を最小限に抑えるようにモデル・パラメーターを調整することで、モデルのアウトプットは正しいアウトプットに近づきます。モデルの推論ステップは最終的なアウトプットに影響を与えるという点で「重要」ですが、通常は従来の損失関数では測定されません。

逆に、知識蒸留では、最適化のソフト・ターゲットとして離散推論ステップを使用する特殊なタイプの損失関数である蒸留損失の追加を通じて、教師モデルの推論プロセスを模倣するように学生モデルをトレーニングします。

あらゆるAIモデルのアウトプットは、予測として理解できます。自己回帰LLMは、指定されたシーケンス内の次の単語を予測します。画像分類に使用されるコンピューター・ビジョン・モデルは、特定の画像のカテゴリーを予測します。この文脈では「ハード・ターゲット」と呼ばれるこれらの最終予測に到達するために、ディープラーニング・モデルは通常、複数の予備予測を行い、ソフトマックス関数を使用して、最も高い確率の予測をアウトプットします。トレーニング中は、クロスエントロピー損失関数を使用して、正しいアウトプットに割り当てられる確率を最大化し、誤ったアウトプットに割り当てられる確率を最小化します。

例えば、画像分類モデルは、モデルが認識するようにトレーニングされた既知の各クラスにインプット画像が属する確率を予測し、最も高い確率値を持つクラスをアウトプットします。機械学習の数学的用語では、これらの個別のクラスごとの予測はロジットと呼ばれます。同様に、自己回帰LLMは、次の単語ごとに複数の可能性を予測し、（その温度設定に応じて）それらの可能性の1つをアウトプット用にサンプリングします。

知識蒸留では、教師モデルによって生成されたこれらの中間予測（「ソフト・ターゲット」）が、多くの場合、生徒モデルの主要なトレーニング・データを提供します。これらの予備予測に割り当てられた相対確率は、教師モデルがどのように一般化される傾向があるかについての貴重な洞察を提供します。例えば、画像分類モデルでは、キツネの画像を「サンドイッチ」と誤分類するよりも「犬」と誤分類する可能性が何倍も高くなります。したがって、ソフト・ターゲットは、ハード・ターゲットのみの場合よりも、トレーニング・ケースごとにはるかに多くの情報を提供します。

また、ソフト・ターゲットはハード・ターゲットよりも一貫性が高くなるため、モデルの最終的な予測は、最終的には 2 つのロジット値のわずかな差に左右される可能性がありますが、ロジット値自体には、各トレーニング例間の勾配のばらつきははるかに少なくなります。

ソフト・ターゲットによって提供される情報の豊富さと安定性により、元の教師モデルのトレーニングに使用されたものよりも少ないトレーニング例で、より高い学習率を使用して、学生モデルをトレーニングできます。

学生ネットワークの汎化傾向を教師ネットワークの汎化傾向に近づけるために、知識の蒸留は通常、2つの損失関数を使用します。1つ目は「ハード損失」で動作する標準損失関数で、グラウンド・トゥルース・ラベル（教師あり学習の場合）または元のデータ・サンプル（自己教師あり学習の場合）に対して学生モデルの最終アウトプットを測定します。2つ目は蒸留損失で、生徒モデルのソフトターゲットを教師のそれに対して測定する「ソフト・ロス」です。

各トレーニング例には複数のソフト・ターゲットが存在する可能性があるため、蒸留損失は教師ネットワークのソフト・ターゲットの確率分布と 生徒の確率分布の差を測定します。この目的には、カルバック・ライブラー・ダイバージェンス（または「KLダイバージェンス」）が一般的に使用されます。

ロジットは教師と生徒の間の知識移転の典型的な焦点ですが、ディープ・ニューラル・ネットワークでは「知識」がさまざまな方法で表現されます。その他の知識蒸留方法は、ネットワークの隠れ層の重みとアクティベーション、またはネットワークのさまざまな部分間の関係に焦点を当てています。

これらの異なる形式の知識は、通常、応答ベースの知識、機能ベースの知識、または関係ベースの知識の3つのカテゴリーのいずれかに分類されます。

応答ベースの知識は、知識蒸留の最も一般的なジャンルであり、教師モデルの最終アウトプット層からの情報の転送に重点を置いています。典型的な応答ベースのKDでは、生徒モデルは教師モデルの予測と一致するロジットをアウトプットするようにトレーニングされます。

教師モデルのソフト・ターゲットのエントロピーが低い場合、つまり、予測が非常に「確実性が高い」場合（例：分類モデルが1つのクラスに対して1に非常に近いロジット（確実性を表す）をアウトプットし、他のすべてのクラスに対して0に近いロジットをアウトプットする場合）、それほど多くの情報は提供されません。したがって、応答ベースの方法では、モデル・アウトプットに高い温度設定が使用されることが多く、これによりモデル予測のエントロピーが増加します。これにより、より可変的な確率分布が保証され、各トレーニング例からの情報量が増加します。

機能ベースの知識は、ニューラル・ネットワークの中間層、つまり「隠れ層」で伝達される情報に焦点を当てます。ここでニューラル・ネットワークは特徴抽出を行う傾向があり、手元のタスクに関連するインプット・データの明確な特徴やパターンを識別します。

例えば、画像セグメンテーションなどのコンピューター・ビジョン・タスクに主に使用される畳み込みニューラル・ネットワークでは、データがネットワークを介して送信されるにつれて、連続する各隠し層が徐々に豊富な詳細を取得します。動物の画像を種類別に分類するために使用されるモデルでは、最初の隠れ層は単に写真の一部に動物の形があることを識別し、中間の隠れ層は動物が鳥であることを認識し、アウトプット層の直前の最後の隠れ層は、ある種類の鳥と他の近縁種を区別する微妙な詳細を識別します。

したがって、特徴ベースの知識蒸留法の目標は、生徒モデルをトレーニングして、教師ネットワークと同じ特徴を学習させることです。特徴ベースの蒸留損失関数は、2つのネットワークの特徴アクティベーション間の差を測定し、最小化するために使用されます。

関係ベースの知識と機能ベースの知識は、どちらも特定のモデル層のアウトプットに焦点を当てるのに対し、関係ベースの知識による蒸留法は、異なる層間、または異なる層や位置の活性化を表す特徴マップ間の関係に焦点を当てます。

本質的に、関係ベースの知識は、おそらく、教師モデルの「思考プロセス」を模倣するように生徒のネットワークを訓練するための包括的なアプローチを表しています。これらの関係と相関は、特徴マップ間の相関、異なるレイヤー間の類似性を表すマトリックス、特徴の埋め込み、特徴表現の確率分布など、さまざまな方法でモデル化できます。

知識蒸留スキームは、教師のネットワークへの影響によって分類することもできます。Hintonらによって最初に提案された蒸留方法と、その後のそのアプローチの進化のほとんどでは、生徒ネットワークのトレーニングのみを目的としていますが、他の蒸留スキームでは、教師ネットワークの重みの同時更新も必要になります。

オフライン蒸留では、教師ネットワークはすでに事前トレーニングされており、モデルの重みはそれ以上の変更を防ぐために固定されています。オフライン蒸留は、LLMの多くのKDアプローチの典型であり、教師は多くの場合、モデルの重みを変更できない大規模な独自モデルです。

状況によっては、適切に事前トレーニングされ、適切に機能する教師モデルが利用できない場合や、データ・サイエンティストが教師ネットワークを特定の使用ケースに合わせて調整したい場合があります。オンライン蒸留スキームは、教師と生徒の両方のネットワークを同時にトレーニングすることを目的としています。

例えば、Cioppaらは、試合中に視覚的な状況が変化する可能性がある生中継のスポーツ・イベントで使用されるセマンティック・セグメンテーション・モデルのオンライン蒸留スキームを提案しました。これは、生中継の試合データで低速だがパフォーマンスの高いモデルを継続的にトレーニングし、同時にその大規模モデルの知識をより小型で高速なモデルに蒸留してリアルタイムでアウトプットすることで、小規模ネットワークの速度と大規模ネットワークの精度のトレードオフを回避することを目的としていました。⁶

自己蒸留では、1つのネットワークが教師と生徒の両方の役割を果たします。従来の知識蒸留は、あるモデルから別のモデルへの知識の転送を伴うのに対し、自己蒸留は、ネットワークのより深い層から同じネットワークの浅い層への知識の転送として理解できます。⁷

自己蒸留では、複数の注意ベースの「浅層分類ツール」がさまざまな深さでモデルの中間層に追加されます。トレーニング中、より深い階層の分類器は教師モデルとして機能し、アウトプットに対するKLダイバージェンス・メトリクスの損失と特徴マップに対するL2正則化損失という2種類の蒸留損失を通じて、他の注意ベースのモジュールのトレーニングをガイドします。

モデルがトレーニングされ、推論の準備が整うと、これらの浅層分類ツールはすべてモデルから外されます。本質的には、これにより、トレーニング中にモデルが大きくなり、パターン認識能力が向上しますが、デプロイ時にはモデルが小さくなり、結果としてより高速で効率的になります。

LLMの登場により、知識蒸留は、大規模で多くの場合独自のモデルの高度な機能を、小規模で多くの場合オープンソースのモデルに転送する重要な手段として登場しました。そのため、生成AIの民主化において重要なツールとなっています。

最も高い機能を備えたLLMは、ほとんどの場合、コストが高すぎて計算量も多すぎるため、愛好家、スタートアップ企業、研究機関などの多くの潜在的なユーザーがアクセスすることができません。さらに、高度なパフォーマンスと独自の機能にもかかわらず、独自のLLMはその性質上、ニッチなアプリケーションや特定のユースケースに合わせてカスタマイズすることはできません。

さらに、商業的に実現可能なLLMのほとんどは、携帯電話やその他のエッジ・デバイス上でローカルに使用するには大きすぎて計算量も多くなりすぎます。これにより、ロジスティクス、計算、プライバシーに関するさまざまな複雑さが生じますが、携帯端末で直接実行できる小規模なモデルを使用すれば、これらの複雑さを回避できます。したがって、KDのモデル圧縮は、大規模なモデルの出現特性をデバイス上で実行できるほど小さいモデルに転送するための有望な手段を提供します。

LLMにおける知識蒸留のその他の一般的な用途は次のとおりです。

• それぞれ別の言語に特化した複数の教師モデルを使用して、言語知識を単一の生徒モデルに転送したり⁸、または別の言語のモデルを共同トレーニングして、同じ文に対して同様の埋め込みを生成したり⁹することなどにより、LLMを多言語化します。
  
  
• 大規模な独自のLLMを使用して、小規模モデルの命令チューニング用のデータセットを生成します。例えば、Microsoft社のOrcaモデルは、「説明の痕跡、段階的な思考プロセス、その他の複雑な命令を含む、GPT-4の豊富なシグナルから学習」します。¹⁰
  
  
• 教師モデルを使用して生徒のアウトプットをランク付けし、その好みと調整設定を、人間からのフィードバックによる強化学習（RLHF）の一種であるAIからのフィードバックによる強化学習（RLAIF）を通じて蒸留します。¹¹