论文阅读|Shunted Self-Attention

Shunted Self-Attention via Multi-Scale Token Aggregation

代码

目录

Abstract

Introduction

Contributions

Related Work

Self-Attention in CNNs

Vision Transformer

Efficient ViT Variants

Method

3.1. Shunted Transformer Block

3.1.1 Shunted Self-Attention

3.1.2 Detail-specific Feedforward Layers

3.2. Patch Embedding

3.3. Architecture Details and Variants

Experiments

Image Classification on ImageNet-1K

Object Detection and Instance Segmentation

​Semantic Segmentation on ADE20K

Ablation Studies

Patch Embedding

Token Aggregation Function

Detail-specific Feed-Forward

Conclusion

Abstract

最近的 Vision Transformer (ViT) 模型在各种计算机视觉任务中展示了令人鼓舞的结果，这要归功于它们通过自我注意对图像块或标记的远程依赖关系进行建模的能力。然而，这些模型通常指定每一层内每个标记特征的相似感受野。这种约束不可避免地限制了每个自注意力层捕获多尺度特征的能力，从而导致处理具有不同尺度的多个对象的图像的性能下降。为了解决这个问题，我们提出了一种新颖的通用策略，称为shunted self-attention分流自我注意 (SSA)，它允许 ViT 在每个注意层的混合尺度上对注意进行建模。 SSA的关键思想是将异构感受野大小注入tokens：在计算self-attention矩阵之前，它选择性地合并tokens以表示更大的对象特征，同时保留某些tokens以保留细粒度的特征。这种新颖的合并方案使self-attention能够学习不同大小的对象之间的关系，同时减少token数量和计算成本。

Introduction

与专注于局部建模的卷积神经网络不同，ViTs 将输入图像划分为一系列patches（tokens），并通过全局自注意力逐步更新token特征。自注意力可以有效地建模token的长期依赖关系和通过聚合来自其他token的信息，逐渐扩大其感受野的大小，这在很大程度上是 ViT 成功的原因。

然而，自注意力机制也带来了昂贵的内存消耗成本，该成本与输入令牌的数量成二次方。 因此，最先进的 Transformer 模型采用了各种下采样策略来减少特征大小和内存消耗。 例如，[7]的方法在第一层进行16×16的下采样投影，并在得到的粗粒度和单尺度特征图上计算self-attention； 因此，所产生的特征信息丢失不可避免地会降低模型性能。

其他方法努力计算高分辨率特征的自注意力，并通过将token与token的空间缩减合并来降低成本。 然而，这些方法往往会在一个自注意力层中合并太多的token，从而导致来自小对象和背景噪声的token混合在一起。反过来，这种行为会使模型在捕获小对象时效率降低。

此外，先前的 Transformer 模型在很大程度上忽略了注意力层内场景对象的多尺度性质，使得它们在涉及不同大小对象的野外场景中变得脆弱。 从技术上讲，这种无能归因于它们潜在的注意力机制：现有方法仅依赖于token的静态感受野和一个注意力层内的统一信息粒度，因此无法同时捕获不同尺度的特征。

为了解决这个限制，我们引入了一种新颖的通用自注意力方案，称为分流自注意力（SSA），它明确允许同一层内的自注意力头分别解释粗粒度和细粒度特征。 与合并太多token或无法捕获小对象的先前方法不同，SSA 有效地在同一层内的不同注意力头上同时对各种尺度的对象进行建模，使其具有良好的计算效率以及保留细粒度细节的能力。

 *图2：比较了具有相同特征图大小的Vision Transformer(VIT)、金字塔视觉Transformer(PVT)和SSA中的不同注意机制。圆圈数代表了自我注意计算所涉及的token数，并反映了计算成本。圆圈的大小表示相应token的接受域大小。与VIT和PVT方法不同，我们的方法自适应地合并大对象上的圆，以提高计算效率，并同时考虑不同尺度的对象

