技術概要：IndexMem - 潜在メモリを備えた学習型 KV キャッシュエビクション

1. 問題定義

大規模言語モデル（LLM）は、長いコンテキストを処理する際に重大なボトルネックに直面します。自己回帰生成に必要なキー・バリュー（KV）キャッシュは、シーケンス長に比例して線形に増加し、すぐに GPU メモリ容量を超え、メモリ帯域幅を飽和させます。標準的な KV キャッシングはアテンションの再計算を回避しますが、展開可能なコンテキスト長を制限する「メモリウォール」を生み出します（例：700 億パラメータのモデルは、100 万トークンのコンテキストに対して約 320GB を必要とします）。

既存の解決策は 2 つのカテゴリに分類され、いずれにも限界があります：

1. 表現圧縮：量子化や低ランク近似などの技術は、トークンあたりのフットプリントを削減しますが、コンテキスト長に伴うキャッシュの線形成長という本質的な問題には対処していません。
2. トークンエビクション：メモリを制限するために重要度の低いトークンを破棄する手法です。しかし、現在のエビクションポリシーは主にヒューリスティック（キーの大きさやローカルアテンションウィンドウなどの静的統計に依存）であり、トレーニングフリーです。これらは、入力に依存する複雑なトークンの重要度パターンを捉えられないことが多くあります。さらに、標準的なエビクションは情報を恒久的に破棄するため、不可逆的な忘却を引き起こし、遠くのトークンの検索を必要とするタスク（例：Haystack 内の針）のパフォーマンスを低下させます。

2. 手法

著者らは、正確なトークン選択のための学習可能なインデクサーと、エビクションされたトークンから情報を回復するための潜在メモリモジュールを組み合わせたフレームワークIndexMemを提案します。

A. トークン重要度のための学習可能インデクサー

ヒューリスティックなスコアの代わりに、IndexMem は KV トークンの重要度を予測するための軽量な学習可能なインデクサーを採用します。

• アーキテクチャ：インデクサーは、マルチクエアアテンション（MQA）スタイルのモジュールです。マルチヘッドクエリを低次元空間にダウンプロジェクションし、隠れ状態から導出された単一の共有キーを使用します。安定性のために QK ノルムを統合し、ヘッドの寄与を調節する学習可能なゲーティング機構を備えています。
• 機能：隠れ状態 $X$ とクエリ $Q$ が与えられると、インデクサーは密なクエリからキーへのスコア行列 $A$ を出力します。ここで、$A_{s,t}$ はクエリ $s$ に対するキートークン $t$ の重要度を測定します。
• トレーニング：ベースラインの LLM は固定されたままです。インデクサーは、ベースラインのアテンション動作に一致させるために知識蒸留を通じてトレーニングされます。目的関数は、教師（ベースラインアテンションのログオッズ）と学生（インデクサーのスコア）のプールされたキーごとの重要度分布間の KL 発散を最小化します。$O(L^2)$ 行列によるメモリボトルネックを回避するため、トレーニングでは、実行中の最大値集計を計算するためのストリーミング型の FlashAttention 風アプローチが使用されます。
• エビクション前：すべての KV 対をエビクション前に具体化する手法とは異なり、インデクサーはエビクション前の予測を可能にし、完全なプリフィル計算の前に保持するエントリを予測することで、メモリと計算を節約します。

B. エビクションされたトークンのための潜在メモリ

トークンを破棄することによる不可逆的な情報損失に対処するため、IndexMem は潜在メモリモジュールを導入します。

• 設計：エビクションされたトークンを追加の KV 対として注入する（これによりアテンションシンクの問題を引き起こす可能性がある）のではなく、メモリはアテンションストリームに対する明示的な残差補償として機能します。
• メカニズム：
  • 高速重み（オンライン更新）：エビクションされたトークンは、推論中に固定サイズの状態行列 $M$ と安定化ベクトル $b$ をオンラインで更新します。これは減衰を伴う外積集積を使用し、メモリが特定のコンテキストに適応できるようにします。
  • 低速重み（学習済み）：投影層とゲーティング関数 $g(q)$ が、メモリ状態 $M$ から読み取り、与えられたクエリに対してどの程度それに依存するかを決定するようにトレーニングされます。
