PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Di He, Pengxiang Li, Da Chang, Qingyan Meng, Lu Yin, Shiwei Liu

arXiv:2610.01153v1cs.LGcs.AI

TL;DR

Deeper recurrence in MoE language models often brings diminishing returns or degradation, leaving its benefits under matched compute unclear. LOOM combines state stabilization with loop-specific expert routing and a Looping Residual. Experiments report stable scaling to 9–12 loops and lower loss at five loops under matched FLOPs.

  • Problem

    Deeper looped MoE often yields diminishing returns or loss degradation, leaving its benefit under matched FLOPs unclear and leading prior approaches to use two loops.

    Source: §1 · p. 2

  • Method

    LOOM stabilizes recurrent states with residual scaling and embedding re-injection, and diversifies computation with per-loop routers and a Looping Residual.

    Source: §1 · pp. 2–3

  • Results

    LOOM scales stably to 9–12 loops, and under matched FLOPs scales to five iterations while consistently achieving lower loss.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The matched-FLOP results provide evidence that recurrent depth can be a compute-efficient scaling axis for MoE LLMs.

    Source: §1 · p. 3

Abstract

from arXiv · show

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.

1. Introduction

Looped MoE can gain effective depth without increasing parameter count, but deeper recurrence often yields diminishing returns or degradation under matched compute. LOOM targets instability and expert-selection collapse with a recipe for stable, diverse computation across loops.

1. Introduction

  • Looped Transformers reuse shared blocks to increase effective depth without increasing parameter count, adding a scaling axis alongside model size and training data.Prior studies report benefits for algorithmic computation, length generalization, and latent reasoning.Source: §1 · p. 2
  • Deeper looped MoE often degrades or yields diminishing returns, so extra iterations can add computation without corresponding performance gains and prior approaches often stop at two loops.The identified obstacles are residual accumulation that grows variance and causes latent drift, and routers repeatedly selecting nearly identical experts across loops.Source: §1 · p. 2
  • LOOM is designed so each loop contributes new computation while keeping the recurrent state stable, enabling shared experts to progressively refine representations beyond two iterations.The contribution summary attributes stable scaling to 9–12 loops across experiments spanning 100M–1.7B parameters and up to 60B training tokens.Source: §1 · pp. 2–3
  • LOOM stabilizes recurrence with residual scaling and input-embedding re-injection, while loop-specific routers diversify expert selection and a Looping Residual carries earlier outputs forward for refinement.The injected embedding is recalibrated across iterations so its anchoring signal remains well-scaled.Source: §1 · p. 3
  • Under matched FLOPs, LOOM scales to five iterations while consistently achieving lower loss, supporting recurrent depth as a compute-efficient scaling axis.Across models from 100M to 1.7B parameters, experiments show stable scaling to 9–12 loops; the 1.7B model scales to 9 loops.Source: §1 · p. 3

2. Related Work

Prior work explores recurrent computation in dense Transformers and several forms of MoE recurrence, but scalable deep recurrence under a fixed MoE architecture remains underexplored. LOOM builds on residual stabilization and input re-injection while addressing routing and cross-loop information flow.

2. Related Work

  • Dense Transformer studies link recurrence to test-time looping, latent-space reasoning, and deeper computation, while Chain-of-Experts iterates experts within layers rather than reusing a full shared MoE block.Universal Transformer and ALBERT are identified as early precedents for reusing Transformer blocks along depth.Source: §2 · p. 3
  • Residual-stability work addresses variance growth and representation drift through residual scaling, normalization changes, or redesigned residual pathways; input re-injection also alleviates drift.Prior recurrent approaches include a loop-aware scaling factor and inter-loop mixing.Source: §2 · p. 3
  • In the 9-loop stability experiment, unscaled residual accumulation produces activation-variance growth, residual scaling controls it, and scaling alone has higher training loss than the complete recipe.The figure compares Native, Only Res Scale, and Only Embed Inject variants across training loss and activation variance.Source: Figure 3 · p. 4
  • Existing recurrent MoE studies use limited forms of recurrence, and compute-matched SMELT identifies two loops as best; scalable deep recurrence under a fixed architecture remains underexplored.Other approaches include iteration-adaptive normalization or reducing layer count as loop count increases.Source: §2 · p. 4

3. Methodology

LOOM reuses an M-layer MoE decoder across H loops, combining recurrent-state stabilization with mechanisms for diverse, cumulative computation. Segmented backpropagation limits gradient depth while retaining information across loops.

3.1. Stabilizing the recurrent state

  • At H=9, removing embedding re-injection raises validation perplexity from 18.62 to 26.05; LOOM pairs it with residual scaling to control variance and preserve token input.Residual scaling accounts for loop count and physical depth, while re-injection begins at loop 2.Source: §3.1 · pp. 4–5

3.2. Diversifying computation across loops

  • Per-loop routers decouple expert choices, while global cross-loop and local within-loop EMA memories carry prior attention computations forward for refinement.The global memory persists across recurrence; the local memory resets at each loop, and each stores only one accumulator tensor and one scalar.Source: §3.2 · p. 5

3.3. Training with segmented backpropagation

  • At H=12, segmented backprop cuts training time from 26.1 to 12.6 hours; across H=3–12, peak memory stays near 12.6 GiB.Segmented runs also achieve lower evaluation loss and higher average accuracy across seven tasks; at H=9, unsegmented training reaches approximately 4.4 loss and approximately 34% accuracy.Source: Figure 5 · p. 6, §3.3 · p. 6

4. Empirical results

