PaperlayerView exact revision on arXiv

Source-linked AI summary

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

Shaowen Wang, Ge Zhang, Kairong Luo, Yuhao Wu, Shaofan Liu, Jiaheng Liu, Wenhao Huang, Shen Yan, Jian Li

arXiv:2609.01343v1cs.LG

TL;DR

Whether looping improves Transformers beyond the extra compute often entailed by recurrent depth has remained unclear. This paper evaluates budget-matched MoE looping and finds that SMELT beats the Baseline across scales, saving 6.8–18.0% of frontier training compute.

  • Problem

    Whether looping offers an architectural advantage beyond extra computation remains unknown, despite prior evaluations often increasing compute or comparing parameter efficiency against larger untied models.

    Source: §1 · p. 1

  • Method

    The paper matches looped and unlooped MoE models under budget constraints, then selects SMELT, which loops the middle half of layers twice.

    Source: §1 · p. 2, §7 · p. 22

  • Results

    6.8–18.0% less training compute on the frontier; SMELT consistently beats the Baseline across four scales and four sparsity levels, with downstream gains exceeding validation-loss predictions.

    Source: §7 · p. 22

  • Takeaways & Limitations

    The findings support depth reuse as a practical design option for MoE Transformers under budget matching.

    Source: §7 · p. 22

    Due to compute constraints, design ablations were conducted at 200M scale, and the optimal loop span or count may differ at larger scales.

    Source: §7 · p. 22

Abstract

from arXiv · show

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.

1 Introduction

The introduction argues that looping should be evaluated under matched compute, parameter, and KV-cache budgets, and presents SMELT as a budget-matched MoE design with improved scaling and downstream performance over a Baseline.

1 Introduction

  • Looping reuses a shared block to increase effective depth, but prior fixed-parameter or FLOP-matched evaluations leave architectural gains confounded or reduce unique parameters.The introduction frames a clean evaluation as matching per-token FLOPs, total parameters, and KV cache, which respectively constrain cost, capacity, and servable context length.Source: §1 · pp. 1–2
  • 14.7% less compute at C = 1021 FLOPs reaches the same loss on SMELT’s compute-optimal frontier; across four model sizes up to 54B non-embedding parameters and four sparsity levels, it beats Baseline loss at every grid point.SMELT’s loss drops faster with compute and data, and its CE Gain increases across the fitted compute range.Source: Figure 1 · p. 2, §1 · p. 2
  • SMELT loops the middle 50% of layers twice while matching Baseline per-token FLOPs, total non-embedding parameters, and KV cache.The recipe narrows hidden dimension, raises expert count, scales looped residuals by 1/2, and adjusts head size and GQA ratio; ablations favor looping the middle half, a larger effective depth-to-width ratio, and two rather than three or four visits.Source: Figure 1 · p. 2, §1 · p. 2
  • On DCLM and MMLU, SMELT outperforms the Baseline in nearly all matched pairs beyond validation-loss predictions, with the largest training-domain gain on Code and larger advantages on longer samples and with more demonstrations.At the 1.6B scale, the figure reports that SMELT leads the Baseline on all six plotted axes.Source: §1 · p. 2, Figure 1 · p. 2
  • On the second visit, experts and attended tokens overlap, aligned residual updates grow, and the attention sink falls as mass shifts toward content-relevant tokens.Values change more than queries and keys; the authors suggest the reduced sink may relate to the observed in-context-learning advantage.Source: §1 · pp. 2–3

2 Related Work

Prior looped-Transformer studies suggest middle-layer repetition can help, but often leave compute or token budgets unmatched. This work uses MoE models to compare looped and unlooped architectures across multiple scales while matching per-token FLOPs, total parameters, and KV cache.

Looped Transformers

  • Prior work often finds middle-layer repetition strongest; Saunshi et al. [54] report worse perplexity than iso-FLOP baselines but stronger reasoning and better perplexity for middle-span than full-stack looping.Universal Transformer introduced iterative weight sharing; Huginn [22] loops the middle 50% in a prelude–recur–coda layout, whereas Ouro [68] loops the full stack. RINS [3] matches training compute using fewer tokens, while most prior comparisons leave repetition's extra FLOPs unmatched.Source: §2 · p. 3

