JA EN
体系推論・高速化
·★ 会員·論文·16分で読めます

論文解説 SMELT — 「同じ層をもう一周」は、予算を揃えても得なのか

層を積むかわりに同じ層をもう一度通す「ループ型Transformer」を、FLOPs・総パラメータ・KVキャッシュの3予算を揃えてMoEで検証した論文。中央半分を2周するSMELTが学習FLOPsを6.8〜18.0%節約したと報告する。

対象textタスクinference

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

一次資料 — この記事の根拠

論文の発表 2026-09-01この解説の公開 2026-09-03同月

SMELT: Scaling Laws for Compute-Matched MoE Looped TransformersShaowen Wang, Ge Zhang, Kairong Luo ほか · 2026-09-01 · v1arXiv:2609.01343論文ページ·PDF
原文の要旨(Abstract)を読む

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.


「もう一周させる」という古くて新しい手

深いモデルが欲しいとき、素直な方法は層を増やすことです。もう一つの手は、層を増やさず同じ層をもう一度通すこと。10階建てのビルの5〜7階だけを2回歩いてから外に出るイメージで、フロアの数(重み)は変わらないのに歩いた階数(実行される層)は増える。これがループ型Transformerです。

今回読むのは "SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers"(arXiv:2609.01343、Shaowen Wang ほか、2026年9月1日)。

要旨はこうです。ループは共有ブロックの反復で実効深さを稼ぐが、従来の評価はモデルサイズ固定の比較が多く、「アーキテクチャの優位」と「余分にFLOPsを使っただけ」を分離できていなかった。そこでMixture-of-Experts(MoE)上でループを検討し、1トークンあたりFLOPs・非埋め込み総パラメータ・KVキャッシュの3つを近い値に揃える。得たレシピが SMELT(Sparse MoE Transformer, middle layers Loop Twice)で、中央半分の層を2回まわしつつ3予算すべてをBaselineに合わせる。非埋め込み54Bまでスケールしてアーキテクチャごとにスケーリング則を当てはめると、SMELTは損失が計算量に対して速く下がり、計算最適フロンティア上で学習FLOPsを6.8〜18.0%節約した。優位は検証損失が予測する以上に下流へ伝わり、Codeで最大、サンプル長と文脈内例示が増えるほど大きくなる。機構解析では、2周目がattention sinkを弱めて注意の質量を内容に関係するトークンへ振り向けていた。

なぜ「ループは得」と言い切れなかったのか

論文はまず釘を刺します — パラメータが半分でも、コストは半分ではない(§1)。12層をループして24層分実行すれば、保存する重みは12層分でも、1トークンFLOPsはおよそ24層分、KVキャッシュも24層分必要です。

だから土俵が問題になります(§1, Table 1)。サイズ固定の評価ではFLOPsとKVが野放しになり、FLOPsだけ固定するとループ側のユニークパラメータが減って原因を切り分けられない。公平な判定には3つを同時に押さえるしかない、というのが論文の立場です。

  1. 1トークンあたりFLOPs — 学習と推論のコスト
  2. 総パラメータ数 — 覚えられる知識量の上限
  3. KVキャッシュ — 配信できる文脈長の上限

ここでMoEが効きます。総パラメータと1トークンFLOPsを切り離せるので、2周目のFLOPsは隠れ次元 HH を細くして支払い、失った容量はエキスパート数で取り戻せる。KVキャッシュはヘッドサイズとGQA比で別途合わせる。独立した3つのつまみがあるから3予算を同時に揃えられます(MoE自体はMixture of Experts を1から理解する)。

予算を「揃える」とはどういう操作か

まず疎さの物差しです(§3.1)。全エキスパートを活性化する対照モデルの1トークンFLOPsを F0F_0、総非埋め込みパラメータを N0N_0、比較対象を FFNN として

Nacteq=FF0N0,S=1NacteqNN_{\mathrm{act}}^{\mathrm{eq}}=\frac{F}{F_0}\,N_0,\qquad S=1-\frac{N_{\mathrm{act}}^{\mathrm{eq}}}{N}
(1)

要するに「対照の半分のFLOPsしか使わない構成は、パラメータも半分しか使っていないとみなす」換算です。SS使われていないパラメータの割合で、対照は定義上 S=0S=0。論文はこれを計算等価スパース度と呼び、S85%,95%,97%S\approx85\%,95\%,97\% で比較します。

MoEのルータは各エキスパートに点数をつけ、softmaxで割合に変えて上位8個を選びます。分布が尖るか平らかで挙動が変わるので、動かしてみてください。

FIG 1MoEのルータは点数をsoftmaxで割合に変え、上位8個だけを使う。尖れば毎回ほぼ同じエキスパートに、平らになれば選択が散らばる。2周目でこの選択がどれだけ重なるかは記事後半で出てくる

ループのコストは、LL 層のうち連続する mm 層を rr 回まわすと実行層数が

Leff=L+(r1)mL_{\mathrm{eff}}=L+(r-1)\,m
(2)

つまり元の LL 層に (r1)×m(r-1)\times m 層ぶんが上乗せされるだけ。FLOPsも比例して増えるので、どこかを削らないと予算に収まりません。

§3.2の例です。200M・S95%S\approx95\% のBaselineは L=12L=12H=1280H=1280、各層192エキスパート、1トークンFLOPs 1.33×1091.33\times10^{9}、総非埋め込みパラメータ 3.87×1093.87\times10^{9}。中央6層を2回まわすと実行18層になるので、HH を1056へ細くしてFLOPsを戻す。今度はエキスパートのFFNが痩せるので、エキスパート数を288へ増やしてパラメータを戻す。結果はFLOPs 1.37×1091.37\times10^{9}(+2.9%)、総パラメータ 3.89×1093.89\times10^{9}(+0.4%)、KVキャッシュ4%以内。疎な12セル全体での最大ずれはFLOPs 3.9%、パラメータ1.0%、KV 3.6%です(付録A)。

レシピを決めた3つのアブレーション

どの層をまわすか(§3.3)。 L=12L=12 で中央の連続区間を0層(Baseline)から12層(全層ループ)まで振ると、検証損失は S85%S\approx85\% で 1.9445 → 1.9257(6層=50%)→ 1.9322、S95%S\approx95\% で 1.8735 → 1.8517 → 1.8601。どちらも50%付近が最小でした。ただしこの掃引でDCLMのスコアは検証損失と同順位にならず、選択は検証損失で行いDCLM列は整合性チェック扱いにした、と断っています。誤差棒は10評価シードのばらつきで、各構成は1回しか学習していません。

最後の `/ r` は飾りではありません。ループ区間の各サブ層の残差更新を でスケールし(§3.1、式8–9)、重みを共有した相関の強い更新が周回ごとに残差ストリームを膨らませるのを防ぎます。

この先にあるもの

§

ここから先は会員限定です

解説記事371本・教科書26章・学生モード48単元・論文精読6本が、月額¥490ですべて読み放題になります。新しい解説は毎日3本ずつ増えます。いつでも解約でき、解約後も期間の終わりまで読めます。

会員の方はログインすると続きが表示されます

参考文献

  1. Shaowen Wang, Ge Zhang, Kairong Luo, Yuhao Wu et al.. (2026-09-01) SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers. arXiv:2609.01343論文ページ·PDF

本記事は上記論文の本文にもとづいて執筆しています。数値・主張は原典を優先してください。

コメント

コメントにはログインが必要です