Looped mixture of experts improves model efficiency by reshaping experts
How to Loop MoE: Flatten the Experts, Untie the Attention
Machine LearningArtificial IntelligenceComputation and Language
Summary
Large AI models sometimes reuse parts of themselves repeatedly to do more with less. This paper shows a way to reshape and reorganize these reused parts, called experts, so the system can choose from a bigger and better pool each time. The authors found that this approach leads to better learning and more balanced use of the parts, without increasing the model size or computation needed. This means AI models can perform better by simply looping through a clever arrangement of their own components.
What this means in practice
- •For machine learning engineers: Improve efficiency of large language models by deploying looped MoE architectures that better use experts without increasing compute.
- •For cloud ai infrastructure teams: Optimize resource allocation on AI hardware by implementing flattened and looped MoE layers for faster training and inference at fixed cost.
Authors
Shouren Wang, Chuang Ma, Mohsen Hariri, Debargha Ganguly, Wang Yang, Xiaoqing Tong, Qianying Liu, Xiaotian Han, Vipin Chaudhary
Abstract
Looped Transformers reuse one block of layers several times: by spending extra computation they push a model of fixed size further, and so use its parameters more fully; while sparse mixture-of-experts (MoE) models activate only a few of many experts for each token. Looped MoE bridges these two design philosophies and gives MoE models new potential for better expert usage, but it raises a question: how to loop a MoE? We answer it with Foil. With the expert parameters and the expert compute per token held fixed, Foil (1) flattens the experts, halving the expert layers, doubling the experts per layer and doubling the passes, so that every routing decision chooses from a larger pool, and (2) unties the attention, giving each pass its own attention parameters while the experts and routers stay shared. Experiments show that Foil clearly outperforms the unflattened looped baseline: at 20B tokens every Foil model has lower pretraining loss than the baseline; at 100B tokens the loss improves monotonically with the degree of flattening, the most flattened Foil ending 0.012 nat below the baseline at equal parameters and compute, with downstream accuracy on par or better; untying the attention also yields more balanced and more confident routing at equal shape. Our ablations analyse why Foil works and turn the findings into design guidance for looped MoE: the returns of looping and of widening the expert layers amplify each other, routing confidence tracks healthy expert use better than load balance, and a sparse looped MoE should therefore use more experts per layer and more passes. Code and configurations are available at https://github.com/SR-A-W/how-to-loop-moe.