Section 4 covers experimental setup, loop-depth results from approximately 100M to 1.7B parameters, and ablations of the LOOM recipe.

  • Section 4 presents the experimental setup, loop-depth results from approximately 100M to 1.7B parameters, and recipe ablations.Source: §4 · p. 6

4.1. Experimental setup

The experiments use Llama-style MoE decoders pretrained on FineWeb-Edu and evaluate both loop-depth settings and recurrent-method baselines. The setup also specifies model scales, optimization, and seven zero-shot reasoning tasks.

  • Models use a Llama-style pre-norm decoder with grouped-query attention, RoPE, and sparse SwiGLU MoE layers, pretrained on FineWeb-Edu with sequence length 1,024.The vocabulary has 65,536 tokens.Source: §4.1 · p. 6
  • The near-iso-FLOP comparison uses an M=10, E=80 backbone; all models have 700M parameters and train on 10B tokens.Source: Table 2 · p. 7
  • Per-token layer cost is f = H(6 + 3k) in units of d2, anchored at f1 = 84.Source: §4.1 · p. 7
  • Across the three scales, loop-specific coefficients remain fixed at λ=0.5, β=0.5, and K=3, while width and expert dimensions scale with model size.Settings within each scale stay fixed when varying H, isolating recurrence depth in those comparisons.Source: §4.1 · p. 7, Table 1 · p. 7
  • Training uses AdamW, gradient clipping at 1.0, a 10/80/10 warmup–stable–decay schedule, and global batch size 1,024.The peak learning-rate floor is 0.1×.Source: §4.1 · p. 7
  • Baselines comprise a non-looped 1× model and three recurrent variants: unmodified looping, residual scaling alone, and embedding re-injection alone.Source: §4.1 · p. 7
  • Zero-shot evaluation uses seven commonsense tasks: PIQA, SIQA, HellaSwag, WinoGrande, ARC-c, ARC-e, and OBQA.The setup reports length-normalized accuracy for OBQA, ARC-c, ARC-e, HellaSwag, and PIQA, and raw accuracy for WinoGrande and SIQA.Source: §4.1 · p. 7

4.2. Loop-depth experiments

The experiments test loop depth under near-iso-FLOP and non-iso-FLOP settings, then evaluate scaling at approximately 1.7B parameters. LOOM remains effective beyond two loops, with the preferred depth varying by compute setting.

4.2.1. ISO-FLOP COMPARISON

  • Under near-iso-FLOP compute, validation perplexity falls from the non-looped baseline through Loop4, and Loop5 and Loop6 maintain Loop4 performance.The experiment adjusts top-k so compute remains approximately matched to the non-looped anchor.Source: §4.2.1 · pp. 7–8

4.2.2. NON-ISO-FLOP LOOP EXPERIMENTS

  • In the non-iso-FLOP sweep, the highest seven-task average occurs at 12 loops, while LOOM improves validation perplexity and downstream accuracy over width-matched non-looped baselines across both model scales.Gains gradually saturate as recurrence deepens.Source: §4.2.2 · p. 8, Table 3 · p. 8

4.2.3. SCALING UP TRAINING

  • At ∼1.7B parameters, 9 loops achieve 7.77 validation perplexity and 47.7% average zero-shot accuracy, versus 9.62 and 42.4% for the non-looped baseline.The model was trained on ∼60B tokens; its unrolled scale is roughly 15B parameters.Source: Table 4 · p. 9, §4.2.3 · p. 9

4.2.2. NON-ISO-FLOP LOOP EXPERIMENTS

  • The unmodified, residual-scaling-only, and embedding-re-injection-only recurrent baselines underperform non-looped baselines even at shallow recurrence; residual scaling alone still encounters unrecoverable loss spikes at larger loop counts.Residual scaling mitigates perplexity degradation, while embedding re-injection alone exhibits substantial degradation as loop depth increases.Source: Table 3 · p. 8, §4.2.2 · p. 9

4.3. Ablation

The ablations examine which LOOM components support deep recurrence and whether segmented backpropagation improves training. They support the full recipe and report benefits from restricting gradient propagation to segments.

4.3.1. COMPONENT ABLATION

  • Removing residual scaling, embedding re-injection, or MoE RMSNorm substantially increases perplexity and reduces average downstream accuracy; removing per-loop routing or the Looping Residual worsens both metrics.The ablations support stable recurrent states, computational diversity, and information preservation as design principles.Source: §4.3.1 · p. 9

4.3.2. SEGMENTED BACKPROPAGATION

  • Compared with full backpropagation, segmented backpropagation with K=3 improves training stability, validation perplexity, and average downstream accuracy.It restricts gradient propagation to the current segment while preserving forward information across loops; deepest full-backpropagation runs use a smaller micro-batch and are not matched in activation memory.Source: §4.3.2 · pp. 9–10, Table 6 · p. 10

5. Conclusion

The conclusion presents LOOM as a recipe for scaling recurrent depth in MoE LLMs beyond two loops. Its experiments show stable scaling across model sizes, including under matched training FLOPs.

5. Conclusion

  • LOOM extends recurrent depth in MoE LLMs beyond the commonly used two-loop regime.The experiments span approximately 100M to 1.7B parameters and up to 60B training tokens, with stable scaling to 9–12 loops.Source: §5 · p. 10
  • Residual scaling and embedding re-injection stabilize recurrent states, while per-loop routers and a Looping Residual diversify computation and preserve information across iterations.Source: §5 · p. 10
  • Under matched training FLOPs, LOOM scales to five loops while consistently achieving lower loss, supporting recurrent depth as a compute-efficient scaling axis for MoE language models.Source: §5 · p. 10
Loading 2610.01153v1…