蒸留

蒸留とは

蒸留(distillation)は、手間と計算をかけて訓練した大きなモデルの知識を、小さなモデルへ移すための訓練の方法です。2015年にGoogleのヒントン(Hinton)、ヴィニャルズ(Vinyals)、ディーン(Dean)が論文 Distilling the Knowledge in a Neural Network で、この方法に蒸留という名前を付けました。

知識を渡す側の大きなモデルについて、論文は次の2つを挙げています。

• 別々に訓練した複数のモデルを組み合わせたアンサンブル
• ドロップアウトのような強い正則化(過学習を抑えるための工夫)をかけて訓練した、1つの非常に大きなモデル

アンサンブルは、同じデータで複数のモデルを訓練し、それぞれの予測を平均する仕組みです。論文は、これがほとんどどんな機械学習アルゴリズムでも性能を上げる、とても単純な方法だと述べています。一方で、予測のたびにモデルをすべて動かすのは手間がかかり、特に1つ1つが大きなニューラルネットワークのときは、多くの利用者へ提供するには計算の負担が大きすぎることがある、とも指摘しています。蒸留は、アンサンブルが持つ知識を、提供しやすい1つのモデルへ圧縮するための方法です。

訓練段階と展開段階では求められるものが違う

論文は、蒸留が必要な理由を、訓練の段階と展開(実際に利用者へ提供する段階)の違いから説明しています。

• 訓練段階: 音声認識や物体認識では、非常に大きく重複の多いデータセットから構造を取り出す必要がある。ただし実時間で動く必要はなく、膨大な計算を使ってよい
• 展開段階: 多くの利用者に提供するため、応答の遅れ(レイテンシ)と計算資源について、ずっと厳しい制約がある

それにもかかわらず、大規模な機械学習では、訓練と展開でよく似たモデルを使うのが普通でした。

論文はここで昆虫の例えを使います。多くの昆虫は、環境からエネルギーと栄養を取り込むのに適した幼虫の形と、移動と繁殖に適したまったく違う成虫の形を持っています。これと同じように、訓練ではデータから構造を取り出しやすいなら手に負えないほど大きなモデルを使ってよく、訓練が終わってから別の種類の訓練で知識を展開向きの小さなモデルへ移せばよい、という考え方です。この別の種類の訓練が蒸留です。

大きなモデルの確率を目標にする

大きなモデルの汎化能力(未知のデータにも正しく答える力)を小さなモデルへ移すわかりやすい方法は、大きなモデルが出力した各クラスの確率を、小さなモデルを訓練するときの目標として使うことです。論文はこの目標をソフトターゲット(soft target、柔らかい目標)と呼びます。

これと対になるのがハードターゲット(hard target)で、正解のクラスだけを1とし、ほかのクラスを0とする通常の正解ラベルです。ソフトターゲットは、正解以外のクラスにもどれくらいの確率を割り振ったかという情報を含みます。

論文によると、ソフトターゲットのエントロピーが高い(確率が1つのクラスに集中せず散らばっている)ときは、次の利点があります。

• 訓練データ1件あたりに含まれる情報がハードターゲットよりずっと多い
• 訓練データの間で勾配(パラメータを更新する方向と大きさ)のばらつきがずっと小さい

そのため、小さなモデルは元の大きなモデルよりずっと少ないデータで、しかもずっと高い学習率で訓練できることが多い、と論文は述べています。

ごく小さな確率にも情報がある

論文は、手書き数字の画像データセットMNISTを例に挙げています。MNISTのように、大きなモデルがほぼいつも高い確信度で正解を出す課題では、学習した関数についての情報の多くが、ソフトターゲットの中のごく小さな確率どうしの比に含まれています。

たとえば、数字の2を書いたある画像について、3である確率が10のマイナス6乗、7である確率が10のマイナス9乗と出たとします。別の2の画像では、この大小が逆になることがあります。これは、どの2が3に似ていて、どの2が7に似ているかを表しており、データの間の豊かな類似の構造を示す価値のある情報です。

ところが、これらの確率は0に非常に近いため、知識を移す段階で使う交差エントロピー(予測と目標のずれを測る損失の関数)にはほとんど影響しません。せっかくの情報が、そのままでは小さなモデルの訓練に生かされないのです。

softmaxの温度を上げる

この問題への論文の解決策が、蒸留の中心となる手順です。

1. 大きなモデルの最後にあるsoftmax(出力を各クラスの確率に変える関数)の温度を上げる。大きなモデルが十分に柔らかい目標、つまり正解以外のクラスにもはっきりと確率が割り振られた目標を出すようになるまで上げる
2. 小さなモデルを訓練するときにも同じ高い温度を使い、そのソフトターゲットに合わせるように訓練する

温度を上げると、0に近かった小さな確率が相対的に大きく出るため、どのクラスとどのクラスが似ているかという情報が、小さなモデルの訓練に効くようになります。