MLP-Mixer：一种无卷积、无注意力、纯MLP构成的视觉架构

摘要

目前计算机视觉领域多采用卷积神经网络（CNN）和基于自注意力的网络（如ViT）。近日，谷歌的研究团队（原 ViT 团队）提出了一种舍弃卷积和自注意力且完全使用多层感知机（MLP）的视觉网络架构，在设计上非常简单，并且在 ImageNet 数据集上实现了媲美 CNN 和 ViT 的性能表现。
论文地址：MLP-Mixer: An all-MLP Architecture for Vision

文章结构

第一部分介绍了MLP-Mixer的宏观结构，介绍了
Mixer采用的两种类型的MLP层。
第二部分介绍了Mixer的详细结构以及原理。
第三部分介绍了所采用的实验和实验结果。

方法

MLP-Mixer是一种区别于CNN和Transformer的新架构，它无需卷积与自注意力。相反，MLP-Mixer仅仅依赖于在空域或者特征通道上重复实施的多层感知器；Mixer仅依赖于基础矩阵乘操作、数据排布变换(比如reshape、transposition)以及非线性层。

图1 Mixer的宏观建构示意图

Mixer的大致结构如图1，
它以一系列图像块的线性投影(其形状为patches x channels)作为输入。Mixer采用了两种类型的MLP层(注:这两种类型的层交替执行以促进两个维度间的信息交互)：

• channel-mixingMLP：用于不同通道前通讯，每个token独立处理，即采用每一行作为输入
• token-mixingMLP：用于不同空域位置通讯，每个通道图例处理，即采用每一列作为输入。

在极端情况下，MLP-Mixer架构可以看作一个特殊的CNN，它使用1×1通道混合的卷积，全感受域的单通道深度卷积以及token混合的参数共享。

Mixer 混合器架构

一般来讲，当今深度视觉体系结构采用三种方式进行特征混合：
（i）在给定的空间位置；
（ii）不同的空间位置之间；
（iii）将上述两种方式组合。
在CNNs中，（ii）是采用N× N进行卷积和池化，其中N>1；（i）采用1×1卷积；较大的核则同时执行（i）和（ii）。通常更深层次的神经元有更大的感受野。
在Transformer和其他注意力架构中，自注意力层允许同时执行（i）和（ii），而MLP只执行（i）。Mixer架构背后的思想是：通过MLP实现每个通道混合操作（i）和 token混合操作（ii）的显著分离。
在图1体系架构中，Mixer将序列长度为S的非重叠的图像块作为输入，每个图像块都投影到所需的隐层维度C，并产生一个二维实值输入X∈ RS×C。如果原始图像的分辨率为（H x W），每个图像块的分辨率为（P x P），那么图像块的数量则为S=HW/P2。所有的块都采用相同的投影矩阵进行线性投影。
Mixer由等尺寸的多层组成，每层有两个MLP块。第一个是token mixing MLP块：它作用于X的列，从RS映射到R S，可在所有列中共享。第二个是Channel-mixing MLP块：它作用于X的行，从Rc映射到 R C，可在所有行中共享。每个MLP块包含两个全连接层和一个独立于输入的非线性层。其基本方程如图2 ：

图2

Ds Dc分别代表token-mixing与channel-mixing MLP中隐层宽度。由于Ds的选择独立于输入图像块的数量，因此，网络的计算复杂度与输入块的数量成线性关系；此外，Dc独立于块尺寸，整体计算量与图像的像素数成线性关系，这类似于CNN。
相同的通道混合MLP（或令牌混合MLP）应用于X的每一行和列，在每一层内绑定通道混合MLP的参数都是一种自然选择，它提供了位置不变性，这是卷积的一个显著特征。
不过，跨通道绑定参数的情况在CNN中并不常见。例如CNN中可分离卷积，将不同的卷积核独立应用于每个通道。而Mixer中的token 混合MLP可以对所有通道共享相同的核（即获得完全感受野）。通常来讲，当增加隐层维数C或序列长度S时，这种参数绑定可以避免体系架构增长过快，并且节省内存。但这种绑定机制并没有影响性能。
Mixer中的每个层（除了初始块投影层）接收相同大小的输入。这种“各向同性”设计最类似于使用固定宽度的Transformer和RNN。这与大多数CNN不同，CNN具有金字塔结构：越深的层具有更低的分辨率，更多的通道。
此外，与ViTs不同，Mixer不使用位置嵌入，因为token混合mlp对输入token的顺序敏感，因此可以学习表示位置。最后，Mixer使用一个标准的分类head和一个线性分类器。

实验

主要是对三个问题进行了分析：

• 在下游任务上的精度
• 总计预训练计算量，这对于在上游数据上从头开始训练模型非常重要
• 推理耗时，这对于实际应用非常重要

文章的目的不在于达成SOTA结果，而在于表明：简简单单的MLP模型即可取得与当前最佳CNN、注意力模型相当的性能。

图3

图3给出了最大Mixer模型与SOTA模型在ImageNet、ReaL数据集上的性能对比，从中可以看到：

• 当在ImageNet-21k+额外正则技术预训练后，Mixer在ImageNet数据集取得非常强的性能：84.15%top1，比其他模型稍弱。此时，类似于ViT，正则技术很有必要，有助于Mixer避免过拟合。
• 当提升上游数据集尺寸后，Mixer的性能可以进一步提升。具体来说，Mixer-H/14取得了87.94%的top1精度，比BiT-ResNet152x4高0.5%，比ViT-H/14低0.5%。值得一提的是，Mixer-H/14的推理速度要比ViT-H/14快2.5倍，比BiT快2倍。
  

图4

图4对比了不同模型大小、不同分辨率输入时模型的性能对比，从中可以看到：

• 当在ImageNet上从头开始训练时，Mixer-B/16取得了一个比较合理的精度：76.44%，。这要比ViT-B/16低3%。也就是说，在ImageNet上从头开始训练时，Mixer-B/16出现了过拟合问题。
• 随预训练数据集变大，Mixer的性能稳步提升。值得一提的是，JFT-300M数据上预训练、224尺寸微调的Mixer-H/14取得了86.32%的精度，仅以0.3%弱于ViT-H/14，但推理速度快2.2倍。