Reformulating KV Cache Eviction Problem for Long-Context LLM Inference

📄 arXiv: 2605.07234v1 📥 PDF

作者: Tho Mai, Joo-Young Kim

分类: cs.CL, cs.AI

发布日期: 2026-05-08

💡 一句话要点

提出LaProx框架：通过输出感知矩阵近似重构KV Cache驱逐策略，实现长文本推理的高效压缩

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文推理 KV Cache压缩 注意力机制 模型推理优化 矩阵近似 大语言模型

📋 核心要点

  1. 现有KV Cache驱逐方法多基于局部注意力权重，未能充分考虑值表示、输出投影及头间交互，导致重要Token识别不准确。
  2. 提出LaProx框架，将驱逐问题建模为输出感知的矩阵乘法近似，通过显式建模注意力与值状态的交互来量化Token贡献。
  3. 实验证明该方法在长文本任务中仅需5%缓存即可保持高性能，在极端压缩下精度损失较现有SOTA方法降低达2倍。

📝 摘要（中文）

大型语言模型（LLM）在支持长上下文推理时，因KV Cache的线性增长面临巨大的内存与计算开销。现有的KV Cache驱逐方法主要依赖于局部注意力权重，忽略了值表示（Value representations）、输出投影及头间交互的影响。本文将KV Cache驱逐问题从传统的基于头部的权重平均法，重构为输出感知的层级矩阵乘法近似问题。我们提出了LaProx，一种新型驱逐策略，通过显式建模注意力图与投影值状态之间的乘法交互，在考虑头间依赖的同时准确量化Token贡献。基于该度量标准，我们提出了首个统一的驱逐策略，能够为Token分配全局可比的重要性分数，从而实现模型范围内的选择而非局部的头部决策。在LongBench和Needle-In-A-Haystack等19个数据集上的实验表明，该方法在仅保留5% KV Cache的情况下仍能保持模型性能，并显著优于现有基线，在极端压缩场景下精度损失降低达2倍。

🔬 方法详解

问题定义：现有KV Cache驱逐方法通常采用“头内独立”的启发式策略，仅关注注意力分数，忽略了后续线性投影层对最终输出的影响，导致在长上下文场景下无法准确捕捉对模型预测贡献最大的Token。

核心思路：论文将驱逐问题重新定义为输出感知的矩阵乘法近似问题。核心逻辑在于：Token的重要性应由其对最终输出向量的贡献决定，而非仅仅取决于注意力权重。通过建模注意力图与投影值状态的乘法交互，能够更全面地评估Token的价值。

技术框架：LaProx框架通过计算层级输出的近似值来评估Token贡献。它不再对每个注意力头进行独立评分，而是将所有头的输出投影整合，构建一个统一的度量空间，从而实现跨头、跨层的全局Token重要性排序。

关键创新：最重要的创新在于引入了“输出感知”的度量标准，并实现了全局可比的评分机制。与传统方法相比，它不仅考虑了注意力权重，还纳入了值向量及其投影矩阵的影响，有效捕捉了模型内部的复杂依赖关系。

关键设计：该方法通过数学推导将注意力机制的输出近似为一系列矩阵乘积，利用这些乘积的范数或特定投影值作为重要性分数。这种设计允许在推理过程中以极低的计算开销动态计算Token得分，并支持在模型全层范围内进行统一的驱逐决策。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在LongBench和Needle-In-A-Haystack等19个基准测试中，LaProx表现卓越。在仅保留5% KV Cache的极端压缩比下，模型性能依然稳健。对比当前主流的驱逐算法，LaProx在极端场景下的精度损失降低了2倍，且计算开销极小，证明了其在长上下文推理优化中的领先地位。

🎯 应用场景

该技术主要应用于长文本LLM推理场景，如超长文档分析、长篇代码库理解及复杂对话系统。通过显著降低KV Cache的显存占用，它能使消费级硬件运行更长的上下文，提升推理吞吐量，并降低大规模部署LLM的硬件成本，在金融、法律及科研文献处理等领域具有极高的实用价值。

📄 摘要（原文）

Large language models (LLMs) support long-context inference but suffer from substantial memory and runtime overhead due to Key-Value (KV) Cache growth. Existing KV Cache eviction methods primarily rely on local attention weights, neglecting the influence of value representations, output projection, and inter-head interactions. In this work, we reformulate KV Cache eviction from a conventional head-wise, weight-averaging approach into an output-aware, layer-wise matrix multiplication approximation problem. We introduce LaProx, a novel eviction strategy that explicitly models the multiplicative interaction between attention maps and projected value states to accurately quantify token contributions while accounting for inter-head dependencies. Building on this metric, we propose the first unified eviction strategy that assigns globally comparable importance scores to tokens, enabling model-wide selection instead of local, head-wise decisions. Experimental results across 19 datasets on long-context benchmarks LongBench and Needle-In-A-Haystack demonstrate that our approach maintains model performance with only 5\% of the KV cache and consistently outperforms prior works across all configurations. Notably, our method achieves up to 2$\times$ accuracy loss reduction under extreme compression scenarios compared to existing state-of-the-art baselines with minimal overhead.