通用 Vision Backbone 超详细解读 ：S2-MLP：空间 Shift 操作实现通用基础视觉 MLP

作者丨科技猛兽

编辑丨极市平台

极市导读

本文介绍了一种新式的视觉 MLP 类模型，主要基于空间移位操作。它不需要 token-mixing MLP。相反，不同空间位置之间的通信是通过所提出的空间移位模块实现的。

加入极市CV技术交流群，走在计算机视觉的最前沿

本文目录

1 -MLP：空间 Shift 操作实现通用基础视觉 MLP
(来自百度)
1 -MLP 论文解读
1.1 背景和动机
1.2 空间 Shift 操作 MLP 架构
1.3 空间移位操作 (Spatial-shift operation)
1.4 -MLP 实验结果
1.5 改进的空间 Shift 操作
1.6 -MLPv2 实验结果

1 -MLP：空间 Shift 操作实现通用基础视觉 MLP

论文名称：

-MLP: Spatial-Shift MLP Architecture for Vision

-MLPv2: Improved Spatial-Shift MLP Architecture for Vision

论文地址：

https://arxiv.org/pdf/2106.07477.pdf

https://arxiv.org/pdf/2108.01072.pdf

这篇文章诞生于 MLP 类模型诞生的初期，是当时一种新式的视觉 MLP 类模型。视觉 MLP 类模型的开山是 MLP-Mixer，其设计范式遵循了传统的 ViT 类架构，但是通过多层感知机 (multi-layer perceptrons, MLP) 来代替较为复杂的注意力机制，同时可以进一步减少归纳偏置 (西瓜书解释 ：机器学习算法在学习过程中对某种类型假设的偏好。归纳偏好可以看作学习算法自身在一个庞大的假设空间中对假设进行选择的启发式或价值观)。

MLP-Mixer 中的基本模块由两个部分组成：channel-mixing MLP 和 token-mixing MLP 。channel-mixing MLP 沿通道维度 进行特征映射，用于不同 channel 之间的信息交互。同时，token-mixing MLP 沿空间维度 进行特征映射，，用于不同空间位置之间的信息交互。这种架构在 JFT-300M 等大规模数据集上训练后，MLP-Mixer 获得了良好的图像识别精度。但在中等规模数据集 ImageNet1K 和 ImageNet-22K 上，MLP-Mixer 和 ViT 的精度仍然存在差距。具体而言，Mixer-Base-16 在 ImageNet-1K 上的 Top-1 精度仅为 76.44%，而 ViT-Base-16 的 Top-1 精度为 79.67%。

MLP-Mixer 在 ImageNet-1K 和 ImageNet-22K 上性能的不理想促使作者重新思考 MLP- mixer 中的 token-mixing MLP 的设计。给定形状为 的特征, token-mixing MLP 的过程相当于是 , 其中 是可学习的权重矩阵。token-mixing MLP 的过程本质上是一个加权求和的过程，类似于 ViT 中的 attention。但是不同的是，ViT 中的 attention 具体的值是与输入图片的内容有关的，而 token-mixing MLP 可学习参数是与输入图片的内容无关的。token-mixing MLP 里加权求和的过程也很像 depthwise convolution，但卷积的感受野只是局部的。而相反，token-mixing MLP 具有全局的感受野。而且，depthwise convolution 的核权重在不同的空间位置之间是共享的，而 token-mixing MLP 权值在不同的空间位置是不同的。由于没有局部感受野的限制和空间位置共享权重的限制，token-mixing MLP 理论上应该比 depthwise convolution 具有更好的灵活性和更强的拟合能力。但是也同样伴随着容易过拟合的风险，需要大型数据集来支持。

本文提出一种空间 Shift 操作 MLP 架构，在空间位置之间进行通信时，作者采用无参数且计算效率高的空间移位操作。所构成的架构称之为 -MLP，它不需要 token-mixing MLP。相反，不同空间位置之间的通信是通过所提出的空间移位模块实现的。 -MLP 概念上很简单，在参数和计算量相当的情况下，在 ImageNet-1K 数据集上实现了比 MLP-Mixer 更高的精度。

-MLP 的基本架构如下图1所示。spatial-shift MLP 的骨干模型包含1个 patch-wise 的全连接层 (FC 层)， NN 个 -MLP 块。和一个最后的分类头 (1个 FC 层)。

图1：S2MLP 的基本架构

Patch-wise 的全连接层就是对于输入图片 , 首先分成大小为 的 Patch, 数量是 。然后把每个 Patch unfold 成 的向量, 并且通过一个 FC 层和 LN 归一化得到图片的 Embedding。

式中,

是全连接层的参数。

