One-sentence Summary

Researchers from Tsinghua University, ByteDance Seed, and collaborators propose SMELT, a compute-matched Mixture-of-Experts Looped Transformer that loops the middle half of layers twice while matching per-token FLOPs, total non-embedding parameters, and KV cache, yielding $6.8\text{--}18.0\%$6.8–18.0% training FLOPs savings on the compute-optimal frontier and strong downstream gains, especially on code, via an attention-sink-reducing inductive bias.

Key Contributions

• SMELT, a looped Mixture-of-Experts Transformer, repeats the middle half of layers twice while strictly matching per-token FLOPs, non-embedding parameters, and KV cache, isolating the architectural effect of depth reuse.
• Scaling experiments and Chinchilla-style laws demonstrate that SMELT saves 6.8–18.0% of training FLOPs on the compute-optimal frontier, with downstream benchmark gains that surpass validation loss predictions, particularly on code, long samples, and many-shot in-context tasks.
• Mechanistic analysis reveals that the second loop visit consistently reduces the attention sink and redirects mass toward content-relevant tokens, suggesting looping functions as a refinement step rather than merely adding capacity.

Introduction

Looped Transformers repeat a shared block of layers to increase effective depth without adding parameters, offering a parameter-efficient path to stronger reasoning. However, prior evaluations typically compared looped models against parameter-matched baselines while allowing extra per-token FLOPs and KV cache, leaving it unclear whether looping provides an architectural advantage beyond the additional compute it expends. Attempts to control FLOPs often reduced unique parameters, conflating the effect of weight sharing with capacity loss. The authors resolve this by using Mixture-of-Experts to simultaneously match per-token FLOPs, total parameters, and KV cache between looped and unlooped models. They search the loop design space under these budgets and identify SMELT, a sparse MoE Transformer that loops the middle half of its layers twice at an elevated depth-to-width ratio. Across multiple scales and sparsity levels, SMELT consistently achieves lower loss than the matched baseline, yields compute savings on the Chinchilla frontier, and shows disproportionate gains on structured data and in-context learning, with mechanistic evidence that the second visit amplifies the first and redirects attention away from the sink token.

Method

The authors design a family of sparse Mixture-of-Experts Transformers that reuse a contiguous block of layers through a simple loop mechanism. They refer to the standard unlooped model as Baseline and to the variant that repeats a span of $m$m layers $r$r times as the Looped Transformer. The final recipe, SMELT, emerges from a controlled set of ablations under a strict compute-matching protocol.

Training protocol. All models are decoder-only Transformers with sparse MoE feed-forward layers, each routing every token to its top‑8 experts. Attention uses grouped-query attention (GQA). The authors adopt an internal Baseline family and label four matched scales (100M, 200M, 600M, 1.6B) by their active non-embedding parameter count. In a Looped Transformer, the residual update of each sublayer inside the looped span is scaled by $1/r$1/r to prevent correlated weight-tied updates from inflating the residual stream across repeated visits.

To compare models fairly, the authors introduce a compute-equivalent sparsity metric. Given a fully active control Baseline that activates all experts per token, with per-token training FLOPs $F_0$F0​ and total non-embedding parameters $N_0$N0​, any configuration with FLOPs $F$F and parameters $N$N is characterized by

Here $N_{\mathrm{act}}^{\mathrm{eq}}$Nacteq​ scales the control’s parameter count by the FLOPs ratio, and $S$S is the inactive share, analogous to parameter-ratio sparsity but computed from FLOPs to account for attention cost. A Looped Transformer matched to $S=0$S=0 is not dense; it still routes tokens to a top‑8 subset of a larger expert pool. The experiments report matched Baseline/Looped pairs at $S \approx 85\%$S≈85%, $95\%$95%, and $97\%$97%.

All runs use AdamW with a warmup-stable-decay (WSD) schedule. The stable phase trains at a constant learning rate for 196,075 steps (≈1M tokens per step, global batch size 256). From each stable run, six cosine-decay branches are launched at different steps, each decaying over 10B additional tokens, yielding six token horizons per configuration. Pretraining data is an internal corpus; the stable phase consumes 205B tokens, and each branch adds 10B fresh tokens, so the longest horizon reaches ≈215B tokens without repetition. Matched pairs are trained on identical token sequences. Evaluation uses validation loss across 39 sources, the DCLM Core suite (22 tasks, mean centered accuracy), DCLM Completion (token-weighted cross-entropy on free-form answer tasks), and MMLU 5-shot. Known evaluation items are filtered from the pretraining corpus.

Matching as a compute-allocation problem. To attribute performance differences to architecture rather than resource budgets, the authors closely match three quantities between the Looped Transformer and the Baseline: total parameters (knowledge capacity), per-token FLOPs (training/inference cost), and KV cache (deployment context length). Looping a span of $m$m layers $r$r times increases effective depth to $L_{\mathrm{eff}} = L + (r-1)m$Leff​=L+(r−1)m and raises FLOPs proportionally, while total parameters remain unchanged. To keep FLOPs fixed, the model must be narrowed (reducing hidden dimension $H$H) or shortened (reducing $L$L). Narrowing also shrinks the FFN of each expert, reducing total parameters; to compensate, the per-layer expert count is increased. These adjustments are applied jointly so that per-token FLOPs, total parameters, and KV cache remain within a few percent of the Baseline. The matching procedure is repeated at each sparsity level $S$S.

Ablations and the SMELT recipe. With the matching framework in place, the authors conduct three ablations at 200M scale to determine the optimal loop configuration.

