人工智能 - SMCA：港中文提出注意力图校准的DETR加速方案 | ICCV 2021 - 晓飞的算法工程笔记

为了加速DETR收敛，论文提出了简单而有效的Spatially Modulated Co-Attention（SMCA）机制，通过在初始边界框位置给予较高的协同注意力响应值的约束来构建DETR的回归感知协同注意力。此外，将SMCA扩展为多头注意力和尺度选择注意力后，对比DETR可以实现更好的性能（108周期45.6 mAP vs 500周期43.3 mAP）
来源：晓飞的算法工程笔记公众号

论文: Fast Convergence of DETR with Spatially Modulated Co-Attention

论文地址：https://arxiv.org/abs/2108.02404
论文代码：https://github.com/gaopengcuhk/SMCA-DETR

Introduction

最近提出的DETR通过消除手工设置的锚框和非极大值抑制（NMS），显著简化了目标检测流程。但与两阶段或一阶段检测器相比，DETR的收敛速度较慢（500 vs 40 周期），导致算法设计周期变长。研究人员难以进一步扩展该算法，从而阻碍了其广泛使用。

在DETR中，对象查询向量负责检测不同空间位置的对象。每个对象查询都与卷积神经网络 (CNN) 编码的空间视觉特征交互，通过协同注意力机制自适应地从空间位置收集信息，估计边界框位置和对象类别。但在DETR的解码器中，每个对象查询的协同关注视觉区域可能与查询要预测的边界框无关。因此，DETR的解码器需要长时间的训练来搜索适当的协同关注视觉区域，以准确识别相应的对象。

受这一观察的启发，论文提出了一种名为空间调制协同注意（SMCA）的新模块。取代DETR中现有的协同注意力机制，实现更快的收敛和更高的性能。

SMCA动态预测每个对象查询对应的框的初始中心和尺寸，生成2D空间的类高斯权重图。权重图逐元素与对象查询和图像特征产生的协同注意力特征图相乘，可以有效地聚合来自视觉特征图的查询对象的相关信息。通过这种方式，空间权重图有效地调节每个对象查询的协同注意力的搜索范围，使其适当地围绕最初预测对象的中心和尺寸。通过利用预测的高斯分布空间先验，SMCA可以显著加快DETR的训练速度。

虽然将SMCA机制简单地融入DETR中可以加速收敛速度，但与DETR相比性能较差（50 个周期时为41.0 mAP，108 个周期时为 42.7 mAP vs. 500个周期时为43.3 mAP）。受到多头注意力和多尺度特征的启发，论文将其与SMCA集成来进一步增强：

对于编码器中的多尺度视觉特征，不是简单地对CNN主干的多尺度特征进行缩放以形成联合多尺度特征图，而是引入了尺度内和多尺度混合的自注意力机制来进行多尺度视觉特征之间的信息传播。
在解码器中，每个对象查询可以通过尺度选择注意力自适应地选择适当尺度的编码特征。对于解码器中的多个协同注意头，都会估计特定的对象中心和尺度，生成不同的空间权重图用于调整协同注意力特征。

论文的贡献如下：

提出了一种新颖的空间调制协同注意力（SMCA），通过位置约束的目标回归来加速DETR的收敛。没有多尺度特征和多头注意力的基础版SMCA已经可以在 50 个周期时达到41.0 mAP，在 108 个周期时达到42.7 mAP。
完整版SMCA进一步集成了多尺度特征和多头空间调制，可以通过更少的训练迭代进一步显着改进和超越DETR。SMCA在 50 个周期时可以达到43.7 mAP，在 108 个周期时可以达到45.6 mAP。
在COCO 2017数据集进行了广泛的消融研究，以验证SMCA模块和网络设计。

Spatially Modulated Co-Attention

A Revisit of DETR

DETR将目标检测转化为集合预测问题。对DETR不熟悉的，可以看看之前的文章，【DETR：Facebook提出基于Transformer的目标检测新范式 | ECCV 2020 Oral】。

