SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

AI summary

Overview

Research area: Efficient Transformer architecture and neural scaling laws — specifically Mixture-of-Experts (MoE) models with weight-tied, "looped" layers trained on language modeling at scale.

Technical level: Advanced. The paper assumes familiarity with Chinchilla-style scaling laws, MoE routing, grouped-query attention (GQA), KV cache budgets, and per-token FLOP accounting.

Scope (one sentence): The paper isolates whether looping layers in an MoE Transformer provides a genuine architectural advantage — rather than just extra computation — by matching per-token FLOPs, total non-embedding parameters, and KV cache across a 4×4 grid of scales and sparsity levels, fitting a separate scaling law for each architecture, and quantifying the resulting compute savings.

What This Paper Is About

Looped Transformers reuse a shared block of layers multiple times, which raises effective depth without adding stored parameters. Most prior evaluations compare looped models at fixed model size, so the reported gains mix true architectural benefit with the extra FLOPs and KV cache that repetition costs. This paper asks whether looping still helps once per-token FLOPs, total non-embedding parameters, and KV cache are all held fixed at the same time, and how that advantage scales from small models up to 54B non-embedding parameters.

Key Contributions

1. A budget-matched experimental protocol for looping. The authors define compute-equivalent sparsity S (Eq. 1) and match a Looped Transformer against an unlooped Baseline on three budgets simultaneously — per-token FLOPs, total non-embedding parameters, and KV cache — so performance gaps can be attributed to architecture rather than uncontrolled computation.

2. The SMELT recipe. Through three ablations they derive a concrete design: loop the middle 50% of layers twice, narrow the hidden dimension to pay for the extra visits, raise expert count to recover total parameters, shrink head size / adjust the GQA ratio to hold KV cache nearly unchanged, and scale the looped residual updates. They call it SMELT (Sparse MoE Transformer, middle Layers Loop Twice).

3. Separate scaling laws per architecture across a 4×4 grid. The paper fits a Chinchilla-style surface for each architecture on a grid of four scales (100M / 200M / 600M / 1.6B, physical depths 10 / 12 / 20 / 30) crossed with a dense-reference control (S = 0%) and three sparse levels (S ≈ 85%, 95%, 97%), reporting 32 runs, 96 matched Baseline/SMELT pairs, and 192 evaluation endpoints.

4. A mechanistic account of the second visit. Analysis of routing and attention shows the second pass amplifies the first rather than overwriting it, and that it reduces the attention sink and redirects attention mass toward content-relevant tokens.

Main Findings

• SMELT reaches lower loss at every grid point. Across all four scales and all sparsity levels, SMELT's validation loss lies below the matched Baseline's at comparable measured compute.

• Compute savings on the compute-optimal frontier: 6.8–18.0% of training FLOPs. At C = 10^21 FLOPs, the paper reports SMELT reaching the same loss with 14.7% less compute. The paper also reports that its CE Gain increases across the fitted compute range.

• Downstream gains exceed what validation loss predicts. On DCLM and MMLU 5-shot, SMELT outperforms the Baseline in nearly all matched pairs, with an improvement larger than the validation-loss gap alone would imply. At the 1.6B scale (54B non-embedding parameters, S ≈ 97%), SMELT leads the Baseline on all six reported axes: the five DCLM Core categories (Reading/Comprehension, Knowledge/World Knowledge, Commonsense/Reasoning, Language/Understanding, Symbolic/Problem Solving) and MMLU.

• Code benefits the most among training domains. Among the pretraining domains, the improvement is largest on Code.

• The advantage grows with sample length and in-context examples. Longer samples and more in-context demonstrations both amplify SMELT's edge.

• Three ablation rules define the recipe. (1) Loop the middle half of layers rather than the full stack — validation loss is minimized near 50% span at both S ≈ 85% and S ≈ 95%. (2) Give the looped model a larger effective depth-to-width ratio — the Baseline peaks at physical depth 12 while the Looped Transformer peaks at 12/18. (3) Loop twice rather than three or four times — 2× gives the lowest validation loss (1.9257) and highest DCLM Core (27.57 ± 0.13) at 200M, S ≈ 85%; 3× and 4× regress because budget matching forces a thinner model.

• Mechanistic signature of the second visit. The experts selected and the tokens attended overlap substantially across the two visits, while residual updates grow substantially larger and stay aligned. Within attention, values change more than queries and keys, and the attention sink is broadly reduced.

• Budget matching is tight. In the worked 200M example at S ≈ 95%, the Baseline has L = 12, H = 1280, and 192 experts per layer (1.33×10^9 training FLOPs per token, 3.87×10^9 total non-embedding parameters); the matched SMELT narrows H to 1056 and raises experts to 288, landing at 1.37×10^9 FLOPs (+2.9%), 3.89×10^9 parameters (+0.4%), and KV cache within 4%.

• Positioning against prior work. The paper's Table 1 compares looped studies on architecture, loop span, and whether FLOPs/token, stored parameters, and KV cache are matched. It identifies SMELT as the only entry matching all three budgets while fitting a scaling ladder spanning 0.1B–54B non-embedding parameters; it contrasts with Schwethelm et al. (iso-depth law at 25M–1.6B, where r recurrences count as roughly r^0.46 unique-block equivalents) and Prairie et al. (compute-optimal recurrence at fixed parameter count, 100M–1.3B).

Methodology in Plain English

The problem with comparing looped and unlooped models is that looping is cheap in storage but expensive at runtime: a 12-layer model run to 24 executed layers stores half the weights of a 24-layer model but spends roughly a 24-layer model's per-token FLOPs and needs its full KV cache.

