COCO新记录：60.6AP！微软提出采用注意力机制进行检测头统一的Dynamic Head

来源丨极市平台

编辑丨极市平台

极市导读

本文提出一种新颖的动态头框架，它采用注意力机制将不同的目标检测头进行统一。COCO数据集上实验验证了所提方案的有效性与高效性。以ResNeXt-101-DCN为骨干，将目标检测的性能提升到了54.0AP。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

paper: https://arxiv.org/abs/2106.08322

code: https://github.com/microsoft/DynamicHead

Abstract

目标检测中定位与分类合并的复杂性衍生出了各式各样的算法，然而这些方法从不同的角度出发进行目标检测性能的提升，难以从一个统一的角度进行分析度量。

本文提出一种新颖的动态头框架，它采用注意力机制将不同的目标检测头进行统一。通过特征层次之间的注意力机制用于尺度感知，空间位置之间的注意力机制用于空间感知，输出通道内的注意力机制用于任务感知，该方法可以在不增加计算量的情况显著提升模型目标检测头的表达能力。

COCO数据集上实验验证了所提方案的有效性与高效性。以ResNeXt-101-DCN为骨干，我们将目标检测的性能提升到了54.0AP，取得了一个新的高度；更进一步，采用最新的Transformer骨干与额外数据，我们可以将COCO的指标推到一个新记录：60.6AP。

Method

为将尺度感知、空间位置感知以及任务感知同时嵌入一个目标检测头中，我们需要了解已有方法对目标检测头的改进。

给定以金字塔形式拼接的L级特征

mathcal{F}_{in} = {F_i}_{i=1}^L

，我们可以通过上采样/下采样将连续特征朝中间级特征缩放。尺度重调整的特征金字塔可以描述为4D张量

mathcal{F} in R^{L times H times W times C}

。我们进一步定义

S = H times W

并将上述特征reshape为3D张量

mathcal{F} in R^{Ltimes S times C}

。基于该表达，我们将对每个张量维度的作用进行探索：

目标尺度的差异对应了不同尺度的特征，改进

mathcal{F}

不同级的表达能力可以有效提升目标检测器的尺度感知能力；

不相似目标形状的不同几何变换对应了特征的不同空间位置，改进

mathcal{F}

不同空间位置的表达能力可以有效提升目标检测器的空间位置感知能力；

目标表达与任务的多样性对应了不同通道特征，改进

mathcal{F}

不同通道的表达能力可以有效提升目标检测的任务感知能力。

我们发现：上述方向可以统一到一个高效注意力学习问题中。本文也是首个尝试采用多注意力将三个维度组合构建统一头并最大化其性能的工作。

Dynamic Head: Unifying with Attentions

给定特征张量

mathcal{F} in R^{L times S times C}

，自注意力的广义形式可以描述如下：

W(mathcal{F}) = pi(mathcal{F}) cdot mathcal{F}

一种最朴素的方案是采用全连接层，但是在所有维度上直接学习注意力函数会导致过多的计算量，且由于高维度问题而不切实际。

相反，我们将注意力函数转换为三个序列注意力，每个注意力仅聚焦一个维度：

W(mathcal{F}) = pi_C(pi_S(pi_L(mathcal{F})cdot mathcal{F}) cdot mathcal{F}) cdot mathcal{F}

Scale-aware Attention

pi_L

我们首先引入尺度感知注意力基于其语义重要性对不同尺度特征进行融合：

pi_L(mathcal{F}) cdot mathcal{F} = sigma(f(frac{1}{SC}sum_{S,C}mathcal{F})) cdot mathcal{F}

其中，

f(cdot)

为线性函数，采用

1times 1

卷积近似，

sigma(x)

为hard-sigmoid激活函数。

Spatial-aware Attention

pi_S

我们接下来引入另一个空间位置感知注意力模块以聚焦不同空间位置的判别能力。考虑到S的高纬度，我们对其进行解耦：首先采用形变卷积对注意力学习稀疏化，然后进行特征跨尺度集成：

