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

What it claims

Foil restructures a looped MoE at fixed expert parameters and fixed expert compute per token. The baseline, Base (8 experts × 8 layers × 2 passes), already has per-pass attention, so the 0.012 nat gain at 100B tokens comes purely from flattening, not from untying.

Untying attention is a separate, larger effect measured against shared-attention controls. At the most flattened shape, untied attention beats tied by:

• 0.042 nat at 20B tokens
• 0.049 nat at 100B tokens
• 3.3 points of mean zero-shot accuracy at 100B

The tied controls have fewer parameters (490.8–520.2M vs 553.7M), so that comparison is not at equal parameters. It is at equal compute.

How it works

The skeleton is borrowed from Huginn (Geiping et al., 2025). A prelude (embedding plus one ordinary 8-expert MoE layer) feeds a looped core of D layers run L times, followed by a coda (one 8-expert MoE layer plus the LM head). Only the core is reshaped, written as (E, D, L).

• Flattening. One step maps (E, D, L) to (2E, D/2, 2L). This keeps the following fixed:
• real experts (E×D = 64)
• expert calls per token (k×D×L = 32)
• effective depth (D×L = 16)

It enlarges each routing pool. The variants are Foil-3 (16,4,4), Foil-2 (32,2,8) and Foil-1 (64,1,16). "Equivalent experts" (E×D×L) rise from 128 to 1024.

• Untying. Experts and routers are shared across passes. Each pass gets its own attention set, D×L = 16 sets in total, selected by pass index. There is no pass embedding, and this adds no compute.
• Routing. Top-2 routing with a linear softmax router at temperature 1, renormalised weights, no noise or bias terms. Experts are SwiGLU.

flowchart LR
  P[Prelude: embed + MoE] --> C[Core layer d]
  C --> A{pass index ℓ}
  A --> AT[attention set ℓ, untied]
  AT --> X[shared experts + router]
  X -->|ℓ < L| C
  X -->|ℓ = L| K[Coda: MoE + LM head]

Evidence

All models share the same data order, seed and LR schedule. They were trained for 20B tokens (35–40 tokens/param), and the main models were continued to 100B.

Loss vs Base, untied attention throughout:

• At 20B, the gain levels off after two flattening steps, and Foil-2 beats Foil-1.
• At 100B, the loss falls monotonically with flattening.

Downstream (LAMBADA, HellaSwag, XWinograd, zero-shot):

• At 20B, the Foils are at or above Base, within standard error.
• At 100B, they are slightly above Base, by 1–2 standard errors. At 100B, Foil-1 vs Base:

Tied vs untied attention: with tied attention, loss rises again after the first flattening step. At every shape, the untied model has both higher load balance (B2, Jain's fairness index) and higher routing confidence (median margin ratio: top-1 vs median-ranked expert probability).

Looping motivation:

• Looping twice at equal parameters lowers loss by 0.064 nat.
• Distinct experts reached per token grow from 16 (no looping) to 36 (eight passes).

Compute: about 12.6k H200 GPU-hours for the reported runs, or about 18.2k including control and failed runs. Each run used at most one 8×H200 node.

How to use it

Code and configurations: https://github.com/SR-A-W/how-to-loop-moe. The page gives no install or run commands.

The authors' design guidance:

• Use more experts per layer and more passes in a sparse looped MoE.
• Give each pass its own attention.
• Judge expert health by routing confidence rather than load balance alone.

Limits

• Scale. Results come from roughly 554M-parameter models only, with a maximum of 100B tokens.
• Size of gains. The headline gain is 0.012 nat, and downstream gains over Base are within 1–2 standard errors.
• Routing metrics. B2 and MMR come from a single-seed probe without error bars. Both depend on E, so they are only comparable at equal shape.
• Tied-attention comparison. It confounds untying with extra attention parameters.
• Unverified parts. The fetched text cuts off mid-Section 3.2. The ablations behind the design guidance (Section 4), the seed-change robustness experiment and the masking test (appendices) were not seen.