关于人工智能:CVPR22-Oral｜通过多尺度token聚合分流自注意力代码已开源

【写在后面】

最近的视觉Transformer（ViT）模型在各种计算机视觉工作中获得了令人鼓舞的后果，这得益于其通过自注意力建模图像块或token的长期依赖性的能力。然而，这些模型通常指定每个层内每个token特色的相似感触野。这种束缚不可避免地限度了每个自注意力层捕获多尺度特色的能力，从而导致解决具备不同尺度的多个对象的图像的性能降落。为了解决这个问题，作者提出了一种新的通用策略，称为分流自注意力（SSA），该策略容许VIT在每个注意力层的混合尺度上对注意力进行建模。SSA的要害思维是将异质感触野大小注入token：在计算自注意力矩阵之前，它选择性地合并token以示意较大的对象特色，同时保留某些token以保留细粒度特色。这种新的合并计划使自注意力可能学习不同大小对象之间的关系，同时缩小了token数和计算成本。跨各种工作的大量试验证实了SSA的优越性。具体来说，基于SSA的Transformer达到了84.0%的Top-1精度，优于ImageNet上最先进的Focal Transformer，模型尺寸和计算成本仅为其一半，在类似的参数和计算成本下，在COCO上超过了Focal Transformer 1.3 mAP，在ADE20K上超过了2.9 mIOU。

1. 论文和代码地址

Shunted Self-Attention via Multi-Scale Token Aggregation

论文地址：https://arxiv.org/abs/2111.15193

代码地址：https://github.com/oliverrensu/shunted-transformer

2. Motivation

最近的视觉Transformer（ViT）模型在各种计算机视觉工作中体现出了卓越的性能。与专一于部分建模的卷积神经网络不同，ViTs将输出图像划分为一系列patch，并通过全局自注意力逐渐更新token特色。自注意力能够无效地模仿token的长期依赖性，并通过聚合来自其余token的信息来逐渐扩充其感触野的大小，这在很大水平上解释了VIT的胜利。

然而，自注意力机制也带来了低廉的内存耗费老本，即输出token数量的平方比。因而，最先进的Transformer模型采纳了各种降采样策略来缩小特色大小和内存耗费。一些办法致力计算高分辨率特色的自注意力，并通过将token与token的空间缩减合并来降低成本。然而，这些办法偏向于在一个自注意力层中合并过多的token，从而导致来自小对象和背景噪声的token的混合。这种行为反过来会升高模型捕捉小对象的效率。

此外，以前的Transformer模型在很大水平上疏忽了注意力层内场景对象的多尺度性质，使得它们在波及不同大小对象的家养场景中变得软弱。从技术上讲，这种能干归因于其潜在的留神机制：现有办法仅依赖token的动态感触野和一个留神层内的对立信息粒度，因而无奈同时捕捉不同尺度的特色。

为了解决这一局限性，作者引入了一种新的通用自注意力计划，称为分流自注意力（SSA），该计划明确容许同一层内的自注意力头别离思考粗粒度和细粒度特色。与以前合并过多token或捕获小对象失败的办法不同，SSA无效地在同一层的不同留神头上同时对不同规模的对象建模，使其具备良好的计算效率和保留细粒度细节的能力。

作者在上图中展现了自注意力（来自ViT）、下采样辅助注意力（来自PVT）和SSA之间的定性比拟。当对雷同大小的特色映射利用不同的留神时，ViT捕获细粒度的小对象，但具备计算成本极高（上图（a））；PVT升高了计算成本，但其注意力仅限于粗粒度较大的对象（上图（b））。相比之下，提出的SSA放弃了较轻的计算负载，但同时思考了混合尺度留神（上图（c））。无效地，SSA不仅准确地关注粗粒度的大对象（如沙发），而且还关注细粒度的小对象（如灯光和风扇），甚至一些位于角落的对象，这些对象可怜被PVT疏忽。作者还在下图中展现了注意力图的视觉比拟，以突出SSA的学习尺度自适应注意力。

SSA的多尺度注意力机制是通过将多个注意力头分成几个组来实现的。每个组都有一个专用的注意力粒度。对于细粒度组，SSA学习聚合大量token并保留更多部分细节。对于残余的粗粒度头组，SSA学习聚合大量token，从而升高计算成本，同时放弃捕捉大型对象的能力。多粒度组独特学习多粒度信息，使模型可能无效地建模多尺度对象。

