MLP-Mixer: 并肩卷积与自注意，多层感知机的神奇魔法

0.引言

MLP-Mixer模型是谷歌AI团队于2021年初发表的文章，题为MLP-Mixer: An all-MLP Architecture for Vision。在计算机视觉领域的历史上，卷积神经网络一直是首选的模型。然而最近，注意力机制网络（例如Vision Transformer）也变得非常流行。在MLP-Mixer工作中，研究人员表明，尽管基于卷积或注意力算法的模型都足以获得良好的性能，但它们都不是必需的，纯MLP+非线性激活函数+层归一化也能取得不错的性能，其预训练和推理成本可与最新模型相媲美。

那么，为什么要用全连接层？它有什么优势呢？

全连接层具有更低的归纳偏置（Inductive Bias）。归纳偏置指的是机器学习算法在学习过程中对某些类型的假设和解决方案的偏好。换句话说，归纳偏置就是算法在面对不确定性时，根据其内在设计选择某种解决方案的倾向。例如，卷积神经网络的归纳偏置在于卷积操作的局部相关性和平移不变性。而时序网络的归纳偏置则在于时间维度上的连续性和局部性。实际上，注意力机制网络已经延续了旨在摒弃神经网络中手工视觉特征和归纳偏置的趋势，让模型完全依赖原始数据进行学习。MLP则更进一步推动了这一发展。

论文名称：MLP-Mixer: An all-MLP Architecture for Vision

论文下载：https://arxiv.org/abs/2105.01601

1.Per-patch全连接层

相较于卷积神经网络层，全连接神经网络层无法获取输入数据局部区域间的信息。为了解决这个问题，MLP-Mixer增加了Per-patch全连接层（Per-patch Fully-connected Layer，FC），对该问题进行了更有效的处理。具体来说，在处理图像时，MLP-Mixer首先将输入图像划分为多个固定大小的小块（Patches）。接下来，每个Patch被展平，并通过一个Per-patch全连接层，以便将每个Patch转换为一个具有固定维度的向量。Per-patch全连接层的目的是学习一个从图像Patch到嵌入空间的映射，从而捕捉每个Patch内的局部信息。这个全连接层在MLP-Mixer的整个训练过程中都会进行更新和优化，以便更好地捕捉和表示图像Patch的信息。

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

总之，Per-patch全连接层实现了将输入图像的维度从(W,H,3)映射到(S,C)的特征向量维度。

2.Mixer-Layer代替自注意力机制

在上一节中，维度为16×16×3的Patch被展平成长度为768的向量，之后经过MLP的映射作为二维矩阵的行，二维矩阵的行实际上代表了同一空间位置（某个Patch）在不同特征通道上的信息，而列代表了不同空间位置（不同Patch）在同一特征通道上的信息。换句话说，对矩阵的每一行进行操作可以实现特征通道域的信息融合，对矩阵的每一列进行操作可以实现空间域的信息融合。

在传统卷积神经网络中，可以通过1×1的卷积来实现特征通道域的信息融合，如果使用尺寸大于1的卷积核，则可以同时实现空间域和特征通道域的信息融合。

在Transformer中，通过自注意力机制实现空间域的信息融合，通过MLP同时实现特征通道域的信息融合。

而在MLP-Mixer模型中，通过Mixer-Layer，使用MLP先后对输入数据的列和行进行映射，实现空间域和特征通道域的信息融合。与传统基于卷积神经网络的模型不同，由于Mixer-Layer对行、列的操作是存在先后顺序的，这意味着Mixer-Layer将空间域和特征通道域分开操作，这种思想与Xception和MobileNet中的深度可分离卷积相似，也类似于Transformer中对空间信息和特征通道信息的分开处理。

Mixer架构采用两种不同类型的MLP层：Token-mixing MLP（即图1中的MLP1）和Channel-mixing MLP（即图1中的MLP2），如图1所示。Token-mixing MLP允许不同空间位置之间进行通信，可以融合空间信息；Channel-mixing MLP允许不同通道之间进行通信，可以融合通道信息。Token-mixing MLP模块作用在每个Patch的列上，即先对Patches部分进行转置，并且所有列参数共享MLP1，得到的输出再重新转置一下。Channel-mixing MLP模块则作用在每个Patch的行上，所有行参数共享MLP2。这两种类型的层交替执行，可以促进两个维度间的信息交互。

最后，Mixer-Layer还加入了残差连接和层归一化来提高模型性能。MLP1和MLP2都采用了相同的结构，如图2所示。

3.MLP-Mixer模型结构

MLP-Mixer的整体思路为：先将输入图片拆分成多个Patch（每个Patch之间不重叠），通过Per-patch全连接层的操作将每个Patch转换成特征嵌入，然后送入N个Mixer-Layer。最后，输出结果经过全局平均池化层后送入全连接层进行分类预测。

