技術的概要：変分線形注意機構（VLA）

1. 問題定義

長文脈シーケンスモデリングは、計算効率とメモリ安定性の間で根本的なトレードオフに直面している。標準的なトランスフォーマーの注意機構は最適な連想検索を実現するが、$O(T^2)$ の時間計算量と $O(Td)$のメモリ使用量を伴い、$10^5$トークンを超えるシーケンスには非現実的である。線形注意機構は、カーネル特徴マップと再帰的状態更新を通じて計算量を$O(T)$に削減することで計算ボトルネックに対処する。しかし、これは決定的な失敗モードを導入する：無条件の加法的蓄積により、内部メモリ状態$S_t$がフロベニウスノルムにおいて ($O(T)$) 無制限に成長することである。この無制限な成長は、格納された連想間での漸進的な干渉を引き起こし、文脈長が増加するにつれて検索精度を劣化させる。

DeltaNet は、状態を均一に減衰させるスカラーゲーティング機構を通じてこれを緩和しようとするが、方向選択的な制御を欠いている。スカラーゲートは、最近上書きされた方向を優先的に忘却しつつ、より古く安定した連想を保持することができず、高メモリ負荷下で急速な精度の崩壊を招く。著者らは、核心的な問題が計算コストではなく、メモリ安定性、具体的にはメモリ更新の幾何学にあると主張する。

2. 手法：変分線形注意機構

本論文は、線形注意機構のメモリ更新をオンライン正則化最小二乗問題として再定義する**変分線形注意機構（VLA）**を提案する。

2.1 中核的な定式化

無条件の蓄積 ($S_t = S_{t-1} + v_t k_t^\top$) の代わりに、VLA はペナルティ化された予測誤差を最小化するメモリ行列 $S_t$ を求める：
$$S^*_t = \arg \min_S \sum_{s=1}^t \|v_s - S\hat{k}_s\|^2 + \text{tr}(S M_t S^\top)$$
ここで、$\hat{k}_s$ は正規化されたキーであり、$M_t$ は過去に観測されたキーの幾何学を符号化する時間変化するペナルティ行列である。

2.2 適応的ペナルティ幾何学

VLA は、逆ペナルティ行列 $A_t = M_t^{-1}$ を、ステップあたり $O(d^2)$ のコストでSherman-Morrison 階数 1 公式を用いて厳密に維持する。

• メカニズム: $M_t$ は学習されたペナルティ方向の外積を蓄積する。その結果、$A_t$ は最近「ペナルティ質量」を受け取った（すなわち書き込まれた）方向において収縮する。
• 効果: 最近書き込まれた部分空間はより小さな更新を受け取り、未接触の部分空間は利用可能なままとなる。これは、方向選択的なメモリ更新を可能にする、DeltaNet のスカラーゲートに対する行列値のアナログとして機能する。

2.3 再帰的更新則

この定式化から導出される厳密な再帰的更新則は以下の通りである：

1. 逆ペナルティの更新: $A_t = A_{t-1} - \frac{z_t z_t^\top}{1 + u_t^\top z_t}$ (Sherman-Morrison)。
2. ゲーティングベクトルの計算: $\hat{\alpha}_t = \frac{A_t \hat{k}_t}{\|A_t \hat{k}_t\|}$。
3. 状態の更新: $S_t = S_{t-1} + (v_t - S_{t-1}\hat{k}_t) \hat{\alpha}_t^\top$。

重要なのは、キー $\hat{k}_t$ とゲーティングベクトル $\hat{\alpha}_t$ の両方が単位長に正規化されている点である。出力は、標準的な線形注意機構と同様の再帰の比率として計算され、ELU+1 特徴マップが利用される。

2.4 実装効率

Sherman-Morrison 更新は定数倍のオーバーヘッド（標準的な線形注意機構より約 $5\times$ 多くの演算）をもたらすが、著者らは並列 Blelloch-scan 定式化と融合 Triton カーネルを実装している。これにより、すべての $T$ ステップが単一の GPU カーネル起動に融合され、トークンごとのディスパッチオーバーヘッドが排除され、逐次的な Python 実装に対して $14\times$ の高速化が達成される。

3. 主要な貢献

3.1 理論的保証

本論文は、2 つの安定性特性について厳密な証明を提供する：

