技術概要：大規模並列シーケンス生成のための構造化再帰ミキサー

問題定義

言語モデルの分野は、トレーニングの効率性と安定性を達成するために、再帰型アーキテクチャ（RNN、LSTM）から非再帰型のシーケンス並列モデル（Transformer）へと移行してきました。しかし、この移行は推論時に重大なコストを伴います。Transformer は $O(n)$ のメモリと逐次的なトークン生成を必要とするため、再帰型モデルと比較して推論スループットと並行性が低下します。一方、再帰型モデルはサンプルあたり $O(1)$ のメモリと効率的な逐次生成を提供しますが、歴史的にトレーニングの効率性と安定性に課題を抱えており、情報圧縮の限界により長いコンテキスト長へのスケーリングに失敗することが多いです。

ハイブリッドアプローチ（Mamba、Gated Delta Net など）はこのギャップを埋めようと試みていますが、しばしば専用デバイスカーネルに依存し、キャッシュのスケーリング問題に依然として直面しています。さらに、出力が迅速に検証可能なタスク（コード生成、数学的証明など）においては、「計算あたりの精度」というパラダイムは、より大きく精度の高いモデルから少数のサンプルを生成するよりも、小さく高速なモデルから多数のサンプルを生成する方が効率的であることを示唆しています。本論文は、情報豊富な入力に対してシーケンス長をスケーリングするには再帰型モデルが不適切である一方、サンプルあたりの一定のメモリフットプリントにより、サンプル（バッチ）次元でのスケーリングには理想的に適しているとの見解を示しています。

手法：構造化再帰ミキサー（SRM）

著者は、専用カーネルやデバイス固有のメモリ管理を必要とせずに、シーケンス並列表現（トレーニング用）と再帰表現（推論用）を代数的に変換するように設計されたアーキテクチャである**構造化再帰ミキサー（SRM）**を導入します。

中核アーキテクチャ

SRM は、アテンションがマスク付き多層パーセプトロン（MLP）に置き換えられた Transformer 類似のアーキテクチャであるマスクドミキサーに基づいています。主要な革新は、トークン混合行列のパラメータ化にあります。

1. 定性的特徴マッピング: 学習済みのマスクドミキサーの分析により、効率的なトークン混合行列は特定の定性的特徴を示すことが判明しました。すなわち、ほぼ同一の値を持つ列、ほぼ同一の値を持つ行、および主対角線からの距離に応じて増加する重み減衰です。
2. 構造的制約: 代数的に変換して再帰形式を可能にするため、SRM はトークン混合行列を行反復または列反復構造に制限します。
   • 行反復: 行が同一の要素を含みます。
   • 列反復: 列が同一の要素を含みます。
   • 減衰: 情報保持を管理するために、対角要素または非対角要素に学習可能な減衰定数 $\lambda$ が適用されます。
3. 代数的変換: これらの制約により、行列乗算 $Y = XM + B$を固定メモリの再帰操作として書き換えることができます（論文の式 2 および 3）。再帰隠れ状態（キャッシュ）は単純な和項となり、シーケンス長に関係なく一定のメモリ（$O(d)$）のみを必要とします。
4. ハイブリッドヘッド: このアーキテクチャは「混合」ヘッドを採用しており、半分のヘッドは行反復行列を使用し、残りの半分は列反復行列を使用し、線形投影を介して結合されます。この構成は、トレーニング効率と情報保持の間の最適なバランスを提供することが判明しました。

トレーニングと推論

• トレーニング: SRM は、Transformer と同様にトークンをシフトさせるシーケンス並列性を使用してトレーニングされます。トレーニング時の二次的な複雑さは、メモリ転送のオーバーヘッドによって実質的に隠蔽されるため、標準的なコンテキスト長においてはスループットが Transformer と同等になります。
• 推論: モデルは、トークンを 1 つずつ生成し、一定サイズのキャッシュを更新する再帰表現に切り替わります。これにより、大規模な KV キャッシュの必要性が排除され、メモリ帯域幅の要件が劇的に減少します。

主要な貢献

