S2-MLP: Spatial-Shift MLP Architecture for Vision

Abstract最近，Visual Transformer (ViT)及其后续作品放弃了卷积，采用了自注意运算，达到了与CNN相当甚至更高的精度。最近，MLP-Mixer放弃了卷积和自注意操作，提出了一个只包含MLP层的架构。为了实现跨patch communication，在channel-mixing MLP的基础上，设计了一个额外的token-mixing MLP。当在极其大规模的数据集上进

HHzdh

2276人浏览 · 2021-10-21 15:38:39

HHzdh · 2021-10-21 15:38:39 发布

Abstract

最近，Visual Transformer (ViT)及其后续作品放弃了卷积，采用了自注意运算，达到了与CNN相当甚至更高的精度。最近，MLP-Mixer放弃了卷积和自注意操作，提出了一个只包含MLP层的架构。为了实现跨patch communication，在 channel-mixing MLP的基础上，设计了一个额外的 token-mixing MLP。当在极其大规模的数据集上进行训练时，它获得了很好的结果。但在ImageNet1K和ImageNet21K等中等规模的数据集上进行训练时，它的表现不如CNN和ViT。MLP-Mixer的性能下降促使我们重新思考 token-mixing MLP。我们发现，MLP-Mixer中的 token-mixing操作是一个具有全局接收域和空间特定配置的深度卷积变体。但是全局接收域和空间特性使得 token-mixing MLP容易过度拟合。在本文中，我们提出了一种新的纯MLP结构——空间移位MLP (S2-MLP)。与MLP-Mixer不同，我们的S2-MLP只包含通道混合MLP。我们设计了一个空间移位操作spatial-shift来实现patch之间的通信。它有一个当地的接收域，是空间不可知的。同时，该方法无参数，计算效率高。在ImageNet-1K数据集上训练时，所提出的S2-MLP比MLP-Mixer具有更高的识别精度。与此同时，S2-MLP在ImageNet-1K数据集上实现了与ViT一样出色的性能，架构更简单，FLOPs和参数更少。

1. Introduction

在自然语言处理NLP的[31]基础上，提出了visual Transformer (ViT)[6]。它用Transformer中的自注意操作代替了CNN中的卷积操作，对图像不同空间位置的局部patch之间的视觉关系进行建模。ViT和后续工作[29,34,32,20,9,33,30]取得了与CNN模型相当甚至更好的效果。相比CNN一般需要对卷积核进行细致的设计，ViT只是简单地堆叠了一系列标准的Transformer块，减少了手工操作，减少了归纳偏差。

最近，MLP-Mixer[27]提出了一个完全基于多层感知器(MLP)的更简单的替代方案，以进一步减少归纳偏差。MLP-Mixer的基本模块由两部分组成:channel混合MLP和token混合MLP。信道混合MLP沿着信道维度投射特征映射，从而实现不同信道之间的通信。同时，token混合MLP沿着空间维度投射特征映射，利用空间位置之间的通信。在超大数据集(如JFT-300M[24])上训练时，MLP-Mixer获得了良好的识别精度。但在ImageNet-1K和ImageNet21K[5]等中型数据集上，MLP-Mixer和ViT之间仍然存在精度差距。具体来说，Mixer-Base-16[27]在ImageNet-1K上仅达到76.44%的top-1精度，而vite-base-16[6]达到79.67%的top-1精度。

MLP- Mixer在ImageNet-1K和ImageNet-21K上的不理想性能促使我们重新思考MLP-Mixer中的token-mixing MLP。给定矩阵形式的 $N$ 个patch特征， $X=[x_1,...,x_N]$ ，token-mixing会将 $XW$ ，其中 $W\in \mathbb{R}^{N\times M}$ 作为token-mixing MLP的学习权矩阵。很容易观察到 $XW$ 的每一列，即token-mixing MLP的输出，都是patch特征(输入x中的列)的加权求和。求和中的权重可以看作是Transformer中需要注意的问题。但Transformer中的自关注是依赖于数据的，而token-mixing MLP中的总和权重与输入无关。加权求和在某种程度上更类似于深度卷积[2,14,15]。但深度卷积只有一个局部接收域。相反，token-mixing MLP有一个全局接收字段。此外，深度卷积核在不同的空间位置之间共享，而token-mixing MLP的求和权值在不同的空间位置是不同的。不受局部接收域和空间不可知约束的限制，token-mixing MLP比深度卷积更灵活，具有更强的拟合能力。但这种免于打破链条的自由也伴随着过度拟合的风险。（token-mixing这种全局性质需要很多数据做训练，香则容易过拟合）为了避免token-mixing MLP中存在过拟合问题，需要提供大量的训练样本。说明在超大数据集JFT-300M上进行预处理后，MLP-Mixer与ViT之间的识别精度差距缩短。

