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

SMELT improves Transformer efficiency by looping middle layers twice under matched compute, parameter, and memory budgets.

Can we improve Transformer performance by looping layers in a Mixture-of-Experts (MoE) model while strictly controlling for compute and parameter budgets?

Looped Transformers increase effective depth by repeating layers, but prior evaluations conflate this architectural gain with the extra compute and memory costs of the repeated steps. SMELT (Sparse MoE Transformer, middle layers Loop Twice) matches per-token FLOPs, total parameters, and KV cache between looped and unlooped models by narrowing hidden dimensions and increasing expert counts. This recipe saves 6.8–18.0% of training compute on the compute-optimal frontier, with downstream gains on code and structured reasoning that exceed validation loss predictions.

Paper Primer

The core move is a design recipe that treats looping as a budget-matching problem: it loops the middle 50% of layers twice, giving the model a higher effective depth-to-width ratio than a standard unlooped baseline. By narrowing the hidden dimension to offset the cost of the second pass and adding experts to recover parameter capacity, SMELT isolates the architectural benefit of depth reuse from the confounding variables of compute and memory.

SMELT achieves superior compute efficiency across all tested scales.

Fitted Chinchilla-style scaling laws show SMELT's loss drops faster with compute than the baseline, yielding a 6.8–18.0% training FLOPs saving on the compute-optimal frontier.

Downstream performance gains exceed validation loss improvements.

Residual analysis on DCLM and MMLU benchmarks shows SMELT consistently outperforms baselines matched for validation loss, with the advantage growing at larger model scales. Positive residuals across all benchmarks, with Code domain gains reaching 20.4%.

Why does looping the middle layers work better than looping the entire stack?

The first and last layers of a Transformer serve specialized roles that benefit from independent parameters; looping only the middle 50% preserves these specialized entry and exit representations while gaining the reasoning benefits of increased effective depth.

Does the second visit simply repeat the first visit's computation?

No; while the second visit reuses a core subset of experts and attention patterns, it writes larger updates to the residual stream and redirects attention mass toward content-relevant tokens, effectively amplifying the first visit's output.

Introduction and Motivation

We expose how prior looped‑Transformer studies conflate architectural gains with extra compute.

Looped Transformers increase a model’s effective depth by iterating a shared block of layers, yet most prior evaluations keep the stored parameter count fixed and thus conflate the architectural benefit with the extra per‑token FLOPs required for the additional passes.

A Looped Transformer re‑uses the same layer block multiple times, so each token experiences more sequential transformations without adding new parameters for each pass.

Effective depth counts how many times a token is processed through a layer stack, i.e., the total number of sequential computation steps it undergoes.

To isolate the true architectural advantage of looping, we must hold three budgets fixed simultaneously: per‑token FLOPs (which dictate training and inference cost), total non‑embedding parameters (which bound knowledge capacity), and KV‑cache size (which limits context length).

Mixture‑of‑Experts (MoE) Transformers make this budget matching feasible: by narrowing the hidden dimension for the looped pass and increasing the number of experts, we can keep FLOPs constant while recovering lost capacity, and we adjust the attention head size and GQA ratio to preserve KV‑cache parity.

The core problem is that prior looped‑Transformer research conflates architectural gains with extra compute, making it unclear whether looping itself provides a genuine advantage.

**Figure 1.** SMELT overview and main results. (a) The SMELT recipe on a compute-matched MoE Baseline: loop the middle 50% of layers twice, narrow the hidden dimension, raise the expert count to recover total parameters, and scale the looped residuals by 1/2, alongside a smaller attention head size at a higher GQA ratio that holds KV cache nearly unchanged across the extra layer executions. Together they closely match per-token FLOPs, total non-embedding parameters, and KV cache (Section 3). (b) Compute-optimal loss under separate Chinchilla-style scaling fits at compute-equivalent sparsity $S \approx 97\%$ (defined in Section 3.1). SMELT's frontier drops faster; at $C = 10^{21}$ FLOPs it reaches the same loss with 14.7% less compute (Section 4). (c) Downstream accuracy at the 1.6B scale (54B non-embedding parameters, $S \approx 97\%$), over the five DCLM Core categories (Appendix E), abbreviated on the axes as Reading (Comprehension), Knowledge (World Knowledge), Commonsense (Reasoning), Language (Understanding) and Symbolic (Problem Solving), with MMLU on its own axis. Each axis is the unweighted mean accuracy over its tasks. SMELT leads the Baseline on all six axes.

