巨大な AI モデルが知っていることすべてを小さな AI モデルに教える方法

目次

1. 導入
2. ダークナレッジ
3. 蒸留手順
4. MNISTでの実験
5. MNISTに関する追加実験

1. はじめに

この記事では、AI における知識蒸留プロセスについて、その一般的な仕組み、その重要性、そしてそれを使用する理由について説明します。

パフォーマンスをあまり低下させずに、より大きなモデルまたはモデルのアンサンブル（データから構造を抽出するために非常に大きなデータセットでトレーニングされたもの）から単一の小さなモデルに知識を圧縮して転送するにはどうすればよいでしょうか。

しかし、なぜこれを実行するのでしょうか? 大規模なモデルやアンサンブル モデルがすでにテスト データで優れた結果を示しているのに、なぜ小規模なモデルが必要なのでしょうか?

トレーニング時には、主な目標が非常に大きなデータセットから構造を抽出することであるため、通常は大規模なモデル/アンサンブルをトレーニングします。また、トレーニング時にドロップアウトやデータ拡張などの多くのものを適用して、これらの大規模なモデルにあらゆる種類のデータを供給することもできます。

しかし、予測時の私たちの目的はまったく異なります。できるだけ早く結果を得たいのです。そのため、より大きなモデルやモデルのアンサンブルを使用するとコストが非常に高くなり、多数のユーザーへの展開が妨げられます。そこで、この大きなモデルからの知識を、簡単に展開できる小さなモデルに圧縮するにはどうすればよいかという疑問が生じます。

Google の Geoffrey Hinton、Oriol Vinyals、Jeff Dean は論文を通じて、この知識をより小さなモデルに転送するための蒸留と呼ばれる異なる種類のトレーニングを考案しました。これは、Hugging Face がDistill BERT実装で使用したのと同じ手法です。

この小さなモデルを大きなモデルと同じ方法で一般化するように訓練できれば、この方法で訓練されたこの小さなモデルは、同じデータで通常の方法で訓練された小さなモデルよりもはるかに優れたパフォーマンスを発揮します。これが蒸留の背後にある主要な原則の1つです。

2. ダークナレッジ

通常、機械学習では、多数のクラスを区別することを学習するモデルの主なトレーニング目標は、正解の平均対数確率を最大化することです。たとえば、MNIST データセットの例を考えてみましょう。このデータセットの目標は、画像を 1 か 2 か、あるいは 9 かに分類することです。したがって、実際の画像が 2 の場合、モデルの目標はP(その 2/画像) (特定の画像が与えられた場合に 2 である確率と読み替えることができます) を最大化することです。ただし、モデルはすべての不正解に対しても確率を与えますが、それらの確率は非常に小さく、そのうちのいくつかは他の確率よりもはるかに大きいです。重要なのは、これらの確率が小さくても、不正解の相対的な確率から、モデルがどのように一般化できるかについて多くのことがわかるということです。これを理解するために、以下の例を見てみましょう。

上の図では、このバージョンの 2 が 3 になる確率は 10-6、7 になる確率は 10-9 ですが、別のバージョンではその逆になる場合があります。これは、データ上の豊富な類似性構造を定義する貴重な情報です (つまり、どの 2 が 3 に見え、どの 2 が 7 に似ているかを示します)。ただし、確率がゼロに非常に近いため、転送段階でのクロスエントロピー コスト関数にはほとんど影響しません。

しかし、蒸留手順に進む前に、モデルが実際に出力確率をどのように生成したかについて少し考えてみましょう。ここで、ソフトマックス活性化が登場します。モデル処理の最後のステップはソフトマックスであり、このコンポーネントが出力確率を生成します。ソフトマックスへの入力はロジットと呼ばれ、NN の最終層は、隠れユニットの数 = 分類するクラスの数となるように設計します。

ソフトマックスを計算する式は次のように与えられる。

上記の式は各iの確率を示し、 i全体のすべての確率の合計は 1 になります。トレーニング時間中、これらのソフトマックス確率をハード ターゲット (ラベル) と比較し、バックプロパゲーション係数を使用して損失が最小になるまで更新することで、単一のトレーニング例の損失が計算されます。

上で見たように、このソフトマックスは正しいラベルに高い確率を与え、間違ったラベルに低い確率を与えます。また、間違った答えの確率は小さくても、モデルの一般化に役立つ多くの情報が隠されていることもわかります。これをダークナレッジと呼びます。

3. 蒸留手順

論文によると、大規模モデルの一般化機能を小規模モデルに移行する最良の方法は、扱いにくいモデルによって生成されたクラス確率を、小規模モデルをトレーニングするためのソフトターゲットとして使用することです。

プロセスは次のようになります。

• より大きなモデルをトレーニングするために使用された元のトレーニング セットを取得し、そのトレーニング データをより大きなモデルに渡して、さまざまなクラスにわたるソフトマックス確率を取得します。上記のように、正しいラベルは高い確率を取得し、誤ったラベルは低い確率を取得します。しかし、これらの低い確率には多くの情報が隠されていることがわかりました。そのため、これらの確率の重要性を拡大するために、論文の著者は、ソフトマックスを通過する前にすべてのロジットを分割するために Temperature(T) と呼ばれる変数を使用しました。これにより、クラス間のよりソフトな確率分布が生成されます。以下を参照してください。