MLP-Mixer的第一步实际上与注意力机制网络是一致的。Mixer-Layer替换了Transformer Block，最后的输出直接连接到全连接层，无须使用Class Token。假设有一个大小为224×224×3的输入图像，首先将其划分成大小为32×32×3的多个Patch，共得到(224/32)×(224/32)=49个Patch，每个patch是32×32×3=3072维的。通过一个Per-patch全连接层进行降维，如降至512维，就可以得到49个维度为512的Token。然后将它们送入Mixer-Layer。

Mixer架构采用两种不同类型的MLP层：Token-mixing MLP和Channel-mixing MLP。Token-mixing MLP指的是跨位置操作（Cross-location Operation），即针对49个512维的Token，将每个Token内部进行自我融合，实现49维到49维的映射，即“混合”空间信息。Channel-mixing MLP则指的是位置先验操作（Pre-location Operation），即针对49个512维的Token，将每一维进行融合，实现512维到512维的映射，即“混合”每个位置的特征。为了简化实现，实际上仅需对矩阵进行转置即可。这两种类型的层交替执行，可以促进两个维度间的信息交互。

单个MLP是由两个全连接层和一个GELU激活函数组成的。此外，Mixer还使用了残差连接和层归一化。以上操作的公式如下：

\(\begin{aligned}U_{*,i}&=X_{*,i}+W_2\sigma(W_1\text{LayerNorm}(X)_{*,i}),\mathrm{for~}i=1,\cdots,C\\Y_{j,*}&=U_{j,*}+W_4\sigma(W_3\text{LayerNorm}(U)_{j,*}),\mathrm{for~}j=1,\cdots,S\end{aligned}\)

为什么要采用Token-mixing MLP和Channel-mixing MLP两种操作呢？首先，在图像中的信息融合分为空间维度的融合和特征通道维度的信息融合。对比卷积操作，卷积神经网络的卷积核大小为K×K×C，其实是同时考虑了这两个：当K=1时，是逐点卷积，考虑的融合特征通道维度信息；当C=1时，是逐层卷积，其实就是考虑了融合空间维度信息。有的卷积神经网络变体，例如MobileNet、EfficientNet等，为了计算复杂度降低采用通道可分离卷积，先做逐点卷积，再做逐层卷积，也得到了不错的效果，说明两种卷积操作可以不一起做。所以MLP-Mixer也干脆将这两个正交分开来处理。

MLP-Mixer由多个Mixer-Layer拼接而成，其中每一层采用相同尺寸的输入（Token维度保持不变），这种“各向同性”设计类似于Transformer和循环神经网络中的定宽结构，这与卷积神经网络中的金字塔结构（越深的层，特征图的分辨率越低，特征的通道数越多）有所不同。

至于最后的全局池化和全连接层，则是相对常规的操作。按照通道对每个Token进行池化，最终得到通道维度的向量，然后通过一个全连接层输出目标分类分数即可。MLP-Mixer模型的结构图如图3所示。

最终MLP-Mixer的网络配置见表1。

4.小结

对于模型性能方面，文中使用了ImageNet、ImageNet21K和JFT-300M三个数据集进行MLP-Mixer的预训练，并在ImageNet数据集上进行微调。最高的Top-1精度可以达到88.55%。尽管略低于ViT模型，但其性能仍然很出色。

以下是关于MLP-Mixer的一些反思：

（1）卷积神经网络和Transformer已经在目标检测和分割方面展现出出色的性能。然而MLP-Mixer能否在分割、识别等方面提供更大的帮助尚不得而知。MLP-Mixer的特征在分类任务中表现有力，但在其他下游任务的有效性仍需验证。此外，由于MLP-Mixer的计算量较大，因此它在处理高分辨率图像输入方面可能并不理想。最后，尽管MLP思路看似具有创新性，但残差结构和层归一化的设计同样具有重要作用。如果没有它们，MLP-Mixer是否仍能正常工作呢？

（2）对于MLP-Mixer的一些展望：MLP-Mixer能否像Swin Transformer一样增加Patch内信息的交流呢？从混合模型的角度来看，CNN+MLP-Mixer的结构是否比CNN+Transformer的结构更好呢？实际上，MLP-Mixer有许多微小的修改和改进，例如RepMLP、ResMLP、gMLP、CycleMLP等，这些值得在后续学习中探索。

总的来说，从卷积神经网络到ViT，再到MLP-Mixer，研究者们在减少算法的归纳偏置的道路上一路狂奔。甚至可以用全连接层来直接替换自注意力机制。然后，研究者们绕了一个大圈又回到了十几年前的MLP。实际上，最初对于归纳偏置的设计是希望算法少走弯路，以人的经验为算法提供尽可能多的引导，帮助深度学习模型快速地在计算机视觉、自然语言处理等领域落地。而现在，我们拥有更多的数据，计算资源上限也提高了，预算也更充裕，于是说：“不需要归纳偏置，让算法自己去学习吧！”

这样的方式是否更好呢？目前还没有一个统一的说法。据笔者观察，机器自己学到的规律往往比人为赋予的归纳偏置更合理，这就像自己悟出的知识通常比别人传授的记忆更深刻一样。前提是你有大量的数据可供机器来学习，并且有足够的资金购买计算资源。因此，关键在于根据场景和问题选择合适的算法，算法没有好坏，时势造英雄。