Search papers, labs, and topics across Lattice.
This paper investigates the performance of Looped Transformers, specifically focusing on Mixture-of-Experts (MoE) architectures while maintaining equivalent per-token FLOPs and parameter counts. The authors introduce SMELT, a method that loops the middle layers of the transformer twice, resulting in significant reductions in training FLOPs and enhanced performance on downstream tasks, particularly in code-related benchmarks. Their findings reveal that this looping mechanism not only accelerates convergence but also improves the model's ability to focus on relevant tokens, thus providing a practical approach to leveraging depth reuse in transformer architectures.
Looping the middle layers of MoE Transformers can save up to 18% in training FLOPs while boosting performance beyond validation loss predictions.
Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size, conflating architectural advantage with extra FLOPs. We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache. Through a series of ablations, we arrive at a recipe we call SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets. We scale SMELT across four sizes up to 54B non-embedding parameters and fit a separate Chinchilla-style scaling law for each architecture. SMELT's loss drops faster with compute, saving 6.8--18.0\% of training FLOPs on the compute-optimal frontier. 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. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens, an inductive bias that may underlie the observed performance gains. These results show that looping can improve Transformers even under budget matching, offering a practical recipe that turns depth reuse into measurable gains.
