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
cs.LG
2026-09-01
SMELT loops the middle half of an MoE Transformer twice while matching per-token FLOPs, parameters, and KV cache, saving 6.8-18.0% training compute on the Chinchilla frontier.
Looped Transformers buy extra sequential depth by reusing a shared block. Most papers compare at fixed stored size, so the extra FLOPs and KV cache that come with more executed layers are never charged. Schwethelm et al. pin per-token FLOPs instead, and then the looped model loses unique parameters. The clean question is whether looping still helps when three budgets move together: per-token FLOPs, total non-embedding parameters, and KV cache.
Mixture-of-Experts makes that comparison possible. Extra layer visits are paid for by a narrower hidden size; lost capacity is recovered by adding experts; KV cache is held nearly flat with a smaller head size and a higher GQA ratio. ByteDance Seed, Tsinghua, and TokenWave lock that matching into a recipe they call SMELT.
Three 200M-scale ablations fix the recipe:
Residual updates inside the looped span are scaled by 1/r so two correlated writes do not inflate the residual stream. Every MoE layer routes to top-8 experts. Training uses an internal corpus and a WSD schedule up to about 215B tokens, across four active-parameter scales (100M to 1.6B, up to 54B non-embedding parameters) and three sparsity levels near 85%, 95%, and 97%. That yields 96 matched endpoints. Each architecture gets its own sparsity-aware Chinchilla surface.
Inside the fitted window from 10^20 to 10^21 FLOPs, SMELT has a steeper capacity exponent (a=0.3892 vs 0.3703) and data exponent (c=0.7011 vs 0.6594). The frontier exponent γ is 5.5% higher. At matched loss, SMELT saves 6.8-10.0% of training FLOPs at 10^20 and 14.7-18.0% at 10^21.
| Setting | Metric | Result |
| 10^21 FLOPs, S≈97% | training compute | 14.7% less |
| DCLM Completion | 96 matched pairs | 96 wins |
| DCLM Core | 96 pairs | 83 wins |
| MMLU (Baseline ≥10 pp above chance) | 30 pairs | 29 wins |
Downstream gains exceed what the validation-loss gap predicts. Code is the largest domain saving. Relative loss improvement on 512-4096 token documents is 1.52× that on 32-256 token documents. On 16 few-shot tasks the gap is 0.9 pp at k=0 and 1.9 pp once demonstrations are present. On Dyck bracket matching, the second visit drops BOS sink mass from 0.60 to 0.02 and raises mass on demonstration answers from 0.24 to 0.85. Internally, routers and attention positions overlap across visits while residual writes grow and stay aligned: a refinement step, not a second computation.
For MoE pretraining this is a depth-reuse path that does not grow the serving cache: parameters, FLOPs, and KV stay matched. The gains concentrate on code, long samples, and in-context learning, which matches the retrieval-refinement picture inside the second pass. It is an incremental architecture change. The two scaling slopes have already diverged, so the paper-compute saving grows with scale.
The design ablations ran only at 200M (up to 3.9B non-embedding parameters); the best span and loop count may shift at larger scales. Matching arithmetic FLOPs is not matching wall-clock: serial re-execution and sparse routing can eat part of the saving on real hardware. The looping itself is the simplest form, a contiguous fully shared block, with no per-visit adapters or token-wise halting. The 10^22 extrapolation already has confidence intervals that cross zero. Each configuration is trained once, so DCLM error bars reflect eval seeds, not training noise.