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

Looped Transformers may gain effective depth without adding unique parameters, but fixed-size evaluations can mistake extra FLOPs for architectural benefit. This paper uses MoE models with matched FLOPs, parameters, and KV cache to develop and scale SMELT, finding compute-efficient and transferable gains. The second visit also reduces attention-sink mass toward content-relevant tokens, suggesting a refinement-like mechanism.

  • Problem

    Prior looping evaluations often compare fixed model sizes while leaving extra FLOPs uncontrolled, so the architectural advantage beyond additional computation remains unknown.

  • Method

    The paper searches loop configurations in MoE Transformers while matching per-token FLOPs, total parameters, and KV cache, then scales SMELT and fits separate Chinchilla-style laws.

  • Results

    SMELT consistently beats the Baseline across four scales and four sparsity levels, saving 6.8–18.0% of training compute on the frontier and transferring gains beyond validation loss.

  • Takeaways & Limitations

    Looping can be a net win for MoE Transformers under budget matching, with gains concentrated on structured data, long samples, and in-context learning.

  • Takeaways & Limitations

    Design ablations run only at the 200M scale, and arithmetic FLOPs, parameters, and KV-cache matching does not establish wall-clock cost parity.

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

Looping increases effective depth through repeated shared layers, but prior comparisons often confound architectural benefits with extra computation. This work uses MoE budget matching to identify SMELT and shows gains across scaling, downstream tasks, and mechanistic measures.

  • Looped Transformers increase effective depth by repeating shared layers, potentially providing more sequential computation without additional parameters.
  • Prior evaluations often leave extra computation uncontrolled, making looping’s architectural advantage beyond additional FLOPs unknown.
  • MoE enables matching per-token FLOPs, total parameters, and KV cache while searching loop span, depth-to-width ratio, and loop count.
  • SMELT loops the middle half of layers twice and reaches lower loss than the Baseline across four scales and four sparsity levels.
  • SMELT outperforms the Baseline on nearly all downstream matched pairs, with especially strong gains on Code, longer samples, and more in-context demonstrations.
  • The second visit reuses retrieval coordinates, amplifies residual updates, and reduces attention-sink mass toward content-relevant tokens, which may relate to in-context-learning gains.

2 Related Work

Related work establishes looping as a way to decouple effective depth from unique parameters and explores MoE recurrence, scaling laws, and attention sinks. SMELT extends this literature through simultaneous budget matching and dynamic analysis of repeated visits.

  • Universal Transformer introduced shared weights across depth, while Huginn established a prelude–recur–coda layout that loops the middle 50% of layers.
  • MoE decouples total parameters from per-token FLOPs, enabling looped models to recover capacity with more experts after narrowing hidden dimensions.
  • Prior looped-MoE studies vary in matched budgets, recurrence structure, and evaluation focus, motivating comparisons across architecture, span, budgets, scaling analysis, and findings.
  • Scaling-law research relates loss to model size and training tokens, with MoE extensions incorporating routing variables such as expert count and sparsity.
  • Attention-sink work characterizes concentration on initial tokens, while SMELT reports that a second tied visit reduces sink mass and redirects it toward content-relevant tokens.

3 Design Recipe under Matched Compute

The design recipe matches looped and unlooped MoE Transformers on per-token FLOPs, total parameters, and KV cache, then selects looping choices through matched ablations. The resulting SMELT recipe loops the middle half twice, with a larger effective depth-to-width ratio than the Baseline.

  • 3.2 Matching as a compute-allocation problem: Matched comparisons hold total parameters, per-token FLOPs, and KV cache approximately equal between the Baseline and Looped Transformer.The budgets represent capacity, computation cost, and deployable context length.
  • 3.2 Matching as a compute-allocation problem: MoE matching narrows hidden width to pay for extra loop executions, then increases expert count to recover total parameters.Compute-equivalent sparsity is defined from FLOPs ratios, accounting for context-dependent attention cost.
  • 3.3 Looping the middle half beats full looping: Validation loss is minimized near a 50% loop span at both tested sparsity levels, so subsequent experiments loop the middle half.Full-stack looping performs worse in the span sweep, while DCLM rankings do not consistently track validation loss.
  • 3.4 The Looped Transformer prefers a larger effective depth-to-width ratio: The Looped Transformer peaks at effective depth-to-width ratio 12/18, whereas the Baseline peaks at physical depth 12.The looped model uses physical depth 12 and executes 18 layers.
  • 3.5 Two loops beat three or four: Two visits outperform one, three, and four on all three metrics; additional visits force a thinner model under matched FLOPs.The regression after two visits is attributed to the width reduction required by extra computation.
  • 3.5 Two loops beat three or four: SMELT combines middle-half looping, a larger effective depth-to-width ratio, and exactly two visits.The recipe is selected from the three matched ablations.

4 Scaling Laws and Compute Savings

