[论文解读] 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）损失和一个因概率尺度不匹配而扭曲优化的附加损失，重新思考了知识蒸馏。为解决此问题，作者提出了一种用于非目标分布知识的分布式损失，以及一种使用教师模型目标输出作为平滑标签的软损失，从而形成一种新型损失（NKD），在CIFAR-100和ImageNet上实现了最先进性能。无教师模型变体（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），以生成稳定软目标。
• 应用特定平滑函数 $ S_t + V_t - \text{mean}(S_t) $ 以稳定学生预测并提升tf-NKD中软目标的质量。

实验结果

研究问题

• RQ1 经典KD损失与原始交叉熵损失在优化动态方面有何关系？
• RQ2 为何标准KD损失导致次优优化？学生与教师概率尺度不匹配的根本原因是什么？
• RQ3 通过重构KD损失，能否更有效地传递非目标分布知识？
• RQ4 能否将教师的目标输出用作软标签以提升蒸馏性能？
• RQ5 是否可以通过使用自身平滑预测作为软目标，实现无需教师的高效学生模型训练？

主要发现

• 当使用ResNet-34作为教师时，所提出的NKD损失将ResNet-18在ImageNet上的Top-1准确率从69.90%提升至71.96%。
• 无教师模型变体tf-NKD在ResNet-18上实现了70.76%的ImageNet Top-1准确率，相比基线提升0.86%。
• 在CIFAR-100上，NKD损失在ResNet-18和ResNet-34作为教师时均达到80.76%的Top-1准确率，优于先前方法。
• tf-NKD变体采用平滑函数 $ S_t + V_t - \text{mean}(S_t) $ 在ImageNet上实现了最高0.86%的增益，甚至优于使用教师输出作为软目标。
• 可视化结果表明，平滑后的学生目标在训练各阶段更一致地反映真实类别隶属关系，提升了标签稳定性。
• 该方法无需额外训练时间，因为所有权重均在训练过程中由学生预测直接得出。