技術概要：Parallax - 言語モデリングのためのパラメータ化された局所線形アテンション

問題定義

大規模言語モデル（LLM）は、その中核的な計算プリミティブとして Softmax アテンションに依存しており、このメカニズムはトランスフォーマーの導入以来、構造的に変化していません。線形アテンションや状態空間モデル（SSM）のような効率的な代替手段は、二次以下の計算量を提供しますが、コンテキスト内情報検索タスクにおいて Softmax アテンションに一貫して劣ります。この性能格差は、それらの設計における根本的なトレードオフを示唆しています。

理論的枠組み、特にテスト時回帰枠組みは、アテンションをコンテキスト内回帰ソルバーとして解釈します。この枠組みの中で、**局所線形アテンション（LLA）**は、Softmax アテンションに対する優れた代替案として提案されました。LLA は、Softmax アテンションで使用される局所定数推定量（Nadaraya-Watson）を局所線形推定量に置き換えることで、理論的に厳密に優れたバイアス - 分散のトレードオフと連想記憶容量を実現します。しかし、LLA は以下の 3 つの主要な課題により、LLM の事前学習への成功したスケーリングがなされていませんでした：

1. 計算オーバーヘッド： 正確な LLA は、トークンごとに並列共役勾配（CG）ソルバーを用いて線形システムを解く必要があり、大幅な I/O および計算コストをもたらします。
2. 数値的不安定性： CG ソルバーは精度形式に敏感であり、低精度ハードウェア環境での管理が困難です。
3. 正則化のトレードオフ： 数値的安定性を確保しつつ、モデルが Softmax アテンションへ退化しないように正則化パラメータ $\lambda$ を調整することは、実用的な事前学習設定において容易ではありません。

手法：Parallax

著者は、LLM の事前学習に対してスケーラブルで効率的かつ数値的に安定したパラメータ化された局所線形アテンションであるParallaxを導入します。

中核メカニズム

Parallax は、LLA を Softmax アテンションに対する加法的な補正として再定式化します。各クエリに対して最適な補正ベクトル $\rho^*_i$ を求める（共分散行列の逆行列を必要とする）代わりに、Parallax は層入力 $x_i$ からプローブベクトル $\rho_i = W_R x_i$ への直接マッピングを学習します。
出力は以下のように計算されます：
$$o^{PLX}_i = o^{SA}_i - \Sigma^{(i)}_{KV} \rho_i$$
ここで、$o^{SA}_i$ は標準的な Softmax アテンション出力、$\Sigma^{(i)}_{KV}$ は KV 共分散行列です。

主要な設計選択には以下が含まれます：

• ソルバーの排除： $\rho_i$ をパラメータ化することで、トークンごとの共役勾配ソルブの必要性を排除しました。
• 境界増幅の除去： 著者は境界増幅係数 $\eta_i$ をゼロに設定しました。正確な LLA において $\eta_i$ は境界バイアスを補償しますが、パラメータ化された設定ではトレーニングの不安定化を引き起こす可能性があります。これを除去することで、共分散補正能力を維持しつつフォワードパスを簡素化しました。
• 統合フレームワーク： Parallax は、帯域幅、プローブ構築、アフィン構造によって接続されたアテンションメカニズムのファミリー内で位置づけられています。これは、非パラメトリック手法（LLA）とパラメトリック手法（線形アテンション、MesaNet）の間のギャップを埋めます。

ハードウェア対応アルゴリズム

効率性の課題に対処するため、著者は FlashAttention（FA）と互換性のあるストリーミングアルゴリズムを開発しました。

• 演算強度： このアルゴリズムは、同じ KV ストリームを再利用しながら計算量を約 2 倍に増やすことで、FlashAttention と比較して演算強度（メモリトラフィック 1 バイトあたりの FLOPs）を向上させます。これにより、操作はより計算束縛的な領域へと移行し、現代の GPU アーキテクチャにとって有利となります。
• カーネル最適化： NVIDIA Hopper GPU 向けに CuTeDSL でカスタムデコードカーネルを実装しました。これは、2 つのブランチ（Softmax と共分散）が同じ KV ストリームを共有するという事実を利用し、同じ Tensor Core 命令（WGMMA）内で $QK$と$RK$ の積の同時計算を可能にすることで、HBM トラフィックを最小化します。

