蒸留：大きなモデルから小さなモデルへ教える

最も大きく最も有能なモデルは、最も実行コストが高くもあります。遅く、クエリあたりのコストが高く、しばしば実際に手元にあるハードウェアに収まりません。知識蒸留は、その能力の大部分を保ちながらコストの大部分を削ぎ落とす技術です。一行で言えば、大きなモデルを真似るよう小さなモデルを訓練する。 面白いのは、「真似る」が何を意味することになるかです。

この技術は、扱いにくく正確なモデルを、コンパクトで展開可能なものに圧縮する方法として、ディープラーニングの時代に広まりました。その構図には覚えやすい名前があります。大きなモデルが教師、小さなモデルが生徒です。生徒は、タスクをゼロから再発見するのではなく、教師の振る舞いを再現するよう訓練されます。

なぜ小さなモデルを直接訓練しないのか

明白な代替案は、教師が見たのと同じラベル付きデータで小さなモデルを訓練し、教師を完全に飛ばすことです。それでうまくいくこともあります。うまくいかないこともよくあり、その理由は示唆に富みます。

実際の訓練ラベルはたいていハードラベルです。この画像は猫、以上。その単一の答えは、よく訓練された教師が知っていることの多くを捨ててしまいます。良い教師は単に「猫」と言うのではなく——「ほぼ確実に猫、わずかに犬の可能性、飛行機では断じてない」と言います。すべての選択肢にわたるその自信の分布は、一語のラベルよりはるかに豊かな訓練信号です。それは、どの間違いが妥当でどれが馬鹿げているかを符号化しています。ハードラベルで訓練された小さなモデルはそれを決して見ません。教師の完全な出力で訓練された生徒は見るのです。

ソフトターゲット：発想の核心

起こりうる答えにわたる教師の完全な確率分布は、しばしばソフトターゲットと呼ばれます（単一の正解ラベルというハードターゲットに対して）。これらのソフトターゲットは、研究者が時にダークナレッジと呼ぶもの——教師が学んだ、ラベル自体には見えない関係——を運びます。

数字認識を考えてみましょう。手書きの7は、わずかな確率を1へ引くかもしれません。7と1は似て見えうるからです。そして8へはほとんど引きません。1へのその小さな傾きは、入力の形と問題の構造についての本物の情報です。トップの答えだけでなく分布全体に合わせるよう生徒を訓練することが、その構造を転移させます。生徒は教師の世界観を学ぶのであって、結論だけではありません。

これらのソフトターゲットをさらに情報豊かにするため、蒸留はしばしば分布をさらに柔らかくし、確率を広げて、次点の選択肢の間の小さな差をより際立たせ学びやすくします。生徒は、この柔らかくした絵に密接に合わせるよう求められます。

何が転移し、何が転移しないか

蒸留は理解ではなく振る舞いを転移させます。生徒は、訓練された種類の入力に対して、教師の出力に似た出力を生み出すことを学びます。それは強力であり、同時に限定的です。

• カバレッジの良さまでしか良くならない。 生徒は、見た例において教師を真似ます。蒸留データのどれとも似ない入力では、生徒には写すべき教師がなく、なんとか一般化できたものに頼ります。
• 教師の欠点を受け継ぎうる。 教師が偏っていたり、過信していたり、体系的に間違っていたりすれば、生徒もそれを写します。蒸留は忠実な模倣であり、間違いの忠実な模倣も含みます。
• 蒸留されたタスクで教師を超えることはまれ。 生徒は教師の振る舞いを追いかけており、その特定の信号については、たとえ生徒のほうが効率的でも、教師が天井です。

これらは蒸留を有用でなくしません。期待値を設定するだけです。あなたが買っているのは効率であって、新しい能力ではありません。

言語モデルにおける蒸留

同じ発想が大規模言語モデルにも適用されます、いくつかのひねりを伴って。言語モデルは次のトークンを語彙にわたる確率分布として予測するので、そのソフトターゲットはまさに蒸留が栄える種類の豊かな信号です。生徒モデルは、大量のテキストにわたって教師の次トークン分布に合わせるよう訓練できます。

第二の、ますます一般的になっている流儀があり、これは教師の内部確率へのアクセスをまったく必要としません。ここでは教師が単に出力を生成し——答え、説明、解いた解——生徒はその生成されたテキストを通常の訓練データのように訓練します。これはシーケンスレベルまたは生成ベースの蒸留と呼ばれることがあり、モデルが生み出したデータで訓練するというより広い実践へと溶け込んでいきます。問い合わせられる教師なら、テキストを返すインターフェース越しにしか届かない教師でも使えるので便利です。

両方の流儀は、核心の賭けを共有します。生のラベルだけでなく大きなモデルの振る舞いで訓練すれば、小さなモデルが大きなモデルの能力の驚くほどの一部を担いうる、という賭けです。

なぜこれが実務で重要なのか

蒸留は、有能なAIが安く、必要とされる場所の近くで動ける主な理由の1つです。蒸留されたモデルは、大量に提供できるほど小さく、インタラクティブな利用に十分速く、控えめなハードウェアで動くほどコンパクトでありえます。多くの実際の展開では、問いは「存在する中で最も有能なモデルは何か？」ではなく「1日に100万回実行する余裕のある、最も有能なモデルは何か？」です。蒸留はそのフロンティアを動かします。

それはまた、有用な分業を可能にします。大きく高コストな1つの教師に重く投資し、それを、異なるコストとレイテンシの予算に合わせた一群の小さな生徒へと蒸留する。難しい仕事の代金を一度払い、多くのより安いモデルにわたって償却するのです。

正直なトレードオフ

蒸留はタダではなく、無損失でもありません。

• いくらかの品質を手放す。 生徒はより小さく、最も難しい入力では教師と生徒の差が現れます。技は、その損失があなたのユースケースで許容できる生徒のサイズを選ぶことです。
• 適切なデータが要る。 生徒は、教師が実演する場所でしか学びません。何の上で蒸留するか——あなたが実際に気にかける入力をカバーすること——を選ぶことは、アルゴリズムと同じくらい重要です。
• 静かな失敗を増幅しうる。 生徒が教師を無批判に写すため、教師の微妙な偏りが、その後広く出荷するモデルに焼き込まれうるのです。

これらの限界を知ることが、信頼できるエンジニアリングの道具としての蒸留と、希望的な近道としての蒸留を分けます。

まとめ

知識蒸留は、大きな教師を真似るよう小さな生徒を訓練します——そして鍵となる洞察は、写すべき最も価値あるものが教師の最終的な答えではなく、自信の完全な分布、すなわち教師が不確かさをどう捉えるかを明かすソフトターゲットだということです。その豊かな信号が、コンパクトなモデルに、大きなモデルの能力の多くをわずかなコストで担わせます。蒸留は教師を超えはせず、教師の欠点を受け継ぎますが、高コストな能力を展開可能な能力に変える方法として、現代の機械学習で最も静かに重要な技術の1つです。