AS-MLP:注意力驱动下的多层感知机升级

0.引言

AS-MLP模型出自上海科技大学和腾讯优图实验室共同合作发表的文章，题为AS-MLP: AN AXIAL SHIFTED MLP ARCHITECTURE FOR VISION。纯MLP网络架构专注于全局的信息交流，却忽略了局部信息的收集。然而在计算机视觉中，局部特征非常重要，它们可以提供传统计算机视觉中定义的底层特征（边缘、纹理、拐点等），这对于目标检测、语言分割等计算机视觉任务意义重大。

为此，该论文提出了轴向转移（Axial Shift）操作，通过将特征图沿着空间轴向进行偏移，使得不同空间位置的特征位于同一通道中，随后进行通道投影来融合不同位置的特征，从而实现局部信息的交流（其思想与ShuffleNet有一定相似之处）。这样的操作能够利用一个纯粹的MLP架构来实现与类似卷积神经网络架构相同的局部感受野。AS-MLP也是首个迁移到下游任务的MLP架构，其骨干网络的设计思想和Swin Transformer模型基本一致。AS-MLP在ImageNet数据集上取得了优于其他MLP架构的性能，并在COCO数据集上进行目标检测任务与ADE20K数据集上进行语义分割任务上，实现了与Swin Transformer作为骨干网络相当的性能。

论文名称：AS-MLP: an axial shifted mlp architecture

下载地址：https://arxiv.org/pdf/2107.08391

1.AS-MLP模块

AS-MLP（Adaptive-Structured Multi-Layer Perceptron）Block和Transformer Block的整体结构相似，但将多头注意力部分替换为Axial Shift块。单个Axial Shift模块包含四个全连接层，分别对应四个Channel Projection块，如图1所示。

首先，当特征图输入Axial Shift模块后，对通道信息进行一个全连接层的映射，这个全连接层其实是通过1×1卷积操作实现的，目的是实现通道维度上的信息交互。中间使用了一个并行结构，分别对特征图进行水平和竖直方向的循环移位。对水平方向进行举例，具体操作如图2所示。

如图2所示，当shift_size=3，dilation=1时，意味着相邻三行特征图分别位移像素{+1,0,-1}，之后的每三个特征图也按照这种方式进行循环位移，直到最后一行特征图。若shift_size=5，dilation=2，则意味着相邻五个特征图分别位移像素{+4,+2,0,-2,-4}；之后的每五个特征图也按照这种方式进行循环位移，直到最后一行特征图。

竖直方向的循环移位操作与此类似，不再赘述。这样不同位置的特征被重新对齐在一个通道上，然后对通道进行全连接，这里是通过1×1卷积实现不同位置通道信息的交流。接着将水平和竖直方向移动后的结果相加，便得到局部交流之后的结果。最后对通道进行一个全连接，将融合后的信息进行整合，这里同样是通过1×1卷积实现的。这便是Axial Shift模块的数据处理流程。

关于Axial Shift模块中对特征图进行水平和竖直方向的循环移位，为什么采用并行处理而非串行处理？笔者认为这是因为竖直和水平在语义上是一致的，因此我们没有理由为它们安排先后处理顺序。而且，通过实验结果发现，并行处理的效果更好。实验结果见表1。

接下来将进一步分析模型的感受野，具体示例如图3所示：

（1）传统的稠密卷积具有方形感受野，仅具有局部依赖性。

（2）Swin Transformer在窗口内使用自注意力机制，使得感受野涵盖整个窗口，具有局部依赖性。

（3）MLP-Mixer采用全连接层，相当于一个超大的卷积，其感受野涵盖整个特征图，具有长距离依赖性，因此对于局部特征的把握能力较差。

（4）AS-MLP将水平和垂直方向的移动解耦，并将它们相加，从而获得十字形感受野。通过调整感受野的大小和扩张因子，可以获得不同的感受野。不同的设置使得采样位置包含局部依赖和长距离依赖的信息。

最后，关于循环移位过程中，出现的空位是采用零填充（Zero-padding）还是循环填充（Cycle-padding）好呢？笔者认为直观上是零填充好。在循环位移（Circular Shift）过程中，当向水平方向移动时，最左边一列会移动出去，最右边一列会空出来，导致一列的信息丢失。循环填充可以将最左边移动出去的那一列补充到右侧的缺失一列上。然而，这并不是一个好的选择，因为图像最左边和最右边的信息未必相关，或者说在大多数情况下，它们是无关的。强制要求图像最左边和最右边的信息交流不一定会起到作用。论文中的实验结果也表明零填充的方法对模型效果更好，具体的数据见表2。