温度(T)でソフトマックスを適用した結果は、ソフトターゲットと呼ばれるものです。このプロセスは、著者が蒸留と呼んでいるものです。温度を上げて水中の不純物を除去するのと似ています。

• 大規模モデルから学習された関数に関する情報の多くは、ソフトターゲット内の非常に小さな確率の比率に存在します。
• いくつかの用語:

1. ソフトターゲット- ソフトマックス中に温度Tが適用された後の大規模モデルからの出力
2. ソフト予測- ソフトマックス中に温度Tが適用された後の小さいモデルからの出力
3. ハード予測- 温度 T = 1 の場合の小さいモデルからの出力 (通常のソフトマックス)
4. 真のターゲット- トレーニング セットからの実際のターゲット

以下はトレーニングプロセス全体のフローチャートです。

したがって、小さなモデルのトレーニング プロセスには 2 つの損失関数があります。最初の損失関数は、ソフト予測とソフト ターゲットの両方を取り、クロスエントロピー損失関数です。これは、ソフト ターゲットを一致させることによって、一般化能力が大きなモデルから小さなモデルに転送される方法です。この損失関数では、両方のソフトマックスが 'T' の温度を使用します。

著者らは、小さなモデルを使用して真のターゲットを一致させることが役立つことも発見しました。これは 2 番目のコスト関数に組み込まれています。最終的なコストは、ハイパーパラメータ アルファとベータを使用してこれら 2 つのコスト関数の加重平均になります。

4. MNISTでの実験

著者らはMNISTデータセットを使用してこのアプローチをテストした。このために、中間層の隠れユニットの数だけが異なる2つのアーキテクチャを使用した。著者らはどちらの場合も2つの隠れ層ニューラルネットワークを使用した。

1. 784 -> 800 -> 800 -> 10 と見ることができるより小さなモデル (784 は画像の展開された次元、800 は RELU アクティベーションを持つ隠れユニットの数、10 は予測するクラスの数)。このモデルは、正規化なしで 146 のテスト エラーを生成しました。
2. 784 -> 1200 -> 1200 -> 10 と見ることができるより大きなモデルです (784 は画像の展開された次元、1200 は RELU アクティベーションを持つ隠れユニットの数、10 は予測するクラスの数です)。このモデルは、ドロップアウト、重み制約、およびジッタ入力画像を使用して MNIST でトレーニングされ、このネットは 67 のテスト エラーを達成しました。

より大きなモデルでのこの改善を小さなモデルに転送できますか?

著者らは、大きなネットから取得したソフト ターゲットと、ドロップアウトや画像のジッタのない実際のターゲットの両方を使用しました。つまり、小さなネットは、温度 20 で大きなネットによって生成されたソフト ターゲットを一致させるという追加タスクを追加するだけで正規化され、結果は次のとおりです。

784 -> 800 -> 800 -> 10 を使用した74 回のテストエラー

これは、ソフト ターゲットが、変換されたトレーニング データから学習した一般化方法に関する知識を含め、大量の知識を小さなモデルに転送できることを示しています。言い換えると、入力を変換することで得られた利点は、小さなネットの入力を変換していなくても、小さなネットに転送されます。

入力をさまざまな変換で変換すると、モデルの一般化が大幅に改善されることはよく知られている事実です。この場合、一般化の方法に関する情報はダーク ナレッジに表示され、ソフト ターゲットに隠れています。この情報はいずれも真のターゲットには存在しません。したがって、ソフト ターゲットからの情報を使用することで、私たちの小さなネットのパフォーマンスは大幅に向上します。

ソフトターゲットを使用するビッグネットは、「何が何に似ているか」を学習した類似性メトリックを学習し、この知識転送により、リトルネットに「何が何に似ているか」を伝えています。

MNISTに関する上記の実験はすべて以下に要約されている。

5. MNISTに関する追加実験

さらに、著者らは蒸留を使用して小型モデルをトレーニングする際に、数字 3 の例を省略することも試みました。そのため、小型モデルの観点からは、3 は見たことのない架空の数字です。小型モデルはトレーニング中に 3 を見たことがないため、テスト セットで 3 に遭遇すると多くのエラーが発生することが予想されます。それにもかかわらず、蒸留モデルは 206 のテスト エラーしか発生せず、そのうち 133 はテスト セット内の 1010 個の 3 に対するものでした。これは、蒸留中に大型モデルの一般化機能が小型モデルに転送され、これにより小型モデルがほとんどの場合に 3 を正しく予測していることを明確に示しています。

これがこの物語の教訓です。

入力画像を変換すると、一般化が大幅に改善されます。ターゲットを変換すると、同様に大きな効果があり、ソフトターゲットをどこかから取得できれば、より小さなモデルで同じパフォーマンスを得ることができるため、コストがはるかに安くなります。

参考文献:

1. Hinton、Geoffrey、Oriol Vinyals、Jeff Dean。「ニューラル ネットワークで知識を抽出する」
2. intellabs による知識の抽出。