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
要旨
Looped Transformerは共有された層ブロックを反復することで実効深さを増加させるが、ほとんどの評価は固定モデルサイズで比較しており、アーキテクチャ上の利点と追加のFLOPsが混同されている。我々はトークンあたりのFLOPs、非埋め込みパラメータ総数、KVキャッシュを厳密に一致させながら、Mixture-of-Experts Transformerに対するループ化を研究する。一連のアブレーションを通じて、我々はSMELT（Sparse MoE Transformer, middle layers Loop Twice）と呼ぶ手法に到達した。これは中央半分の層を2回ループさせつつ、3つの予算すべてにおいてループなしのベースラインと一致するものである。我々はSMELTを非埋め込みパラメータ最大54Bの4つの規模に拡張し、各アーキテクチャに対して個別のChinchilla型スケーリング則を適合させた。SMELTの損失は計算量に対してより速く低下し、計算最適フロンティア上で学習FLOPsの6.8--18.0%を節約する。この利点は検証損失が予測する以上に下流ベンチマークへと転移し、Codeで最も顕著であり、サンプル長と文脈内例の数とともに増大する。機構的分析により、2回目の訪問がアテンションシンクを低減し、内容に関連するトークンへ重みを再配分することが示され、この帰納バイアスが観察された性能向上の根底にある可能性がある。これらの結果は、予算一致の条件下でもループ化が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.