Hugging Face、計算資源に最適化された MoE ループトランスフォーマー「SMELT」のスケール法則を発表

Looped Transformer は、共有されたレイヤーブロックを反復実行することで実効的な深さを増大させますが、多くの評価ではモデルサイズを固定した比較が行われ、アーキテクチャ上の利点と追加の FLOPs（演算量）が混同されています。本研究では、Mixture-of-Experts (MoE) Transformer におけるループ化を検討し、トークンあたりの FLOPs、埋め込み層を除く全パラメータ数、および KV キャッシュ量を厳密に一致させた上で評価を行いました。

一連の消融実験（アブレーションスタディ）を通じて、私たちは「SMELT (Sparse MoE Transformer, middle layers Loop Twice)」と呼ばれるレシピを導き出しました。これは、ループ化しないベースラインとすべての計算リソース予算において同等となるよう設計され、中間層の半分だけを 2 回ループさせる手法です。

SMELT は非埋め込みパラメータが最大 54B に達する 4 つのサイズ規模でスケーリングし、各アーキテクチャに対して Chinchilla スタイルのスケーリング則を個別に導出しました。その結果、SMELT の損失は計算リソースに対してより急速に低下し、計算最適化フロンティアにおいてトレーニングに必要な FLOPs を 6.8--18.0% 削減できることが示されました。

この利点は、検証時の損失が予測する範囲を超えて後続のベンチマークにも波及し、特に Code（コード）タスクで顕著に現れます。また、サンプル長やコンテキスト内での例数が増えるほどにその効果は拡大します。

メカニズム分析によると、2 回目の訪問によりアテンションシンクが軽減され、計算リソースがコンテンツに関連するトークンへと再配分されることが明らかになりました。これは観測された性能向上の背後にある帰納的バイアス（インダクティブ・バイアス）である可能性があります。

これらの結果は、予算を厳密に合わせた条件下でもループ化が Transformer の性能向上に寄与することを示しており、深さの再利用を実際の利益に変換する実用的なレシピを提供するものです。