MLP-Mixer: An all-MLP Architecture for Vision

摘要

        卷积神经网络(cnn)是计算机视觉的首选模型。最近，以注意力为基础的网络，如Vision Transformer也开始流行起来。在本文中，我们证明卷积和注意力对于良好的性能都是足够的，但它们都不是必要的。我们提出了MLP-Mixer，一个专门基于多层感知器(MLPs)的架构。MLP-Mixer包含两种类型的层:一种是独立应用于图像斑块的MLPs(即“混合”每个位置的特征)，另一种是跨斑块应用的mlp(即“混合”空间信息)。在大数据集上或使用现代正则化方案进行训练时，MLP-Mixer在图像分类基准测试中获得了有竞争力的分数，其预训练和推理成本可与最先进的模型相媲美。我们希望这些结果能激发进一步的研究，超越已建立的CNN和Transformer的领域。

1  介绍

        正如计算机视觉的历史所证明的那样，更大的数据集的可用性加上不断增加的计算能力常常会导致范式的转变。虽然卷积神经网络(CNNs)已经成为计算机视觉的事实上的标准，但最近基于自我注意层的替代方案Vision Transformer (ViT)获得了最先进的性能。ViT延续了去除模型中手工制作的视觉特征和归纳偏差的长期趋势，并进一步依赖于从原始数据中学习。

        我们提出了MLP-Mixer网络结构(简称“Mixer”)，这是一种有竞争力但概念和技术上都很简单的替代方案，不使用卷积或自我注意。相反，Mixer的架构完全基于多层感知器(MLPs)，这些感知器在空间位置或特征通道上重复应用。Mixer只依赖于基本的矩阵乘法例程，数据布局的更改(重塑和转置)，和标量非线性。

        图1描述了Mixer的宏观结构。它接受线性投影的图像块序列(也称为token)形状为“patches×channels”作为输入，并保持这个维度。Mixer使用两种类型的MLP层:通道混合MLP和token混合MLP（channel-mixing MLPs and token-mixing MLPs）。通道混合mlp允许不同信道之间的通信;它们独立地对每个token进行操作，并将表中的各个行作为输入。token混合mlp允许不同空间位置(token)之间的通信;他们独立操作每个通道，并将表的各个列作为输入。这两种类型的层是交错的，以支持两个输入维度的交互。

图1:MLP-Mixer由每patch线性嵌入，Mixer层，和分类器头组成。Mixer层包含一个token混合MLP和一个信道混合MLP，每个MLP由两个完全连接的层和一个GELU非线性组成。其他组件包括:跳过连接，dropout，通道上的层规范，和线性分类器头。

         在极端情况下，我们的架构可以被视为一个非常特殊的CNN，它使用1×1卷积来进行信道混合，并使用一个完全接受域和参数共享的单信道深度卷积来进行token混合。然而，事实并非如此，因为典型的CNN并非Mixer的特殊情况。此外，在MLPs中，卷积比简单的矩阵乘法更复杂，因为它需要额外的代价来减少矩阵乘法和/或专门的实现。

        尽管因为它的简单，Mixer取得了竞争性的结果。在大数据集(即100M的图像)上进行预先训练时，它在精确度/成本权衡方面接近CNN和Transformer 此前声称的最先进的性能。这包括87.94%的ILSVRC2012“ImageNet”的top-1验证精度。当预先训练更适度的尺度数据(即~ 1-10M图像)，再加上现代正则化技术，Mixer也取得了强大的性能。然而，与ViT类似，它略低于专门的CNN架构。

2  Mixer Architecture

        MLP-Mixer主要包括三部分：Per-patch Fully-connected、Mixer Layer、分类器。

2.1 Per-patch Fully-connected

        MLP-Mixer通过Per-patch Fully-connected将输入图像转化为2D的Table，方便在后面进行局部区域间的信息融合。

        具体来说，MLP-Mixer将输入图像相邻无重叠地划分为S个Patch，每个Patch通过MLP映射为一维特征向量，其中一维向量长度为C，最后将每个Patch得到的特征向量组合得到大小为S*C的2D Table。需要注意的时，每个Patch使用的映射矩阵相同，即使用的MLP参数相同。

        实际上，Per-patch Fully-connected实现了三维图像（W，H，C）的向量空间到二维（S，C）的向量空间的映射。

        例如，假设输入图像大小为240*240*3，模型选取的Patch为16*16，那么一张图片可以划分为（240*240）/（16*16）= 225个Patch。结合图片的通道数，每个Patch包含了16*16*3 = 768个值，把这768个值做Flatten（拉平）作为MLP的输入，其中MLP的输出层神经元个数为128。这样，每个Patch就可以得到长度的128的特征向量，组合得到225*128的Table。MLP-Mixer中Patch大小和MLP输出单元个数为超参数。

2.2  Mixer Layer

        观察Per-patch Fully-connected得到的Table会发现，（如上组合得到225*128的Table）Table的行代表了同一空间位置在不同通道上的信息，列代表了不同空间位置在同一通道上的信息。换句话说，对Table的每一行进行操作可以实现通道域的信息融合，对Table的每一列进行操作可以实现空间域的信息融合。

        从Layer Norm层出来的为Patches*Channel(即为table)，每个patch即上述对同一位置的所有通道（如16*16*3）进行展开，通过T转为Channel*Patches，token-mixing MLPs（MLP1）对table的列进行映射，所以对Channel*Patches的行进行映射，对不同空间位置在同一通道上的信息进行操作，实现空间域的信息融合。channel-mixing MLPs（MLP2）对table的行进行映射，对同一空间位置在不同通道上的信息进行映射，实现通道域的信息融合。

        在传统CNN中，可以通过1*1 Conv来实现通道域的信息融合，如果使用更大一点的卷积核，可以同时实现空间域和通道域的信息融合。在Transformer中，通过Self-Attention实现空间域的信息融合，通过MLP同时实现空间域和通道域的信息融合。

        而在MLP-Mixer中，通过Mixer Layer使用MLP先后对列、行进行映射，实现空间域和通道域的信息融合。与传统卷积不同的是，Mixer Layer将空间域和通道域分开操作。

        根据上述内容，MLP-Mixer在Mixer Layer中使用分别使用token-mixing MLPs（图中MLP1）和channel-mixing MLPs（图中MLP2）对Table的列和行进行映射，与Per-patch Fully-connected相同，MLP1和MLP2在不同列、行中的映射过程中共享权重。除此之外，Mixer Layer还加入了LN和跳接来提高模型性能。

2.3  整体结构