SMELT loops the middle half of an Mixture of Experts Transformer’s layers twice while narrowing width and adding experts to hold Per-token FLOPs, parameters, and KV cache constant, saving 6.8–18.0% of training compute on the Chinchilla scaling law-style frontier.
You’re training an MoE language model and someone on your team proposes “just run the middle layers twice, we get more reasoning for free.” The catch: past evaluations of Looped Transformer designs kept parameter count fixed but let per-token FLOPs and KV cache balloon with the extra passes, so you couldn’t tell whether looping was actually a better architecture or just a bigger one in disguise. The dominant baseline to beat is a plain unlooped MoE Transformer sized to spend the same compute. The paper’s contribution is a clean apples-to-apples test: match all three budgets that matter for training and serving, then check whether looping still wins.
The intuition is that MoE decouples storage from compute, which gives you three knobs to rebalance when you add loop passes. Looping the middle 6 of 12 layers twice costs 6 extra layer executions of FLOPs, shrinks total parameters if you do nothing, and would grow the KV cache. So the recipe pays for each: narrow the hidden dimension to give FLOPs back, add more experts per layer to restore total parameters, and shrink the attention head size while raising the Grouped-Query Attention (GQA) ratio to keep KV cache flat. Each residual write inside the loop is also scaled by 1/r so weight-tied updates don’t blow up the residual stream.
The authors search this design space with three ablations at 200M active parameters: which layers to loop (middle 50% wins over full-stack), how deep to make the model (looped models tolerate a larger effective-depth-to-width ratio than the baseline), and how many passes (2 beats 3 or 4, because extra passes force a thinner model under the FLOPs cap). Those three rules define SMELT. They then fit a separate Chinchilla scaling law-style scaling law per architecture across four model sizes up to 54B non-embedding parameters and four sparsity levels, replacing the usual parameter axis with per-token FLOPs and a compute-equivalent sparsity term.
# SMELT recipe vs matched MoE Baseline
L, H, E = baseline.depth, baseline.hidden, baseline.experts_per_layer
span = L // 2 # loop the middle half
r = 2 # loop twice
H_new = shrink_to_match_flops(H, span, r) # narrow width
E_new = raise_to_match_params(E, H, H_new) # add experts
head_size, gqa = adjust_to_match_kv_cache(...) # hold KV flat
for layer in stack_with_loop(L, span, r, scale_residual=1/r):
x = layer(x) # middle layers execute twice
The prevailing default when you want more reasoning depth from a fixed compute budget is to spend the budget on width, more experts, or more tokens, treating weight sharing across depth as a parameter-efficiency trick with no architectural upside. This paper shows the opposite: at matched FLOPs, parameters, and KV cache, running the middle layers a second time is a genuinely better use of compute than pouring the same budget into a wider or shallower stack. The load-bearing evidence is not the headline benchmark lift but the separately-fitted scaling law showing SMELT’s compute exponent γ is 5.5% higher than the Baseline’s, so the gap compounds with scale.
The scaling law fit is the finding that makes the thesis true. SMELT’s compute exponent γ (0.250 vs 0.237) is larger than the Baseline’s, so its frontier drops faster with compute at every sparsity level. Translated into concrete savings via the Compute Efficiency Gain metric: at 10^20 training FLOPs SMELT reaches the Baseline’s loss with 6.8–10.0% less compute, and at 10^21 FLOPs the saving widens to 14.7–18.0%, with the gap compounding across the fitted range.
•
Downstream exceeds validation-loss prediction. SMELT wins 96/96 matched pairs on DCLM Completion, 83/96 on DCLM Core, and 29/30 above-chance pairs on MMLU. Fitting a sigmoid from validation loss to each benchmark on Baseline points, SMELT sits above the fitted curve on every benchmark at every scale, and the residual grows with scale.
•
Structured data benefits most. By-category CE Gain at 10^21 FLOPs: Code 20.4%, Finance 16.8%, Math/STEM 16.6%, Knowledge 14.9%, Web 14.8%.
•
Long samples and in-context learning amplify the gain. Normalized loss improvement on the four longest document buckets is 1.52× the gain on the four shortest. On 16 few-shot tasks the SMELT-vs-Baseline accuracy gap is 0.9 pp at zero-shot and widens to 1.9 pp once demonstrations appear.
•
Mechanistic signature. The second visit reuses Q/K (cosine 0.89–0.93) while V and downstream outputs diverge more; residual writes on visit 2 are 1.2–3.5× larger than visit 1 and point in a similar direction; on Dyck bracket-matching, the Attention sink on the target head collapses from 0.60 to 0.02 while mass on demonstration answers rises from 0.24 to 0.85.
Reach for this when you’re already training an MoE language model at ≥100M active parameters and looking for a compute-efficiency win that doesn’t require touching the data pipeline or the routing algorithm. The change is local to the model definition: shrink the hidden dimension, add experts, adjust head size and GQA to hold KV cache, and re-execute the middle half of layers twice with a 1/r residual scale inside the loop. The rest of your training stack (optimizer, schedule, tokens) stays the same. The gain is largest when your workload is code-heavy or leans on long contexts and in-context examples.
The paper does not release code, weights, or the training corpus. The model family is described as proprietary and Appendix A explicitly notes that the full architecture and training stack cannot be released. What’s reproducible from the paper is the recipe itself and the exact matched configurations (hidden dim, expert count, depth, head geometry) for the 16 grid cells, which is enough to reimplement on top of any MoE Transformer stack you already run.
When you’re deciding between wider, deeper, or more experts under a fixed compute budget, running the middle layers a second time is a real fourth axis, not just a parameter-sharing trick. It pays off most on code and long contexts, and the mechanism seems to be that the second pass keeps the retrieval coordinates fixed while rewriting values and freeing attention mass from the sink.
•
The design ablations (loop span, depth-to-width, loop count) were all run at 200M active parameters. The scaling ladder confirms the locked recipe holds up to 54B non-embedding parameters, but whether “middle 50%, loop twice” is still optimal at frontier scale is untested.
•
Budget matching equates arithmetic FLOPs, parameters, and KV cache, not wall-clock time. Serial re-execution of a shared block and sparse expert routing may leave hardware-utilization gaps that eat into the 6.8–18.0% saving in practice.
•
The scaling law is fit on runs spanning 1.3×10^19 to 2.2×10^21 FLOPs. The 10^22 CE Gain row is flagged as extrapolation and its bootstrap intervals reach zero at the higher sparsity levels, so “the gap keeps compounding” is a claim about the fitted window, not a guarantee at frontier training budgets.