纯MLP在下游任务上欠佳？Meta AI等提出稀疏MLP，超越transformer

2022 年 4 月 16 日 机器之心

机器之心报道

机器之心编辑部

来自 Meta AI 和纽约州立大学布法罗分校的研究者分析了 MLP 在表达能力方面的局限性，并提出了在特征和输入（token）维度上带有混合专家系统（MoE）的稀疏激活 MLP。

作为基于注意力模型的替代方案，纯 MLP 架构吸引了越来越多的关注。在 NLP 中，gMLP 等近期工作表明，纯 MLP 在语言建模方面可以达到与 transformer 相当的性能，但在下游任务中弱于 transformer。

来自 Meta AI 和纽约州立大学布法罗分校的研究者分析了 MLP 在表达能力方面的局限性，并提出了在特征和输入（token）维度上带有混合专家系统（MoE）的稀疏激活 MLP。这种稀疏的纯 MLP 显著提高了模型容量和表达能力，同时保持计算不变。该研究解决了将条件计算与两种路由策略结合起来的关键挑战。

论文地址：https://arxiv.org/pdf/2203.06850.pdf

与基于 transformer 的 MoE、密集 Transformer 和纯 MLP 相比，该研究提出的稀疏纯 MLP（sMLP）改进了语言建模的困惑度并获得了高达 2 倍的训练效率提升。最后，研究者在六个下游任务上评估了稀疏纯 MLP 的零样本上下文学习性能，发现它超过了基于 transformer 的 MoE 和密集的 transformer。

方法

稀疏激活的纯 MLP

sMLP 的整体架构如下图 2 所示，包含 N_1 个密集块和 N_2 个稀疏块。N_1 和 N_2 都是超参数。每个稀疏块包含两个模块：

tMoE 模块：该研究采用 Base Layers 的 MoE (Lewis et al., 2021) 来替换密集 transformer 中的 FFN 模块 (Vaswani et al., 2017b)；
sMoE 模块：该研究设计了 sMoE 模块来替代 transformer 中的自注意力模块（Vaswani et al., 2017b）和 gMLP 中的空间门控单元（SGU，Liu et al., 2021a）。

tMoE 模块和 sMoE 模块都包含两个元素：

专家模块

专家模块负责处理输入。对于 tMoE 模块，每个专家都包含一个 FFN，如上图 2 所示。对于 sMoE 模块，每个专家都包含空间门控单元，如下图 6（右）所示。

门控函数

该模块决定哪个专家应该处理输入的每个部分，因此需要设计一种独特的路由方法来将 MoE 结构扩展到特征维度。

图 3（左）显示了现有基于 transformer 的 MoE 的门控函数示例（Lepikhin et al., 2020；Fedus et al., 2021；Lewis et al., 2021；Roller et al., 2021）。x_ij 表示 i_th token 中 j_th 隐藏维度的值。

如下等式 (3) 所示：

tMoE 使用由

参数化的等式（3）中描述的学习门控函数将这 4 个 token 发送给 FFN 层的 3 个专家。与已有的一些 MoE 不同，在稀疏纯 MLP 架构中，该研究提出沿隐藏维度对隐藏表示进行分块，并将分块向量发送给不同的专家，如图 3（右）所示。

与这些现有的 MoE 不同，在稀疏的全 MLP 架构中，该研究提出沿隐藏维度对隐藏表示进行分块，并将分块向量发送给不同的专家，如图 3（右）所示。

特征空间中的路由

与路由 token 相比，路由隐藏维度在自回归模型中面临着一个独特的挑战，如果简单地预测未来的 token，信息会泄漏。此外，与具有 selfattention 的基于 Transformers 的 MoE 不同，此处不能直接应用适当的掩码来防止信息泄露，因此不能在基于 transformer 的 MoE 中采用现有的路由方法进行语言建模。该研究比较了以下两种解决方案：确定性路由（deterministic routing）和部分预测（partial prediction）。

实验及结果

token 操作比较

该研究将 sMLP 模型与两个密集模型进行比较：Transformer (Vaswani et al., 2017b) 和 gMLP (Liu et al., 2021a)。基于全 MLP 和基于 transformer 的模型之间的主要区别在于 token 操作。该研究比较了这三种 token-wise 操作：Transformers 中的 self-attention 模块、gMLP 中的 Spatial Gating Unit 和 sMLP 模型中的 sMoE 模块。表 3 比较了三种 token 操作及其各自的头部机制：

下图 4 将模型与不同头数的密集模型进行了比较。Transformer 模型极大地受益于多头机制。然而，gMLP 模型虽然增加了参数量，但并没有通过多头机制提高性能。sMLP 模型也可以看作是 gMLP 的一种多头解决方案，显著提高了基于 MLP 模型的性能，并且优于 transformer 模型。

稀疏 MLP 的结果

下图 5 给出了质量（有效困惑度）和训练效率，通过训练步骤数（顶部）和训练时间（底部）来衡量。研究者发现，具有两种路由策略变体的 sMLP 优于具有大致相同数量的 FLOP 的最先进的基于 Transformer 的 MoE 模型。

下表 4 总结了主要实验中的详细比较结果。研究者将所有模型的 FLOPs 控制为约 0.8T。除了模型层数不同，它们的嵌入维数为 1024，隐藏维数为 4096。可以看到，sMLP 模型在 25k 训练步骤时实现了最好的泛化，同时实现了最高的训练速度。HASH 层在所有 Transformer 基线中具有最佳性能，并且需要的时间最少。