How to Loop MoE: Flatten the Experts, Untie the Attention

AI summary

Overview

Research area: Machine learning — efficient large-language-model architecture, specifically the intersection of looped (recurrent-depth) Transformers and sparse mixture-of-experts (MoE) models, validated by large-scale pretraining runs.

Technical level: Advanced. The paper assumes familiarity with Transformer blocks, top-k routing, expert parallelism, pretraining loss accounting, and MoE utilisation metrics.

Scope: A single design study proposing Foil, a recipe for laying out experts and attention across a looped MoE block at fixed expert parameters and fixed expert compute per token, evaluated at 20B and 100B tokens of pretraining.

What This Paper Is About

Looped Transformers reuse one block of layers several times to buy extra computation without storing more parameters, while sparse MoE models store many experts per layer but send each token to only a few of them. The two ideas can be combined, but nobody had asked how to arrange the experts over the looped block — how many experts per layer, how many layers in the block, and how many passes — when the expert parameters and the compute per token are held fixed. The paper answers that question with Foil, which flattens the experts (fewer, wider layers looped more often) and unties the attention (each pass gets its own attention parameters while experts and routers stay shared).

Key Contributions

1. The Foil design. A method that flattens a looped MoE — halving the layers in the looped block, doubling the experts per layer, and doubling the passes per step — and unties the attention so every pass owns its own attention parameters, all while holding the expert parameters and the expert compute per token fixed.
2. Large-scale validation. Pretraining at 20B tokens and continued training to 100B tokens showing that Foil lowers pretraining loss relative to the unflattened looped baseline, matches or slightly exceeds it on downstream zero-shot tasks, and that untying attention yields more balanced and more confident routing than tying it at the same shape.
3. Systematic ablations of looped MoE. Experiments isolating how widening expert layers and adding passes interact, how load balance relates (or fails to relate) to healthy expert use, and how a routing-confidence metric called MMR behaves pass by pass.
4. Design guidance and release. A distilled recommendation that a sparse looped MoE should use more experts per layer and more passes, plus public code and configurations at https://github.com/SR-A-W/how-to-loop-moe.

Main Findings

• Flattening with more passes lowers loss at 20B tokens. Holding the real expert count E_real = 64 and the expert calls per token E_comp = 32 fixed, all three Foils reach a lower loss than Base. The loss falls through the second flattening step and then levels off: Foil-1 ends 0.007 nat below Base, slightly above Foil-2. On the three representative downstream tasks the Foils are at or above Base, within their standard errors.
• At 100B tokens the gain grows monotonically with flattening. Every flattening step lowers the loss and Foil-1 is the strongest, 0.012 nat below Base. All three Foils score slightly above Base on the three representative tasks, by one to two standard errors.
• Why loop a MoE at all (motivating measurements). At equal parameters, looping twice lowers the loss by 0.064 nat, and the number of distinct experts a token reaches grows from 16 without looping to 36 with eight passes. A looped block flattened from (8,8,2) through (16,4,4) and (32,2,8) to (64,1,16) raises the equivalent expert count E_eq from 128 to 1024 while E_real, E_comp and the effective depth D_eff stay fixed.
• Untying attention helps at every shape, more so when flatter. At 20B tokens the untied model wins in all four tied/untied pairs, most of all when fully flattened: Foil-1 is 0.042 nat below proto-Foil-1. After continued training to 100B tokens the gap grows steadily with flattening, to 0.049 nat for the Foil-1 pair. Downstream, the untied model is ahead in all four pairs at both token budgets; at 20B the difference exceeds two standard errors in every pair except the Foil-2 pair, and at 100B it reaches 3.3 points for Foil-1 over proto-Foil-1.
• Untied attention routes more evenly and more confidently at equal shape. The untied model has both a higher load balance B_2 and a higher routing confidence MMR than its tied counterpart in all four pairs, with the smallest gap at the Base pair.
• Widening and looping amplify each other. From four to eight passes, (16,4) gains 0.041 nat against at most 0.033 for (8,4), which has half its real experts, and (8,8), which has twice its layers. Widening (8,4,L) to (16,4,L) gains 0.038 nat at L = 2 but 0.049 at L = 8. Widening shows diminishing returns: going from (8,8,2) to (16,8,2) gains only 1.2 times as much as going from (4,8,2) to (8,8,2), although it adds eight experts per layer instead of four.
• Load balance alone does not indicate healthy specialisation. Along the flattening sequence with untied attention, B_2 falls steadily from Base to Foil-1 while MMR rises, yet the most flattened model reaches a lower loss than the unflattened one. Masking the experts carrying the least 10% of routing traffic raises the loss more than traffic-matched random groups in seven of the eight models compared, and in none does it fall below the random range; the exception, proto-Foil-2, is on par with the random median.
• Routing confidence tracks the looping gain. In the (8,8,L) series the model-level MMR falls from 4.15 to 3.36 and the average gain per pass over the non-looped (8,8,1) falls from 0.064 to 0.022 nat as L goes from 2 to 8; the two decline together at every step.
• A per-pass MMR peak signals that flattening is nearly exhausted. Models such as (8,8,8), proto-Foil-2, proto-Foil-1 and Foil-1 peak at some pass and fall afterwards. After its peak, flattening proto-Foil-2 further to proto-Foil-1 raises the loss by 0.017 nat; Foil-1, which peaks, is 0.0020 nat above Foil-2 at 20B, about three standard errors; and (8,8,8) has the smallest average gain per pass of its series. Foil-2, which shows no peak, has the lowest loss of the Foil family including Base. The authors state this as "very likely" rather than as a rule.
• Compute cost of the study. The runs consumed approximately 12.6k NVIDIA H200 GPU-hours (about 18.2k including earlier control and failed runs); each run used at most one node with eight H200 GPUs.

