生活的情况越艰难，我越感到自己更坚强，甚而也更聪明。——高尔基

Pay Attention to MLPs

摘要

Transformers[1]已经成为深度学习领域最重要的架构创新之一，并在过去几年中取得了许多突破性进展。 在这里，我们提出了一种简单的网络架构gMLP，它基于带有门控的MLP，并表明它在关键的语言和视觉应用中可以与Transformer实现相同的性能。我们的比较表明，自注意力对于 Vision Transformers 来说并不重要，因为 gMLP 可以达到相同的精度。对于BERT，我们的模型在预训练复杂度方面与Transformers相当，在一些下游NLP任务中表现更好。在gMLP表现较差的微调任务中，大幅扩大gMLP模型可以缩小与Transformers的差距。总体而言，我们的实验表明，在数据和计算量增加的情况下，gMLP的扩展能力不亚于Transformers。（gMLP可以像Transformer一样在增加的数据和计算上进行扩展。）

一、引言

Transformer已经在自然语言处理（例如[2, 3, 4, 5, 6]）方面取得了许多突破，并且在计算机视觉（例如[7, 8, 9, 10]）方面表现出色。得益于这一成功，Transformer在很大程度上取代了LSTM-RNN[11]，成为NLP领域的默认架构，并在计算机视觉领域成为ConvNets[12, 13, 14, 15, 16, 17]的替代方案。

Transformer架构结合了两个重要概念：(1)无递归架构，并行计算每个单个标记的表征；(2)多头自注意块，聚合各标记的空间信息。一方面，注意机制[18]引入了归纳偏差，即空间交互应该根据输入表征动态参数化。另一方面，众所周知，静态参数化的MLP可以表示任意函数[19]。因此，自我注意中的归纳偏差是否是Transformers取得显著效果的关键，仍然是一个悬而未决的问题。

在此，我们研究了自注意模块在Transformers的关键语言和视觉应用中的必要性。具体而言，我们提出了一种基于MLP的无自我注意的Transformers替代方案，该方案由通道投影和静态参数化的空间投影组成。我们对该架构的几种设计选择进行了实验，发现当空间投影是线性的并与乘法门控配对时，效果很好（图1）。我们将该模型命名为gMLP，因为它是由带有门控的基本MLP层构建而成的。

我们将gMLP应用于图像分类，并在ImageNet上取得了很好的结果。在类似的训练设置下，gMLP的性能与DeiT[8]，即改进了正则化的Vision Transformer（ViT）[7]相当。在参数减少66%的情况下，gMLP模型的准确度比MLP-Mixer[20]高3%。与Tolstikhin等人[20]、Melas-Kyriazi[21]、Touvron等人[22]和Ding等人[23]的研究结果一样，我们的研究结果也质疑了视觉转换器中自注意力模块的必要性。

我们将gMLP应用于BERT [2]中的掩码语言建模（MLM），这是Transformer最成熟的应用之一，并发现在预训练过程中，它与Transformer一样能够最小化困惑度。我们的实验证明，困惑度只与模型容量相关，对自注意力的存在不敏感。随着容量的增加，我们观察到gMLP和Transformer在预训练和微调指标上的改善速度相同。这是非常显著的，因为它表明尽管没有自注意力，gMLP与Transformer一样能够进行有效的扩展，而且通过增加数据和计算训练更大的模型，可以弥补任何性能差距。在与原始BERT相同的256批次大小×1M步骤的训练设置下，大型gMLP模型在MNLI上达到87.7%的准确率，在SQuAD v2.0上达到82.1%的F1分数。请注意，这些结果优于Devlin等人[2]在使用Transformer时报告的\(BERT_{large}\)结果。

图1：带有空间门控单元（SGU）的gMLP架构概述。该模型由L个具有相同结构和大小的块堆叠而成。所有投影操作都是线性的，"\(\odot\)"表示逐元素相乘（线性门控）。输入和输出协议遵循NLP中的BERT和视觉中的ViT。与Transformer不同，gMLP不需要位置编码，并且在NLP微调期间也不需要屏蔽填充。

左边是：gmlp架构概述 ；右边是：gMLP模块伪代码