Mixture-of-Experts and loop recurrence

  • Concurrent MoE loop studies show gains at individual scales, while this work first matches all three budgets across block-level models at multiple scales and fits separate scaling surfaces.The matched budgets are per-token FLOPs, total parameters, and KV cache; the scaling analysis uses MoE models.Source: §2 · p. 3

Scaling laws

  • This work fits architecture-specific Chinchilla-style surfaces for MoE while matching per-token FLOPs, total non-embedding parameters, and KV cache to isolate looping's scaling effect.Schwethelm et al. [55] estimate recurrent-block capacity as Nonce + r^0.46 Nrec at fixed effective depth; Prairie et al. [47] hold unique parameters fixed while depth, inference FLOPs, and KV cache grow with r.Source: §2 · p. 4

Attention sinks

  • Prior work treats attention sinks as a trained-model property, whereas this analysis finds that a second tied visit reduces sink mass and redirects attention toward content-relevant tokens.Attention sinks are attention concentration on initial tokens regardless of semantic content; Xiao et al. [63] exploit them for streaming inference, and Qiu et al. [48] show that a query-dependent sigmoid gate mitigates them.Source: §2 · pp. 4–5

3 Design Recipe under Matched Compute

SMELT is selected by looping the middle half of layers twice, at a larger effective depth-to-width ratio than the Baseline. The comparisons match total parameters, per-token FLOPs, and KV cache to isolate architectural differences.

3.1 Training protocol

  • Configurations share the architecture template, optimizer, data, and evaluation; loop settings and matched width, expert, and head hyperparameters differ.Within looped spans, residual updates are scaled by 1/r; compute-equivalent sparsity is defined using FLOPs ratios, and evaluation includes validation loss, DCLM Core and Completion, and MMLU 5-shot.Source: §3.1 · pp. 5–6

3.2 Matching as a compute-allocation problem

  • Matching jointly controls total non-embedding parameters, per-token FLOPs, and KV cache by adjusting width and expert count as looping increases effective depth.For the 200M, S ≈95% example, the matched looped model has 1.37 × 10^9 FLOPs per token (+2.9%), 3.89 × 10^9 total parameters (+0.4%), and KV cache within 4% of Baseline.Source: §3.2 · p. 6

3.3 Looping the middle half beats full looping

  • Validation loss is minimized near a 50% loop span at both tested sparsity levels, so the recipe loops the middle half rather than the full stack.At S ≈95%, DCLM Core instead peaks at 67% span and dips at 50%; the authors select span by validation loss because the DCLM metrics do not track it in this sweep.Source: §3.3 · p. 7, Table 3 · p. 7

3.4 The Looped Transformer prefers a larger effective depth-to-width ratio

  • The Looped Transformer’s preferred depth is 12/18, versus physical depth 12 for the Baseline; the 12/18 variant wins on all three metrics.The 12/18 variant has physical depth 12 and executes 18 layers, supporting the choice to match Looped Transformer physical depth to the Baseline in later experiments.Source: Table 4 · p. 7, §3.4 · pp. 7–8

3.5 Two loops beat three or four

  • Two visits win on all three metrics; third and fourth visits regress under the matched-FLOPs width constraint, defining SMELT’s loop count.The recipe combines a middle-half span, a larger effective depth-to-width ratio, and two visits rather than three or four.Source: Table 5 · p. 8, §3.5 · p. 8

4 Scaling Laws and Compute Savings

SMELT is scaled to 54B non-embedding parameters and evaluated with a separate scaling law for each architecture. Its compute-optimal frontier yields lower loss per unit compute than the Baseline, while requiring a similar compute allocation.

4.1 The SMELT recipe scales across the grid

  • At S ≈97%, SMELT’s training-loss curve stays below the Baseline through most of training at both 22B and 54B non-embedding parameters.Across the 4 × 4 grid, matched Baseline and SMELT pairs show lower validation loss for SMELT at comparable measured compute in every cell.Source: §4.1 · p. 8, Figure 2 · p. 9, Figure 3 · p. 9

