【论文原文】:Decoupled Knowledge Distillation
获取地址:https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=9879819
博主关键词: 解耦,目标知识蒸馏,非目标知识蒸馏
0. Abstract
目前的蒸馏方法主要基于从中间层蒸馏深层特征,而logit蒸馏的意义被大大忽视。为了为研究logit蒸馏提供一个新的观点,我们将经典的KD损失重新表述为两部分,即目标类知识蒸馏(TCKD)和非目标类知识蒸馏(NCKD)。我们通过实证研究和证明了这两部分的影响:TCKD转移了关于训练样本的“困难”的知识,而NCKD是logit蒸馏工作有效的主要原因。更重要的是,我们揭示了经典的KD损失是一个耦合公式,它(1)抑制了NCKD的有效性,(2)限制了平衡这两部分的灵活性。为了解决这些问题,我们提出了解耦知识蒸馏(DKD),使TCKD和NCKD能够更有效和灵活地发挥它们的作用。与复杂的基于特征的方法相比,我们的DKD在CIFAR-100、ImageNet和MSCOCO数据集上实现了图像分类和目标检测任务的结果,并具有更好的训练效率。本文证明了logit蒸馏的巨大潜力,希望它对今后的研究有帮助。该代码可以在https://github.com/megviiresearch/ mdistiller上找到。
1. Introduction
在过去的几十年里,深度神经网络(DNN)使计算机视觉领域发生了革命性的变化,它成功地促进了各种真实场景任务,如图像分类[9,13,21]、目标检测[8,27]和语义分割[31,45]。然而,强大的网络通常受益于较大的模型容量,这引入了较高的计算和存储成本。在广泛部署轻量级模型的工业应用程序中,这种成本并不可取。在文献中,降低成本的一个潜在方向是知识蒸馏(KD)。KD代表了一系列聚焦于将知识从重模型(教师)转移到轻模型(学生)的方法,其可以在不引入额外成本的情况下提高轻模型的性能。
KD的概念首先在[12]中提出,通过最小化教师和学生的预测logits之间的Kl-div来转移知识(图 1a)。自[28]以来,研究主要集中在从中间层的深层特征中提取知识。与基于logits的方法相比,特征蒸馏在各种任务上的性能更为优越,因此对logit蒸馏的研究几乎不深入。然而,基于特征的方法的训练成本并不令人满意,因为在训练期间引入了额外的计算和存储使用(如网络模块和复杂的操作)来提取深度特征。
logit蒸馏需要边际的计算和存储成本,但性能较差。直观地说,logit蒸馏应该达到与特征蒸馏相当的性能,因为logit比深度特征处于更高的语义级别。我们认为,logit蒸馏的潜力受到未知原因的限制,导致性能不理想。为了重振基于logits的方法,我们通过深入研究KD的机制来开始这项工作。首先,我们将一个分类预测分为两个层次:(1)对目标类和所有非目标类的二值预测,(2)对每个非目标类的多类别预测。在此基础上,我们将经典的KD损失[12]重新表述为两部分,如图1b所示。一种是针对目标类的二进制logit蒸馏,另一种是针对非目标类的多类logit蒸馏。为了简化问题,我们将它们分别命名为目标分类知识蒸馏(TCKD)和非目标分类知识蒸馏(NCKD)。重新制定的方案使我们能够独立地研究这两部分的影响。
Figure 1. 经典的KD[12]和我们的DKD的插图。我们将KD重新表示为两部分的加权和,即TCKD和NCKD。第一个方程表明KD (1)将NCKD与ptT(教师对目标类别的置信度)耦合,(2)将两部分的importance耦合。此外,我们还证明了第一个耦合抑制了有效性,第二个耦合限制了知识转移的灵活性。我们提出DKD来解决这些问题,它为TCKD使用超参数α,为NCKD使用超参数β,一举两得。
TCKD通过二进制logit蒸馏来传递知识,这意味着只提供了对目标类的预测,而对每个非目标类的具体预测都是未知的。一个合理的假设是,TCKD转移了关于训练样本的“难度”的知识,即,该知识描述了识别每个训练样本的困难程度。为了验证这一点,我们从三个方面设计实验来增加训练数据的“难度”,即更强的增强、更嘈杂的标签和固有具有挑战性的数据集。
NCKD只考虑非目标logits之间的知识。有趣的是,我们通过经验证明,只有应用NCKD才能获得与经典KD相当甚至更好的结果,这表明非目标logits中包含的知识至关重要,这可能是突出的“暗知识”。
更重要的是,我们的重新表述证明了经典的KD损失是一个高度耦合的公式(如图 1b所示),这可能是为什么logit蒸馏的潜力有限的原因。首先,NCKD损失项由一个与教师对目标类别的预测置信度呈负相关的系数进行加权。因此,预测分数越大,权重就会越小。耦合显著抑制了NCKD对预测良好的训练样本的影响。这种抑制并不可取,因为教师对训练样本越有信心,它就能提供越可靠和越有价值的知识。其次,TCKD和NCKD的significance是耦合的,即不允许分别加权TCKD和NCKD。这种限制并不可取,因为TCKD和NCKD应该单独考虑,因为它们的贡献来自不同的方面。
为了解决这些问题,我们提出了一种灵活且高效的logit蒸馏方法,称为解耦知识蒸馏(DKD,图1b)。DKD用常数值代替NCKD损失与与教师可信度负相关的系数解耦,提高了对良好预测样本的蒸馏效率。同时,NCKD和TCKD也被解耦了,因此可以通过调整每个部分的权重来分别考虑它们的重要性。
总的来说,我们的贡献总结如下:
- 我们通过将经典的KD划分为TCKD和NCKD来为研究logit蒸馏提供了一个深刻的观点。并分别对这两部分的影响效果进行了分析和验证。
- 我们揭示了由其高度耦合公式引起的经典KD损失的局限性。NCKD与教师预测分数的耦合会抑制知识转移的有效性。TCKD与NCKD的耦合限制了平衡这两个部分的灵活性。
- 我们提出了一种有效的logit蒸馏方法DKD来克服这些限制。DKD在各种任务上实现了最先进的性能。与基于特征的蒸馏方法相比,DKD具有更高的训练效率和更好的特征可转移性。
2. Related work
知识蒸馏(KD)的概念是由Hinton等人在[12]一书中首次提出的。KD定义了一种学习方式,使用一个更大的教师网络来指导训练一个较小的学生网络的许多任务[12,17,18]。“暗知识”通过老师的软标签转移到学生身上。为了提高人们对negative logits的关注,超参数温度被引入。以下工作可分为两种类型,logit[3,6,22,40,44 ]蒸馏和中间特征[10,11,14,15,23,25,28,33,34,41,43]蒸馏。
以往的logit蒸馏工作主要集中在提出有效的正则化和优化方法,而不是提出新的方法。DML [44]提出了一种相互学习的方式来同时训练学生和教师。TAKD [22]引入了一个名为“教师助理”的中等规模的网络来弥合教师和学生之间的差距。此外,一些工作也集中在解释经典的KD方法[2,26]。
目前最先进的方法主要基于中间特征蒸馏,它可以直接将表示(知识)从教师转移到学生[10,11,28],或者将在教师中捕获的样本之间的相关性转移到学生[23,33,34]。大多数基于特征的方法可以获得更好的性能(明显高于基于logits的方法),但涉及到相当高的计算和存储成本。
本文主要分析了是什么限制了基于logits的方法和振兴logit蒸馏的潜力。
3. Rethinking Knowledge Distillation
在本节中,我们将深入研究知识蒸馏的机制。我们将KD损失重新表述为两部分的加权和,一部分与目标类相关,另一部分不与目标类相关。我们探讨了知识蒸馏框架中各部分的影响,并揭示了经典KD的一些局限性。受研究结果的启发,我们进一步提出了一种新的logit蒸馏方法,在各种任务上都取得了显著的性能。
3.1. Reformulating KD
Notations.
对于来自第t类的一个训练样本,分类概率可以记为 p = [p1,p2,…,pt,…,pC] ∈ R1✖C,其中pi是第i类的概率,C是类别总数,p中的每个元素都可以通过softmax函数得到:
其中zi表示第i类的logit。
为了分离与目标类相关和无关的预测,我们定义了以下符号。b=[pt,p\t]∈ R1✖2表示目标类(pt)和所有其他非目标类(p\t)的二进制概率,其计算方法为:
同时,我们将 p = [p1,…,pt-1,pt+1,…,pC] ∈ R1✖C-1声明为非目标类之间的独立建模概率(即不考虑第t类)。每个元素的计算方法如下:
Reformulation.
在这部分中,我们试图用二进制概率b和非目标类的概率ˆp来重新表示KD。T和S分别表示老师和学生。经典的KD使用KL-div作为损失函数,可以写为:
根据Eqn.(1)和Eqn.(2)我们有因此我们可以将Eqn.(3)重写为:
那么,Eqn.(4)可以被重写为:
如Eqn.(5)所示。KD损失被重新表述为两项的加权和。KL(bT||bS)表示教师和学生的目标类的二进制概率之间的相似性。因此,我们将其命名为目标类知识蒸馏(TCKD)。同时,KL(ˆpT||ˆpS)表示教师和学生在非目标类间的概率的相似性,称为非目标类知识蒸馏(NCKD)。Eqn.(5)可重写为:
显然,NCKD的权重与ptT是耦合的。
上面的重新表述启发我们研究TCKD和NCKD的个体效应,这将揭示经典耦合公式的局限性。
3.2. Effects of TCKD and NCKD
Performance gain of each part.
我们分别研究了TCKD和NCKD对CIFAR-100 [16]的影响。选择ResNet [9]、WideResNet(WRN)[42]和ShufflfleNet [21]作为训练模型,其中考虑了相同和不同的体系结构。
实验结果见表1。对于每一对师生,我们展示了(1)学生基线(venilla普通 训练),(2)经典的KD(其中TCKD和NCKD都被使用),(3)单一TCKD,(4)单一NCKD的结果。每个损失的权重设为1.0(包括默认的交叉熵损失)。其他实现细节与第4节中的实现细节相同。
直观地说,TCKD专注于与目标类相关的知识,因为相应的损失函数只考虑二进制概率。相反,NCKD侧重于非目标类之间的知识。我们注意到,单独应用TCKD可能对学生是无用的(例如,TCKD-v1增加0.02%和0.12%),甚至是有害的(例如,WRN-16-2下降2.30%,ResNet8×4下降3.87%)。然而,NCKD的蒸馏性能与经典的KD相当,甚至更好(例如,在ResNet8×4上为1.76% vs. 1.13%)。
消融结果表明,目标类相关的知识没有非目标类相关的知识重要。为了深入研究这一现象,我们提供了以下进一步的分析。
TCKD transfers the knowledge concerning the“difficulty”of training samples.
根据Eqn.(5),TCKD通过二值分类任务转移“暗知识”,这可能与样本的“难度”有关。例如,一个使用ptT = 0.99的训练样本与另一个使用ptT = 0.75的训练样本相比,学生可以“更容易”学习。由于TCKD传达了训练样本的“难度”,我们假设当训练数据变得具有挑战性时,其有效性就会被揭示出来。然而,CIFAR-100训练集很容易适合。因此,教师提供的“困难”的知识是没有信息的。
在这一部分中,我们从三个角度进行了实验来验证:训练数据越困难,TCKD所能提供的好处就越多。
(1)应用Strong Augmentation技术是增加训练数据难度的一种直接的方法。我们在CIFAR-100数据集上使用自动增强[5]训练了ResNet32×4模型,达到了81.29%的top-1验证精度。对于学生,我们训练有/没有TCKD的ResNet8×4和shuffleNetV1模型。表2中的结果显示,如果应用Strong Augmentation,TCKD将获得显著的性能收益。
(2)噪声标签也会增加训练数据的难度。我们在CIFAR-100上,将ResNet32×4模型作为教师,ResNet8×4作为学生,对称噪声比为{0.1,0.2,0.3},遵循[7,35]。如表3所示,结果表明,TCKD在噪声更大的训练数据上获得了更多的性能提升。
(3)还考虑了具有挑战性的数据集(例如,ImageNet [29])。结果表明,TCKD在ImageNet上带来0.32%的性能提高。
最后,我们通过实验各种策略来增加训练数据的难度(如强增强、噪声标签、困难任务),证明了TCKD的有效性。结果表明,在对更具挑战性的训练数据上提取知识时,关于训练样本的“难度”的知识可能更有益。
NCKD is the prominent reason why logit distillation works but is greatly suppressed.
有趣的是,我们在表1中注意到,当只应用NCKD时,其性能与经典的KD相当,甚至更好。这表明,非目标类之间的知识对logit蒸馏至关重要,这可以是突出的“暗知识”。然而,通过回顾Eqn.(5),我们注意到NCKD的损失与(1−ptT)相结合,其中ptT代表教师对目标类的预测分数。因此,更高的预测分数会导致更小的NCKD权重。我们假设教师在训练样本中的表现越大,它就能提供越可靠和越有价值的知识。然而,这种自信的预测高度抑制了损失的权重。我们假设这一事实会限制知识转移的有效性,这首先是由于我们在Eqn.(5)中对KD的重新表述而被研究的。
我们设计了一个消融实验来验证良好预测的样本确实比其他样本传递更好的知识。首先,我们根据ptT对训练样本进行排序,并将其平均分成两个子集。为了清晰起见,一个子集包括ptT为前50%的样本,而剩余的样本在另一个子集中。然后,我们在每个子集上用NCKD训练学生网络,以比较性能增益(而交叉熵损失仍然在整个集合上)。
从表5可以看出,对前50%的样本使用NCKD,获得了更好的性能,表明预测良好的样本的知识比其他样本更丰富。然而,教师的高预测置信度抑制了预测良好的样本的损失权重。
3.3. Decoupled Knowledge Distillation
到目前为止,我们已经将经典的KD损失重新表述为两个独立部分的加权和,并进一步验证了TCKD的有效性,并揭示了NCKD的抑制。具体来说,TCKD传递了关于训练样本的“难度”的知识。TCKD可以在更具挑战性的训练数据上获得更显著的改进。NCKD在非目标类之间转移知识,在权重(1−ptT)相对较小的情况下,这些知识会被抑制。
本能地说,TCKD和NCKD都是不可或缺的和关键的。然而,在经典的KD公式中,TCKD和NCKD是从以下几个方面耦合起来的:
- 首先,NCKD与(1−ptT)相结合,这会抑制预测良好的样本上的NCKD。由于表5中的结果显示,预测良好的样本可以带来更多的性能增益,耦合形式可能会限制NCKD的有效性。
- 另一方面,NCKD和TCKD的权重在经典的KD框架下进行耦合。不允许改变每个术语的权重来平衡其重要性。
由于TCKD和NCKD的贡献来自不同的方面,我们认为应该加以考虑。
得益于我们对KD的重新制定,我们提出了一种新的logit蒸馏方法,即解耦知识蒸馏(DKD)来解决上述问题。我们提出的DKD在一个解耦的公式中独立地考虑了TCKD和NCKD,如图1b所示。具体来说,我们引入了两个超参数α和β,分别作为TCKD和NCKD的权重。DKD的损失函数可写成如下:
在DKD中,抑制NCKD有效性的权重(1−ptT)被β所取代。更重要的是,我们可以调整α和β,以平衡TCKD和NCKD的重要性。通过解耦NCKD和TCKD,DKD为logit蒸馏提供了一种高效、灵活的方式。算法1以类似于pytorch[24]的风格提供了DKD的伪代码。
4. Experiments
我们主要实验了两个具有代表性的任务,即图像分类和目标检测,包括:
CIFAR-100 [16]是一个著名的图像分类数据集,包含100个类别的32像素×32像素的图像。训练和验证集由5万张和1万张图像组成。
ImageNet [29]是一个由1000个类组成的大规模分类数据集。训练集包含128万张图像,验证集包含5万张图像。
MS-COCO [20]是一个80类的通用目标检测数据集。train2017包含118k张图像,val2017包含5k张图像。
由于页面的限制,所有的实现细节都附在补充文件中。
4.1. Main Results
首先,我们分别演示了由解耦(1) NCKD和ptT以及(2) NCKD和TCKD所贡献的改进。然后,我们在图像分类和目标检测任务上对该方法进行了基准测试。
Ablation:α和β.
下面两个表报告了不同α和β的学生准确率(%)。ResNet32×4和ResNet8×4分别设置为老师和学生。首先,我们在第一个表中证明了解耦(1−ptT)和NCKD可以带来合理的性能增益(73.63% vs. 74.79%)。然后,我们证明了NCKD和TCKD的解耦权重可能有助于进一步的改进(74.79% vs. 76.32%)。此外,第二个表表明TCKD是不可缺少的,在不同α下,TCKD的提升稳定在1.0左右。
CIFAR-100图像分类。我们讨论了在CIFAR-100上的实验结果来检验我们的DKD。验证精度见表6和表7。表6包含了教师和学生具有相同网络架构的结果。表7显示了教师和学生来自不同系列的结果。
值得注意的是,与baseline和经典的KD相比,DKD在所有的师生对上都取得了一致的改进。我们的方法对相同系列和不同系列的师生对分别取得了1~2%和2~3%的改进。它强烈支持DKD的有效性。此外,DKD实现了与基于特征的蒸馏方法相比相当甚至更好的性能,显著提高了蒸馏性能和训练效率之间的权衡,这将在第4.2节中进一步讨论。
ImageNet image classifification.
表8和表9报告了ImageNet上图像分类准确率的top-1和top-5。我们的DKD取得了显著的改进。值得一提的是,DKD的性能优于最先进的特征蒸馏方法的结果。
MS-COCO object detection.
正如在之前的工作中所讨论的,目标检测任务的性能在很大程度上取决于定位感兴趣目标的深度特征的质量。这一规则也适用于在检测器[17,37]之间传递知识,即特征模拟是至关重要的,因为logits不能为目标定位提供知识。如表10所示,单应用DKD很难取得优异的性能,但预期超过经典的KD。因此,我们引入了基于特征的蒸馏方法,ReviewKD[1],以获得满意的结果。可以观察到,我们的DKD可以进一步提高AP指标,即使ReviewKD的蒸馏性能也相对较高。最后,通过将我们的DKD与基于特征的蒸馏方法相结合,得到了新的最先进的结果。
4.2. Extensions
为了更好地理解DKD,我们从四个角度提供了Extensions。首先,我们综合比较了DKD与具有代表性的先进方法的训练效率。然后,我们提供了一个新的视角来解释为什么更大的模型并不总是更好的教师,并通过利用DKD来缓解这个问题。此外,在[33]之后,我们研究了由DKD学习到的深度特征的可转移性。我们还提出了一些可视化方法来验证DKD的优越性。
Training effificiency.
我们评估了最先进的蒸馏方法的训练成本,证明了DKD的高训练效率。如图2所示,我们的DKD在模型性能和训练成本(例如,训练时间和额外参数)之间实现了最好的权衡。由于DKD是由经典的KD重新表述出来的,因此它需要与KD几乎相同的计算复杂度,当然也不需要额外的参数。然而,基于特征的蒸馏方法需要额外的训练时间来蒸馏中间层特征,以及GPU内存成本。
Figure 2. 训练时间(每批)vs. CIFAR-100的准确性。我们设置ResNet32×4为老师,ResNet8×4为学生。该表显示了每种方法的额外参数的数量。
Improving performances of big teachers.
我们关于更大的模型并不总是更好的教师问题提供了一个新的潜在的解释。具体来说,期望更大的教师,但不能转移更多有益的知识,甚至取得比小教师更差的表现。
[3,36]之前的工作用大教师与小学生之间的大容量差距来解释这一现象。然而,我们认为主要原因是NCKD的抑制,即(1−ptT)会随着教师的增大而变小。
更重要的是,关于这个问题的相关工作也可以从这个角度来解释,例如,ESKD [3]使用早期停止的教师模型来缓解这个问题,这些教师会收敛不足,产生更小的ptT。
为了验证我们的猜想,我们在一系列的教师模型上进行了DKD测试。表11和表12中的实验结果一致表明,我们的DKD缓解了较大的模型并不总是更好的教师问题。
Feature transferability.
我们进行了实验来评估深度特征的可转移性,以验证我们的DKD转移了更一般化的知识。在[33]之后,我们使用从WRN-40-2中提取出来的WRN-16-2作为特征提取器。然后在下游数据集上执行线性检测任务,即STL-10 [4]和Tiny-ImageNet6。
结果报告在表13中,证明了用我们的DKD学习到的特征具有突出的可转移性。实施细节见补充部分。
Visualizations.
我们从两个角度呈现可视化(在CIFAR-100上将老师设置为ResNet32x4,学生设置为ResNet8x4)。(1)t-SNE(图3)结果表明,DKD的表示比KD更可分离,证明DKD有利于深度特征的可鉴别性。(2)我们还可视化了学生和教师logits的相关矩阵的差异(图4)。与KD相比,DKD帮助学生输出与老师更相似的logits,即获得更好的蒸馏性能。
5. Discussion and Conclusion
本文通过将经典的KD损失重新定义为目标类知识蒸馏(TCKD)和非目标类知识蒸馏(NCKD)两部分,提出了一种新的观点来解释logit蒸馏。对这两部分的影响分别进行了研究和证明。更重要的是,我们揭示了KD的耦合公式限制了知识转移的有效性和灵活性。为了克服这些问题,我们提出了解耦知识蒸馏技术(DKD),该方法对CIFAR-100、ImageNet和MS-COCO数据集进行了显著改进,用于图像分类和目标检测任务。此外,还证明了DKD在训练效率和特征可转移性方面的优越性。希望本文对今后的logit蒸馏研究有贡献。
limitations and future works.
所讨论的明显的限制如下。DKD在目标检测任务上不能优于最先进的基于特性的方法(例如,ReviewKD [1]),因为基于logits的方法不能传递关于定位的知识。此外,我们还在我们的附录中提供了一个关于如何调整β的直观指导。然而,蒸馏性能与β之间的严格相关性还没有得到充分的研究,这将是我们未来的研究方向。