SMELT：計算匹配MoE循環Transformer的擴展定律

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

September 1, 2026

作者: Shaowen Wang, Ge Zhang, Kairong Luo, Yuhao Wu, Shaofan Liu, Jiaheng Liu, Wenhao Huang, Shen Yan, Jian Li

cs.AI

摘要

循環Transformer透過反覆執行共享的層區塊來增加有效深度，但多數評測是在固定模型大小下進行比較，將架構優勢與額外的FLOPs混為一談。我們在緊密匹配每個token的FLOPs、非嵌入參數總數與KV快取的前提下，研究循環機制應用於混合專家（Mixture-of-Experts, MoE）Transformer的效果。透過一系列消融實驗，我們提出一個稱之為SMELT（Sparse MoE Transformer, middle layers Loop Twice，即稀疏MoE Transformer中間層循環兩次）的方法：將中間一半的層循環兩次，同時在上述三項預算上皆與未循環的基線匹配。我們將SMELT擴展至四種規模，最大達540億非嵌入參數，並為每種架構分別擬合了一條Chinchilla式縮放定律。SMELT的損失隨計算量增加而下降得更快，在計算最優前沿上節省了6.8%至18.0%的訓練FLOPs。此優勢會轉移到下游基準測試，且幅度超過驗證損失所預測的程度；在程式碼任務上最為明顯，並隨樣本長度與上下文範例數量增加而增強。機制分析顯示，第二遍通過會減少注意力匯集點（attention sink），並將權重質量重新導向至與內容相關的token；這種歸納偏置可能是觀察到性能提升的底層原因。這些結果表明，即使在預算匹配的條件下，循環仍能改善Transformer，提供了一個將深度重用轉化為可量測收益的實用方法。

English

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.