通过学习预测关键记忆,实现长文本推理的高效压缩与性能提升
Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

- 基于全局保留率设计可学习的键值缓存淘汰策略
- 在多个长文本任务上用更少内存达到甚至超过全缓存效果
- 适合需要长上下文理解的对话、多模态推理场景
键值(KV)缓存是长序列推理中的主要瓶颈,其内存和计算量随序列长度增长。现有KV淘汰方法虽降低开销,但通常导致性能下降。我们发现:全缓存注意力并非总是最优——长上下文中无关令牌会稀释注意力,而有选择性的可学习淘汰反而能提升生成质量。本文提出一种基于全局保留率的KV淘汰方法,在统一内存预算下学习每个令牌未来的有用性。轻量级保留门为缓存的键值条目分配效用评分,共享的最终评分投影将这些评分跨层、跨头进行校准,从而实现所有层、头和模态间直接竞争缓存空间的全局淘汰策略。理论分析表明,优先保留有用令牌可减少注意力稀释,且几何保留可作为未来效用的查询无关代理。在多种长上下文语言与视觉-语言推理及多轮对话基准上,该方法显著降低KV内存,同时匹配或超越全缓存推理表现。结果表明,学习的全局校准淘汰不仅是压缩手段,更是提升长上下文推理能力的机制。
原文摘要 · Abstract (English)
The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。