CrossKD: Cross-Head Knowledge Distillation for Dense Object Detection

CrossKD：用于密集目标检测的交叉头知识蒸馏

论文链接：https://arxiv.org/abs/2306.11369v2

项目链接：https://github.com/jbwang1997/CrossKD

Abstract

知识蒸馏(Knowledge Distillation, KD)是一种有效的学习紧凑目标检测器的模型压缩技术。现有最先进的KD检测方法大多基于特征模仿。在本文中，我们提出了一种通用且有效的预测模拟蒸馏方案，称为CrossKD，它将学生检测头的中间特征传递给教师检测头。由此产生的十字预测，然后被迫模仿老师的预测。这种方式使学生头从注释和老师的预测中接收到相互矛盾的监督信号，大大提高了学生的检测性能。此外，由于模仿老师的预测是KD的目标，与特征模仿相比，CrossKD提供了更多面向任务的信息。在MS COCO上，仅应用预测模拟损失，我们的CrossKD将GFL ResNet-50在1倍训练计划下的平均精度从40.2提高到43.7，优于所有现有的KD方法。此外，我们的方法也可以很好地用于异构主干检测器的提取。

1. Introduction

知识蒸馏(Knowledge Distillation, KD)作为一种模型压缩技术，在目标检测中得到了深入的研究[5,13,29,31,56,60,61,74,75]，近年来取得了优异的成绩。根据检测器的蒸馏位置，现有的KD方法大致可分为预测模仿和特征模仿两类。预测模仿(见图1(a))最早是在[24]中提出的，它指出，教师预测的平滑分布比ground truth的狄拉克分布更有利于学生学习。换句话说，预测模仿迫使学生模仿老师的预测分布。不同的是，特征模仿(见图1(b))遵循了FitNet[54]提出的思想，该思想认为中间特征包含比教师预测更多的信息。它旨在加强师生对之间的特征一致性。

预测模拟在提取目标检测模型中起着至关重要的作用。然而，长期以来，人们观察到它比特征模仿效率低。最近，Zheng等[74]提出了一种定位蒸馏(localization distillation, LD)方法，通过转移定位知识来改进预测模拟，将预测模拟推向了一个新的高度。尽管刚刚赶上了先进的特征模仿方法，例如PKD[5]，LD表明预测模仿具有转移特定任务知识的能力，这使学生从正交方面受益于特征模仿。这促使我们进一步探索和改进预测模拟。

通过调查，我们观察到传统的预测模仿可能会受到学生分配者的真实目标与教师预测的升华目标之间的冲突的影响。当用预测模拟训练检测器时，学生的预测被迫同时模拟真实目标和老师的预测。然而，老师预测的蒸馏目标通常与分配给学生的真实目标有很大的差异。如图2(a)和图2(b)所示，教师在绿色圈出的区域生成班级概率，这与分配给学生的真实目标相冲突。因此，学生检测器在蒸馏过程中经历了一个矛盾的学习过程，严重干扰了优化。

为了缓解上述冲突，以往的预测模拟方法[13,19,74]倾向于在师生差异中间的区域内进行蒸馏。然而，我们认为高度不确定的区域通常容纳更多对学生有益的信息。本文提出了一种新的交叉头知识蒸馏管道，简称CrossKD。如图1©所示，我们建议将中间特征从学生的头部馈送到教师的头部，从而产生交叉头部预测。然后，在新的十字预测和老师的预测之间进行KD运算。

尽管它很简单，但CrossKD提供了以下两个主要优点。首先，由于交叉头预测和教师的预测都是通过共享教师的部分检测头产生的，因此交叉头预测与教师的预测相对一致。这缓解了师生对之间的差异，提高了预测模拟训练的稳定性。此外，由于模仿老师的预测是KD的目标，所以相对于特征模仿，CrossKD在理论上是最优的，可以提供更多面向任务的信息。这两个优点使我们的CrossKD能够有效地从教师的预测中提取知识，从而产生比以前最先进的特征模仿方法更好的性能。