为了缓解MLP-Mixer在只有中等规模训练数据可用时的过拟合问题，我们提出了一个空间移位MLP (S2-MLP)体系结构，一个概念简单的体系结构，只包含通道混合MLP。为了实现空间位置间的通信，我们采用了空间移位操作，该操作无参数且计算效率高。同时，空间的转移是不可知的，同时保持了局部的接待场。图1说明了提议的S2-MLP的体系结构。它在 $w\times h$ 的图像上裁剪成patch。对于每个patch，它通过全连通层得到embedding。这个 $w\times h$ 的patch进一步通过S2-MLP块堆栈进行特征提取。每个S2-MLP区块包含三个完全连接的层。在每个S2-MLP块中使用的全连接层与MLP-Mixer中使用的通道混合MLP具有相同的功能。但是我们的S2-MLP不需要token-mixingMLP。相反，不同空间位置之间的通信是通过提出的空间移位模块实现的。它是无参数的，只是简单地将通道从一个patch转移到其毗邻的patch。尽管空间移位模块只支持相邻斑块之间的通信，但通过堆叠一系列S2-MLP块，可以实现远程通信。

图1 提出的空间移位多层感知(S2-MLP)模型的体系结构。从图像中裁剪出的非重叠patch是模型的输入。它们通过一个S2-MLP块堆栈，这些块通过全局平均池化进一步聚合成一个单一的特征向量。然后，将特征向量送入全连接层进行标签预测。一个S2-MLP块包含四个全连接层，两个GELU层[11]，两层归一化[1]，两个跳跃连接[10]和一个空间移位模块。所提出的空间移位模块将信道分组为若干组。然后它将不同的频道组向不同的方向移动。

所提出的S2-MLP在架构上看似简单。与ImageNet1K数据集上的MLP-Mixer相比，它在参数和FLOPs规模相当的情况下获得了相当高的识别精度。与此同时，它在ImageNet1K数据集上以相当简单的结构、更少的参数和FLOPs实现了与ViT相当的识别精度。

S2-MLP 的改动主要集中在空间 MLP 的位置，由原来的Spatial-MLP(Linear->GeLU->Linear)转变为Spatial-Shifted Channel-MLP(Linear->GeLU->Spatial-Shift->Lienar)

2. Related Work

Transformer-based视觉模型。Visual Transformer (ViT)[6]是构建纯粹基于Transformer的视觉骨干的第一个工作。通过在超大数据集JFT-300M[24]上的训练，与人工视觉主干、卷积神经网络相比，取得了较好的效果。DeiT[29]采用高级的训练和增强策略，仅在ImageNet-1K上进行训练时，表现优异。最近，一些工作从多个角度进一步提高了visual Transformer的性能。例如，PVT[32]使用渐进式收缩金字塔来减少对大型特征图的计算。T2T[34]逐步标记图像，以模拟图像的局部结构信息。TNT[9]在外部级Transformer中构建另一个Transformer，以对本地patch进行建模。CPVT[4]提出了一种条件位置编码方法，可以有效地对patch的空间位置进行编码。Visual Longformer[35]采用全球代币提升效率。PiT[12]研究了空间维度转换，并集成了自注意块之间的池化层。Swin - transformer[20]采用高灵活性的分层拱形架构，可以在不同的比例下建模图像。Twins[3]采用了一个由局部分组自注意和全局下采样注意组成的层次结构。CaiT[30]构建并优化了用于图像分类的更深层次的transformer网络。

MLP-based视觉模型。MLP-Mixer[27]提出了一个概念和技术上简单的架构，完全基于MLP层。为了建模空间位置之间的通信，它提出了一个token-mixing MLP。尽管MLP-Mixer在大型数据集JFT-300M上训练时取得了令人满意的结果，但在中等规模的数据集(包括ImageNet-1K和ImageNet-21K)上训练时，它不如它的对手visual Transformer。FF[22]采用了类似的架构，但继承了ViT的全局[CLS]令牌和位置嵌入。Res-MLP[28]还设计了一个纯MLP架构。提出了一种仿射变换层，可将大量的MLP块进行叠加。使用比MLP-Mixer更深层次的架构，Res-MLP实现了比MLP-Mixer更高的精度和类似ViT的识别精度。gMLP[19]设计了一个门控操作，以增强空间位置之间的通信，并实现了与DeiT相当的识别精度。EA[8]通过从训练数据中学习的外部记忆，将自我注意模块替换为外部注意模块。它是由两个线性层的级联实现的。Container[7]提出了一个广义上下文聚合构建块，它将静态关联矩阵组合为token-mixing，将动态关联矩阵组合为可视变形器。

3. Method

在本节中，我们将描述空间移位MLP (S2-MLP)。

