论文

Moment-KV:面向长生成的基于动量的解码时 KV 缓存压缩

Moment-KV: Momentum-Based Decode-Time KV Cache Compression for Long Generation

模型推理KV Cache

摘要

键值(KV)缓存仍然是在长生成任务中部署大语言模型(LLM)的主要瓶颈。先前工作常对预填充(prefill)与解码两个阶段的缓存施加统一压缩,但压缩预填充缓存会破坏关键上下文,从而降低性能。尽管保留预填充缓存至关重要,解码阶段的压缩却仍缺乏研究,现有方法依赖僵化的近因窗口或瞬时注意力。我们对注意力动态的分析揭示了强的时间模式:关键 token 在长时程上持续获得注意力,而局部推理则表现为短促的爆发。静态启发式无法捕捉这种行为,导致重要 token 被过早逐出或陈旧 token 被保留。我们提出 Moment-KV,一种基于动量驱动的时间注意力聚合的解码时 KV 缓存压缩方法。我们的方法将 token 重要性建模为一个持续演化的状态,对注意力进行带衰减的聚合,从而同时捕获长期影响与近期相关性。实验表明,Moment-KV 显著提升了长生成任务中的生成保真度(2.3–3.2%),同时保持解码延迟不变。

Moment-KV:面向长生成的基于动量的解码时 KV 缓存压缩:论文配图
图 1:GSM8K 上的 LLaMA-3.1-8B-Instruct 的注意力热图,显示了第 16 层和第 25 层的预填充和解码阶段的令牌注意力模式。缩放区域突出显示了解码期间注意力权重的暂时下降,其中重要的令牌暂时受到较低的关注。