arXiv:2606.01563cs.LG2026-06被引 2

提出新方法缓解长文本推理中缓存淘汰的定向偏差问题。

MomentKV: Closing the Directional Gap in KV Cache Eviction for Long-Context Inference

论文配图:MomentKV: Closing the Directional Gap in KV Cache Eviction for Long-Context Inference
图 1 · 摘自论文原文
  • 通过统计被淘汰令牌的几何特性,优化淘汰策略
  • 在不同缓存预算下均优于现有方法,极端压缩时提升显著
  • 适合需要高效长文本推理的场景,如智能写作与知识问答

基于Transformer的语言模型自回归解码依赖KV缓存,其内存随序列长度线性增长,成为长上下文推理的主要瓶颈。通过保留固定大小的关键值对子集并丢弃其余内容来解决这一问题。我们发现输出质量下降的主要原因并非已淘汰令牌的残余注意力质量(现有方法已最小化),而是保留与淘汰令牌集合之间的方向不匹配——实际中被淘汰的令牌往往与保留的令牌近似正交。因此,即使少量被淘汰的注意力质量也可能对最终方向分布产生巨大影响,导致显著输出误差。这揭示了现有策略的根本局限。为此,我们提出MomentKV,维护被淘汰令牌集合的小型矩统计量,包括数量、键均值、值均值及值-键协方差。在淘汰阶段,利用这些统计量识别已与累积摘要良好对齐的令牌,保持淘汰集合的几何规律性;在推理阶段,生成被淘汰注意力输出的一阶闭式近似,形成选择性淘汰与精确修正的良性循环。在LongBench和RULER数据集上,使用LLaMA-3.1-8B-Instruct和Qwen3-4B-Instruct模型,MomentKV在所有缓存预算下均超越所有基线方法,尤其在高压缩比下表现最优。

原文摘要 · Abstract (English)

Autoregressive decoding in Transformer-based language models relies on the KV cache, whose memory footprint grows linearly with sequence length and becomes the primary bottleneck for long-context inference. KV cache eviction addresses this by retaining a fixed-size subset of key-value pairs and discarding the rest. We identify that a primary source of output degradation is not the residual attention mass on evicted tokens, which existing methods already minimize, but a directional mismatch between the retained and evicted token sets. Specifically, the evicted tokens in practice are often near-orthogonal to the retained ones. Thus, even a small evicted mass could have an oversized impact on the resulting direction distribution and amplify into substantial output error. This reveals a fundamental limit in existing strategies. To address this, we propose MomentKV, which maintains compact, small-size moment statistics over the evicted token set, including a count, key mean, value mean, and value-key covariance. During eviction, the moment statistics is leveraged to identify tokens already well aligned with and captured by the accumulated summary, keeping the evicted set geometrically regular. During inference, they yield a closed-form first-order approximation of the evicted attention output, forming a mutually reinforcing loop between selective eviction and accurate correction. On LongBench and RULER with LLaMA-3.1-8B-Instruct and Qwen3-4B-Instruct, MomentKV outperforms all baselines at every cache budget, with the largest gains under aggressive compression.

长文本推理缓存优化注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。