Looping middle layers of MoE transformers cuts training FLOPs by up to 18%

Researchers introduced a recipe called SMELT that shows reusing the middle layers of a Mixture-of-Experts (MoE) Transformer improves performance even when the FLOPs, non-embedding parameters and KV cache budgets are matched exactly.

Previous comparisons of looped Transformers increased effective depth by repeatedly executing a shared block of layers, but most tests kept model size constant, mixing architectural benefit with pure compute addition. The new paper isolates the architectural contribution by locking three budgets—token-level FLOPs, non-embedding parameters and KV cache size—while applying loops to MoE models.

After a series of ablations the authors arrived at SMELT (Sparse MoE Transformer, middle layers Loop Twice): the loop spans half of the middle layers and runs them twice. The non-looped baseline and SMELT share exactly the same number of non-embedding parameters, the same token-level compute cost and the same KV cache size. Experiments were conducted on four model scales up to 54 billion non-embedding parameters.

Each architecture was fitted with a Chinchilla-style scaling law. The SMELT curve declines more rapidly: on the optimal compute frontier it saves 6.8%–18.0% of the training FLOPs required to reach the same validation loss. The gap widens as the model grows, indicating an advantage that amplifies with scale.

The benefit transfers to downstream tasks beyond what the validation loss predicts, most noticeably in code. The improvement grows with context length and the number of examples in few-shot settings. A mechanistic analysis shows that the second visit to the looped layers reduces the “attention sink” and redirects attention mass toward relevant tokens, an inductive bias that likely explains the observed gain.

The loop is not a trick for cheap model enlargement; when budgets are fully matched it alters attention dynamics and translates into measurable compute savings. SMELT provides a practical recipe that turns depth reuse into a quantitative gain without increasing parameters, FLOPs or KV cache.