论文:DETRDistill: A Universal Knowledge Distillation Framework for DETR-families
ICCV2023
代码:https://github.com/BIVLab-USTC/DETRDistill
abtract
基于 Transformer 的检测器 (DETR) 因其简单的框架而日益流行,但庞大的模型规模和耗时阻碍了其在现实世界中的部署。知识蒸馏 (KD) 是一种颇具吸引力的技术,可以将巨型检测器压缩成小型检测器,从而获得相当的检测性能和较低的推理成本。由于 DETR 将目标检测表述为一个集合预测问题(a set prediction problem),因此现有为经典卷积检测器(classic convolution-based)设计的 KD 方法可能并不直接适用。
本文提出了(propose)一种专用于(dedicated) DETR 系列(families)的新型(novel)知识蒸馏(knowledge distillation)方法 DETRDistill。(Specifically) 具体而言,我们首先设计了一个匈牙利匹配(Hungariam-matching)逻辑向量蒸馏方法,以鼓励(encourage)学生模型获得与教师 DETR 完全相同的预测结果。然后,我们提出了一种目标感知特征蒸馏方法,帮助学生模型从教师模型以目标为中心的特征(object-centric)中学习。最后,为了提高学生 DETR 的收敛速度,我们引入了一个查询先验分配蒸馏方法(query-prior assignment distillation),以加快学生模型从训练良好的查询中学习的速度以及教师模型的稳定分配。在 COCO 数据集上进行的大量实验结果验证了我们方法的有效性。值得注意的是,DETRDistill 持续将各种 DETR 的 mAP 提升了 2.0 以上,甚至超越了其对应的教师模型。
1. Introduction
目标检测旨在从输入图像中定位和分类视觉目标。早期的研究通常使用卷积神经网络 (CNN) 来处理输入图像的区域特征 [26, 19],其中包含一系列归纳偏差,例如锚点 [26]、标签分配 [41] 和重复项删除 [1]。近期,基于 Transformer 的目标检测器(例如 DETR [2])被提出,将检测视为一个集合预测任务,这显著简化了检测流程,并帮助用户摆脱了繁琐的手动调整工作,例如调整锚点大小和比例 [26]。
尽管基于 Transformer 的检测器已经达到了最佳性能 [21, 17, 11],但它们存在计算成本高昂的问题,难以部署在实时应用中。为了获得快速准确的检测器,知识蒸馏 [14] (KD) 是一项颇具吸引力的技术。通常,知识蒸馏方法通过模仿预测(mimicking the predicitons) [14] 或特征分布(feature distributions) [27],将知识从权重高但性能强大的教师模型迁移到小型高效的学生网络。
在目标检测研究领域,已经发表了各种知识蒸馏 (KD) 方法 [37, 36, 39, 27, 4, 18, 12, 42, 38, 32, 29]。然而,这些方法大多是为基于卷积的检测器设计的,由于检测框架的差异,可能无法直接应用于基于 Transformer 的 DETR。至少存在两个挑战:❶ 对数级蒸馏方法(logits-level distillation methods) [14, 42] 不适用于 DETR。无论是基于锚点 [41] 还是无锚点 [31] 的卷积检测器,框预测都与特征图网格紧密相关,因此自然地确保了在教师和学生之间进行知识蒸馏时,框预测具有严格的空间对应关系。然而,对于 DETR,解码器生成的框预测是无序的,并且对于对数级蒸馏,教师和学生之间预测框之间不存在自然的一一对应关系。 ❷ 特征级蒸馏(Feature-level )方法可能不适用于 DETR。由于卷积和 Transformer [27] 之间的特征生成机制不同,感兴趣目标的特征激活区域差异很大。如图 2 所示,基于卷积的检测器的活动区域几乎被限制在真实框内,而 DETR 检测器则进一步激活了背景区域。因此,直接将以前的特征级知识蒸馏 (KD) 方法用于 DETR 并不一定能带来性能提升,有时甚至会损害学生检测器的性能,如表 1 所示。
为了应对上述挑战,我们提出了 DETRDistill,一个专为 DETR 家族检测器设计的知识蒸馏框架。具体来说,DETRDistill 主要由三个组件组成:
(1) 匈牙利匹配对数回归蒸馏(Hungarian-matching logits distillation):为了解决挑战❶,我们使用匈牙利算法在学生和教师的预测之间找到最优的二分匹配,然后在对数回归级别进行知识提取 (KD)。然而,由于教师模型中预测为正值的边界框数量非常有限,仅对正值预测进行知识提取 (KD) 并不能带来显著的性能提升。因此,我们建议在教师和学生模型之间大量的负值预测上引入蒸馏损失,以充分利用教师检测器中的知识。此外,考虑到 DETR 方法通常包含多个解码器层来进行级联预测细化,我们还在每个阶段创建了知识提取损失,以实现渐进式蒸馏。
(2) 目标感知特征蒸馏(Target-aware feature distillation):根据挑战❷中的分析,我们建议利用目标查询和教师模型特征来生成软激活掩码。由于训练良好的教师查询与各种目标密切相关,因此生成的软掩码将以目标为中心,从而使基于软掩码的特征级蒸馏具有目标感知能力。
(3) 查询先验分配蒸馏(Query-prior assignment distillation):由于学生模型中的查询和解码器参数是随机初始化的,学生模型中不稳定的二分分配会导致收敛速度缓慢,如 [17] 中所述。虽然我们通过经验发现,教师模型中训练良好的查询始终能够生成一致的二分分配(如图 7 所示),但因此我们建议让学生模型将教师查询作为一组额外的先验查询,并鼓励其基于教师网络稳定的二分分配进行预测。这种蒸馏成功地帮助学生模型快速收敛并实现更佳性能。
总而言之,我们的贡献体现在三个方面:
• 我们详细分析了与传统的基于卷积的检测器相比,DETR 在蒸馏任务中遇到的困难。
• 我们分别从对数级、特征级和收敛速度的角度,提出了多种针对 DETR 的知识蒸馏方法。
• 我们在 COCO 数据集上进行了不同设置下的大量实验,结果证明了我们提出的方法的有效性和泛化能力。
2. Related Work
2.1. Transformer-based Object Detectors
随着 Transformer [33] 在自然语言处理中的优异表现,研究者们也开始探索将 Transformer 结构应用于视觉任务 [23, 7, 24]。然而,DETR 的训练过程极其低效,因此许多后续工作试图加速收敛。其中一项工作尝试重新设计注意力机制。例如,Dai 等人 [43] 提出了可变形 DETR,通过仅与参考点周围的可变采样点特征进行交互来构建稀疏注意力机制。SMCA [10] 引入高斯先验来限制交叉注意力。AdaMixer [11] 设计了一种新的自适应三维特征采样策略,无需任何编码器,然后将通道和空间维度上的采样特征与自适应权重进行混合。
另一项工作重新思考了查询(Query)的含义。Meng 等人[25] 可视化地表明 DETR 依靠交叉注意中的内容嵌入来定位物体末端是无效的,因此建议将查询分离为内容部分和位置部分。Anchor-DETR [35] 直接将查询的 2D 参考点作为其位置嵌入来引导注意。DAB-DETR [22] 除了位置之外,还在注意机制中引入了宽度和高度信息来建模不同尺度的物体。DNDETR [17] 引入了查询去噪任务来加速训练。Group-DETR [6] 和 H-DETR [15] 通过在解码器训练中增加正样本作为辅助组来提高性能。与以前的工作不同,我们期望通过蒸馏来提高小模型的性能。
2.2. Knowledge Distillation in Object Detection
知识蒸馏是一种常用的模型压缩方法。[14] 首次提出这一概念并将其应用于图像分类。他们认为,与独热编码相比,教师输出的软标签包含类别间相似性的“暗知识”,这有助于模型的泛化。注意力转移 [39] 将蒸馏重点放在特征图上,通过缩小教师和学生的注意力分布(而不是蒸馏输出逻辑值)来迁移知识。
FitNet [27] 提出通过隐藏层模仿教师模型的中级提示。[4] 首次将知识蒸馏应用于解决多类目标检测问题。[18] 认为背景区域会引入噪声,并提出对 RPN 采样的区域进行蒸馏。DeFeat [12] 分别对前景和背景进行蒸馏。FGD [37] 分别在焦点区域和特征的全局关系方面模仿教师模型。 LD [42] 将软标签蒸馏扩展到位置回归,使学生模型拟合老师模型的边界预测分布。MGD [38] 使用蒙版图像建模 (MIM) 将模仿任务转化为生成任务。除了上述基于 CNN 的蒸馏之外,一些研究还涉及视觉转换器。DeiT [32] 通过蒸馏标记将归纳偏差从 CNN 老师模型转移到 ViT [9],并在分类任务中取得了颇具竞争力的表现。ViDT [29] 对块标记执行知识挖掘 (KD),并提出了一种转换器检测器的变体。然而,这种蒸馏不能直接应用于 DETR 系列。我们的工作分析了 DETR 不同组成部分的独特现象,并提出了一种通用的蒸馏策略。
3. A Review of DETR
DETR [2] 是一个端到端目标检测器,包含主干网络、Transformer 编码器、可学习的查询嵌入和解码器。给定图像 I,CNN 主干网络提取其空间特征,然后 Transformer 编码器(某些变体不需要编码器 [11])将增强特征表示。利用更新后的特征 F ∈ RHW×d,将查询嵌入 Q ∈ RN×d 输入到多个 Transformer 解码器(通常为 6 个)中,其中 d 是特征维度,N 是固定的查询数量。每个解码器阶段中的操作类似:首先,利用自注意力机制建立查询之间的关系以捕获互信息。其次,通过灵活的交叉注意力机制将查询与图像特征进行交互,从而将查询与有价值的语义信息聚合在一起。第三,前馈网络 (FFN) 将每个查询解码为 ˆyi = (ˆci, ˆbi),其中包含预测类别和边界框。
在训练阶段,标签分配的原则是最小化模型预测与真实值(GT)之间的匹配成本,利用匈牙利算法[16]进行二分匹配。最优匹配求解如下:
因此,DETR 中的每个 GT 仅对应一个正样本查询,其余所有查询均视为负样本。
其中,负样本的位置回归不受任何监督。
4. Our Approach: DETRDistill
在本节中,我们将介绍我们提出的 DETRDistill 的细节,它由三个部分组成:(1)匈牙利匹配 Logits 蒸馏;(2)目标感知特征蒸馏;(3)查询优先分配蒸馏。图 3 展示了 DETRDistill 的整体架构。
4.1. Hungarian-matching Logits Distillation
知识蒸馏最常见的策略之一是直接在对数级对齐两个模型的预测。然而,基于查询的集合形式的预测 [17] 使得 DETR 难以将教师的成绩与学生的成绩进行有序的对应。为了实现这一目标,我们重用了匈牙利算法,将教师的预测与学生的预测进行一对一的匹配。
形式化地,令 ˆyT 和 ˆyS 分别表示教师模型和学生模型的预测,满足 ˆyT ={ˆiyT pos }Mpos i=1, {ˆjyT neg }Mneg j=1 和 ˆyS = {ˆyS i }N i=1,其中Mpos 和 Mneg 分别表示教师模型的正向预测和负向预测的数量。M = Mpos +Mneg 和 N 分别是教师和学生的解码器查询总数。M 通常大于或等于 N。由于教师模型的正向预测与目标密切相关,一种直接的想法是将它们视为知识型伪知识图谱 (GT),并利用匈牙利算法在这些正向预测 yˆT pos 和学生模型 ˆyS 的正向预测之间找到一个匹配的 ˆσpos。这样就可以实现对数级知识发现 (KD)
然而,我们通过经验发现,这种简单的知识发现 (KD) 仅能带来微小的性能提升,如表 8 所示。我们假设正向预测的数量非常有限(平均每幅图像只有 7 个,而查询总数通常超过 100 个),并且提取的信息与 (GT) 高度一致。另一方面,教师模型的大量负向预测被忽略,我们认为这些预测是有价值的。
负位置蒸馏。由于教师模型通常经过良好优化,生成的正预测和负预测可能存在明显差异,因此匈牙利算法可以生成合理的分配,即这些负预测的框将远离目标。而随机初始化的学生网络可能不会产生这样的效果,学生的负预测可能会与正预测纠缠在一起。因此,我们建议创建一种蒸馏方法,以利用教师模型负预测中蕴含的知识。*
其中 ˆσneg 表示 ˆyT neg 和 ˆyS 之间的赋值。渐进式蒸馏。考虑到 DETR 解码器通常包含多个阶段,并且默认包含分阶段监督 [2],我们进一步建议将公式 3 和公式 4 中的 KD 损失引入到每个解码器阶段,以实现渐进式蒸馏。
其中 K 是解码器阶段的数量。Lpos (5) logitsKD[k] 和 Lneg logitsKD[k] 分别表示第 k 个解码器阶段的正 KD 损失和负 KD 损失。
请注意,我们将知识从教师模型的逐阶段输出迁移到学生模型的相应阶段,而不是简单地使用教师模型的最后一个阶段的输出来监督学生模型的所有阶段。这是因为我们认为教师模型在不同阶段会包含不同的知识,正如在最近的研究[3]中观察到的那样,而之前的蒸馏策略可以充分利用教师模型中的知识,表9中的实证结果验证了我们的论点。
4.2. Target-aware Feature Distillation
检测性能在很大程度上取决于特征金字塔网络 (FPN) 生成的特征表示,这可以归因于其与目标相关的丰富语义信息。因此,我们认为有必要在特征层面提炼教师模型的知识。模拟教师模型空间特征的典型方式可以计算为
其中 FT ∈ RH×W×d 和 FS ∈ RH×W×dS 分别表示由教师和学生模型生成的特征表示。H 和 W 表示特征的高和宽,d 是教师特征的通道数。φ 是一个可学习的维度自适应层,用于将学生特征转换为 d 维。⊙ 是两个矩阵的 Hadamard 积。ψ ∈ RH×W 表示用于在各种知识发现 (KD) 方法中选择知识区域的软掩码,例如,Romero 等人 [27] 将掩码视为一个填充 1 的矩阵。Wang 等人 [34] 根据锚框和预测框 (GT box) 之间的 IoU 分数生成掩码。Sun 等人 [30] 利用高斯掩码覆盖 GT box。与上述方法不同,我们提出通过计算**查询嵌入(query embeddings )**和特征表示(feature representations)之间的相似度矩阵来构建 DETR 的软掩码。形式上,给定教师模型的一整套查询 QT ∈ RM×d,可以得到选择掩码
Qi 是第 i 个教师查询,M 表示教师模式的解码器查询数量
然而,我们通过实证研究发现,这种香草(vanilla)蒸馏方法效果不佳,如表 2 所示。我们推测其原因在于,并非所有教师模型的对象查询都应被视为有价值的线索。基于教师模型查询生成的预测,图 4 展示了一些基于查询的掩码 {ψi} 的可视化结果,我们发现预测分数较低的掩码关注的是对象区域之外的区域。
根据观察结果,我们建议选择性地利用教师模型查询来生成掩码 ψ。具体而言,我们使用 [8] 中提出的质量分数作为衡量指标
ci 和 bi 分别表示第 i 个教师查询的分类分数和预测框。bGTi 是相应的二分匹配 GT 框。γ =0.5 是用于平衡分类分数和框 IoU 权重的超参数。然后,目标感知质量分数作为指标,指导哪个查询应该对知识蒸馏做出更多贡献,公式 6 中的 KD 损失可以扩展为
4.3. Query-prior Assignment Distillation
由于DETR中的查询和解码器参数通常为了模型优化而随机初始化,一个查询可能会在不同的训练周期被分配给不同的对象,从而导致二分图匹配不稳定且收敛速度缓慢[17]。在(KD)的设置下,学生DETR的训练也存在同样的问题。然而,我们通过经验观察到,教师模型中经过良好优化的查询能够在不同的解码器阶段之间持续实现稳定的二分图分配(如图7所示),并且利用教师模型的知识来提高学生模型训练的稳定性是直观的。基于此,我们提出了查询优先分配蒸馏(Query-prior assignment distillation)。
具体来说,给定教师查询集 QT,我们可以从教师那里获得任何给定input-GT pairs
对应的分配排列(assignment permutation) ˆσT。我们将教师查询嵌入 QT 作为一组额外的先验查询输入到学生模型中,并直接使用教师的分配 ˆσT 生成检测结果并进行损失计算。
提出的 KD 损失将有助于学生模型将教师查询视为先验,并鼓励学生检测器尽可能地实现稳定的分配。如图 7 所示,采用提出的蒸馏损失后,学生模型的匹配稳定性得到了显著提升。请注意,这种额外的教师查询组仅在训练期间使用,学生模型将使用其默认查询集进行最终评估。
4.4. 总损失
综上所述,训练学生 DETR 的总损失是公式 2、公式 5、公式 9 和公式 1 的加权组合。
其中 λ1 = 1,λ2 = 20,λ3 = 1 是所提出的三个 KD 损失项的平衡权重。由于我们的蒸馏方法遵循常见的 DETR 范式,因此可以轻松应用于 DETR 系列中的各种检测器。
7. 结论
本文介绍了一种用于 DETR 类检测器的通用知识蒸馏框架,名为 DETRDistill。
我们的方法包含三个蒸馏模块:匈牙利匹配对数蒸馏、目标感知特征蒸馏和查询优先分配蒸馏。在竞争性 COCO 基准上进行的大量实验证明了我们方法的有效性和泛化能力。我们希望 DETRDistill 能够为未来基于 DETR 的知识蒸馏研究奠定坚实的基础。