大道至简 | 设计 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 整体架构

图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

在图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

结果如表 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

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

图 2

此外，在图 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至关重要。

图 3

为了更好地理解这一现象，在不移除任何 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

长按扫描下方二维码添加小助手并加入交流群， 群里博士大佬云集， 每日讨论话题有目标检测、语义分 割、 超分辨率、模型部署、数学基础知识、算法面试题分享的等等内容，当然也少不了搬砖人的扯犊子

长按扫描下方二维码添加小助手。

可以一起讨论遇到的问题

声明：转载请说明出处

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