arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.26875cs.CLcs.AI

面向长推理的信息感知KV缓存压缩

Information-Aware KV Cache Compression for Long Reasoning

Shanghai Jiao Tong University(上海交通大学)

Jushi Kai, Zhuiri Xiao, Alexandra Birch, Zhouhan Lin

AI总结:

提出InfoKV框架,通过结合预测不确定性与注意力分数进行KV缓存压缩,在长上下文推理中优于现有方法。

摘要:

推理能力在大语言模型(LLMs)中快速发展,导致键值(KV)缓存在预填充和解码阶段规模不断增大。现有的KV缓存压缩方法主要依赖注意力权重来估计token重要性。虽然注意力能有效捕捉上下文相关性,但它忽略了与预测不确定性和token信息量相关的互补信息论信号。本文从前瞻性视角重新审视token重要性,并引入\textit{前向影响}(Forward Influence)这一度量,衡量压缩后的token如何影响未来上下文。我们的分析表明,由注意力分数选择的token主要影响邻近上下文,而与高预测不确定性相关的token对遥远未来上下文具有显著更强的影响。基于这一观察,我们提出\textbf{InfoKV},一种融合信息论信号的熵感知KV缓存压缩框架。它结合了token级预测不确定性与层级表示演化,并在推理过程中将生成的熵分数与注意力分数整合。在Llama-3.1、Llama-3.2和DeepSeek-R1上的长上下文推理基准实验表明,InfoKV在长预填充和解码场景中均持续优于现有的基于注意力的KV压缩方法。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