JA EN
LearnInference & Serving
·★ MEMBER·PAPER·15 min read

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.

ModalitytextTaskinference

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

Primary source — what this article is built on

undefined2026-09-01undefined2026-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·PDF
undefined

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):

  1. Per-token FLOPs — sets training and inference cost
  2. Total parameters — bounds how much knowledge the model can memorise
  3. 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 HH, 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 F0F_0 and its total non-embedding parameters N0N_0. For any configuration with per-token FLOPs FF and NN total parameters:

Nacteq=FF0N0,S=1NacteqNN_{\mathrm{act}}^{\mathrm{eq}}=\frac{F}{F_0}\,N_0,\qquad S=1-\frac{N_{\mathrm{act}}^{\mathrm{eq}}}{N}
(1)

In words: a model that spends half the control's FLOPs is treated as activating half the parameters, and SS is the share of parameters that go unused. The control maps to S=0S=0 by construction. The paper calls this compute-equivalent sparsity and compares matched pairs at S85%S\approx85\%, 95%95\%, and 97%97\%.

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.

FIG 1An MoE router turns per-expert scores into shares with a softmax and keeps only the top 8. A peaked distribution means the same experts get picked every time; a flat one spreads the choice. How much SMELT's two visits overlap in that choice is measured later in the paper

Now the cost of looping. Repeat a contiguous span of mm layers out of LL a total of rr times, and the executed depth becomes:

Leff=L+(r1)mL_{\mathrm{eff}}=L+(r-1)\,m
(2)

That is, "the original LL layers, plus (r1)×m(r-1)\times m 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 S95%S\approx95\%, the Baseline has L=12L=12, hidden dimension H=1280H=1280, and 192 experts per layer — 1.33×1091.33\times10^{9} training FLOPs per token and 3.87×1093.87\times10^{9} total non-embedding parameters. Loop the middle 6 layers twice and 18 layers execute, costing six extra layers of FLOPs. So narrow HH 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 1.37×1091.37\times10^{9} (+2.9% vs. the Baseline), total parameters at 3.89×1093.89\times10^{9} (+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 L=12L=12 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 S85%S\approx85\%, and 1.8735 → 1.8517 → 1.8601 at S95%S\approx95\%. 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 S95%S\approx95\%, 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.

Those three rules define SMELT. The implementation sketch is short:

What's behind this

§

Members-only from here

371 walkthroughs, 26 textbook chapters, 48 student units and 6 close readings — all included for $4.99/mo, with three new explainers every day. Cancel any time; access runs to the end of the period.

Already a member? Sign in to keep reading

References

  1. Shaowen Wang, Ge Zhang, Kairong Luo, Yuhao Wu et al.. (2026-09-01) SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers. arXiv:2609.01343Paper page·PDF

This article is written from the source paper above. Where they differ, the original is authoritative.

Comments

Sign in to comment