我们的方法可以显著提高学生检测器的性能，实现更快的训练收敛。本文在COCO[40]数据集上进行了综合实验，阐述了CrossKD的有效性。具体来说，在只应用预测模拟损失的情况下，CrossKD在GFL上实现了43.7 AP，比基线高3.5 AP，超过了之前所有最先进的目标检测KD方法。此外，实验还表明我们的CrossKD与特征模仿方法是正交的。通过将CrossKD与PKD等最先进的特征模仿方法相结合[5]，我们进一步在GFL上实现了43.9 AP。此外，我们还证明了我们的方法可以用于具有异构主干的检测器的提取，并且性能优于其他方法。

2. Related Work

2.1. 目标检测

目标检测是最基本的计算机视觉任务之一，它需要同时识别和定位目标。现代目标探测器可简单分为两类:一级[3、10、11、35、39、51、57、70]探测器和两级[8、17、18、20、21、38、52、58、73]探测器。其中，一级探测器，又称密集探测器，由于其出色的速度精度权衡，已成为检测的主流趋势。自YOLOv1以来，密集目标探测器受到了极大的关注[49]。通常，YOLO系列检测器[2,16,44,49 - 51]试图平衡模型大小和精度，以满足实际应用的要求。anchor-free探测器[27,57,76]试图放弃anchor box的设计，以避免耗时的box操作和繁琐的超参数调谐。提出了动态标签分配方法[15,47,70]，以便更好地定义模型学习的正样本和负样本。GFL[34,35]引入了质量焦点损失(Quality Focal Loss, QFL)和一个分布导向的质量预测器来提高分类分数和定位质量之间的一致性。它还将边界框表示建模为概率分布，以便它可以捕获框边缘的定位模糊性。近年来，由于变压器块编码表达特征的能力较强，DETR族[4,6,30,42,45,68,77]已成为目标检测界的新趋势。

2.2. 面向目标检测的知识蒸馏

知识蒸馏(Knowledge Distillation, KD)是一种将知识从大规模教师模型转移到小规模学生模型的有效方法。在分类任务中已经进行了广泛的研究[12,23,26,36,37,46,48,54,63,67,71,72]，但由于极端的背景比，提取检测模型仍然具有挑战性。先驱工作[7]通过简单地结合特征模仿和预测模仿，提出了用于目标检测的第一个蒸馏框架。自此，特征模仿引起了越来越多的研究关注。通常，一些工作[13,25,33,61]侧重于选择有效的蒸馏区域以更好地模仿特征，而其他工作[19,31,75]旨在更好地衡量模仿损失。也有方法[5,65,66,69]试图设计新的师生一致性函数，旨在探索更多的一致性信息或释放MSE损失的严格限制。

预测模拟作为[24]中最早提出的蒸馏策略，在分级蒸馏中起着至关重要的作用。近年来，人们提出了一些改进的预测模拟方法来适应目标检测。例如Rank miming[31]将教师的分数排名视为一种知识，目的是迫使学生将实例作为教师进行排名。LD[74]提出提取边界盒的定位分布[35]来传递定位知识。为了缓解预测模拟中的目标冲突问题，我们构建了一个将检测和蒸馏分离到不同头部的CrossKD管道。值得注意的是，HEAD[59]将学生特征传递给独立的助理HEAD，以弥合异构师生对之间的差距。相反，我们观察到，仅仅将学生特征传递给教师就足以有效地实现SOTA结果。这使得我们的方法非常简洁，与HEAD不同。我们的方法也与[1,28,32,64]相关，但它们都旨在提取分类模型，而不是为目标检测量身定制的。

3. Methodology

3.1. 目标冲突问题分析

目标冲突是传统预测模拟方法中经常遇到的问题。与为每个图像分配特定类别的分类任务不同，高级检测器中的标签通常是动态分配的，而不是确定性的。通常，检测器依赖于手工制作的原理，即分配器，以确定每个位置的标签。在大多数情况下，检测器不能准确地复制分配者的标签，这导致KD中教师-学生目标之间的冲突。此外，在现实场景中，学生和教师的分配者的不一致扩大了基础真理和蒸馏目标之间的距离。

为了定量衡量目标冲突的程度，我们在COCO最小数据集中统计了不同师生差异下冲突区域与积极区域的比例，结果如图3所示。

我们可以看到，即使教师(ATSS[70]和GFL[35])和学生(GFL)都有相同的标签分配策略，仍然有许多位置的基础真值和蒸馏目标之间的差异分别大于0.5。当我们使用具有不同分配者(RetinaNet)的教师来提炼学生(GFL)时，冲突区域会大大增加。4.5节中的更多实验也表明，目标冲突问题严重阻碍了预测模拟的性能。