Methodology in Plain English

The authors start from the looped skeleton of Huginn: a token embedding, a prelude of one ordinary MoE layer, a looped block of D layers applied L times, a coda of one ordinary MoE layer, and the output layer. The prelude and coda have 8 experts each, run once, and are never flattened, so every configuration differs only in the looped block, written as (E, D, L).

Flattening is a single step that maps (E, D, L) to (2E, D/2, 2L): half the layers, twice the experts per layer, twice the passes. That keeps the number of stored experts (E_real = E × D), the expert calls per token (E_comp = k × D × L) and the effective depth (D_eff = D × L) unchanged, while enlarging the pool each routing decision chooses from and the equivalent expert count E_eq = E × D × L. Three steps from (8,8,2) produce (16,4,4), (32,2,8) and (64,1,16), raising E_eq from 128 to 1024. The three flattened models are called Foil-3, Foil-2 and Foil-1; the unflattened looped model (8,8,2) is Base.

A conventional looped model shares the whole block across passes, attention included, so flattening would also discard attention parameter sets. The authors therefore keep experts and routers shared across passes but give every pass its own attention, giving D × L = 16 attention sets along the whole sequence, with no pass embedding. Flattened models with shared attention are called proto-Foil; those with untied attention are called Foil. Base and the Foils each have 553.7M parameters; the shared-attention models have 490.8M to 520.2M.

Every model routes each token to the k = 2 most probable experts of a layer under a bias-free linear router with softmax at temperature 1 and renormalised weights, using SwiGLU experts. All models are trained on the same data in the same order from the same initialisation seed under the same learning-rate schedule for 20B tokens, which is 35–40 tokens per parameter for the eight compared models, and the main models are scaled to 100B tokens.

To see how experts are actually used, the authors probe the recurrent core over 255,500 tokens with three metrics: the number of distinct experts reached per token (U_t), load balance B_2 (Jain's fairness index, equivalently the effective number of experts N_2 divided by E), and routing confidence MMR — the median margin ratio p_(1)/p_(E/2), comparing the router's top choice with the median-ranked expert of the whole pool. Because B_2 and MMR both depend on E, they are only read at equal shape or as relative trends.

Downstream, the main text reports three zero-shot tasks: LAMBADA (standard split) for word prediction, HellaSwag for sentence continuation, and XWinograd (English) for coreference resolution. Loss differences are reported as paired differences with standard errors.

Why This Matters

Impact on research. The paper treats the layout of experts and attention over a looped block as a first-class design variable and shows it can be changed at no cost in parameters or compute. It also questions the standard practice of reading load balance as evidence of healthy expert specialisation, proposing routing confidence (MMR) and a traffic-matched masking test instead, and showing that the least-used experts are not the least useful. For the looped-MoE literature — which previously combined looping and experts without comparing layouts at fixed budget — it supplies the missing controlled comparison and a concrete recommendation: more experts per layer, more passes.

Potential real-world applications (implications of the design, not claims tested in the paper):

• Serving large models under memory-bandwidth limits, where trading stored parameters for repeated computation is attractive as compute capacity grows faster than memory capacity and bandwidth.
• Deploying a fixed-size model that can be made stronger by computing more, which suits latency-tolerant batch workloads and iterative reasoning-style computation.
• Building sparse MoE systems where more of the stored experts are actually exercised per token, which matters for utilisation and hardware efficiency of expert-parallel infrastructure.
• Informing architecture choices for pretraining

Authors’ abstract

Looped Transformers reuse one block of layers several times: by spending extra computation they push a model of fixed size further, and so use its parameters more fully; while sparse mixture-of-experts (MoE) models activate only a few of many experts for each token. Looped MoE bridges these two design philosophies and gives MoE models new potential for better expert usage, but it raises a question: how to loop a MoE? We answer it with Foil. With the expert parameters and the expert compute per token held fixed, Foil (1) flattens the experts, halving the expert layers, doubling the experts per layer and doubling the passes, so that every routing decision chooses from a larger pool, and (2) unties the attention, giving each pass its own attention parameters while the experts and routers stay shared. Experiments show that Foil clearly outperforms the unflattened looped baseline: at 20B tokens every Foil model has lower pretraining loss than the baseline; at 100B tokens the loss improves monotonically with the degree of flattening, the most flattened Foil ending 0.012 nat below the baseline at equal parameters and compute, with downstream accuracy on par or better; untying the attention also yields more balanced and more confident routing at equal shape. Our ablations analyse why Foil works and turn the findings into design guidance for looped MoE: the returns of looping and of widening the expert layers amplify each other, routing confidence tracks healthy expert use better than load balance, and a sparse looped MoE should therefore use more experts per layer and more passes. Code and configurations are available at https://github.com/SR-A-W/how-to-loop-moe.