pi_S(mathcal{F}) cdot mathcal{F} = frac{1}{L} sum_{l=1}^L sum_{k=1}^K w_{l,k} cdot mathcal{F}(l;p_k Delta_k;c) cdot Delta m_k

其中，K为稀疏采样位置数。其他参数信息与形变卷积中的参数信息类似，故略过。

Task-aware Attention

pi_C

为促进联合学习与目标表达能力的泛化性，我们设计了一种任务感知注意力。它可以动态开关特征通道以辅助不同任务：

pi_C(mathcal{F}) cdot mathcal{F} = max(alpha^1(mathcal{F})cdot mathcal{F}_c beta^1(mathcal{F}), alpha^2(mathcal{F})cdot mathcal{F}_C beta^2(mathcal{F}))

其中，

[alpha^1,alpha^2,beta^1,beta^2]^T = theta(cdot)

为超参数，用于控制激活阈值，而

theta(cdot)

类似于DyReLU。

最后，将上述注意力机制以序列方式实施，我们可以堆叠实施多次。下图给出了DynamicHead的配置示意图。

总而言之，基于所提DynamicHead的目标检测框架图见下图。任意类型的骨干网络均可用于提取特征金字塔并进一步缩放到统一尺度构建3D张量，然后将其作为DynamicHead的输入。接下来，多个包含尺度感知、空间位置感知以及任务感知的DyHead模块串行堆叠。动态头的输出可以用于不同的任务，比如分类、中心/方框回归等等。

以上图底部为例，我们可以看到不同类型注意力的输出。可以看到：

源自骨干的初始特征由于域差异存在噪声；
经过尺度感知注意力模块处理后，特征就对不同尺度的目标变的更敏感；
经过空间位置感知注意力模块处理后，特征变的更稀疏，聚焦于不同位置的前景目标；
经过任务感知注意力模块处理后，特征会基于不同下游任务形成不同的激活。

以上可视化结果很好的验证了每个注意力模块的有效性。

Generalizing to Existing Detectors

接下来，我们将验证如何将所提DynamicHead集成到现有检测器中以提升其性能。

One-stage Detector 单阶段检测器通过在特征图上稠密采样方式预测目标位置，这种方式简化了目标检测器的设计。经典的单阶段检测器(如RetinaNet)包含一个骨干网络提取稠密特征、多个任务相关子网络处理不同任务。正如DyReLU一文所提到：目标分类子网络的行为与回归子网络的差异甚大。不同于这种经典方案，我们仅在骨干之后附加一个统一分支，受益多注意力机制，它可以同时处理多个任务。通过这种方式，架构可以进一步简化，同时提升了效率。将DyHead与单阶段检测器集成的方式可参考见下图b。

Two-stage Detector 两阶段检测器采用Region Proposal与ROI-pooling层从特征金字塔中提取中间表达。为耦合该特性，我们首先在ROI-pooling添加实施尺度感知注意力与空间位置感知注意力，然后采用任务感知注意力替换原始的全连接层。具体实现方式可参见上图c。

Experiments

我们在MS-COCO数据集上对所提方案的有效性进行了验证，检测器的架构以Mask R-CNN为基准，动态头的训练则参考了ATSS框架。

Main Results

我们先来看一下总体结果，见上面两个表，从中可以看到：

所提方案以较大的优势超越了其他已有方案；
相比同配置的最佳检测器BorderDet，采用ResNet101骨干时所提方案取得了1.1AP指标提升；采用ResNeXt-x4d-1-1骨干时取得了1.2AP指标提升；
相比仅采用多尺度寻览的方法，所提方案仅需2x训练即可取得新的SOTA指标 52.3AP；
相比EfficientDet与SpineNet，所提方法训练时间更少(仅1/20);
当同时采用多尺度训练与测试时，所提方法取得了新的SOTA指标54.0AP，以1.3AP指标优于此前最佳。