Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts
Di He, Pengxiang Li, Da Chang, Qingyan Meng, Lu Yin, Shiwei Liu
cs.LG, cs.AI
2026-10-01
LOOM uses residual scaling, input re-injection, and per-loop routers. At matched FLOPs, 700M perplexity falls from 18.36 to 16.54 at 5 loops; a 1.7B model hits 7.77 at 9.
Reusing one stack of Transformer blocks adds effective depth without copying the weights. That is a looped Transformer. Dense models treat it as a scaling axis beside parameter count and data, a way to buy more computation along depth. MoE looks like a fit: each token already touches only a few experts, so another pass could refine the same pool.
The gains fade fast, and more loops often make the loss worse. Published looped MoE systems usually stop at two iterations. Two mechanisms set that ceiling. Looping amplifies the curse of depth. Pre-LN puts normalization on the residual branch, residual updates pile up, hidden-state variance grows, and the representation drifts. The other is expert selection collapse. A router shared across loops, the gate that assigns experts, keeps picking nearly the same experts, so the extra FLOPs do not buy a new mix of computation.
Each loop should add computation while the recurrent state stays stable. The backbone is a Llama-style pre-norm MoE decoder: an M-layer block runs, then repeats for H loops. Attention and expert weights are shared. Each loop has its own router. Effective depth is M×H, and the added parameters are mostly routers.
λ is fixed at 0.5 and not trained. Residual branches are scaled by γ = λ/√(H×M), on both the attention output and the MoE output, and the MoE branch also goes through RMSNorm. The 1/√H factor limits accumulation across loops. The 1/√M factor limits variance across physical layers. Variance control is not sufficient, because deep recurrence still washes out the input. From loop 2 on, the input embedding is mixed back with gt = λ/√(t×M): the new state is (1−gt) times the previous loop output plus gt times the input. gt shrinks as the loop index grows.
Per-loop routers are there to stop expert selection collapse. Every layer has a sigmoid router for each loop, computed in float32, while the experts stay shared. On a roughly 350M model at 9 loops, expert loads then become less similar across loops and fluctuate more. A Looping Residual keeps earlier attention outputs around. Scaled attention outputs enter an exponential moving average with β = 0.5. One memory is global across loops. One is local and resets every loop. The two summaries replace the direct attention residual, and the MoE residual runs after that.
Backprop is cut into segments of K = 3. On every third loop, and on the final loop, the model takes a language-modeling loss, backpropagates, steps the optimizer, and detaches the hidden state and the global memory. Nine loops therefore put three supervised updates on one sequence. The forward pass still visits every loop.
Pretraining uses FineWeb-Edu at sequence length 1024. Zero-shot scores cover OpenBookQA, WinoGrande, ARC-Challenge, ARC-Easy, HellaSwag, SIQA, and PIQA.
The matched-FLOP backbone is a separate 700M model: M = 10, hidden size 512, 80 routed experts, no shared experts, 10B tokens, full backpropagation. Per-token layer cost is f = H×(6+3k) in units of d². The non-looped anchor uses top-k = 26 and f = 84. Deeper runs shrink top-k to stay near that anchor.
| Setup | Val. perplexity | Seven-task avg. |
| No loop, k=26 | 18.36 | 38.84% |
| 2 loops, k=12 | 17.37 | 39.00% |
| 3 loops, k=8 | 16.91 | 39.34% |
| 4 loops, k=5 | 16.58 | 39.53% |
| 5 loops, k=4 | 16.54 | 39.53% |
| 6 loops, k=3 | 16.57 | 39.50% |
Five loops reach 16.54, below the non-looped 18.36 and the two-loop 17.37. Average accuracy is 39.53% at both four and five loops, and 39.50% at six, which is 0.69 points above 38.84%. At five loops WinoGrande falls from 51.38% to 50.20%, and SIQA from 35.72% to 34.65%. Loops 3, 5, and 6 sit at f = 90, about 7% more multiply-adds than the anchor.
If top-k stays fixed, every extra loop is extra FLOPs. At about 350M parameters (about 141M activated) and 10B tokens, the non-looped run is 20.07 / 37.96%, nine loops are 14.80 / 41.15%, and twelve loops are 14.86 / 41.42%. At about 100M (about 70M activated) and 5B tokens, the non-looped run is 26.26 / 36.73%, and twelve loops are 19.26 / 39.11%. On the 350M model, a loop with no extra technique hits perplexity 1312 at six loops. Residual scaling alone, or embedding re-injection alone, is already worse at three loops (23.80 and 42.76, against a 20.07 baseline) and later runs into unrecoverable loss spikes.
The 1.7B model, about 0.63B activated, is trained for 60B tokens with top-k fixed at 8.
| Loops | Effective depth | Val. perplexity | Seven-task avg. |
| 1 | 15 | 9.62 | 42.4% |
| 6 | 90 | 7.91 | 46.7% |
| 9 | 135 | 7.77 | 47.7% |
| 12 | 180 | 7.84 | 47.5% |
Nine loops is the peak: perplexity moves from 9.62 to 7.77, and average accuracy from 42.4% to 47.7%. HellaSwag goes from 37.1% to 48.8%, ARC-Easy from 50.2% to 57.7%. Twelve loops slips on both.
At step 5,000, the 350M model with 9 loops has perplexity 18.62 with the full recipe. Dropping residual scaling, input re-injection, or the MoE RMSNorm moves that to 24.63, 26.05, and 26.50. Dropping the Looping Residual leaves 19.25. Sharing one router across loops leaves 19.83.
Segmented backpropagation is the default when FLOPs are not matched. On the 350M model with microbatch 1, NPU memory stays near 12.6 GiB from 3 to 12 loops, while full backpropagation grows from 12.5 GiB to 19.3 GiB. At 12 loops, training time falls from 26.1 hours to 12.6 hours. At 9 loops, full backpropagation reaches perplexity 81.04 and 34.3% accuracy, against 14.80 and 41.1% with segments. The deepest full-backprop runs also cut the microbatch, so the memory numbers are not strictly matched.
Matched FLOPs are possible here because the baseline activates 26 experts and can trade that width for loops. A top-2 or top-8 MoE has no such budget. Looping it multiplies FLOPs by the loop count. Inside the trade that was actually run, four or five loops beat both the non-looped model and the two-loop run on perplexity. Downstream accuracy gains 0.69 points and is flat after that.
Nine loops on the 1.7B model add 5.3 accuracy points and take effective depth from 15 to 135. The writeup treats 1.7B×9 as an unrolled scale of about 15B. The parameter count is still 1.7B, and no non-looped model was trained at the same FLOPs.
λ, β, and K stay at 0.5, 0.5, and 3 from the 100M models through 1.7B. The training piece that transfers most cleanly is segmented backpropagation: at 12 loops, wall-clock time is about half, and memory stops tracking depth. A deployed MoE that already uses a small top-k does not get cheaper FLOPs from this recipe. It gets more FLOPs, unless expert width is cut to pay for the loops.
There is no separate limitations section. Under matched FLOPs, accuracy rises only 0.69 points, and WinoGrande and SIQA fall. Seven short commonsense tasks are weak support for calling recurrent depth an efficient scaling axis. Perplexity carries that claim. Context length is 1024. Math, code, long context, algorithmic tasks, and length extrapolation are not measured, though the introduction cites algorithmic computation and length generalization as reasons to loop.
The 5.3-point gain arrives with roughly 9× the depth compute, and with a language-modeling loss plus an optimizer step every three loops. The matched-FLOP runs turn segmentation off, and the gain shrinks. Depth, intermediate supervision, and update frequency stay tangled.
Per-loop routers and the Looping Residual are ablated only at step 5,000, where perplexity moves by about 0.6 to 1.2, against 6 to 8 for the stabilization terms. Figure 4 shows that routing becomes more diverse. It does not show that diversity is the main source of the gain. λ, β, and K have no sensitivity sweep, and there is no same-setup comparison with Mix-LN, mHC, or fixed-point recurrence. Scale stops at 1.7B parameters and 60B tokens.