尽管目标冲突的影响很大，但在以往的预测模拟方法中，这一问题长期被忽视[24,31]。这些方法旨在直接减少师生预测之间的差异。其目标可以描述为：
$${\mathcal{L}}_{\mathrm{KD}}=\frac{1}{|\mathcal{S}|}\sum _{r\in \mathcal{R}}\mathcal{S}(r){\mathcal{D}}_{\mathrm{pred}}({\mathbf{p}}^s(r),{\mathbf{p}}^t(r)),\text{\hspace{1em}(1)}$$
其中，${\mathbf{p}}^s$和${\mathbf{p}}^t$分别是由学生和老师的检测头生成的预测向量。${\mathcal{D}}_{\mathrm{pred}}(\cdot )$是计算${\mathbf{p}}^s$和${\mathbf{p}}^t$之间差异的损失函数，如分类用KL Divergence[24]，回归用L1 loss [7]， LD[74]。$\mathcal{S}(\cdot )$是区域选择原则，在整个图像区域r的每个位置$\mathcal{R}$处产生一个权值。

值得注意的是，$\mathcal{S}(\cdot )$在一定程度上可以通过降低师生差异较大区域的权重来缓解目标冲突问题。然而，高度不确定的区域通常比无争议的区域为学生提供更多的信息。忽略这些区域可能会对预测模拟方法的有效性产生很大影响。因此，为了推动预测模拟的极限，有必要优雅地处理目标冲突问题，而不是直接降低权重。

3.2. 交叉头知识蒸馏

如3.1节所述，我们观察到，直接模仿教师的预测会遇到目标冲突问题，这阻碍了预测模仿取得令人满意的成绩。为了缓解这个问题，我们在本节中提出了一种新的交叉头知识蒸馏(CrossKD)。总体框架如图4所示。像许多以前的预测模拟方法一样，我们的CrossKD对预测执行蒸馏过程。不同的是，CrossKD将学生的中间特征传递给教师的检测头，并生成交叉头预测来进行蒸馏。

给定一个密集检测器，如RetinaNet[39]，每个检测头通常由一系列卷积层组成，表示为 { C i } \{C_i\} {Ci​}。为了简单起见，我们假设每个检测头总共有$n$个卷积层(例如，在RetinaNet中有5个卷积层，有4个隐藏层和1个预测层)。我们用 f i , i ∈ { 1 , 2 ， ⋅ ⋅ ⋅ ， n − 1 } f_{i}, i∈\{1,2，···，n - 1\} fi​,i∈{1,2，⋅⋅⋅，n−1}表示$C_i$生成的特征映射，$f_0$表示$C_1$的输入特征映射。预测$\mathbf{p}$是由最后一个卷积层$C_n$生成的。因此，对于给定的师生对，教师和学生的预测可以分别表示为${\mathbf{p}}^s$和${\mathbf{p}}^t$。

除了教师和学生的原始预测外，CrossKD还将学生的中间特征 f i s , i ∈ { 1 , 2 ， ⋅ ⋅ ⋅ ， n − 1 } f^{s}_{i}, i∈\{1,2，···，n - 1\} fis​,i∈{1,2，⋅⋅⋅，n−1}传递到教师检测头的$(i+1)$卷积层$C_{i+1}^t$，从而得到交叉头预测${\hat{\mathbf{p}}}^s$。给定${\hat{\mathbf{p}}}^s$，我们不计算${\mathbf{p}}^s$和${\mathbf{p}}^t$之间的KD损失，我们建议使用十字头预测${\hat{\mathbf{p}}}^s$和教师${\mathbf{p}}^t$的原始预测之间的KD损失作为CrossKD的目标，其描述如下:
$${\mathcal{L}}_{\mathrm{CrossKD}}=\frac{1}{|\mathcal{S}|}\sum _{r\in \mathcal{R}}\mathcal{S}(r){\mathcal{D}}_{\mathrm{pred}}({\hat{\mathbf{p}}}^s(r),{\mathbf{p}}^t(r)),\text{\hspace{1em}(2)}$$
式中$\mathcal{S}(\cdot )$、$|\mathcal{S}|$分别为区域选择原则和归一化因子。我们没有设计复杂的$\mathcal{S}(\cdot )$，而是在整个预测图上平等地对${\hat{\mathbf{p}}}^s$和${\mathbf{p}}^t$进行蒸馏。具体来说，S(·)在我们的CrossKD中是一个常数函数，其值为1。根据每个分支的不同任务(例如，分类或回归)，我们执行不同类型的${\mathcal{D}}_{\mathrm{pred}}(\cdot )$，以有效地向学生传递特定任务的知识。