The study scales SMELT across matched sparse MoE grids and fits separate Chinchilla-style laws for SMELT and the Baseline. SMELT consistently reaches lower loss, has a steeper compute frontier, and saves 6.8–18.0% of training compute at matched loss.

  • 4.1 The SMELT recipe scales across the grid: The experiment evaluates matched Baseline/SMELT pairs across four scales up to 54B non-embedding parameters and three sparse levels plus a dense control.Six cosine-decay branches per run provide variation along the data axis for scaling-law fits.
  • 4.1 The SMELT recipe scales across the grid: SMELT remains below the Baseline in training loss across the highlighted large-scale cells and their cosine-decay branches.Figure 2 shows the 600M and 1.6B scales at S ≈97%, while Appendix B covers all 16 grid cells.
  • 4.1 The SMELT recipe scales across the grid: At comparable measured compute, SMELT reaches lower validation loss than the Baseline in every sparsity-and-scale cell.Figure 3 uses cumulative training FLOPs, with panels faceted by sparsity and colors marking scale.
  • 4.2 Scaling form and compute efficiency gain: The fitted law models loss as irreducible, capacity-limited, and data-limited terms, with measured per-token FLOPs F and training tokens D as key variables.Separate six-coefficient fits are used for the Baseline and SMELT; the dense control is excluded from the sparse-grid fit because it belongs to a distinct regime.
  • 4.2 Scaling form and compute efficiency gain: SMELT has a higher frontier exponent, γSMELT = 0.250 versus γbase = 0.237, so its reducible loss decreases faster with compute.The corresponding 5.5% higher γ is confirmed visually across scales and sparsity levels.
  • 4.3 SMELT saves training compute on the frontier: 6.8–10.0% of training compute is saved at C = 10^20, increasing to 14.7–18.0% at C = 10^21 for matched loss.The gain compounds with compute; sparsity changes the gain level rather than its growth rate.

5 Downstream Performance and Domain Analysis

SMELT outperforms the Baseline across downstream evaluations under matched budgets, and its advantage exceeds what validation-loss improvements predict. Gains are strongest for structured domains, longer samples, and tasks with more in-context demonstrations.

  • 5.1 Downstream performance: SMELT wins 96 of 96 DCLM Completion pairs, 83 of 96 DCLM Core pairs, and 29 of 30 eligible MMLU pairs.The MMLU win-rate summary excludes pairs whose Baseline is within 10 percentage points of chance.
  • 5.1 Downstream performance: SMELT’s downstream advantage exceeds the improvement predicted by validation loss on all three benchmarks.Residuals are positive at every scale, with excess increasing monotonically with scale on DCLM Completion and DCLM Core.
  • 5.2 Domain analysis: 20.4% CE Gain makes Code the strongest category, followed by Finance at 16.8% and Math/STEM at 16.6%.Knowledge and Web have lower gains of 14.9% and 14.8%; the ranking tracks internal structure, including syntax and long-range dependencies.
  • 5.2 Domain analysis: 18.0% CE Gain occurs in Q1, declining to 13.8% at Q3 before rebounding to 15.9% at Q4.The high-loss rebound is unresolved and may mix noisy sources with hard-but-structured sources.
  • 5.3 Length and in-context examples: 1.52× is SMELT’s normalized-gain ratio for the four longest sample-length buckets versus the four shortest.Baseline parameter scaling has a ratio of 0.98, while adding experts has a ratio of 0.88.
  • 5.3 Length and in-context examples: The architecture gap widens from 0.9 percentage points at k = 0 to 1.9 percentage points at k = 1 and persists through k = 8.A demonstration-sensitive Dyck evaluation also reports 29.8% for SMELT at k = 32 versus 26.4% for the Baseline.

6 Inside the Second Pass

The second visit reuses substantial routing and attention structure while producing larger, aligned residual updates. Mechanistically, it preserves retrieval coordinates but changes values and reduces attention-sink mass, redirecting attention toward relevant content.

  • 6.1 Expert routing: The second visit reuses a core subset of experts while diversifying the remainder as sparsity increases.At S ≈97%, overlap falls to 2–3 experts but remains above independent-random-routing expectations.
  • 6.2 Residual stream: Visit 2 writes larger residual updates than visit 1 at every layer across four scales.Relative updates are also consistently larger at the three larger scales; the 100M model has a few comparable or reversed early layers.
  • 6.2 Residual stream: The second visit writes in a similar direction to the first, amplifying rather than overriding the established signal.All four norm ratios exceed 1, ranging from 1.2–3.5×, and both attention and MoE sub-layers contribute.
  • 6.3 Attention patterns: Q and K remain near the shared-input similarity level, while V and downstream projections diverge more strongly across visits.Q and K have cosine similarity 0.89–0.93, versus 0.65–0.74 for V and 0.49–0.64 for O.
  • 6.3 Attention patterns: 56–66% of top-8 attended tokens overlap across visits, compared with 28–34% for the different-head control.The overlap is therefore head-specific rather than attributable only to positional context.
  • 6.4 Dyck case study: In the Dyck case study, BOS attention falls from 0.60 to 0.02 while demonstration-answer attention rises from 0.24 to 0.85 on visit 2.The second visit redirects mass from the attention sink toward demonstration answers; the broader held-out analysis finds sink reduction beyond Dyck.