First, they ask which layers should loop. Fixing physical depth at $L=12$L=12 and using two passes, they sweep the length of the contiguous middle segment that loops, from 0 (Baseline) to 12 (full looping). The best performance comes from looping the middle half of the layers, and this result holds across sparsity levels.

Second, they sweep the depth-to-width ratio for both architectures while keeping the loop span at 50% and using two passes. The Looped Transformer achieves its best validation loss at the same physical depth as the Baseline, but it sustains a larger effective depth-to-width ratio because the shared layers receive gradient contributions from multiple visits, making additional serial computation easier to optimize.

Third, they vary the number of loop passes from one to four. Extra visits cost FLOPs, forcing a narrower width under the matching budget. Two loops yield the best performance; more loops degrade results.

These three findings define SMELT, a Sparse MoE Transformer whose middle layers Loop Twice: (1) loop the middle half of layers rather than the full stack, (2) give the looped model a larger effective depth-to-width ratio than the Baseline, and (3) loop exactly twice. This recipe is then scaled to larger model sizes for the main experiments.

Experiment

The SMELT recipe, which loops the middle half of layers twice with a larger effective depth-to-width ratio, is scaled to 54B parameters and compared against an unlooped baseline under matched compute budgets. Fitted scaling laws reveal a steeper compute-loss frontier for SMELT, yielding 6.8–18% training compute savings that grow with scale. Downstream benchmarks amplify the validation loss improvement, especially on structured data like code, and the gain increases with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reuses a core expert subset, writes larger but aligned residual updates, preserves attention patterns while changing value projections, and reduces the attention sink to better retrieve relevant demonstrations.

Looped Transformer studies predominantly loop the middle layers, with spans of 33–50%, and find that weight sharing improves reasoning or perplexity. Dense models typically match only parameter count, leaving FLOPs and KV cache uncontrolled, while MoE-based looped models can match both parameters and per-token FLOPs, achieving strong benchmark results. Scaling analyses show that repeated layers contribute sublinearly (r^0.46 effective contribution) and that optimal loop count grows with compute, while mechanistic experiments reveal that a second loop visit reduces attention sink mass, redirecting focus to informative tokens. Middle-layer looping is the dominant design, with spans of 33–50% across all compared studies. Only MoE looped models match both total parameters and per-token FLOPs, winning 8/9 benchmarks at 3B and 9/9 at 9B. Dense looped models match parameters but not FLOPs or KV cache, making the extra compute from looping uncontrolled. A scaling law for dense looped models finds that r recurrences contribute like r^0.46 unique layers, and the optimal loop count r increases with compute budget. A second weight-tied visit sharply reduces attention sink mass, shifting attention from initial tokens to demonstration answers in in-context learning tasks.

To isolate architectural effects, the Looped Transformer is matched to a Baseline on three budgets: per-token FLOPs are controlled by shrinking hidden dimension or layer count, total parameters are recovered by adding experts, and KV cache is aligned via head size or GQA ratio. After these adjustments, typical residual mismatch stays under 4% for FLOPs and KV cache, and under 1% for total parameters. Adjusting hidden dimension or layer count keeps per-token FLOPs mismatch typically under 4%. Adding experts recovers total parameters to within 1% of the baseline.

Validation loss is minimized when loop span is near 50% for both sparsity levels, so 50% span is adopted as the default. DCLM metrics do not consistently track validation loss, making validation loss the more reliable selection criterion. Looping a middle block outperforms full-stack looping, consistent with specialized roles for initial and final layers. Validation loss reaches its lowest point at 50% loop span for both the 85% and 95% sparsity levels. At 95% sparsity, DCLM Core peaks at 67% span while validation loss is best at 50% span, illustrating a metric disagreement.

Under matched compute and parameter budgets, the Looped Transformer with physical depth 12 and effective depth 18 outperforms all Baseline configurations on validation loss, DCLM Core accuracy, and DCLM Completion perplexity. The Baseline achieves its best validation loss at physical depth 12, and the Looped Transformer's optimal physical depth matches this, while its larger effective depth indicates a preference for a higher effective depth-to-width ratio. This suggests that reusing layers with multiple passes enables more serial computation than an equally deep stack of independent layers. The Looped Transformer with 12 physical layers executed 18 times achieves the best validation loss, highest DCLM Core, and lowest DCLM Completion among all tested configurations. The Baseline's validation loss optimum occurs at physical depth 12; the Looped Transformer matches this physical depth but reaches its peak performance with an effective depth of 18, showing a larger effective depth-to-width ratio.

Under matched FLOPs, looping the middle layers twice yields the lowest validation loss and highest DCLM Core score, outperforming a single visit and deeper loop counts. Adding a third or fourth visit forces a thinner model, causing performance to regress on all metrics. Two visits achieve the best validation loss and DCLM Core, surpassing the baseline and both three- and four-visit configurations. Three and four visits both regress relative to two visits, with no consistent ordering between them, indicating that the thinner model forced by extra loops harms performance.

The experiments evaluate Looped Transformers under controlled compute, parameter, and KV cache budgets, comparing dense and mixture-of-experts variants. Looping the middle 50% of layers with two visits yields the best validation loss and downstream performance, outperforming both single-visit and deeper recurrence while reducing attention sink mass. Mixture-of-experts looped models match both total parameters and per-token FLOPs, winning nearly all benchmarks at 3B and 9B scales, whereas dense looped models only match parameters and exhibit sublinear scaling where repeated layers contribute like r^0.46 unique layers. Overall, the optimal loop count increases with compute budget, and reusing layers enables a higher effective depth-to-width ratio than stacking independent layers.