Looped mixture of experts model improves performance with uneven compute use
LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation
Machine Learning
Summary
Large AI models often try to reuse parts of themselves to save computing power. This paper shows that when reusing the same part twice, doing less work the first time and more the second time can improve results. The authors created LatentLoop, which compresses the input during the first reuse and fully processes it the second time, while adjusting how much each reuse updates the data. This approach led to better performance and efficiency across multiple model sizes and tests.
What this means in practice
- •For machine learning engineers: Build large language models that use computation unevenly to save resources while maintaining or improving accuracy.
- •For ai system architects: Design efficient looped MoE architectures by compressing partial computations and reallocating savings to increase model capacity.
Authors
Zefeng He, Yunkai Chai, Xiaoyang Chen, Tong Zhu, Zongliang Wu, Yun Luo, Yafu Li, Yu Cheng
Abstract
Looped Mixture-of-Experts (MoE) models combine parameter reuse with sparse computation to improve performance under matched budgets. Prior work demonstrates that different loop visits can play different roles even when they share parameters. To probe the contributions of different loop visits, we learn separate residual coefficients in a model that applies a shared block twice, observing a smaller scale in the first visit. Fixed coefficient experiments show that reducing the first visit's scale preserves or improves performance, whereas the same reductions to the second hurt it. This asymmetry motivates us to explore asymmetric compute allocation. We propose \textbf{LatentLoop}, which compresses the first visit's sequence, retains full resolution in the second, and reallocates saved computation to increase the hidden size. We further learn how strongly the first visit updates the full sequence, and share MoE weights across visits while using separate attention parameters to accommodate their different sequence resolutions. Across 15 configurations spanning 1.3B--7B parameters, with the largest budget exceeding $10^{21}$ FLOPs, LatentLoop consistently lowers validation NLL. At a fitted NLL of 1.988, predicted compute savings are 19.0\% and 7.9\% relative to vanilla MoE and the looped baseline. It also improves average scores on nine benchmarks at all sizes in 0-shot and 5-shot settings. These results suggest asymmetric compute allocation is a promising direction for looped MoE models.