ECCV2022_MaxViT+：Multi-Axis Vision Transformer

先上图：

一、问题提出

如果没有广泛的预训练，ViT在图像识别方面表现不佳。这是由于Transformer具有较强的建模能力，但是缺乏归纳偏置，从而导致过拟合。有效的解决方法就是控制模型容量并提高其可扩展性，如Twins、LocalViT、Swin Transformer等，通常重新引入层次结构以弥补非局部性的损失，比如Swin Transformer在移位的非重叠窗口上self attention。由于这些模型普遍失去了类似于ViT的非局部性，即具有有限的模型容量，导致无法在更大的数据集上扩展（ImageNet-21K、JFT等）

如何实现对不同数据量的适应，如何有效结合局部与全局计算的优势成为本文要解决的目标。

二、模型方法

引入了一种新型的注意模块，称为dubbed blocked multi-axis self-attention (Max-SA)，通过将完全密集的注意机制分解为window attention和 grid attention两种稀疏形式，将self attention的二次复杂度降低为线性，而不损失任何非局域性non-locality

由于Max-SA灵活性和可伸缩性，通过简单地将Max-SA层与MBConv堆叠，构建一个新的backbone，称为MaxViT。

1、Multi-axis Attention

与局部卷积相比，全局交互作用是自注意的主要优势之一。通过分解空间轴得到局部（block attention）与全局（grid attention）两种稀疏形式。

尽管local-attention绕过full self-attention的繁重计算，但在大数据集上拟合不足。受block attention启发，提出了一种简单有效的方法来获得稀疏的全局注意力，称之为grid attention。

（Max-SA模块可以直接替换Swin注意模块，具有完全相同数量的参数和flops）。

具有全局交互能力，不需要masking、padding或cyclic-shifting，优于shifted window scheme。多轴注意可以很容易地用einops实现，而不需要修改原始的注意操作。

怎么理解呢？如下：block attention中设W为2，则在每个窗口进行attention。

grid attention是在每一个grid中，不同grid的相同颜色块进行attention。为什么能这么做：因为像素信息是高冗余的，因此可以采用这种方式。（图来自MobileViT网络讲解_哔哩哔哩_bilibili）

Multi-axis Attention不同于 Axial attention。如下图所示, 在 Axial attention 中 首先使用列注意力（column-wise attention），然后使用行注意力（ row-wise attention） 来计算全局 注意力, 相当于 O(N*N**0.5) 的计算复杂度。然而 Multi-Axis attention 则先采用局部注意力 (block attention), 再使用稀疏的全局注意力 (grid attention), 这样的设计充分考虑了图像的 2D 结构，并且仅具有 O(N) 的线性复杂度。

2、MaxViT block

将这两种类型的注意按顺序堆叠，以在单个块中获得局部和全局交互。在transformer中也采用了LayerNorm、FFNs和skip-connections。在多轴注意之前，还添加了一个MBConv块（包含squeeze-and-excitation(SE)模块（在mobileNet中有），可增加了网络的泛化和可训练性），深度卷积可以被视为条件位置编码(CPE)，使模型没有显式的位置编码层。独立多轴注意可以一起使用，也可以单独用于不同的目的——块注意用于局部交互，网格注意用于全局混合。这些元素可以很容易地插入到许多视觉体系结构中，特别是在高分辨率任务中，这些任务可以通过可负担的计算的全局交互获益。

MaxViT block：

MBConv(mobileNetv2中的结构):

不带下采样：

带下采样：

3、模型架构尺寸参数：

在SE、inverted bottleneck 中expansion rate=4和shrinkage rate=0.25，B和C表示每一级的block 数和channel数。attention heads=32。

怎么看呢？如MaxViT-T：

S0：Conv-stem层（纯卷积层），2个卷积层。第一层Conv：kernel_size=3*3，padding='same',outChannel=64，stride=2，输出feature map为(112,112,64)。第二层Conv：kernel_size=3*3,padding='same',outChannel为64，输出feature map为(112,112,64)。对输入进行下采样

S1：经过两个MaxViT block，in feature map为（112,112,64），out feature map为(56,56,64)

S2：经过两个MaxViT block，in feature map为（56,56,64），out feature map为(28,28,128)

S3：经过两个MaxViT block，in feature map为（28,28,128），out feature map为(14,14,256)

S3：经过两个MaxViT block，in feature map为（14,14,256），out feature map为(7,7,512)

Head：Pool层降维，FC层分类

模型细节：

三、实验

在各种视觉任务中的有效性:ImageNet分类、图像目标检测和实例分割、图像美学/质量评估和图像生成

1、ImageNet分类

在224×224px下，MaxViT在整个FLOPs上的表现比最新的强混合模型CoAtNet要好很多

2、Object Detection and Instance Segmentation

3、Image Aesthetic Assessment.

4、Image Generation

5、Ablation Studies

Table7：Global grid-attention。Table8：Block order study 。Table9：MBConv layer。Table10：Sequential vs. parallel

垂直布局设计，即每个阶段的块的数量（严格遵循Swin垂直布局设计进行比较）：

参考资料：MaxViT: Multi-Axis Vision Transformer - 知乎

代码：https://github.com/google-research/maxvit/

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