4.2 Scaling form and compute efficiency gain

  • The scaling law uses per-token FLOPs F and compute-equivalent sparsity S for capacity, training tokens D for data, and separate fitted surfaces for Baseline and SMELT.Compute efficiency gain is the fraction of compute SMELT saves over the Baseline at a common loss.Source: §4.2 · pp. 9–10

4.3 SMELT saves training compute on the frontier

  • SMELT’s frontier exponent is 0.250 versus the Baseline’s 0.237, a 5.5% increase, so its reducible loss decreases faster with compute.Figure 4 shows steeper SMELT reducible-loss curves at every scale and sparsity.Source: §4.3 · p. 11, Figure 4 · p. 11
  • At C = 10^20, SMELT saves 6.8–10.0% of training compute; at C = 10^21, it saves 14.7–18.0% across sparsity levels.The C = 10^21 estimates are extrapolated beyond the fitted compute window.Source: Table 7 · p. 11, §4.3 · p. 11
  • SMELT’s compute-optimal tokens per dense-equivalent parameter are within 6% of the Baseline’s at every sparsity, with nearly overlapping bootstrap intervals.Larger data coefficient K and exponent c pull the allocation in opposite directions and roughly cancel, so savings come from lower loss at the same allocation.Source: §4.3 · p. 12, Table 8 · p. 12

5 Downstream Performance and Domain Analysis

SMELT’s downstream advantage appears across independent benchmarks and exceeds what its validation-loss improvement predicts. Its gains are largest on structured data, concentrate on long samples, and grow when in-context examples are available.

5.1 Downstream performance amplifies the validation loss gap

  • SMELT wins on DCLM Completion in 96 of 96 matched pairs, DCLM Core in 83 of 96, and MMLU in 29 of 30 qualifying pairs.The MMLU win rate includes pairs where the Baseline scores at least 10 percentage points above chance.Source: §5.1 · p. 12, Figure 5 · p. 12, Figure 6 · p. 13
  • SMELT’s downstream residual is positive on all three benchmarks at every scale, growing with scale on DCLM Completion and DCLM Core; on MMLU, the two larger scales exceed the two smaller ones.The residual compares SMELT with a curve calibrated from Baseline validation-loss and benchmark-score pairs; the sigmoid fits have R2 values of 0.997 for DCLM Completion, 0.974 for DCLM Core, and 0.911 for MMLU. Reading Comprehension has the largest pooled DCLM category residual, but Symbolic Problem Solving overtakes it at 1.6B.Source: §5.1 · pp. 13–14, Figure 8 · p. 14, Figure 9 · p. 14

5.2 SMELT gains most on structured data

  • 20.4% Code leads CE Gain, followed by Finance at 16.8%, Math/STEM at 16.6%, Knowledge at 14.9%, and Web at 14.8%.By Baseline-loss quartile, CE Gain falls from Q1 (18.0%) through Q2 (16.3%) to Q3 (13.8%), then rebounds at Q4 (15.9%).Source: §5.2 · p. 15, Figure 10 · p. 15

5.3 Gain grows with sample length and in-context examples

  • 1.52× is SMELT’s mean normalized gain on the four longest sample buckets relative to the four shortest; adding parameters yields 0.98 and adding experts 0.88.The buckets span document lengths from 32 to 4096 tokens.Source: §5.3 · p. 16, Figure 11 · p. 16
  • 0.9 pp at k = 0 widens to 1.9 pp once demonstrations are provided; on Dyck Languages at k = 32, SMELT scores 29.8% against the Baseline’s 26.4%.The average is over 16 few-shot tasks; on Dyck Languages, both architectures score near 0% without demonstrations.Source: §5.3 · p. 16, Figure 11 · p. 16

6 Inside the Second Pass

The second visit reuses part of the first visit’s expert set, makes larger residual updates, and preserves attention retrieval patterns while changing values. In a Dyck case study and held-out pretraining data, it redirects attention away from the sink token toward content-relevant tokens.

