Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts
Looped Transformers reuse shared blocks to scale depth without adding parameters, but MoE variants degrade beyond two loops. This paper diagnoses two bottlenecks—hidden-state drift and expert selection collapse—and proposes LOOM, combining residual scaling, embedding re-injection, per-loop routers, and a Looping Residual. It matters because it attempts to establish recurrent depth as a genuine, compute-efficient scaling axis for sparse LLMs.
LOOM is a coherent engineering recipe backed by thorough ablations. The central design principle—"each loop should contribute new computation while keeping the recurrent state stable" (Section 1)—is well executed, and the paper convincingly shows that no single ingredient suffices. The 1.7B model "peaks at 9 loops, reducing perplexity from 9.62 to 7.77 and improving average zero-shot accuracy from 42.4\% to 47.7\%" (Abstract). However, the empirical validation stops at 1.7B parameters and 60B tokens on sequences of length 1{,}024, leaving open whether these mechanisms survive frontier-scale training and long-context regimes.
The ablation studies are the strongest evidence. Table 3 shows that native looping and partial stabilizations (residual scaling only or embedding re-injection only) suffer "unrecoverable loss spikes" at deeper recurrence, while the full recipe scales smoothly to 12 loops. Component ablation (Table 5) confirms that every proposed mechanism matters: removing residual scaling raises perplexity from 18.62 to 24.63, and dropping embedding re-injection raises it to 26.05 at $H=9$. The empirical demonstration of "expert selection collapse" (Figure 4) is also compelling: shared routers yield near-identical expert utilization across loops, whereas per-loop routers produce measurably diverse load distributions.
Evaluation scale and coverage are limited. The largest experiment uses only $\sim$1.7B parameters and 60B tokens with a 1{,}024-token sequence length, which is far smaller than current production MoE训练和长上下文设置. The iso-FLOP protocol in Table 2 is only approximately matched: the anchor is $f_1=84$, yet the 5-loop configuration has $f=90$, so the best-reported depth is not strictly compute-equivalent to the baseline. All loop-specific coefficients are fixed across scales: "Loop-specific coefficients are fixed across scales ($\lambda=0.5$, $\beta=0.5$, $K=3$); width and expert dimensions scale with model size, while these coefficients do not" (Section 4.1). The absence of tuning or scaling laws for these hyperparameters raises doubts about optimality at larger sizes. Finally, downstream evaluation is restricted to lightweight zero-shot commonsense tasks; no coding, mathematical reasoning, or long-context benchmarks are reported, leaving unclear whether deeper looping helps or hurts on harder reasoning workloads.
The internal evidence strongly supports the claim that LOOM stabilizes deep recurrence. Comparisons to disabled baselines are fair and internally consistent: "Native looping is insufficient... 'No tech.' loops the unmodified backbone without additional techniques and rapidly becomes unstable" and "Residual scaling alone does not ensure stability... Embedding re-injection alone cannot sustain deep recurrence" (Section 4.2.2). However, external comparisons to prior looped MoE work such as SMELT or LoopMoE are sparse and confounded by differing scales and training budgets. The paper also does not compare against a non-looped MoE with matched effective depth and total parameters, leaving ambiguous whether the gains come from recurrence per se or simply from additional compute capacity enabled by the stabilization tricks.
Reproduction is theoretically feasible but practically demanding. The authors release code at a linked GitHub repository ("Code: [LOOM](https://github.com/hed-ucas/LOOM)"), use the public FineWeb-Edu dataset, and disclose architecture dimensions, optimizer settings, and fixed coefficients in Table 1. The training procedure is summarized in Algorithm 1, including the segmented backpropagation logic. The main barriers for independent reproduction are computational: training a $\sim$1.7B-parameter model on 60B tokens requires substantial NPU/GPU resources. Additionally, while the algorithm is described, the low-level implementation of the loop-specific routing and Looping Residual together with segmented backprop and gradient detaching may hide subtle stability-critical details not fully captured in the pseudocode.
Looped Transformers introduce recurrent depth as a new scaling axis for LLMs: by repeatedly applying shared Transformer blocks, they increase effective depth without increasing parameter count. However, the benefits of looping remain unclear for large MoE LLMs under FLOPs-matched comparisons. The main reason is that the gains from additional iterations diminish quickly and can even turn into degradation, so the extra FLOPs spent on looping yield little substantial improvement. Consequently, prior work typically settles on two loops. We identify two main obstacles to scaling looped MoE. First, looping inherits and amplifies the curse of depth: hidden-state variance grows with each iteration as residual updates accumulate, which destabilizes deep recurrence and causes representations to drift. Second, looped MoE suffers from expert selection collapse: routers repeatedly select the same experts across loops, so extra iterations add computation without adding computational diversity. Guided by this diagnosis, we propose LOOM, built on a single principle: each loop should contribute new computation while keeping the recurrent state stable. LOOM stabilizes recurrence by scaling residual updates to bound variance growth and re-injecting the input embedding at every loop, and diversifies it through per-loop routers that engage different experts and a Looping Residual that carries earlier outputs forward. Experiments across 100M-1.7B models show stable scaling to 9-12 loops. Under near-iso-FLOP, the 700M model performs best at 5 loops, reducing perplexity from 18.36 to 16.54 and improving average zero-shot accuracy from 38.84% to 39.53% over the non-looped baseline. Without FLOP matching, the 1.7B model trained on 60B tokens peaks at 9 loops, reducing perplexity from 9.62 to 7.77 and improving average zero-shot accuracy from 42.4% to 47.7%. Code is available https://github.com/hed-ucas/LOOM.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.