[論文レビュー] From Knowledge Distillation to Self-Knowledge Distillation: A Unified Approach with Normalized Loss and Customized Soft Labels

Zhendong Yang, Ailing Zeng+4 more|arXiv (Cornell University)|Mar 23, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

本稿では、正規化知識蒸留（NKD）と汎用自己知識蒸留（USKD）を導入することで、知識蒸留（KD）と自己知識蒸留（self-KD）を統一するフレームワークを提案する。NKDは非ターゲットスコアを正規化することで、ソフトラベルの利用効率を向上させる。一方、USKDは学生モデルの特徴量とZipfの法則を用いて、教師モデルを必要とせず、ターゲットおよび非ターゲットクラスのカスタマイズされたソフトラベルを生成する。この手法により、CNNおよびViTアーキテクチャで最先端の性能を達成しつつ、計算コストを最小限に抑える。

ABSTRACT

Knowledge Distillation (KD) uses the teacher's prediction logits as soft labels to guide the student, while self-KD does not need a real teacher to require the soft labels. This work unifies the formulations of the two tasks by decomposing and reorganizing the generic KD loss into a Normalized KD (NKD) loss and customized soft labels for both target class (image's category) and non-target classes named Universal Self-Knowledge Distillation (USKD). We decompose the KD loss and find the non-target loss from it forces the student's non-target logits to match the teacher's, but the sum of the two non-target logits is different, preventing them from being identical. NKD normalizes the non-target logits to equalize their sum. It can be generally used for KD and self-KD to better use the soft labels for distillation loss. USKD generates customized soft labels for both target and non-target classes without a teacher. It smooths the target logit of the student as the soft target label and uses the rank of the intermediate feature to generate the soft non-target labels with Zipf's law. For KD with teachers, our NKD achieves state-of-the-art performance on CIFAR-100 and ImageNet datasets, boosting the ImageNet Top-1 accuracy of ResNet18 from 69.90% to 71.96% with a ResNet-34 teacher. For self-KD without teachers, USKD is the first self-KD method that can be effectively applied to both CNN and ViT models with negligible additional time and memory cost, resulting in new state-of-the-art results, such as 1.17% and 0.55% accuracy gains on ImageNet for MobileNet and DeiT-Tiny, respectively. Our codes are available at https://github.com/yzd-v/cls_KD.

研究の動機と目的

• 知識蒸留におけるソフトラベルの利用を向上させるために、KD損失を再定式化すること。
• 学生と教師の分布をよりよく一致させるために、共通の損失フレームワークを用いてKDとself-KDを統合すること。
• 教師モデルを必要とせず、一般かつ効率的なカスタマイズされたソフトラベル生成手法を構築すること。
• CNNおよびビジョン変換器（ViT）アーキテクチャの両方に対して、効果的なself-KDを可能にすること。
• 推論および学習のオーバーヘッドを最小限に抑えつつ、最先端の性能を達成すること。

提案手法

• 標準的なKD損失を、ターゲット損失（交差エントロピーに類似）と非ターゲット損失（交差エントロピー形式）に分解する。
• 非ターゲットスコアの合計を正規化することで、教師と学生の分布を均等化する正規化KD（NKD）を提案する。
• 学生モデル自身の特徴統計を用いて、教師モデルを必要とせずソフトラベルを生成する汎用自己知識蒸留（USKD）を導入する。
• 一般化を向上させるために、学生のターゲットスコアを二乗し、滑らかにしたものをソフトターゲットラベルとして使用する。
• 中間特徴のランクにZipfの法則を適用し、類似性に基づいて非ターゲットラベルを生成する。
• 弱い特徴と最終特徴を正規化して組み合わせ、非ターゲットクラスの順位付けに用いることで、ラベル生成の精度を向上させる。

実験結果

リサーチクエスチョン

• RQ1 標準的なKD損失をどのように再定式化すれば、蒸留におけるソフトラベルの利用効率を向上させられるか？
• RQ2 非ターゲットスコアを正規化することで、KDにおける学生と教師の分布をよりよく一致させられるか？
• RQ3 特にViTアーキテクチャにおいて、教師モデルを必要とせず、効果的なソフトラベルを生成できるか？
• RQ4 self-KDにおける学生のターゲットスコアのスムージング戦略が、性能に与える影響は何か？
• RQ5 特徴のランク（弱い特徴 vs. 最終特徴）の選択が、非ターゲットラベル生成の性能に与える影響は何か？

主な発見

• NKDはCIFAR-100およびImageNetで最先端の性能を達成し、ResNet-34を教師とする場合、ResNet18の精度を69.90%から71.96%まで向上させた。
• USKDはImageNetにおけるself-KDで新たな最先端を記録し、MobileNetでは1.17%、DeiT-Tinyでは0.55%の精度向上を達成した。
• 学生のターゲットスコアを$S_t + V_t - \text{mean}(S_t)$としてスムージングする手法は、ベースライン比で0.86%の精度向上をもたらした。
• 弱い特徴と最終特徴の両方を正規化して順位付けに用いることで、非ターゲットラベル生成の性能が最良となり、最終特徴のみを使用する場合に比べ0.09%の向上を達成した。
• USKDは追加の時間とメモリをほとんど必要としないため、実世界の展開において実用的である。
• 本手法はCNNおよびViTモデルの両方で有効であり、従来のself-KD手法がCNNに限定されていたという制限を克服した。