**Figure 20.** Training loss at the 100M scale across all four sparsity levels. The Baseline is shown in blue and SMELT in red.

**Figure 21.** Training loss at the 200M scale across all four sparsity levels. The Baseline is shown in blue and SMELT in red.

**Figure 22.** Training loss at the 600M scale across all four sparsity levels. The Baseline is shown in blue and SMELT in red.

**Figure 25.** Same-physical-depth attention-sink profiles on the 1M-token held-out sample ($S \approx 85\%$). Blue circles show the Baseline; open red squares and filled red diamonds show SMELT visits 1 and 2. Bottom axes index Baseline physical layers and top axes index SMELT execution layers. Across all four scales, the second visit reduces segment-start mass even though the unlooped Baseline tends to accumulate more sink at later depth.

**Table 1.** Comparison of related work on looped architectures. We categorize works by their architecture (Dense vs. MoE), the span of the loop, whether they match FLOPs, parameters, or KV cache budgets, the type of scaling analysis performed, and their key findings.

The SMELT Design Recipe

We define the SMELT design recipe, a compute‑matched looped MoE that boosts performance.

Looped Transformers increase effective depth, but comparing them at fixed model size mixes architectural gains with extra FLOPs. To isolate the architectural advantage we need a recipe that keeps compute constant while varying the loop configuration.

We treat per‑token FLOPs as a budget and reshape the model (hidden size, layer count, expert count) so that a looped design spends the same amount of compute as an unlooped baseline.

Extra executions add $6$ layer passes, increasing FLOPs by roughly $6/12 = 0.5$ of the original cost.

To keep $F$ within $+3\%$, shrink $H$ from $1280$ to $1056$, which reduces per‑layer FLOPs by $(1056/1280)^{2}\approx0.68$.

Reduced $H$ also cuts parameters; raise the expert count from $192$ to $288$ to recover the lost parameters.

Re‑compute $N_{\text{eq}}^{\text{act}} = \frac{1.37\times10^{9}}{1.33\times10^{9}} \times 3.87\times10^{9}\approx3.99\times10^{9}$, yielding $S = 1 - 3.99/3.99 \approx 0$ (i.e., effectively matched).

Even though the loop adds extra computation, a modest hidden‑size shrink and extra experts keep the FLOPs budget almost unchanged while preserving total parameter count.

Why can we shrink the hidden dimension without hurting model capacity?

Because the looped architecture re‑uses the same layers multiple times; the extra depth compensates for the reduced width, preserving overall representational power while staying within the FLOPs budget.

SMELT (Sparse MoE Looped Transformer) prescribes three concrete rules that together give a compute‑matched looped model with consistent gains.

Rule 1: Loop layers 4–9 (six layers) twice, giving an effective depth of $12 + (2-1)\times6 = 18$ layers.

Rule 2: Keep physical depth $L=12$ but accept the larger effective depth; no change to $H$ is needed because the FLOPs increase is offset by the loop.

Rule 3: Use exactly two visits; the FLOPs increase is $+2.9\%$, so we shrink $H$ to $1056$ to stay within the budget.

After shrinking $H$, we add experts per layer from $192$ to $288$ to recover the parameter count, achieving $S\approx95\%$.

The toy shows how the three rules interact: looping adds depth, shrinking width controls FLOPs, and adding experts restores parameters, yielding a model that matches the baseline compute while being deeper.

How does SMELT differ from simply looping the entire stack twice?