The authors' fix exploits a property of MoE models: total parameters and per-token FLOPs are decoupled. So they narrow the hidden dimension to buy back the FLOPs spent on the extra visits, then add experts (top-8 routing in every MoE layer) to restore the lost parameter capacity. KV cache is restored separately by shrinking head size and adjusting the GQA ratio. They quantify the leftover mismatch — under 4% on FLOPs, under 1% on parameters, under 4% on KV cache — and report exact ratios per configuration in an appendix.

Training uses AdamW with a warmup-stable-decay schedule: a stable phase of 196,075 steps at batch size 256 sequences (about 1M tokens per step, 205B tokens total), from which they branch six cosine-decay schedules at steps 10,000, 20,000, 50,000, 100,000, 150,000, and 196,075, each decaying over 10B additional tokens of fresh data. The longest branch reaches approximately 215B tokens with no repeated data. These six branch points per run supply the variation along the data axis for fitting the scaling law.

Evaluation uses token-weighted cross-entropy on a held-out set of 39 sources grouped into Code, Math/STEM, Knowledge, Finance, and Web, plus the 22-task DCLM Core suite (10 random few-shot seeds per few-shot task, 7 deterministic zero-shot tasks), a DCLM Completion metric defined on the nine free-form-answer tasks, and MMLU 5-shot. Known evaluation items were filtered from the pretraining corpus.

For the scaling law, the authors replace total parameters N with measured per-token FLOPs F and scale the capacity term by compute-equivalent sparsity S: ℒ(F,S,D) = E + A(1−S)^b / F^a + K / D^c. They fit this separately for the Baseline and for SMELT on the full 4×4 grid, then read off the compute saving along the compute-optimal frontier.

Why This Matters

Impact on research. Prior looped-Transformer results were ambiguous: Saunshi et al. found a matched-FLOPs perplexity deficit for the prefix-loop-suffix layout, and Schwethelm et al.'s iso-depth law puts recurrence at roughly r^0.46 unique-block equivalents, implying looping loses at matched compute — but that protocol shrinks the looped model's unique parameters, so the deficit may reflect lost capacity rather than a flaw in looping. By matching FLOPs, parameters, and KV cache at once and fitting separate scaling surfaces, this paper converts scattered single-scale observations into an attributable, extrapolable claim across 0.1B–54B non-embedding parameters. It also reframes looping as a design axis rather than a parameter-efficiency trick, and its mechanistic finding that the second visit dampens the attention sink and shifts mass to content-relevant tokens adds a dynamic observation to a literature that has treated sinks as a static property of trained models.

Real-world applications (implications of the reported results, not claims the paper makes):

• Pretraining cost reduction. A 6.8–18.0% saving in training FLOPs on the compute-optimal frontier translates directly into reduced accelerator time and energy for a fixed target loss.
• Code-oriented models. Code is the domain where the advantage is largest, so code assistants and code-focused pretraining are natural beneficiaries.
• Long-context and retrieval-heavy serving. Because KV cache is matched, the recipe does not trade away servable context length, and the advantage grows with sample length — relevant for long-document and retrieval-augmented workloads.
• Few-shot / in-context-learning systems. The gain grows with the number of in-context demonstrations, which matters for prompting-based pipelines and evaluation suites with many examples.

Industry relevance. The recipe is a practical modification to an existing MoE training stack: narrow width, add experts, adjust head size and GQA ratio, loop the middle half twice, and scale the looped residual updates. It requires no new hardware primitives and no change to per-token serving cost or cache footprint, which makes it deployable within existing MoE training and inference infrastructure. The paper also reports tight residuals on all three budgets (under 4% on FLOPs, under 1% on parameters, under 4% on KV cache), so the engineering accounting a production team would need to verify is small.

Future Directions

• Extend the scaling ladder. The fitted grid runs four scales up to 54B non-embedding parameters; whether the same exponent advantage continues at larger scales is an extrapolation, not a measured result.

• Reconcile with the iso-depth law. The paper notes that Schwethelm et al.'s r^0.46 exponent and Prairie et al.'s fixed-parameter compute-optimal recurrence bracket the question from opposite sides. A unified law covering matched-depth, matched-parameter, and matched-FLOPs regimes is not yet available.

• Explain the downstream gap. SMELT's benchmark improvements exceed what the validation loss gap predicts, and the paper does not report a mechanism that accounts for the excess. The reported links to longer samples, more in-context examples, and reduced attention sink are suggestive but not a complete account.

• Explore loop configurations beyond the locked recipe. The ablations tested span (0–100% at L = 12), effective depth, and loop count (1–4) at 200M; whether the three rules — middle-half span, larger effective depth-to-width ratio, two passes — hold at other scales, or whether designs such as per-layer routing to variable recursion depth (as in Mixture-of-Recursions) interact with the budget-matching protocol, is left open.

• The single-layer, many-recurrence regime. The paper classifies DREAMER-style single-layer recurrence (16–32 applications under matched FLOPs, parameters, and memory) as closer to implicit-depth models than to block-level looping; bridging the two is an open question the paper raises but does not test.

Target Audience

Researchers and engineers working on efficient Transformer architectures, MoE training, and neural scaling laws; teams deciding whether to invest in depth reuse for large-scale language model pretraining; and readers who want a methodological template for running architecture comparisons that control for compute rather than confound it. The paper is most valuable to those already comfortable with scaling-law fitting, MoE routing mechanics, and FLOP/KV-cache accounting; readers looking for an introductory treatment of looped Transformers will find the prior work it surveys (Universal Transformer, Huginn, Ouro) a better entry point.

Authors’ abstract

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.