SMELT loops MoE middle layers, cutting training FLOPs 6.8-18.0%

Looping a transformer's middle layers twice cuts 6.8 to 18.0 percent of training FLOPs at compute-optimal scale, according to a new arXiv paper that scales the recipe across four Mixture-of-Experts models up to 54B non-embedding parameters.

The setup is called SMELT (Sparse MoE Transformer, middle layers Loop Twice), and the point of the paper is the comparison, not the trick. Most prior looped-Transformer evaluations "compare at fixed model size, conflating architectural advantage with extra FLOPs," the abstract states. SMELT instead matches its unlooped baseline on per-token FLOPs, total non-embedding parameters, and KV cache, then fits a separate Chinchilla-style scaling law to each architecture.

The gains do not stop at validation loss. "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," the authors write. A mechanistic look attributes it to what changes on the second pass: "the second visit reduces the attention sink and redirects mass toward content-relevant tokens."

The paper reports no inference-time latency numbers for the looped setup, and 54B non-embedding parameters remains below the scale of published frontier MoEs.