蒸留(Knowledge Distillation)
1. 蒸留とは何か?
蒸留(Knowledge Distillation)とは、「高い性能を持つ巨大なモデル(親/Teacher)の知識を、軽量なモデル(子/Student)に効率よく引き継がせる技術」です。
モデルを小型化・高速化しつつ、単純にゼロから学習させるよりも高い精度を維持することを目的としています。
2. なぜ通常の学習(教師データのみ)だけでは不十分なのか?
通常の学習では、人間が作成した [犬: 1, 猫: 0, 狐: 0] という 100% デジタルな正解データ(ハードターゲット)のみを目指します。
しかし、この方法では親モデルが学習プロセスで獲得した「この画像は犬だけど、20% くらいは狐にも似ている(猫には全く似ていない)」という、クラス間の相関関係や判断のニュアンス(暗黙知 / 暗黒知識:Dark Knowledge)を子モデルに移植することができません。
蒸留では、親モデルが出力する確率分布(ソフトターゲット)をターゲットにすることで、この高度な構造を直接子モデルに教え込みます。
3. 蒸留のメカニズムと重み行列の最適化
重み最適化の本質
通常の学習と蒸留において、各レイヤの重み行列(パラメータ)を誤差逆伝播法(Backpropagation)によって更新・最適化するというアルゴリズムは同じです。
違いは、重み行列の最適化に使用する損失関数の設計にあります。
- 通常の学習: モデルの出力が正解ラベル(Hard Label)に近づくように重み行列を最適化します
- 蒸留の学習: 子モデルの出力が、親モデルの確率分布(Soft Label)に近づくように重み行列を最適化します
つまり正解データの確率分布を $p$ 、モデルが算出した確率分布を $q$ とした場合の交差エントロピー誤差は以下で求めらます。
\[H(p,q) = - \sum_{x} p(x) \log_2 q(x)\]このときの $p$ に使う分布が正解データか親モデルの予想した確率分布かの違いになります。
損失関数(Loss Function)のハイブリッド構成
実際の蒸留では、親モデルの分布だけでなく、元の正解データも合わせて学習させるのが一般的です。
全体の損失関数 $\mathcal{L}_{\text{total}}$ は以下のようになります。
- 第1項:通常の分類誤差(Cross Entropy Loss)
- 元の正解データ($y$)と子モデルの出力($z_s$)の誤差。基本的な正解のラインを担保します
- 第2項:蒸留誤差(Knowledge Distillation Loss / KL Divergence)
- 親モデルの出力($z_t$)と子モデルの出力($z_s$)の分布間のズレ(カルバック・ライブラー情報量)を計算します
- 温度(Temperature, $T$): 確率分布をあえて滑らか(マイルド)にし、隠れた特徴(狐に似ている度合いなど)を子モデルが学びやすくするためのハイパーパラメータです。
- ハードターゲット損失 ($L_{\text{hard}}$): 通常の正解ラベルとの交差エントロピー。このとき温度 $T=1$。
- ソフトターゲット損失 ($L_{\text{soft}}$): 先生モデルの出力(予測結果)との交差エントロピー。このとき温度 $T > 1$。
- スケールの調整: $L_{\text{soft}}$ の勾配は $1/T^2$ に比例して小さくなる性質があるため、2つの損失のスケールを揃えるために $L_{\text{soft}}$ に $T^2$ を掛け算します。
4. 蒸留のメリット
- 軽量化と高速化: 少ないレイヤ数・次元数(小さな重み行列)で構成される子モデルへ圧縮できるため、スマートフォンやエッジデバイス、リアルタイム性が求められる環境(自動運転など)での推論が可能になります
- 精度の維持: 親モデルの「思考プロセス」をショートカットして吸収するため、小さなモデルを最初から普通に学習させるよりも圧倒的に高い精度に到達します
- コストの削減: サーバの計算コストや推論時の消費電力を大幅に抑えられます