SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

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.

• The paper introduces SMELT, a sparse decoder-only MoE Transformer that loops the middle layers twice to enhance performance, finding that looping remains beneficial under matched compute and memory budgets, leading to significant reductions in validation loss and training FLOPs.
• The empirical study across varied scales and sparsity levels reveals that SMELT achieves a 6.8–18.0% reduction in training FLOPs while demonstrating advantages in downstream tasks, particularly those with structured data and long documents.
• The benefits observed in SMELT suggest enriched computational reuse, stable attention retrieval across loops, and sensitivity to structured data tasks, offering practical insights for designing more efficient and effective transformer models.

Research question and contribution

“SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers” (2609.01343) studies whether weight sharing across depth provides an architectural advantage when the usual confounds of looping are removed. Prior looped-Transformer comparisons often hold parameter count fixed while allowing repeated execution to increase per-token FLOPs, or hold FLOPs fixed while reducing the number of unique parameters. In either case, performance differences cannot be cleanly attributed to recurrence itself. SMELT instead matches three deployment- and training-relevant budgets: per-token training FLOPs, total non-embedding parameters, and KV-cache size.

The central construction is a sparse MoE decoder-only Transformer in which the middle half of the physical layer stack is executed twice with shared weights. The additional execution is paid for by reducing hidden width; the resulting parameter deficit is recovered by increasing the number of experts, while attention head dimensions and GQA ratios are adjusted to preserve KV-cache capacity. The resulting recipe is named SMELT: Sparse MoE Transformer, middle layers Loop Twice.

The empirical study spans four active-parameter scales—100M, 200M, 600M, and 1.6B—and four compute-equivalent sparsity levels, including a dense-reference control and approximately 85%, 95%, and 97% sparse regimes. Because the largest sparse models store up to 54B non-embedding parameters while activating substantially fewer parameters per token, the experimental design separates active computation from stored capacity. The paper’s principal claim is that looping remains beneficial under these matched budgets: SMELT reaches lower validation loss at every scale and sparsity level, and its fitted compute-optimal frontier implies 6.8–18.0% training-FLOP savings over the measured range (2609.01343).

Budget matching and the SMELT design

The matching problem is nontrivial. Repeating a span of $m$ layers $r$ times increases effective depth from $L$ to $L+(r-1)m$. Since Transformer compute scales approximately with the number of executed layers and the square of hidden width, the looped model must narrow its hidden dimension to remain near the Baseline’s FLOP budget. Narrowing the model also reduces expert and attention parameters, so the authors increase the expert-pool size to restore total parameter count. KV-cache parity is handled separately through head-size and GQA adjustments.

The residual mismatches are small but not zero. Across the sparse configurations used in the scaling analysis, the maximum absolute mismatches are 3.9% for per-token FLOPs, 1.0% for total parameters, and 3.6% for KV cache. Importantly, the authors do not treat these pairs as exactly equal in the scaling-law analysis: each run enters the fit using its measured FLOP value, and comparisons are made between separately fitted surfaces at common compute and sparsity coordinates. This is methodologically stronger than simply labeling configurations “matched” while ignoring the remaining budget differences.

Three ablations determine the recipe. First, at 200M and 12 physical layers, validation loss is minimized near a 50% loop span at both approximately 85% and 95% sparsity. Full-stack looping is inferior, supporting the use of a prelude–recurrent-middle–coda layout. The authors select the span based on validation loss because downstream metrics do not rank the spans consistently; at approximately 95% sparsity, for example, DCLM Core favors a 67% span while validation loss favors 50%. This divergence is an important qualification: the architectural selection criterion is not uniformly aligned with downstream accuracy.

Second, the looped model benefits from a larger effective depth-to-width ratio than the Baseline. At approximately 85% sparsity, the Baseline’s best validation loss occurs at 12 physical layers, whereas the looped model performs best at 12 physical layers and 18 effective executed layers. The authors hypothesize that shared layers receive gradient signals from multiple execution depths, making additional serial computation easier to optimize than an equally deep untied stack under the same budget. This interpretation remains conjectural, since the ablation does not establish a causal relationship between multi-visit gradients and optimization stability.