Looping the full stack inflates FLOPs dramatically, forcing a severe hidden‑size reduction that harms performance. By looping only the middle half, SMELT preserves the specialized roles of the first and last layers while still gaining depth, allowing a modest width shrink and a modest increase in expert count.

Table 3 shows that validation loss is minimized when the looped span covers roughly 50 % of the layers, for both sparsity levels $S\approx85\%$ and $S\approx95\%$. This validates Rule 1: a middle‑half loop outperforms full‑stack looping.

Table 4 demonstrates that the Looped Transformer achieves its best validation loss when the effective‑depth‑to‑width ratio is larger than the baseline’s optimum, confirming Rule 2. The physical depth remains equal to the baseline (12), while the loop adds extra executed layers.

Table 5 reveals that two loop visits give the lowest validation loss and highest DCLM scores; adding a third or fourth visit forces a thinner model under the matched FLOPs budget and degrades performance, establishing Rule 3.

Scaling Laws and Compute Savings

SMELT achieves up to 23.5% compute savings while consistently lowering validation loss at matched compute.

SMELT saves up to 23.5 % of training compute compared to the Baseline at the same validation loss.

Table 7 reports a 23.5 % CE Gain (95 % interval [8, 42]) for S ≈ 97 % at a compute budget of 10²² FLOPs.

The portion of the validation loss that can be lowered by improving the model, after removing the irreducible baseline component.

How does reducible loss differ from the total validation loss?

Reducible loss removes the irreducible term $E$, which is a fixed floor that all models share. The remaining value $L_{\text{reducible}}$ reflects how much loss can still be lowered by changing model capacity, sparsity, or data usage.

SMELT consistently reaches lower validation loss at matched compute.

Downstream Performance and Domain Analysis

SMELT consistently outperforms the baseline on downstream benchmarks across scales and sparsity.

Recall that SMELT follows a compute‑matched recipe that isolates the looped‑layer advantage.

SMELT wins on DCLM Completion in every matched pair (96 / 96).

All 96 compute‑matched comparisons show lower completion loss for SMELT.

SMELT wins on DCLM Core in 83 of 96 compute‑matched pairs.

Core scores are higher for SMELT in the majority of comparisons.

SMELT wins on MMLU in 29 of the 30 pairs whose baseline exceeds chance by 10 pp.

Across the high‑signal subset, SMELT outperforms the baseline.

Residuals are positive on all three benchmarks and increase monotonically with model scale.

Mean residual \bar$\delta$ rises from the 100 M to the 1.6 B scale for DCLM Completion and DCLM Core, and the two larger scales dominate on MMLU.

SMELT’s normalized gain on the four longest length buckets is 1.52 × the gain on the four shortest buckets.

Controls that add experts or parameters show ratios near 1.0 (0.98 and 0.88), indicating the effect is specific to SMELT.

With a single in‑context demonstration (k = 1) SMELT improves accuracy by 1.9 percentage points over the baseline.

The gap widens from 0.9 pp at k = 0 to 1.9 pp at k = 1 and then remains stable up to k = 8.

On the Dyck Languages task SMELT reaches a 29.8 % accuracy gap versus 26.4 % for the baseline at k = 32.

Both models start near 0 % without examples; the gap widens as more demonstrations are provided.

**Figure 5** DCLM Completion (gold-completion micro loss) vs. cumulative training FLOPs across compute-equivalent sparsity levels. Lower is better; SMELT wins all 96 matched pairs.

**Figure 6** DCLM Core vs. cumulative training FLOPs across compute-equivalent sparsity levels. Higher is better; SMELT wins 83 of 96 matched pairs.

**Figure 7** MMLU accuracy vs. cumulative training FLOPs across compute-equivalent sparsity levels; higher is better. The dotted line marks chance accuracy (25%). All 96 pairs are shown; because pairwise comparisons are unstable near the chance floor, the win-rate summary is restricted to the 30 pairs whose Baseline is at least 10 percentage points above chance, of which SMELT wins 29.