3.1. Preliminary

Layer Normalization (LN)[1]被广泛应用于使用Transformer和BERT架构的模型。给定一个 $c$ 维的向量 $x= [x_1,...,x_c]$ ，层归一化计算出均值：

$\mu =\frac{1}{c}\sum_{i=1}^{c}x_i$

和标准差

$\sigma =\sqrt{\frac{1}{c}\sum_{i=1}^{c}(x_i-\mu )^{2}}$

它通过

$\bar{x}_i=\frac{x_i-\mu }{\sigma},\forall i \in [1,c]$

标准化x中的每一项。

Gaussian Error Linear Units (GELU)[11]是Transformer和BERT模型中广泛使用的激活函数。定义为 $GELU(x) =x \Phi (x)$ ，其中 $\Phi (x)$ 标准高斯累积分布函数。是定义为：

$\Phi (x)=\frac{1}{2}[1+erf (\frac{x}{\sqrt{2}})]$

MLP-Mixer[27]由N个相同大小和结构的基本块堆叠而成。每个基本块由两种类型的MLP层组成:channel-mixing MLP和token-mixing MLP。我们设表示一个图片由n个patch组成，即 $P=[p_1,...,p_n]\in \mathbb{R}^{c\times n}$ 。通道混合MLP沿通道维度投影：

$\hat{P}=P+W_2GELU(W_1LN(P))$ ……（1）

其中 $W_1\in \mathbb{R}^{\hat{c}\times c}$ ， $W_2\in \mathbb{R}^{c\times \hat{c}}$ 。同时，token-mixing MLP沿着空间维度投射通道混合patch特征：

$\bar{P}=\hat{P}+GELU(LN(\hat{P})W_3)W_4$ ……（2）

其中 $W_3\in \mathbb{R}^{N\times \bar{N}}$ 和 $W_4\in \mathbb{R}^{ \bar{N}\times N}$ 。

3.2. Spatial-Shift MLP Architecture

如图1所示，我们的空间移位MLP主干由一个路径全连接层（path-wise fully-connected layer）、N个S2-MLP块和一个用于分类的全连接层组成。由于我们已经引入的全连接层来分类是众所周知的，所以我们只引入了patch-wise全连接层和提出的空间移位块。

Patch-wise fully-connected layer。用 $I\in \mathbb{R}^{W\times H\times 3}$ 表示图像。将图片均匀分割成 $w\times h$ 的patch, $P=\left \{ P_i \right \}_{i=1}^{wh},where P_i \in \mathbb{R}^{p\times p\times 3},w=\frac{W}{p},h=\frac{H}{p}$ 。对于每一个patch $P_i$ ，我们将其展开为一个向量 $p_i\in \mathbb{R}^{3p^2}$ ，并通过一个完全连通的层，再进行一层归一化将其投影到一个嵌入向量 $e_i$ 中， $e_i=LN(W_0 p_i+b_0)$ ，其中 $W_0\in \mathbb{R}^{c\times 3p^2},b_0\in \mathbb{R}^c$ 是全连通层的参数，LN(·)表示上面我们介绍的层归一化。

S2-MLP block.。我们的架构堆栈N个具有相同的大小和结构的S2-MLP。每个空间移位块包含四层全连接层、两层归一化层、两层GELU层、两层跳接和提出的空间移位模块。值得注意的是，我们S2-MLP中使用的所有全连接层仅用于混合通道（channel-mixing）。我们在MLPMixer中不使用token-mixing MLP。由于全连接层是众所周知的，而且我们已经在上面介绍了层归一化和GELU，所以我们在这里只关注提出的空间移位模块。我们用 $\tau\in \mathbb{R}^{w\times h \times c }$ 来表示在我们的spatial-shift模块输入的feature map，其中w表示宽度，h表示高度，c表示通道数量。空间移位操作可分解为两个步骤:1)将信道分成若干组，2)将每组信道向不同方向移位。

Group。我们沿通道维均匀分割 $\tau$ ，得到更薄的张量 $\left \{ \tau_\tau \right \}_{\tau =1}^{g}$ ，其中 $\tau_\tau \in \mathbb{R}^{w\times h\times c/g}$ 。值得注意的是，组的数量g取决于第二步中移动方向的设计。例如，默认情况下，我们只沿着四个方向（四个角的方向）移动，因此在这个配置中设置为4。

Spatial-shift operation。我们在不同的方向上转换不同的通道组。对于第一组通道， $\tau _1$ ，我们将它沿宽维度移动+1。同时，我们将第二组通道 $\tau_2$ 沿宽维度移动−1。类似地，在 $\tau_3$ 中，我们将它沿高度维度移动+1，在 $\tau_4$ 中沿高度维度移动−1。我们用式子(3)中空间移位运算的表达式，并在Algorithm1中演示了伪代码。

