sMLP:稀疏全mlp进行高效语言建模

这是一篇2022由纽约州立大学布法罗分校和Meta AI发布的论文，它主要的观点如下：

具有专家混合(MoEs)的稀疏激活mlp在保持计算常数的同时显着提高了模型容量和表达能力。此外gMLP表明，所有mlp都可以在语言建模方面与transformer相匹配，但在下游任务方面仍然落后。所以论文提出了sMLP，通过设计确定性路由和部分预测来解决下游任务方面的问题。

sMLP

sMLP模型包含N1个密集块和N2个稀疏块。在每个稀疏块中，包含两个模块:

tMoE模块:采用Base Layers (Lewis et al.， 2021)中的MoE代替transformer中的FFN模块。

sMoE模块:这是轮文提出的sMoE模块，目的是为了取代gMLP的自注意模块和空间门控单元。

稀疏激活 all-MLP

左:一个来自现有基于transformer的moe (tMoE)的门控功能示例。tMoE使用学习到的门控函数将这四个令牌发送给FFN层的3位专家。

右图:在稀疏的all-MLP体系结构中，论文提出沿隐藏维度对隐藏表示进行分块，并将分块向量发送给不同的专家。

特征空间中的路由

与路由令牌相比，路由隐藏维度在自回归模型中面临着独特的挑战，如果只预测未来的令牌，信息就会泄露。

所以采用确定性路由，将向量在隐藏维度上进行分块，并将隐藏向量确定性地发送给专家。并且令牌X1的前20%用于决定路由，令牌X2的其余80%用于预测。

训练方法是不在整个序列长度T上训练语言模型，而是训练它来预测X2。而X1用来学习门控权值Wr。

结果

sMLP显著地提高了全基于mlp的模型的性能，并且也优于Transformer模型。

sMLP在25k训练步长时达到了最好的泛化效果，同时达到了最高的训练速度。

对于2.0 TFLOPs训练，模型大小增加。嵌入从1024调整为2048，隐藏维数从4096调整为8192。

sMLP仍然优于Switch Transformer，并且后者有更多的flop。

sMLP在平均精度方面优于所有稀疏Transformers 。显著的改进来自常识推理任务，如COPA、StoryCloze和HellaSwag。

论文地址：

Efficient Language Modeling with Sparse all-MLP

https://arxiv.org/abs/2203.06850

作者：Sik-Ho Tsang