[論文レビュー] Rethinking Knowledge Distillation via Cross-Entropy

Zhendong Yang, Zhe Li+5 more|arXiv (Cornell University)|Aug 22, 2022

Advanced Neural Network Applications被引用数 13

ひとこと要約

この論文は、古典的KD損失を交差エントロピー（CE）損失と、確率スケールの不一致による最適化の歪みを引き起こす余分な損失に分解することで、知識蒸留を見直している。これを解消するため、非ターゲット分布の知識を分散的に伝達する損失と、教師のターゲット出力を滑らかなラベルとして使用するソフト損失を提案し、CIFAR-100およびImageNetで最先端の精度を達成する新しい損失（NKD）を構築した。教師なしバージョン（tf-NKD）は、学生の予測を平滑化することで性能をさらに向上させ、ResNet-18を用いたImageNetでTop-1精度を+0.86%向上させた。

ABSTRACT

Knowledge Distillation (KD) has developed extensively and boosted various tasks. The classical KD method adds the KD loss to the original cross-entropy (CE) loss. We try to decompose the KD loss to explore its relation with the CE loss. Surprisingly, we find it can be regarded as a combination of the CE loss and an extra loss which has the identical form as the CE loss. However, we notice the extra loss forces the student's relative probability to learn the teacher's absolute probability. Moreover, the sum of the two probabilities is different, making it hard to optimize. To address this issue, we revise the formulation and propose a distributed loss. In addition, we utilize teachers' target output as the soft target, proposing the soft loss. Combining the soft loss and the distributed loss, we propose a new KD loss (NKD). Furthermore, we smooth students' target output to treat it as the soft target for training without teachers and propose a teacher-free new KD loss (tf-NKD). Our method achieves state-of-the-art performance on CIFAR-100 and ImageNet. For example, with ResNet-34 as the teacher, we boost the ImageNet Top-1 accuracy of ResNet18 from 69.90% to 71.96%. In training without teachers, MobileNet, ResNet-18 and SwinTransformer-Tiny achieve 70.04%, 70.76%, and 81.48%, which are 0.83%, 0.86%, and 0.30% higher than the baseline, respectively. The code is available at https://github.com/yzd-v/cls_KD.

研究の動機と目的

• 古典的知識蒸留（KD）損失と元の交差エントロピー（CE）損失との間の根本的関係を分析すること。
• 学生の相対的確率と教師の絶対的確率の不一致によって引き起こされるKDにおける最適化の課題を特定すること。
• 非ターゲット分布の知識伝達の最適化を分離・改善する新しい損失定式化を開発すること。
• 教師のターゲット出力を滑らかなラベルとして扱うソフト損失を提案し、蒸留性能を向上させること。
• 学生の自身の平滑化された予測を用いて安定なソフトターゲットを生成することで、事前学習済み教師が不要な教師なし知識蒸留を可能とすること。

提案手法

• 古典的KD損失を元のCE損失と、学生が教師の非ターゲットクラスの絶対的確率に一致するよう強制する余分な損失に分解する。
• 確率スケールを正規化することで、非ターゲット分布の知識伝達をより効果的にするための分散損失を提案する。
• 教師のターゲット出力をソフトラベルとして使用するソフト損失を導入し、ハードラベルよりも滑らかな指導を提供する。
• 分散損失とソフト損失を元のCE損失と組み合わせ、性能向上を図る新しい知識蒸留（NKD）損失を構築する。
• 学生自身のターゲット出力を訓練エポックにわたって平滑化することで、安定なソフトターゲットを生成する教師なしバージョン（tf-NKD）を開発する。
• tf-NKDにおける学生予測の安定化とソフトターゲット品質の向上を図るため、特定の平滑化関数 $ S_t + V_t - \text{mean}(S_t) $ を適用する。

実験結果

リサーチクエスチョン

• RQ1 最適化ダイナミクスの観点から、古典的KD損失と元のCE損失との関係は何か？
• RQ2 標準KD損失が最適化を劣悪に導く理由は何か？ また、学生と教師の確率スケールの不一致はどのような原因によって生じるか？
• RQ3 KD損失を再定式化することで、非ターゲット分布の知識をより効果的に伝達できるか？
• RQ4 教師のターゲット出力をソフトラベルとして使用することで、蒸留性能を向上させられるか？
• RQ5 学生自身の平滑化された予測をソフトターゲットとして用いることで、教師が不要な状況でも効果的な学生モデルの訓練が可能か？

主な発見

• ResNet-34を教師として用いた場合、提案されたNKD損失により、ResNet-18のImageNet Top-1精度が69.90%から71.96%に向上した。
• 教師なしのtf-NKD手法は、ResNet-18を用いてImageNetで70.76%のTop-1精度を達成し、ベースライン比0.86%の向上を示した。
• CIFAR-100では、NKD損失がResNet-18およびResNet-34を教師として用いた場合に80.76%のTop-1精度を達成し、先行手法を上回った。
• 平滑化関数 $ S_t + V_t - \text{mean}(S_t) $ を用いたtf-NKDバージョンは、ImageNetで最大0.86%の向上を達成し、教師の出力をソフトラベルとして使用する場合でさえもそれを上回った。
• 可視化結果から、平滑化された学生ターゲットが訓練エポックにわたって真のクラス所属関係をより一貫して反映しており、ラベルの安定性が向上していることが確認された。
• すべての重みが訓練中に学生の予測から導出されるため、ベースラインと比較して追加の訓練時間が不要である。