MLP-Mixer: An all-MLP Architecture for Vision

原文链接：

MLP-Mixer: An all-MLP Architecture for Vision

文章概述

        CNN已成为计算机视觉的首选方案，近来兴起的基于注意力的网络(比如Vision Transformer)也逐渐变得流行。在本文中我们发现卷积和注意力虽然会得到更好的表现，但都不是必须的。本文提出一种基于MLP的结构-MLP Mixer，该结构包括两种类型的层，一种用于混合每个位置的特征，一种用于混合空间信息。实验表示该网络不仅能在数据集上取得较好的成绩，而且在预训练和预测时花费的成本与目前最先进的模型相当。

介绍

        作者提出了MPL-Mixer，一种具有竞争力并且在概念和技术上更简单的结构，仅依赖于在空间位置和特征通道上反复使用多层感知机，Mixer只使用基本的矩阵乘法来改变数据的布局以及其非线性。Mixer利用了两种类型的MLP层，channel-mixing MLPs和token-mixing MLPs。前者用于不同通道间的交互，后者用于不同token之间。

方法

        该网络的结构如下图所示：

       一共有三个部分Per-patch Fully-connected、Mixer以及最后的分类器。分类器由一个全局池化层+全连接层构成。

        Per-patch Fully-connected层，主要是将输入的图片分patch进行映射，最后转换为X=S*C的矩阵。假设输入图像为H*W，将其划分为大小为P*P的patch（共有S = HW/个patch），并由相同参数的矩阵对每个patch映射为一维特征，长度为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的矩阵。其中Patch的值和MLP输出单元个数为超参数。

        Mixer Layer由多个大小相同的层组成，每一层包括两个MLP块。第一个是token-mixing MLP（即图中的MLP1）,作用于X的列向量。第二个是channel-mixing MLP（图中的MLP2），作用于X的行向量。每个MLP包括两个全连接层和一个GELU激活函数。此外，Mixer 还用了跳跃连接（Skip-connection）和层归一化（Layer Norm），这里的跳跃连接即残差链接，将输入输出相加。

        可以发现X的行代表了同一空间位置在不同通道上的信息，列代表了同一通道不同空间位置的信息。即对X的每一行进行操作可以实现通道域的信息融合，对Table的每一列进行操作可以实现空间域的信息融合。在CNN中，是通过卷积来核来实现通道空间上的信息融合，在Transformer中，是通过Multi-Head Attention来实现空间和通道的信息融合。

         除了最初patch映射层，每一层输出大小都是一样的，这种”各向同性“的设计与transformer或RNN中使用固定的宽度相似，与大多数越深层网络输入越小，通道数越多的CNN不同。与ViTs不同，Mixer不使用位置编码，因为token-mixing MLPs对输入的顺序很敏感。

实验

        作者评估了MLP-Mixer模型的性能，这些模型用中到大规模的数据集进行了预训练，然后应用在一系列中小型的下游分类任务上。如下表所示，展示了本文训练的不同规模下MLP-Mixer的具体架构。

总体实验结果 

        在不同数据集上进行预训练后的模型，在ILSVRC2012 “ImageNet” 下游任务上的结果如下表所示，其中Avg.5表示在ImageNet,CIFAR-10, CIFAR-100, Pets, Flowers这五个下游任务的平均性能。

        表2可以看出Mixer在成本相近的情况下具有与最先进的模型相当的精度。ImageNet-21k上进行预训练时，Mixer实现了整体上的强大性能（在ImageNet上有84.15%的top-1），尽管略逊于其他模型，但其运算速度比ViT和BiT快很多。 当上游数据集的规模增加时，（使用JFT进行预训练）Mixer的性能明显提高，Mixer-H/14在ImageNet上达到了87.94%的top-1准确率，比BiTResNet 152x4好0.4%，只比ViT-H/14低0.6%，并且其运行速度是ViT-H/14的2.5倍，几乎是BiT的两倍。

验证模型规模的作用 

         使用不同规模的MLP-Mixer进行对比，实验结果如下所示：

         由上图可以看出，Mixer与这些性能极佳的ResNets、ViTs和混合模型一样好，并且与HaloNet、ViT、NFNet和MPL处于同一水平。且随着数据规模的增长，Mixer的精度赶上或超过了BiT和ViT，即大型Mixer和ViT模型之间的差距缩小了。 

         ImageNet 验证不同规模的 ViT、BiT 和 Mixer 模型的 top-1 准确度与总预训练计算（左）和吞吐量（右）的结果如下图所示。

        可以看出，尽管Mixer在模型规模的较小时略低于最好的结果，但它在规模较大的时候却明显地表现出了最好的水平。 

 验证预训练数据集的大小对模型的影响

       

        从上表可以看出，随着数据集的增加，Mixer-L/32和Mixer-L/16的性能都比BiT增长得快；Mixer-L/16不断改进，而BiT模型则趋于平稳。Mixer-L/16和ViT-L/16之间的性能差距随着数据规模的扩大而缩小，Mixer比ViT更受益于数据集规模的增长。这是由于ViT中的自我关注层导致所学函数的某些属性与Mixer架构发现的真实基础分布的兼容性较差。  

输入序列的不变性

        这里对输入的图像进行处理，1）对patch进行打乱，并在每个patch中使用相同的排列方式对像素进行排列 2）对整个图像中的像素进行全局排列，所有图像使用相同的排列方式。并在Mixer和Resnet上进行对比实验，具体结果如下图所示：

         可以看出 Mixer对于输入patch的顺序和patch内像素的顺序是没有影响的（蓝色和绿色曲线完全一致），ResNet在打乱后的数据上 表现出明显的下降 即其性能依赖于图像中像素的特定顺序。另外，在对像素进行全局排列时，Mixer的性能与ResNet相比下降得更少。 

总结 

        本文提出了一个非常简单的结构，实验证明该网络在权衡精确度和训练预测所需要计算的同时，达到了跟目前最好的模型差不多的效果。作者认为研究该模型与CNN、transformer学到的特征的主要区别可能是有用的，并且希望引发进一步研究，超越基于卷积和自我注意力模型，并对该设计在其他领域中（如NLP）的作用进行探索。