Why this matters Sharing layers (looping) can increase a model's effective depth at low memory cost, but prior comparisons often conflate depth reuse with extra compute. This paper isolates the architectural effect by matching three critical budgets—per-token FLOPs, total non-embedding parameters, and KV cache—and finds a practical loop recipe (SMELT) for Mixture-of-Experts Transformers that improves scaling behavior and downstream performance.
Key Findings
- SMELT recipe: loop the middle 50% of layers twice, narrow hidden dimension, increase expert count to recover parameters, scale looped residuals by 1/2, and adjust head size/GQA to preserve KV cache. This keeps per-token FLOPs, non-embedding parameters, and KV cache nearly equal to the unlooped Baseline while adding effective depth.
- Scaling and compute gains: across four model sizes up to 54B non-embedding parameters and at multiple sparsity levels, SMELT reaches lower validation loss and fits a separate Chinchilla-style scaling law that drops faster with compute and data. On the compute-optimal frontier SMELT saves about 6.8–18.0% of training FLOPs depending on sparsity and compute budget.
- Downstream and mechanistic evidence: gains transfer beyond validation loss, are largest on code benchmarks, grow with sequence length and in-context examples, and persist across tasks. Mechanistic probes show the second pass reduces an attention sink, redirects attention mass to content-relevant tokens, reuses retrieval coordinates, and amplifies residual writes — suggesting the second visit acts as a refinement step.
Who it's for and trade-offs
Great fit if you build or research large MoE-based language models and can modify architecture and routing: SMELT is practical when you can adjust hidden widths, expert counts, attention head sizing, and GQA settings to enforce budget parity. Look elsewhere if you must stick to dense (non-MoE) stacks, cannot change KV cache/head-size constraints, or prefer architectures where routing divergence across loops is undesirable. Engineering trade-offs include added design complexity (routing, GQA/head adjustments) and reliance on MoE benefits; dense looped models do not show the same scaling advantages.