[论文解读] 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通过归一化非目标logits，提升软标签的利用效率；USKD则利用学生模型的特征和齐夫定律，为目标类与非目标类生成定制化软标签，实现对CNN和ViT模型的SOTA性能，且计算开销极低。

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与视觉Transformer（ViT）架构的有效自知识蒸馏。
• 在极低推理与训练开销下实现SOTA性能。

提出的方法

• 将标准KD损失分解为目标损失（类似于交叉熵）与非目标损失（以交叉熵形式表示）。
• 通过将非目标logits之和归一化，使教师与学生分布对齐，提出归一化知识蒸馏（NKD）。
• 提出通用自知识蒸馏（USKD），利用学生模型自身特征统计量生成软标签，无需教师模型。
• 将学生的目标logit进行平方与平滑处理，作为软目标标签，以提升泛化能力。
• 利用中间特征的排序应用齐夫定律，基于相似性生成软非目标标签。
• 通过归一化结合弱logits与最终logits，确定非目标类别在标签生成中的排序。

实验结果

研究问题

• RQ1 如何重构标准KD损失，以提升蒸馏过程中软标签的利用效率？
• RQ2 能否通过归一化非目标logits，使KD中学生与教师分布对齐更优？
• RQ3 能否在无教师模型的情况下，为self-KD生成有效软标签，特别是针对ViT架构？
• RQ4 对自知识蒸馏中学生目标logit的平滑策略选择，会产生何种影响？
• RQ5 在软非目标标签生成中，特征排序（弱logits vs. 最终logits）的选择如何影响性能？

主要发现

• NKD在CIFAR-100与ImageNet上均达到SOTA性能，将ResNet18准确率从69.90%提升至71.96%（使用ResNet-34教师模型）。
• USKD在ImageNet上为self-KD设立了新SOTA，使MobileNet准确率提升1.17%，DeiT-Tiny提升0.55%。
• 将学生目标logit平滑为$S_t + V_t - \text{mean}(S_t)$的方法，相比基线实现0.86%的准确率提升。
• 结合归一化的弱logits与最终logits进行排名赋值，用于非目标标签生成，性能最佳，相比仅使用最终logits提升0.09%。
• USKD所需额外时间与内存可忽略，适用于实际部署。
• 该方法在CNN与ViT模型上均有效，克服了以往self-KD方法仅适用于CNN的局限性。