全MLP结构模型学习笔记

全MLP结构模型学习笔记

MLP-Mixer: An all-MLP Architecture for Vision

代码：GitHub - google-research/vision_transformer

Abstract

卷积神经网络（CNN）是计算机视觉的常用模型。最近，基于注意力的网络，如Vision Transformer，也变得很流行。在本文中，我们表明，虽然卷积和注意力都足以获得良好的性能，但它们都不是必需的。我们提出了MLPMixer，一种专门基于多层感知器（MLP）的体系结构。MLP混合器包含两种类型的层：一种是将MLP独立应用于图像patches的层（即“混合”每个位置的特征），另一种是跨patches应用MLP的层（即“混合”空间信息）。当在大型数据集上或使用现代正则化方案进行训练时，MLP混合器在图像分类基准上获得了有竞争力的分数，预训练和推理成本与最先进的模型相当。我们希望这些结果能够激发出成熟的CNN和变压器领域之外的进一步研究。

1 Introduction

正如计算机视觉的历史所证明的那样，更大数据集的可用性加上计算能力的增加往往会导致范式的转变。虽然卷积神经网络（CNN）已经成为计算机视觉的事实标准，但最近，基于自我注意层的另一种视觉转换器（ViT）获得了最先进的性能。ViT延续了从模型中去除手工制作的视觉特征和归纳偏见的长期趋势，并进一步依赖于从原始数据中学习。

我们提出了MLP-Mixer体系结构（简称“Mixer”），这是一种具有竞争力但在概念和技术上都很简单的替代方案，它不使用卷积或自我关注。相反，Mixer的体系结构完全基于多层感知器（MLP），这些感知器在空间位置或特征通道上重复应用。混合器仅依赖于基本矩阵乘法例程、数据布局的更改（重塑和换位）以及标量非线性。

图1描述了混合器的宏观结构。它接受一系列线性投影的图像patches（也称为标记），形状为“patches×通道”表，作为输入，并保持此维度。混合器使用两种类型的MLP层：信道混合MLP和令牌混合MLP。信道混合MLP允许不同信道之间的通信； 他们独立地对每个令牌进行操作，并将表中的各行作为输入。令牌混合MLP允许不同空间位置（令牌）之间的通信；它们独立地对每个通道进行操作，并将表中的各个列作为输入。这两种类型的层是交错的，以实现两个输入维度的交互。

在极端情况下，我们的体系结构可以看作是一种非常特殊的CNN，它使用1×1卷积进行信道混合，使用全感受野的单信道深度卷积和参数共享进行令牌混合。然而，情况并非如此，因为典型的CNN不是混合器的特例。此外，在MLP中，卷积比普通矩阵乘法更复杂，因为它需要对矩阵乘法和/或专用实现进行额外的代价降低。

尽管它很简单，Mixer还是取得了有竞争力的效果。在对大型数据集进行预培训时（即。，∼100万张图像），在精确度/成本权衡方面达到了CNNs和Transformers之前声称的接近最先进的性能。这包括ILSVRC2012“ImageNet”上87.94%的top-1验证准确率【13】。当对更适度规模的数据进行预培训时（即。，∼1–10M图像），再加上现代正则化技术[49，54]，Mixer也实现了强大的性能。然而，与ViT类似，它略低于专门的CNN架构。

2 Mixer Architecture

现代深度视觉体系结构由以下各层组成：（i）在给定的空间位置，（ii）在不同的空间位置之间，或同时在两者之间混合特征。在CNN中，（ii）使用N×N卷积（对于N>1）和池来实现。更深层次的神经元有更大的感受野[1，28]。同时，1×1卷积也执行（i），较大的核同时执行（i）和（ii）。在视觉转换器和其他基于注意的架构中，自我注意层允许（i）和（ii），MLP块执行（i）。混频器体系结构背后的思想是明确区分每位置（信道混合）操作（i）和跨位置（令牌混合）操作（ii）。这两种操作都是通过MLP实现的。图1总结了体系结构。

