论文翻译--MLP-Mixer: An all-MLP Architecture for Vision

MLP-Mixer: An all-MLP Architecture for Vision

MLP-Mixer:一个全MLP的视觉架构

Google Research, Brain Team

 

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

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

摘要

CNNs是CV的首选模型。最近，attention-based networks也开始流行了，例如：Vision Tranformer。

本文章中，我们表明convolutions 和attention尽管都足以获得良好的性能，但是他们都不是必须的。我们提出了 MLP-Mixer，一个专门基于多层感知机的架构（MLPs）。MLP-Mixer包含两种类型的layers：

一个是MLPs独立应用于image patches（i.e. “mixing” the per-location features)，

另一个是MLPs跨patches应用(i.e. “mixing” spatial information)

当在大数据及或带有正则化方案进行训练时，MLP-Mixer在图像分类的benchmarks上获得了具有竞争力的score，预训练和推理成本可与最先进的模型相媲美。我们希望这个结果可以激发超越地位稳固的CNNs和Transformers的进一步研究。

引言

正如计算机视觉的历史所示，更大的数据集的可用性加上计算能力的增加，通常会导致范式的转变。虽然CNNs已经是计算机视觉事实上的标准，但是最近的Vision Transformers【14】（ViT）,一种基于self-attention的替代品获得了SOTA的性能。ViT延续了从模型中去除手工制作的视觉特征和可归纳的偏置的长期趋势，并进一步依赖于从原始数据中心学习。

我们提出了一种MLP-Mixer架构（简称“Mixer”）,这是一种具有竞争力但在概念和技术上比较简单的替代方案，这里边不用convolution 或self-attention.相反，Mixer的架构完全基于多层感知机（MLP），它是一种跨空间位置或特征通道重复使用的MLP。Mixer仅仅依赖于基本的矩阵乘法routines,对data layout(reshapes和transpositions)的更改和scalar non-linearities（标量非线性）。

图1描述了Mixer的宏观结构。它接受a sequence of linearly projected image patches(also referred tokens)形状为“patches*channels”的表作为一个输入，and 保持这个维度。Mixer 使用了两种类型的MLP layers: channel-mixing MLPs和token-mixing MLPs。这个channel-mixing MLPs允许不同channels之间的通信；他们独立地操作每一个token并且将表中的各个行作为输入。这个token-mixing MLPs允许不同的spatiallocations（tokens）之间的通信；他们独立的操作每一个channel并且将表中的各个列作为输入。这两种类型的layes是相互交错的，以实现两个输入维度的交互。

在极端的情况下，我们的架构可以看做一个非常特殊的CNN，它使用1*1的卷积用于channel mixing,以及single-channel depth-wise convolutions of a full receptive fifield（感受野） and parameter sharing for token mixing。但是相反却不然，因为典型的CNNs并不是Mixer的特殊情况。此外，卷积比MLPs中的脸是矩阵相乘更复杂，因为它需要额外地减少矩阵乘法和或者特殊的实现。

尽管它很简单，但是Mixer取得了有竞争力的结果。挡在大型数据集上进行预训练时，它在accuracy 和cost的权衡上达到了CNNs和Transformer之前声称的SOTA的性能。这包括在ILSVRC2012,"ImageNet"【13】上最顶级验证的87.94%的准确率。当在更适度规模的数据（即~1-10M images）上进行预训练时，加上modern的正则化家属【48,53】，Mixer 也取得了很强的性能。但是和ViT类似，性能略低于CNN架构。

 

2 Mixer 架构

现代深度视觉框架由混合特征的layers组成（i）在一个给定的空间位置上（ii）在两个不同的空间位置之间，或者两者同时进行。

• 在CNNs中，（ii）的实现是使用N*N卷积（for N >1）和pooling来实现的。深层的神经元有较大的感受野【1,28】。同时，1*1的卷积也实现（i）,而较大的kernels同时实现（i）和（ii）.
• 在ViT和其他attention-based的架构中，self-attention layers实现（i）和（ii）并且MLP-blocks 实现（i）.
• Mixer架构背后的想法是清晰地分离开per-location(channel-mixing)（i）和cross-location(token-mixing)(ii).这两部分都使用MLPs实现。
•  

图1中总结了这个结构。Mixer以一系列不重叠的imgage patches  S作为输入，每一个都投影到一个所需的hidden dimension C上。就得到了一个二维实数的input table,。如果原始图像的分辨率（H，W），并且每一个patch的分辨率为（P，P）,那么path的数量为。所有的patche都用相同的映射矩阵进行线性映射。Mixer由多个相同size的layers组成，并且每一层都由两个MLP blocks组成。第一个是token-mixing MLP block：它作用于X的列上（i.e 它被应用于一个输入的转置上），，在所有列中共享。第二个是channel-mixing MLP block：它作用于X的行上，maps ,在所有行中共享。每个MLP block 包含两个全连接层并且一个非线性独立用于输入数据tensor的每一行。Mixer layers可以被写成如下（省略layer指数）：