我们在图 2 中展示了 vanilla self-attention（来自 ViT）、下采样辅助注意力（来自 PVT）和提议的 SSA 之间的定性比较。 当对相同大小的特征图应用不同的注意力时，ViT 捕获细粒度的小对象，但计算成本非常高（图 2（a））； PVT 降低了计算成本，但它的注意力仅限于粗粒度的较大对象（图 2（b））。相比之下，所提出的 SSA 保持了轻计算负载，但同时考虑了混合尺度的注意力（图 2（c）） . 实际上，SSA 不仅可以精确地处理粗粒度的大物体（例如沙发），还可以精确地处理细粒度的小物体（例如灯和风扇），甚至一些位于角落的物体，不幸的是 PVT 错过了这些物体。 我们还在图 3 中展示了注意力图的视觉比较，以突出 SSA 的学习尺度自适应注意力。

 *图3：PVT和SSA的attention map，PVT-倾向于只捕捉像沙发和床这样的大物体，而我们的模型则相反，准确地捕捉到像灯这样的小物体。

SSA的多尺度注意力机制是通过将多个注意力头分成几个组来实现的。每个组占一个专门的注意力粒度。对于细粒度的组，SSA学习聚合很少的token并保留更多的局部细节。 对于剩余的粗粒度头部组，SSA 学习聚合大量的token，从而降低计算成本，同时保留捕获大对象的能力。 多粒度组共同学习多粒度信息，使模型能够有效地建模多尺度对象。

 *图1：最新的Sota CNN和transformer模型在ImageNet上的精度排名第一。我们建议的Shunted Transformer超越了所有的baseline，包括最近的SOTA Focal Transformer(基本尺寸)。值得注意的是，它在模型尺寸小2倍的情况下达到了与DeiT-S相当的精度

如图 1 所示，我们展示了通过堆叠多个基于 SSA 的块获得的Shunted Transformer 模型的性能。 在 ImageNet 上，我们的Shunted Transformer 优于最先进的Focal Transformer  [29]，同时将模型大小减半。 当缩小到很小的尺寸时，Shunted Transformer 实现了与 DeiT-Small [20] 相似的性能，但只有 50% 的参数。 在对象检测、实例分割和语义分割方面，Shunted Transformer 在 COCO 和 ADE20K 上以相似的模型大小始终优于 Focal Transformer。

Contributions

• 我们提出了Shunted Self-Attention(SSA)，通过多尺度token聚合，将多尺度特征提取统一在一个self-attention layer内。SSA自适应地合并大对象上的标记以提高计算效率，并保留小对象上的标记。
• 基于SSA，我们构建了Shunted Transformer，它能够有效地捕获多尺度物体，尤其是小型和远程孤立物体。
• 我们在包括分类、对象检测和分割在内的各种研究中评估我们提出的Shunted Transformer。 实验结果表明，我们的Shunted Transformer在相似的模型尺寸下始终优于以前的vision transformer。

Related Work

Self-Attention in CNNs

卷积层的感受场通常较小且固定。虽然膨胀卷积[30]扩大了接收范围，可变形卷积允许核的某些偏移量[5]，但它们很难自适应和灵活地扩展到整个特征映射。受NLP领域首创的自我注意层[21]transformer的启发，一些工作引入了self-attention或non-local block来增强计算机视野中的卷积神经网络。这种注意通常适用于特征图尺寸较小且由多个卷积层进行预处理的深层。因此，它们不会招致太多的额外计算代价，但带来的性能改进有限。

Vision Transformer

自开创性的 ViT 模型以来，开发了许多后续工作，以通过更复杂的数据增强或知识蒸馏来提高模型的分类性能 [17、20]。 由于 self-attention 的计算复杂度与 token 的数量成二次方，因此很难直接应用于大量的 token。 因此，这些 ViT 模型通常将图像划分为不重叠的大尺寸patch（token）。 但是这样的划分过于粗糙，丢失了很多细粒度的信息。 为了保留细粒度的特征，这些模型通常对特征图进行下采样并对低分辨率特征进行操作。 然而，这种妥协阻碍了它们在分割和检测等密集预测任务中的部署。