Third, two visits outperform three or four. At 200M and approximately 85% sparsity, the two-pass model obtains validation loss 1.9257, compared with 1.9385 for three passes, 1.9360 for four passes, and 1.9445 for the Baseline. Its DCLM Core score is 27.57, versus 27.15, 27.30, and 24.92, respectively. Additional visits consume the fixed FLOP budget and force a thinner representation, so the study does not show that two is intrinsically optimal independent of width reduction. It shows that two is optimal within this matched compute-allocation regime.

Scaling-law results

The scale-up experiment trains 32 stable runs and derives six cosine-decay branches from each run, producing 96 matched Baseline–SMELT endpoint pairs. The sparse-grid scaling fits use 72 endpoints per architecture across approximately 85%, 95%, and 97% sparsity. The dense-reference control is excluded because it lies in a distinct regime and degrades the joint fit.

The authors fit separate Chinchilla-style surfaces of the form

$$\mathcal{L}(F,S,D) = E+ \frac{A(1-S)^b}{F^a} + \frac{K}{D^c},$$

where $F$ is measured per-token training FLOPs, $S$ is compute-equivalent sparsity, and $D$ is the number of training tokens. This formulation treats sparsity as a capacity variable while retaining FLOPs as the primary compute axis. Separate coefficients are essential: imposing a shared scaling surface would prevent the architecture from changing either the capacity scaling or data scaling behavior.

SMELT has a larger fitted capacity exponent and data exponent than the Baseline:

Quantity	Baseline	SMELT
Capacity exponent $a$	0.3703	0.3892
Data exponent $c$	0.6594	0.7011
Frontier exponent $r$0	0.237	0.250
Validation-loss RMSE	0.00554	0.00952

The frontier exponent increases by approximately 5.5%, indicating that reducible loss declines more rapidly with compute for SMELT. The fitted irreducible loss is slightly higher for SMELT, 1.4493 versus 1.4439, but the difference is smaller than either fit’s error scale and is therefore not interpreted as a genuine asymptotic disadvantage. The evidence supports faster improvement over the observed compute range, not a lower ultimate loss floor.

At a target budget of $r$1 FLOPs, estimated savings range from 6.8% to 10.0% depending on sparsity. At $r$2 FLOPs, the range increases to 14.7–18.0%. The approximately 97% sparse condition yields 14.7% savings at $r$3 FLOPs, while approximately 85% sparsity yields 18.0%. The extrapolated $r$4-FLOP estimates are larger—19.6–23.5%—but are outside the fitted compute window and have wide bootstrap intervals, including intervals reaching zero at the higher sparsities.

Compute allocation and interpretation of the frontier

The separate scaling surfaces also permit comparison of compute-optimal token-to-parameter allocations. At $r$5 FLOPs, SMELT and the Baseline select nearly identical tokens-per-dense-equivalent-parameter ratios. The point estimates differ by no more than 6% across sparsity levels, and the bootstrap intervals overlap substantially.

This result matters because it locates SMELT’s advantage in the achieved loss surface rather than in a radically different allocation strategy. At approximately 97% sparsity, for example, the Baseline and SMELT prefer 91 and 86 tokens per dense-equivalent parameter, respectively. The authors attribute the similarity to offsetting changes: SMELT has a larger data coefficient, which penalizes small token budgets, but also a larger data exponent, which causes the data-limited term to decline more quickly as training data increases. Thus, the architecture improves compute efficiency primarily by reducing loss at a similar allocation, not by shifting strongly toward either larger models or more data.

The fitted curves nevertheless remain empirical interpolations over a limited range. The reported savings at $r$6 FLOPs are explicitly extrapolations, and predictions at much larger budgets would require additional training runs. This limitation is especially relevant because scaling-law comparisons can be sensitive to functional form, endpoint selection, and fitting procedure.

Downstream transfer and structured-data effects

The validation-loss advantage transfers to independent evaluations. Across 96 matched pairs, SMELT wins DCLM Completion in all 96, DCLM Core in 83, and MMLU in 29 of the 30 pairs for which the Baseline is at least ten percentage points above the 25% chance level. The restriction on MMLU win-rate reporting is appropriate: near the chance floor, small absolute changes are statistically and substantively unstable.

