Information-Aware KV Cache Compression for Long Reasoning

2026-06-25

AI Summary

This work addresses the challenge of excessive KV cache memory consumption in large language models during long-sequence inference. While existing compression methods rely solely on attention weights, they overlook the informational content and predictive uncertainty of individual tokens. To overcome this limitation, the paper proposes InfoKV, a novel framework that introduces an information-theoretic perspective to KV cache compression for the first time. InfoKV quantifies each token’s influence on subsequent context through a Forward Influence metric, integrating token-level prediction uncertainty, cross-layer representation evolution, and attention scores into a unified, information-aware compression strategy. Experimental results demonstrate that InfoKV consistently outperforms state-of-the-art attention-based baselines across both prefill and decoding stages on Llama-3.1, Llama-3.2, and DeepSeek-R1 models.

Abstract

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 \textit{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 \textbf{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.

Problem

• KV cache compression
• long reasoning
• information-theoretic signals
• predictive uncertainty
• token importance

Innovation

• KV cache compression
• information theory
• predictive uncertainty
• Forward Influence
• long-context reasoning