当前位置：首页 > news >正文

用于快速训练收敛的 Conditional DETR

news 2025/7/3 10:20:38

摘要

最近提出的 DETR 方法将 transformer 的 encoder-decoder 架构应用于目标检测，并取得了可观的性能表现。本文关注其关键问题——训练收敛速度慢，并提出了一种用于快速 DETR 训练的条件交叉注意力机制。我们的方法的动机在于，DETR 中的交叉注意力在定位物体的四个边界点以及预测边框时高度依赖 content embedding，这对高质量 content embedding 的要求较高，从而增加了训练难度。

我们的方法称为 Conditional DETR，它从 decoder embedding 中学习一个条件空间查询，用于 decoder 的多头交叉注意力。其优点在于，通过条件空间查询，每个交叉注意力头可以聚焦于包含特定区域的带状区域，例如物体的某个边缘点或物体框内部的区域。这样可以缩小进行目标分类与边框回归时需要定位的空间范围，从而降低对 content embedding 的依赖，简化训练过程。实验结果表明，Conditional DETR 在 ResNet-50 和 ResNet-101 骨干网络上收敛速度提升了 6.7 倍，在更强的 DC5-R50 和 DC5-R101 骨干网络上加速达到 10 倍。代码已开源，地址为：https://github.com/Atten4Vis/ConditionalDETR。

1. 引言

DEtection TRansformer（DETR）方法 [3] 将 transformer 的 encoder-decoder 架构应用于目标检测任务，并取得了良好的性能。它有效地消除了对许多手工设计组件的依赖，包括非极大值抑制（non-maximum suppression）和 anchor 的生成。然而，DETR 方法存在训练收敛缓慢的问题，需要长达 500 个训练周期才能获得较好的性能。最近的工作 deformable DETR [53] 通过将全局稠密注意力（即 self-attention 和 cross-attention）替换为只关注少量关键采样点的可变形注意力，并引入高分辨率和多尺度的 encoder，有效地缓解了这一问题。与此不同的是，我们仍然保留全局稠密注意力机制，并提出了一种改进的 decoder cross-attention 机制，用于加速训练过程。

在这里插入图片描述

我们的方法受到以下现象的启发：cross-attention 中对 content embedding 的高度依赖，以及 spatial embedding 所起的作用相对较小。DETR [3] 的实证结果表明，如果从第二个 decoder 层中移除 key 中的位置嵌入和 object query，仅使用 key 和 query 中的 content embedding，检测的 AP 仅略有下降¹。

图 1（第二行）展示了在训练 50 个 epoch 后，DETR 中 cross-attention 的空间注意力权重图。从中可以看到，有两个图没有正确地突出对应目标边缘位置的带状区域，因此无法有效收缩 content query 在空间上的关注范围，以精确定位目标边缘。其原因有两个：(i) 空间查询（即 object query）只能给出一般性的注意力权重图，未能利用图像的具体信息；(ii) 由于训练轮数较少，content query 本身尚不够强，难以很好地与 spatial key 匹配，而 content query 同时还要与 content key 匹配。这使得模型对高质量 content embedding 的依赖增强，从而加大了训练难度。

在这里插入图片描述

我们提出了一种 Conditional DETR 方法，它为每个 query 从对应的上一个 decoder 输出 embedding 中学习一个条件空间嵌入（conditional spatial embedding），用于构建所谓的条件空间查询（conditional spatial query），以用于 decoder 的多头 cross-attention。这个条件空间查询是通过将用于回归目标框的信息映射到 embedding 空间中来预测的，该空间与 key 的二维坐标也被映射到的空间一致。

我们通过实证观察到，在使用空间查询和 key 的情况下，每一个 cross-attention head 会在空间上关注于包含目标边缘或目标框内部区域的带状区域（见图 1 第一行）。这使得 content query 用于定位类别与框预测相关区域时，其空间范围被有效收缩。由此一来，模型对 content embedding 的依赖被减弱，从而使训练过程更容易。

实验结果表明，Conditional DETR 在使用 R50 和 R101 骨干网络时收敛速度提升了 6.7 倍，在使用更强的骨干网络 DC5-R50 和 DC5-R101 时收敛速度提升了 10 倍。图 2 展示了 Conditional DETR 与原始 DETR [3] 的收敛曲线。

2. 相关工作

基于 anchor 和无 anchor 的目标检测。目前大多数目标检测方法都是从经过精心设计的初始假设出发进行预测，主要有两类初始假设：anchor boxes（锚框）和目标中心点。基于 anchor 的方法继承自基于候选框的方法，如 Fast R-CNN。代表性方法包括 Faster R-CNN [9]、SSD [26]、YOLOv2 [31]、YOLOv3 [32]、YOLOv4 [1]、RetinaNet [24]、Cascade R-CNN [2]、Libra R-CNN [29]、TSD [35] 等。

无 anchor 的检测器则是在靠近目标中心的点上预测边界框。典型方法包括 YOLOv1 [30]、CornerNet [21]、ExtremeNet [50]、CenterNet [49, 6]、FCOS [39]，以及其他方法 [23, 28, 52, 19, 51, 22, 15, 46, 47]。

DETR 及其变种。DETR 成功地将 transformer 应用于目标检测，消除了许多人工设计的模块，如非极大值抑制和初始框生成。但由于 encoder 中全局 self-attention 带来的高计算复杂度问题，后续工作如 adaptive clustering transformer [48] 和 deformable DETR [53] 分别采用了聚类和稀疏注意力机制来缓解这一问题。