Figure 1: DCLM Completion loss across matched compute conditions; SMELT wins all 96 pairs.

Figure 2: DCLM Core accuracy across matched compute conditions; SMELT wins 83 of 96 pairs.

Figure 3: MMLU accuracy across matched compute conditions, with win-rate analysis restricted to non-floor Baseline cases.

The paper makes a stronger claim than simple downstream transfer. It fits monotone sigmoid calibrations from validation loss to each benchmark using the 96 Baseline endpoints, then measures whether SMELT scores exceed the Baseline-predicted score at the same validation loss. The calibrations fit DCLM Completion, DCLM Core, and MMLU with $r$7 values of 0.997, 0.974, and 0.911, respectively. SMELT’s residual is positive on all three metrics at every scale. The excess grows monotonically with scale for DCLM Completion and DCLM Core; on MMLU, the two larger scales exceed the two smaller scales.

Figure 4: Benchmark performance as a function of validation loss; SMELT lies favorably off the Baseline calibration curve.

This residual analysis suggests that looping changes the mapping from language-modeling loss to task performance, particularly for structured tasks. The domain breakdown is consistent with that interpretation. At $r$8 FLOPs and approximately 95% sparsity, the estimated CE gains are 20.4% for Code, 16.8% for Finance, 16.6% for Math/STEM, 14.9% for Knowledge, and 14.8% for Web. Code therefore benefits most, although the per-domain analysis shares scaling exponents with the aggregate fit and estimates only domain-specific intercepts. The numerical ranking should consequently be interpreted as a structured decomposition of the fitted model rather than as five fully independent scaling-law studies.

Figure 5: SMELT’s downstream residual by aggregate benchmark and model scale.

Figure 6: SMELT’s downstream residual by DCLM domain category.

The source-level analysis is nonmonotonic with respect to Baseline difficulty. CE gain falls from 18.0% in the lowest-loss quartile to 13.8% in the third quartile, then rises to 15.9% in the highest-loss quartile. The authors do not resolve why the hardest sources rebound; the high-loss group may mix noisy data with difficult but highly structured material. This is a genuine open interpretation rather than evidence for a simple “harder data benefits more” law.

Length, in-context learning, and the second visit

SMELT’s gain is disproportionately concentrated on long documents. The normalized improvement over the four longest buckets, 512–4096 tokens, is 1.52 times the improvement over the four shortest buckets, 32–256 tokens. Two Baseline controls do not show the same pattern: adding parameters produces a long-to-short ratio of 0.98, while adding experts produces 0.88. The result distinguishes the effect from generic capacity growth and is consistent with a mechanism that benefits from additional retrievable structure.

In-context learning shows a related effect. Averaged over 16 few-shot tasks, the SMELT–Baseline gap is 0.9 percentage points at zero examples and 1.9 points once demonstrations are present. On the demonstration-sensitive Dyck Languages task, SMELT reaches 29.8% at 32 examples compared with 26.4% for the Baseline. The task-level appendix reports significant maximum-shot gains for 14 of 16 tasks under a paired permutation test. These results do not establish that SMELT performs a qualitatively different learning algorithm, but they show that its advantage increases when the input contains demonstrations that must be retrieved and transformed.

Figure 7: SMELT’s gain increases with document length and in-context examples, with the largest contrast on Dyck Languages.

Mechanistic analysis of the repeated block

The mechanistic probes characterize the second visit as partially redundant in retrieval structure but nonredundant in its residual update. Expert routing has substantial cross-visit overlap. At approximately 97% sparsity, the same token reuses roughly two to three of its eight experts on the second visit—far above the overlap expected under independent random routing—while also diversifying the remainder of its expert set. At the dense-reference level, reuse is close to all eight experts. This behavior is compatible with a shared computational core whose second pass modifies rather than completely reassigns the token’s processing route.

The second visit also produces larger residual updates. Across all 16 scale–sparsity cells, the norm ratios for the full residual update, attention write, MoE write, and RMSNorm output all exceed one, with ratios ranging from 1.2 to 3.5. At 1.6B, same-layer cross-visit residual-update cosine ranges from 0.42 to 0.65, with a mean of 0.56, compared with 0.16 for nonmatching cross-visit pairs. The second visit therefore writes in a direction aligned with the first visit while increasing update magnitude. The authors describe this as amplification, but the measurements remain descriptive: they do not show that amplification is necessary for the downstream gain.