Efficient ViT Variants

为了使自注意力注意适用于大尺寸特征图，最近的工作开发了两种解决方案策略 [2, 14, 25, 26, 29] 以降低计算成本：

（1）将特征图分割成区域并执行局部区域内的自注意力

（2）合并token以减少token的数量

局部注意力的代表工作是 Swin Transformer [14]，它将特征图分割成不重叠的平方区域并在局部进行自注意力。然而，为了通过自注意力对全局依赖进行建模，这些局部注意力需要在图像上移动窗口或堆叠很多层以获得全局感受野。

关于token合并的策略，PVT（Pyramid Vision Transformer）[26]设计了一种空间缩减注意力来合并key和query的token。然而，PVT 和类似的模型倾向于在这种空间缩减中合并太多的token。这使得小物体的细粒度信息与背景混合在一起，损害了模型的性能。

因此，我们提出了Shunted Self-attention，它可以同时保留粗粒度和细粒度的细节，同时保持对图像标记的全局依赖建模。

Method

 我们提出的Shunted Transformer的整体架构如图 4 所示。它建立在新颖的Shunted Self-Attention(SSA) block之上。 我们的 SSA 块与 ViT 中的传统自注意力块有两个主要区别：

• SSA 为每个自注意力层引入了Shunted Self-Attention机制，以捕获多粒度信息和更好的不同大小的模型对象，尤其是小物体；
• 它通过增强跨token交互来增强在逐点前馈层中提取局部信息的能力。

此外，我们的Shunted Transformer部署了一种新的patch嵌入方法，以便为第一个注意块获得更好的输入特征图。 下面，我们将一一详述这些创新点。

3.1. Shunted Transformer Block

在所提出的Shunted Transformer的第i级中，有Li个transformer block。 每个transformer块都包含一个self-attention层和一个前馈层。为了降低处理高分辨率特征图时的计算成本，PVT [26]引入了空间缩减注意力（SRA）来代替原来的多头自注意力（ MSA）。 然而，SRA 倾向于在一个自我注意层中合并太多的token，并且只提供单一尺度的token特征。 这些限制阻碍了模型捕捉多尺度物体的能力，尤其是小尺寸物体。 因此，我们引入了我们的Shunted Self-attention，并在一个自注意力层中并行学习多粒度。

3.1.1 Shunted Self-Attention

首先将输入序列投影到Q、K、V张量中，然后MHSA采用H个独立的注意力头来并行计算自注意力。为了降低计算成本，我们遵循PVT并减少K和V的长度，而不是像Swin Transformer中那样将{Q,K,V}拆分为区域。

*图5：将我们的Shunted Self-Attention与ViT、Swin、PVT 中的Self-attention进行比较。 ViT在小尺寸特征图上全局应用SA。 Swin Transformer 在小区域内的大尺寸特征图上应用局部SA。 PVT 通过large stride了key和value。 不同的是，我们的shunted self-attention进行多尺度token聚合以获得各种大小的k和v。

如图 5 所示，我们的 SSA 与 PVT 的 SRA 不同之处在于，K、V 的长度在同一自注意力层的注意力头上并不相同。 取而代之的是，不同head的长度不同，用于捕获不同粒度的信息。 这给出了多尺度token聚合（MTA）。 具体来说，对于由i索引的不同head，Key和Value被下采样到不同的大小：

 这里的是第i个head中的多尺度token聚合层，具有下采样率。在实际应用中，我们采用一个kernel和stride大小为的卷积层来实现下采样。, ,是第i个head的线性投影层的参数。在注意力头的一层中存在变量。 因此，key和value可以在self-attention中捕获不同的尺度。 LE(·) 是MTA通过深度卷积对值V进行的局部增强分量。 与PVT中的空间缩减相比，保留了更多细粒度和低级细节。