6.1 Do the two visits activate the same experts?

  • At S ≈97%, the router reuses only 2–3 of the top-8 experts between visits, still far above independent-random overlap; at S = 0%, it reuses nearly all 8.The router therefore reuses a core subset while diversifying the rest as the expert pool grows.Source: §6.1 · p. 17, Figure 12 · p. 17

6.2 Does the second visit write more or less into the residual stream?

  • Across all 16 scale–sparsity cells, visit-2/visit-1 norm ratios exceed 1 for every component, ranging from 1.2–3.5×; relative updates are higher at the three larger scales.Same-layer cross-visit residual updates are aligned: cosine ranges from 0.42 to 0.65 (mean 0.56), versus a mean of 0.16 for nonmatching cross-visit pairs, so visit 2 amplifies rather than overrides visit 1.Source: §6.2 · pp. 18–19, Figure 14 · p. 19, Figure 15 · p. 19

6.3 Does the second visit look at the same places or different ones?

  • Across the 4 × 4 grid, Q and K cosine similarities are 0.89–0.93, while V is 0.65–0.74 and attention-output O is 0.49–0.64.Top-8 attended-token sets overlap by 56–66% across visits versus 28–34% for the different-head control, indicating head-specific retrieval-pattern retention.Source: §6.3 · pp. 20–21, Figure 16 · p. 20, Figure 17 · p. 21

6.4 Case study: Dyck language

  • On the Dyck head at 1.6B, BOS attention mass falls from 0.60 to 0.02 while demonstration-answer mass rises from 0.24 to 0.85 on visit 2.The attention map is for a k = 4 example and averages over three prediction targets.Source: §6.4 · p. 21, Figure 18 · p. 22
  • On held-out pretraining data, sink mass falls from visit 1 to visit 2 throughout the repeated block, reversing the Baseline’s increase toward later layers.The measurement covers a 1M-token sample across all five source categories; the Baseline depth trend is consistent with prior findings [52].Source: §6.4 · p. 22

7 Conclusion and Future Work

The paper presents SMELT as evidence that reusing depth can improve MoE Transformers under matched compute. It reports gains across scales, sparsity levels, and downstream evaluations, while identifying limits of the current study and directions for future work.

Conclusion

  • SMELT consistently beats the Baseline across four scales and four sparsity levels, saving 6.8–18.0% of training compute on the frontier.The recipe loops the middle half of layers twice.Source: §7 · p. 22
  • On downstream benchmarks, SMELT’s gain exceeds what validation loss predicts, concentrates on structured data and long samples, and grows with in-context examples.Source: §7 · p. 22
  • The second visit reuses retrieval coordinates, amplifies residual writes, and frees attention mass from the sink.On a 1M-token held-out sample (S ≈85%), it reduces the sink across nearly all heads and layers.Source: §7 · p. 22, Figure 19 · p. 23

Future Work

  • Design ablations were limited to 200M scale, up to 3.9B non-embedding parameters, so optimal loop span or count may differ at larger scales.The study tests only contiguous-block repetition with fully shared weights; whether richer variants preserve or amplify CE Gain under budget matching remains open.Source: §7 · p. 22
  • Budget matching equates arithmetic FLOPs, parameters, and KV-cache size, not wall-clock cost, and re-execution and sparse routing may require systems-level optimization.The causal mechanisms behind the gain also remain open; current probes are a descriptive starting point.Source: §7 · pp. 22–23

Conclusion

  • Weight sharing across depth is a viable axis for improving language models at matched compute, complementary to scaling width, depth, and expert count.Source: §7 · p. 23

Appendix · A Model Configurations

The appendix documents the Baseline and SMELT configuration grids, the conventions used to describe their parameters and sparsity, and the residual budget mismatches handled in scaling analysis.

