CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM

技術サマリー：CONF-KV

問題定義
長_horizon_の大規模言語モデル（LLM）推論は、キー・バリュー（KV）キャッシュが GPU メモリの主要な消費源となり、トークンあたりのアテンション遅延がシーケンス長に比例して線形に増加するという重大なボトルネックに直面しています。既存の排除ポリシー（スライディングウィンドウや過去の注意に基づく手法（例：H2O、SCISSORHANDS）など）は、どのトークンを保持するかを決定するために、静的な最近性や過去の信号に依存しています。これらのアプローチは、各デコーディングステップで利用可能な信号、すなわちモデルの現在の不確実性を活用できていないことが多く、その結果、静的なポリシーは困難な生成区間中にコンテキストを早期に排除したり、モデルが確信を持っている際に不要なトークンを保持したりして、メモリ使用量と生成品質の間の最適でないトレードオフを招いています。

手法：CONF-KV
著者らは、モデルの現在の次のトークン分布に基づいてキャッシュ予算を動的に調整する、信頼度認識型 KV キャッシュマネージャーであるCONF-KVを提案します。このシステムは以下のメカニズムを通じて動作します。

1. 信頼度推定: 各デコーディングステップにおいて、本手法はログitをスカラー信頼度スコア（$c$）に変換します。このスコアは、正規化エントロピー、上位 2 つのトークン間のログ確率マージン、および上位トークンの確率の重み付き組み合わせです。
   $$c = w_H(1 - \hat{H}) + w_m\sigma(m) + w_p p_{(1)}$$
   ここで、$\hat{H}$は正規化エントロピー、$m$はログマージンであり、重みは$(0.4, 0.3, 0.3)$に設定されます。

2. 適応的予算選択: 閾値$\tau$が、現在のステップにおけるキャッシュ予算（$N$）を決定します。

   • $c \ge \tau$（高信頼度）の場合、厳格な予算（$N_{high}$）が選択され、積極的な排除を可能にします。
   • $c < \tau$（低信頼度/不確実性）の場合、より多くのコンテキストを保持するために緩い予算（$N_{low}$）が選択されます。
     これにより、不確実な区間中はキャッシュが拡張し、確信のある区間中は縮小する「ノコギリ歯状」のメモリプロファイルが作成されます。

3. トークンランキングと排除: 選択された予算内では、トークンが指数移動平均（EMA）アテンション質量と最近性を組み合わせた複合スコアでランキングされます。ハードな保護ウィンドウ（$P$）により、最も最近のトークンが排除されることはなく、局所的な一貫性が維持されます。

4. システム統合:

   • 混合精度ストレージ: 最も最近の$W$トークンは FP16 で保存され、より古い保持トークンはヘッドごとおよびチャネルごとに INT8 に量子化されます。これにより、より低いビット数の量子化と比較してパープレキシティの劣化を最小限に抑えつつ、メモリフットプリントを削減します。
   • ピラミッド型予算（CONF-KV-L）: 任意のバリエーションとして、層全体に予算を非一様に割り当てます。深い層ほど情報を少数のトークンに集中させる傾向があるという観察に基づき、深い層にはより小さな予算を割り当てます。
   • 実装: マネージャーは連続的な圧縮を使用して、アテンションカーネルのための密なレイアウトを維持し、間接テーブルを回避するとともに、ブロック単位のアテンション読み取りに非量子化を融合させます。

主な貢献

• 信頼度認識型ポリシー: 歴史的な信号のみに依存するのではなく、現在の出力分布を使用してステップごとのキャッシュ予算の決定を駆動する、新規の排除戦略。
• システム設計: 適応的排除とブロック単位オンラインソフトアテンション、混合 FP16/INT8 ストレージ、およびオプションのピラミッド型層予算を組み合わせた完全な実装。
• メカニズム的検証: 最近のコンテキストを除去することによって誘発される KL シフトと信頼度スコアが強く負の相関（ピアソン$r = -0.77$）することを示すテスト。これは、低信頼度がより多くの保持コンテキストの必要性を示すという仮説を検証するものです。
• 包括的評価: 4 つのモデルファミリー（GPT-2、Qwen-14B/32B、gpt-oss-20b）およびパープレキシティ、長コンテキスト検索、ウェブエージェントタスクなどを含む多様なワークロードにわたる広範なベンチマーク。

実験結果

• メモリ - 品質トレードオフ: 2048 トークンの生成長を持つ GPT-2 において、CONF-KV+INT8 は固定 512 トークンのスライディングウィンドウ（約 38.7 MB）と同等のメモリフットプリントを達成しますが、スライディングウィンドウと比較してパープレキシティを 3.11 ポイント削減します。CONF-KV-L（ピラミッド型バリエーション）はさらにメモリを 34.2 MB に削減し、パープレキシティを 3.89 ポイント改善することで、スライディングウィンドウと完全 KV の間の品質ギャップの 74% を埋めます。
• 長コンテキスト検索: 最大 32K トークンの「干し草の山の中の針（NIAH）」テストにおいて、CONF-KV は 91.4% の検索精度を達成し、スライディングウィンドウ（53.8%）や H2O（80.6%）を大幅に上回ります。
• エージェントタスク: 75 件の VisualWebArena タスクにおいて、CONF-KV はピークメモリを 2.8 倍削減しながら、完全 KV の成功率の 95.3% を維持します。
• スループットと遅延: GPT-2 および Qwen-32B において、CONF-KV は完全 KV と比較してステップあたりの遅延を 1.8 倍削減します。スループットはバッチサイズに比例してスケーリングし、バッチサイズ 8 において CONF-KV は完全 KV の 2.06 倍のスループットを達成します。
• アブレーション: 信頼度信号を分離して検討した結果、同様のレートでのランダム排除はパフォーマンスを 36.54 PPL に劣化させますが、完全な CONF-KV は 30.92 PPL を達成し、信頼度ゲーティングとランキングメカニズムの両方が不可欠であることを確認しました。

意義と主張
本論文は、CONF-KV が現在の不確実性を、メモリ - 品質のパレートフロンティアを改善するための有用なシステムメタデータとして実証していると主張しています。この手法は先を見据えたものであるため、排除されたトークンを回復することはできませんが、困難な区間中の早期排除を防ぐ点で特徴的です。このアプローチはトレーニング不要に設計されており、モデルの重みを変更せず、ヘッド単位割り当て、層単位割り当て、精度選択などの既存の技術と直交しているため、それらと組み合わせることが可能です。

著者らは、この手法が排除閾値を超えることのない短いコンテキストでは「ノーオペレーション（no-op）」となること、および連続的な圧縮には収集コストがかかること（ただし、アテンションの節約によって上回る）を限界として指摘しています。また、高温サンプリング下では信頼度信号が有用性が低下することも認めています。最後に、論文は二重利用の考慮事項を強調しています。すなわち、推論コストの低下は効率を向上させる一方で、望ましくない長_horizon_自動化の障壁を低下させる可能性もあり、基盤モデルに対する安全性の制御は依然として必要であるとしています。