**Figure 8** Benchmark scores against validation loss. DCLM Completion is a loss (lower is better); DCLM Core and MMLU are scores (higher is better). The curve fitted to the 96 Baseline endpoints uses the monotone sigmoid family in Eq. 6 for all three metrics ($R^2 = 0.997, 0.974, 0.911$ from left to right). At every scale on average, SMELT lies below the Baseline curve on DCLM Completion and above it on DCLM Core and MMLU: its downstream advantage exceeds what its validation-loss improvement alone predicts.

**Figure 9** Downstream residual of SMELT beyond what the validation loss improvement predicts (Eq. 7), broken down by model scale. In panel (a), the mean residual is positive on every benchmark at every scale; it increases monotonically with scale on DCLM Completion and DCLM Core, while the two larger scales exceed the two smaller ones on MMLU. mn = millinats ($10^{-3}$ nats); pp = percentage points; DCLM Core and the category breakdown use centered accuracy (the DCLM-v2 reference baseline maps to zero), while MMLU uses raw accuracy. Panel (b) fits the same sigmoid separately to each DCLM task before category aggregation. The “All” row pools residuals over endpoints from all four scales rather than fitting a separate row. Positive residuals mean SMELT outperforms the Baseline-fitted curve; negative residuals mean it underperforms that curve.

**Figure 10.** CE Gain on the compute-optimal frontier at $C = 10^{21}$ and $S \approx 95\%$. (a) By domain category. (b) By Baseline-loss quartile. From per-category intercepts: all five categories are positive and Code leads. Within each matched pair, the 39 validation sources are ranked by Baseline per-source loss and split into four near-equal groups (10/10/10/9), from Q1 (lowest loss) to Q4 (highest). Averaged per-source CE Gain is highest in Q1, and the profile is U-shaped with its minimum at Q3.

**Figure 11** SMELT’s extra gain concentrates on long samples and on in-context learning. (a) Loss improvement by document length, normalized to mean 1 within each contrast so that the curves compare shape rather than magnitude. The three families are SMELT against its matched Baseline, and two Baseline-only controls that add experts or add active parameters. SMELT’s gain on the four longest buckets is 1.52× its gain on the four shortest; the two controls stay flat (0.98 and 0.88). (b) Mean accuracy over 16 few-shot tasks (15 DCLM tasks plus MMLU) against the number of in-context examples $k$. The gap is 0.9 pp at $k = 0$ and 1.9 pp once demonstrations are available. (c) Dyck Languages task is demonstration-sensitive: both architectures score near 0% without examples, and the only way to improve is to read the demonstration answers. Here the gap keeps widening with $k$ rather than plateauing, reaching 29.8% against 26.4% at $k = 32$.

Mechanistic Analysis of Looped Layers

We dissect how the second loop visit reshapes routing, updates, and attention.

To understand why the second loop visit improves performance we probe three aspects: expert routing, residual‑stream updates, and attention patterns.

Each Transformer layer adds a correction $\Delta x_{\ell}$ to the residual stream; looping the same physical layer twice lets the second correction build on the first, potentially amplifying or redirecting the signal.

How does the residual update in a looped block differ from a standard pre‑norm layer?

In a standard pre‑norm layer the residual is updated once per layer, so $\Delta x_{\ell}$ is computed from $x_{\ell}$. In a looped block the same physical layer runs twice; the second $\Delta x_{\ell}^{(2)}$ is computed from $x_{\ell}^{(1)}+\Delta x_{\ell}^{(1)}$, giving it a different input and often a larger norm, as the loop amplifies the correction.

Expert overlap between the two visits drops from nearly all 8 experts at $S=0\%$ to only 2–3 experts at $S\approx97\%$, a reduction of about 75%.

Figure 12 shows the distribution of shared experts across sparsity levels.