块的基本结构包含四个 FC 层、两个 LN 层、两个 GELU 激活函数、两个残差连接和 spatial-shift 块, 如上图1所示。用 表示 Spatial-Shift 模块输入中的特征图, 其中 表示宽度, 表示高度, 是通道数。Spatial-Shift 模块可分解为两个步骤：1) 将信道分成若干组。2) 将每组信道按不同方向移位。

如何分组

把特征 按照 channel 维度分为 组, 每组符号是 , 式中, 。默认情况下, 只沿着四个方向移动, 因此设置 为 4 。

把不同的组向着不同的方向移动。对于第1组通道 , 将它沿着宽度的方向移动 。对于第 2 组通道 , 将它沿着宽度的方向移动 。对于第 3 组通道 , 将它沿着高度的方向移动 。对于第4组通道 , 将它沿着高度的方向移动 。

空间移位操作的 PyTorch 伪代码：

          
def spatial_shift(x):  
    w,h,c = x.size()  
    x[1:,:,:c/4] = x[:w-1,:,:c/4]  
    x[:w-1,:,c/4:c/2] = x[1:,:,c/4:c/2]  
    x[:,1:,c/2:c*3/4] = x[:,:h-1,c/2:c*3/4]  
    x[:,:h-1,3*c/4:] = x[:,1:,3*c/4:]  
    return x   

      

在空间移位操作之后，每个 Patch 从相邻的 Patch 中获得了相应的视觉信息。空间移位操作通过一种无参数的手段，使得不同空间位置之间的通信成为可能。上述空间移位方式是最简单、最直接的移位方法之一。使用空间移位操作，使得整个模型不再需要 token mixer 作为混合不同空间位置信息的操作。模型只需要沿着通道维度投影特征。但是，一个 Block 中的空间移位操作只能获得相邻 Patch 的视觉内容，不能访问图像中所有 Patch 的视觉内容。所以作者这里堆叠 个 S2-MLP 块，全局的视觉内容会逐渐扩散到每个 Patch。

与 Depth-wise Convolution 的关系

空间移位操作等价于具有固定权重的 Depth-wise Convolution。设一组 Depth-wise Convolution 的核为。空间移位操作等价于下面的这个 Depth-wise Convolution，它的 Groups 为4，Kernel 的参数如下式所示：

也就是说，空间移位操作本质上是固定权重的深度卷积的变体。空间移位操作在组内共享权重，对不同的组采用不同的权重。下图2总结了它们之间的联系和区别。

图2：Token Mixer，空间移位操作和 Depth-wise Convolution 的区别与联系

-MLP 的超参数如下图3所示，分为 Wide 模型和 Deep 模型。

图3：S2-MLP 的超参数

ImageNet-1K 实验结果

ImageNet-1K 图像分类结果如下图4所示。与 ViT 相比，Mixer-B/16 在精度上没有竞争力。相比之下，-MLP 获得了与 ViT 相当的精度。通过引入一些精心设计，ViT 之后的 PVT-Large、TNT-B、CaiT、Swin-B、Nest-B 等工作的性能得到了很大的改善。但是 S2-MLP 不能达到与最先进的基于 ViT 的视觉模型一样高的识别精度。

作者将 -MLP 架构与最近提出的 MLP 对应架构进行了比较，包括 MLP- mixer, ResMLP-36 等。它们都采用了类似的结构。ResMLP-36 和 MLP-Mixer 的不同之处在于 ResMLP36 开发了一个仿射转换层来取代 LN，从而实现更稳定的训练。如下图4所示，与 MLP-Mixer 和 FF 相比，所提出的 -MLP-Wide 具有更高的识别精度。具体而言，MLP-Mixer 的 Top-1 准确率仅为 76.4%，FF 的 Top-1 准确率仅为 74.9%。相比之下，-MLP-Wide 的 Top-1精度为 80.0%。与 ResMLP-36 相比，-MLP-Deep 具有更高的识别精度。

图4：S2-MLP ImageNet-1K 实验结果

以上是模型 -MLP 的介绍，作者团队又相继推出了更强的 -MLPv2 模型。在本节中，介绍 -MLPv2 模型。与 -MLP 模型类似，-MLPv2 模型也由 Patch-wise 的全连接层、-MLPv2 基本模块和分类头组成。

-MLPv2 基本模块的架构如下图5所示。可以看到，它包含了两个 Spatial Shift 操作，和一个 Split Attention 操作。

图5：S2MLPv2 的基本架构

-MLPv2 和 v1 基本模块的对比如下图6所示。可以发现，v1 是直接采用 Shift 操作，但是 v2 先对每个分支分别采用 Shift 操作，再通过 Split Attention 把特征图汇合在一起。

图6：S2MLPv2 和 v1 基本模块的对比

