技術要約：すべてを読み取らない：線形アテンションのための曲率条件付きクエリ

1. 問題提起

線形アテンション・メカニズムは、標準的なソフトマックス・アテンションの二次的な計算コスト（$O(T^2)$）とメモリ・オーバーヘッドを解決する手段として登場しました。ソフトマックスの内積をカーネル化された内積に置き換えることで、線形アテンションは、行列値の状態を持つ線形RNNとして読み取り操作を再定義し、線形時間計算量（$O(T)$）と一定のキャッシュ（$O(1)$）を実現します。

しかし、書き込み側のメカニズム（GLA、Mamba、RWK Vにおけるデータ依存型ゲーティングや、Gated DeltaNetにおけるデルタ・ルール更新など）の進展にもかかわらず、線形アテンション・モデルは、標準的なソフトマックス・アテンションと比較して、インコンテキスト・リトリーバル（文脈内検索）および長文脈タスクにおいて一貫して劣った性能を示しています。著者らは、この性能差は情報のメモリへの書き込み方だけでなく、決定的な要因として、どのように読み取るかに起因すると主張しています。

標準的なソフトマックス・アテンションでは、log-sum-exp正規化項が競争環境を作り出し、最も特徴的なキーが出力を支配します。これにより、メモリ内の高密度領域に位置する一致（マッチ）は効果的に抑制されます。一方、線形アテンションは効率性を得るためにこの正規化項を破棄します。その結果、過去のすべてのキーが出力に対して加算的に寄与し続けます。これにより、「有用なターゲット」（針）が、保存されたベクトルの「塊」（干し草）によって希釈されてしまいます。特に、クエリがメモリの密度が高い方向と重なる場合に顕著です。既存の改善策は書き込み側（何が状態に入るかを決定すること）に焦点を当てていますが、将来のクエリがどの方向に対して読み取りを行うかに対して反応することはできません。

2. 手法：曲率条件付きクエリ (Curvature-Conditioned Query: CCQ)

本論文は、欠落しているソフトマックス正規化項の局所的な幾何学から着想を得た、軽量な読み取り時補正メカニズムである曲率条件付きクエリ (CCQ) を提案しています。

理論的導出

著者らは、ソフトマックスの競争項 $A_t(q_t) = \log \sum_{j \le t} \exp(k_j^\top q_t)$ が、キーが固定されている場合、クエリ $q_t$ のみの関数であることを指摘しています。彼らは、等方的アテンション点（$q=0$）の周りで、この対数分配関数の二次のテイラー展開を行います。

$$A_t(q) \approx \text{const} + \mu_t^\top q + \frac{1}{2} q^\top (\bar{C}_t - \mu_t \mu_t^\top) q$$

ここで、展開のヘッシアン（ヘッセ行列）は、キーの実行時サンプル共分散 $\Sigma_t = \text{Cov}_{\le t}[k]$ に対応します。この共分散行列は、メモリにおける高密度の方向を捉えます。

メカニズム

CCQはこの幾何学的知見を利用して、バックボーンのステートと相互作用する前に「クリーン」なクエリを構築します。生のクエリ $q_t$ でステートを読み取る代わりに、モデルは線形収縮（linear contraction）を適用します。

$$q_t^{\text{clean}} = (I - \lambda_t \Sigma_t) q_t = q_t - \lambda_t \Sigma_t q_t$$

• $\Sigma_t$: 実行中のキー共分散。バックボーンのステートと同じ再帰的またはチャンク単位のメカニズムを通じて維持されます（各ヘッドにつき、追加の $d_k \times d_k$ 行列と $d_k$ ベクトルを必要とします）。
• $\lambda_t$: クエリに条件付けられたスカラー・ゲート（$\lambda_t = \sigma(W_\lambda q_t + b_\lambda)$）であり、収縮の激しさを制御します。
• 効果: この演算子は、メモリの高分散（高密度）方向に沿ってクエリを収縮させます。これにより、保存されたベクトルの塊と整列するキーにペナルティを与え、効果的にクエリを混雑した方向から引き離し、孤立したターゲットを優先させます。

統合

CCQは読み取りステップのみを修正します。これは、書き込みルールやステートの再帰 $S_t = \text{Update}(S_{t-1}; k_t, v_t, \dots)$ を変更しないため、あらゆる線形アテンション・バックボーン（例：GLA、Gated DeltaNet）と完全に構成可能です。最終的な出力は $o_t = S_t q_t^{\text{clean}}$ として計算されます。

