以下文章来源于极市平台 ，作者CV开发者都爱看的
导读
本文引入局域性的方法是受到以往 Shift 操作的启发，提出了一种轴向移位策略，在水平和垂直方向上对特征进行空间移位。基于这种轴向移位策略，得到了一个轴向移位 MLP 视觉骨干架构 AS-MLP，也是第1个适用下游任务 MLP 类视觉骨干架构。
本文目录
1 AS-MLP：首个适用下游任务的轴向移位 MLP 视觉骨干架构
(来自上海科技大学)
1 AS-MLP 论文解读
1.1 背景和动机
1.2 AS-MLP 总体设计
1.3 AS-MLP Block 的设计
1.4 Axial Shift 操作，卷积，Transformer，和 MLP-Mixer 的区别
1 AS-MLP：首个适用下游任务的轴向移位 MLP 视觉骨干架构
论文名称：AS-MLP: An Axial Shifted MLP Architecture for Vision
论文地址：https://arxiv.org/pdf/2107.08391.pdf
1.1 背景和动机
这篇文章诞生于 MLP 类模型诞生的中期，是当时一种高性能的新式视觉 MLP 类骨干网络。视觉 MLP 类模型的开山是 Tolstikhin 等人首次提出的 MLP-Mixer，其中几乎所有的网络参数都是从 MLP (线性层)中学习的，其设计范式遵循了传统的 ViT 类架构，但是通过多层感知机 (multi-layer perceptrons, MLP) 来代替较为复杂的注意力机制，同时可以进一步减少归纳偏置。在 MLP-Mixer 中，Channel-Mixing MLP 沿通道维度进行信息的建模，Token-Mixing MLP 沿空间维度进行信息的建模。
然而，MLP-Mixer 中的 Token-Mixing MLP 本质上是一个矩阵乘的操作，它对全部的空间信息建模，而难以充分利用局部信息，而局部信息在 CNN 架构中是非常重要的，因为不是所有像素都需要长期依赖关系。而且，局部信息更侧重于提取底层特征。
为了把局部信息引入 Transformer 类架构中，最简单和最直观的想法之一是向 MLP-Mixer 添加一个类似 Window 的限制，然后在 Window 内执行局部特征的 Token Mixing，典型的做法是 Swin Transformer[1]。其他一些论文也已经强调了局部感受野的优势，并在 Transformer 中引入了局部信息，比如 Localvit[2]，NesT[3]等。
要引入局部性，就需要在任意位置，直接对特征点与其周围特征点之间的关系进行建模。本文引入局域性的方法是受到以往 Shift 操作的启发，提出了一种轴向移位策略，在水平和垂直方向上对特征进行空间移位。基于这种轴向移位策略，得到了一个轴向移位 MLP 视觉骨干架构 AS-MLP，也是第1个适用下游任务 MLP 类视觉骨干架构，因为其权重不随输入
分辨率的变化而变化，使得用固定图像大小训练的模型权重能够适应不同输入大小的下游任务。AS-MLP 在 88M parameters 和 15.2 GFLOPs 的复杂度下在 ImageNet 上获得了 83.3% 的 Top-1 精度。通过 ImageNet-1K 数据集中的预训练模型，AS-MLP 在 COCO 上获得 51.5 mAP，在 ADE20K 上获得 49.5 MS mIoU，与基于 Transformer 的架构相比也有竞争力。
1.2 AS-MLP 总体设计
AS-MLP 总体设计遵循 Swin Transformer。给定输入
为 , AS-MLP 首先进行 Patch Embedding 操作, 将原始图像分成多个大小为 的 Patch。Patch Embedding 操作输出的特征维度是 。AS-MLP 共 4个阶段, 如下图1所示。各个阶段特征的分辨率减半的同时, 宽度加倍。
图1：AS-MLP 总体架构
AS-MLP 不同型号的模型各个 Stage 的详细配置如下。AS-MLP-S 相比于最小的 AS-MLP-T 主要是增加了模型的深度，最小的 AS-MLP-B 相比于 AS-MLP-S 主要是增加了宽度。每一种模型的 Params 和 FLOPs 如下图2所示。
AS-MLP-T: C = 96, 4个 stage 的 Block 数分别为：{2, 2, 6, 2}
AS-MLP-S: C = 96, 4个 stage 的 Block 数分别为：{2, 2, 18, 2}
AS-MLP-B: C = 128, 4个 stage 的 Block 数分别为：{2, 2, 18, 2}
图2：AS-MLP 不同型号模型的配置
1.3 AS-MLP Block 的设计
AS-MLP Block 的内部结构如下图3所示，主要由 Norm 层、Axial Shift 操作、MLP 和残差连接组成，和之前的工作不一样的就是大图中的 Axial Shift 操作。
图3：AS-MLP Block 的内部结构
以图3中的 Horizontal Shift 为例, 如图4所示, 假设输入维度是 , 假设图4中的 , 将输入特征分为3部分, 分别沿水平方向移动 个单位, 并使用 zero padding, 之后, 虚线框中的特征将被取出, 用于下一次的 Channel Projection。Vertical Shift 执行类似的操作, 可以将来自不同空间位置的信息组合在一起。
图4：AS-MLP Block 的设计
AS-MLP Block 的 PyTorch 伪代码如下：
# norm: normalization layer
# proj: channel projection
# actn: activation layer
import torch
import torch.nn.functional as F
pad = shift_size // 2
def shift(x, dim):
x = F.pad(x, (pad, pad, pad, pad) , "constant", 0)
x = torch.chunk(x, shift_size, 1)
x = [ torch.roll(x_s, shift, dim) for x_s, shift in zip(x, range(-pad, pad+1))]
x = torch.cat(x, 1)
return x[:, :, pad:-pad, pad:-pad]
def as_mlp_block(x):
shortcut = x
x = norm(x)
x = actn(norm(proj(x)))
x_lr = actn(proj(shift(x, 3)))
x_td = actn(proj(shift(x, 2)))
x = x_lr + x_td
x = proj(norm(x))
return x + shortcut
torch.chunk(input,chunks,dim=0) 是把输入分块的 PyTorch 函数。
torch.cat(tensors,dim=0,*,out=None) 是把输入 Concatenate 起来的 PyTorch 函数。
torch.roll(input,shifts,dims=None) 是把输入沿着给定方向滚动的 PyTorch 函数。
1.4 Axial Shift 操作，卷积，Transformer，和 MLP-Mixer 的区别
首先下图5所示是卷积、Swin Transformer、MLP-Mixer 和 AS-MLP 的不同采样位置，图中 为 shift size， 为 dilation。
图5：卷积、Swin Transformer、MLP-Mixer 和 AS-MLP 的不同采样位置
Convolution： 对于卷积操作，对输入 执行 的滑动卷积核，得到输出：
Swin Transformer： 在基于 Transformer 的架构中引入了 Window，以覆盖局部注意力：
MLP-Mixer： 放弃了注意力机制, 首先对输入 进行转置, 然后附加一个 Token Mixing MLP, 以获得输出 :
其中，为 Token Mixing MLP 中的可学习权重。MLP-Mixer 仅通过矩阵转置来感知全局信息。
AS-MLP： 给定输入 , shift size , dilation , 首先将 在水平和垂直方向上分成 段, 并轴向移位后, 输出 为:
其中, 分别为水平方向和垂直方向的可学习的权重。
与 MLP-Mixer 相比，AS-MLP 在水平和垂直方向上使用轴向位移，更关注这两个方向上的信息交换。
1.5 实验结果
ImageNet-1K 图像分类
实验结果如下图6所示。输入分辨率为224 × 224，在相似的参数和 FLOPs 时，本文提出的 AS-MLP 优于其他的 MLP 类架构。例如，AS-MLP-S 比 Mixer-B/16 和 ViP-Medium/7 以更少的参数获得更高的 Top-1 精度 (83.1%)。此外，与基于 Transformer 的架构相比，也实现了具有竞争力的性能，例如 AS-MLP-B (83.3%) 和 Swin-B (83.3%) 的精度相当，这表明了AS-MLP 架构的有效性。
图6：ImageNet-1K 图像分类实验结果
除了标准实验，作者还比较了 AS-MLP 在小模型下的结果，如图7所示。Swin (mobile) 模型和 AS-MLP(mobile) 模型有相似的参数 (约10M)。实验结果表明，Swin (mobile) 模型在移动环境下显著超过 Swin Transformer (76.05% vs. 75.11%)。
图7：小模型图像分类实验结果
COCO 目标检测和实例分割
对于目标检测和实例分割，作者使用 mmdetection 作为框架，COCO 作为评估数据集，该数据集由 118K 训练集和 5K 验证集组成。作者考虑了两种典型的目标检测框架：Mask R-CNN 和 Cascade R-CNN。所有的骨干都使用 ImageNet-1K 上预训练的权重进行初始化，所有模型都使用 3x schedule (36 epochs)。
如下图8所示是 COCO 目标检测实验结果。AS-MLP 是第一个将 MLP 类的架构应用于目标检测的工作。在类似的资源限制下，AS-MLP 的性能与 Swin Transformer 相当。其中，Cascade Mask R-CNN + Swin-B 在 145M 参数和 982 GFLOP s下实现了 51.9 的 APb，Cascade Mask R-CNN + AS-MLP-B 在 145M 参数和 961 GFLOPs下实现了 51.5 的 APb。
图8：COCO 目标检测和实例分割实验结果
ADE20K 语义分割
作者在具有挑战性的语义分割数据集 ADE20K 上进行了实验，该数据集包含 20,210 张训练集和 2,000 张验证集，并用 UperNet 作为分割头。该框架基于 mmsegmentation。在 ImageNet-1K 上预训练权重初始化所有模型，并训练 160K iterations。
在 FLOPs 略低的情况下，AS-MLP-T 比 Swin-T 获得了更好的结果 (46.5 MS mIoU vs. 45.8 MS mIoU)。对于大型模型，UperNet + Swin-B 在 121M 参数和 1188GFLOPs 下获得了 49.7 MS mIoU。而 UperNet + AS-MLP- B 在 121M 参数和 1166GFLOPs 下获得 49.5 MS mIoU，这也表明了 AS-MLP 在处理下游任务方面的有效性。
可视化结果
作者将从 Swin Transformer 和 AS-MLP 中学到的特征的 heatmap 可视化在图9中，其中：
第1列：来自 ImageNet 的图像。
第2列：Swin-B 最后一层的 activation heatmap。
第3，4，5列：最后一层的 horizontal shift，vertical shift，以及二者的加和之后的 activation heatmap。
可以看出：
AS-MLP 与 Swin 相比，可以更好地聚焦对象区域
AS-MLP (h) 可以更好地聚焦物体的垂直部分，因为它在水平方向上作了特征的偏移，水平方向的偏移可以覆盖垂直部分的边缘，更有利于物体的识别。
AS-MLP (v) 可以更好地聚焦物体的水平部分，因为它在垂直方向上作了特征的偏移，垂直方向的偏移可以覆盖水平部分的边缘，更有利于物体的识别。
图9：Swin Transformer 和 AS-MLP 的特征可视化
总结
AS-MLP，一种轴向移位的 MLP 视觉骨干架构，诞生于 MLP 类架构中期，它是第1个适用下游任务 MLP 类视觉骨干架构，因为其权重不随输入
分辨率的变化而变化，使得用固定图像大小训练的模型权重能够适应不同输入大小的下游任务。AS-MLP 在 88M parameters 和 15.2 GFLOPs 的复杂度下在 ImageNet 上获得了 83.3% 的 Top-1 精度。通过 ImageNet-1K 数据集中的预训练模型，AS-MLP 在 COCO 上获得 51.5 mAP，在 ADE20K 上获得 49.5 MS mIoU，与基于 Transformer 的架构相比也有竞争力。
参考
^Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
^Localvit: Bringing locality to vision transformers
^Aggregating nested transformers
转自：“arXiv每日学术速递”微信公众号
如有侵权，请联系本站删除！