主要な貢献

1. アーキテクチャ： トークンごとのソルバーを排除し、学習可能なプローブを導入することで、局所線形アテンションの理論的利点を大規模事前学習にスケーリングする Parallax の導出。
2. 効率性： 多様なバッチサイズとコンテキスト長において FlashAttention 2/3 と同等かそれ以上の性能を発揮する、ハードウェア対応のストリーミングアルゴリズムおよびカスタムデコードカーネルの開発。
3. オプティマイザ - アーキテクチャ相互作用： Parallax アーキテクチャとMuonオプティマイザ間の重要な相互作用の特定。本論文は、Parallax の補正ブランチがオプティマイザの幾何学構造に強く依存することを示しており、Muon 下では大幅な改善が見られる一方、AdamW 下では Softmax アテンションと同等の性能しか発揮しないことを明らかにしています。

実験結果

著者は、合成ベンチマークおよび 0.6B および 1.7B パラメータ規模での LLM 事前学習を通じて Parallax を検証しました。

• 合成ベンチマーク： MAD-Benchmark において、Parallax はリコール指向タスク（コンテキスト内リコール、ファジーリコール、ノイズリコール、選択的コピー）において Softmax アテンション、Mamba、Gated DeltaNet を一貫して上回り、最高の全体精度を達成しました。
• 言語モデリング：
  • パープレキシティ： Muon でトレーニングされた Parallax モデルは、LAMBADA および WikiText においてトランスフォーマーベースラインよりも低いパープレキシティを達成しました。
  • ダウンストリームタスク： Parallax は、ゼロショット QA および常識推論ベンチマークにおいて一貫した改善を示しました。
  • コントロール： パラメータ数一致（ベースラインでクエリヘッド数を増加）および計算量一致（Parallax でヘッド次元を減少）の両方のコントロール条件下で改善が持続し、これが追加のパラメータや計算量ではなく、メカニズム自体に起因することを確認しました。
• オプティマイザ依存性： 重要な発見として、Parallax の優位性はオプティマイザに依存することが明らかになりました。Muon 下では、補正対出力比（COR）が層の深さとともに増加し、モデルは共分散補正を利用することを学習します。一方、AdamW 下では、モデルは補正ブランチを抑制するように学習し（ゲート値 $\approx 0.26$）、トランスフォーマーベースラインと同等の性能となります。
• スペクトル分析： 重み行列の分析により、Muon は AdamW に比べて投影重み（特に $W_R$ および $W_R K$ 回路）においてより高い安定ランクを維持することが示されました。これにより、プローブが主要な共分散方向と効果的に整列することを可能にしています。

意義と主張

本論文は、文献においてアテンションメカニズムに対する強力なアーキテクチャ - オプティマイザの共設計の最初の経験的実証を提供すると主張しています。Parallax のような高度なアテンションメカニズムの成功は、単にアーキテクチャ設計の関数ではなく、最適化の景観と深く絡み合っていることを強調しています。

著者は、Parallax をパレート改善として位置づけています。正しいオプティマイザ（Muon）を使用すれば、効率性を犠牲にすることなく Softmax アテンションよりも優れた性能を提供します。また、Parallax は事前学習済みトランスフォーマーからの事後学習適応（$W_R=0$ で初期化）を可能にする点も指摘しており、これは線形アテンションファミリーには共有されていない特徴です。

この研究は、将来のアテンションメカニズムに関する研究において、特に学習されたプローブの大きさと整列が重要なアフィン構造において、「オプティマイザ - アーキテクチャ相互作用」を主要な設計制約として考慮する必要があることを示唆しています。