CVPR 2022 | 解耦知识蒸馏
CVPR 2022 | 解耦知识蒸馏
|
近年来，SOTA蒸馏方法多基于网络中间层的深层特征，而基于logit的KD(知识蒸馏)则被忽略。 旷视科技、早稻田大学、清华大学的研究者研究了传统KD的局限之处后，给logit蒸馏的研究提出了新的方法。 作者证明传统KD的损失函数存在耦合，该耦合限制了logit蒸馏的性能。 作者提出将耦合的损失函数解耦成TCKD(target class KD)和NCKD(non-target class KD)的加权和。 这样，1)释放了NCKD部分对模型性能的影响； 2) 解耦了TCKD和NCKD，可以分别配置权重，平衡二者对模型性能的影响。 通过一系列的实验，作者也证明了本文的DKD(Decoupled Knowledge Distillatio) 方法可以在更高的训练效率下达到相当于甚至优于SOTA蒸馏方法的效果。 |
论文地址：
https://arxiv.org/abs/2203.08679
代码地址：
https://github.com/megvii-research/mdistiller
一、本文方法
按理说，logits在语义水平上要比深层特征更高级，那么logit蒸馏应该是可以达到与特征蒸馏相当的水平，但是现实却是：logit蒸馏虽然在算力和存储消耗上有优势，但是性能却远远比不上特征蒸馏。
是不是有什么“暗黑力量”封禁了logit的“洪荒之力”？
那还得从KD的机制出发，反思一下，有没有什么不合理的地方。
分类预测可以被看成两个层次：1）目标类+所有非目标类的二元预测； 2）每个非目标类的多元预测。
传统的KD训练机制如下图：
作者通过详细的推导(原文中)，将损失函数整理为：
其中T,S分别表示教师和学生网络；b表示目标类的二元概率； $p_t$ 表示目标类概率；$\hat{p}$ 表示非目标类的多元概率。
从上式可以看出，KD损失函数被重构成两部分：
$KL(bT||bS)$ 代表教师和学生网络的目标类二元概率的相似性，命名为TCKD；$KL(\hat{p}T||\hat{p}S)$ 代表教师和学生网络非目标类多元概率的相似性，命名为NCKD。所以KD损失函数可以表示为：
可以看出，NCKD的权重被 $p_t^T$ 耦合。
这种耦合是否正是抑制NCKD“洪荒之力”的“暗黑力量” ？
作者通过实验研究了TCKD和NCKD两个部分的作用，以及耦合的影响：
上表可以看出一个有趣的现象：只用NCKD一个部分时的性能甚至比baseline还要好，足以说明NCKD部分对于整个网络的重要性，“洪荒之力”确实被耽误了！
上表也证明了NCKD在预测情况较好的情况下作用更大，应该有更大的权重才对，但是从上边的损失函数公式可以看出，现实是NCKD的权重却随着 $p_t^T$的增大而减小
但是到此处还看不出TCKD的明显优势，作者从上述公式推导过程推测，TCKD部分表征的是识别训练样本的难度，并通过实验证明了这一推测。
作者通过对数据集加入AutoAugment、噪音以及用更具有挑战性的数据集等方法来提升训练难度，结果如下：
这一部分实验证明了TCKD在训练难度提升情况下的有效作用。
研究到此，将损失公式重构解耦势在必行。
解耦的好处有两点：
1）将NCKD从$(1-p_t^T)$ 中解耦，释放其“洪荒之力”；
2）将TCKD和NCKD两个部分解耦，分别配置权重，平衡二者的作用。
作者将新的logit蒸馏命名为DKD(Decoupled Knowledge Distillatio)
损失函数如下：
这里的$\alpha$和$\beta$ 是可以自行调整的超参数，用于平衡TCKD和NCKD两个部分的重要性。
二、实验分析
1 解耦对于性能的提升
2 图像分类-CIFAR100
3 图像分类-ImageNet
4 目标检测-MS-COCO
logits无法提供目标定位的知识，所以单独将DKD用于目标检测效果不佳，但是DKD可以辅助用于目标定位的模型，提升其性能。
5 DKD vs SOTA --训练效率
6 更大的教师网络
更大的教师网络与小的学生网络之间的容量差增大，反而会使蒸馏效果变差。本文DKD方法通过解耦，可以缓解这个问题。
7 特征转化性能
8 可视化
三、总结
本文给了logit蒸馏一种新的解释，将传统的KD损失重构成两部分--TCKD和NCKD。两部分的重要性都得到证明，耦合的KD方式限制了知识迁移的有效性的灵活性也被证明。
本文提出的解耦蒸馏(DKD)被用于不同数据集，用于完成图像分类和目标检测任务，效果超越了传统logitKD，达到甚至超越特征蒸馏的SOTA方法。