DETR先用卷积神经网络 (CNN) 从图像 $I\in\mathbb{R}^{3\times H_{0}\times W_{0}}$ 中提取视觉特征图 $f\in\mathbb{R}^{C\times H\times W}$，其中 $H,W$ 和 $H_{0},W_{0}$ 分别是输入图像和特征图的高度和宽度。

经过位置嵌入增强的视觉特征 $f_{pe}$ 将被输入到__Transformer__的编码器中，对 $f_{p e}$ 生成的键、查询和值特征进行自注意力计算，在所有空间位置的特征之间交换信息。为了增加特征多样性，这些特征将沿着通道维度分为多个组进行多头自注意力：

$$ \begin{array}{l c r}{{E_{i}=\mathrm{Softmax}(K_{i}^{T}Q_{i}/\sqrt{d})V_{i},}}\ {{E=\mathrm{Concat}(E_{1},\cdot\cdot\cdot\cdot\cdot E_{H}),}}\end{array} \quad\quad (1) $$

其中 $K_{i}$, $Q_{i}$, $V\_{i}$ 表示键、查询和值特征的第 $i$ 个特征组。每种类型的特征都有 $H$ 组，输出的编码器特征 $E$ 经过进一步变换后输入到Transformer的解码器中。

对于编码器编码的视觉特征 $E$，DETR在对象查询 $O\_q\in\mathbb{R}^{N\times C}$ 和 $E\in\mathbb{R}^{L\times C}$ 之间执行协同注意力（$N$ 表示预先指定的对象查询的数量，$L$ 是空间视觉特征的数量）：

$$ \begin{array}{l} {{Q=\mathrm{FC}(O\_{q}),\ K,V=\mathrm{FC}(E)}} \ {{C\_{i}=\mathrm{Softmax}(K\_i^T Q\_i/\sqrt{d})V\_i,}} \ {{C_{i}=\mathrm{Concat}(C_{1},\ldots\cdot C\_{H}),}} \end{array} \quad\quad (2) $$

其中 $\mathrm{FC}$ 表示单层线性变换，$C\_i$ 表示第 $i$ 个协同注意力头关于对象查询 $O\_q$ 的协同特征。解码器的每个对象查询的输出特征由多层感知器 (MLP) 处理，输出每个对象的类分数和框位置。

给定预测的框和类别，在预测结果和GT之间应用匈牙利算法来识别每个对象查询的学习目标。

Spatially Modulated Co-Attention

DETR中的原始协同注意力不知道预测的边界框，因此需要多次迭代才能为每个对象查询生成正确的注意力图。论文提出的SMCA的核心思想是，将可学习的协同注意力图与手工设计的查询空间先验相结合，使得关注的特征限制在对象查询的初始输出位置周围，从而与最终的对象预测更加相关。

Dynamic spatial weight maps

每个对象查询先动态预测其负责对象的中心和比例，用于生成类似2D高斯的空间权重图。对象查询 $O\_q$ 的类高斯分布归一化中心 $c\_h^{norm}$、$c\_w^{norm}$ 和尺度 $s\_h$、$s\_w$ 的计算为：

$$ \begin{array}{c}{{c_{h}^{\mathrm{norm}},c_{w}^{\mathrm{norm}}={\mathrm{sigmoid}}(\mathrm{MLP}(O_{q})),}}\ {{s_{h},s_{w}=\mathrm{FC}(O_{q}),}}\end{array} \quad\quad (3) $$

由于物体有不同的尺度和高宽比，预测与宽度和高度无关的尺度$s_{h}$、$s_{w}$可以更好地处理现实场景中的复杂物体。对于大或小物体，SMCA动态生成不同的 $s_{h}$、$s_{w}$ 值，以便通过空间权重图 $G$ 来调整协同注意力图，从而聚合来自大物体所有部分的足够信息或抑制小物体的背景杂乱。

