当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

Transformer 近来在各种视觉任务上表现出卓越的性能,感受野赋予 Transformer 比 CNN 更强的表征能力。然而,简单地扩大感受野会引起一些问题。一方面,使用密集注意力(例如 ViT)会导致过多的内存和计算成本,并且特征可能会受到超出兴趣区域的无关部分的影响;另一方面,PVT 或 Swin Transformer 中采用的稀疏注意力与数据无关,可能会限制对远程(long range)关系建模的能力。

为了缓解这些问题,清华大学、AWS AI 和北京智源人工智能研究院的研究者提出了一种新型可变形自注意力模块,其中以数据相关的方式选择自注意力中键值对的位置。这种灵活的方案使自注意力模块能够专注于相关区域并捕获更多信息特征。

在此基础上,该研究提出了可变形注意力 Transformer(Deformable Attention Transformer,DAT),一种具有可变形注意力的通用主干网络模型,适用于图像分类和密集预测任务。该研究通过大量基准测试实验证明了该模型的性能提升。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

论文地址:https://arxiv.org/abs/2201.00520v1


可变形注意力 Transformer

现有的分层视觉 Transformer,特别是 PVT 和 Swin Transformer 试图解决过多注意力的挑战。前者的下采样技术会导致严重的信息损失,而后者的 Swin 注意力导致感受野的增长要慢得多,这限制了对大型物体进行建模的潜力。因此,需要依赖于数据的稀疏注意力来灵活地对相关特征进行建模,从而导致首先在 DCN [9] 中提出可变形机制。

然而,在 Transformer 模型中实现 DCN 是一个不简单的问题。在 DCN 中,特征图上的每个元素单独学习其偏移量,其中 H ×W ×C 特征图上的 3 × 3 可变形卷积具有 9 HWC 的空间复杂度。如果在注意力模块中直接应用相同的机制,空间复杂度将急剧上升到 N_qN_kC,其中 N_q、N_k 是查询和键的数量,通常与特征图大小 HW 具有相同的比例,带来近似于双二次的复杂度。

尽管 Deformable DETR [54] 已经设法通过在每个尺度上设置较少数量的 N_k = 4 的键来减少这种开销,并且可以很好地作为检测头,但由于不可接受的信息丢失(参见附录中的详细比较),在骨干网络中关注如此少的键效果不佳。与此同时,[3,52] 中的观察表明,不同的查询在视觉注意力模型中具有相似的注意力图。因此,该研究选择了一个更简单的解决方案,为每个查询共享移位键和值,以实现有效的权衡。


模型架构

该研究在 Transformer(等式 (4))中的可变形注意力替换了 vanilla MHSA,并将其与 MLP(等式 (5))相结合,以构建一个可变形的视觉 transformer 块。在网络架构方面, DAT 与 [7, 26, 31, 36] 共享类似的金字塔结构,广泛适用于需要多尺度特征图的各种视觉任务。如下图 3 所示,形状为 H × W × 3 的输入图像首先被步长为 4 的 4 × 4 非重叠卷积嵌入,然后一个归一化层获得当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT


为了构建分层特征金字塔,主干包括 4 个阶段,步幅逐渐增加。在两个连续的阶段之间,有一个步长为 2 的非重叠 2×2 卷积,对特征图进行下采样,将空间大小减半并将特征维度加倍。

在分类任务中,该研究首先对最后阶段输出的特征图进行归一化,然后采用具有池化特征的线性分类器来预测对数;在对象检测、实例分割和语义分割任务中,DAT 在集成视觉模型中扮演主干的角色,以提取多尺度特征。该研究为每个阶段的特征添加一个归一化层,然后将它们输入到以下模块中,例如对象检测中的 FPN [23] 或语义分割中的解码器。


实验

该研究在 3 个数据集上进行了实验,以验证提出的 DAT 的有效性。该研究展示了在 ImageNet-1K [10] 分类、COCO 目标检测和 ADE20K 语义分割任务上的结果。此外,该研究提供了消融研究和可视化结果,以进一步展示该方法的有效性。


ImageNet-1K 分类

ImageNet-1K [10] 数据集有 128 万张用于训练的图像和 5 万张用于验证的图像。研究者在训练分割上训练 DAT 的三个变体,并报告验证分割上的 Top-1 准确度,并与其他 Vision Transformer 模型进行比较。

该研究在下表 2 中给出了有 300 个训练 epoch 的结果。与其他 SOTA 视觉 Transformer 模型相比, DAT 在具有相似计算复杂性的情况下在 Top-1 精度上实现了显著提高。DAT 在所有三个尺度上都优于 Swin Transformer [26]、PVT [36]、DPT [7] 和 DeiT [33]。没有在 Transformer 块 [13, 14, 35] 中插入卷积,或在补丁嵌入 [6, 11, 45] 中使用重叠卷积,DAT 比 Swin Transformer [26] 实现了 +0.7、+0.7 和 +0.5 的增益。在 384 × 384 分辨率下进行微调时,该模型继续比 Swin Transformer 性能好 0.3。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

COCO 目标检测

COCO 目标检测和实例分割数据集有 118K 的训练图像和 5K 的验证图像。该研究使用 DAT 作为 RetinaNet [24]、Mask R-CNN [17] 和 Cascade Mask R-CNN [2] 框架中的主干,以评估该方法的有效性。该研究在 ImageNet-1K 数据集上对该模型进行 300 个 epoch 的预训练,并遵循 Swin Transformer [26] 中类似的训练策略来公平地比较该方法。该研究在 1x 和 3x 训练计划中报告在 RetinaNet 模型上的 DAT。如下表 3 所示,在微型和小型模型中,DAT 的性能优于 Swin Transformer 1.1 和 1.2 mAP。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

当在两阶段检测器(例如 Mask R-CNN、Cascade Mask R-CNN)中实现时,DAT 模型在不同尺寸的 Swin Transformer 模型上实现了一致的改进,如下表 4 所示。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

下表 5 给出了在验证集上各种方法的 mIoU 分数。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT


消融实验

为了验证 DAT 模型中关键组件设计的有效性, 该研究进行了消融实验,报告了基于 DAT-T 的 ImageNet-1K 分类结果。对于几何信息开发,该研究首先评估了所提可变形偏移和可变形相对位置嵌入的有效性,如下表 6 所示。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT

对于不同阶段的可变形注意力,该研究用不同阶段的可变形注意力替换了 Swin Transfomer [26] 的移位窗口注意力。如下表 7 所示,仅替换最后阶段的注意力提高了 0.1,替换最后两个阶段导致性能增益为 0.7(达到 82.0 的整体准确度)。然而,在早期阶段用更多可变形注意力替换会略微降低准确性。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT


可视化

该研究在 DAT 中可视化学习变形位置的示例,以验证该方法的有效性。如下图 4 所示,采样点描绘在对象检测框和实例分割掩码的顶部,从中可以看到这些点已转移到目标对象。

当Swin Transformer遇上DCN,清华可变形注意力Transformer模型优于多数ViT


上一篇:将执行器「缝」进衣袖,斯坦福无接触设备,远程传递触摸信息


下一篇:迈入全自主手术的第一步!机器人独立完成猪腹腔镜肠吻合,效果好于人类