Paper Walkthrough — SMELT: Is Looping the Same Layers Twice Actually a Win When the Budget Is Matched?
A study of looped Transformers that finally controls the comparison: per-token FLOPs, total parameters, and KV cache are all held close. The resulting recipe, SMELT, loops the middle half twice and reports 6.8–18.0% training-FLOPs savings on the compute-optimal frontier.
SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers
Primary source — what this article is built on
undefined2026-09-01→undefined2026-09-03same month
SMELT: Scaling Laws for Compute-Matched MoE Looped TransformersShaowen Wang, Ge Zhang, Kairong Luo et al. · 2026-09-01 · v1arXiv:2609.01343Paper page·PDFundefined
Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size, conflating architectural advantage with extra FLOPs. We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache. Through a series of ablations, we arrive at a recipe we call SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets. We scale SMELT across four sizes up to 54B non-embedding parameters and fit a separate Chinchilla-style scaling law for each architecture. SMELT's loss drops faster with compute, saving 6.8--18.0\% of training FLOPs on the compute-optimal frontier. The advantage transfers to downstream benchmarks beyond what validation loss predicts, is largest on Code, and grows with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens, an inductive bias that may underlie the observed performance gains. These results show that looping can improve Transformers even under budget matching, offering a practical recipe that turns depth reuse into measurable gains.
An old idea, revisited: run the same layers again
If you want a deeper model, the obvious move is to stack more layers. More layers means more weights, more storage, and more arithmetic per token.
There is another move. Don't add layers — send the activations through the same layers a second time. Picture a ten-storey building where you walk floors 5 through 7 twice before leaving. The number of floors (the weights you store) is unchanged, but the number of floors you actually walked through (the layers executed) goes up. That is a Looped Transformer, an idea that has been proposed repeatedly since the Universal Transformer.
The paper we are reading is "SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers" (arXiv:2609.01343, Shaowen Wang et al., 1 September 2026).
Its claim, in the authors' own framing: Looped Transformers increase effective depth by iterating a shared block of layers, but most prior evaluations compare at fixed model size, which conflates an architectural advantage with simply spending extra FLOPs. This paper studies looping on Mixture-of-Experts (MoE) Transformers while closely matching three things at once — per-token FLOPs, total non-embedding parameters, and KV cache. A series of ablations produces a recipe called SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of the layer stack twice while matching the unlooped Baseline on all three budgets. The recipe is scaled across four sizes up to 54B non-embedding parameters, with a separate Chinchilla-style scaling law fitted to each architecture. SMELT's loss drops faster with compute, saving 6.8–18.0% of training FLOPs on the compute-optimal frontier. The advantage transfers to downstream benchmarks beyond what validation loss predicts, is largest on Code, and grows with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens — an inductive bias that may underlie the gains.
Why "looping wins" could not be stated cleanly before
This is the paper's starting point. Reports that a looped model "matches a model several times its size with half the weights" sound wonderful, but the paper puts the obvious objection first: a fraction of the parameter count is not a fraction of the cost (§1). Looping a 12-layer model to 24 executed layers stores half the weights, yet spends roughly a 24-layer model's per-token FLOPs and needs a 24-layer KV cache.
So the terms of the comparison decide the answer. In the survey of prior work (§1, Table 1), most studies hold model size fixed while increasing recurrence, which leaves FLOPs and KV cache uncontrolled. Fix FLOPs instead, and the looped model loses unique parameters — so when it loses, you cannot tell whether looping is bad or whether the parameter cut is.
A fair verdict requires pinning down three budgets simultaneously (§1):
- Per-token FLOPs — sets training and inference cost
- Total parameters — bounds how much knowledge the model can memorise
- KV cache — bounds the context length a deployment can serve
This is where MoE earns its place. MoE decouples total parameters from per-token FLOPs (§1), so a looped model can pay for its second visit by narrowing the hidden dimension , then recover the lost capacity by adding experts — and KV-cache parity is restored separately through head size and the GQA ratio. Three independent knobs, three budgets. If MoE routing itself is unfamiliar, see Mixture of Experts from scratch.
What "matching the budget" actually means
The paper first defines a ruler for sparsity (§3.1). Take a fully active control model that routes every token to all its experts; call its per-token training FLOPs and its total non-embedding parameters . For any configuration with per-token FLOPs and total parameters:
In words: a model that spends half the control's FLOPs is treated as activating half the parameters, and is the share of parameters that go unused. The control maps to by construction. The paper calls this compute-equivalent sparsity and compares matched pairs at , , and .
An MoE router scores every expert for each token and keeps the top 8. Turning scores into shares is a softmax, and whether that distribution is peaked or flat changes the behaviour. Play with it before we use the idea again later.
Now the cost of looping. Repeat a contiguous span of layers out of a total of times, and the executed depth becomes:
That is, "the original layers, plus extra layer executions." FLOPs grow in proportion, so something else has to shrink to stay inside the budget.
The worked example in §3.2 makes it concrete. At the 200M scale with , the Baseline has , hidden dimension , and 192 experts per layer — training FLOPs per token and total non-embedding parameters. Loop the middle 6 layers twice and 18 layers execute, costing six extra layers of FLOPs. So narrow from 1280 to 1056 to bring FLOPs back. That thins every expert's FFN and cuts total parameters, so raise the per-layer expert count from 192 to 288 to recover them. The result: per-token FLOPs land at (+2.9% vs. the Baseline), total parameters at (+0.4%), and KV cache within 4%. Across the twelve sparse grid cells, Appendix A reports maximum mismatches of 3.9% in FLOPs, 1.0% in total parameters, and 3.6% in KV cache.
The three ablations that fixed the recipe
With the comparison controlled, the paper runs three sweeps at the 200M scale.
Which layers to loop (§3.3). Fixing and varying the contiguous middle span from 0 layers (the Baseline) to 12 (full-stack looping), validation loss goes 1.9445 (0) → 1.9257 (6 layers, 50%) → 1.9322 (all 12) at , and 1.8735 → 1.8517 → 1.8601 at . The minimum sits near 50% span at both sparsity levels. The paper is candid about a wrinkle here: the DCLM metrics do not track validation loss in this sweep (at , DCLM Core peaks at 67% span and dips at 50%), so the span was selected on validation loss — which averages over billions of tokens — and the DCLM columns were treated as a consistency check, not a selection criterion. The reported error bars reflect variability across the 10 evaluation seeds only; each configuration was trained once, so training-run variability is not captured.
Comments
Sign in to comment