如上图所示，作者展现了从重叠多个基于SSA的块取得的分流Transformer模型的性能。在ImageNet上，本文的分流Transformer优于最先进的聚焦Transformer，同时将模型尺寸减半。当放大到渺小尺寸时，分流Transformer实现了与DeiT Small类似的性能，但只有50%的参数。对于对象检测、实例宰割和语义宰割，在模型大小类似的COCO和ADE20K上，分流Transformer始终优于聚焦Transformer。

本文的奉献如下：

作者提出了分流自注意力（SSA），它通过多尺度token聚合在一个自注意力层内对立多尺度特征提取。本文的SSA自适应地合并大对象上的token以进步计算效率，并保留小对象的token。

基于SSA，作者构建了分流Transformer，可能无效捕捉多尺度对象，尤其是小型和近程孤立对象。

作者评估了提出的分流Transformer的各种钻研，包含分类，指标检测和宰割。试验结果表明，在类似的模型尺寸下，本文的分流Transformer始终优于以前的视觉Transformer。

3. 办法

本文提出的分流Transformer的整体架构如上图所示。它建设在新型分流自注意力（SSA）块的根底上。本文的SSA块与ViT中的传统自注意力块有两个次要区别：1）SSA为每个自注意力层引入了分流注意力机制，以捕捉多粒度信息和更好地建模不同大小的对象，尤其是小对象；2） 它通过加强跨token交互，加强了在逐点前馈层提取部分信息的能力。此外，本文的分流Transformer部署了一种新的patch嵌入办法，用于为第一个注意力块取得更好的输出特色图。在下文中，作者将逐个论述这些新鲜之处。

3.1. Shunted Transformer Block

在提出的分流Transformer的第i个阶段，有$L_i$个Transformer块。每个transformer块蕴含一个自注意力层和一个前馈层。为了缩小解决高分辨率特色图时的计算成本，PVT引入了空间归约注意力（spatial-reduction attention，SRA）来取代原来的多头自注意力（multi-head self attention，MSA）。然而，SRA偏向于在一个自注意力层中聚合太多token，并且仅在单个尺度上提供token个性。这些局限性妨碍了模型捕获多尺度对象尤其是小尺寸对象的能力。因而，作者在一个自注意力层内并行引入了具备学习多粒度的分流自注意力。

3.1.1 Shunted Self-Attention

输出序列$F \in \mathbb{R}^{h \times w \times c}$首先投影到查问（Q）、键（K）和值（V）张量中。而后，多头自注意力采纳H个独立留神头并行计算自注意力。为了升高计算成本，作者遵循PVT并缩小K和V的长度，而不是像在Swin Transformer中那样将{Q，K，V}宰割为多个区域。

如上图所示，本文的SSA不同于PVT的SRA，因为K，V的长度在同一个自注意力层的留神头之间不雷同。相同，长度在不同的头部不同，用于捕捉不同粒度的信息。这提供了多尺度token聚合（MTA）。具体而言，对于由i索引的不同头部，将键K和值V下采样到不同大小：

$$\begin{aligned} Q_{i} &=X W_{i}^{Q} \\ K_{i}, V_{i} &=M T A\left(X, r_{i}\right) W_{i}^{K}, M T A\left(X, r_{i}\right) W_{i}^{V} \\ V_{i} &=V_{i}+L E\left(V_{i}\right) \end{aligned}$$

这里，$\operatorname{MAT}\left(; ; r_{i}\right)$是第i个头中的多尺度token聚合层，下采样率为$r_i$。在实践中，作者应用卷积核和步长为$r_i$的卷积核进行卷积层。$W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}$是第i个头中线性投影的参数。在注意力头的一层中有变体$r_{i}$。因而，键和值能够在自注意力中捕获不同的尺度。$L E(\cdot)$是MTA的部分加强重量，用于深度卷积的V值。与空间缩减相比，保留了更多细粒度和低层次的细节。

而后通过以下公式计算分流的自注意力：

$$h_{i}=\operatorname{Softmax}\left(\frac{Q_{i} K_{i}^{\top}}{\sqrt{d_{h}}}\right) V_{i}$$

其中$d_{h}$是尺寸。多亏了多尺度键和值，本文分流的自注意力在捕获多尺度对象时更弱小。计算成本的升高可能取决于r的值，因而，能够很好地定义模型和r，以衡量计算成本和模型性能。当r变大时，K，V中合并了更多token，并且K，V的长度较短，因而，计算成本较低，但仍放弃捕捉大型对象的能力。相同，当r变小时，保留了更多细节，但带来了更多计算成本。在一个自注意力层中集成各种r使其可能捕捉多粒度特色。