Split Attention 操作

设 个尺寸为 的特征为 , 其中 是 Patch 的数量, 是 channel 的 数量, Split-Attention 操作首先对它们进行平均, 得到:

其中, 是一个 维的行向量, 这样得到的 就是一个 维的列向量。再把 通过一个 MLP, 得到：

其中, 为 GELU 实现的激活函数, 和 为 MLP 的权值, 输出的 。然后 Reshape 为矩阵 , 再由一个 Softmax 函数沿第一维进一步处理, 生成 。然后它生成参与的特征图:

-MLPv2 Block 的架构

Block 的架构遵循 MLP 类模型的一般设计的范式, 每个 Block 内部依次执行 tokenmixing 操作和 Channel-mixing 操作。token-mixing 操作是使用图6所示的移位方式实现的。

设输入特征为 Block 首先把 Channel 维度扩展3倍, 得到：

再沿着 Channel 的维度分成3份：

对于 , 执行2式的操作。对于 , 执行下式的操作:

对于，不执行任何操作。把这3部分的输出结果通过 Split Attention，最后再通过一个 MLP 层得到最终的输出。

-MLPv2 的移位操作的 PyTorch 伪代码如下：

          
def spatial_shift1(x):  
     b,w,h,c = x.size()  
     x[:,1:,:,:c/4] = x[:,:w-1,:,:c/4]  
     x[:,:w-1,:,c/4:c/2] = x[:,1:,:,c/4:c/2]  
     x[:,:,1:,c/2:c*3/4] = x[:,:,:h-1,c/2:c*3/4]  
     x[:,:,:h-1,3*c/4:] = x[:,:,1:,3*c/4:]  
     return x  
  
def spatial_shift2(x):  
     b,w,h,c = x.size()  
     x[:,:,1:,:c/4] = x[:,:,:h-1,:c/4]  
     x[:,:,:h-1,c/4:c/2] = x[:,:,1:,c/4:c/2]  
     x[:,1:,:,c/2:c*3/4] = x[:,:w-1,:,c/2:c*3/4]  
     x[:,:w-1,:,3*c/4:] = x[:,1:,:,3*c/4:]  
     return x  
  
class S2-MLPv2(nn.Module):  
    def __init__(self, channels):  
        super().__init__()  
        self.mlp1 = nn.Linear(channels,channels*3)  
        self.mlp2 = nn.Linear(channels,channels)  
        self.split_attention = SplitAttention()  
    def forward(self, x):  
        b,w,h,c = x.size()  
        x = self.mlp1(x)  
        x1 = spatial_shift1(x[:,:,:,:c/3])  
        x2 = spatial_shift2(x[:,:,:,c/3:c/3*2])  
        x3 = x[:,:,:,c/3*2:]  
        a = self.split_attention(x1,x2,x3)  
        x = self.mlp2(a)  
        return x  

      

-MLP V2 的金字塔结构

遵循[1]的设计方式，-MLP V2 也设计了一个两阶段的金字塔式架构以提升性能。各个阶段的参数配置如下：

图7：S2-MLPv2 的超参数

ImageNet-1K 实验结果

ImageNet-1K 图像分类结果如下图8和图9所示。作者首先在图8中将 -MLPv2 与 MLP-Mixer，EAMLP，ResMLP，gMLP，S2-MLP-deep，CCSMLP，GFNet，AS-MLP，CycleMLP，ViP 等进行对比。如图8所示，-MLPv2 在使用相当数量的参数时，在小模型和中模型的设置中都始终优于其同类模型，甚至比 AS-MLP，CycleMLP 和 ViP 的大模型表现更好。

图8：S2-MLPv2 ImageNet-1K 实验结果 (与 MLP 模型对比)

如下图9所示是 -MLPv2 模型与 CNN 和 Transformer 的对比。 -MLPv2-Medium 在使用更少的参数的情况下实现了与 Transformer 差不多的精度。与 Vision Transformer 相比，-MLPv2 模型的设计要简单得多，需要的参数也少得多，使其在实际部署中具有竞争力。

总结

本文介绍了一种新式的视觉 MLP 类模型，主要基于空间移位操作。它不需要 token-mixing MLP。相反，不同空间位置之间的通信是通过所提出的空间移位模块实现的。空间移位操作是一种无参数，零计算量，且能够帮助视觉基础架构完成空间位置通信的高效操作手段。基于它构建的模型在概念上简单，不需要注意力机制，同时在视觉识别任务上实现了与 CNN 和 Transformer 相比很有竞争力的精度，同时也超过了当时主流的视觉 MLP 类架构。

参考

1. ^Vision permutator: A permutable mlp-like architecture for visual recognition