Press "Enter" to skip to content

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

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

作为基于注意力模型的替代方案,纯 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 基线中具有最佳性能,并且需要的时间最少。

 

 

 

扩展

 

为了测试模型的可扩展性,该研究增加了 2.0 TFLOPs 的模型大小训练。表 4(底部)总结了结果。

 

与表 4(顶部)中的模型相比,该研究扩大了所有模型,将嵌入从 1024 更改为 2048,并将隐藏维度从 4096 调整为 8192,如表 5 所示。该研究还增加了预训练数据大小,如表 2 所示。

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。