[论文解读] Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation

Taehyeon Kim, Jaehoon Oh+3 more|arXiv (Cornell University)|May 19, 2021

Advanced Neural Network Applications参考文献 21被引用 13

一句话总结

本文提出用教师与学生logit向量之间的均方误差（MSE）损失替代知识蒸馏中标准的Kullback-Leibler（KL）散度损失，实证表明通过MSE进行直接logit匹配优于KL散度，在多个基准测试中表现更优。关键贡献在于MSE能够实现更精确的logit匹配，从而提升泛化性能，尤其当学生模型小于教师模型时效果更显著。

ABSTRACT

Knowledge distillation (KD), transferring knowledge from a cumbersome teacher model to a lightweight student model, has been investigated to design efficient neural architectures. Generally, the objective function of KD is the Kullback-Leibler (KL) divergence loss between the softened probability distributions of the teacher model and the student model with the temperature scaling hyperparameter tau. Despite its widespread use, few studies have discussed the influence of such softening on generalization. Here, we theoretically show that the KL divergence loss focuses on the logit matching when tau increases and the label matching when tau goes to 0 and empirically show that the logit matching is positively correlated to performance improvement in general. From this observation, we consider an intuitive KD loss function, the mean squared error (MSE) between the logit vectors, so that the student model can directly learn the logit of the teacher model. The MSE loss outperforms the KL divergence loss, explained by the difference in the penultimate layer representations between the two losses. Furthermore, we show that sequential distillation can improve performance and that KD, particularly when using the KL divergence loss with small tau, mitigates the label noise. The code to reproduce the experiments is publicly available online at https://github.com/jhoon-oh/kd_data/.

研究动机与目标

• 探究KL散度损失中的温度缩放超参数τ如何影响知识蒸馏的性能。
• 解决当前对软标签为何能提升知识蒸馏泛化性能的理论与实证理解不足的问题。
• 提出一种基于MSE损失的直接logit匹配方法，作为替代KL散度的改进方案，以提升学生模型性能。
• 评估不同损失函数在标签噪声设置下的鲁棒性。
• 探索序列蒸馏策略，以在教师与学生模型容量差距较大时提升性能。

提出的方法

• 理论分析表明，当τ → ∞时，KL散度损失近似于logit匹配；当τ → 0时，其近似于标签匹配。
• 提出一种基于学生与教师logit向量之间MSE损失的直接logit匹配框架，绕过软概率分布的使用。
• 采用可视化技术（Müller et al., 2019）对比使用KL与MSE损失训练的模型在倒数第二层特征表示的差异。
• 通过先使用KL损失（τ = 3）训练，再用MSE损失微调的方式实施序列蒸馏，以改善大容量差距下的性能。
• 通过在CIFAR-100中随机翻转标签，评估标签噪声下的鲁棒性，并测量不同损失函数下的测试准确率。
• 使用GitHub上公开的代码复现实验，以确保可复现性。

实验结果

研究问题

• RQ1 KL散度损失中的温度缩放超参数τ如何影响知识蒸馏的泛化行为？
• RQ2 通过MSE损失实现的直接logit匹配是否能在知识蒸馏中超越KL散度损失？
• RQ3 使用KL散度与MSE损失训练的模型在倒数第二层表示学习方面有何差异？
• RQ4 序列蒸馏（先使用KL损失，后使用MSE损失）是否能在教师与学生模型容量差距较大时提升性能？
• RQ5 在标签噪声设置下，不同损失函数的表现如何，特别是在缓解对错误标签的过拟合方面？

主要发现

• 在所有评估的数据集中，MSE损失的表现均优于KL散度损失，在CIFAR-100上达到75.78%的测试准确率，而KL损失（τ=20）为75.47%。
• 当使用相同的教师与学生架构时，MSE损失的性能高于任何KL散度配置，包括τ=20。
• KL散度损失（τ=20）达到75.47%的准确率，而MSE损失达到75.78%，表明直接logit匹配比软标签匹配更有效。
• 序列蒸馏（先KL后MSE）在大容量差距下能提升性能，优于单独使用任一损失的直接蒸馏。
• 在40%对称噪声的标签噪声设置下，τ=0.5的KL损失表现最佳，优于MSE和高τ的KL损失，原因在于其隐式地实现了标签平滑。
• 使用MSE损失训练的模型比使用KL损失训练的模型更忠实地保留了教师模型的倒数第二层表示，后者因拉长表示而造成失真。