这里边是一个元素的非线性（GELU【16】）。分别是在token-mixing和channel-mixing MLPs中可调的隐藏宽度。注意，是不依赖于输入patches的数量来选择的。因此，网络的计算复杂度在输入的patch数量上是线性的。不像ViT的复杂度是二次的。因为与patch的大小无关，整体复杂性与一个图像的像素数量是线性的关系，对于一个典型的CNN。

综上所述，相同的channel-mixing MLP（token-mixing MLP）应用于X的每一row(xolumn)。固定channel-mixing MLPs的阐述（在每一层内）是一个自然的选择——它提供了位置不变性，这是卷积的一个特征。但是跨通道绑定参数则不那么常见。例如，在一些CNNs中使用的分离的卷积【9,40】，将卷积用于每个信道，而不是独立于其他信道。然而，在可分离卷积中，每个通道都应用了不同的卷积核，不像Mixer中的token-mixing MLPs，对所有通道共享相同的kernel(完全感受野的)。参数固定可以防止在增加隐藏维度C或序列长度S时，体系结构增长太快，并可以节省显存。令人惊讶的是，这个选择并不影响性能，见补充A.1.

在Mixer中的每一层（初始化batch  projection layer除外）接受相同大小的输入。这个“isotropic”的设计与Transformers或其他领域的deep RNNs最为相似,他们也使用了一个固定的宽度。这不像大多数的CNNs有一个金字塔结构：较深的层的输入分辨率较低但是通道数更多。请注意虽然这些都是典型的设计，但也存在其他的组合，例如isotropic ResNet[38],和pyramidal ViT[51]。

除了MLP layers之外，Mixer 也用了其他标准的人工智能组件：skip-connection【15】和Layer Normalization【2】。此外，不像ViTs，Mixer不使用position embeddings因为token-mixing MLPs对输入 tokens的顺序敏感，因此可以学习表示位置。最后Mixer使用一个标准的分类头带有GAP（全局平均池化层）后面是一个线性分类器。总的来说，这个结构可以在JAX/Flax上完整的编写，并且代码在补充E中给出。

3.实验

我们评估MLP-Mixer models的性能，使用中奥大规模数据集的预训练，在一系列中小型的下游分类任务上。我们对三个主要的数量很感兴趣：（1）在下有任务上的Accuracy（2）预训练的总计算成本，这在上游数据及上从头训练模型时很重要（3）Throughput at inference time ,这对从业者很重要。我们的目标不是战士最SOTA的结果，而是表明，值得注意的事，一个简单的基于MLP的模型与当今最佳的卷积和attention-based模型具有竞争力。

Downstream tasks

我们使用了多个受欢迎的下游任务例如ILSVRC2012"ImageNet"(1.3M training examples,1k classes)原始的验证label和已经清理过的真实标签，CIFAR-10/100（50k examples,10/100classes）[23],Oxford-IIIT Pets(3.7k examples,36 classes)[32],和Oxford Flowers-102(2k examples,102classes))[31].我们还评估了视觉任务适应基准测试（VTAB-1K）,它由19个不同的数据集组成，每个数据集都有1k个训练示例【57】。

Pre-training data

我们遵循标准的迁移学习的设置：进行预训练，然后在下游任务上进行微调。我们在两个共欧诺个数据集上对所有模型进行预训练：ILSVR2021ImageNet和ImageNet-21k，这是ILSVRC2012的超级包含21k类和1400万幅图像【13】。

为了评估更大规模的性能，我们还在JFT-300M上进行训练，这是一个带有3亿个示例和18K类【44】专有数据集。我们就像Dosovitskiy等人所做的那样，对下游任务的测试集删除所有与预训练数据集重复的数据。

Pre-training details

我们预训练所有的模型使用Adam优化器with ,，并且batch size 为4096，使用weight decay,和在global norm 1 下的梯度剪切。我们使用10ksteps的warmup和线性衰减。我们预训练所有模型使用分辨率为224.对于JFT-300M，我们预处理图像除了随机的水平翻转还使用Szegedy et al[45]等人的cropping technique。对于ImagesNet和ImageNet-21K，我们使用了额外的数据增强和正则化技术。尤其，we use RandAugment [12], mixup [58], dropout [43], and stochastic depth [19].这套技术的灵感来自于timm library[53]和Touvron [47]。更多的关于超参数的细节在补充部分B中给出。

Fine-tuning details

我们用SGD with momentum batch size 512 gradient clipping at global norm1,a cosine learning rate schedule with a linear warmup.进行fine-tune.我们没有使用weight decay在fine-tune的时候。按照常见的实验【22,27】我们还对比预训练时更大的分辨率进行微调。因为我们保持patch的分辨率固定，所以这就增加了输入patch的数量（例如从S到S‘）因此需要修改Mixer的token-mixing MLP block的shape。形式上，在式子（1）的输入是向左乘以权重矩阵并且这个运算在更改输入尺寸S时是必须调整的。为此，我们将隐藏层宽度从Ds增加到D's,并初始化W’2

 

 

........未完待续