**Figure 12** Distribution of experts shared between the two visits for the same token and physical gate. Four panels show the $S = 0\%$ dense-reference control and $S \approx \{85\%, 95\%, 97\%\}$. Dashed purple line: expected overlap under independent random routing ($8^2/n$ for top-8 routing over a candidate pool of $n$ experts). At $S = 0\%$, the same token reuses nearly all 8 selected experts; at $S \approx 97\%$, it reuses only 2–3, still far above chance.

The second visit writes a larger residual update, with the norm ratio $\|\Delta x^{(2)}_{\ell}\| / \|\Delta x^{(1)}_{\ell}\|$ averaging about 1.2 across layers.

Figure 13 reports the RMS of residuals, updates, and relative updates for visits 1 and 2.

**Figure 13** Residual flow inside the repeated block (matched $S \approx 85\%$, four scales). Row 1: residual RMS $||x_{\ell}||_2/\sqrt{H}$; row 2: update RMS $||\Delta x_{\ell}||_2/\sqrt{H}$; row 3: relative update $||\Delta x_{\ell}||_2/||x_{\ell}||_2$. Visit 2 (red) writes larger updates than visit 1 (blue) absolutely at every layer (row 2) and relative to the residual at the three larger scales (row 3).

Residual‑update directions between visits are more aligned in SMELT, with mean cosine similarity around 0.56 versus 0.16 for the unlooped baseline.

Figure 15 compares cosine‑similarity heatmaps for the baseline and SMELT models.

**Figure 15** Pairwise residual-update cosine for the 1.6B model. The Baseline panel (left; 30 layers) provides the unlooped reference for ordinary layer-to-layer cosine structure. The SMELT panel (right) uses flattened execution order; dashed lines mark the repeated block. Along the same-physical-layer cross-visit diagonal, cosine ranges from 0.42 to 0.65 (mean 0.56), compared with a mean of 0.16 for nonmatching cross-visit pairs: the two visits write in more closely aligned directions at the same physical layer.

Top‑8 attended‑token sets overlap by roughly 60% between visits, far above the 30% expected for random heads.

Figure 17 reports token‑overlap percentages for same‑head versus control pairings.

**Figure 17** Attention token overlap between the two visits of the repeated block ($S \approx 85\%$, four scales): (a) top-1 retention, (b) top-8 set intersection, and (c) overlap of the full attention distributions. Blue compares the same head across visits; grey is a control that pairs each head with a different head in the same layer and visit, sharing the positional context but not the learned weights. The top-8 sets overlap by 56–66% across visits against 28–34% for the control, so the two visits attend to nearly the same tokens and the overlap is head-specific.

The BOS attention mass drops from 0.60 to 0.02 between visits, a reduction of 0.58.

Figure 18 visualizes the Dyck head’s attention maps for visits 1 and 2.

Demonstration‑answer attention mass rises from 0.24 to 0.85 between visits, an increase of 0.61.

Figure 18 shows the shift of mass onto the demonstration tokens.

**Figure 18** Token-level attention map for the Dyck head at 1.6B, $S \approx 85\%$, Dyck $k = 4$. The Dyck head is the head in the repeated block whose ablation raises Dyck perplexity the most, identified separately at each scale. Rows are the three prediction targets; columns are key positions. Top: visit 1; middle: visit 2; bottom: difference. Orange marks are the 11 demonstration-answer tokens of the four shots, the completions that follow each shot's "Output:" marker. Averaged over the three targets, the second visit cuts BOS mass from 0.60 to 0.02 and lifts demonstration-answer mass from 0.24 to 0.85, so the sink loses almost exactly what the demonstrations gain.

**Figure 14** Eight cross-visit diagnostics over the 4 × 4 grid, with scale on the horizontal axis and compute-equivalent sparsity S on the vertical axis. (a)–(d) Similarity between the two visits: top-8 attended-token overlap, attention-mass overlap, expert-set overlap divided by 8, and residual-update cosine. (e)–(h) Visit-2 / visit-1 norm ratios: the full update $\Delta x_\ell$, the attention write $\Delta x_\ell^{attn}$, the MoE write $\Delta x_\ell^{moe}$, and the RMSNorm output that feeds the next sub-layer. All four norm ratios exceed 1 in all 16 cells: the second visit writes more in every component, at every scale and sparsity.