Mixer将S个不重叠的图像面片序列作为输入，每个面片投影到所需的隐藏维度C。这将生成一个二维实值输入表X∈ R^{S×C}。如果原始输入图像具有分辨率（H，W），并且每个面片具有分辨率（P，P），则面片数为S=HW/P^2。所有面片都使用相同的投影矩阵进行线性投影。混合器由多个尺寸相同的层组成，每层由两个MLP块组成。第一个是令牌混合MLP：它作用于X的列（即，它应用于转置的输入表X^T），映射R^S → R^S，并在所有列中共享。第二个是通道混合MLP：它作用于X行，映射R^C→ R^C，并在所有行中共享。每个MLP块包含两个完全连接的层和非线性独立应用于其输入数据张量的每一行。混合器层可编写如下（省略层索引）：

这里σ是元素非线性（GELU【16】）。D_S和D_C分别是令牌混合和信道混合mlp中的可调隐藏宽度。请注意，D_S的选择与输入面片的数量无关。因此，网络的计算复杂度在输入面片的数量上是线性的，不像ViT的复杂度是二次的。由于D_C与面片大小无关，因此与典型的CNN一样，图像中像素数的总体复杂性是线性的。

如上所述，相同的信道混合MLP（令牌混合MLP）应用于X的每一行（列）。绑定信道混合MLP的参数（在每一层内）是一种自然选择，它提供了位置不变性，这是卷积的一个显著特征。然而，跨通道绑定参数的情况要少得多。例如，在一些CNN中使用的可分离卷积（9，40）将卷积独立于其他信道应用于每个信道。然而，在可分离卷积中，不同的卷积核应用于每个信道，这与混合器中的令牌混合MLP不同，混合器中的令牌混合MLP对所有信道共享相同的核（全感受野）。当增加隐藏维度C或序列长度S时，参数绑定可以防止体系结构增长过快，并导致显著的内存节省。令人惊讶的是，这种选择并不影响实证绩效，见补充A.1。

混合器中的每一层（初始面片投影层除外）都接受相同大小的输入。这种“各向同性”设计最类似于变压器，或其他领域的深RNN，也使用固定宽度。这与大多数具有金字塔结构的CNN不同：更深的层具有更低的分辨率输入，但有更多的通道。请注意，虽然这些是典型的设计，但也存在其他组合，如各向同性网状结构[38]和棱锥状VIT[52]。

除了MLP层之外，Mixer还使用其他标准体系结构组件：跳过连接[15]和层规范化[2]。与VIT不同，Mixer不使用位置嵌入，因为令牌混合MLP对输入令牌的顺序很敏感。最后，混合器使用一个标准的分类头，该分类头具有全局平均池层，后跟一个线性分类器。总体而言，该体系结构可以用JAX/Flax紧凑地编写，代码在补充E中给出。

3 Experiments

我们在一系列中小型下游分类任务中，评估了使用中大型数据集预先训练的MLP混合器模型的性能。我们对三个主要数量感兴趣：（1）下游任务的准确性；（2）预训练的总计算成本，这在上游数据集上从头开始训练模型时很重要；（3） 测试时间吞吐量，这对从业者很重要。我们的目标不是展示最先进的结果，而是要表明，值得注意的是，一个简单的基于MLP的模型与当今最好的卷积模型和基于注意力的模型具有竞争力。

Downstream tasks

我们使用流行的下游任务，如ILSVRC2012“ImageNet”（1.3M训练示例，1k类），带有原始验证标签[13]和清理过的真实标签[5]，CIFAR-10/100（50k示例，10/100类）[23]，Oxford IIIT Pets（3.7k示例，36类）[32]，以及Oxford Flowers-102（2k示例，102类）[31]。我们还使用了视觉任务适应基准（VTAB-1k），该基准由19个不同的数据集组成，每个数据集都有1k个训练示例【58】。

Pre-training