通过执行CrossKD，检测损失和蒸馏损失分别应用于不同的分支。如图4所示，检测损失的梯度穿过学生的整个头部，而蒸馏损失的梯度通过冻结的教师层传播到学生的潜在特征，这启发式地增加了教师和学生之间的一致性。与直接关闭师生对之间的预测相比，CrossKD允许学生的部分检测头仅与检测损失相关，从而导致对地面实况目标的更好的优化。定量分析在我们的实验部分介绍。

3.3. 优化目标

训练总损失可表示为检测损失与蒸馏损失的加权和，表示为：
$$\begin{array}{rl}\mathcal{L}& ={\mathcal{L}}_{\mathrm{cls}}({\mathbf{p}}_{\mathrm{cls}}^s,{\mathbf{p}}_{\mathrm{cls}}^{gt})+{\mathcal{L}}_{\mathrm{reg}}({\mathbf{p}}_{\mathrm{reg}}^s,{\mathbf{p}}_{\mathrm{reg}}^{gt})\\ & +{\mathcal{L}}_{\mathrm{CrossKD}}^{\mathrm{cls}}({\hat{\mathbf{p}}}_{\mathrm{cls}}^s,{\mathbf{p}}_{\mathrm{cls}}^t)+{\mathcal{L}}_{\mathrm{CrossKD}}^{\mathrm{reg}}({\hat{\mathbf{p}}}_{\mathrm{reg}}^s,{\mathbf{p}}_{\mathrm{reg}}^t),\end{array}\text{\hspace{1em}(3)}$$
其中${\mathcal{L}}_{\mathrm{cls}}$ 和${\mathcal{L}}_{\mathrm{reg}}$表示在学生预测$p_{\mathrm{cls}}^{\mathrm{s}},p_{\mathrm{reg}}^{\mathrm{s}}$和相应的ground truth目标$p_{\mathrm{cls}}^{gt},p_{\mathrm{reg}}^{gt}$之间计算的检测损失。额外的CrossKD损失表示为${\mathcal{L}}_{\mathrm{CrossKD}}^{\mathrm{cls}}$和${\mathcal{L}}_{\mathrm{CrossKD}}^{\mathrm{reg}}$，它们在交叉头预测${\hat{p}}_{\mathrm{cls}}^s,{\hat{p}}_{\mathrm{reg}}^{\mathrm{s}}$和教师的预测$p_{\mathrm{cls}}^t,p_{\mathrm{reg}}^t$之间执行。

我们使用不同的距离函数${\mathcal{D}}_{\mathrm{pred}}$在不同的分支中传递特定任务的信息。在分类分支中，我们将教师预测的分类分数作为软标签，直接使用GFL[35]中提出的质量焦点损失(Quality Focal Loss, QFL)来拉近师生距离。对于回归，在密集检测器中主要有两种回归形式。第一种回归形式直接从anchor框(例如，RetinaNet [39]， ATSS[70])或点(例如，FCOS[57])回归边界框。在这种情况下，我们直接使用GIoU[53]作为${\mathcal{D}}_{\mathrm{pred}}$。在另一种情况下，回归形式预测一个向量来表示框位置的分布(例如GFL[35])，它比边界框表示的Dirac分布包含更丰富的信息。为了有效地提取位置分布知识，我们采用了KL散度，如LD[74]来转移定位知识。关于损失函数的更多细节在补充材料中给出。

4. Experiments

4.1. 实现细节

我们在大规模MS COCO[40]基准上对所提出的方法进行了评估，与大多数先前的工作一样。为了确保与标准实践的一致性，我们使用trainval135k集(115K图像)进行训练，使用minival集(5K图像)进行验证。为了进行评估，使用了标准的coco风格测量，即平均精度(AP)。我们还报告了IoU阈值为0.5和0.75的mAP，以及小型、中型和大型对象的AP。我们提出的方法CrossKD是在Python的MMDetection[9]框架下实现的。为了公平的比较，所有的实验都是使用8个Nvidia V100 GPU开发的，每个GPU有两个图像的小批量。除非另有说明，所有的超参数都遵循相应的训练和测试学生模型的默认设置。

