Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts

Summary (Overview)

• Problem: Looped Transformers introduce recurrent depth as a scaling axis for LLMs, but looped MoE models fail to scale beyond 2 loops—additional iterations yield diminishing returns or degradation under FLOPs-matched comparisons.
• Key Diagnosis: The authors identify two main obstacles: (1) the curse of depth—hidden-state variance grows with each iteration as residual updates accumulate, destabilizing deep recurrence; and (2) expert selection collapse—shared routers repeatedly select the same experts across loops, adding computation without adding computational diversity.
• Proposed Solution: LOOM (Looping Mixture-of-Experts), built on the principle that each loop should contribute new computation while keeping the recurrent state stable. It combines residual scaling, embedding re-injection, per-loop routers, and a Looping Residual mechanism.
• Results: LOOM scales to 9–12 loops across 100M–1.7B parameter models. Under near-iso-FLOP conditions, a 700M model performs best at 5 loops (perplexity: 18.36→16.54; avg. zero-shot accuracy: 38.84%→39.53%). Without FLOP matching, a 1.7B model scales to 9 loops (perplexity: 9.62→7.77; avg. accuracy: 42.4%→47.7%).

---

Introduction and Theoretical Foundation

Background

Looped Transformers scale computational depth by repeatedly applying shared Transformer blocks, increasing effective depth without increasing parameter count (Dehghani et al., 2018; Giannou et al., 2023; Saunshi et al., 2025; Geiping et al., 2025). This introduces a scaling axis orthogonal to model size and training data. MoE architectures offer a natural pathway for extending looped computation by activating only a sparse subset of experts per token, providing additional computation capacity under a fixed budget.

The Two Obstacles

1. The Curse of Depth: Repeatedly reusing shared Transformer blocks causes residual accumulation and activation variance growth, making hidden states drift and destabilizing optimization in deep recurrence.

2. Expert Selection Collapse: Shared routers produce nearly invariant expert assignments across loops, so later iterations repeatedly activate similar experts instead of refining earlier representations. This means increasing loop depth adds computation without adding useful computation.

Central Research Question

Is there an effective way to improve cross-loop information flow in looped MoE, so that recurrence can scale well beyond two iterations?

---

Methodology

Architecture Overview

LOOM repeatedly applies an M-layer, pre-norm MoE decoder for H passes. Attention and expert weights are shared across loops, while each loop has its own routers. The effective depth is $MH$; increasing H adds router parameters but does not replicate attention or expert weights.

3.1 Stabilizing the Recurrent State

Residual Scaling: LOOM scales residual updates with a factor that accounts for both loop count and physical depth:

$$\gamma = \frac{\lambda}{H \sqrt{M}}, \quad \widetilde{h}_{\ell}^{t} = h_{\ell-1}^{t} + \gamma \operatorname{Attn}_{\ell}(\mathrm{RMS}(h_{\ell-1}^{t})), \quad h_{\ell}^{t} = \widetilde{h}_{\ell}^{t} + \gamma \operatorname{RMS}(\mathrm{MoE}_{\ell}^{t}(\mathrm{RMS}(\widetilde{h}_{\ell}^{t})))$$

The $1/H$ factor limits residual accumulation across loops, while $1/\sqrt{M}$ controls variance growth across physical layers. LOOM fixes $\lambda = 0.5$ as a non-learnable hyperparameter.

Embedding Re-injection: To preserve access to the original token representation and prevent representation drift, LOOM re-injects the input embedding at the beginning of each loop $t \geq 2$:

$$g_t = \frac{\lambda}{t\sqrt{M}}, \quad h_0^t = (1 - g_t)h_M^{t-1} + g_t x$$

The decreasing coefficient $g_t$ places progressively greater weight on the recurrent representation while convex mixing regulates the embedding contribution.

3.2 Diversifying Computation Across Loops

Loop-Specific Routers: Each layer uses a separate router for each loop, while expert weights remain shared. This decouples routing decisions across iterations, enabling different stages of recurrence to access different expert combinations without replicating expert parameters.

Looping Residual: A fixed-decay exponential moving average (EMA) with two complementary memories:

• Global memory ($A = H$): spans layers and loops
• Local memory ($A = L$): aggregates attention outputs within the current loop

$$N_A \leftarrow \beta N_A + o_{\ell}^{t}, \quad D_A \leftarrow \beta D_A + 1, \quad r_A = \frac{N_A}{D_A}, \qquad A \in \{H, L\}$$