我们遵循标准的迁移学习设置：先进行预培训，然后对下游任务进行微调。我们在两个公共数据集上对模型进行预训练：ILSVRC2021 ImageNet和ImageNet-21k，这是ILSVRC2012的超集，包含21k类和1400万图像[13]。为了更大规模地评估性能，我们还对JFT-300M进行了培训，这是一个拥有300M示例和18k类的专有数据集[44]。正如Dosovitskiy等人【14】、Kolesnikov等人【22】所做的那样，我们针对下游任务的测试集消除了所有预训练数据集。我们使用Adam对所有模型进行分辨率为224的预训练，其中β1=0.9，β2=0.999，线性学习率预热10k步，线性衰减，batch size为4096，重量衰减，梯度裁剪为全局范数1。对于JFT-300M，除了随机水平翻转外，我们还通过应用Szegedy等人的裁剪技术对图像进行预处理。对于ImageNet和ImageNet-21k，我们采用了额外的数据增强和正则化技术。特别是，我们使用RandAugment[12]、mixup[60]，dropout【43】和随机深度【19】。这套技术的灵感来自timm library【54】和Touvron等人【48】。补充B中提供了有关这些超参数的更多详细信息。

Fine-tuning

我们使用动量SGD、批次大小512、全局范数1下的梯度剪裁和具有线性预热的余弦学习速率计划进行微调。微调时不使用权重衰减。按照常见做法【22，48】，我们也会对训练前使用的分辨率进行微调，以获得更高的分辨率。由于我们保持面片分辨率不变，这会增加输入面片的数量（例如从S到S0），因此需要修改混合器的令牌混合MLP块的形状。形式上，等式（1）中的输入左乘以权重矩阵W_1∈ R^{D_S×S}，在更改输入维度S时，必须调整此操作。为此，我们将隐藏层宽度从D^S增加到D^{S^{'}}，与面片数成比例，并初始化（现在更大的）权重矩阵\mathbf{W}_{2}^{\prime} \in \mathbb{R}^{D_{S^{\prime}} \times S^{\prime}}，具有块对角矩阵，其对角线上包含W_2的副本。该特定方案仅允许S^{'}=K^2S和K∈N、有关更多详细信息，请参见补充C。在VTAB-1k基准上，我们分别在具有小输入图像和大输入图像的数据集上遵循位超规则[22]和分辨率为224和448的微调混频器模型。

Metrics

我们评估了模型计算成本和质量之间的权衡。对于前者，我们计算了两个指标：（1）TPU-v3加速器上的总预训练时间，它结合了三个相关因素：每个训练设置的理论失败率、相关训练硬件上的计算效率和数据效率。（2） TPU-v3上的吞吐量（以映像/秒/核心为单位）。由于不同大小的模型可能受益于不同的批次大小，我们将扫描批次大小，并报告每个模型的最高吞吐量。对于模型质量，我们关注微调后的top-1下游精度。在两种情况下（图3，右图和图4），微调所有模型的成本太高，我们报告了通过解决图像和标签的冻结学习表示之间的“2-正则化线性回归问题”获得的少数镜头精度。

 

Models

我们将表1中总结的混合器的各种配置与最新的、最先进的、CNN和基于注意的模型进行了比较。在所有的图和表中，基于MLP的Mixer模型用粉红色标记，基于卷积的模型用黄色标记，基于注意力的模型用蓝色标记。视觉变换器（VIT）具有与混合器相似的模型比例和面片分辨率。HaloNets是一种基于注意力的模型，它使用一种类似ResNet的结构，具有局部自我注意层，而不是3×3卷积[51]。我们专注于特别有效的“HaloNet-H4（base 128，Conv-12）”模型，该模型是更广泛的HaloNet-H4架构的混合变体，其中一些自我注意层被卷积所取代。注意，我们用蓝色标记了带有注意力和卷积的HaloNets。Big Transfer（BiT）[22]模型是为迁移学习而优化的网络。NFNets【7】是无规范化器的RESNET，具有针对ImageNet分类的若干优化。我们考虑NFNet-F4+型号变体。我们考虑MPL[34]和ALIGN[21]来实现高效的网络体系结构。MPL是在JFT-300M图像上进行大规模预训练的，使用ImageNet的元伪标签代替原始标签。我们将其与EfficientNetB6宽模型变体进行比较。在有噪声的web图像-文本对上以对比方式对齐预训练图像编码器和语言编码器。我们将其与效率最高的ET-L2图像编码器进行比较。

3.1 Main results

表2给出了最大混合器模型与文献中最先进模型的比较。“ImNet”和“ReaL”列指原始ImageNet验证[13]和清理过的ReaL[5]标签。“Avg.5”代表所有五个下游任务（ImageNet、CIFAR-10、CIFAR-100、Pets、Flowers）的平均性能。图2（左）显示了精度计算边界。当在ImageNet-21k上进行预训练并进行额外的正则化时，混频器实现了整体强大的性能（在ImageNet上排名前1的84.15%），尽管略低于其他模型2。在这种情况下，正则化是必要的，没有正则化，混频器会过度拟合，这与ViT的类似观察结果一致【14】。当从ImageNet上的随机初始化训练混合器时，同样的结论成立（见第3.2节）：混合器-B/16在分辨率224时达到76.4%的合理分数，但往往过度拟合。该分数与香草ResNet50相似，但落后于ImageNet“从头开始”设置的最先进CNN/混合网络，例如84.7%的僵尸网络[42]和86.5%的NFNet[7]。

当上游数据集的大小增加时，混合器的性能会显著提高。特别是，Mixer-H/14在ImageNet上实现了87.94%的top-1精度，比BiTResNet152x4高0.5%，比ViT-H/14低0.5%。值得注意的是，Mixer-H/14的运行速度是ViT-H/14的2.5倍，几乎是BiT的两倍。总的来说，图2（左）支持我们的主要主张，即在精度方面，计算折衷混合器与更传统的神经网络架构具有竞争力。该图还显示了培训前总成本与下游精度之间的明显相关性，即使是跨架构类的。

表中的BiT-ResNet152x4是使用SGD进行预训练的，具有动量和长的时间表。由于Adam趋向于更快地收敛，我们使用Dosovitskiy等人的BiT-R200x3模型完成了图2（左）中的图片。[14]使用Adam对JFT-300M进行了预训练。该ResNet的精确度略低，但训练前计算量要低得多。最后，图中还报告了较小ViT-L/16和Mixer-L/16模型的结果。

3.2 The role of the model scale

上一节概述的结果侧重于计算谱上端的（大型）模型。我们现在将注意力转向更小的Mixer型号。

我们可以通过两种独立的方式缩放模型：（1）在预训练时增加模型大小（层数、隐藏维度、MLP宽度）；（2） 在以下情况下提高输入图像分辨率微调，前者影响训练前计算和测试时吞吐量，后者只影响吞吐量。除非另有说明，我们对第224号决议进行微调。

我们将混频器的各种配置（见表1）与类似规模的ViT模型和Adam预先训练的BiT模型进行了比较。结果总结在表3和图3中。当在ImageNet上从头开始训练时，Mixer-B/16达到了76.44%的合理top-1精度。这比ViT-B/16型号落后3%。训练曲线（未报告）表明，两个模型的训练损失值非常相似。换句话说，Mixer-B/16比ViT-B/16更适合。对于Mixer-L/16和ViT-L/16型号，这种差异更为明显。

随着预训练数据集的增长，Mixer的性能稳步提高。值得注意的是，Mixer-H/14在JFT-300M上预先训练并在224分辨率下微调，仅比ImageNet上的ViT-H/14慢0.3%，同时运行速度快2.2倍。图3清楚地表明，尽管混合器在模型比例尺的低端略低于前沿，但它自信地位于高端的前沿。

3.3 The role of the pre-training dataset size

迄今为止的结果表明，在较大数据集上进行预训练可以显著提高混音器的性能。在这里，我们将更详细地研究这种效应。

为了研究Mixer利用越来越多的训练示例的能力，我们在JFT-300M的随机子集上对Mixer-B/32、Mixer-L/32和Mixer-L/16模型进行预训练，其中包含233、70、23和7个时期所有训练示例的3%、10%、30%和100%。因此，每个模型都针对相同数量的总步骤进行预训练。我们还在一个更大的JFT-3B数据集[59]上预训练Mixer-L/16模型，该数据集包含大约3B个图像，30k个类，总步长相同。

虽然没有严格的可比性，但这使我们能够进一步推断规模的影响。我们使用ImageNet上的线性5拍top-1精度作为传输质量的代理。对于每次训练前的跑步，我们都会根据最佳的上游验证性能提前停止。结果如图2（右）所示，其中我们还包括ViT-B/32、ViT-L/32、ViT-L/16和BiT-R152x2模型。

当在JFT-300M的最小子集上进行预训练时，所有混频器模型都严重过拟合。BiT模型也过拟合，但程度较小，可能是由于与卷积相关的强感应偏差。随着数据集的增加，Mixer-L/32和Mixer-L/16的性能增长速度都快于BiT；Mixer-L/16不断改进，而BiT模型停滞不前。

同样的结论也适用于ViT，与Dosovitskiy等人的结论一致。然而，大型混频器模型的相对改进更为明显。Mixer-L/16和ViT-L/16之间的性能差距随着数据规模的缩小而缩小。Mixer似乎比ViT更能从不断增长的数据集大小中获益。人们可以通过归纳偏差的不同来推测和解释它：ViT中的自我注意层导致学习功能的某些特性，这些特性与真实的底层分布不太兼容，而与混音器结构中发现的特性相比。

3.4 Invariance to input permutations

在本节中，我们将研究混频器和CNN结构的电感偏置之间的差异。具体而言，我们在预培训后在JFT-300M上培训Mixer-B/16和ResNet50x1型号设置如第3节所述，并使用两种不同的输入变换之一：（1）将16×16面片的顺序和每个面片内的像素置换为共享置换；（2） 在整个图像中全局排列像素。所有图像都使用相同的排列。我们在图4（下图）中报告了ImageNet上经过训练的模型的线性5拍top-1精度。图4（顶部）显示了一些原始图像及其两个转换版本。正如所料，混合器对面片和面片内像素的顺序保持不变（蓝色和绿色曲线完美匹配）。另一方面，ResNet的强感应偏差依赖于图像中特定顺序的像素，当面片排列时，其性能会显著下降。值得注意的是，当全局排列像素时，混频器的性能下降幅度要小得多(∼与ResNet相比，下降了45%(∼下降75%）。

3.5 Visualization

通常可以观察到，CNN的第一层倾向于学习作用于图像局部区域像素的类Gabor检测器。相反，Mixer允许在令牌混合MLP中进行全局信息交换，这回避了它是否以类似方式处理信息的问题。图5显示了在JFT-300M上训练的混合器前三个令牌混合MLP的隐藏单元。回想一下，令牌混合MLP允许不同空间位置之间的全局通信。一些学习到的特征操作于整个图像，而其他特征操作于较小的区域。深层似乎没有明确可识别的结构。与CNN类似，我们观察到许多相位相反的特征检测器对【39】。学习单元的结构取决于超参数。第一个嵌入层的图如补充D的图7所示。

4 Related work

MLP混频器是一种新的计算机视觉体系结构，它不同于以前成功的体系结构，因为它既不使用卷积层，也不使用自我注意层。然而，设计选择可以追溯到CNN【24、25】和变压器【50】文献中的想法。

自从AlexNet模型【24】超越了基于手工图像特征的主流方法【35】以来，CNN已经成为计算机视觉的事实标准。许多工作集中于改进CNN的设计。Simonyan和Zisserman【41】证明，仅使用具有小3×3核的卷积就可以训练最先进的模型。he等人【15】引入了跳跃连接和批量归一化【20】，这使得能够训练非常深入的神经网络并进一步提高性能。一项重要的研究已经调查了使用稀疏卷积的好处，例如分组（57）或深度（9，17）变体。本着与我们的标记混合MLP类似的精神，Wu等人[55]在自然语言处理的深度卷积中共享参数。Hu et al.（18）和Wang et al.（53）建议使用非局部操作增强卷积网络，以部分缓解CNN的局部处理限制。Mixer将使用小核卷积的想法发挥到极致：通过将核大小减小到1×1，它将卷积转换为独立应用于每个空间位置的标准密集矩阵乘法（信道混合MLP）。仅此一点不允许空间信息的聚合，为了进行补偿，我们将密集矩阵乘法应用于所有空间位置的每个特征（令牌混合MLP）。在Mixer中，矩阵乘法在“patches×features”输入表上按行或列进行，这也与稀疏卷积的工作密切相关。混合器使用跳过连接[15]和规范化层[2，20]。

在计算机视觉中，基于自我注意的变压器架构最初用于生成建模[8，33]。它们对图像识别的价值后来得到了证实，尽管是与卷积式局部偏差相结合[37]，或是在低分辨率图像上[10]。Dosovitskiy等人【14】介绍了ViT，这是一种纯变压器模型，具有较少的局部性偏差，但可以很好地扩展到大数据。ViT在流行的视觉基准上实现了最先进的性能，同时保持了CNN的健壮性[6]。Touvron等人[49]使用广泛的正则化在较小的数据集上有效地训练ViT。Mixer借鉴了最近基于transformer的架构的设计选择。混合器MLP块的设计源于V aswani等人【50】。将图像转换为一系列补丁并直接处理这些补丁的嵌入源于Dosovitskiy等人【14】。

最近的许多工作都致力于为vision设计更有效的体系结构。Srinivas等人[42]用自我注意层取代了resnet中的3×3卷积。Ramachandran等人【37】、Tay等人【47】、Li等人【26】和Bello【3】设计了具有新的注意力样机制的网络。混合器可以看作是正交方向上的一个步骤，不依赖于位置偏差和注意机制。

Lin等人的工作密切相关。它在CIFAR-10上通过使用全连接网络、大量数据扩充和使用自动编码器进行预训练获得了合理的性能。Neyshabur【30】设计了定制的正则化和优化算法，并训练了一个完全连接的网络，在小规模任务中获得了令人印象深刻的性能。相反，我们依赖令牌和信道混合MLP，使用标准的正则化和优化技术，并有效地扩展到大数据。

传统上，在ImageNet上评估的网络【13】是使用初始风格的预处理从随机初始化中训练出来的【46】。对于较小的数据集，ImageNet模型的传输很流行。然而，现代最先进的模型通常使用在更大数据集上预先训练的权重，或者使用更新的数据扩充和训练策略。例如，Dosovitskiy等人【14】、Kolesnikov等人【22】、Mahajan等人【29】、Pham等人【34】、Xie等人【56】都利用大规模预训练提升了图像分类的技术水平。增强或正则化变化带来的改进示例包括Cubuk等人【11】，他们通过学习数据增强获得了优异的分类性能，以及Bello等人【4】，他们表明，如果使用最新的训练和增强策略，规范重新网仍然接近最先进水平。

5 Conclusions

我们描述了一个非常简单的vision架构。我们的实验表明，就训练和推理所需的准确性和计算资源之间的权衡而言，它与现有的最先进的方法一样好。我们认为，这些结果引发了许多问题。在实践方面，研究模型学习到的特征并确定主要与CNN和Transformers学习到的差异（如有）。在理论方面，我们想了解隐藏在这些不同特征中的归纳偏差，以及它们在推广中的作用。最重要的是，我们希望我们的结果能够激发进一步的研究，超越基于卷积和自我注意的已建立模型的领域。看看这样的设计是否适用于NLP或其他领域，这将是非常有趣的。