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

📄 arXiv: 2605.29873v1 📥 PDF

作者: Soumyadeep Jana, Sagar Nishad, Sanasam Ranbir Singh

分类: cs.AI

发布日期: 2026-05-28


💡 一句话要点

提出Moment-KV,一种基于动量的解码时KV缓存压缩方法,用于提升长文本生成质量。

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

关键词: KV缓存压缩 长文本生成 注意力机制 动量算法 解码优化

📋 核心要点

  1. 现有KV缓存压缩方法在长文本生成中存在瓶颈,尤其是在解码阶段,缺乏对token重要性的动态建模。
  2. Moment-KV通过动量驱动的时间注意力聚合,动态建模token的重要性,兼顾长期影响和近期相关性。
  3. 实验表明,Moment-KV在长文本生成任务中,显著提升了生成质量,同时保持了解码速度。

📝 摘要(中文)

Key-Value (KV) 缓存是大型语言模型 (LLM) 在长文本生成任务中部署的主要瓶颈。现有方法通常对预填充 (prefill) 和解码缓存进行统一压缩,但压缩预填充缓存会损害关键上下文,从而降低性能。虽然保留预填充缓存至关重要,但解码阶段的压缩仍未得到充分探索,现有方法依赖于严格的近因窗口或瞬时注意力。我们对注意力动态的分析揭示了强大的时间模式:关键 token 在很长的时间范围内受到持续关注,而局部推理涉及短暂的爆发。静态启发式方法无法捕捉到这种行为,导致重要 token 过早被驱逐或陈旧 token 被保留。我们提出了 Moment-KV,一种基于动量驱动的时间注意力聚合的解码时 KV 缓存压缩方法。我们的方法将 token 重要性建模为一个不断演变的状态,其中注意力以衰减的方式聚合,从而捕捉长期影响和近期相关性。实验表明,Moment-KV 在保持解码延迟的同时,显著提高了长文本生成任务中的生成保真度 (2.3-3.2 %)。

🔬 方法详解

问题定义:长文本生成任务中,KV缓存占用大量内存,成为部署大型语言模型的瓶颈。现有方法要么统一压缩预填充和解码缓存,导致关键上下文信息丢失;要么采用静态启发式规则,无法准确捕捉token的重要性,导致重要token被过早删除或无用token被保留。

核心思路:Moment-KV的核心思想是动态地评估每个token的重要性,并根据其重要性决定是否保留在KV缓存中。它通过模拟物理中的动量概念,利用时间注意力聚合来捕捉token的长期影响和近期相关性,从而更准确地反映token的重要性。

技术框架:Moment-KV主要在解码阶段进行KV缓存压缩。其整体流程如下:1)计算每个token的注意力权重;2)使用动量更新公式,将当前注意力权重与历史注意力权重进行加权平均,得到token的重要性评分;3)根据重要性评分,决定是否保留该token在KV缓存中。

关键创新:Moment-KV的关键创新在于使用动量来聚合时间注意力。与现有方法相比,它不是简单地基于近因性或瞬时注意力来判断token的重要性,而是考虑了token在整个生成过程中的影响,从而更准确地捕捉token的重要性。

关键设计:Moment-KV的关键设计包括:1)动量更新公式:使用一个衰减因子来控制历史注意力权重的影响,从而平衡长期影响和近期相关性;2)重要性评分阈值:设置一个阈值,只有当token的重要性评分高于该阈值时,才将其保留在KV缓存中。具体参数设置需要根据实际任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,Moment-KV在长文本生成任务中,相比于基线方法,生成保真度提升了2.3-3.2%,同时保持了解码延迟。这表明Moment-KV能够在不牺牲性能的情况下,有效地压缩KV缓存,提高长文本生成的质量。

🎯 应用场景

Moment-KV适用于各种需要长文本生成的场景,例如机器翻译、文本摘要、对话生成、代码生成等。通过降低KV缓存的内存占用,它可以帮助在资源受限的设备上部署大型语言模型,并提高长文本生成的效率和质量。未来,该方法可以进一步扩展到其他类型的序列生成任务中。

📄 摘要(原文)

Key-Value (KV) cache remains a major bottleneck for deploying Large Language Models (LLMs) in long-generation tasks. Prior work often applies uniform compression across both prefill and decoding caches, but compressing the prefill cache degrades performance by corrupting critical context. While preserving the prefill cache is essential, decoding-phase compression remains underexplored, with existing methods relying on rigid recency windows or instantaneous attention. Our analysis of attention dynamics reveals strong temporal patterns: critical tokens receive sustained attention over long horizons, while local reasoning involves short-lived bursts. Static heuristics fail to capture this behavior, leading to premature eviction of important tokens or retention of stale ones. We propose Moment-KV, a decoding-time KV cache compression method based on momentum-driven temporal attention aggregation. Our method models token importance as a continuously evolving state, where attention is aggregated with decay, capturing both long-term influence and recent relevance. Experiments show that Moment-KV significantly improves generation fidelity in long-generation tasks (2.3-3.2 %) while maintaining decoding latency.