对于BERT的微调，在需要跨句子对齐的任务上，Transformer在实践中可能比gMLP更具优势（例如，在300M参数范围内在MNLI-m上提高0.8%），即使预训练困惑度相似。这个问题可以通过使gMLP模型大幅增大，——是Transformers的3倍来解决。一个更实用的解决方案是仅混合少量的自注意力——单头自注意力的大小最多为128，足以使gMLP在我们评估的所有NLP任务上表现优于Transformer，且参数利用效率更高。在某些情况下，这种改进非常显著（例如，在SQuAD v2.0上超过BERTlarge的+4.4%）。

总体而言，gMLPs在视觉和NLP领域令人惊讶的有效性表明，自我关注并不是扩展机器学习模型的必要因素，尽管根据任务的不同，它可能是一个有用的补充。通过增加数据和计算，具有更简单的空间交互机制（如gMLP）的模型可以与Transformer一样强大，并且分配给自注意力的容量可以被删除或大幅减少。

二、模型

我们的模型gMLP由L个具有相同大小和结构的块堆叠而成。让\(X∈\R^{n×d}\)是序列长度为n、维度为d的标记表示： \[ Z=\sigma(XU), \\ \widetilde{Z}=s(Z),\\ Y=\widetilde{Z}V \] 其中，σ是一个激活函数，例如GeLU [24]。U和V定义了沿着通道维度的线性投影，与Transformer中的前馈神经网络（FFN）中的投影相同（例如，对于\(BERT_{base}\)，它们的形状为\(768×3072\)和\(3072×768\)）。为简洁起见，省略了快捷方式、归一化和偏置。 上述公式中的一个关键因素是\(s(·)\)，它是一个捕捉空间交互的层（见下文）。当\(s\)是一个恒等映射时，上述转换退化为一个常规的FFN，其中单个标记被独立处理，没有任何跨标记通信。因此，我们的主要关注点之一是设计一个能够捕捉跨标记之间复杂空间交互的良好的\(s\)。整体的块布局受到反向瓶颈（inverted bottlenecks）[25]的启发，其中s(·)被定义为一个空间深度卷积。请注意，与Transformer不同，我们的模型不需要位置嵌入，因为\(s(·)\)中会捕捉到这些信息。

我们的模型与BERT（用于NLP）和ViT（用于视觉）完全相同的输入和输出协议。例如，在语言任务的微调中，我们将多个文本段连接在一起，然后进行填充，预测是通过保留的<cls>符号的最后一层表示推断出来的。尽管许多这些协议是为了Transformer而引入的，因此对于gMLP来说可能不是最优的，但严格遵循它们有助于避免我们实验中的混淆因素，并使我们的层与现有的Transformer实现更兼容。

2.1 空间门控单元

为了实现跨标记交互，层s(·)必须包含空间维度的收缩运算。最简单的方法是线性投影： \[ f_{W,b}(Z) = W Z + b \] 其中\(W ∈ \R ^{n×n}\)是一个矩阵，其大小与序列长度\(n\)相同，\(b\)是特定标记的偏置。

例如，如果填充的输入序列有128个标记，则\(W\)的形状将是\(128×128\)。与自注意力不同，自注意力的\(W(Z)\)是从\(Z\)动态生成的，而在公式（2）中的空间投影矩阵\(W\)与输入表示无关。

在本文中，我们将层\(s(·)\)表述为线性门控的输出： \[ s(Z)=Z\odot f_{W,b}(Z) \] 其中\(\odot\)表示逐元素乘法。为了训练的稳定性，我们发现将W初始化为接近零的值，将b初始化为1是至关重要的，这意味着\(f_{W,b}(Z) ≈ 1\)，因此在训练开始时\(s(Z) ≈ Z\)。这种初始化确保了每个gMLP块在训练的早期阶段行为类似于常规的FFN，即每个标记都是独立处理的，并且在学习过程中逐渐注入跨标记的空间信息。

我们进一步发现，在门控函数和乘法旁路中，将Z沿着通道维度分为两个独立的部分\(（Z_1，Z_2）\)是有效的： \[ s(Z)=Z_1\odot f_{W,b}(Z_2) \] 我们还将输入归一化为\(f_{W,b}\)，这从经验上提高了大型NLP模型的稳定性。这就得到了图1所示的单元，我们在本文中将其称为空间门控单元（SGU）。在表3中，我们提供了消融研究，将SGU与\(s(·)\)的其他几种变体进行了比较，结果表明SGU工作得更好，缩小了与自我关注的性能差距。

