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

2026-06-03

AI Summary

This work addresses the challenge of fairly evaluating iterative computation in recurrent architectures, which has been confounded by reliance on dense backbones that prevent isolation of iteration efficacy under matched parameter counts and per-token compute. The authors propose LoopMoE, the first architecture to integrate sparse Mixture-of-Experts (MoE) with iterative weight-sharing computation, enabling rigorous comparison under strictly aligned total parameters, per-token FLOPs, and activated sublayer ratios. Key innovations include IterAdaLN to break weight-sharing symmetry, a capacity-balancing mechanism to preserve attention-to-FFN activation proportions, and a sparse routing strategy. Experiments show that at 3B scale, LoopMoE outperforms standard MoE on 8 of 9 downstream tasks, with an average gain exceeding 1 point; this advantage persists at 9B scale, demonstrating the scalability of the architectural gains.

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.

Problem

• Mixture-of-Experts
• looped architectures
• iterative computation
• parameter-FLOPs coupling
• language modeling

Innovation

• LoopMoE
• Mixture-of-Experts
• Iterative Computation
• Sparse Routing
• Weight Sharing