在得到上述对象信息后，SMCA生成类高斯权重图：

$$ G(i,j)=\mathrm{exp}(-\frac{(i-c\_w)^2}{\beta s^2\_w} - \frac{(j-c\_h)^2}{\beta s^2\_h}) \quad\quad (4) $$

其中 $(i,j)\in0,W\;\times0,H$ 是权重图 $G$ 的空间索引，$\beta$ 是用于调整类高斯分布bandwidth的超参数。一般来说，权重图 $G$ 对靠近中心的空间位置分配高重要性，对远离中心的位置分配低重要性。$\beta$ 可以通过手动调整来确保 $G$ 在训练开始时覆盖较大的空间范围，使得网络能够接收到更多信息的梯度。

Spatially-modulated co-attention

给定动态生成的空间先验 $G$，用于调整对象查询 $O\_Q$ 和自注意力编码特征 $E$ 之间的每个协同注意图 $C\_i$（在基础版的SMCA中，$G$ 是所有协同注意力头共享的）：

$$ C_{i}=\mathrm{softmax}\bigl(K_{i}^{T}Q_{i}/\sqrt{d}+\mathrm{log}G\bigr)V_{i} \quad\quad (5) $$

SMCA在空间图 $G$ 的对数和协同注意力点积 $K_{h}^{T}Q_{h}/\sqrt{d}$ 之间执行元素加法，然后对所有空间位置进行softmax归一化。这样，解码器的协同注意力在预测的边界框位置周围的权重更大，可以限制协同注意力的搜索空间，从而提高收敛速度。类高斯权重图如图 2 所示，限制协同注意力更多地关注预测边界框位置附近的区域，从而显着提高DETR的收敛速度。在SMCA基本版中，多个注意力头的协同注意力图 $C\_{i}$ 共享相同的类高斯权重图 $G$。

SMCA with multi-head modulation

论文还研究了针对不同的协同注意头使用不同的类高斯权重图调整协同注意特征。每个头从共享的中心 $c_{w}, c_{h}$ 开始，预测特定的相对于中心的偏移 $\Delta c_{w,i},\Delta c_{h,i}$ 和特定的尺度 $s_{w,i},s_{h,i}$。基于上述信息生成特定的类高斯空间权重图 $G\_{i}$，随后得到协同注意特征图 $C\_1\cdots C\_H$：

$$ C_{i}=\mathrm{softmax}(K_{i}^{T}Q_{i}/\sqrt{d}+\log G_{i})V\_{i}\quad\mathrm{for}\,i=1,\cdot\cdot,H. \quad\quad (6) $$

多个空间权重图可以强调不同的上下文并提高检测精度。

SMCA with multi-scale visual features

特征金字塔在目标检测框架中很流行，对比单尺度特征编码能带来显着的性能提升。因此，论文还将多尺度特征集成到SMCA中，通过在Transformer的编码器中用多尺度特征编码替换单尺度特征编码来进一步改进框架。

给定图像，从CNN主干网络提取采样率为16、32、64多尺度视觉特征 $f_{16}$、$f_{32}$、$f\_{64}$，编码器中的自注意力机制在不同尺度的所有特征像素之间传播和聚合信息。但由于多尺度特征的像素数量都非常大，自注意力操作的计算成本很高。为了解决这个问题，论文引入了尺度内自注意力编码，仅在每个尺度内的特征像素之间计算自注意力，而Transformer块的权重（自注意力和前馈子网络）在不同尺度上共享。论文的实证研究表明，跨尺度的参数共享增强了尺度内自注意力编码的泛化能力。

对于SMCA编码器的最终设计，采用2个尺度内自注意力编码块，然后是1个多尺度自注意力块，再接另外2个尺度内自注意力块。该设计具有与5个多尺度自注意力编码块非常相似的检测性能，但计算量要小得多。

