大規模モデルは小規模モデルに正確にフィードバックし、知識の蒸留はAIアルゴリズムのパフォーマンスを向上させるのに役立ちます。

01 知識蒸留の誕生の背景

近年、ディープニューラルネットワーク (DNN) は、特にコンピュータービジョンのタスクにおいて、産業界と学術界の両方で大きな成功を収めています。ディープラーニングの成功は、主に、データをエンコードするための数十億のパラメータを備えたスケーラブルなアーキテクチャによるものです。そのトレーニング目標は、既存のトレーニングデータセットで入力と出力の関係をモデル化することです。そのパフォーマンスは、ネットワークの複雑さと、ラベル付けされたトレーニングデータの量と品質に大きく依存します。

コンピュータビジョン分野の従来のアルゴリズムと比較して、ほとんどの DNN ベースのモデルは、過剰パラメータ化により強力な一般化機能を備えています。この一般化機能は、モデルが、特定の問題のすべてのデータ入力（トレーニングデータ、テストデータ、問題に属する未知のデータなど）に対して、より優れた予測結果を提供できるという事実に反映されています。

現在のディープラーニングの状況では、アルゴリズムエンジニアはビジネスアルゴリズムの予測効果を向上させるために、次の 2 つのソリューションを採用することがよくあります。

過剰パラメータ化によるより複雑なネットワークを使用します。これは学習には非常に強力ですが、トレーニングには大量のコンピューティングリソースが必要であり、推論速度が遅くなります。

統合モデルは、通常、パラメータと結果の統合を含め、効果が弱い多くのモデルを統合します。

これら 2 つのソリューションは既存のアルゴリズムの有効性を大幅に向上させることができますが、どちらもモデルの規模が大きくなり、計算負荷が大きくなり、大量の計算リソースとストレージリソースが必要になります。

実際には、さまざまなアルゴリズムモデルの最終的な目標は、特定のアプリケーションを提供することです。売買と同様に、収入と支出を管理する必要があります。産業用アプリケーションでは、モデルに優れた予測能力が求められるだけでなく、コンピューティングリソースの使用も厳密に制御する必要があり、効率を考慮せずに結果だけを考慮することはできません。入力データのエンコード量が多いコンピュータービジョンの分野では、コンピューティングリソースがより制限されており、制御アルゴリズムのリソース使用がより重要になります。

一般的に、モデルが大きいほど予測結果は良くなりますが、トレーニング時間が長く推論速度が遅いため、リアルタイムで展開することが難しくなります。特に、ビデオ監視、自動運転車、高スループットのクラウド環境など、コンピューティングリソースが限られているデバイスでは、応答速度が明らかに不十分です。モデルが小さいほど推論速度は速くなりますが、パラメータが不十分なため、推論効果と一般化パフォーマンスはそれほど良くない可能性があります。大規模モデルと小規模モデルのバランスをどのように取るかは、常にホットな話題となっています。現在のほとんどのソリューションは、展開環境の端末デバイスのパフォーマンスに基づいて、適切なサイズの DNN モデルを選択することです。

より高速な推論速度を維持しながら、大規模モデルと同じまたは同様の結果を達成できる小規模モデルが必要な場合、どうすればよいでしょうか?

機械学習では、入力と出力の間に潜在的なマッピング関数の関係があると想定されることが多く、新しいモデルをゼロから学習することは、入力データと対応するラベル間の未知のマッピング関数を近似することです。入力データが変更されないという前提では、小さなモデルをゼロからトレーニングして大きなモデルの効果に近づけることは困難です。一般的に言えば、小規模モデルアルゴリズムのパフォーマンスを向上させるための最も効果的な方法は、より多くの入力データにラベルを付ける、つまり、より多くの監督情報を提供することです。これにより、学習したマッピング関数がより堅牢になり、パフォーマンスが向上します。 2 つの例を挙げると、コンピュータービジョンの分野では、インスタンスセグメンテーションタスクは追加のマスク情報を提供することでターゲット境界ボックス検出の効果を向上させることができます。また、転移学習タスクは、大規模なデータセットで事前トレーニング済みのモデルを提供することで、新しいタスクの予測効果を大幅に向上させることができます。したがって、より多くの監督情報を提供することが、小規模モデルと大規模モデルの間のギャップを狭める鍵となる可能性があります。

前述のように、より多くの教師情報を取得するには、より多くのトレーニングデータにラベルを付ける必要があり、多くの場合、膨大なコストがかかります。では、教師情報を取得するための低コストで効率的な方法はあるのでしょうか? 2006年の論文[1]では、新しいモデルは元のモデル（モデルは関数）に近似できることが指摘されました。元のモデルの機能はわかっているので、新しいモデルをトレーニングすることは、自然にさらに多くの監督情報を追加することと同等であり、明らかにより実現可能です。

