大道至简 | 设计 ViT 到底怎么配置Self-Attention才是最合理的？

Transformer已成为深度学习中的主要架构之一，尤其是作为计算机视觉中卷积神经网络 (CNN) 的强大替代品。然而，由于Self-Attention在长序列表示上的二次复杂性，特别是对于高分辨率密集预测任务，先前工作中的Transformer训练和推理可能非常昂贵。为此，我们提出了一种新颖的少注意力视觉Transformer（LIT），基于Transformers中早期的Self-Attention仍然专注于局部模式，并在最近的分层视觉Transformers中带来较小的好处。

具体来说，提出了一个分层 Transformer，使用纯多层感知器（MLP）在早期阶段对丰富的局部模式进行编码，同时应用Self-Attention模块在更深层捕获更长的依赖关系。此外，进一步提出了一种 Learned Deformable Token Merging Module，以非均匀方式自适应地融合信息Patch。

所提出的 LIT 在图像识别任务（包括图像分类、对象检测和实例分割）上取得了可观的性能，可作为许多视觉任务的强大支柱。

开源地址：https://github.com/zhuang-group/LIT

1简介

Transformers在自然语言处理(NLP)和最近在计算机视觉(CV)领域取得了长足的进步。受CNN中金字塔设计的启发，最近的分层视觉转换器(HVT)将Transformer Block划分为多个阶段，并随着网络的深入逐渐缩小特征图。然而，早期的高分辨率特征图导致了很长的token序列，由于Self-Attention的二次复杂度，带来了巨大的计算成本和内存消耗。例如，大小为56×56×96 的特征图在一个Multi-Head Self-Attention（MSA）中需要2.0G FLOPs，而ResNet-18的整个模型只需要1.8G FLOPs。如此巨大的计算成本使得将Transformer应用到广泛的计算机视觉任务中变得非常困难。

在HVT的早期阶段，已经做出了一些努力来降低计算成本。例如，一些工作减少了MSA层中Self-Attention Head的数量或进一步减少了Transformer Block的数量。另一行工作建议通过启发式近似来权衡MSA的准确性和效率，例如空间缩减注意力(SRA)和基于Shift Window Multi-Head Self-Attention(SW-MSA)。也有研究当特征图的分辨率相当大时使用卷积层。然而，早期采用Self-Attention层对最终性能有多大贡献仍不清楚。

在本文中提出了一种Less attention Vision Transformer（LIT）来解决上述HVT问题。具体来说，建议在早期Stage使用MLP层来捕获局部信息，同时引入具有足够数量的Head来处理的MSA层后期的长期依赖关系。

作者的动机来自2个方面。首先，先前的CNN和Transformer相关研究表明，模型浅层会专注于局部信息，而深层倾向于捕获高级语义或全局关系，这就产生了在早期Stage是否有必要使用Self-Attention的问题。其次，从理论的角度来看，一个具有足够Head的Self-Attention应用于图像可以表达任何卷积层。

然而，MSA层中较少的Head理论上阻碍了逼近具有大kernel-size的卷积层的能力，其中极端情况与1×1卷积一样具有表现力，可以被视为独立应用于每个像素的标准FC层。虽然最近的HVT在早期阶段采用很少的Head来提供金字塔表示，但作者认为这不是最优的，因为这样的配置会引入高计算和内存成本，但带来的收益却很小。

需要强调的是，通过在前期Stage利用MLP Block可以避免Self-Attention在高分辨率特征图上产生的巨大计算成本和内存占用。此外，在后期Stage应用Self-Attention来捕获远程依赖关系。综合实验结果表明，经过如此简单的架构设计后可以带来模型性能和效率之间的最佳平衡点。

此外，最近的HVT要么采用标准卷积层，要么采用线性投影层来合并邻近的Token，旨在控制特征图的规模。然而，考虑到并非每个像素对输出单元的贡献都相同，这种方法阻碍了Vision Transformer对几何变换建模的表示能力。

为此，作者提出了一个受Deformable Convolutions启发的Deformable Token Merging(DTM) 模块，在该模块中，学习了一个偏移网格以自适应地增加空间采样位置，进而合并来自特征图中子窗口的相邻Patch。通过这种方式可以获得更多信息的下采样Token用于后续处理。

主要贡献：