3.1.2 Detail-specific Feedforward Layers

在传统的前馈层中，全连贯层是逐点的，无奈学习穿插token信息。在这里，作者的指标是通过指定前馈层中的细节来补充部分信息。如上图所示，作者通过在前馈层的两个齐全连贯的层之间增加数据特定层来补充前馈层中的部分细节：

$$\begin{aligned} x^{\prime} &=F C\left(x ; \theta_{1}\right) \\ x^{\prime \prime} &=F C\left(\sigma\left(x^{\prime}+D S\left(x^{\prime} ; \theta\right)\right) ; \theta_{2}\right) \end{aligned}$$

其中$D S(\cdot ; \theta)$是具备参数的细节特定层，在实践中通过深度卷积实现。

3.2. Patch Embedding

Transformer首先设计用于解决程序数据。如何将图像映射到序列对于模型的性能很重要。ViT间接将输出图像宰割为16×16非重叠patch。最近的一项钻研发现，在patch嵌入中应用卷积能够提供更高质量的token序列，并有助于transformer比传统的大步非重叠patch嵌入“看得更好”。因而，一些文献应用7×7卷积进行重叠patch嵌入。

在本文的模型中，作者依据模型大小采纳不同的重叠卷积层。作者将步长为2且零填充的7×7卷积层作为patch嵌入的第一层，并依据模型大小减少步长为1的额定3×3卷积层。最初，应用步长为2的非重叠投影层生成大小为$\frac{H}{4} \times \frac{W}{4}$的输出序列。

3.3. Architecture Details and Variants

给定一个大小为H×W×3的输出图像，作者采纳上述patch嵌入计划取得长度为$\frac{H}{4} \times \frac{W}{4}$、token维数为C的信息量更大的token序列。依据之前的设计，本文的模型中有四个阶段，每个阶段蕴含几个分流Transformer块。在每个阶段，每个块输入雷同大小的特色图。作者采纳带步长2（线性嵌入）的卷积层来连贯不同的阶段，在进入下一阶段之前，特色图的大小将减半，但维数将加倍。因而，有每个阶段输入的特色映射$F_{1}, F_{2}, F_{3}, F_{4}$，$F_{i}$的大小为$\frac{H}{2^{i+1}} \times \frac{W}{2^{i+1}} \times\left(C \times 2^{i-1}\right)$。

作者提出了本文模型的三种不同配置，以便在相似参数和条件下进行偏心比拟。如上表所示，head和$N_{i}$示意一个块中的头数和一个阶段中的块数。变体仅来自不同阶段的层数。具体来说，每个块中的头数设置为2,4,8,16。patch嵌入中的卷积范畴为1到3。

4.试验

上表展现了本文办法在ImageNet-1K上的试验后果，，能够看出，本文的办法在各个模型大小上都能达到SOTA后果。

上表展现了用Mask R-CNN进行指标检测和语义宰割时，各个模型的性能，能够看出，本文的办法在性能上具备显著劣势。

上表展现了用RetinaNet进行指标检测和语义宰割时，各个模型的性能。

后果如上表所示。本文的分流Transfomer在所有框架中都具备较高的性能和较少的参数，优于以前最先进的Transfomer。

作者还以SegFormer为框架，在Segformers中比拟了本文的骨干与MiT的骨干。后果见上表。

上表展现了不同patch embedding的试验后果。

作者提出了一种新的token聚合函数，用于合并多尺度对象的令牌，同时保留全局和部分信息。从上表中能够看出，本文的新token聚合函数具备与卷积空间归约类似的计算，但取得了更多改良。

上表展现了本文提出的FFN的性能改良后果。

5. 总结

在本文中，作者提出了一种新的分流自注意力（SSA）计划来明确阐明多尺度特色。与以往只关注一个注意力层中动态特色图的工作不同，作者保护了多尺度特色图，这些特色图关注一个自注意力层中的多尺度对象。大量试验表明，本文的模型作为各种上游工作的骨干是无效的。具体来说，该模型优于先前的Transformers，并在分类、检测和宰割工作上实现了最先进的后果。

已建设深度学习公众号——FightingCV，欢送大家关注！！！

ICCV、CVPR、NeurIPS、ICML论文解析汇总：https://github.com/xmu-xiaoma...

面向小白的Attention、重参数、MLP、卷积外围代码学习：https://github.com/xmu-xiaoma...

退出交换群，请增加小助手wx：FightngCV666

本文由mdnice多平台公布