MaxViT: Multi-Axis Vision Transformer

论文：https://arxiv.org/abs/2204.01697
代码地址：https://github.com/google-research/maxvit

在本文中，介绍了一种高效且可扩展的注意力模型，称之为多轴注意力，该模型由两个方面组成：分块的局部注意力和扩张的全局注意力。这些设计选择允许在仅具有线性复杂度的任意输入分辨率上进行全局-局部空间交互。本文还通过有效地将注意力模型与卷积相结合，提出了一种新的架构元素，并相应地提出了一个简单的分层视觉主干，称为MaxViT，通过简单地在多个阶段重复基本构建块。值得注意的是，MaxViT能够在整个网络中“看到”全局，甚至在早期的高分辨率阶段。

一、引言

在本文中，提出了一种新型的Transformer模块，称为多轴自注意力（Max SA），它能够作为基本的架构组件。其可以在单个块中执行局部和全局空间交互。与完全自注意力相比，Max SA具有更大的灵活性和效率，即自然适应具有线性复杂性的不同输入长度；与（平移）窗口/局部注意力相比，Max SA通过提出一个全局感受野，加强的模型能力。此外，仅具有线性复杂性，Max SA可以用作网络任何层的一般独立注意力模块，甚至在早期的高分辨率阶段。
为了证明其有效性和通用性，作者进一步设计了一种简单但有效的视觉主干，称为多轴视觉Transformer（MaxViT），通过分层堆叠由Max SA和卷积组成的重复块。虽然该模型属于混合视觉Transformer的范畴，但MaxViT与以前的方法的区别在于，本文通过设计一个统一卷积、局部和全局注意力的基本块，然后简单地重复它，力求简单。实验表明，MaxViT在各种视觉任务（包括分类、对象检测和分割、图像美学评估和图像生成）的所有数据体制下都显著提高了最先进的（SOTA）性能。本文的贡献包括：
1、一个通用的Transformer主干MaxViT，可以在网络的每个阶段捕获局部和全局空间交互。
2、一种新颖的独立多轴注意力模块，由分块的局部注意力和扩张的全局注意力组成，在线性复杂性中享受全局感知
3、通过广泛的消融研究，展示了大量的设计选择，包括层数、布局、MBConv的使用等，汇聚到最终模块化设计MaxViT块。

二、实现细节

受稀疏方法的启发，本文引入了一种新型的注意力模块，称为阻塞多轴自注意力（Max SA），通过将完全密集的注意力机制分解为两种稀疏形式——窗口注意力和网格注意力，这将普通注意力的二次复杂度降低为线性，而不会损失任何非局部性。顺序设计提供了简单性和灵活性，同时比以前的方法执行得更好——每个单独的模块可以单独使用或以任何顺序组合使用，而并行设计没有这些好处。由于Max SA的灵活性和可扩展性，能够通过简单地将Max SA的替代层与MBConv堆叠在分层结构中来构建一个新的视觉主干，称之为MaxViT，如下图所示。MaxViT受益于整个网络的全局和局部感受野，从浅层到深层，在模型容量和泛化能力方面表现出优异的性能。

自注意力

自注意力允许整个空间（或序列）所有位置的空间混合，同时还受益于基于标准化成对相似性的内容相关权重。一般的模型中定义的标准自注意力是不知道位置的， 平移不变性是ConvNets中的一个重要归纳偏差。相对自注意力通过在注意力权重中引入相对学习偏差来提高普通注意力，这已被证明在许多视觉任务中始终优于原始注意力。在本文中，主要采用预归一化相对自注意力作为MaxViT中的关键算子。

多轴注意力