SGU与现有层的联系

SGU的整体结构类似于门控线性单元（GLUs）[26, 27, 28]以及早期的工作，以及包括Highway Networks[29]和LSTM-RNNs[11]在内的早期工作。关键区别在于，我们的门控是基于对空间（跨标记）维度的投影进行计算，而不是基于通道（隐藏）维度。SGU 也与挤压-激发（Squeeze-and-Excite，SE）块[30]在元素乘法方面相关。然而，与SE块不同，SGU根本不包含跨通道的投影，并且也不强制要求置换不变性(以内容为基础的激励模块的关键特征），因为它对空间变换采用了静态参数化。理论上，SGU中的空间投影可以学习表达表面的深度卷积--与典型的具有特定通道滤波器的深度卷积不同，SGU只学习跨通道共享的单一变换。最后，我们注意到SGU提供了除了自注意力以外的捕获高阶关系的替代机制。具体而言，方程（3）的输出包含高达二阶的相互作用（例如\(z_iz_j\)），而自注意力（假设没有非线性）的输出最多包含三阶相互作用（例如\(q_ik_jv_k\)）。在计算成本方面，SGU具有\(n²e/2\)个乘加操作，与点积自注意力的\(n²d\)相当。两者在输入通道大小上都是线性的，在序列长度\(n\)上都是二次的。

三、图像分类

在这里，我们通过将gMLP应用于ImageNet的图像分类任务中来研究其在视觉领域的表现，而无需使用额外的数据。我们将我们的类似MLP的模型与基于传统Transformer的最新注意力模型进行比较，包括Vision Transformer (ViT) [7]、DeiT [8]（改进了正则化的ViT）以及其他几种代表性的卷积网络模型。

表格1总结了我们的gMLP图像分类模型的配置。输入和输出协议遵循ViT/B16的方式，即将原始图像转换为16×16的图块。深度和宽度的选择使得这些模型在容量上与ViT/DeiT相当。与Transformer类似，我们发现gMLP模型很容易对训练数据过拟合。因此，我们采用了与DeiT中相似的正则化方法。为了避免过多的调整，我们只在表格1中从较小模型转移到较大模型时调整了随机深度的强度。所有其他超参数在我们的三个模型中保持不变。详细信息请参见附录A.1。

表1：用于视觉的gMLP模型的架构规格。从较小模型到较大模型，随机深度的存活概率是唯一的超参数变化。

	#L	\(d_{model}\)	\(d_{ffn}\)	参数量 (百万)	FLOPs (十亿)	存活概率
gMLP-Ti	30	128	768	5.9	2.7	1.00
gMLP-S	30	256	1536	19.5	8.9	0.95
gMLP-B	30	512	3072	73.4	31.6	0.80

我们的ImageNet结果总结在表格1和图2中。有趣的是，gMLP模型与DeiT [8]（即使用改进的正则化训练的ViT [7]）具有可比性。这些结果表明，没有自注意力的模型在图像分类任务上可以与Transformers一样高效。事实上，当模型得到适当的正则化时，它们的准确性似乎更与容量相关，而不是与自注意力的存在相关。此外，gMLP模型的准确性-参数/FLOPs的权衡超过了所有同时提出的类似MLP架构[20, 21, 22]，我们将其归因于我们的空间门控单元（请参见下一节的表3进行消融实验）。我们还注意到，虽然gMLP模型在与原始Transformers的比较中表现竞争力，但其性能仍落后于最好的现有卷积网络模型（例如[33, 34]）或混合模型（例如[35, 36, 37, 38, 10]）。

图2：ImageNet准确率与模型容量的关系。

表 2：没有额外数据的 ImageNet-1K 结果

图3展示了gMLP-B中的空间投影矩阵。值得注意的是，学习后的空间权重展现了局部性和空间不变性。换句话说，每个空间投影矩阵有效地学会了使用数据驱动的、不规则（非方形）的卷积核进行卷积操作。

tu-3

图3：gMLP-B中的空间投影权重。每一行显示了同一层中一组选定标记的滤波器（经过重塑成2D形状）。

四、使用BERT进行掩码语言建模

在这里，我们对掩码语言建模（MLM）任务进行了实证研究。预训练和微调的输入/输出协议遵循BERT[2]。与基于Transformer的模型不同，我们不使用位置编码。在gMLP块的微调过程中，我们发现无需屏蔽<pad>标记，因为模型可以快速学会忽略它们。对于剖析和案例研究，所有模型都使用批量大小2048，在RealNews-like的C4子集上进行了125K步的训练。对于主要结果，模型使用批量大小256，在完整的英文C4数据集上进行了100万步的训练。详细信息请参见附录A.2。

我们的初步MLM实验显示，gMLP始终学习到类似Toeplitz矩阵的空间权重（附录C）。这意味着gMLP能够从数据中学习到平移不变性的概念，这个属性在MLM任务中是自然隐含的，因为输入序列的任何偏移都不会影响插槽填充的结果。在这种情况下，学得的fW,b(·)的作用类似于一维卷积，其核大小等于整个序列长度（与具有通道特定滤波器的深度卷积不同，在这里相同的W在通道间是共享的）。在接下来的MLM实验中，我们将W限制为Toeplitz矩阵，以避免冗余的模型参数化（因为在学习后，W无论如何都会变成类似Toeplitz的形式）。请注意，这个约束在实践中对模型质量没有影响。

4.1 消融实验 ：gMLP中门控对BERT预培训的重要性

在下面的表3中，我们建立了我们消融实验的基线。包括：

1. 使用Transformer架构和可学习的绝对位置编码的BERT。
2. 使用Transformer架构和T5风格的可学习相对位置偏置的BERT [5]。我们发现这种偏置既与层次有关，又与头部有关，可以得到最佳结果。
3. 与上述相同，但我们移除了softmax内部所有依赖于内容的项，只保留了相对位置偏置。这个基线是Transformer的一个直接变体，没有自注意力，也可以看作是一个随机合成器 [40]。
4. MLP-Mixer [20]，它用两层空间MLP替代了Transformer中的多头自注意力模块。该模型是为图像分类而开发的，这里我们使用与BERT和gMLP相同的训练设置来研究它在MLM任务中的表现。

我们在表3中将这些基线与几个类似规模的gMLP版本进行了比较。请注意，乘法、分割（最后一行）是我们在方法部分中描述的空间门控单元（Spatial Gating Unit），并在本文的其余部分中使用。首先，SGU在困惑度方面优于其他变体。令人惊讶的是，带有SGU的gMLP也实现了与Transformer相当的困惑度。请注意，最强基线（困惑度=4.26）与我们的结果（困惑度=4.35）之间的差异相对于模型规模变化时的困惑度变化来说是微不足道的（请参见下一部分的表4）。gMLP学习到的空间投影权重在图4中进行了可视化。

表 3：Transformer 基线模型和四个版本 gMLPs 的 MLM 验证困惑度。\(f\) 指式（2）中输入归一化后的空间线性投影。MLP-Mixer基线模型有\(L=24\)层，\(d_{model}=768\)，\(d_{spatial}=384\)，\(d_{ffn}=3072\)。每个gMLP模型有L=36层，\(d_{model}=512\)，\(d_{ffn}=3072\)。Mixer和gMLP不使用位置编码。

图4：在MLM任务中学习的gMLP空间过滤器的可视化。对于模型中的每一层，我们绘制了W中与序列中间标记相关的行。每个子图的X轴长度为128，等于序列中的标记数。学习到的滤波器看起来很平滑，并且有几种类型：前向滤波器（例如，第2行中的第1个）、后向滤波器（例如，第2行中的第5个）和双向滤波器（例如，最后一行中的倒数第2个）。

tu-4

4.2 案例研究：gMLP在模型规模增大时的行为

在表4中，我们研究了BERT中Transformer和gMLP的规模性能，随着模型容量的增长。具体来说，我们通过{0.5，1，2，4}×的因子来扩展这些模型的深度，并报告它们在验证集上的预训练MLM困惑度以及在GLUE [41]的两个任务的微调结果。注意，每个独立的Transformer层实际上是两个连续的块：一个用于自注意力，一个用于FFN。在下表中，我们使用12 + 12的表示方法来表示Transformer基线中的12个自注意力块和12个FFN块。

表 4：模型容量增加时的预训练和开发集微调结果。我们使用表 3 中表现最好的 Transformers 中使用了相对位置编码方案。

biao-4

上述结果表明，足够深的gMLP能够与相同容量的Transformer相媲美甚至超越其困惑度。此外，两种体系结构的困惑度和参数之间的关系大致遵循幂律（图5左侧）。这意味着最初观察到的基于Transformer的语言模型的经验缩放规律可能在不同的模型家族中广泛适用。

tu-5

图 5：展示了在困惑度和微调准确率方面的缩放特性。图表显示，在预训练中，gMLP和Transformer在优化困惑度方面同样出色。在微调中，尽管存在任务特定的差异，但两种模型家族表现出相当的可扩展性。

根据表4，还可以观察到一个有趣的现象，不同模型家族在微调方面的预训练困惑度并不相等。虽然gMLP在SST-2任务上表现优于Transformer，但在MNLI任务上表现较差。这些结果表明，NLP任务的微调性能不仅取决于困惑度，还取决于架构中的归纳偏差。图5显示，尽管预训练和微调之间存在架构特定的差异，gMLP和Transformer在微调任务上的可扩展性（斜率）是可比较的。这意味着可以通过增加模型容量来弥补差距。换句话说，这些结果表明，相对于下游指标，模型的可扩展性与是否存在自注意力是独立的。

4.3消融研究：BERT微调中微小注意力的作用

到目前为止，我们发现自注意力不是实现强大的MLM困惑度或可扩展性所必需的组成部分。同时，我们还发现在NLP微调任务中，gMLP的迁移效果不如Transformer（见表4）。我们特别注意到我们的MLP-like模型在SST-2任务上具有优势，但在MNLI任务上表现较差——前者是单句子任务，而后者涉及句子对（前提和假设）[43]。我们怀疑在微调过程中，自注意力的作用与跨句子对齐有关。 为了分离出自注意力的影响，我们在gMLP的门控函数上添加了一个微小的自注意力块的混合模型（图6）。由于gMLP本身已经能够捕捉空间关系，我们假设这个额外的自注意力模块不必太重，而它的存在与其容量相比更重要。在我们的实验中，典型的微小注意力模块只有一个头，大小为64，远小于Transformer中典型的多头自注意力模块，后者有12个头和总大小为768。以下，我们将这种混合模型，即带有微小自注意力的gMLP，称为aMLP（“a”代表attention）。

tu-6

图 6：带有微小自我注意模块的混合空间门控单元。我们使用 gMLP 模块的归一化输入（输入归一化后、通道扩展前的端点）作为微小自注意力的输入。对于 SGU，由于通道分割，我们使用$ d_{out} = d_{ffn}/2$。

在图 7 中，我们通过 MLM 模型的预训练困惑度和微调指标之间的校准图，研究了这些模型的可移植性。评估的模型包括 BERTbase、gMLP 及其混合版本 aMLP（64-d 单头自注意）（图 6）。数据点是通过改变模型深度{0.5, 1, 2}×或数据{1, 2, 4, 8}×收集的。可以看出，无论是否存在 自注意力，gMLPs 在 SST-2 上的传输效果都优于 Transformers。在附录 D 中，我们将 aMLP 中的微小自我注意模块与 MNLI 示例进行了可视化对比，显示它们主要负责句对之间的对齐。

tu-7

图 7：从 MLM 预训练perpexity 到 GLUE 上微调精度的可转移性。aMLP 是指 gMLP，如图 6 所示，增强了 64 个单头自我注意。相比之下，BERT 基线中的每个自我注意模块包含 12 个头，总大小为 768。

在图 8 中，我们将这三种模型的缩放特性汇总在一起，显示 aMLP（gMLP + 微小注意力）在两项微调任务中的表现始终优于 Transformer。

tu-8

图 8：比较Transformer、gMLPs 和 aMLPs 的缩放特性（64-d，单头关注）。结果使用第 4.2 节中的相同设置得出。

4.4 BERT设置下的MLM的主要结果

下面我们展示了在完整的BERT设置下的预训练和微调结果。与消融和案例研究不同，这里我们使用完整的英文C4数据集，并采用常见的MLM设置，批大小为256，最大长度为512，训练步数为1M。为了公平比较，我们调整了gMLPs的深度和宽度，以确保与Transformer基线具有可比的模型容量。模型的规格见表5，超参数详见附录A.2。对于微调，我们报告了GLUE [41]中SST-2和MNLI的开发集性能，每个结果条目是通过对五次独立运行结果取中值得到的。此外，我们还报告了在SQuAD [44, 45]上的微调结果，以测试模型在推理较长上下文方面的能力。

结果如表6所示。与我们在4.1节和4.2节中的发现一致，gMLPs在困惑度方面与Transformers具有竞争力，尤其是在较大规模的设置中。关于微调结果，有几点观察结果：

首先，在 gMLP 性能低于 Transformers 的微调任务上，随着模型容量的增加，性能差距趋于缩小。例如，虽然 gMLP 在 SQuAD-v2.0 基准规模上的表现差了 \(8.5\%\)，但在更大的规模上，相对于基准的性能差距降至 \(2.7\%\)。值得注意的是，我们的 gMLPlarge 在 SQuAD-v1.1 上实现了 \(89.5\%\) 的 F1，而没有任何自我关注或动态空间参数化[28]，远高于 Devlin 等人[2]报告的 \(BERT_{base}\) 的 \(88.5\%\)，与 BERTlarge 的原始结果仅相差 \(1.4\%\)。我们还通过进一步扩展 gMLP 增加了一个数据点。由此产生的模型 gMLPxlarge 在 SQuAD-v2.0 （一个涉及问题-答案对的困难任务）上的表现优于 \(BERT_{large}\)，而没有使用任何自我关注。虽然由于模型大小不同，这种比较并不公平，但它证明了在具有挑战性的 NLP 任务中，类 MLP 模型可以与 Transformers 竞争。

此外，我们展示了通过融合大小为64或128的微小单头自注意力，gMLPs能够优于具有类似容量的Transformers，有时优势显著。例如，我们的混合模型aMLPlarge在SQuAD-v2.0上的F1分数比Transformers高出4.4%。这些结果表明，Transformers中的多头自注意力的容量很大程度上是多余的，而gMLPs中的空间门控单元可以捕捉到大多数其功能。这些结果还表明，gMLPs的空间门控单元和微小的自注意力的归纳偏差是互补的。尽管随着计算量的增加 ，体系结构的归纳偏差的好处可能会消失，但微小的自注意力确实提高了gMLPs在我们研究中的实际价值。

表 5：完整 BERT 设置中的模型规格

biao-5

表 6：预训练困惑度和微调的开发集结果。 “ours”表示使用我们的设置训练的模型。我们报告 SST-2 和 MNLI 的准确度，以及 SQuAD v1.1/2.0 的 F1 分数。

biao-6

五、 结论

自从Vaswani等人的开创性工作[1]以来，Transformers已经广泛应用于NLP和计算机视觉领域。这种应用使得在NLP领域取得了许多令人印象深刻的结果。然而，到目前为止，仍不清楚是什么赋予了Transformers如此成功的力量：是Transformers的前馈特性，还是Transformers中的多头自注意力层？

我们的工作提出了一种更简单的替代方案来取代Transformers中的多头自注意力层。我们展示了gMLPs，一种带有门控机制的简单MLP变体，可以在BERT的预训练困惑度和ViT的准确性方面与Transformers竞争。在数据和计算量增加的情况下，gMLPs的可扩展性也与Transformers相当。在BERT的微调任务中，我们发现gMLPs在一些具有挑战性的任务（如SQuAD）中可以在没有自注意力的情况下取得令人满意的结果，并且在某些情况下明显优于Transformers。我们还发现，在需要跨句子对齐的下游任务中，Transformer中的多头自注意力的归纳偏差非常有用。然而，在这些情况下，通过使gMLP变得更大可以缩小与Transformers之间的差距。从更实际的角度来看，将一个小的单头自注意力融合到gMLP中可以实现更好的架构，而无需增加模型的大小。