[論文レビュー] Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation

Taehyeon Kim, Jaehoon Oh+3 more|arXiv (Cornell University)|May 19, 2021

Advanced Neural Network Applications参考文献 21被引用数 13

ひとこと要約

本稿では、知識蒸留における標準的なKullback-Leibler（KL）ダイバージェンス損失を、教師モデルと生徒モデルのログィットベクトル間の平均二乗誤差（MSE）損失に置き換える手法を提案する。MSEを用いた直接的なログィットマッチングが、複数のベンチマークでKLダイバージェンスを上回ることを実証している。主な貢献は、MSEがより正確なログィットマッチングを可能にし、特に生徒モデルが教師モデルより小さい場合に一般化性能の向上に寄与することである。

ABSTRACT

Knowledge distillation (KD), transferring knowledge from a cumbersome teacher model to a lightweight student model, has been investigated to design efficient neural architectures. Generally, the objective function of KD is the Kullback-Leibler (KL) divergence loss between the softened probability distributions of the teacher model and the student model with the temperature scaling hyperparameter tau. Despite its widespread use, few studies have discussed the influence of such softening on generalization. Here, we theoretically show that the KL divergence loss focuses on the logit matching when tau increases and the label matching when tau goes to 0 and empirically show that the logit matching is positively correlated to performance improvement in general. From this observation, we consider an intuitive KD loss function, the mean squared error (MSE) between the logit vectors, so that the student model can directly learn the logit of the teacher model. The MSE loss outperforms the KL divergence loss, explained by the difference in the penultimate layer representations between the two losses. Furthermore, we show that sequential distillation can improve performance and that KD, particularly when using the KL divergence loss with small tau, mitigates the label noise. The code to reproduce the experiments is publicly available online at https://github.com/jhoon-oh/kd_data/.

研究の動機と目的

• KLダイバージェンス損失における温度スケーリングハイパーパrameter τ が知識蒸留性能に与える影響を調査すること。
• ソフトラベルが知識蒸留において一般化を向上させる理由について、理論的・実験的Understandingの不足を解消すること。
• KLダイバージェンスの代替として、MSE損失を用いた直接的なログィットマッチング手法を提案し、生徒モデルの性能を向上させること。
• ノイズのあるラベル設定下での異なる損失関数のロバストネスを評価すること。
• 教師モデルと生徒モデルの容量差が大きい場合に、性能を向上させるための逐次的蒸留戦略を検討すること。

提案手法

• 理論的分析により、τ → ∞ のときKLダイバージェンス損失はログィットマッチングに近づき、τ → 0 のときにはラベルマッチングに近づくことが示された。
• 教師モデルと生徒モデルのログィットベクトル間のMSE損失を用いた直接的ログィットマッチングフレームワークを提案し、ソフト確率分布を経由しない。
• 可視化技術（Müller et al., 2019）を用いて、KL損失とMSE損失で訓練されたモデルの最終層の直前特徴量を比較した。
• 最初にKL損失（τ = 3）で訓練し、その後MSE損失でファインチューニングする逐次的蒸留を実施し、容量差が大きい場合の性能向上を検証した。
• CIFAR-100でラベルをランダムに反転させ、異なる損失関数下でのテスト精度を測定することで、ラベルノイズに対するロバストネスを評価した。
• GitHubで公開されているコードを再現することで、再現性を確保した。

実験結果

リサーチクエスチョン

• RQ1 KLダイバージェンス損失における温度スケーリングハイパーパrameter τ は、知識蒸留の一般化行動にどのように影響を与えるか？
• RQ2 MSE損失による直接的ログィットマッチングは、知識蒸留においてKLダイバージェンス損失を上回る性能を達成できるか？
• RQ3 KLダイバージェンスとMSE損失で訓練されたモデルの、最終層の直前特徴量学習にはどのような差があるか？
• RQ4 最初にKL損失、次にMSE損失を用いた逐次的蒸留は、教師と生徒の容量差が大きい場合に性能を向上させるか？
• RQ5 異なる損失関数は、ノイズのあるラベル設定下で、誤ったラベルへの過学習をどのように緩和できるか、特に一般化性能に与える影響は？

主な発見

• MSE損失は、評価したすべてのデータセットでKLダイバージェンス損失を上回り、CIFAR-100では75.78%のテスト精度を達成したのに対し、KL損失（τ=20）では75.47%であった。
• 同じ教師・生徒アーキテクチャを用いた場合、MSE損失はKLダイバージェンスのあらゆる設定（τ=20を含む）よりも高い性能を達成した。
• KLダイバージェンス（τ=20）では75.47%の精度、MSEでは75.78%の精度を記録し、直接的ログィットマッチングがソフトラベルマッチングよりも効果的であることが示された。
• 逐次的蒸留（KLの次にMSE）は、容量差が大きい場合に性能を向上させ、直接蒸留（いずれの損失関数を単独で用いても）を上回った。
• 対称的ノイズが40%のラベルノイズ設定下では、τ=0.5のKL損失が最も優れた一般化性能を示し、MSEおよび高τのKL損失を上回った。これは、内部的なラベルスムージングの効果によるものである。
• MSE損失で訓練されたモデルは、KL損失で訓練されたモデルよりも、教師の最終層の直前特徴量をより忠実に保持していた。KL損失では、特徴量が伸びる（歪曲する）傾向があった。