技術概要：CriticalKV

問題定義

トランスフォーマーアーキテクチャに基づく大規模言語モデル（LLM）は、自己回帰生成中のキー・バリュー（KV）キャッシュに関連する二次的な計算コストおよびストレージコストにより、重大な展開上の課題に直面しています。シーケンス長が増加するにつれて KV キャッシュは劇的に拡大し、GPU メモリオーバヘッドの高騰と I/O レイテンシの増大を招きます。

近年のキャッシュエビクション手法（H2O、Scissorhands、SnapKV など）は、非重要なエントリを剪定することでこれを緩和しようとしていますが、これらは経験的なヒューリスティクスに大きく依存しています。具体的には、これらの手法は、より高いアテンション重みを持つエントリが最も重要であると仮定しています。著者らは、この仮定に形式的な根拠が欠如しており、以下の 2 つの重要な問いを提起すると主張しています：

1. KV キャッシュエントリの重要性を真に決定する基準とは何か？
2. 重要なエントリを特定するために、アテンション重みのみに依存することは十分か？

本論文は、現在の手法が値状態（value states）と事前学習済みパラメータ行列の決定的な役割を軽視しており、これにより最適でないエビクション判断や著しい品質低下（圧縮損失）を招く可能性があると指摘しています。

手法

本論文は、出力摂動分析に基づいて重要な KV キャッシュエントリを特定するための形式的枠組みであるCriticalKVを導入します。

理論的枠組み

著者らは、重要なエントリ特定の問題を、完全な KV キャッシュを重要なエントリの部分集合に置き換えた際の出力摂動を最小化する問題として定義します。

• 摂動の定義: 出力摂動 $L$ は、完全なキャッシュの出力 ($o$) と選択された部分集合 ($\hat{o}$) の間の $L_1$ 距離として定義されます：$L = \|o - \hat{o}\|_1$。
• 上限の導出: 理論的分析を通じて、著者らはこの摂動に対する上限 $\theta$ を導出します。重要なことに、この上限は摂動がアテンション重み ($A$) だけでなく、投影された値状態（$V W_O$、ここで $W_O$ は出力投影行列）によっても影響を受けることを明らかにしています。
• 重要な洞察: 導出された上限は、摂動を最小化するには、アテンション重みとパラメータ行列を通じて投影された値状態の大きさの両方を考慮する必要があることを示しています。アテンション重みのみに依存することは、理論的に最適ではない「緩和された」上限を最小化するに過ぎません。

アルゴリズム：摂動制約付き選択

この理論を実践化するため、著者らは worst-case 摂動の導出された上限を最小化するように設計された 2 段階の貪欲選択アルゴリズム（アルゴリズム 1）を提案します：

1. ステージ 1（高アテンション重み）: キャッシュ予算の割合 ($b' = b \times \alpha$、ここで $\alpha=0.5$) を割り当て、最も高いアテンション重みを持つエントリを選択します。この段階により、選択された部分集合の累積アテンション重みが全体の 50% を超えることを保証し、理論的限界のための重要な仮定を満たします。
2. ステージ 2（摂動制約付き）: 残りの予算 ($b''$) を用いて、アテンション重みと投影された値状態の $L_1$ ノルムの積 ($A_i \|V_{i,:}\|_1$) を最大化するエントリを選択します。このステップは、導出された摂動上限の最小化を直接目指します。

本アルゴリズムは、既存のキャッシュエビクションパイプライン（SnapKV、AdaKV、HeadKV など）の基盤となる蓄積プロセスを変更することなく統合可能なプラグ＆プレイ型の拡張として設計されています。

主要な貢献

1. 形式的な問題定義: 本論文は、経験的ヒューリスティクスを超えて、出力摂動の観点から重要なキャッシュエントリ選択の問題を初めて形式化しました。
2. 理論的洞察: アテンション重みだけでは重要なエントリを特定するには不十分であり、パラメータ行列を通じて投影された値状態も同様に不可欠であることを特定しました。
3. 新規アルゴリズム: worst-case 出力摂動を最適化する摂動制約付き選択アルゴリズムを提案しました。
4. 汎用的な拡張: この手法は、既存の最先端（SOTA）エビクション技術に対する汎用的な拡張として機能し、計算オーバヘッドはほとんど不要です。

実験結果

著者らは、CriticalKV を 3 つの最先端エビクション手法（SnapKV、AdaKV、HeadKV）に統合し、RulerおよびLongBenchベンチマークの 29 データセット上で 3 つの異なる LLM（Llama-3.1-8B、Mistral-7B-v0.3、Qwen2.5-32B）を用いて評価しました。

• 性能向上: 平均して、アルゴリズムはすべてのテスト構成において圧縮損失を半分以上削減しました。
  • Ruler ベンチマーク（キャッシュサイズ 40%）において、Qwen2.5-32B に対して HeadKV に本手法を統合した結果、平均損失は 13.7% から 3.4% に減少しました。
  • LongBench において、本手法は長依存ドメインの 90 のテストケースのうち 88 で性能を向上させました。
• ロバスト性: 改善は、キャッシュサイズ（2.5% から 100%）や異なるモデルアーキテクチャにわたって一貫していました。
• 効率性: アルゴリズムは計算オーバヘッドをほとんど導入しません。32K コンテキスト長でバッチサイズ 1 の場合、プレフィルリングの TTFT（先頭トークンまでの時間）は約 0.06 秒のみ増加し、デコーディングレイテンシはベースラインのエビクション手法と同一でした。
• 摂動分析: 経験的分析により、本アルゴリズムが Llama モデルの 92% 以上のアテンションヘッドにおいて出力摂動を削減し、最終層の隠れ状態における摂動を大幅に低下させることが確認されました。

意義と主張

本論文は、キャッシュエビクションに対する形式的に根拠のある新たな視点を提供すると主張しています。直感的なアテンション重みのヒューリスティクスから、出力摂動の理論的 minimization へと焦点を移すことで、CriticalKV は現在のエビクション戦略における根本的な限界に対処します。

著者らは、自らの研究が以下の点に重点を置いていることを強調しています：

• キャッシュエビクションにおいて値状態がなぜ重要なのかについての理論的根拠を提供する。
• メモリ制約下で生成品質を大幅に向上させる、実用的で低オーバヘッドの解決策を提供する。
• 出力摂動の理論的分析を通じて重要なキャッシュ選択を最適化する将来の研究への有望な道を開く。

コードは公開されており、著者らはリポジトリに後続の研究におけるアプローチの拡張版をホストしていることに言及しています。