二十九章：Shunted Self-Attention via Multi-Scale Token Aggregation ——通过多尺度令牌聚合的抑制型自注意力

&原文信息

原文名：《Shunted Self-Attention via Multi-Scale Token Aggregation 》

引用：Ren S, Zhou D, He S, et al. Shunted self-attention via multi-scale token aggregation[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 10853-10862.

原文链接：https://openaccess.thecvf.com/content/CVPR2022/papers/Ren_Shunted_Self-Attention_via_Multi-Scale_Token_Aggregation_CVPR_2022_paper.pdfhttps://openaccess.thecvf.com/content/CVPR2022/papers/Ren_Shunted_Self-Attention_via_Multi-Scale_Token_Aggregation_CVPR_2022_paper.pdf

0.摘要

        最近的视觉Transformer（ViT）模型在各种计算机视觉任务中取得了令人鼓舞的结果，这要归功于其通过自注意力机制对图像补丁或标记建模长程依赖性的能力。然而，这些模型通常在每个层内为每个标记特征指定相似的感受野。这种约束不可避免地限制了每个自注意力层在捕捉多尺度特征方面的能力，从而导致在处理具有多个不同尺度对象的图像时性能下降。为了解决这个问题，我们提出了一种新颖的通用策略，称为抑制型自注意力（SSA:Shunted Self-Attention），它允许ViTs在每个注意力层上对混合尺度的注意力进行建模。SSA的关键思想是在计算自注意力矩阵之前，将不同的感受野大小注入到标记中：选择性地合并标记以表示更大的对象特征，同时保留某些标记以保留细粒度的特征。这种新颖的合并方案使得自注意力能够学习不同尺寸对象之间的关系，同时减少了标记数量和计算成本。在各种任务上进行的大量实验证明了SSA的优越性。具体而言，基于SSA的Transformer模型在ImageNet上实现了84.0%的Top-1准确率，在模型大小和计算成本仅为现有Focal Transformer的一半的情况下超越了现有技术，在COCO上超过Focal Transformer 1.3 mAP，在ADE20K上超过Focal Transformer 2.9 mIOU，参数和计算成本相似。代码已经在https://github.com/OliverRensu/Shunted-Transformer上发布。

1.引言

        最近的视觉Transformer（ViT）模型在各种计算机视觉任务中展示了卓越的性能，例如分类、目标检测、语义分割和视频动作识别。与侧重于局部建模的卷积神经网络不同，ViTs将输入图像划分为一系列补丁（标记），并通过全局自注意力逐步更新标记特征。自注意力可以有效地对标记之间的长程依赖进行建模，并通过从其他标记中聚合信息逐步扩大它们的感受野大小，这在很大程度上解释了ViTs的成功。

        然而，自注意力机制也带来了昂贵的内存消耗，其与输入标记数量呈二次关系。因此，最先进的Transformer模型采用了各种下采样策略来减小特征大小和内存消耗。例如，[7]的方法在第一层进行16×16的下采样投影，并在得到的粗粒度和单尺度特征图上计算自注意力；这样做会导致特征信息的损失，从而不可避免地降低模型性能。其他方法则努力在高分辨率特征上计算自注意力，并通过对标记进行空间缩减来减少计算成本。然而，这些方法往往会在一个自注意力层内合并过多的标记，从而导致小物体和背景噪声的标记混合。这种行为反过来使得模型在捕捉小物体方面不够有效。

        此外，先前的Transformer模型在注意力层内很大程度上忽视了场景对象的多尺度特性，使它们对涉及不同尺寸对象的真实场景变得脆弱。这种不足在技术上归因于它们的底层注意力机制：现有方法仅依赖于标记的静态感受野和在一个注意力层内的统一信息粒度，因此无法同时捕捉不同尺度的特征。

        为了解决这个限制，我们引入了一种新颖且通用的自注意力方案，称为抑制型自注意力（SSA），它明确允许同一层内的自注意力头分别考虑粗粒度和细粒度特征。与之前将太多标记合并或无法捕捉小物体的方法不同，SSA在同一层的不同自注意力头中同时有效地对各种尺度的对象进行建模，既具有有利的计算效率，又能保留细粒度的细节。

        我们在图2中展示了普通自注意力（来自ViT）、辅助下采样注意力（来自PVT）和提出的SSA之间的定性比较。当不同的注意力应用于相同大小的特征图时，ViT可以捕捉到细粒度的小物体，但计算成本非常高（图2(a)）；PVT减少了计算成本，但其注意力仅限于粗粒度的较大物体（图2(b)）。相比之下，提出的SSA保持了较轻的计算负荷，同时考虑了混合尺度的注意力（图2(c)）。有效地，SSA不仅精确关注粗粒度的大物体（例如沙发），还关注细粒度的小物体（例如灯和风扇），甚至一些位于角落的物体，而这些物体被PVT遗漏了。我们还在图3中显示了注意力图的视觉比较，以突出SSA学习到的尺度自适应注意力。

        SSA通过将多个注意力头分成几个组来实现多尺度的关注机制。每个组负责特定的注意力粒度。对于细粒度的组，SSA学习聚合少量标记并保留更多的局部细节。对于其余的粗粒度头组，SSA学习聚合大量标记，从而降低计算成本，同时保留捕捉大物体的能力。多尺度组共同学习多粒度信息，使得模型能够有效地建模多尺度对象。

        如图1所示，我们展示了通过堆叠多个基于SSA的块获得的Shunted Transformer模型的性能。在ImageNet上，我们的Shunted Transformer在减小模型大小的同时，优于最先进的Focal Transformers [29]。当缩小到微小尺寸时，Shunted Transformer的性能与DeiT-Small [20]相似，但参数量仅为其50%。对于目标检测、实例分割和语义分割，在COCO和ADE20K上，Shunted Transformer在模型大小相似的情况下持续优于Focal Transformer。

总结起来，我们的贡献如下：

• 我们提出了抑制型自注意力（SSA），通过多尺度标记聚合，将多尺度特征提取统一到一个自注意力层中。我们的SSA能够自适应地在大物体上合并标记以提高计算效率，并保留小物体的标记。
• 基于SSA，我们构建了我们的抑制型Transformer，能够高效地捕捉多尺度对象，尤其是小型和远离的孤立对象。
• 我们在包括分类、目标检测和分割在内的各种任务上评估了我们提出的抑制型Transformer。实验结果表明，在相似的模型大小下，我们的抑制型Transformer始终优于以前的Vision Transformer模型。

图1. 最近SOTA的CNN和Transformer模型在ImageNet上的Top-1准确率。我们提出的抑制型Transformer在所有基线模型中表现最好，包括最近的SOTA Focal Transformer（基准大小）。值得注意的是，它以两倍较小的模型大小达到了与DeiT-S相竞争的准确性。

 

图2. ViT（Vision Transformer）、PVT（Pyramid Vision Transformer）和我们的SSA在相同特征图大小下的不同注意力机制的比较。圆圈的数量表示自注意计算中涉及的标记数量，并反映了计算成本。圆圈的大小表示相应标记的感受野大小。与ViT和PVT不同，我们的方法通过自适应地合并大物体上的圆圈来增强计算效率，并同时考虑不同尺度的物体。

图3. PVT和我们模型的注意力图。PVT倾向于只关注大物体，如沙发和床，而与之相比，我们的模型能够准确地捕捉到与大物体并存的小物体，如灯光。

2.相关工作

2.1.在CNNs中的自注意力

        卷积层的感受野通常较小且固定。虽然扩张卷积可以扩大感受野，可变形卷积允许在卷积核中进行一些偏移，但它们很难自适应和灵活地扩展到整个特征图。受自注意力在NLP领域中的先驱性作用的启发，一些研究将自注意力或非局部块引入计算机视觉领域的卷积神经网络中。这些注意力通常应用于深层，特征图的尺寸较小且经过多个卷积层进行预处理。因此，它们不会带来太多额外的计算成本，但只能带来有限的性能提升。

2.2.视觉Transformer

        Vision Transformer (ViT)模型[7]直接在非常深层次应用自注意力来构建一个无卷积的神经网络模型。自从ViT模型问世以来，许多后续工作通过更复杂的数据增强或知识蒸馏来改善模型的分类性能[17,20]。因为自注意力的计算复杂度与标记数量的平方成正比，所以很难直接应用于大量的标记。因此，这些ViT模型通常将图像分割为非重叠且较大尺寸的块（标记）。但是这种分割太粗糙，会丢失很多细粒度的信息。为了保留细粒度特征，这些模型通常对特征图进行下采样，并在低分辨率特征上操作。然而，这种折衷会阻碍它们在分割和检测等密集预测任务中的部署。

2.3.高效ViT变种

        为了使自注意力能够应用于较大尺寸的特征图，最近的研究提出了两种减少计算成本的解决策略[2,14,25,26,29]：

        (1)将特征图分割为区域，并在区域内进行局部自注意力计算，或者

        (2)合并标记以减少标记数量。局部自注意力的代表性工作是Swin Transformer[14]，它将特征图分割为非重叠的方形区域，并在局部进行自注意力计算。

        然而，为了通过自注意力建模全局依赖关系，这些局部注意力需要在图像上移动窗口或堆叠大量层以获得全局感受野。关于标记合并的策略，PVT（金字塔视觉Transformer）[26]设计了一种空间缩减注意力来合并关键标记和查询标记。然而，PVT和类似的模型往往在这种空间缩减中合并了过多的标记。这使得小物体的细粒度信息与背景混合在一起，损害了模型的性能。因此，我们提出了抑制型自注意力，它能够同时保留粗粒度和细粒度的细节，并在图像标记之间进行全局依赖建模。

3.方法

        我们提出的Shunted Transformer的总体架构如图4所示。它是基于新颖的抑制型自注意力（SSA）块构建的。我们的SSA块与ViT中的传统自注意力块之间有两个主要区别：

        1）SSA引入了一个抑制型注意机制，用于每个自注意力层，以捕捉多粒度信息并更好地建模不同尺寸的对象，特别是小物体；

        2）它通过增强跨标记交互，增强了点对点前馈层提取局部信息的能力。

        此外，我们的Shunted Transformer采用了一种新的块嵌入方法，以获得更好的输入特征图用于第一个注意力块。接下来，我们逐一详细介绍这些创新之处。

3.1.抑制型Transformer块

        在提出的Shunted Transformer的第i阶段，有Li个Transformer块。每个Transformer块包含一个自注意力层和一个前馈层。为了在处理高分辨率特征图时减少计算成本，PVT [26]引入了空间缩减注意力（SRA）来替代原始的多头自注意力（MSA）。然而，SRA往往会在一个自注意力层内合并过多的标记，并且只提供单一尺度的标记特征。这些限制阻碍了模型在捕捉多尺度对象，特别是小尺寸对象方面的能力。因此，我们引入了抑制型自注意力，在一个自注意力层内并行学习多粒度信息。

3.1.1抑制型自注意力

        首先，将输入序列F ∈ Rh×w×c投影到查询（Q）、键（K）和值（V）张量中。然后，多头自注意力采用H个独立的注意力头并行计算自注意力。为了减少计算成本，我们遵循PVT [26]的做法，而不是像Swin Transformer [14]那样将{Q，K，V}分割为区域，而是减少K和V的长度。如图5所示，我们的抑制型自注意力与PVT的空间缩减注意力（SRA）不同，因为同一个自注意力层中的注意力头的K和V的长度不相同。相反，长度在不同的头中变化，以捕捉不同粒度的信息。这实现了多尺度标记聚合（MTA）。具体而言，键K和值V被下采样到由i索引的不同大小，以获得不同的头。

 这里的MAT(·;ri)是第i个头中的多尺度标记聚合层，其下采样率为ri。在实践中，我们采用卷积层，卷积核大小和步长均为ri，来实现下采样。WQi，WKi，WVi是第i个头中线性投影的参数。在注意力头中的一层中，ri有不同的变体。因此，键和值可以在自注意力中捕捉不同的尺度。LE(·)是MTA中通过深度卷积对值V进行本地增强的组件。与空间缩减[26]相比，MTA保留了更细粒度和低级别的细节。然后，通过以下方式计算抑制型自注意力：

其中dh是维度大小。由于多尺度键和值的存在，我们的抑制型自注意力在捕捉多尺度对象方面更加强大。计算成本的减少取决于r的值，因此，我们可以很好地定义模型和r来权衡计算成本和模型性能。当r变大时，K和V中的更多标记被合并，K和V的长度变短，因此计算成本较低，但仍然保留了捕捉大对象的能力。相反，当r变小时，保留了更多细节，但带来了更多的计算成本。在一个自注意力层中整合不同的r使其能够捕捉多粒度特征。

图4左侧：我们Shunted Transformer的总体架构。右侧：我们抑制型自注意力块的详细信息。

 

图5.将我们的抑制型自注意力与ViT、Swin和PVT中的自注意力进行比较。ViT在小尺寸特征图上全局应用自注意力。Swin Transformer在大尺寸特征图上的小区域内应用局部自注意力。PVT使用大步长融合键和值。与之不同的是，我们的抑制型自注意力通过多尺度标记聚合获取不同尺寸的键和值。

 

3.1.2.细节特定的前馈层

        在传统的前馈层中，全连接层是逐点计算的，无法学习到跨标记的信息。在这里，我们旨在通过在前馈层中的细节来补充局部信息。如图6所示，我们通过在前馈层的两个全连接层之间添加我们的数据特定层来补充前馈层中的局部细节：

 其中DS(·;θ)是具有参数θ的细节特定层，在实践中通过深度卷积实现。

表1.我们Shunted Transformer的模型变体。C和N分别表示块的维度和数量。head表示头的数量。

3.2.（图像）块嵌入

        Transformer最初设计用于处理序列数据。如何将图像映射到序列对模型的性能至关重要。ViT直接将输入图像分割成16×16的非重叠补丁。最近的研究发现，使用卷积在补丁嵌入中提供了更高质量的标记序列，并且帮助Transformer"看得更好"，而不是传统的大步长非重叠补丁嵌入。因此，一些工作采用了重叠的补丁嵌入，比如使用7×7的卷积核[14,26]。在我们的模型中，根据模型的大小，我们采用不同的带有重叠的卷积层。我们使用步长为2和零填充的7×7卷积层作为补丁嵌入的第一层，并根据模型的大小添加额外的步长为1的3×3卷积层。最后，使用步长为2的非重叠投影层生成大小为H/4×W/4的输入序列。

3.3.架构细节和变种

        给定大小为H×W×3的输入图像，我们采用上述的补丁嵌入方案，以获得长度为H4×W4和标记维度为C的更具信息量的标记序列。遵循之前的设计[2,14,26,29]，我们的模型分为四个阶段，每个阶段包含多个Shunted Transformer块。在每个阶段中，每个块输出相同尺寸的特征图。我们采用步长为2的卷积层（线性嵌入）来连接不同的阶段，并且在输入到下一个阶段之前将特征图的大小减半，但维度加倍。因此，每个阶段的输出有四个特征图F1、F2、F3、F4，Fi的大小为2iH+1×2Wi+1×(C×2i−1)。

        我们提出了三种不同配置的模型，以便在相似的参数和计算成本下进行公平比较。如表1所示，head和Ni分别表示一个块中头的数量和一个阶段中块的数量。这些变体仅来自不同阶段的层数。具体而言，每个块中头的数量设置为2、4、8、16。补丁嵌入中的卷积层范围从1到3。

图6.比较ViT（左侧）、PVT（右侧）和我们的细节特定的前馈层。我们在前馈层中补充了细粒度的跨标记细节信息。

4.实验

        为了评估我们的Shunted Transformer的有效性，我们将我们的模型应用于ImageNet-1K[6]分类、COCO[13]目标检测和实例分割、ADE20K[36]语义分割。此外，我们通过消融研究评估了我们模型的不同组成部分的效果。

表2.在ImageNet-1K分类任务上不同骨干网络的比较。除了EfficientNet（EffNet-B4*）外，所有模型都在输入尺寸为224×224上进行训练和评估。 

4.1.在ImageNet-1K数据集上的图像分类

        我们首先在ImageNet-1K上评估我们的模型，并将其与最近的SOTA CNN和基于Transformer的模型进行比较。为了公平比较，我们遵循DeiT[20]和PVT[26]的相同训练策略。具体而言，我们使用AdamW作为优化器，权重衰减为0.05。整个训练过程共300个epoch，输入尺寸为224×224，批量大小为1024。学习率采用余弦学习率衰减设置为1×10^(-3)。数据增强和正则化方法遵循DeiT[20]，包括随机裁剪、随机翻转、标签平滑[18]、Mixup[33]、CutMix[32]和随机擦除[35]。

        如表2所示，通过与其他CNN骨干网络在相似的参数和计算成本下进行比较，我们的模型是第一个能够与使用更大输入分辨率的EfficientNet取得可比较结果的基于Transformer的模型。值得注意的是，尽管RegNet和EfficientNet来自神经架构搜索，但我们手动设计的Transformer仍然表现优于它们。

        接下来，我们将我们的模型与Transformer骨干网络进行比较。我们的微型模型在性能上与Transformer基准模型（DeiT-S）相当，但是只需要一半的参数（22M→11M）和计算成本（4.6G→2.1G FLOPs）。当我们的模型尺寸增长到与DeiT-S相似时，性能超过了3%。与最近的SOTA模型（如Swin和Twin）相比，我们的模型始终表现出更好的性能。具体而言，我们的小型模型在减小了20%的模型大小的同时，比现有的最先进模型Focal Transformer Tiny提高了0.7%的性能。当模型尺寸增大时，我们的基础模型相对于Focal Transformer只需要一半的参数和计算成本，就能达到最先进的性能水平。

4.2.目标检测和实例分割

        我们在COCO 2017数据集上评估了目标检测和实例分割的模型。我们将我们提出的Shunted Transformer作为骨干网络，并将其插入到Mask R-CNN中。我们将其与其他最先进的骨干网络进行比较，包括ResNet、Swin Transformer、Pyramid Vision Transformer、Twin和Focal Transformer。我们遵循Swin的相同设置：在ImageNet-1K上进行预训练，然后在COCO上进行微调。在微调阶段，我们使用了两个训练计划：1×训练计划共12个epoch和3×训练计划共36个epoch。在1×训练计划中，输入图像的较短边将被调整为800，同时保持较长边不超过1333。在3×训练计划中，我们采用多尺度训练策略，将较短边的大小调整为480至800之间。我们使用AdamW作为优化器，权重衰减为0.05。批量大小为16，初始学习率为10^(-4)。

        在表3中，我们使用Mask-RCNN进行目标检测，并报告了不同CNN和Transformer骨干网络的bbox mAP（APb）。在可比参数下，我们的模型以显著的差距超越了先前的SOTA模型。对于目标检测，在1×训练计划中，我们的微型模型相比ResNet-50提高了9.1个点，相比Focal Transformer仅使用85%的模型大小，提高了2.3个点。此外，在3×训练计划和多尺度训练下，我们的骨干网络在平均上仍然比CNN骨干网络高出7.7个点，比Transformer骨干网络高出1.6个点。我们在实例分割中也得到了类似的结果。我们在表3中报告了mask mAP（APm）。在1×训练计划和3×训练计划中，我们的模型相比ResNet-50分别提高了8.1个点和1.5个点，相比Focal Transformer分别提高了8.1个点和1.5个点。我们的模型在更小的模型大小下实现了这些优越的性能，清楚地展示了其通过学习多粒度标记和处理多尺度视觉对象存在的有效性的益处。

        我们还在表4中报告了RetinaNet的结果。在1×和3×训练计划中，我们的模型以最少的参数超越了所有先前的模型。与PVT相比，我们的模型在所有小尺寸、中尺寸和大尺寸对象上都带来了改进，这显示了我们的分流自注意力在捕捉多尺度对象方面的强大能力。

表3.在COCO数据集上使用Mask R-CNN进行目标检测和实例分割。只有3×训练计划具有多尺度训练。所有的骨干网络都在ImageNet-1K上进行了预训练。*表示这些方法尚未经过同行评审。

表4.在COCO数据集上使用RetinaNet进行目标检测。只有3×训练计划具有多尺度训练。所有的骨干网络都在ImageNet-1K上进行了预训练。*表示这些方法尚未经过同行评审。

4.3.在ADE20K数据集上的语义分割

        我们在ADE20K [36]数据集上评估了我们模型在语义分割任务上的性能。该数据集包含20,210张训练图片、2,000张验证图片和3,352张测试图片，共涵盖了150个细粒度的语义类别。我们报告了使用多尺度测试和不使用多尺度测试的mIOU。我们选择了UperNet和Semantic FPN作为主要框架，并采用了不同的骨干网络架构。我们遵循Focal Transformer和mmsegmentation [3]的默认设置。对于UperNet，我们使用AdamW优化器进行160K次迭代训练，权重衰减率为0.01。学习率为6×10−5，在训练开始时进行了1500次迭代的预热，并采用线性学习率衰减。数据增强包括随机翻转、随机缩放和随机光度失真。训练时的输入尺寸为512×512，包括单尺度和多尺度(MS)测试。对于SemanticFPN，我们使用AdamW优化器进行80K次迭代训练，权重衰减率为0.0001，学习率也为0.0001。

        结果报告在表5中。我们的Shunted Transformer在所有框架中表现出色，与先前的最先进方法相比，具有更大的差距和更少的参数。具体而言，当使用Semantic FPN时，我们的模型比Swin Transformer高出6.7个mIOU，并且模型大小减少了20%。当框架为UpperNet时，我们的Shunted Transformer比Focal Transformer高出3.1%和2.9%。分割结果显示了我们Shunted Transformer的优越性。

        我们还将SegFormer [28]作为框架，并将我们的骨干网络与SegFormer中的MiT进行了比较。结果报告在表6中。在更少的参数情况下，我们的方法在mIoU上比SegFormer提升了1.8个单位。

表5.在ADE20K数据集上，比较不同骨干网络在Semantic FPN和UpperNet框架中的分割性能。

表6.在ADE20K数据集上，比较SegFormer框架中不同骨干网络的性能。

4.4.消融研究

块编码:图像到标记映射的功能，即图像的分块嵌入头，是许多最近研究[9,10,24]的重点。它们发现，设计良好的头部可以为transformer模型提供更好的输入序列。我们评估了我们的非重叠头在ViT中的分块嵌入以及在Swin和PVT中的重叠头的影响。结果如表7所示。使用像重叠或我们的分块嵌入这样更复杂的头部，计算成本和模型大小只略微增长，但性能显著提高。具体来说，通过有限的额外参数，从传统的非重叠头部或重叠头部转换为分块嵌入，模型分别获得了1.4%和0.3%的性能提升。

标记聚合函数我们提出了一种新的标记聚合函数，用于合并多尺度对象的标记，并同时保持全局和局部信息。从表8可以看出，我们的新型标记聚合函数与卷积空间降维具有类似的计算量，但获得了更多的改进。

细节特定的前馈在前馈层[20]中，所有的操作都是逐点操作，没有跨标记的操作，因此，补充跨标记和局部信息将显著改善前馈层的学习能力。在表9中，我们比较了我们的新的细节特定前馈层、传统的前馈层[20]和卷积前馈层[25]在ViT和我们的模型中的表现。细节特定的前馈层在性能上一直优于传统的前馈层，这表明在前馈层中补充局部细节的有效性。

表8.不同标记聚合函数在ImageNet上的Top-1准确率。

表9.在相似的参数数量和浮点运算量（FLOPs）下，细节特定的前馈层在ImageNet上比传统的前馈层提供了更高的Top-1准确率。

5.总结

        在本文中，我们提出了一种新颖的Shunted Self Attention (SSA)方案，以明确考虑多尺度特征。与以往只关注一个注意力层中的静态特征图的工作不同，我们在一个自注意力层中维护了与多尺度对象相关的各种尺度的特征图。大量实验证明了我们的模型作为各种下游任务的骨干的有效性。具体而言，所提出的模型优于以往的Transformer，并在分类、检测和分割任务上取得了最先进的结果。