给定编码器输出的多尺度编码特征 $E_{16}$、$E_{32}$、$E_{64}$，解码器执行协同注意力的简单解决方案是首先重新缩放并连接多尺度特征以形成单尺度特征图，然后在对象查询和生成的特征图之间计算协同注意力。然而，论文注意到一些对象查询可能只需要特定尺度的信息，并不总是需要所有尺度的信息。例如，低分辨率特征图 $E_{64}$ 中缺少小对象的信息，负责小物体的对象查询应该更有效地仅从高分辨率特征图中获取信息。与将每个边界框显式分配给特定尺度的特征图的传统方法（例如FPN）不同，论文使用可学习的尺度注意力机制自动地为每个框选择尺度：

$$ \alpha_{16},\alpha_{32},\alpha_{64}=\mathrm{Sortmax}(\mathrm{FC}(O_{q})), \quad\quad (7) $$

其中 $\alpha_{16}$、$\alpha_{32}$、$\alpha\_{64}$ 代表选择对应特征的重要程度。

为了在对象查询$O_{q}$和多尺度视觉特征 $E_{16}$、$E_{32}$、$E_{64}$ 之间计算协同注意力，首先通过不同的线性变化从对应的编码特征中获得注意力头 $i$ 的多尺度键 $K_{i,16},K_{i,32},K_{i,64}$ 和值特征 $V_{i,16},V_{i,32},V_{i,64}$，随后执行添加尺度选择权重和多尺度的SMCA计算：

$$ \begin{array}{l} {{C_{i,j}=\mathrm{Softmax(K_{i,j}^T\ Q_i/\sqrt{d}+log G\_i)V_{i,j}\odot \alpha\_{j}}}} \quad\quad\ (8) \ {{C_{i,j}=\underset{\mathrm{all}\ j}{\sum}\mathrm{C}_{i,j},\quad\mathrm{for} j\in{16,32,64},}} \quad\quad\quad\quad\quad (9) \end{array} $$

其中 $C\_{i,j}$ 代表查询和第 $j$ 个尺度视觉特征在第 $i$ 个协同注意力头中的协同注意力特征。通过这种尺度选择注意力机制，与每个对象查询最相关的尺度被平滑地选择，同时抑制其它尺度的视觉特征。

加上尺度内注意力和尺度选择注意力机制后，完整的SMCA可以比基础版更好地处理目标检测。

SMCA box prediction

在对象查询和编码图像特征之间进行协同注意力后，得到对象查询 $O\_q$ 的更新特征 $D\in\mathbb{R}^{N\times C}$。在原始DETR中，分别使用 3 层MLP和 1 个线性层来预测边界框和分类置信度：

$$ \begin{array}{l} {{\mathrm{Box}=\mathrm{Sigmoid}(\mathrm{MLP}(D))}} \quad\quad (10) \ {{\mathrm{Score}=\mathrm{FC}(D)}} \quad\quad\quad\quad\quad\quad\ \ (11) \end{array} $$

其中 $\mathrm{Box}$ 代表预测框在归一化坐标系中的中心、高度、宽度，${\mathrm{Score}}$ 代表分类预测。而在SMCA中，由于协同注意力被限制在初始预测的对象中心 $c_{h}^{\mathrm{norm}},c_{w}^{\mathrm{norm}}$ 周围，将初始中心作为约束边界框预测的先验进行校正预测：

$$ \begin{array}{c} {{\widehat{\mathrm{Box}}=\mathrm{MLP}(D)}} \ {{\widehat{\mathrm{Box}}:2=\widehat{\mathrm{Box}}:2 + c\_h^\widehat{norm},c\_w^\widehat{norm}}} \ {{\mathrm{Box}=\mathrm{Sigmoid(\widehat{Box})}}} \end{array} \quad\quad (12) $$

在Sigmoid之前进行中心坐标的相加，能确保边界框预测与SMCA中突出的协同注意区域高度相关。