Figure 8: Cross-visit similarity and norm ratios across the full scale–sparsity grid.

Figure 9: Same-physical-layer residual updates are substantially more aligned across visits than nonmatching updates.

The attention decomposition isolates where the two visits diverge. Cross-visit cosine similarity is high for queries and keys, 0.89–0.93, comparable to or greater than the similarity of the normalized input itself, 0.84–0.93. Values are less similar, at 0.65–0.74, and the attended context and output diverge further. This pattern indicates that the second visit largely preserves retrieval coordinates while changing the content representation retrieved through those coordinates.

Figure 10: Q and K remain stable across visits, whereas V, attended context, and output change more substantially.

At the token level, the top-eight attended sets overlap by 56–66% across visits, compared with 28–34% for a different-head control. The overlap is therefore head-specific rather than merely a consequence of positional context. This provides a direct link between the projection-level result and the discrete attention pattern.

Figure 11: Same-head attention retains substantially more attended-token overlap across visits than the control pairing.

The Dyck case study gives the clearest mechanistic observation. For a Dyck head in the 1.6B model with four demonstrations, segment-start BOS attention falls from 0.60 on the first visit to 0.02 on the second, while attention to demonstration-answer tokens rises from 0.24 to 0.85. The second visit thus redirects attention away from an attention sink toward the answer spans needed for bracket completion. On general held-out data, sink mass also falls across nearly all heads and layers during the second visit, despite the tendency of ordinary unlooped Transformers to accumulate sink mass with depth.

Figure 12: In the Dyck case study, the second visit transfers attention mass from BOS to demonstration-answer tokens.

Figure 13: Segment-start attention mass decreases broadly on the second visit.

The evidence supports a specific mechanistic hypothesis: the first pass establishes retrieval coordinates and a preliminary representation, while the second pass uses the updated residual state to preserve those coordinates, alter values, strengthen aligned residual writes, and reduce attention-sink behavior. The causal contribution of any individual component—routing reuse, residual amplification, value divergence, or sink reduction—is not identified.

Limitations and open questions

The main design ablations are conducted at 200M active parameters, with total non-embedding parameter counts up to approximately 3.9B. The scale-up therefore validates a fixed recipe rather than showing that the 50% span and two-pass choice remain optimal at every scale. The largest models are evaluated through the scaling ladder, but the loop-design search itself is not repeated at those scales.

The architecture and training stack are proprietary, preventing full reproduction and limiting assessment of whether the result depends on implementation-specific normalization, routing, optimizer, or data choices. The corpus is internal, and the experiments compare only one looped block layout with fully shared weights. Adaptive recursion, per-visit adapters, selective sharing, learned halting, and token-level recurrence are not tested under the same triple-budget protocol.

The budget criterion also matches arithmetic FLOPs, parameter count, and KV-cache size rather than wall-clock latency, communication cost, memory bandwidth, or sparse-kernel efficiency. Looping introduces serial execution, while MoE introduces routing and dispatch overhead; either could alter the practical systems-level tradeoff. Finally, the mechanistic analysis is correlational. The attention-sink reduction and aligned second-pass updates are robust descriptive signatures, but interventions that selectively preserve or suppress these behaviors are required to establish their causal role.

Conclusion

SMELT provides evidence that block-level recurrence can improve MoE LLMs after matching per-token FLOPs, total parameters, and KV-cache size. Its selected configuration—repeating the middle half of the layers twice with residual scaling and compensatory width, expert, and attention adjustments—wins across a 4-by-4 scale–sparsity grid and yields fitted compute savings of 6.8–18.0% within the supported frontier range (2609.01343). The gains are strongest on structured domains, long samples, and demonstration-dependent tasks, while mechanistic measurements associate them with stable retrieval patterns, modified value representations, larger aligned residual writes, and reduced attention-sink mass on the second visit.

The paper leaves a focused question: whether these gains persist, and perhaps increase, when loop span, recursion count, and sharing pattern are jointly re-optimized at larger scales and evaluated under wall-clock rather than arithmetic-compute matching.