Reproduce “LoopMoE: Unifying Iterative Computation with Mixture-of-Experts for Language Modeling” on OpenResearch

About this paper

Wenkai Chen, Tianshu Li, Wenyong Huang, Yichun Yin, Lifeng Shang, Chengwei Qin

The LoopMoE architecture unifies sparse Mixture-of-Experts with iterative computation, outperforming non-looped MoE baselines across language modeling and reasoning tasks and exhibiting positive scaling trends. It addresses challenges of weight-sharing symmetry and imbalanced parameter allocation through IterAdaLN and a capacity balancing strategy.

Key insights

• Iterative computation itself significantly benefits reasoning and mathematics tasks, highlighting the value of increased computational depth even within existing parameter budgets.
• Per-token and per-iteration adaptive modulation is crucial for enabling effective progressive transformation of token states in weight-shared, looped MoE architectures.
• Maintaining an optimal active attention-to-FFN parameter ratio is essential for performance in looped MoEs, as structural distortions can otherwise lead to suboptimal capacity allocation.
• The model exhibits distinct phases in iterative routing dynamics (entry, recurrent core, exit) and an emergent back-loaded update schedule, concentrating effective computation in later attention branches.

Summary by alphaXiv.