🤖 AI Summary
This work addresses the scalability limitations of recurrent language models, which stem from the trade-off between memory efficiency and performance. The authors propose Looped-MoE, a novel architecture that integrates recurrence with a Mixture-of-Experts (MoE) mechanism, dynamically activating distinct experts within shared layers based on recurrent states to enhance model expressivity. They demonstrate for the first time that sparse expert activation is key to improving the scalability of recurrent models. Furthermore, the study reveals that recurrent boundaries naturally serve as high-quality early-exit points, substantially reducing inference costs. Experimental results show that Looped-MoE outperforms standard Transformers in large-scale settings while significantly lowering memory and computational requirements, all without compromising generation quality.
📝 Abstract
Looped language models repeat a set of transformer layers through depth, reducing memory costs and providing natural early-exit points at loop boundaries. However, looped models do not scale as favorably as standard transformers with unique layers. We compare standard and Mixture-of-Experts (MoE) transformers, with and without looping, and find two main results. First, we find Looped-MoE models scale better than the standard baseline while dense looped models do not. We trace this to routing divergence between loops: in Looped-MoE models, different experts are activated on each pass through the same shared layers, recovering expressivity without additional parameters. Our second finding is that looped models have better compute-quality trade-offs with early exits than standard models. Because each loop ends with the same layers that produce the final output, loop boundaries are superior exit points, as confirmed by earlier output convergence at these points. In sum, we provide a clear direction for scaling looped models: a Looped-MoE model with early exits can not only beat standard transformers at scale, but also enable significant memory and inference savings with minimal degradation in quality.