MaxViT : 多轴Vision Transformer

        本文提出了一种高效、可扩展的多轴注意力模型，该模型包括两个方面:局部注意力被阻塞和全局注意力被扩张。这些设计选择允许在任意输入分辨率下的全局-局部空间交互，只有线性复杂度。还通过有效地将注意力模型与卷积混合在一起，提出了一个新的架构元素，并相应地提出了一个简单的分层视觉主干，称为MaxViT，通过简单地在多个阶段重复基本构建块。值得注意的是，即使在早期的高分辨率阶段，MaxViT也能够“看到”整个网络的全局。 

---

        受自然语言处理中Transformer[85]等自注意模型进化的启发，许多研究人员开始将注意机制引入视觉。Vision Transformer (ViT)可能是第一个完全基于Transformer的视觉架构，其中图像补丁被简单地视为单词序列，并且在这些视觉token上应用了一个Transformer编码器。当在大规模数据集上进行预训练时，ViT可以在图像识别上取得令人信服的结果。 

        如果没有广泛的预训练，ViT在图像识别方面表现不佳。这是由于Transformer具有较强的模型能力，具有较少的感应偏置，从而导致过拟合。为了适当地正则化模型容量并提高其可扩展性，许多后续工作研究了为局部注意力等视觉任务量身定制的稀疏Transformer模型。这些方法通常重新引入层次结构来弥补非局部性的损失。

        Swin Transformer通过在移位的非重叠窗口上应用自注意力来修改Transformer。该方法首次在纯视觉Transformer的ImageNet基准测试中优于ConvNets。尽管与ViT中使用的全注意力相比，基于窗口的注意力具有更大的灵活性和泛化性，但由于非局域性的丧失，基于窗口的注意力的模型容量有限，因此在ImageNet-21K和JFT等更大的数据体系上扩展不利。

1. 引言

        提出了一种新型的Transformer模块，称为多轴自注意力(Max-SA)，它能够作为一个基本的架构组件，在单个块中执行局部和全局空间交互。与完全自注意力相比，Max-SA具有更大的灵活性和效率，即对线性复杂度的不同输入长度具有自然适应性; 与(移位的)窗口/局部注意相比，Max-SA通过提出一个全局接受域来允许更强的模型容量。此外，Max-SA仅具有线性复杂性，可以在网络的任何层中用作通用的独立注意力模块，即使在早期的高分辨率阶段也是如此。

        进一步设计了一种简单而有效的视觉骨干，称为多轴视觉变压器(MaxViT)，通过分层堆叠由Max-SA和卷积组成的重复块。虽然提出的模型属于混合Vision Transformer的范畴，但MaxViT与以前的方法不同，因为力求简单，通过设计一个统一卷积、局部和全局注意力的基本块，然后简单地重复它。实验表明，在所有数据体制下，MaxViT显著提高了最先进的(SOTA)性能，适用于广泛的视觉任务，包括分类、目标检测和分割、图像美学评估和图像生成。

MBConv

        MaxViT中的MBConv（Mobile Inverted Bottleneck Convolution）是一个关键组件，继承了MobileNetV2中的倒置瓶颈结构，并且结合了深度可分离卷积来提高效率和减少参数量。MaxViT利用MBConv来处理局部信息，在保持高效计算的同时，还增强了特征提取的能力。

MBConv的操作流程大致如下：

1. 输入与扩展（Expand）

        通道数扩展：MBConv首先通过一个逐点卷积（Pointwise Convolution，也称为1x1卷积）将输入的通道数扩展到一个更高的维度。假设输入的通道数为in_channels，则该卷积通常会扩展成in_channels * expansion_factor，其中expansion_factor是一个倍率因子，比如6。

        激活函数：扩展后的特征通过一个非线性激活函数（如ReLU6或GELU），增加特征的表达能力。

2. 深度可分离卷积（Depthwise Convolution）

        空间卷积：扩展后的特征经过一个深度可分离卷积。深度可分离卷积操作是指对每个通道单独进行卷积操作，而不是传统的卷积对所有通道一起操作。这样可以大大减少计算量和参数量。

        保留空间信息：该卷积主要用于捕获局部的空间信息，同时不会增加额外的通道数。

3. 投射（Projection）

        通道数投射：经过深度可分离卷积后，特征通过另一个逐点卷积，将通道数降回输入的通道数in_channels。这个操作被称为“投射”，用于恢复原始的通道维度。

        跳跃连接（Skip Connection）：如果输入的分辨率和投射后的特征图的尺寸相同（即不涉及下采样），则可以直接将输入添加到输出上。这种跳跃连接（也称为残差连接）可以帮助网络更好地学习到恒等映射，减少梯度消失问题。

4. Squeeze-and-Excitation (SE) 模块

        全局信息捕获：MBConv集成了SE模块。SE模块通过对特征图进行全局池化，获得每个通道的全局特征，然后通过一系列全连接层调整通道的权重，从而动态地调节通道的权重。这有助于网络自动增强对重要特征通道的关注，同时抑制不重要的通道。

• 先是一个全局池化将特征图缩小为1x1大小；
• 然后通过两个全连接层对池化后的特征进行缩放和激活，生成权重；
• 最后将这些权重乘回到特征图上，以动态调整每个通道的贡献。

5. 总结流程

        MaxViT利用MBConv来提取局部特征，结合Transformer的全局建模能力，使得网络能够在多个尺度上对图像进行建模。

2. 相关工作

2.3 混合模型

        纯基于Transformer的视觉模型由于相对较少的归纳偏差而泛化较差。Vision Transformer 也表现出不合格的可优化性。一个有趣的简单改进是采用Transformer层和卷积层的混合设计，例如使用几个卷积来替换粗糙的patchify stem（将图像分割成一系列小块的步骤）。

3. 方法

Blocked Multi-Axis Self-Attention (Max-SA)

        受到稀疏方法的启发，研究者们提出了一种新的注意力模块，称为“Blocked Multi-Axis Self-Attention”（Max-SA）。这个模块通过将全密集（fully dense）的注意力机制分解为两种稀疏形式（窗口注意力和网格注意力）来降低计算复杂度。

稀疏形式的优点

• 降低复杂度：传统的注意力机制（如Transformer中的自注意力）