与局部卷积相比，全局交互是自注意力的关键优势之一。然而，直接沿整个空间应用注意力在计算上是不可行的，因为注意力算子需要二次复杂度。为了解决这个问题，本文提出了一种多轴方法，通过简单地分解空间轴，将全尺寸注意力分解为两种稀疏形式——局部和全局。$X\in R^{H\times W\times C}$是输入特征图。没有将注意力集中在平坦的空间维度HW上，而是将特征分为形状张量$（H/P\times W/P，P\times P，C）$，表示划分为不重叠的窗口，每个窗口大小为P×P。将自注意力集中在局部空间维度（即P×P）上等同于关注一个小窗口。将利用这一块注意力进行局部交互。 尽管绕开了完全自注意力的繁重计算，但在大规模数据集上，局部注意力模型的表现欠佳。受区块注意力的启发，本文提出了一种简单但有效的方法来获得稀疏的全局注意力，称之为网格注意力。代替使用固定窗口大小来划分特征图，使用固定的G×G均匀网格将张量网格化为形状（G×G，H G×W G，C），从而生成具有自适应大小$H/G\times W/G$的窗口 。在分解的网格轴（即G×G）上使用自注意力对应于令牌的扩展的全局空间混合。通过使用相同的固定窗口和网格大小，可以完全平衡局部和全局操作之间的计算，两者在空间大小或序列长度方面仅具有线性复杂度。请注意，本文提出的Max SA模块可以是Swin注意力模块的替代品，其参数和FLOP的数量完全相同。然而，它具有全局交互能力，而不需要掩码、填充或循环移位，使其更易于实现，比移位窗口方案更可取。例如，可以使用einops轻松实现多轴注意力，而无需修改原始注意力操作。

MaxViT block

如下图所示，将两种类型的注意力顺序叠加，以在单个块中获得局部和全局交互。本文还采用了Transformer中的典型设计，包括LayerNorm、前馈网络（FFN和跳过连接。还在多轴注意力之前添加了一个带有（SE）模块的MBConv块，将MBConv与注意力一起使用进一步提高了网络的泛化和可训练性。在注意力之前使用MBConv层提供了另一个优势，因为深度卷积可以被视为条件位置编码（CPE），使得模型没有显式位置编码层。提出的独立多轴注意力可以一起使用，也可以单独使用，用于不同的目的——局部交互的块注意力，以及全局混合的网格注意力。这些元素可以很容易地插入到许多视觉体系结构中，特别是在高分辨率任务中，这些任务可以通过与可负担计算的全局交互而受益。

架构变体

本文设计了一系列极其简单的架构变体来探索本文提出的MaxViT块的有效性。使用了类似于常见ConvNet的分层主干，其中输入首先在stem stage（S0）使用Conv3x3层进行下采样。网络主体包含四个层级（S1-S4），每个层级的分辨率为前一级的一半，通道数加倍（隐藏维度）。在网络中，在整个主干中使用相同的MaxViT块。在每个阶段的第一个MBConv块的深度方向Conv3x3层中应用下采样。反向瓶颈和挤压激励（SE）的膨胀率和收缩率默认为4和0.25。将所有注意力块的注意力头部大小设置为32。通过增加每个阶段的块数B和通道尺寸C来放大模型。在下表中总结了MaxViT变体的架构配置：

三、实验

(一)、分类实验

(2)、目标检测

(3)、消融实验

Global grid-attention:本文的主要贡献之一是网格注意力模块，它允许在线性时间进行稀疏的全局交互，使模型能够在所有阶段捕获全局信息。进行了两次消融以了解其增益：
1）在每个阶段完全消除了全局注意力；
2） 将网格关注替换为块关注，以保持相同的参数计数和FLOP。如下表所示，与仅使用局部注意力或卷积相比，在早期阶段启用全局注意力可以进一步提高性能。

MBConv层:本文还通过移除每个阶段中的所有MBConv来消融MaxViT中MBConv层的使用。请注意，在移除MBConv层时，我们还应考虑参数和FLOP的减少。此外，阶段3有5个块，而其他阶段只有2个块。如下表所示，MaxViT中MBConv层的使用显著提高了性能。