The summaries replace the direct attention residual update:

$$\widetilde{h}_{\ell}^{t} = h_{\ell-1}^{t} + r_H + r_L$$

Each memory stores only one accumulator tensor and one scalar, introducing negligible overhead.

3.3 Segmented Backpropagation

LOOM divides the H recurrent loops into consecutive segments of at most $K = 3$ loops. Each segment receives supervision from a language-modeling loss at its final loop, with gradients restricted to that segment. The residual state and global Looping Residual memory propagate to the next segment with gradients detached.

---

Empirical Validation / Results

Experimental Setup

Models are pretrained on FineWeb-Edu with sequence length 1,024 and vocabulary size 65,536. Three scales are evaluated:

Near-Iso-FLOP Comparison (Table 2)

On the M=10, E=80 backbone (700M parameters, 10B tokens), the per-token layer cost is $f(H,k) = H(6 + 3k) = 3H(2 + k)$ in units of $d^2$, anchored at $f_1 = 84$.

Non-Iso-FLOP Loop Sweeps (Table 3, excerpt)

At ~350M scale, LOOM remains stable and improves with depth:

Partial stabilization baselines ("No tech", "Embed inject", "Res. scale") all exhibit unrecoverable loss spikes (†) at larger loop counts, confirming that the complete recipe is necessary.

Scaling to ~1.7B Parameters (Table 4)

Component Ablation (Table 5, ~350M, 9 loops, step 5,000)

Segmented Backpropagation Benefits

• Memory: Peak memory stays near 12.6 GiB from H=3 to H=12 (vs. 12.5→19.3 GiB unsegmented)
• Speed: Training time decreases from 26.1 to 12.6 hours at H=12
• Stability: At H=9, unsegmented runs exhibit severe degradation (eval loss ≈ 4.4, accuracy ≈ 34%)

---

Theoretical and Practical Implications

Theoretical Significance

1. Recurrent depth as a compute-efficient scaling axis: The near-iso-FLOP results provide strong evidence that recurrent depth is a genuine scaling dimension for MoE LLMs, not merely a parameter-inefficient alternative to depth scaling.

2. Understanding failure modes: The identification of the curse of depth and expert selection collapse as distinct, addressable obstacles provides a diagnostic framework for future recurrent architecture research.

3. Generalizable design principles: The principle that each loop should contribute new computation while keeping the recurrent state stable offers a unified framework for designing deep recurrent architectures.

Practical Implications

1. Cost-effective scaling: LOOM achieves substantial gains (e.g., 1.7B model: perplexity 9.62→7.77) without increasing parameter count, making it attractive for deployment scenarios with memory constraints.

2. Stable training at scale: The recipe's hyperparameters ($\lambda = 0.5$, $\beta = 0.5$, $K = 3$) transfer across model scales without tuning, simplifying adoption.

3. Efficient memory management: Segmented backpropagation enables training at depths that would otherwise be memory-prohibitive, with negligible overhead from the Looping Residual.

---

Conclusion

This paper presents LOOM, a scalable recipe for extending recurrent depth in MoE LLMs beyond the two-loop regime. The key contributions are:

1. Diagnosis: Identification of hidden-state variance growth (curse of depth) and expert selection collapse as the primary obstacles to deep recurrence in MoE models.

2. Recipe: A combination of residual scaling ($\gamma = \lambda/(H\sqrt{M})$), embedding re-injection ($g_t = \lambda/(t\sqrt{M})$), per-loop routers, and a Looping Residual (EMA-based) that together enable stable, diverse recurrent computation.

3. Empirical validation: Stable scaling to 9–12 loops across 100M–1.7B parameters, with the 1.7B model scaling to 9 loops and achieving substantial improvements (perplexity: 9.62→7.77; accuracy: 42.4%→47.7%).

Future Directions

• Exploring even larger model scales and longer training runs
• Investigating whether the recipe extends to other architectures (e.g., dense Transformers, hybrid models)
• Studying the interaction between loop depth and other scaling axes (data, width, expert count)
• Potential applications to test-time compute scaling and inference-time reasoning

The results provide compelling evidence that recurrent depth offers a compute-efficient scaling axis for MoE language models, opening new avenues for scaling LLM capabilities within fixed parameter budgets.