用学习方法动态分配注意力缓存，显著提升长文本推理效率。
LKV: End-to-End Learning of Head-wise Budgets and Token Selection for LLM KV Cache Eviction
- 端到端学习每个头的缓存预算和重要性，不再依赖人工规则。
- 在LongBench上仅保留15%缓存仍接近无损性能。
- 适合追求高精度长文本推理的模型部署场景。
大语言模型在长文本推理时受限于键值（KV）缓存内存的线性增长。现有压缩方法因依赖启发式策略而存在根本局限：启发式预算依赖统计先验而非任务目标，导致资源错配；启发式选择则依赖查询-键耦合或静态归纳偏置（如注意力汇点）。为此，我们提出LKV（学习型KV淘汰），将KV压缩建模为端到端可微优化问题。LKV-H学习任务优化的全局预算，LKV-T在不显式计算注意力矩阵的前提下推导内在KV重要性。该设计规避了启发式代理，严格对齐压缩与任务目标。大量实验表明，LKV在LongBench和RULER基准上均实现顶尖性能，尤其在LongBench上仅需保留15%的KV缓存即可近乎无损。关键分析指出，学习型预算分配是保真度的主要驱动力，证明数据驱动分配对突破人工规则局限至关重要。
原文摘要 · Abstract (English)
Long-context inference in Large Language Models (LLMs) is bottlenecked by the linear growth of Key-Value (KV) cache memory. Existing KV cache compression paradigms are fundamentally limited by heuristics: heuristic budgeting relies on statistical priors rather than task objectives, causing resource misallocation, while heuristic selection relies on coupled query-key interactions or static inductive biases (e.g., attention sinks). To address this limitation, we introduce LKV (Learned KV Eviction), which formulates KV compression as an end-to-end differentiable optimization problem. LKV integrates LKV-H to learn task-optimized global budgets, and LKV-T to derive intrinsic KV importance without materializing attention matrices. This design bypasses heuristic proxies, strictly aligning compression with task objectives. Extensive evaluations demonstrate that LKV achieves state-of-the-art performance on both LongBench and RULER benchmarks at high compression rates. In particular, on LongBench, LKV achieves near-lossless performance with only 15\% KV cache retention. Crucially, our analysis identifies learned budgeting as the dominant driver of fidelity, demonstrating that data-driven allocation is essential to overcome the limitations of hand-crafted heuristics.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。