さらに考えてみると、元のモデルがもたらした監督情報には、異なる次元の知識が含まれている可能性があります。この固有の情報は、新しいモデル自体では捕捉できない可能性があります。ある程度、これは新しいモデルにとって一種の「クロスドメイン」学習でもあります。

2015年にヒントンは論文「ニューラルネットワークにおける知識の蒸留」[2]で同様のアイデアを使用し、知識蒸留（KD）の概念を初めて提案しました。これは、最初に大規模で強力なモデルをトレーニングし、次にそのモデルに含まれる知識を小さなモデルに転送することで、「小さなモデルのより速い推論速度を維持しながら、大規模モデルと同じまたは同様の効果を達成する」という目標を達成するというものです。最初にトレーニングされる大きなモデルは教師モデルと呼ばれ、後でトレーニングされる小さなモデルは生徒モデルと呼ばれます。トレーニングプロセス全体を比喩的に「教師と生徒の学習」と表現することができます。その後の数年間で、大量の知識の蒸留と教師と生徒の学習作業が生まれ、業界にさらに多くの新しいソリューションが提供されました。現在、KDはモデル圧縮と知識転送の2つの異なる分野で広く使用されています[3]。

02 知識の蒸留

導入

知識蒸留は、「教師と生徒のネットワーク」の考え方に基づいたモデル圧縮手法です。そのシンプルさと有効性から、業界で広く使用されています。目的は、トレーニングされた大規模モデルに含まれる知識を別の小規模モデルに抽出すること、つまり蒸留することです。では、大規模モデルの知識、つまり一般化能力を小規模モデルに転送するにはどうすればよいでしょうか? KD 論文では、大規模モデルのサンプル出力の確率ベクトルをソフトターゲットとして小規模モデルに提供し、小規模モデルの出力をソフトターゲット (元々はワンホットエンコーディング) に近づけて、大規模モデルの動作を近似できるようにしています。

従来のハードラベルトレーニングプロセスでは、すべてのネガティブラベルが均一に扱われますが、この方法ではカテゴリ間の関係が分離されます。たとえば、手書きの数字を認識する場合、同じラベル「3」の画像でも、一部は「8」に似ている一方、他の部分は「2」に似ている場合があります。ハードラベルではこの情報を区別できませんが、十分にトレーニングされた大規模モデルでは区別できます。大規模モデルのソフトマックス層の出力には、正の例に加えて、負のラベルにも多くの情報が含まれています。たとえば、一部の負のラベルに対応する確率は、他の負のラベルに対応する確率よりもはるかに高いなどです。近似学習の動作により、各サンプルは従来のトレーニング方法よりも多くの情報を学生ネットワークに提供します。

そのため、著者は学生ネットワークをトレーニングする際の損失関数を変更し、トレーニングデータのグラウンドトゥルースラベルを適合させながら、小さなモデルが大きなモデル出力の確率分布に適合できるようにしました。この方法は、知識蒸留トレーニング(KD トレーニング) と呼ばれます。知識蒸留プロセスで使用されるトレーニングサンプルは、大規模モデルのトレーニングに使用されるトレーニングサンプルと同じにすることも、独立した転送セットを見つけることもできます。

メソッドの詳細

具体的には、知識蒸留では教師-生徒モデルが使用されます。このモデルでは、教師は「知識」の出力であり、生徒は「知識」の受信者です。知識蒸留のプロセスは 2 つの段階に分かれています。

教師モデルのトレーニング: Net-T と呼ばれる「教師モデル」をトレーニングします。その特徴は、モデルが比較的複雑であり、個別にトレーニングされた複数のモデルから統合することもできることです。「教師モデル」には、モデルを展開する必要がないため、モデルアーキテクチャ、パラメーターの数、統合されているかどうかに関する制限はありません。唯一の要件は、入力 X に対して Y を出力できることです。Y はソフトマックスによってマッピングされ、出力値は対応するカテゴリの確率値に対応します。
学生モデルのトレーニング:少数のパラメータと比較的単純なモデル構造を持つ単一モデルである、Net-S と呼ばれる「学生モデル」をトレーニングします。同様に、入力Xに対してYを出力することができ、Yはソフトマックスマッピング後の対応するカテゴリに対応する確率値を出力することもできます。

ソフトマックスを使用したネットワークの結果は簡単に極端になる可能性があるため、つまり、1 つのクラスの信頼度が非常に高く、他のクラスの信頼度が非常に低い場合、学生モデルが注目する肯定的な情報は依然として 1 つのクラスにのみ属する可能性があります。さらに、異なるカテゴリの否定的な情報も相対的な重要性を持っているため、すべての否定的なカテゴリに同様のスコアを持つことは良くなく、知識の蒸留の目的を達成できません。この問題を解決するために、温度の概念が導入され、高温を利用して小さな確率値によって運ばれる情報を抽出します。具体的には、ロジットはソフトマックス関数を通過する前に温度 T で除算されます。