7 Conclusion and Future Work

The paper finds that depth reuse improves MoE Transformers under matched budgets, with the SMELT recipe delivering compute savings and downstream gains. It also identifies open questions about scale-dependent designs, hardware cost, richer looping variants, and the causal mechanisms behind the second visit.

  • SMELT loops the middle half of layers twice and consistently beats the Baseline across four scales and four sparsity levels, saving 6.8–18.0% of training compute on the frontier.
  • The advantage extends beyond validation loss, concentrating on structured data and long samples while growing with the number of in-context examples.
  • Mechanistic probes associate the second visit with reused retrieval coordinates, larger aligned residual updates, and reduced attention-sink mass.
  • Future work: The design ablations are limited to the 200M scale, so the optimal loop span or count may differ at larger scales.
  • Future work: The study leaves open whether richer looping variants preserve or amplify the CE Gain and whether hardware-efficiency gaps arise from serial re-execution and sparse routing.
  • Future work: The authors suggest weight sharing across depth as a viable axis complementary to scaling width, depth, and expert count.

8 Contributions

The supplied contribution-section passage lists the paper’s authors.

  • Kairong Luo is listed as an author.
  • Yuhao Wu, Shaofan Liu, and Jiaheng Liu are listed as authors.
  • Wenhao Huang, Shen Yan, and Jian Li are listed as authors.

A Model Configurations

The model-configuration appendix documents matched Baseline and SMELT configurations across a four-by-four grid, with implementation and matching details.

  • Tables 9 and 10 cover four model scales crossed with four sparsity levels and compare SMELT against the Baseline at matched compute tiers.
  • The implementation uses a proprietary Transformer family and reports quantities relevant to matched comparisons, while both architectures share scaling rules and training infrastructure.
  • The configurations report hidden dimension, physical depth, expert counts, active parameters, and non-embedding parameter counts; SMELT additionally specifies the repeated span and number of passes.
  • Matching is coarser at the least sparse reference level, where hardware-aligned changes produce comparatively larger budget shifts.

B Training Curves for Every Grid Cell

The appendix presents training-loss curves for every cell in the four-by-four Baseline–SMELT grid, spanning four scales and four sparsity levels.

  • The appendix contains all 16 grid cells, formed by four scales crossed with a dense-reference control and three sparse levels.
  • Each panel plots training loss against cumulative training tokens for one matched Baseline–SMELT pair.
  • Figures 20–23 show training loss at the 100M, 200M, 600M, and 1.6B scales across all four sparsity levels, with Baseline in blue and SMELT in red.

C Per-Task ICL Accuracy Curves

The shot sweep evaluates SMELT and Baseline across 16 few-shot tasks, focusing on tasks with significant positive maximum-shot gains. Accuracy is averaged across eligible scale–sparsity grid cells and random seeds.

  • 16 few-shot tasks are covered: 15 DCLM Core demonstration tasks plus MMLU.Figure 24 reports the 14 tasks whose maximum-shot SMELT gain is positive and significant at p < 0.05.
  • Curves are averaged first within each eligible scale–sparsity grid cell and then across cells, with significance tested over 10 random seeds.Only cells where both architectures score above chance contribute to the averaging.
  • Figure 24 plots per-task accuracy against the number of in-context examples k for the 14 qualifying tasks.Baseline is shown with dashed blue curves and SMELT with solid red curves.

D Attention-Sink Profiles Across Scales

Across four model scales, SMELT’s second visit produces lower attention-sink mass than its first visit, while the Baseline generally accumulates more sink at later depth.

  • SMELT’s visit 2 remains below visit 1 throughout the repeated middle block at every scale.Baseline and SMELT have the same physical depth in this comparison.
  • The Baseline profile generally rises toward later layers, unlike SMELT’s lower second-visit profile.Figure 25 indexes Baseline physical layers on the bottom axis and SMELT execution layers on the top axis.
  • Across all four scales, SMELT’s second visit reduces segment-start attention-sink mass.The profiles use a 1M-token held-out sample with S ≈85%.

E Downstream Evaluation Protocol

The downstream evaluation combines DCLM Core and separately reported MMLU, using centered accuracy, gold-completion loss, multi-seed few-shot evaluation, and explicit task categorization.

  • Task suite: The evaluation uses the 22-task DCLM Core suite and reports MMLU 5-shot separately rather than folding it into aggregate metrics.DCLM Core tasks are grouped into five official categories.
  • Accuracy metric: DCLM Core is the unweighted mean of centered task accuracies, where 0 is the DCLM-v2 reference baseline and 1 is perfect accuracy.The centered score uses each task’s raw accuracy and DCLM-v2 reference baseline.
  • Completion metric: DCLM Completion is a token-weighted micro-average of gold-answer-token cross-entropies across nine free-form tasks, with lower loss better.This complements Core with a continuous signal less affected by accuracy discretization.
  • Evaluation protocol: Few-shot tasks use 10 independent seeds for demonstration selection and ordering, while deterministic zero-shot tasks use one run.Context truncation preserves the beginning-of-sequence token.
Loading 2609.01343v1…