技術的概要：スペクトルクープマン作用素を用いた KV キャッシュ不要の連想想起「Echo」

1. 問題定義

長文脈推論およびエージェント型ツール呼び出しは、数万トークンにわたるトレースを生成する。Transformer ベースのアーキテクチャは、コンテンツアドレス指定による検索を維持するためにキー・バリュー（KV）キャッシュに依存しているが、これらのキャッシュはシーケンス長に比例して線形に増大し、汎用ハードウェア上でメモリボトルネックを生じさせる。一方、Mamba などの状態空間モデル（SSM）は、シーケンスを固定サイズの再帰状態に圧縮することで、線形時間推論とステップごとの定数メモリを実現する。しかし、この圧縮は損失を伴う：SSM は「メモリ断崖（memory cliff）」に悩まされ、格納された事実とそのクエリとの間隔が再帰状態の実効地平を超えると、検索精度が指数関数的に崩壊する。

現在の対策には、SSM 層と因果的アテンションブロックを交互に配置するハイブリッドアーキテクチャが含まれる。これは検索能力を回復させるが、KV キャッシュの線形メモリ増大を再導入してしまう。本論文は、コンテンツアドレス指定による検索をカーネルリッジ回帰として再構成できると主張する。すなわち、十分な統計量を固定メモリ内のrunning sum（累積和）として蓄積することで、検索性能を犠牲にすることなく KV キャッシュを完全に排除できるというものである。

2. 手法：Echo とスペクトルクープマンアテンション（SKA）

著者は、標準的なアテンション層の代わりとして導入可能なスペクトルクープマンアテンション（SKA）を中核とした、KV キャッシュ不要のアーキテクチャEchoを提案する。Echo は局所的なシーケンシャル処理のために SSM（具体的には Mamba-2）を維持しつつ、グローバルな検索メカニズムを SKA に置き換える。

2.1 中核メカニズム：スペクトルクープマンアテンション

SKA は、ストリーミングされる十分な統計量からコンパクトな線形動力学モデルを推定し、$T \times T$ のアテンション行列を具体化したり KV キャッシュを維持したりすることなく、クエリ位置における値を検索する。

• 十分な統計量の蓄積: SKA は、$O(r^2 + rP)$ のメモリ（ここで $r$ は小さな射影ランク、$P$ はヘッド次元）で 3 つの加法的統計量を蓄積する：

  1. 正則化されたグラム行列 $G = \sum \hat{z}_t \hat{z}_t^\top + \epsilon I$。
  2. 時間横断共分散 $M = \sum \hat{z}_t \hat{z}_{t-1}^\top$。
  3. 値・キー共分散 $C_v = \sum v_t \hat{z}_t^\top$。
     これらの統計量は検索計算に対する正確な十分統計量であり、つまりシーケンスが長くなるほど推定の条件数が悪化するのではなく、改善されることを意味する。

• 作用素推定と検索:

  • システムは、リッジ回帰を通じて、正則化グラム行列のチョレスキー因子 $L$ を用いて、白化されたクープマン作用素 $A_w = L^{-1} M L^{-\top}$ を適合させる。
  • 値の読み出しは $B_v = C_v \tilde{G}^{-1}$ として計算され、これは正則化最小二乗目的関数 $\min_B \sum \|v_t - B\hat{z}_t\|^2 + \epsilon \|B\|_F^2$ の一意な最小化解である。
  • 線形アテンション（ゼロからの単一の勾配ステップに対応）や softmax アテンション（非凸最適化を通じてこの解を暗黙的に近似する）とは異なり、SKA はチョレスキー分解を介して閉形式のリッジ解を直接計算する。

• スペクトルフィルタリング:

  • 白化された作用素 $A_w$ は、そのスペクトル半径を制限するためにスペクトル正規化される。
  • 冪乗反復フィルタ（power-iterated filter）（$A_w^K$）が白化されたクエリに適用される。このフィルタは、永続的な固有モード（$|\lambda| \approx 1$、安定した結合を表す）を増幅し、過渡的なモード（$|\lambda| \ll 1$、ノイズや妨害を表す）を抑制する。このスペクトル選択性は、標準的な回帰目的関数単独では持たない能力である。

2.2 アーキテクチャ構成要素