• 首先，确定了早期MSA层在最近的HVT中的微小贡献，并在早期Stage提出了一个简单的HVT结构，其中包含纯MLP Block;
• 其次，提出了一个Deformable Token Merging模块，以自适应地合并更多信息Patch以提供分层表示，并具有增强建模能力;
• 最后，进行了广泛的实验，以表明所提出的LIT在计算复杂度和内存消耗方面的有效性。

2本文方法

2.1 整体架构

LIT的整体架构如图1所示。设为输入的RGB图像，其中H和W分别表示高度和宽度。首先将I拆分为不重叠的Patch，Patch大小为4×4，因此每个Patch的初始特征维度为4×4×3=48。

接下来，利用线性嵌入层将每个Patch投影到维度，用作以下过程的初始输入。整个模型分为4个Stage。让s∈[1,2,3,4] 作为Stage的索引，在每个Stage使用个Block，其中前2个Stage仅使用MLP Block来编码局部信息，后两个Stage使用标准Transformer Block来编码处理更长的依赖关系。

在每个Stage，将输入特征图缩放为，其中和分别表示第s个Stage的块大小和隐藏维度。对于最后2个Stage，在每个Transformer Block中设置个Self-Attention heads。

2.2 LIT中的Block设计

如图1所示，LIT采用2种类型的模块:MLP模块和Transformer模块。在前期Stage应用MLP Block。具体来说，MLP Block建立在MLP之上，该MLP由2个FC层组成，中间有GELU。对于第s阶段的每个MLP，使用的扩展比。

具体来说，第1个FC层将token的维度从扩展到，另一个FC层将维度缩减回。形式上，令为第s阶段的输入，l为块的索引，MLP Block可以表示为：

其中LN表示层归一化。在最后阶段，ViT中描述的Transformer Block包含一个MSA层和一个MLP，可以表示为：

使用这种架构有2个主要优点：

• 首先，避免了在早期Stage由长序列引入的巨大计算成本和内存占用；

• 其次，与最近使用子窗口缩小注意力图或减少注意力图的工作不同在key和value矩阵的空间维度上，在最后2个Stage保留标准MSA层，以保持LIT处理远程依赖关系的能力。

备注

在这里通过考虑卷积层、FC层和MSA层之间的关系来证明在前期Stage应用纯MLP Block的合理性。

这里建议参考一下字节跳动近期新出的TRT-ViT，有对应的结论给出，同时给出了非常详细的模型设计准则，对应推文的链接如下：

建议背诵 | 字节用4大准则教你设计一个拥有CNN的速度，Transformer精度的模型！

首先，从回顾标准卷积层开始。令为输入特征图，令为卷积权重，其中K为Kernel-size，和分别为输入和输出通道维度。为简单起见，省略了偏置项，并用表示，其中(i,j)表示像素索引，。给定K×K个采样位置的卷积核，像素p的输出可以表示为：

其中是采样索引到预先指定的偏移量ΔK的双射映射。

例如，令ΔK={(−1,−1),(−1,0),...,(0,1),(1,1)}为3×3的kernel，膨胀率为1，则g(0)=(-1,-1)表示第1个采样偏移量。当K=1时，权重张量W等价于一个矩阵，使得$。在这种情况下，Eq(4)可以表示FC层，像素p的输出定义为：

最后，让是MSA层中的Head数，$是第h个Head的可学习参数。在特定的相对位置编码方案下，Cordonnier等人证明了像素p处的MSA层的输出可以表示为：

其中是Head到像素位移的双射映射。在这种情况下，Eq(6)可以看作是Kernel-size为 的卷积层近似。

从Eqs(4)-(6)观察到，虽然具有足够数量的Head的MSA层能够逼近任何卷积层，但理论上更少的Head限制了这种逼近的能力。作为极端情况，具有一个Head的MSA层是只能逼近FC层。

请注意，MSA层在实践中肯定不等同于卷积层。然而，d'Ascoli 等人观察到，早期的MSA层可以在训练时学会卷积的类似表征。考虑最近的HVT在早期Stage采用很少的Head，这种卷积行为可能会限制在小的感受野内。

在图3中，在可视化中显示PVT-S中的早期MSA层确实只关注query像素周围的一小块区域，而删除它们会导致性能轻微下降，但会显着降低模型复杂性。这证明了在前2个stage应用纯MLP Block是合理的。