• 出力：最終的なアテンション出力は拡張されます：$o = o_{attn} + g(q) \cdot m(q)$。ここで、$o_{attn}$ は保持されたトークンに対する標準的なアテンションであり、$m(q)$ は失われた情報を補償するメモリの読み出しです。

C. トレーニングプロトコル

システムは 2 段階のトレーニングスキームを使用します：

1. インデクサーのトレーニング：信頼性の高いトークン重要度推定を学習するために、インデクサーのみをトレーニングします。
2. 共同トレーニング：インデクサーとメモリモジュールを共同でトレーニングします。メモリは、完全なアテンション出力と、圧縮された出力 plus メモリ残差との間の平均二乗誤差（MSE）を最小化するように最適化されます。

3. 主要な貢献

1. 学習可能インデクサー：コンテキスト依存のパターンに基づいてトークンの重要度を正確に予測することで、ヒューリスティックなエビクションポリシーを上回るトレーニングベースのアプローチ。
2. 潜在メモリモジュール：エビクションされたトークンをコンパクトでオンライン更新される状態に圧縮し、エビクション中に失われた情報を回復するための残差読み出しを提供するメカニズム。これにより不可逆的な忘却を軽減します。
3. 包括的な評価：複数のモデルファミリー（Qwen、Mistral、Llama）とベンチマーク（RULER、LongBench、Needle-in-a-Haystack）にわたって一貫した改善を実証しました。

4. 実験結果

実験は、RULER および LongBench スイートを使用して、Qwen3-8B、Mistral-7B、Llama-3.1-8B で行われました。

• RULER パフォーマンス：IndexMem はベースライン（SnapKV、PyramidKV、TOVA、Expected Attention）を大幅に上回りました。積極的なエビクション（90% 圧縮）の下でも、IndexMem ははるかに高いスコアを維持します。例えば、90% のエビクションを伴う Llama-3.1-8B において、IndexMem は RULER-16K で 56.0 のスコアを達成しましたが、PyramidKV は 24.1 に低下しました。
• Haystack 内の針（NIAH）：IndexMem は、コンテキスト長（1K–12.5K）と針の深さにわたって優れた安定性を示しました。ヒューリスティックな手法は位置依存の失敗（最近のトークンへのバイアス）を示すのに対し、IndexMem は位置に関係なく針を確実に検索します。潜在メモリは、インデクサーのみのバリアンで見られる「壊滅的な見落とし」を特に減少させます。
• LongBench：IndexMem は、多様なタスク（HotpotQA、MultiFieldQA、TriviaQA）にわたって最も堅牢なパフォーマンスを達成しました。ヒューリスティックな手法と比較して、圧縮が増加してもより優雅に劣化します。特に、一部のタスクでは、中程度の圧縮が邪魔なトークンを除去することで、かえってパフォーマンスを向上させました。
• 効率性：この手法は、最小限のパラメータオーバーヘッド（8B モデルの場合、インデクサーで約 20M、メモリで約 0.5M）を導入します。メモリモジュールはわずかな更新オーバーヘッドを追加しますが、全体のデコーディング時間はフルキャッシュベースラインと競争力があり、メモリ使用量は圧縮率に比例して線形にスケーリングします。

5. 意義と主張

本論文は、IndexMem が長いコンテキスト推論に対して有利な精度とメモリのトレードオフを提供すると主張しています。ヒューリスティックから学習型エビクションへ移行し、破棄された情報を回復するメカニズムを導入することで、この手法は制限された KV 予算下での正確な長いコンテキスト推論を可能にします。

著者らは、この作業を効率的なアテンションのための学習可能アーキテクチャへの一歩として位置付けています。現在のアプローチは、エビクションとメモリモジュールを固定されたベースラインへのアドオンとして扱っていますが、将来的な研究では、ベースラインがネイティブなエビクションとメモリ機能を学習するエンドツーエンドのトレーニングを探求できることに言及しています。論文は、この技術がハードウェアコストを下げ、長いコンテキストモデルへのアクセスを改善することを意図した一般的な効率化技術であると強調しつつも、重要な詳細を見逃すことが有害な結果を招きうる高リスクな分野では注意が必要であることを認めています。