1. 有界な状態成長（命題 1）: 有界な入力下において、フロベニウスノルム $\|S_t\|_F$ は自己制限的である。$\|S_T\|_F = O(T)$ となる標準的な線形注意機構とは異なり、VLA の状態は予測残差が消失するにつれて有限のプラトーに収束する。
2. 単位ヤコビアンスペクトルノルム（命題 2）: 再帰ヤコビアン $J_t = \partial S_t / \partial S_{t-1} = I - \hat{\alpha}_t \hat{k}_t^\top$ のスペクトルノルムは、ベクトルが単位正規化されている限り、すべてのシーケンス長とヘッド次元に対して正確に 1 である。これは、任意の深さにおいて安定した勾配フロー（爆発も消失もなし）を保証する。正規化がない場合、スペクトルノルムは $d_h/\lambda_0$ として成長し、勾配の増幅（例：$d_h=96$ の場合、25 ステップ後に $10^{32}$）を招く。

3.2 連想メモリ容量

VLA は、ヘッドあたり $d_h$ 個の連想（ここで $d_h$ はヘッド次元）という理論的容量を維持する。しかし、加法的蓄積とは異なり、VLA は新しい書き込みを、最近使用された部分空間に直交する方向へルーティングする。これにより、均一な希釈なしに古い連想を保持し、容量の境界まで正確な想起を維持することが可能となる。

3.3 効率性

提案された Triton 融合カーネルは $O(T)$ スケーリングを達成し、約43,000 トークンで softmax 注意機構のレイテンシを下回るため、長文脈推論において実用的となる。

4. 実証結果

実験は、Multi-Query Associative Recall (MQAR) タスクとコピータスクを用いて、共有された 2 層トランスフォーマーバックボーン上で実施された。

• 状態ノルムの安定性: $T=1,000$ において、標準的な線形注意機構は $\|S_T\|_F \approx 1,600$ に達するのに対し、VLA は15未満に留まる（$109\times$ の削減）。DeltaNet も有界であるが、VLA のような方向選択性を欠いている。
• 連想想起（容量）: ペア数 $n_{pairs} < d_h$ の領域（具体的には $n=24, d_h=32$）において：
  • VLA: 1.000 の完全一致精度を維持。
  • DeltaNet: 0.010 まで崩壊。
  • 標準線形注意機構: 0.074 まで低下。
  • Softmax: 0.26 まで低下。
• 長文脈保持: 固定 $n=8$ で $T$ を最大 512 まで増加させた場合、VLA は1.000 の精度を維持する。DeltaNet は、クエリが到達する前に古い連想を減衰させるスカラーゲートにより、すべての長さで $\approx 0.01$ まで崩壊する。
• 過負荷領域: $n > d_h$ の場合、VLA は適応的ペナルティ行列が最近書き込まれた方向を優先的に上書きし、古い連想を部分的に保持するため、加法的線形注意機構よりも優雅に劣化する。

5. 意義と主張

本論文は、長文脈の信頼性は、単なる計算問題ではなく、メモリ幾何学の問題であると提唱する。

• 幾何学的制御: 主な貢献は、状態が更新を許されるどの方向を制御することが、単に更新の計算コストを削減することよりも重要であることを実証した点にある。無条件の蓄積を、適応的ペナルティ逆行列によって支配される残差誤差更新に置き換えることで、VLA は数値的に安定した、定数メモリを要するシーケンスモデリングを実現する。
• 古典的理論との関連: VLA は線形注意機構を**再帰的最小二乗法（RLS）**適応フィルタと結びつけ、従来の「高速重み」アプローチに欠けていた原理的な導出と安定性理論を提供する。
• 実用的範囲: 著者らは、VLA が $d_h$ 個の最も最近強化された連想の固定サイズで数値的に安定した要約を提供し、KV キャッシュインフラなしでストリーミング推論に適すると主張する。VLA はヘッドあたり $d_h$ 個の連想という根本的な容量制限を克服するものではないが、その制約内で優れた安定性と検索を提供することを明示的に指摘している。

この研究は、VLA が定数倍のオーバーヘッドを伴うものの、理論的安定性（単位ヤコビアン、有界状態）と実証的パフォーマンス（容量までの完全な想起）の組み合わせが、標準的な線形注意機構が干渉により失敗する長文脈トランスフォーマーに対する堅牢な解決策を確立すると結論づけている。