Efficient Language Modeling with Sparse all-MLP

TL;DR

This work tackles expressiveness gaps in all-MLP language models by introducing sparsely activated Mixture-of-Experts (sMLP) that route computations in both feature and token dimensions. It presents two routing strategies—Deterministic Routing and Partial Prediction—implemented via tMoE and sMoE modules to avoid autoregressive leakage while enabling conditional computation. Empirically, sMLP delivers lower perplexity and up to 2x faster pretraining than Transformer-based MoEs and dense transformers, scales to large parameter counts, and achieves superior zero-shot in-context learning on six tasks. Overall, sparse all-MLP emerges as a competitive alternative for large-scale NLP pretraining with practical efficiency and downstream benefits.

Abstract

All-MLP architectures have attracted increasing interest as an alternative to attention-based models. In NLP, recent work like gMLP shows that all-MLPs can match Transformers in language modeling, but still lag behind in downstream tasks. In this work, we analyze the limitations of MLPs in expressiveness, and propose sparsely activated MLPs with mixture-of-experts (MoEs) in both feature and input (token) dimensions. Such sparse all-MLPs significantly increase model capacity and expressiveness while keeping the compute constant. We address critical challenges in incorporating conditional computation with two routing strategies. The proposed sparse all-MLP improves language modeling perplexity and obtains up to 2$\times$ improvement in training efficiency compared to both Transformer-based MoEs (GShard, Switch Transformer, Base Layers and HASH Layers) as well as dense Transformers and all-MLPs. Finally, we evaluate its zero-shot in-context learning performance on six downstream tasks, and find that it surpasses Transformer-based MoEs and dense Transformers.

Figures (8)

• Figure 1: The proposed sparse all-MLP architecture (sMLP) achieves 2$\times$ training efficiency improvement compared to state-of-the-art sparse Transformer-based MoEs: Gshardgshard, Switch Transformer switch, Base Layers baselayer, HASH Layers hashlayer. Full comparison is provided in Fig. \ref{['fig:model_comparision']}.
• Figure 2: Illustration of sMLP Model Architecture with $N_1$ Dense Blocks (gMLP layers) and $N_2$ Sparse Blocks. The arrangement of these blocks will be discussed in Section \ref{['sec:expe_setup']}. Each sparse block contains a tMoE module and a sMoE module. The tMoE module sparsely activated FFN (feed-forward) and the sMoE module replaces the self attention in Transformers with sparse token-wise SL (spatial linear projection) operations.
• Figure 3: Left: tMoE Gating: sending 4 tokens to 3 experts; Right: sMoE Gating: sending 4 hidden vectors to 3 experts.
• Figure 4: Comparison with Dense Models: We compare our model (the orange dashed line) with dense models (solid lines), including both Transformer transformer2017 and gMLP gmlp, with different heads. Our model (sMLP) only has one head activated on each device.
• Figure 5: Top: Valid Perplexity v.s. Training Steps; Bottom: Valid Perplexity v.s. Training Time. We compare our two methods (sMLP-deterministic and sMLP-partial) with sparse Transformer MoEs with 0.8T FLOPs.