LoopMoE: Unifying Iterative Computation with Mixture-of-Experts for Language Modeling

2026-06-03Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors created LoopMoE, a language model that combines two ways to grow models: adding more parameters and repeating calculations multiple times. They designed a system to fairly compare LoopMoE with regular MoE models by keeping the total computing effort and parameter use the same. LoopMoE showed better performance on most tests at both medium and large model sizes. This work shows how mixing sparsity and repeated computations can improve language models.

Mixture-of-Experts (MoE)looped architecturesparameter capacityeffective depthweight sharingsparse routingIterAdaLNattention mechanismfeed-forward network (FFN)language models
Authors
Wenkai Chen, Tianshu Li, Wenyong Huang, Yichun Yin, Lifeng Shang, Chengwei Qin
Abstract
Mixture-of-Experts (MoE) and looped architectures scale models along two orthogonal axes, namely parameter capacity and effective depth. However, mainstream looped architectures rely on dense backbones that couple parameter count with per-token FLOPs, which makes it impossible to isolate the effect of iterative computation under matched budgets. To this end, we present LoopMoE, a looped MoE language model that integrates sparse routing with iterative weight-shared computation through two designs. The first is IterAdaLN, which resolves weight-sharing symmetry via a modulation signal jointly conditioned on the iteration index and the per-token hidden state. The second is a capacity-balancing strategy that recovers the attention-to-FFN active parameter ratio of well-tuned non-looped references. Together, these designs enable the first strictly controlled, head-to-head evaluation of a looped MoE against a Vanilla MoE under identical total parameters, per-token FLOPs, and active sublayer ratios. At the 3B scale, LoopMoE outperforms the Vanilla MoE on 8 of 9 downstream benchmarks with an average improvement exceeding 1 point. At the 9B scale, LoopMoE continues to outperform the matched Vanilla MoE, indicating that the architectural gain persists at larger scale. Our work establishes a controlled synthesis of sparsity and recurrence, and suggests a promising direction for looped language models.