在AS-MLP Block中，除了以上介绍的Axial Shift模块，剩余的操作与Transformer Block的结构是非常相似的。在进入Axial Shift模块之前和Axial Shift模块计算之后，都会经过一个层归一化操作，再经过一个MLP层级结构，整个AS-MLP Block中，存在两个残差结构。AS-MLP Block的数据计算流程代码如下：

x = Axial_Shift(LayerNorm(input)) + input  
result = MLP(LayerNorm(x)) + x

2.AS-MLP模型结构

AS-MLP和Swin Transformer类似，其模型结构如图4所示。在每个阶段中，AS-MLP Block被重复了多次。

首先，从全局角度来看AS-MLP的工作原理：对于一个3×H×W的输入RGB图像，首先进行大小为4×4的图像块切片。这相当于将3×4×4大小的子图像像素展平并叠加在一起，形成长度为48的向量。此时，输入图像的形状变为48×H/4×W/4。选择4×4的Patch大小是为了实现输入数据的四倍空间下采样，这与ResNet等过去的卷积神经网络类似，它们首先通过卷积核尺寸为7×7且stride=2的卷积进行了一次下采样。然后经过BN和ReLU操作后再经过最大池化处理，又进行了一次下采样，随后输入多个Stage中。所以这里一开始也进行了两倍的下采样，与过去的Backbone保持一致。

随后， 48×H/4×W/4的特征图被视为H/4×W/4个Token，每个Token的维度为48。将其经过一个线性嵌入层，将通道数变为C=96、128等，对应不同大小的模型。之后经过AS-MLP Block。AS-MLP Block延续了卷积神经网络和Transformer等的思想，并不对Token的维度进行修改，输入和输出保持一致。每经过一个Stage（多个AS-MLP Block）后，进行一次二倍下采样。因此，第一个Stage后的特征图为4C×H/8×W/8。后续的Stage与此类似，不再赘述。最后的输出结果经过全局池化和全连接层，就可以得到输出了。

基于ViT和Swin等工作的设置，AS-MLP也推出了Tiny、Small和Base三种不同大小配置的缩放模型。这主要涉及Token在一开始经过线性嵌入层时的维度，以及每个Stage中AS-MLP Block的重复次数。这样的话，不同的网络参数量接近，方便进行比较。

（1）AS-MLP-T: C = 96, the number of blocks in four stages = {2, 2, 6, 2}。

（2）AS-MLP-S: C = 96, the number of blocks in four stages = {2, 2, 18, 2}。

（3）AS-MLP-B: C = 128, the number of blocks in four stages = {2, 2, 18, 2}。

3.小结

AS-MLP模型是一种基于MLP的深度学习架构，旨在替代卷积神经网络解决计算机视觉任务。AS-MLP通过使用自适应结构和局部连接，实现了类似于卷积神经网络的性能。AS-MLP模型的主要特点如下：

（1）MLP为基础：AS-MLP使用多层感知机作为基本构建模块，而不是传统的卷积神经网络。MLP是一种全连接的神经网络，通常用于处理序列数据和分类任务。

（2）局部连接：为了降低计算复杂性和参数数量，AS-MLP采用局部连接策略。这意味着网络中的每个神经元仅与相邻神经元相连，类似于CNN中的局部感受野。这种设计有助于提高模型的效率和泛化能力。

（3）适用于计算机视觉任务：尽管AS-MLP是基于MLP构建的，但它在计算机视觉任务上表现出了与卷积神经网络相当的性能，如图像分类、物体检测和语义分割等。

（4）简单且高效：AS-MLP的一个重要优点是其结构相对简单，同时在保持高性能的同时降低了计算复杂性。这使得模型在资源受限的设备上也能表现出良好的性能。

（5）与MLP-Mixer相比，AS-MLP模型更注重局部特征的提取，提出了一个无参数化的方法（Axial Shift模块），对channel信息在水平方向与竖直方向上进行了解耦的信息融合，使得MLP架构拥有与卷积神经网络类似的局部感受野的概念（十字形感受野）。然而，是否十字形感受野可以更好地提取局部依赖关系？实际上，暂无理论证明或者过去的传统计算机视觉表明十字形感受野操作可以更好地提取视觉底层特征。因此，这种方法的贡献效果仍然值得商榷。