2.3 Deformable Token Merging

以前关于 HVT 的工作依赖于Patch Merge来实现金字塔特征表示。然而，从规则网格合并Patch，并忽略了并非每个Patch对输出单元的贡献相同的事实。受可变形卷积的启发，提出了一个Deformable Token Merging模块来学习偏移网格，以自适应地采样更多信息块。形式上，可变形卷积被表述为：

与Eq(4)中的标准卷积运算相比，DC为每个预先指定的偏移量g(k)学习一个偏移量Δg(k)。学习Δg(k) 需要一个单独的卷积层，它也应用于输入特征图X。为了以自适应方式合并Patch，在DTM模块中采用一个DC层，可以表示为：

其中 BN 表示批量归一化，并使用了 GELU。

在消融研究中，当物体的尺度和形状发生变化时，DTM 中的采样位置会进行自适应调整，从而受益于学习到的偏移量。另请注意，与Baseline中的常规网格采样相比，轻量级DTM引入了可忽略的FLOP和参数，因此使其成为最近HVT的即插即用模块。

3实验

3.1 消融实验

1、Effect of the architecture design

结果如表 2 所示。总的来说，LIT在使用更少的 FLOP（3.6G 与 3.8G）的同时，将PVT-S的Top-1精度提高了0.6%。对于Swin-Ti将 FLOPs 减少了0.4G，同时实现了同等性能。

还值得注意的是，PVT-S 和 Swin-Ti 的参数总数都减少了。整体性能证明了所提出架构的有效性，这也强调了早期MSA在PVT和Swin中的微小优势。

2、Effect of deformable token merging

结果如表 3 所示。对于这2个模型，DTM 引入了可忽略的 FLOP 和参数，同时在 Top-1 精度方面分别将 PVT-S 和 Swin-Ti 提高了 0.7% 和 0.3%。

此外，在图 2 中可视化了学习到的偏移量。如图所示，与之前的统一Patch Merge策略不同，之前的统一Patch Merge策略将采样位置限制在绿色矩形内，DTM 根据对象的比例和形状（例如，考拉腿、猫尾巴）。这再次强调了LIT适应各种几何变换的能力。

3、Effect of MSA in each stage

结果如表4所示。首先，在用标准MSA层替换PVT-S中的SRA层后，观察到 Top-1 准确度提高了1.1%，而FLOPs几乎翻了一番。这表明PVT在性能和效率之间进行了权衡。

接下来，通过在前2个阶段逐步去除MSA层，Top-1准确率仅分别下降了0.1%、0.5%。这意味着 PVT 早期阶段的self-attention层对最终性能的贡献低于预期，并且它们的性能并不比纯MLP层好多少。这可以归因于浅层更多地关注编码局部信息的事实。

然而，可以观察到在最后2个阶段移除self-attention时性能大幅下降。结果表明，self-attention层在后期发挥着重要作用，捕获长距离依赖对于表现良好的分层视觉Transformer至关重要。

为了更好地理解这一现象，在不移除任何MSA层的情况下可视化PVT-S的注意力概率，如图3所示。首先，第1阶段的注意力图显示query像素几乎不注意其他位置。

在第2阶段，query像素的感受野略有扩大，但与第1阶段相似。考虑到PVT-S在第1阶段只有1个head，在第2阶段只有2个head，这有力地支持了作者的假设：MSA层中过少的head会导致较小的感受野，此时self-attention几乎等同于FC层。

此外，从最后2个阶段的注意力图中观察到相对较大的感受野。由于大的感受野通常有助于对更长的依赖关系进行建模，这解释了在最后2个阶段移除MSA层后表4中的巨大性能下降。

3.2 分类

3.3 目标检测

3.4 实例分割

3.5 语义分割

4参考

[1].Less is More: Pay Less Attention in Vision Transformers

5推荐阅读

分割冠军 | 超越Swin v2、PvT v2等模型，ViT-Adaptiver实现ADE20K冠军60.5mIoU

DAFormer | 使用Transformer进行语义分割无监督域自适应的开篇之作

Sparse R-CNN升级版 | Dynamic Sparse R-CNN使用ResNet50也能达到47.2AP

扫描下方二维码关注【集智书童】公众号，获取更多实践项目源码和论文解读，非常期待你我的相遇，让我们以梦为马，砥砺前行！