長い推論のための情報認識型KVキャッシュ圧縮

要旨

大規模言語モデル（LLM）における推論能力の急速な向上に伴い、プリフィリング段階とデコーディング段階の両方でキー・バリュー（KV）キャッシュのサイズが増大しています。既存のKVキャッシュ圧縮手法の多くは、トークンの重要度を推定するためにアテンション重みに依存しています。アテンションは文脈上の関連性を効果的に捉える一方で、予測の不確実性やトークンの情報量に関連する相補的な情報理論的シグナルを見落としています。本論文では、将来を見据えた観点からトークンの重要度を再検討し、圧縮されたトークンが将来の文脈にどのように影響するかを測定する指標として、Forward Influenceを導入します。我々の分析により、アテンションスコアによって選択されたトークンは主に近傍の文脈に影響を与えるのに対し、高い予測不確実性を伴うトークンは遠い将来の文脈に対して著しく強い影響を及ぼすことが明らかになりました。この観察に基づき、情報理論的シグナルを統合したエントロピー考慮型KVキャッシュ圧縮フレームワークであるInfoKVを提案します。InfoKVは、トークンレベルの予測不確実性と層別の表現進化を組み合わせ、得られたエントロピースコアを推論中にアテンションスコアと統合します。Llama-3.1、Llama-3.2、DeepSeek-R1を用いた長文脈推論ベンチマークでの実験により、InfoKVが長いプリフィリングとデコーディングの両シナリオにおいて、既存のアテンションに基づくKV圧縮手法を一貫して上回ることが実証されました。

English

Reasoning capability has advanced rapidly in large language models (LLMs), leading to an increasing size of key-value (KV) cache in both prefilling and decoding stages. Existing KV cache compression methods mainly rely on attention weights to estimate token importance. While attention effectively captures contextual relevance, it overlooks complementary information-theoretic signals related to predictive uncertainty and token informativeness. In this paper, we revisit token importance from a forward-looking perspective and introduce Forward Influence, a metric that measures how compressed tokens affect future contexts. Our analysis reveals that tokens selected by attention scores mainly influence nearby contexts, whereas tokens associated with high predictive uncertainty exhibit substantially stronger influence on distant future contexts. Based on the observation, we propose InfoKV, an entropy-aware KV cache compression framework that incorporates information-theoretic signals. It combines token-level predictive uncertainty with layer-wise representation evolution and integrates the resulting entropy scores with attention scores during reasoning. Experiments on long-context reasoning benchmarks with Llama-3.1, Llama-3.2, and DeepSeek-R1 demonstrate that InfoKV consistently outperforms existing attention-based KV compression methods in both long prefilling and decoding scenarios.