• SSM バックボーン: 層の大部分は局所処理のために Mamba-2 を使用する。
• スペクトルクープマン MLP: 標準的な SwiGLU 順伝播層は、入力に対して単位円盤に制約されたブロック対角回転を適用するスペクトル MLP に置き換えられる。この設計は勾配のランクを保持し、SwiGLU と比較してパラメータ数を 3 分の 1 に削減しつつ、ノルム保存を維持する。
• 推論: Echo は、コンテキスト長に関係なく層ごとに $O(1)$ のメモリを維持する。状態は蓄積器（$G, M, C_v$）とスカラー最大ノルムから構成され、5000 万パラメータモデルの場合、合計約 77 KB となる。

3. 主要な貢献

1. メモリ断崖の排除: 本論文は、純粋な SSM の検索失敗が、容量制限ではなく乗法的再帰の構造的帰結であることを実証する。Echo は、ストリーミング統計量に対する閉形式のリッジソルバを使用することで、この断崖を排除する。
2. KV キャッシュ不要の連想想起: Echo は、純粋な SSM が失敗するベンチマークにおいて完璧な検索を達成し、SSM+アテンションのハイブリッドモデルと同等の性能を示すが、定数の推論メモリで実現する。
3. 回帰との理論的接続: この研究は、コンテンツアドレス指定による検索とリッジ回帰の間のリンクを形式化し、SKA がアテンションの非凸最適化ランドスケープを迂回して、チョレスキー分解を介して明示的にこの最適化問題を解くことを示す。
4. 勾配流の保持: 著者は、SKA とスペクトルクープマン MLP が、逆伝播において追加的なランク圧縮をゼロで導入することを理論的に分析して示す（アテンション行列のエントロピー依存の実効ランクを通じて勾配を圧縮する softmax アテンションとは対照的）。これは、データ効率の向上メカニズムを示唆する。

4. 実験結果

著者は、Echo を合成転移タスク、マルチクエリ連想想起（MQAR）、言語モデリングの 3 つの領域で評価した。

• 合成転移（サブ百万規模）: 5 つの保持された検索タスク（システムプロンプトの想起、ツールトレース、干し草の山の中の針など）において、98.2 万パラメータの Echo モデル（SSM+SKA）は、純粋な SSM（111 万パラメータ）および SSM+アテンションのハイブリッド（99.8 万パラメータ）を上回った。Echo は平均精度 81.1% を達成し、純粋な SSM は 54.4%、ハイブリッドは 76.2% だった。重要なのは、Echo が、ベースラインが偶然レベルにまで崩壊した長さ汎化タスクにおいて、4,096 トークンで 65% という高い精度を維持したことである。
• MQAR（5000 万規模）: 最大 4,096 トークンの妨害ギャップと 32 の KV ペアを有する MQAR ベンチマークにおいて、純粋な Mamba-2 は約 3%（偶然レベル）であった。Echo モデル（Mamba-2+SKA）は、すべての設定で100% の検索精度を達成し、メモリ断崖が構造的なものであり、この規模で排除可能であることを確認した。
• 言語モデリング（1 億 8000 万規模）: 100 億トークンで学習された Echo-180M モデルは、5 つのゼロショットベンチマーク（HellaSwag, PIQA, ARC, WinoGrande, LAMBADA）のうち 6 つで、1 億〜4 億パラメータの他のモデル（GPT-2 345M や Mamba-3 変種を含む）と同等かそれ以上の性能を示した。これは、いくつかのベースラインよりも 10 倍少ない学習トークンを使用しながらも達成されたものである。この結果は、このアーキテクチャがメモリ効率だけでなくデータ効率も高いことを示唆する。

5. 意義と主張

本論文は、Echo が SSM において、アテンションの想起能力を回復するために効率性（線形時間、定数メモリ）の利点を犠牲にする必要がないことを実証していると主張する。検索を閉形式のスペクトル回帰問題として再構成することで、Echo は検索タスクにおけるアテンションベースモデルとの性能格差を埋めつつ、固定されたメモリエンベロープを維持する。

著者は、検索の向上がプレフィックスマスキング戦略ではなく、スペクトル作用素と閉形式の解によって駆動されていることを強調する。彼らは、このアプローチが、再帰が局所コンテキストを処理し、軽量な作用素モジュールが疎な長距離ルックアップを処理する、定数メモリ予算内で動作する新しいクラスのアーキテクチャへと向かう道を示していると論じる。この研究は、「メモリ断崖」が状態空間モデルの固有の制限ではなく、解決可能な構造的課題であることを示唆している。