SMELT: 中間層を 2 度巡回する MoE Looped Transformers でトレーニング FLOPs を最大 18% 削減

論文概要

Looped Transformers は共有レイヤーブロックを反復利用して実効的な深さを拡張する手法ですが、従来の評価はモデルサイズを固定した比較が多く、アーキテクチャの優位性と追加 FLOPs（浮動小数点演算数）の混同が生じていました。本研究では、Mixture-of-Experts (MoE) Transformers を対象に、トークンあたりの FLOPs、非埋め込みパラメータ数、KV キャッシュを厳密に一致させた条件下でループ機構を調査しました。その結果、中間層の半分を 2 度巡回させるレシピ「SMELT (Sparse MoE Transformer, middle layers Loop Twice)」を構築し、非埋め込みパラメータ数最大 54B の 4 規模でスケーリングして Chinchilla スケーリング則を検証しました。

関連研究

深層学習における計算効率の向上やパラメータ共有を目的とした再帰型・ループ型 Transformer の研究が多数存在します。しかし、それらの多くはパラメータ数や計算量（FLOPs）の公平な割り当てを欠いており、真のアーキテクチャ上の利点が不明確でした。本研究ではすべてのリソース予算を揃えた厳密な比較を行っています。

新規性と貢献

最大の貢献は、FLOPs、パラメータ数、KV キャッシュを完全に一致させた公平な条件の下で、Looped Transformers がベースラインを上回ることを実証した点です。計算最適フロンティアにおいて 6.8%〜18.0% のトレーニング FLOPs を削減し、その利点がコード生成や長文脈、少ショット学習のタスクにおいて検証ロス以上の規模で現れることを示しました。

提案手法の詳細

SMELT は、Sparse MoE Transformer の中間レイヤーの半分を 2 度ループさせる設計を採用しています。なぜこの設計にしたかという直感として、計算コストを増やさずにパラメータを再利用しつつ、MoE のスパース性と組み合わせることで効率的な表現学習を実現するためです。

評価・考察

最大 54B の非埋め込みパラメータ数を持つ 4 つの規模で評価を行い、独立した Chinchilla スケーリング則をフィッティングしました。計算最適フロンティアにおいて検証ロスの低下が早まり、トレーニング FLOPs の 6.8%〜18.0% 削減を達成しています。また、メカニズム解析により、2 度目の訪問がアテンションシンクを抑制し、コンテンツに関連するトークンへ注意を向ける誘導バイアス（inductive bias）として機能することが示されました。

応用例と今後の展望

実務インパクトとして、国内の AI 開発企業やクラウド事業者（特に大規模言語モデルの事前学習コスト削減を急ぐ開発部門）において、数千億パラメータ規模のモデル学習における 10% 前後の計算コスト削減につながる実用的なレシピとなります。今後の課題としては、より多様なアーキテクチャや長文脈タスクでの汎用性検証が挙げられます。

結論

SMELT は、計算量同等条件の MoE Looped Transformers においてトレーニング効率を大幅に高める有効なレシピであり、深さの再利用が実用的な性能向上をもたらすことを明確に示しました。

注釈

• MoE (Mixture-of-Experts): 入力ごとに異なる少数のサブネットワーク（エキスパート）を選択して処理する仕組み。
• FLOPs: 浮動小数点演算回数。計算量の指標。