基于信息的KV缓存压缩技术,用于长推理任务
Information-Aware KV Cache Compression for Long Reasoning
摘要
大型语言模型中的推理能力得到了迅速发展,这使得关键-值缓存的大小在预填充和解码阶段都不断增大。现有的KV缓存压缩方法主要依赖注意力权重来评估每个词的重要性。虽然注意力机制能够有效捕捉上下文相关性,但它忽略了与预测不确定性和词的重要性相关的信息论信号。本文从前瞻性角度重新审视了词的重要性问题,并提出了“前向影响”这一指标,用于衡量被压缩的词对未来上下文的影响程度。分析表明,由注意力机制确定的词主要影响附近的上下文,而那些与高预测不确定性相关的词则对未来的上下文产生更显著的影响。基于这些发现,我们提出了InfoKV这一基于熵的KV缓存压缩框架,该框架整合了信息论信号,将词级的预测不确定性与层的表示演化相结合,并在推理过程中将熵分数与注意力分数结合起来。在Llama-3.1、Llama-3.2和DeepSeek-R1等长上下文推理基准测试中的实验表明,InfoKV在长文本预填充和解码场景中始终优于现有的基于注意力的KV压缩方法。
English 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 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.