Linearized 2-Simplicial Attention

技術要約：線形化2-単体アテンション（Linearized 2-Simplicial Attention）

問題提起
標準的なソフトマックス・アテンションは、過去のトークンに対してコンテンツベースのアクセスを提供するが、$O(n^2)$ の計算量と増大するKVキャッシュに悩まされ、スケーラビリティが制限される。線形アテンション機構は、過去の情報を固定サイズの状態へと圧縮するが、正確な検索や選択的な更新に苦慮することが多い。対照的に、2-単体（2-simplicial）アテンションは、クエリをトークンのペア（トリプレット）に対してスコアリングすることで高次の相互作用を導入し、理論的には三重依存型のタスクにおいて性能を向上させる。しかし、正確な2-単体アテンションは $O(n^3)$ の計算を必要とし、極めて高コストである。近年の実用的な試みは、両方のトークン軸を固定ウィンドウに制限することに依存しており、効率性のためにグローバルな到達性を犠牲にしている。

手法
著者らは、一方の軸にはグローバルな到達性を、もう一方の軸には局所的な精度を維持しつつ、線形コスト ($O(n)$) を実現する2-単体アテンションの線形化形式である LinSimp を提案する。

1. 三線形書き換えによる線形化:
   核心となる革新は、三線形スコア $\langle q_i, k_j, r_c \rangle$ を、複合クエリとキーとの内積として書き換えることである：
   $$\langle q_i, k_j, r_c \rangle = (q_i \odot r_c)^\top k_j$$
   この再定式化により、トークンインデックス $j$ にわたる総和を、標準的な指数ドット積カーネルと同様の、通常の線形アテンションとして扱うことが可能になる。

2. 正のランダム特徴量（Positive Random Features）:
   指数カーネル $\exp((q_i \odot r_c)^\top k_j)$ を近似するために、著者らは正のランダム特徴量（具体的には直交ランダム特徴量）を採用している。これにより、グローバルなトークン履歴 ($j$) に対する総和を、シーケンス長に依存しない固定サイズのステート ($M_i$ および $a_i$) へと圧縮できる。

   • グローバル軸 ($j$): ランダム特徴量ステートによって処理され、因果的なプレフィックス全体へのアクセスを可能にする。
   • ローカル軸 ($c$): 「アンカー」次元は明示的に保持されるが、短いローリングウィンドウ（例：64トークン）に制限される。これにより、各最近のアンカーが、グローバルステートへと送られる独自の複合クエリを生成するというハイブリッドなメカニズムが構築される。

3. アーキテクチャと実装:

   • モデル構成: 著者らは、ゲート付きデルタ則高速重みモデルである Kimi Delta Attention (KDA) も利用するモデルアーキテクチャの中に、LinSimpレイヤーを統合している。彼らは、18層のKDAレイヤーと6層のLinSimpレイヤーを含む「ノー・ソフトマックス（no-softmax）」モデルを構築した。
   • カスタムカーネル: 本手法はカスタム Tritonカーネル を用いて実装されている。フォワードパスでは、完全なトークンペアを実体化することを避けるために、境界条件からプレフィックスステートを再構成しながらクエリをチャンクごとに処理する。バックワードパスは、大きな中間テンソル ($T \times m \times D_v$) を保存せずに勾配を蓄積するための、2パス（two-sweep）設計を使用している。

主な貢献
論文では、4つの具体的な貢献を挙げている：

1. 導出: 三線形スコアを、線形アテンション近似に適した形式へと変換する、2-単体アテンションの一モード・カーネリゼーション（one-mode kernelization）。
2. 実装: グローバルモード用の正の直交ランダム特徴量バンクと、ローカルモード用の短い明示的アンカーウィンドウを利用した、線形時間の実装。
3. カーネルエンジニアリング: FP32のリファレンスに対する数値的検証を含む、カスタムTritonフォワードおよびバックワードカーネル。
4. 実証的比較: アイソ・トークン（iso-token）およびアイソ-FLOP（iso-FLOP、等価計算量）の設定下での、標準的なアテンション、正確なウィンドウ型2-単体ベースライン、およびKDAハイブリッドとの包括的な評価。

結果
著者らは、ウェブデータ（FineWeb-Edu）および数学データ（FineMath-4+）を用いて、2kおよび16kのコンテキスト長でモデルを評価した。

• Iso-FLOP パフォーマンス: 計算予算が一致している条件下で、「ノー・ソフトマックス」KDA + LinSimp モデルは、3Bトークンの数学実行および16kコンテキスト設定において、比較されたアーキテクチャの中で最高の平均ダウンストリーム精度を達成した。
• 長コンテキスト (16k): 16kコンテキストにおいて、LinSimpベースのモデルは、KDAハイブリッド（標準的なアテンションを使用するもの）よりも平均精度を 0.0079 向上させ、LAMBADAのパープレキシティを 715.6 から 602.6 へと大幅に減少させた。
• 効率性: カスタムカーネルは、ソフトマックスアテンションに近いスループットを実現した（2kコンテキストにおける単一のLinSimpレイヤーにおいて、ソフトマックスレイヤーの時間の1.03倍）が、生の速度において成熟したソフトマックスカーネルをまだ凌駕するには至っていない。
• ベースライン比較: 正確なウィンドウ型2-単体ベースライン（Roy et al., 2025）は、計算量が一致している場合、線形化アプローチと比較して高い検証損失とパープレキシティを示した。

意義と主張
本論文は、LinSimps が高次（三線形）相互作用の理論的利点と、線形時間の推論という実用的な制約との間の溝を埋めることに成功したと主張している。2-単体相互作用を、グローバルなランダム特徴量ステートとローカルな明示的ウィンドウに分解することで、著者らは以下を実証した：

• ソフトマックス・アテンション・レイヤーを一切含まない 高性能な言語モデルを構築することが可能である。
• このアーキテクチャは、等価計算量予算の下で、長コンテキストのシナリオにおいてKDAハイブリッド（標準的なアテンションを一部のレイヤーで使用するもの）を上回ることができる。
• 本手法は、ウィンドウ型2-単体アテンションが欠いている「グローバルな到達性」を提供しつつ、線形計算量を維持している。

限界と謙虚な姿勢
著者らは、実験が予備的なものであり、単一のランダムシードを使用しているため、分散や信頼区間の推定ができないことを明記している。また、現在のバックワードパスカーネルは高度に最適化されておらず、アイソ-FLOP予算は直接的なエネルギーやレイテンシの測定ではなく、分析的な会計上の尺度であることを認めている。さらに、テストされたパラメータ数（約330M）を超えるモデルのスケーリング挙動についても未知である。