3. 主な貢献

1. 読み取りステップのボトルネックの特定: 本論文は、読み取りステップを、書き込み側の改善策を補完する、線形アテンションとソフトマックス・アテンションの性能差を埋めるための、明確かつ未探索のレバーとして特定しました。
2. CCQメカニズム: ソフトマックスの対数分配関数の単一のテイラー展開から導出された、曲率を考慮した線形収縮の提案。これは、書き込みダイナミクスを変更することなく、あらゆる線形アテンションの書き込みルールに追加可能です。
3. 実証的検証: CCQをGLAおよびGated DeltaNetのバックボーンに付加した際の、パープレキシティ、ゼロショット・ダウンストリーム精度、Needle-in-a-Haystack (S-NIAH) リトリーバル、長さ外挿、および LongBench パフォーマンスの向上を示す広範な実験。

4. 実験結果

著者らは、FineWeb-Eduのサブセットを用いて、500M および 1.3B パラメータのモデルでCCQを評価しました。

• 言語モデリングおよびダウンストリーム精度:
  • 1.3Bにおいて、CCQ-GLA は WikiText パープレキシティ (16.92) および LAMBADA 精度 (47.82%) で、再帰モデルの中で最高値を達成し、そのバックボーンを大幅に上回りました。
  • CCQ-Gated DeltaNet は、7つのタスクの平均精度 (56.24%) において、再帰モデルの中で最高値を達成しました。
• 合成リトリーバル (S-NIAH):
  • CCQは、長さ外挿タスクにおいてバックボーンを劇的に救済しました。例えば、1.3Bにおいて、CCQ-GLAは 4K コンテキストにおける S-NIAH-1 (pass-key) 精度を、GLA バックボーンに対して +49.0% 向上させました。
  • 標準的な Transformer は、4K の学習制限により 8K コンテキストで 0% の精度に崩壊しましたが、再帰的な CCQ バリアントは相当な精度を維持しました。
• 長さ外挿 (Length Extrapolation):
  • 6つの長文脈コーパス（GovReport, QMSum など）において、CCQ バリアントは 4K から 20K コンテキストにわたって、バックボーンよりも一貫して低いパープレキシティを達成しました。
  • 20K において、CCQ バリアントはそれぞれのバックボーンを 1〜2.5 ポイントのパープレキシティで上回りました。
• 長文脈理解 (LongBench):
  • CCQ-GLA は、14 のタスクからなる LongBench において、再帰モデルの中で最高の平均スコア (15.77) を記録し、Mamba2 を +2.64 ポイント上回りました。
  • 両方の CCQ バリアントは、長い出力生成を必要とするタスク（GovReport, MultiNews など）および少数の例示による推論（few-shot reasoning）において、顕著な改善を示しました。

5. 意義と主張

本論文は、線形アテンションの性能差は単なる容量の問題ではなく、幾何学的な問題であると主張しています。すなわち、読み取りステップにおける競争の欠如が、メモリ密度によって特定のターゲットをかき消してしまうという点です。

• 幾何学的知見: 著者らは、ソフトマックスの対数分配関数の局所的な曲率（ヘッシアン）によって、ソフトマックスにおける「失われた」競争を近似できることを示しています。これはキーの共分散と一致します。
• メモリ認識型の補正: 入力のみのクエリ投影とは異なり、CCQ は保存されたキーの実行統計を使用するため、「メモリ認識型」です。
• モジュール性: 本手法は、書き込みダイナミクスを変更することなく、読み取りステップのドロップイン・リプレイスメント（置き換え可能な構成要素）として提示されています。
• 限定的な範囲: 著者らは、現在の結果が 500M/1.3B スケール、2 つの特定のバックボーン（GLA, Gated DeltaNet）、および特定のベンチマークに限定されていることを明記しており、CCQ がすべての線形アテンションの問題を解決したり、さらなる調査なしに 7B 以上のスケールや Mamba のような状態空間モデルに一般化したりするとは主張していません。

要約すると、本論文は、「すべてを読み取らない（＝メモリの密度が高い方向に沿ってクエリを収縮させる）」ことにより、ソフトマックス正規化項を破棄したことによるリトリーバルおよび長文脈能力の損失を、最小限の計算コストで回復できると主張しています。