4.2. 方法分析

为了验证我们方法的有效性，我们进行了基于GFL的广泛烧蚀实验[35]。如果没有指定，我们使用带有ResNet-50骨干网的GFL[22]作为教师检测器，在学生检测器中使用ResNet-18骨干网。教师和学生的准确率分别为40.2 AP和35.8 AP。所有实验都遵循默认的1x训练计划(12次)。

要应用CrossKD的位置。如第3.2节所述，CrossKD将学生的第i个中间特征传递给教师头部的一部分。在这里，我们对分类分支和箱回归分支都进行了蒸馏。当i = 0时，CrossKD直接将学生的FPN特征输入到教师的脑海中。在这种情况下，整个学生的头部只受检测损失的监督，而不涉及蒸馏损失。随着i的逐渐增加，学生头部的更多层次共同受到检测损失和蒸馏损失的影响。当i = n时，我们的方法退化为原始预测模拟，其中蒸馏损失将直接在师生对的两个预测之间进行。

在表1中，我们报告了对不同中间特征执行CrossKD的结果。可以看到，我们的CrossKD可以提高i的所有选择的蒸馏性能。值得注意的是，当使用第3个中间特征时，CrossKD达到38.7 AP的最佳性能，比最近最先进的预测模拟方法LD高0.9 AP[74]。这表明，并非学生头中的所有层都需要与蒸馏损失的影响隔离开来。因此，在后续的所有实验中，我们都使用i = 3作为默认设置。

CrossKD v.s. Feature Imitation。我们将CrossKD与先进的特征模仿方法PKD进行比较[5]。为了公平比较，我们在与CrossKD相同的位置执行PKD，包括FPN特征和第三层检测头。结果见表2。可以看出，在师生对的FPN特征之间应用PKD可以达到38.0 AP。在检测头上，PKD甚至显示出性能下降。相比之下，我们的CrossKD达到38.7 AP，比FPN特征上应用的PKD高0.7 AP。

为了进一步研究CrossKD的优势，我们可视化了检测头潜在特征上的梯度，如图5所示。如图所示，PKD生成的梯度对整个特征映射有很大且广泛的影响，这是低效且没有针对性的。相反，CrossKD生成的梯度可以集中在具有感兴趣对象的潜在语义区域。

CrossKD v.s. Prediction Mimicking。我们首先在分类和盒回归分支上分别执行预测模拟和crossskd。结果见表3。可以看到，用CrossKD代替预测模拟可以获得稳定的性能增益，而不考虑分类或回归分支。具体来说，预测模拟在分类和回归分支上分别产生37.3 AP和36.8 AP，而CrossKD产生37.7 AP和37.2 AP，比预测模拟的相应结果有一致的改进。如果在两个分支上执行KD，我们的方法仍然可以比预测模拟高出+0.9 AP。

此外，我们进一步评估了预测模拟和CrossKD的集体效应，如表4所示。有趣的是，我们观察到，同时使用预测模拟和CrossKD产生38.1 AP的最终结果，这甚至低于单独使用CrossKD的结果。我们认为这是因为预测模仿再次引入了目标冲突问题，使得学生模型学习困难。

此外，我们将训练过程中的统计变化可视化，以对CrossKD和预测模拟进行进一步分析。我们首先计算学生的预测 $p^{\mathrm{s}}$和老师的预测$p^t$之间的$L_1$距离，以及每个时期的ground-truth目标$p^{gt}$。如图6(a)所示，距离$L_1({\mathbf{p}}^s,{\mathbf{p}}^t)$可以通过我们的CrossKD显着减少，而预测模拟实现最低距离是合理的，因为蒸馏直接施加在$p^s$上。然而，由于优化目标冲突的存在，预测模拟涉及到一个矛盾的优化过程，因此通常会产生比我们的CrossKD更大的距离$L_1(\dot{{\mathbf{p}}^{\mathrm{s}}},{\mathbf{p}}^{gt})$，如图6(b)所示。在图6©中，我们的方法显示出更快的训练过程，达到了37.8 AP的最佳性能。

