Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

論文「Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction」に基づき、詳細な技術的サマリーを以下に示す。

問題定義

キー・バリュー（KV）キャッシュは、トランスフォーマーベースの大規模言語モデル（LLM）およびビジョン・ランゲージモデル（VLM）における効率的な自己回帰的デコードに不可欠である。しかし、コンテキストウィンドウがテキストおよび視覚トークンを含む数万トークンに拡大するにつれ、KV キャッシュは線形なメモリ増大と二次的なアテンション計算コストにより、主要なボトルネックとなっている。

既存の解決策は、通常 KV 除去を圧縮問題として扱う：メモリ制約下でフルキャッシュ推論を近似するためにトークンを削除する。これらの手法は、頻度やアテンションの大きさなどのヒューリスティック、または層ごと・ヘッドごとの固定予算に依存することが多い。著者らは、このアプローチが最適ではないと主張する。なぜなら、それはフルキャッシュ推論が理想的な目標であると仮定しているからである。実際には、長いコンテキストには多くの無関係な「ディストラクター（妨害）」トークンが含まれており、有用な証拠からのアテンション質量を希薄化させ、モデルの性能を低下させる。本論文は、適切に行われた選択的除去は、メモリを削減するだけでなく、これらのディストラクターを抑制することで推論を「改善」しうることを示唆している。

手法：グローバル保持ベースの除去

著者らは、キャッシュされたトークンの将来の有用性を予測し、すべての層、ヘッド、モダリティに単一のグローバルメモリ予算を割り当てる DBTrimKV（Dynamic Budget TrimKV）を提案する。

1. 理論的洞察：アテンションの希薄化
本論文は、フルキャッシュのアテンションが「アテンションの希薄化」に苦しむことを示す理論的解析を提供する。キャッシュ内に大量のディストラクタートークンが存在する場合、softmax 正規化により有用なトークンがそれらと競合を強いられ、関連する証拠に割り当てられるアテンション質量が消失する。著者らは、有用なトークンを優先的に保持し、ディストラクターを除去することで、この希薄化を軽減し、理論的にアテンションの質を向上させうることを証明する。

2. 幾何学的保持を将来の有用性として
短期の関連性を反映するクエリ依存のアテンションスコアの代わりに、本手法は各トークンに対して幾何学的保持スコア（$\beta_i$）を学習する。

• 概念：トークンの有用性は、持続イベントとしてモデル化される。トークンが将来のデコードステップにおいて有用であり続ける可能性が高い場合、高い保持スコアが与えられる。
• 定式化：ステップ $t$ におけるトークン $i$ の保持重みは $r_{t,i} = \beta_i^{t-i}$ と定義される。
• 正当性：著者らは、長いコンテキストにおいてトークンのアテンション生存率が時間とともに急速に低下するという経験的観察に基づき、この幾何学的減衰を将来の有用性に対するクエリ非依存の代理指標として理論的に正当化する。

3. 重み共有グローバルゲート
統合された除去ポリシーを可能にするため、著者らは重み共有保持ゲートを導入する。

• アーキテクチャ：軽量なゲートが、すべての KV エントリに対してスカラー保持スコア $\beta$ を予測する。重要なのは、これらのゲートの最終投影層が、すべてのトランスフォーマー層およびアテンションヘッド間で共有（結合）されている点である。
• グローバル較正：この重み共有により、異なる層やヘッドからの保持スコアが共通のスケールに強制される。その結果、テキスト中心のヘッドからのトークンと、視覚ヘッドまたは深い層からのトークンが直接比較可能となる。
• グローバル予算：層ごとまたはヘッドごとに固定予算を割り当てる代わりに、システムは単一のグローバル KV 予算（$M_{global}$）を維持する。除去の各ステップにおいて、システムは予測された将来の有用性に基づいてすべてのキャッシュされた KV エントリをグローバルにランク付けし、上位 $M_{global}$ 個のトークンを保持する。

4. 実装
本手法は、ヘッドごとの可変長 KV シーケンスをサポートするPagedAttentionレイアウトを使用し、連続した物理メモリを必要とせずに、動的かつヘッド固有のキャッシュサイズを可能にする。

主な貢献

1. 除去の再定義：本論文は、KV 除去が単なる圧縮技術であるという仮定に挑戦する。選択的除去は、アテンションの希薄化に対する是正メカニズムとして機能し、フルキャッシュ推論よりも推論性能を向上させる可能性があることを示す。
2. グローバル保持ポリシー：単一の有用性シグナルに基づいて層、ヘッド、モダリティ間でメモリを動的に割り当てる、学習可能でグローバルに較正された除去ポリシーを導入し、硬直的で手動設計された予算ルールを置き換える。
3. 理論的および経験的検証：優先的保持がアテンションの希薄化を軽減することを示す理論的限界を提供し、幾何学的保持を将来の有用性の妥当な代理指標として正当化する。

実験結果

本手法は、長期的推論（MathVision, AIME）、多ターン対話（MMDU）、視覚的質問応答（VQAText, MME）を含む多様なベンチマークで評価された。

• フルキャッシュとの性能比較：多くのシナリオ、特に厳しいメモリ制約下では、DBTrimKV はフルキャッシュ（Vanilla）推論の性能と同等以上の性能を示した。例えば、MMDU 多ターン対話ベンチマークでは、512 トークンの予算を持つ DBTrimKV は、フルキャッシュのベースラインを 14% 上回った。
• ベースラインとの比較：DBTrimKV は、既存の KV 除去手法（SnapKV, AdaKV, R-KV など）および視覚トークンプルーニングベースライン（VisionZip, FastV など）を一貫して上回った。
• 効率性：本手法は、フルキャッシュ推論と比較して KV メモリ使用量とデコード時間を大幅に削減し、コンテキスト長および生成長に対して効率的にスケーリングした。
• アブレーション研究：重み共有投影を除去し（局所スコアを強制し）、グローバル較正を欠く場合、性能が大幅に低下した。これは、グローバル較正の必要性を確認するものである。

意義と主張

本論文は、学習された、グローバルに較正された KV 除去は単なる圧縮技術ではなく、長コンテキスト推論を改善するメカニズムであると主張する。

メモリを節約するだけでなく、ディストラクターを抑制する手段として除去を扱うことで、著者らは「無関係なトークンを忘れる」ことがアテンションを鋭くし、モデルの精度を向上させることを示した。結果は、動的で学習されたメモリ管理が、長コンテキスト設定において静的なフルキャッシュ推論を上回る可能性を示唆しており、より効率的で能力の高い長期的 AI システムへの道筋を提供する。著者らは、テストされたモデルの規模（最大 80 億パラメータ）に関する限界と、保持ゲートとベースモデルのエンドツーエンドの共同学習に関する将来の作業の可能性について言及している。