然后shunted self-attention中的计算可以表示为：

  是维度。由于多尺度key和value，我们的shunted self-attention在捕获多尺度对象方面更加强大。 计算成本的降低可能取决于r的值，因此，我们可以很好地定义模型和r来权衡计算成本和模型性能。 当r变大时，K,V 中更多的token被合并，K,V 的长度更短，因此计算成本低，但仍然保留了捕获大对象的能力。 相反，当 r 变小时，保留更多细节但带来更多计算成本。 将各种r集成到一个自注意力层中，使其能够捕获多粒度特征。

3.1.2 Detail-specific Feedforward Layers

 *图6：比较ViT、PVT和我们的Shunted Transformer中的带有细节特定的前馈层，我们在前馈层补充了细粒度的交叉token细节。

在传统的前馈层中，全连接层是逐点的，无法学习到交叉token信息。 在这里，我们旨在通过在前馈层中指定细节来补充本地信息。 如图6所示，我们通过在前馈层的两个全连接层之间添加我们的数据特定层data specific layer来补充前馈层中的局部细节：

 是带有参数的细节特定层，实际上是通过深度卷积实现的。

3.2. Patch Embedding

Transformer 最初是为处理顺序数据而设计的。 如何将图像映射到序列对于模型的性能很重要。 ViT 直接将输入图像分割成 16 × 16 个不重叠的补丁。 最近的一项研究 [23] 发现，在patch embedding中使用卷积可提供更高质量的token序列，并帮助transformer比传统的大步幅非重叠patch embedding“看得更清楚”。 因此，一些作品 (Swin，PVT) 像使用 7×7 卷积一样进行重叠块嵌入。

在我们的模型中，我们根据模型的大小采取不同的卷积层进行重叠。该算法以一个步长为2、零填充的7×7卷积层作为块嵌入的第一层，根据模型大小增加额外的步长为1的3×3卷积层，最后由步长为2的非重叠投影层生成大小为的输入序列。

3.3. Architecture Details and Variants

 给定一幅大小为H×W×3的输入图像，我们采用上述patch embedding方案来获得更多信息的token序列，token长度为，token维度为C。按照前面的设计[Twins,Swin,PVT,Focal]，我们的模型分为四个阶段，每个阶段包含若干个Shunted Transformer块。每一阶段，每个块都输出相同大小的特征图。我们用一个跨度为2(线性嵌入)的卷积层来连接不同的阶段，在进入下一阶段之前，特征图的尺寸将减半，但维度将增加一倍。因此，我们在每个阶段的输出有四个特征图F1，F2，F3，F4，Fi的大小是。

我们提出了三种不同的模型结构，以便在相似的参数和计算代价下进行公平比较。如表1所示，Head和Ni表示一个块中的head的数量和一个阶段中的块的数量。变体只来自不同阶段的层数。具体地说，每个block的head数设置为2、4、8、16。Patch embedding的卷积范围为1到3。

Experiments

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

Image Classification on ImageNet-1K

我们首先评估了我们的模型，并将其与最近在ImageNet-1K上基于SOTA、CNN和Transform的模型进行了比较。为了公平比较，我们遵循了Deit[20]和PVT[26]相同的训练策略。具体而言，我们采用AdamW作为权值衰减率为0.05的优化器，整个训练周期为300次，输入规模为224×224，批规模为1024。在余弦学习速率衰减之后，学习速率被设置为1×10−3。数据增强和正则化方法遵循Deit[20]，包括随机裁剪、随机翻转、标签平滑[18]、混合[33]、CutMix[32]和随机擦除[35]。

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