4.3. 与SOTA KD方法的比较

在本节中，我们在GFL[35]框架上评估了各种最先进的目标检测KD方法，并将它们与我们提出的CrossKD进行了公平比较。我们使用ResNet-101作为教师检测器的主干，该检测器使用2倍时间表和多尺度增强进行训练。对于学生探测器，我们采用了ResNet-50骨干网。我们用1x计划训练学生。教师的预训练检查点直接借用了MMDetection[9]模型zoo。

我们在表5中报告了所有结果。正如我们所看到的，在相同的条件下，CrossKD可以达到43.7 AP，而不需要额外的修饰，这将学生的准确性提高了3.5 AP，优于所有其他最先进的方法。值得注意的是，CrossKD比先进的特征模拟方法PKD高出0.4 AP，比先进的预测模拟方法LD高出0.7 AP，证明了CrossKD的有效性。此外，我们还观察到CrossKD与特征模仿方法也是正交的。在PKD的帮助下，CrossKD达到了43.9 AP的最高成绩，与基线相比提高了3.7 AP。

4.4. 不同检测器上的CrossKD

除了在GFL上进行CrossKD外，我们还选择了三种常用的检测器，即RetinaNet[39]、FCOS[57]和ATSS[70]，来研究CrossKD的有效性。我们严格按照学生设置进行训练，并参考MMDetection模型动物园的师生结果。结果如表6所示。如表6所示，CrossKD显著提高了这三种检测器的性能。具体而言，使用我们的CrossKD的RetinaNet, FCOS和ATSS分别达到39.7 AP, 41.3 AP和41.8 AP，比相应的基线分别高出2.3 AP, 2.8 AP和2.4 AP。蒸馏后的所有结果甚至超过了原始教师，表明CrossKD可以在不同的密度检测器上很好地工作。

4.5. 严重目标冲突下的蒸馏

在本小节中，我们在具有不同分配器的检测器之间执行预测模拟和CrossKD，以探索CrossKD对目标冲突问题的有效性。如表7所示，目标冲突问题对学生的优化影响较大，导致学生成绩较差。具体来说，预测模拟将教师作为RetinaNet的AP降低到30.3，这与GFL有不同的分配者。此外，即使ATSS与GFL有相同的分配者，学生的AP也只有39.7，低于没有KD的表现。相比之下，CrossKD即使在基础真值与蒸馏目标之间存在较大差异的情况下，仍然可以显著提高学生的准确性。当使用ATSS作为教师时，CrossKD将GFL-R50的准确率提高到42.1 (+1.9 AP)。即使在一个较弱的教师ReitnaNet的指导下，CrossKD仍然将GFL-R50的性能提高到41.2 AP，比基线高1.0 AP。这证明了我们的CrossKD是在面对严重的目标冲突时。

4.6. 异构主干之间的蒸馏

在本小节中，我们将探讨我们的CrossKD提炼异质学生的能力。我们在不同骨干网的retanet[39]上执行知识蒸馏，并将我们的方法与最新的最先进的PKD方法[5]进行比较。具体来说，我们选择了两种典型的异构骨干网，即变压器骨干网swing - t[43]和轻量级骨干网MobileNetv2[55]。所有的检测器都用单尺度策略训练了12个epoch。结果见表8。我们可以看到，当从swing - t中提取知识时，CrossKD达到38.0 AP (+1.5 AP)，优于PKD 0.8 AP。CrossKD还将使用mobilenetv2骨干网的RetinaNet的结果提高到34.1 AP，比基线高3.2 AP，超过PKD 0.9 AP。

5. Conclusions and Discussions

在本文中，我们介绍了一种新的KD方法CrossKD，旨在提高密集目标检测器的性能。CrossKD将中间特征从学生的头脑转移到教师的头脑，以产生蒸馏的十字头预测，这是缓解监督目标和蒸馏目标之间冲突的有效方法。我们的研究结果表明，CrossKD可以提高蒸馏效率，达到最先进的性能。在未来，我们将进一步将我们的方法扩展到其他相关领域，例如3D物体检测。

Appendix

6. Details of Distillation Losses

7. The Generalization Ability of CrossKD

8. More Ablations

9. Relation to Previous Works

10. Result Visualization