论文

MomentKV:缩小长上下文推理的 KV 缓存驱逐中的方向差距

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

模型推理KV Cache

摘要

基于 Transformer 的语言模型中的自回归解码依赖于 KV 缓存,其内存占用随着序列长度线性增长,成为长上下文推理的主要瓶颈。 KV 缓存驱逐通过保留固定大小的键值对子集并丢弃其余部分来解决这个问题。我们发现,输出退化的主要来源不是被驱逐词元上的剩余注意力质量(现有方法已经将其最小化),而是保留词元集和被驱逐词元集之间的方向不匹配。具体来说,实践中被逐出的词元通常与保留的词元接近正交。因此,即使是很小的驱逐质量也可能对最终的方向分布产生过大的影响,并放大为显着的输出误差。这揭示了现有策略的根本局限性。为了解决这个问题,我们提出了 MomentKV,它在被逐出的词元集上维护紧凑、小尺寸的矩统计数据,包括计数、键均值、值均值和值键协方差。在逐出过程中,利用时刻统计数据来识别已经与累积摘要很好地对齐并由累积摘要捕获的词元,从而使逐出集保持几何规则。在推理过程中,它们产生被逐出注意力输出的封闭式一阶近似,在选择性逐出和精确校正之间形成相互增强的循环。在使用 LLaMA-3.1-8B-Instruct 和 Qwen3-4B-Instruct 的 LongBench 和 RULER 上,MomentKV 在每个缓存预算下都优于所有基准,在积极压缩下收益最大。