另一个关键问题是训练收敛速度慢，近年来受到了广泛关注。TSP（transformer-based set prediction）方法 [37] 移除了 cross-attention 模块，结合了 FCOS 和类似 R-CNN 的检测头。Deformable DETR [53] 使用 deformable attention 替代 decoder 中的 cross-attention，其通过 content embedding 学习稀疏位置进行关注。

与我们的方法同时提出的还有空间调制协同注意力（SMCA）方法 [7]，它与我们的方法非常相似。SMCA 使用从 decoder embedding 学习得到的一些（偏移的）中心，结合高斯分布生成的注意力图，对 DETR 中的多头 cross-attention 进行调制，使其关注于预测框内部的少数区域。与之不同，我们提出的 Conditional DETR 方法是从 decoder 的 content embedding 中学习条件空间查询（conditional spatial query），并以可学习的方式（而非人为设计）预测空间注意力权重图，该权重图可聚焦于用于框回归的四个边缘区域和用于分类的目标内部显著区域。

条件卷积与动态卷积。我们提出的条件空间查询机制与条件卷积核生成有关。Dynamic filter network [16] 从输入中学习卷积核，在 CondInst [38] 和 SOLOv2 [42] 中被应用于实例分割以学习与实例相关的卷积核。CondConv [44] 和 dynamic convolution [4] 将多种卷积核与从输入中学习得到的权重进行混合。SENet [14]、GENet [13] 和 Lite-HRNet [45] 则从输入中学习通道级别的权重。

这些方法都是从输入中学习卷积核权重，再将其应用于输入。而我们的方法则是从 decoder embedding 中学习一个线性映射，用于表示位移与缩放信息。

Transformers。Transformer [40] 依赖于注意力机制，包括 self-attention 和 cross-attention，用以建立输入与输出之间的全局依赖关系。有若干研究与我们的方法密切相关。Gaussian transformer [11] 和 T-GSA（带高斯加权 self-attention 的 transformer）[18]，以及之后的 SMCA [7]，都依据目标与上下文符号之间的距离，通过学习或人为设计的高斯方差来衰减注意力权重。与我们的方法类似，TUPE [17] 同时从空间注意力权重和内容注意力权重中计算 attention 权重。而我们的工作则更专注于提出一种可学习形式的注意力衰减机制，而不是使用高斯函数，并且可能对语音增强 [18] 和自然语言推理 [11] 等任务带来潜在益处。

3. Conditional DETR

3.1. 概览

Pipeline. 本文提出的方法遵循 detection transformer（DETR）的思路，是一种端到端的目标检测器，可以一次性预测所有目标，无需使用非极大值抑制（NMS）或 anchor 生成。其整体架构包括一个 CNN 主干网络、一个 transformer 编码器、一个 transformer 解码器，以及用于目标类别和边界框位置的预测器。transformer 编码器的目的是增强从 CNN 主干网络输出的内容 embedding。编码器由多个编码层堆叠而成，每一层主要由一个自注意力层和一个前馈层组成。

transformer 解码器是由多个解码层组成的堆叠结构。每个解码层如图 3 所示，包含三个主要部分：(1) 一个用于去除重复预测的自注意力层，该层对用于类别和边界框预测的 embedding（即前一解码层的输出）之间进行交互；(2) 一个交叉注意力层，该层聚合编码器输出的 embedding，以精炼解码器的 embedding，从而提升类别和边界框的预测能力；(3) 一个前馈层。

Box regression. 每个解码器 embedding 用于预测一个候选框，计算如下：
$\mathbf { b } = \mathrm { s i g m o i d } ( \mathrm { F F N } ( \mathbf { f } ) + [ \mathbf { s } ^ { \top } \mathbf { \boldsymbol { 0 } } \mathbf { \boldsymbol { 0 } } ] ^ { \top } )\quad(1)$

其中， $\mathbf{f}$ 是解码器的 embedding， $\mathbf{b}$ 是一个四维向量 $\, b _ { c y } \, b _ { w } \, b _ { h } ] ^ { \top }$ ，分别表示边界框的中心坐标、宽度和高度。 $\mathrm{sigmoid}()$ 用于将预测结果 $\mathbf{b}$ 归一化到 [0, 1] 范围内。 $\mathrm{FFN}(\cdot)$ 用于预测未经归一化的边界框。 $\mathbf{s}$ 是参考点的未归一化二维坐标，在原始 DETR 中设为 (0, 0)。在我们的方法中，我们考虑两种选择：将参考点 $\mathbf{s}$ 作为每个候选框预测的参数来学习，或者从对应的 object query 中生成。

Category prediction. 每个候选框的分类得分也通过一个前馈网络（FNN）从解码器的 embedding 中预测得到，计算如下： $\mathbf{e} = \mathrm{FFN}(\mathbf{f})$ 。

在这里插入图片描述

Main work. 交叉注意力机制旨在定位显著区域，包括用于边界框检测的四个极点以及框内用于目标分类的区域，并聚合对应的 embedding。我们提出了一种条件交叉注意力机制，通过引入条件空间查询（conditional spatial queries）来提升定位能力并加速训练过程。

3.2. DETR 解码器中的交叉注意力

DETR 解码器中的交叉注意力机制接收三个输入：query、key 和 value。每个 key 由内容 key $c_k$ （即编码器输出的 content embedding）和空间 key $p_k$ （对应归一化二维坐标的位置 embedding）相加组成。value 同样由编码器输出的 content embedding 构成，与内容 key 相同。