然后，我们将我们的模型与transformer backbone进行比较。我们的微型模型获得了与transformer baseline(DEIT-S)相似的性能，但只需要一半的参数(22M→11M)和计算成本(4.6G→2.1GFLOPS)。当我们的模型尺寸增长到与Deit-S相似时，它的表现比Deit-S高3%。与Swin和Twin等最新的SOTA模型相比，我们的模型始终表现优于它们。具体地说，我们的小尺寸模型比现有最先进的Focal Transformer Tiny高0.7%，同时将模型尺寸缩小了20%。当模型尺寸变大时，我们的基本模型获得了最先进的性能，与Focal Transformer相比，只需一半的参数和计算成本。

Object Detection and Instance Segmentation

我们在COCO 2017上对目标检测和实例分割模型进行了评估[13]。我们以我们提出的Shunted Transformer为主干，将其接入MASK R-CNN。我们将其与其他SOTA骨干网进行了比较，包括ResNet、Swin Transformer、PVT、Twin和Focal Transformer。我们遵循Swin的相同设置：对ImageNet-1K进行预培训，并对COCO进行微调。

在表3中，我们使用MASK-RCNN进行对象检测，并报告了不同CNN和Transformer骨干的BBox MAP(APB)。在可比较的参数条件下，我们的模型比以前的SOTA模型有很大的差距；在目标检测方面，在1倍调度的情况下，我们的微小模型比ResNet-50提高了9.1个点，在只有85%的模型规模的情况下，我们的微小模型比Focal Transformer提高了2.3个点；此外，在3倍的调度和多尺度训练下，我们的主干仍然简明地超过了CNN主干7.7个点和Transformer主干1.6个点。我们在实例分割中得到了类似的结果。我们报告了表3中的掩码映射(APM)。在1倍调度和3倍调度中，我们的模型比ResNet-50高8.1个点，比Focal Transformer高1.5个点。我们的模型在较小的模型规模下实现了这些优越的性能，清楚地展示了其分流数据在学习多粒度标记方面的好处以及在处理多尺度视觉对象存在时的有效性。

我们还在表 4 中报告了 RetinaNet 的结果。在参数最少的情况下，我们的模型在 1× 和 3× 调度中都优于之前的所有模型。 与 PVT 相比，我们的模型对所有小、中、大尺寸物体都进行了改进，这显示了在我们的shunted self-attention 中捕获多尺度物体的强大能力。 

Semantic Segmentation on ADE20K

Ablation Studies

Patch Embedding

最近的许多工作[9，10，24]研究了图像到token的映射功能，即patch embedding head。他们发现，设计良好的head为transformer模型提供了更好的输入顺序。我们使用 ViT 中的非重叠头、Swin 和 PVT 中的重叠头来评估我们的patch embedding的影响。 结果如表 7 所示。使用更复杂的 head（如重叠）或我们的patch embedding，计算成本和模型大小仅略有增加，但性能提升相对显着。 具体来说，在有限的附加参数下，从传统的非重叠头部或重叠头到patch embedding，该模型分别实现了 1.4% 和 0.3% 的性能增益。

Token Aggregation Function

我们提出了一种新的token聚合函数来合并多尺度对象的token并同时保持全局和局部信息。 从表 8 中可以看出，我们新颖的token聚合函数具有与卷积空间缩减相似的计算，但获得了更多改进。

Detail-specific Feed-Forward

在前馈层[20]中，所有操作都是逐点的，不存在cross token操作，因此，补充cross token和局部信息将显着提高前馈层的学习能力。 在表 9 中，我们比较了 ViT 和我们的模型中新的细节特定前馈层、传统前馈层 [20] 和卷积前馈层 [25]。 与传统的前馈层相比，特定于细节的前馈始终如一地带来性能增益，这表明在前馈层中补充局部细节的实用性。

Conclusion

在本文中，我们提出了一种新的shunted self-attention(SSA)方案，以显式地解释多尺度特征。与以往只关注一个关注层中的静态特征图不同，我们维护了一个关注层中多尺度目标的不变尺度特征图。广泛的实验表明，我们的模型作为各种下游任务的骨干是有效的。具体地说，该模型的性能优于以往的Transformers，并在分类、检测和分割任务方面取得了最先进的结果。