[论文解读] 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损失，可同时提升特征幅值与对齐度，在结合现有知识蒸馏方法时，于ImageNet、CIFAR-100和COCO基准上实现了最先进（SOTA）的准确率。

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}.

研究动机与目标

• 为解决现有知识蒸馏方法直接对齐学生特征与教师特征，但未考虑特征幅值或分类器结构的局限性。
• 通过利用教师模型的类别均值特征作为强而内在的分类器，指导特征学习，从而提升学生模型性能。
• 探究大范数特征在知识蒸馏中的作用，受模型剪枝与域适应研究发现的启发。
• 开发一种简单、可插拔的损失函数（ND损失），同时对特征范数与方向进行正则化，以提升蒸馏效果。

提出的方法

• 该方法从教师模型的预训练特征中构建类别均值特征，并将其作为强分类器，用于指导学生特征的学习。
• 提出一种新颖的ND损失，联合优化两个目标：(1) 最大化学生特征的L2范数；(2) 最小化学生特征与对应教师类别均值特征之间的余弦距离。
• ND损失应用于倒数第二层（嵌入层），该层特征用于计算logits，因此与分类性能直接相关。
• 该方法兼容logits蒸馏与特征蒸馏框架，可无缝集成至现有知识蒸馏流程中。
• 在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 尽管存在架构差异，ND损失在从Transformer蒸馏到CNN时是否仍能提升性能？

主要发现

• 将ND损失应用于标准KD（即KD++），在使用ResNet-18学生与ResNet-50教师时，ImageNet上的top-1准确率达到72.53%，优于原始KD（71.35%）及其他近期方法如ReviewKD（71.10%）和DKD（71.87%）。
• 在ResNet-18学生模型上训练时，KD++相比KD实现1.45%的绝对准确率提升，证明其强大有效性。
• 当从更大容量的教师模型（如ResNet-34至ResNet-152）蒸馏时，KD++表现出一致的性能提升（如71.98% → 72.54%），而基线KD方法则出现性能下降或波动。
• 在COCO目标检测任务中，使用ND损失可提升检测精度，证实其优势不仅限于图像分类任务。
• 可视化结果表明，KD++生成更高范数的特征并实现更好的类别分离，成功恢复了标准KD与朴素学生模型所遗漏的误分类类别（如紫色类别）。
• 即使单独与交叉熵或KL损失结合使用，ND损失仍能将准确率提升约1%，且计算成本可忽略不计，证明其有效性。

(a) teacher (Res56)