1. 定性的特徴マッピング: 行/列反復構造と減衰を強制することにより、二次的な複雑さを持つモデルのトークン混合操作を線形複雑さの再帰モデルにマッピングする手法。
2. デュアルアーキテクチャ設計: トレーニング中はシーケンス並列性を、推論中はバッチ並列性をサポートし、モデルパラメータやキャッシュ構造を変更することなく両者の間で代数的変換を可能にするアーキテクチャ。
3. 最適化された推論エンジン: Mojo/MAXフレームワークを用いた SRM の実装により、標準的な PyTorch 実装に対して顕著な性能向上を実証。
4. 強化学習のためのバランス型リサンプリング: 大規模バッチ強化学習における探索を維持し、トレーニングバッチ内に「良い」サンプルと「悪い」サンプルの混合を確保する、グループ相対方策最適化（GRPO）のためのリサンプリング手法。

実験結果

トレーニング効率

• SRM は、パラメータあたりでは $O(n^2)$ の Transformer やマスクドミキサーよりもわずかにトレーニング効率が劣りますが、多様なデータセット（FineWeb-edu、FineMath）においてMamba モデルよりも効率的です。
• 古いハードウェア（V100 など）でスループットが大幅に低下する Mamba モデルとは異なり、SRM は異なる GPU アーキテクチャ間で一貫したパフォーマンスを発揮します。

ベンチマーク性能

• 精度: 計算同等のベンチマーク（IFEval、SQuAD、GSM8k など）において、Transformer が一般的にトップの成績を収めます。SRM はほとんどのタスクで Mamba モデルを上回りますが、単一サンプルの精度では一般的に Transformer よりも劣ります。
• 情報保持: SRM は Mamba モデルよりも高い情報保持能力と容量を示し、二次的な複雑さを持つモデルに匹敵します。ただし、すべての再帰型モデルと同様に、コンテキスト長が圧縮限界を超えると情報の損失に直面します。

推論スループットと並行性

• スループット: SRM は、メモリ転送の削減により、Transformer に対して莫大なスループット向上を示します。
  • PyTorch（V100）: コンテキスト長に応じて、SRM は Transformer よりも9.66 倍から 43.29 倍高いスループットを達成します。
  • Mojo/MAX（H100）: SRM は vLLM を介して推論された Transformer の12 倍のスループットを達成します。
• 並行性: SRM ははるかに大きなバッチサイズをサポートします。
  • Mojo/MAX（H100）: SRM は Transformer の170 倍の並行性を達成します。
  • Mamba と比較: SRM は、同等の性能を持つ Mamba モデルよりも500 倍以上高い並行性と約 7 倍高いスループットを示します。

強化学習と検証可能な出力

• Pass@k スケーリング: GSM8k において、SRM は Transformer と同様の Pass@k べき乗則に従います。
• 計算効率: 10〜30 倍のスループット増加により、SRM は同じ計算予算でより多くのサンプルを生成できます。計算を一定に保った場合、SRM は Transformer に比べて GSM8k の Pass@k が30% 向上したと報告されています。
• RL トレーニング: バランス型リサンプリングを使用し、大規模バッチ（50 サンプル）で GRPO を用いてトレーニングされた SRM は、小規模バッチ（5 サンプル）でトレーニングされた Transformer よりも高い単一サンプル精度と優れた探索性を達成します。

意義と主張

本論文は、再帰型モデルは情報密度の高い入力に対してシーケンス長のスケーリングには不適切である一方、バッチ次元でのスケーリングには極めて適していると結論付けています。SRM アーキテクチャは、二次モデルのトレーニング効率と再帰モデルの推論効率の間のギャップを成功裏に埋めています。

著者は、SRM が特に迅速な検証が可能なタスクにおいて、「計算あたりの精度」というパラダイムに対する実現可能な道筋を提供すると主張しています。再帰型モデルの一定のメモリフットプリントを活用することで、SRM は Transformer よりも現代ハードウェアの「メモリ壁」の制限をより効果的に克服し、専用カーネルを必要とせずにスループットと並行性において桁違いの改善を達成します。この研究は、将来のマルチスレッドハードウェアの傾向において、メモリアクセスに対する算術演算を最大化する SRM などのアーキテクチャがますます支配的になることを示唆しています。