A Model Configurations

  • Active params counts one physical traversal, including attention, normalization, router weights, and selected top-8 experts; a shared loop block contributes once.S measures compute-equivalent sparsity using FLOPs that count every block execution; sparse levels are labeled S ≈{85%, 95%, 97%}.Source: A Model Configurations · p. 29
  • The proprietary Transformer architecture and training stack cannot be fully released; both architectures nevertheless use the same scaling rules and training stack, differing in looping and budget-matching adjustments.Reported quantities include model scale, hidden dimension, physical depth, expert count, active and non-embedding parameter counts, and compute-equivalent sparsity.Source: A Model Configurations · p. 29
  • For depths L = 10, 12, 20, and 30, the repeated middle-half layers are 3–7, 4–9, 6–15, and 8–22, respectively.For loop-span length m = L/2, the number of preceding layers is ⌊(L −m)/2⌋.Source: A Model Configurations · p. 29
  • At 600M scale, H is 1408 for S = 0 and S ≈85%, but 1512 for S ≈95% and 97%.Budget matching jointly adjusts hidden dimension, expert intermediate dimension, and attention-head geometry to best match per-token FLOPs, total parameters, and KV cache.Source: A Model Configurations · p. 29
  • The 16-cell Baseline grid uses a tuned MoE family, with Scale denoting each matched compute tier by Baseline active-parameter scale.Table 9 reports configurations by hidden dimension and physical depth; parameter counts exclude embeddings.Source: Table 9 · p. 29, A Model Configurations · p. 29
  • Table 10 lists the 16-cell SMELT grid, including loop spans and pass counts, with per-token training-FLOPs and KV-cache ratios relative to corresponding Baselines.Its configurations report hidden dimension and physical depth; parameter counts exclude embeddings.Source: Table 10 · p. 30
  • Across the 12 sparse cells, maximum absolute mismatches are 3.9% in per-token FLOPs, 1.0% in total parameters, and 3.6% in KV cache.Each run enters the scaling fit at measured per-token FLOPs F, and efficiency is compared using separately fitted surfaces at common (C, S) coordinates.Source: A Model Configurations · p. 30

B Training Curves for Every Grid Cell

Across the full 4 × 4 training-curve grid, SMELT ends below its matched Baseline in every cell, extending the two representative cells shown in Section 4.1.

B Training Curves for Every Grid Cell

C Per-Task ICL Accuracy Curves

Figure 24 plots Baseline and SMELT per-task accuracy against the number of in-context examples for the 14 tasks with a significant positive SMELT gain at the maximum shot count.

C Per-Task ICL Accuracy Curves

D Attention-Sink Profiles Across Scales

Across four model scales at matched physical depth, SMELT’s second visit has lower attention-sink mass than its first throughout the repeated block, while the Baseline generally accumulates more sink at later layers.

D Attention-Sink Profiles Across Scales

E Downstream Evaluation Protocol

The appendix defines the downstream benchmark suite, its aggregate metrics, and the evaluation protocol. DCLM Core uses centered accuracy and completion loss, while MMLU is reported separately and few-shot tasks are evaluated across multiple seeds.

Task suite

  • Evaluation covers the 22-task DCLM Core suite and separately reports MMLU 5-shot, which is excluded from Core aggregates.Table 11 maps the tasks to five domain categories and lists the number of in-context demonstrations used.Source: E Downstream Evaluation Protocol · p. 34, Table 11 · p. 35

DCLM Core metric

  • DCLM Core is the unweighted mean of centered task accuracies, where 0 matches the DCLM-v2 reference baseline and 1 is perfect accuracy.The centered-accuracy baselines come from the DCLM project’s v2 release.Source: E Downstream Evaluation Protocol · p. 34

DCLM Completion metric

  • DCLM Completion is the token-weighted mean cross-entropy on gold answer tokens across nine free-form DCLM tasks, with lower loss better.It complements Core with a continuous, fine-grained signal; the nine tasks are CoQA, SQuAD, Jeopardy, LAMBADA, BB QA Wikidata, BB CS Algorithms, BB Dyck Languages, BB Operators, and BB Repeat Copy.Source: E Downstream Evaluation Protocol · pp. 34–35

Multi-seed protocol

  • Few-shot tasks use 10 independent seeds for demonstration selection and ordering, while deterministic zero-shot tasks use one run; truncation preserves BOS.Reported few-shot scores are means over the 10 runs.Source: E Downstream Evaluation Protocol · p. 35

Related papers

Loading 2609.01343v1…