**Figure 16** Cross-visit cosine similarity at six stages of attention, on the full 4 × 4 grid: (a) the shared pre-norm input $X$, (b) Q, (c) K, (d) V, (e) the attended context $PV$, and (f) the output $O$. Q and K stay at 0.89–0.93, at or above the input’s own 0.84–0.93, so the retrieval coordinates carry over. V falls below the input to 0.65–0.74, and $PV$ and $O$ diverge further.

Conclusion

We wrap up the empirical and mechanistic findings and outline open research directions.

Depth reuse proves beneficial for MoE Transformers when the compute budget is matched, and the SMELT recipe—looping the middle half of layers twice—consistently outperforms the baseline across four model scales and sparsity levels while saving 6.8–18.0 % of training compute.

Beyond validation loss, the gains translate to downstream benchmarks, especially on structured data and long sequences, and grow with the number of in‑context examples. Mechanistic probes show that the second visit reuses retrieval coordinates, amplifies residual writes, and frees attention mass from the sink, indicating that looping acts as a refinement step rather than merely adding capacity.

Current limitations include ablations confined to the 200 M scale and the study of only the simplest contiguous‑block looping; richer variants such as low‑rank adapters, adaptive recursion depth, learned halting, block‑selective sharing, and cross‑token state reuse remain to be examined. Moreover, budget matching aligns FLOPs, parameters, and KV‑cache size but not wall‑clock time, so serial re‑execution and sparse routing may introduce efficiency gaps that require systems‑level optimization.

**Figure a.** Per-head segment-start attention mass at 1.6B. Visit 1 (left) concentrates mass on the segment start; visit 2 (right) reduces it across nearly every head.

**Figure b.** Sink mass at matched physical depth. The bottom axis indexes Baseline physical layers; the top axis indexes SMELT execution layers. The Baseline grows with depth; SMELT visit 2 stays below visit 1.

Model Configurations and Reproducibility

Full model configurations, training curves, task accuracy plots, and evaluation protocol details.

The appendix supplies the complete configuration tables for the Baseline (tuned MoE) and SMELT models, all training‑loss curves across the 4 × 4 grid, per‑task accuracy curves, and the full downstream evaluation protocol.

**Figure 23.** Training loss at the 1.6B scale across all four sparsity levels. The Baseline is shown in blue and SMELT in red.

**Figure 24** Per-task accuracy vs. number of in-context examples $k$ for the 14 of 16 tasks with a significant max-shot gain ($p < 0.05$, paired permutation test). Baseline (dashed, blue) and SMELT (solid, red). Each curve averages first within each grid cell, then across cells.

Implementation note: models are built on a proprietary Transformer family; we report only the quantities that affect the compute‑matched comparison—hidden size $H$, depth $L$, expert count, active‑parameter count, and compute‑equivalent sparsity $S$ (from Eq. 1). Loop span $m$ equals $L/2$, and the repeated block occupies layers 3–7, 4–9, 6–15, or 8–22 for $L=10,12,20,30$, respectively.

Matching‑granularity discussion: at the sparsest reference level the expert pool is smallest, so any change to model width, expert intermediate dimension, or head geometry shifts a larger fraction of the compute budget, yielding a coarser feasible configuration grid. Across the 12 sparse cells the maximum absolute mismatches are 3.9 % in per‑token FLOPs, 1.0 % in total parameters, and 3.6 % in KV‑cache.

**Table 11.** Evaluation tasks by domain category, each cited to its original source. The 22 DCLM Core tasks are grouped into five categories following the official DCLM; MMLU 5-shot is listed separately from the Core set. Parentheses give the number of tasks in each category, and the right column the number of in-context demonstrations used.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers