[論文レビュー] Improving Knowledge Distillation via Regularizing Feature Norm and Direction

Yuzhu Wang, Lechao Cheng+4 more|arXiv (Cornell University)|May 26, 2023

COVID-19 diagnosis using AI被引用数 5

ひとこと要約

本論文は、特徴量のノルムと方向を同時に正則化することで、学生モデルの性能を向上させる、新しい知識蒸留法を提案している。具体的には、教師モデルのクラス平均特徴量と整合性のある高ノルム特徴量を生成するよう学生モデルを促進する。主な貢献は、特徴量の大きさと整合性を同時に向上させるシンプルで効果的なND損失であり、既存のKD手法と組み合わせることで、ImageNet、CIFAR-100、COCOベンチマークで最先端の精度を達成している。

ABSTRACT

Knowledge distillation (KD) exploits a large well-trained model (i.e., teacher) to train a small student model on the same dataset for the same task. Treating teacher features as knowledge, prevailing methods of knowledge distillation train student by aligning its features with the teacher's, e.g., by minimizing the KL-divergence between their logits or L2 distance between their intermediate features. While it is natural to believe that better alignment of student features to the teacher better distills teacher knowledge, simply forcing this alignment does not directly contribute to the student's performance, e.g., classification accuracy. In this work, we propose to align student features with class-mean of teacher features, where class-mean naturally serves as a strong classifier. To this end, we explore baseline techniques such as adopting the cosine distance based loss to encourage the similarity between student features and their corresponding class-means of the teacher. Moreover, we train the student to produce large-norm features, inspired by other lines of work (e.g., model pruning and domain adaptation), which find the large-norm features to be more significant. Finally, we propose a rather simple loss term (dubbed ND loss) to simultaneously (1) encourage student to produce large-\emph{norm} features, and (2) align the \emph{direction} of student features and teacher class-means. Experiments on standard benchmarks demonstrate that our explored techniques help existing KD methods achieve better performance, i.e., higher classification accuracy on ImageNet and CIFAR100 datasets, and higher detection precision on COCO dataset. Importantly, our proposed ND loss helps the most, leading to the state-of-the-art performance on these benchmarks. The source code is available at \url{https://github.com/WangYZ1608/Knowledge-Distillation-via-ND}.

研究の動機と目的

• 既存の知識蒸留手法が、特徴量の大きさや分類器の構造を考慮せずに、学生の特徴量を教師の特徴量と直接一致させることによる制限を是正すること。
• 教師のクラス平均特徴量を強力で内蔵された分類器として活用し、学生の特徴量学習をガイドすることで、学生の性能を向上させること。
• モデルのプルーニングやドメイン適応の研究で得られた知見にインspiredされ、大規模なノルム特徴量が知識蒸留に果たす役割を調査すること。
• 特徴量のノルムと方向を同時に正則化するシンプルで即席の損失（ND損失）を開発し、蒸留性能を向上させること。

提案手法

• 教師モデルの既存の特徴量からクラス平均特徴量を構築し、これを強力な分類器として学生の特徴量学習をガイドする。
• 特徴量のL2ノルムを最大化するのと、学生特徴量と対応する教師のクラス平均特徴量とのコサイン距離を最小化するという2つの目的を同時に最適化する、新しいND損失を導入する。
• ND損失は、特徴量がログティスを計算するペンシマット層（埋め込み層）に適用され、分類性能と直接関連する。
• ログティス蒸留および特徴量蒸留の両フレームワークと互換性があり、既存のKDパイプラインにスムーズに統合可能である。
• ImageNet、CIFAR-100、COCOデータセットを用いて実験的に検証され、アブレーションスタディにより各構成要素の有効性が確認されている。
• 実装はKD++として公開されており、計算コストの増加が最小限に抑えられながら、顕著な精度向上が得られている。

Figure 1 : Our main contribution is a simple loss, termed $\mathcal{L}_{nd}$ , that regularizes the n orm and d irection of the student features. $\mathcal{L}_{nd}$ is applicable to different KD methods which are categorized into two types in the context of classification: (left) logit distillation

実験結果

リサーチクエスチョン

• RQ1 学生特徴量のノルムを大きくし、教師のクラス平均特徴量と整合させることで、知識蒸留の性能が向上するか？
• RQ2 単一の統合損失関数が、ノルム最大化と方向整合性の両方を効果的に統合できるか？
• RQ3 ND損失は、ベースラインの正則化手法（例：コサイン損失、ノルム正則化）と比較して、性能向上にどの程度寄与するか？
• RQ4 教師モデルの容量が増加するに従い、ND損失の利点はスケーリングするか？
• RQ5 アーキテクチャの違いがあるにもかかわらず、TransformerからCNNへの蒸留においても、ND損失は性能向上に寄与するか？

主な発見

• 標準的なKD（KD++）にND損失を適用した場合、ResNet-18を学生モデル、ResNet-50を教師モデルとして用いたImageNetでは、トップ-1精度が72.53%に達し、元のKD（71.35%）や最近の手法（ReviewKD：71.10%、DKD：71.87%）を上回った。
• ResNet-18をImageNetで学習する際、KDに比べて1.45%の絶対的精度向上を達成しており、ND損失の強力な有効性を示している。
• より大きな教師モデル（ResNet-34からResNet-152）を用いた場合、KD++は一貫した性能向上（例：71.98% → 72.54%）を示したが、ベースラインKD手法は劣化や変動を示した。
• COCOオブジェクト検出タスクにおいても、ND損失の適用により検出精度が向上し、画像分類を超えた利点が確認された。
• 可視化結果から、KD++はより高いノルム特徴量と優れたクラス分離性を生成しており、標準的なKDや単純な学生モデルが見逃していた誤分類クラス（例：紫色のクラス）を効果的に回復している。
• ND損失は、交差エントロピー損失やKL損失と単独で使用しても有効であり、古典的なKDと比較して約1%の精度向上が得られ、計算コストは無視できるほど低い。

(a) teacher (Res56)