トレーニング中、教師モデルによって学習された知識は、まず小さなモデルに抽出されます。具体的には、サンプル X の場合、大きなモデルの最後から 2 番目のレイヤーが最初に温度 T で除算され、次にソフトマックスによってソフトターゲットが予測されます。同じことが小さなモデルにも当てはまります。最後から 2 番目のレイヤーが同じ温度 T で除算され、次にソフトマックスによって結果が予測されます。この結果とソフトターゲットのクロスエントロピーは、トレーニングの総損失の一部として使用されます。次に、小さなモデルの通常の出力と真の値ラベル (ハードターゲット) 間のクロスエントロピーが、トレーニングの総損失の別の部分として使用されます。合計損失は、これら 2 つの損失の重みを合計して、小さなモデルをトレーニングするための最終的な損失として計算されます。

小さなモデルをトレーニングして予測する必要がある場合、温度 T は必要なく、従来のソフトマックスに従って直接出力できます。

03 フィットネット

導入

FitNet の論文では、学生モデルのトレーニングをガイドするために、蒸留中に中級レベルのヒントが導入されました。広くて浅い教師モデルを使用して、狭くて深い生徒モデルをトレーニングします。ヒントガイダンスを実行する際には、ヒント層とガイド層の出力形状を一致させるための層を使用することが提案されており、後の研究では適応層と呼ばれることが多い。

一般的に、知識蒸留を行う際には、教師モデルのロジット出力だけでなく、教師モデルの中間層特徴マップも教師情報として使用されます。出力端で小さなモデルが大きなモデルを直接模倣するのは難しすぎると考えられます（モデルが深くなるほどトレーニングが難しくなり、最後の層の監督信号を最前線に伝達するのはかなり疲れます）。最も難しいマッピング関数を直接学習するのではなく、途中で監督信号をいくつか追加して、モデルがトレーニング中に層ごとにより難しいマッピング関数を学習できるようにする方がよいでしょう。さらに、ヒントガイダンスは学生モデルの収束を加速し、非凸問題でより良い局所最小値を見つけるため、学生ネットワークをより深く、より速くトレーニングできます。大学入試問題を解かせることが目的なので、まずは中学校の問題を教え（小型モデルにモデルの前半部分を使って画像の根本的な特徴を抽出することを学習させる）、その後、本来の目的に戻って大学入試問題を学習させる（KDを使って小型モデルのすべてのパラメータを調整する）という感じです。

この記事は、中間特徴マップの蒸留の創始者です。提案されたアルゴリズムは単純ですが、そのアイデアは画期的です。

メソッドの詳細

FitNets の具体的なアプローチは次のとおりです。

教師ネットワークを決定し、成熟するまでトレーニングし、教師ネットワークの中間特徴レイヤーのヒントを抽出します。
一般的に教師ネットワークよりも狭く深い生徒ネットワークを構築します。生徒ネットワークは、生徒ネットワークの中間特徴層が教師モデルのヒントと一致するようにトレーニングされます。生徒ネットワークの中間特徴レイヤーは教師ヒントのサイズと異なるため、ヒントレイヤーのサイズに合わせて特徴次元をアップグレードするには、生徒ネットワークの中間特徴レイヤーの後に回帰器を追加する必要があります。回帰変換後の教師ネットワークのヒント層と生徒ネットワークの中間特徴層に一致する損失関数は、平均二乗誤差損失関数です。

実際のトレーニングでは、前節の KD トレーニングと組み合わせて使用されることが多く、2 段階のトレーニング方法が使用されます。最初にヒントトレーニングを使用して、小規模モデルの前半のパラメーターを事前トレーニングし、次に KD トレーニングを使用してすべてのパラメーターをトレーニングします。蒸留プロセスではより多くの監督情報が使用されるため、中間特徴マップに基づく蒸留方法は結果ロジットに基づく蒸留方法よりも優れていますが、トレーニング時間は長くなります。

04 結論

知識蒸留は、アンサンブルまたは高度に正規化された大規模モデルから小規模なモデルに知識を転送するのに非常に効果的です。蒸留モデルのトレーニングに使用される移行データセットに 1 つ以上のクラスのデータが欠落している場合でも、蒸留は非常にうまく機能します。古典的な研究であるKDとFitNetが提案された後、さまざまな蒸留方法が雨後の筍のように現れました。将来的には、モデル圧縮と知識移転の分野でもさらなる探求を行っていきたいと考えています。