SMELT: Looped Transformers Outperform Baselines Under Fair Budget Matching

The SMELT paper (arXiv:2609.01343) addresses a key question: does the Looped Transformer's advantage come from its recurrent architecture or simply from extra compute? Researchers from Tsinghua University and ByteDance Seed introduce a budget-matching framework that simultaneously aligns per-token FLOPs, total parameter count, and KV cache size, enabling the first fair comparison across multiple model scales.

Budget Matching Framework

Fair comparison requires controlling three variables: per-token FLOPs (training/inference cost), total parameters (knowledge capacity), and KV cache (context length limit). The method uses Mixture-of-Experts (MoE) to decouple parameters from compute. Matching proceeds in three steps: (1) narrow or shallow the model to offset the extra compute from looping; (2) increase expert count to restore total parameter count; (3) adjust attention configuration (GQA ratio / head dimension) to match KV cache. Figure 1(a) illustrates the core recipe: repeat the middle 50% of layers twice while matching all three budgets.

Optimal Loop Recipe

Small-scale search over loop range, execution depth-width ratio, and loop count yields three principles:

Loop the middle ~50% of layers. Sweeping loop proportion from 0% to 100% shows minimum validation loss at ~50%. The authors hypothesize that first and last layers handle specialized input encoding and output prediction, so keeping them unshared is beneficial (Table 3).

Looped models prefer a larger execution depth-width ratio. With 50% loop range fixed, varying depth-width ratio reveals Looped Transformers achieve optimum at higher ratios. The authors attribute this to parameter sharing: extra execution depth adds no new parameters, and shared parameters receive gradient signals from both shallow and deep positions, easing training (Table 4).

Two loops are optimal. Scanning 1–4 loops shows two loops perform best; three and four loops degrade performance because FLOPs matching forces the model to become too narrow (Table 5).

The combined recipe—loop middle 50% twice with larger depth-width ratio in a sparse MoE—is named SMELT.

Scaling Law Fitting

SMELT is evaluated at four active-parameter scales (100M, 200M, 600M, 1.6B) and four sparsity levels (increasing total parameters up to 54B via more experts), yielding 16 matched baseline/SMELT pairs. SMELT achieves lower validation loss in all 16 configurations (Figure 3). Chinchilla-style scaling laws are fitted separately for baseline and SMELT: L = E + A/(F^a * S^b) + K/D^c, where F is per-token FLOPs, S is sparsity, D is training tokens. SMELT's capacity exponent a (0.3892 vs 0.3703) and data exponent c (0.7011 vs 0.6594) are both larger, meaning loss drops faster with compute and data. On the compute-optimal frontier, SMELT reaches the same loss with 6.8%–18.0% less training compute, savings growing with budget (Figure 1(b), Table 6).

Downstream Task Gains

On DCLM and MMLU benchmarks, SMELT outperforms baseline in nearly all matched settings: 83 of 96 DCLM configurations (Figure 6) and 29 of 30 above-chance MMLU configurations (Figure 7). Crucially, gains exceed what validation loss alone predicts. Fitting a baseline loss-to-downstream curve and plotting SMELT points shows SMELT consistently scores higher at equal validation loss, with the gap widening at larger scales (Figure 9).

Domain and Length Analysis

Validation loss improvement is largest on structured domains: Code > Finance, Math/STEM > Web. Long-context data shows significantly larger gains than short-context, a pattern not seen when merely increasing experts or parameters, indicating it is a unique property of looping (Figures 10(a), 11(a)). In-context learning benefits grow with the number of provided examples, especially on Dyck Languages (Figures 11(b–c)).

Internal Mechanism Analysis

Three perspectives reveal how the second loop pass differs from the first:

Expert routing: The second pass reuses a subset of experts from the first pass; overlap far exceeds random chance (Figure 12).

Residual stream updates: Second-pass updates have larger magnitude than first-pass updates; directions overlap but are not identical, indicating corrective refinement (Figure 13).

Attention retrieval: Query and key cosine similarities between passes are 0.89–0.93, so attention focuses on similar positions. However, value vectors and subsequent outputs change substantially—the model "looks at the same place but reads different content" (Figure 16).

Case Study: Dyck Language

On a bracket-matching task requiring retrieval of example answers, a key attention head attends to BOS on the first pass but shifts to the answer position on the second pass (Figure 18). Meanwhile, attention sink (BOS focus) strengthens with depth in standard Transformers but weakens on the second loop pass, suggesting the first pass builds initial representations and the second pass repurposes attention for content retrieval (Figure 19(b)).

Conclusion

Under budget-matched conditions, SMELT demonstrates consistent advantages across scaling laws, downstream tasks, and mechanistic analyses. Limitations include limited ablation scale and lack of wall-clock system efficiency optimization. The budget-matching methodology and SMELT recipe are proposed as a standard experimental setup for future looped architecture research.