新坑！谷歌提出MLP-Mixer：一种无卷积、无注意力、纯MLP构成的视觉架构

点击下方卡片，关注“AIWalker”公众号

重磅干货，第一时间送达

paper: https://arxiv.org/abs/2105.01601

code:https://github.com/google-research/vision_transformer(暂未开源)

本文是谷歌大脑的研究员(原ViT团队)在网络架构设计方面挖的新坑：MLP-Mixer。无需卷积、注意力机制，MLP-Mixer仅需MLP即可达到与CNN、Transformer相媲美的性能。比如，在JFT-300M数据集预训练+ImageNet微调后，所提Mixer-H/14取得87.94%的top1精度。尽管所提方法性能并未达到最优，但本文的目的并不在于达成SOTA结果，而在于表明：简简单单的MLP模型即可取得与当前最佳CNN、注意力模型相当的性能。

Abstract

CNN已成为计算机视觉的首选方案，近来兴起的基于注意力的网络(比如Vision Transformer)也逐渐变的流行。这篇文章表明：尽管卷积与注意力有助于取得比较的好的性能，但其并非必要。

本文提出了一种新颖的架构：MLP-Mixer，它仅仅依赖于多层感知器(multi-layer perceptrons, MLP)。MLP-Mixer包含两种类型的层：(1) 独立作用于图像块的MLP，即用于混合每个位置的特性；(2) 跨图像块的MLP，即用于混合空域信息。当在大数据集上训练，或者采用先进正则技术训练后，MLP-Mixer在图像分类基准数据集上取得了极具竞争力的性能。我们希望该工作能激发出超越CNN与Transformer之外的进一步研究。

Introduction

尽管CNN已成为计算机视觉领域的首选方案，近来兴起的Vision Transformer(ViT)打开了另一扇门并取得了优异的性能。本文提出了一种区别于CNN和Transformer的新架构MLP-Mixer，一种极具竞争力且概念与技术简单的新方案，它无需卷积与自注意力。相反，MLP-Mixer仅仅依赖于在空域或者特征通道上重复实施的多层感知器；Mixer仅依赖于基础矩阵乘操作、数据排布变换(比如reshape、transposition)以及非线性层。

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

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

在极端情况下，本文所提架构可视作一种特殊CNN，它采用卷积进行channel mixing，全感受野、参数共享的的单通道深度卷积进行token mixing。

尽管如此简单，Mixer取得极具竞争力的结果，当在大数据(如100M数据量)上预训练时，它可以取得接近SOTA的性能：87.94%top1精度。当在适当尺度数据+先进正则技术进行训练时，所提方法同样取得极强性能。

Mixer Architecture

当前深度视觉架构采用以下三种方式进行特征混合：(i) 在给定空域位置，即channel-mixing；(ii) 不同空域位置，即token-mixing；或者上述两种组合。在CNN中，ii采用卷积或者池化进行，更深的层具有更大的感受野；与此同时，i则采用卷积进行；更大的核则同时进进行i与ii。在Vision Transformer以及其他注意力架构中，自注意力层同时执行i与ii；而MLP则孩子能够i。Mixer背后的思想在于：将上述两种特征混合进行显示分离且均通过MLP进行实现。

前面的Figure1给出了Mixer的架构示意图，Mixer以序列长度为S的非重叠图像块作为输入，每个图像块将投影到期望的隐层维度C。这将得到一个二维实值输入表。如果原始输入图像分辨率为，每个块的分辨率为，就那么序列长度，所有块采用相同的投影矩阵进行线性投影。Mixer包含多个等尺寸的层，每个层包含两种MLP模块：

token-mixing MLP：它作用于X的列，从映射到;
Channel-mixing MLP：它作用于X的行，从映射到。

每个MLP模块包含两个全连接层与一个非线性层，它们可以定义如下：

其中表示GELU激活函数，分别表示token-mixing与channel-mixing MLP中隐层宽度。需要注意的是，的选择独立于输入图像块的数量。因此，网络的计算复杂度与输入块的数量成线性关系，这不同于ViT的二次关系。由于，独立于块尺寸，因此，整体计算量与图像的像素数成线性关系，类似于CNN。

正如上述所提到的，相同的channel(token)-mixing MLP作用于X的每一行(列)。对MLP的参数进行绑定就是一种很自然的选择，它可以提供类似卷积特征的位置不变形。然而，在CNN中进行跨通道参数绑定并不常见。比如，CNN的分离卷积对每个通道采用不同的卷积核，这与本文MLP中的处理机制(所有通道采用相同的核)不相同。这种参数绑定可以避免架构随隐层维度C、序列长度S提升而增长过快，进而导致了显著的显存节省。令人惊讶的是：这种参数绑定机制并不会影响性能。

Mixer中的每一层(除了初始块投影层)采用相同尺寸的输入，这种“各向同性”设计类似于Transformer和RNN中定宽；这与CNN中金字塔结构(越深的层具有更低的分辨率、更多的通道数)不同。

除了MLP外，Mixer还采用其他标准架构成分：跳过连接、LayerNorm。此外，不同于ViT，Mixer并没有采用position embedding，这是因为token-mixingMLP对于输入的顺序极为敏感。最后，Mixer采用了标注分类头，即全局均值池化+线性分类器。基于Flax与jax的参考实现code如下：

Experiments

基于中等与大尺度数据的预训练，我们在不同下游分类任务上对所提MLP-Mixer的性能进行了评估。我们主要对以下三个问题比较感兴趣并进行重点分析。

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

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

Downstream Taks 我们采用了多个主流下游任务，比如ILSVRC2012、清洗后ReaLLabels、CIFAR10/100、Oxford-IIIT Pets、Oxford Flowers-102等。我们同事还在VTAB-1k数据集上进行了评估。

Pre-training data 我们参考标准的迁移学习步骤：预训练+微调。我们在两个公开数据(ILSVRC2012与ImageNet-21k)上进行预训练。为评估在更大尺度上的性能，我们还在JFT-300M数据集上进行了训练。

Pre-training details 训练超参：优化器为Adam，batch=4096，weight decay以及梯度裁剪、线性学习率warmup，输入分辨率均为224。数据增广方面包含RandAugment、mixup、dropout、stochastic depth等等。

Fine-tuning details 微调超参：优化器为momentum-SGD，batch=512，梯度裁剪，cosine学习率机制+线性warmup，并未采用weight decay。微调阶段采用了更高的分辨率。

Metrics 我们从模型性能、计算复杂度之间的均衡性能。对于计算复杂度，我们计算了两个测度：(1) TPU-v3上的合计预训练耗时，它包含理论FLOPs、硬件上的计算效率以及数据加载效率；(2) TPU-v3上的吞吐量。

Models 上表列出了Mixer模型的不同配置以对标CNN与注意力模型。此外，在对标模型方面，ViTs模型包含ViT、HaloNet、BiT等等；CNN模型包含NFNet、MPL、ALIGN等。

Main results

上表给出了最大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倍。

The role of the model scale

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

当在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倍。

全文到此结束，更多消融实验与分析建议查看原文。