Block order study:本文提出了三个不同的模块来构建MaxViT块——MBConv、块和网格注意力——它捕获了从局部到全局的空间交互。为了研究组合它们的最有效方法，使用所有6种排列来评估MaxViT-T模型。总是在第一层应用下采样，这可能会导致较小的模型大小差异。我们可以从下表中观察到，在注意力层之前放置MBConv几乎总是优于其他组合。原因可能是，更适合在早期层中获取局部特征/模式，然后全局聚合它们，这与现有的混合模型相一致，这使Conv层成为人们关注的焦点。然而，在生成对抗实验中，我们发现从全局到局部的最佳顺序是：GA-BA-C。对于生成任务来说，首先使用全局处理块（即网格关注层）获得正确的总体结构，然后使用局部处理块（如MBConv）填充更精细的细节可能是有利的。
Sequential vs. parallel：在本文的方法中，也存在采用并行设计的其他模型。在本次消融中，将顺序Max-SA与分别包含块关注和网格关注的并行分支进行比较。使用一个输入投影来增加通道数量，然后分割注意力头来馈送两个分支，以保持与MaxViT相似的复杂性，并使用一个输出投影来减少级联的分支。本文进行了粗略的参数调整，发现初始学习率为$10^{-3}$。如下表所示，本文的顺序方法以更少的参数和计算显著优于并行方法。原因可能是并行设计学习互补的线索，而它们之间的交互较少，而顺序堆栈能够学习局部和全局层之间更强大的融合。
Vertical layout:我们进一步研究了我们的垂直布局设计，即每个阶段的块数。将本文的设计与Swin/ConvNeXt的选择进行了比较。将MaxViT-T和-S更改为块B=（2，2，6，2），MaxViT-B，-L更改为严格遵循Swin的阶段比的块B=(2, 2, 18, 2)。从下图中可以看出，本文的布局在小型模型中的表现与Swin相当，但在大型模型中的比例明显更好。

四、附录

MBConv:MaxViT利用MBConv块作为主卷积算子。本文还采用了预激活结构，以促进MBConv和Transformer之间的均匀性。具体地，假设x为输入特征，则无下采样的MBConv块被公式化为：

其中，Norm是BatchNorm，Conv是扩展通道的Conv1x1，然后是BatchNom和GELU，这是基于Transformer模型的典型选择。DWConv是Depthwise Conv3x3，后面是BatchNorm和GELU。SE是挤压激励层，而Proj是收缩Conv1x1，以向下投影通道数。对于每个阶段中的第一个MBConv块，通过应用步长为2的深度方向的Conv3x3进行下采样，而shortcut分支也应用l了池化层和通道投影：

Relative Attention:为了简化表达，只使用了多头自注意力的单个头部来演示模型。在实际实现中，总是使用具有相同维度的多头注意力。相对注意力可定义为：

其中$Q、K、V\in R^{（H\times W）\times C}$是查询、键和值矩阵，d是隐藏维度。注意力权重由学习的静态变量以及位置感知矩阵B和缩放输入自适应注意力$Q{K}^T/\surd d$共同决定的, 考虑到2D坐标的差异，相对位置偏差B由矩阵$\hat{B}\in R^{(2H-1)(2W-1)}$进行参数化的，当以更高分辨率（例如$H^{{}^{\prime }}\times W{}^{\prime }$）进行微调时，使用双线性插值来映射相对位置偏差从$R^{(2H-1)(2W-1)}$至$R^{(2H^{{}^{\prime }}-1)(2W^{{}^{\prime }}-1)}$ 这种相对注意力得益于输入适应性、平移不变性和全局交互，这是二维视觉任务中优于普通自注意力的首选。在本文中，所有注意力操作符默认使用上式中定义的相对注意力。
Multi-Axis Attention:假设上式中的相对注意力遵循1维输入序列的惯例，即始终将输入的第二个最后维度（…，L，C）视为空间轴，其中L，C表示序列长度和通道。所提出的多轴注意力可以在不修改自注意力操作的情况下实现。首先使用参数P定义$Block（\cdot ）$操作符作为对输入图像/特征$x\in R^{H\times W\times C}$的分块分为不重叠的块，每个块的大小为P×P。在窗口分割后，块的尺寸被聚集到空间维度（即-2轴）上：

将Unblock（·）操作表示为与上述块划分过程相反的操作。使用参数Grid(·)操作将网格定义为将输入特征划分为均匀的G×G网格，每个网格具有自适应大小$H/G\times W/G$。与块操作不同，此方法需要应用额外的Transpose来将网格维度放置在假定的空间轴（即-2轴）中：

其逆操作Ungrid（·）将网格化输入反转回正常的2D特征空间。 为此，本文准备解释多轴注意力模块。给定输入张量$x\in R^{H\times W\times C}$，局部块注意力可以表示为：

而全局扩展网格注意力模块被公式化为：

为了简单起见，在RelAttention操作中省略了QKV输入格式。LN表示层归一化，其中MLP是由两个线性层组成的标准MLP网络。
其他版本的详细架构：