MemDecay:区域感知的KV缓存淘汰策略,提升LLM代理推理效率
研究提出MemDecay,一种训练无关的区域感知KV缓存淘汰策略,通过为不同语义区域的令牌分配不同优先级和衰减率,在固定缓存预算下保留关键信息。实验表明,系统令牌的半衰期远长于暂存区令牌,且固定策略能在全部设置中保持完美准确率,而现有基线最多仅保留13/24。
近日,一篇来自arXiv的论文提出了MemDecay,一种无需训练的区域感知KV缓存淘汰策略,旨在解决大型语言模型(LLM)代理在推理过程中的内存瓶颈问题。LLM代理在复杂任务中需要累积系统指令、计划、用户轮次、检索文档、工具输出和中间推理等多种异构上下文,这些上下文的键值(KV)缓存可能占据大量显存。传统的缓存淘汰策略通常对所有令牌一视同仁,基于注意力或最近性进行淘汰,但忽略了代理编排器已经拥有的语义结构信息。
MemDecay的创新之处在于为不同语义区域的令牌分配不同的基础优先级和衰减率,并在令牌受到注意力时刷新其保留分数。在固定的缓存预算下,MemDecay会淘汰得分最低的页面,同时允许关键区域(如系统指令)被固定(pinning),从而确保重要信息不被丢失。此外,论文还提供了一种从测量的注意力生命周期中校准衰减率的方法,使得策略可以自适应地调整。
研究人员使用Qwen2.5-1.5B和3B模型,在约450和1,700个令牌的上下文中进行了评估。实验结果表明,不同区域的注意力生命周期相差一个数量级:系统令牌的半衰期范围为148到189个解码步,而暂存区令牌仅为14到16步。通过固定系统区域,在所有设置下系统区域事实的准确率均达到了完全缓存时的水平,而现有的基线方法最多只能保留24个事实中的13个。随着上下文增长,区域感知的保留策略仍然有效,而基于最近性的保留策略则迅速崩溃。不过,在非固定内容上,累积注意力保留策略表现更好,消融实验表明注意力分数归一化是当前公式的主要限制。这些结果清楚地表明,语义提示结构是KV缓存管理的一个稳健信号,同时也阐明了如何将其与基于注意力的重要性相结合。
MemDecay的提出为高效LLM代理推理提供了新的思路,有望在模型部署和推理成本优化中发挥重要作用。未来工作可以进一步探索如何将区域感知与其他注意力机制结合,以提升非固定内容的保留效果。