（四个通道组分别向四个角偏移）

从式子（3）中可以看出，第一组通道， $\tau _1$ 的操作是沿着宽向右移动+1，以此类推。

经过空间移动后，每个patch从相邻的patch中吸收视觉内容。空间移位操作是无参数的，使得不同空间位置之间的通信是可行的。上述的空间移位方式是最简单、最直接的移位方法之一。我们也评估其他的空间转移方式。令人惊讶的是，上述简单的方式相比其他方式取得了优异的性能。使用空间移位操作，我们不再需要像MLP-Mixer那样的token-mixer。我们只需要channel-mixer来沿通道维度投射局部特征。注意，单个块的空间移位操作只能获得相邻块的视觉内容，不能访问图像中所有块的视觉内容。但我们堆叠N个S2-MLP块，全局视觉内容将逐渐扩散到每个patch。

3.3. Relations with depthwise convolution

Depthwise convolution。给定一个定义为tensor $T\in \mathbb{R}^{w\times h\times c}$ 的特征映射，depthwise convolution[2,14,15]在 $\tau [:,:,i]\in \mathbb{R}^{w\times h},i\in [i,c]$ 中的每个二维切片上使用一个二维卷积核 $K_i$ 。深度卷积计算成本低，被广泛应用于高效神经网络的快速推理。

Relations.。提出的空间移动是受到TSM[18]中提出的时间移动的启发。它最初用于建模相邻帧之间的时间关系，使信道沿时间维度移动。在这项工作中，我们将其扩展到二维空间场景。事实上，空间移位操作等于与固定的、特定于组的核权进行深度卷积。设表示一组深度卷积核 $\kappa =\left \{ K_1,...,K_C \right \}$ 。如果我们设置：

则基于 $\kappa$ 核函数组的深度卷积等价于我们的空间移位操作。

也就是说，我们的空间移位操作是与上面定义的固定权值的深度卷积的一种变体。同时，空间移位操作在每组信道中共享核权。正如我们在介绍中提到的，MLP-Mmixer中的token-mixing MLP是一个全局接收和空间特定的深度卷积。同时，与我们的空间移位和vanilla深度卷积相比，token-mixing的和权在特定的空间位置跨通道共享。而vanilla深度卷积对于不同的信道学习不同的卷积核，我们的空间移位操作共享组内权值，对不同的组采用不同的权值。换句话说，MLP-Mixer中的空间移位和token-mixing MLP都是深度卷积的变体。我们在表1中总结了它们之间的关系和差异。通过观察空间移位操作和深度卷积之间的联系，我们认为我们的S2-MLP，在某种程度上，叫好酒(MLP)和卖醋(卷积)（？？？）。我们鼓励研究人员尝试不同设置的深度卷积来构建新的基于MLP的架构。

表1 token-mixing(TM)、空间移位(S2)和深度卷积(DC)之间的关系。

3.4. Complexity Analysis

3.5. Implementation（实现）

我们设置裁剪后的patch大小(p×p) 为16×16。我们将输入图像重塑为224×224的大小。因此，patch的数量为 $M=(224/16)^{2}=196$ 。我们设膨胀率（expansion ratio）为4。我们尝试两种设置:1)宽设置和2)深度设置。宽设置遵循MLP-Mixer[27]的基本模型。宽设置设置S2-MLP块(N)的数量为12，隐藏的大小为768。注意，MLP-Mixer还实现了大模型和大模型。然而，我们有限的计算资源无法承担在ImageNet-1K数据集上研究大型模型的昂贵成本。深度设置遵循Res-MLP-36[28]。深度设置设置S2-MLP块(N)的数量为36，隐藏的大小为cas384。我们总结了表2中两种设置的超参数、参数数量和FLOPs。

表2 超参数，参数的数量和FLOPs。参考MLP-Mixer[27]，参数的数量去除进行分类全连接层的权重

4. Experiments

5. Conclusion

在本文中，我们提出了一个空间转移MLP (S2-MLP)架构。采用纯MLP结构，无卷积和自注意。为了实现空间位置间的通信，我们采用了一种简单、无参数、高效的空间移位操作。在ImageNet-1K数据集上，S2-MLP在参数和FLOPs数量相当的情况下，实现了比MLP-Mixer和ResMLP更高的识别精度。同时，与ViT相比，我们的S2-MLP采用了更简单的架构，更少的参数和FLOPs。此外，我们还讨论了空间移位运算、MLP-Mixer中token-mixing MLP与深度卷积之间的关系。我们发现，token-mixing MLP和提出的空间移位操作都是深度卷积的变体。我们希望这些结果和讨论可以激发进一步的研究，在不久的将来发现更简单和更有效的视觉架构。