AI论文精选
⏶24

KVpop：结合预测性在线剪枝的键值缓存压缩

发表
由 LukasLukas 提交
作者: LukasLukas Hauzenberger, Niklas Schmidinger, Anamaria-Roberta Hartl, David Stap, Thomas Schmied, Sebastian Böck, Günter Klambauer, Sepp Hochreiter

摘要

AI 生成总结
KVpop 通过使用未来注意力目标直接监督保留或丢弃决策，学习最优的键值缓存逐出策略，在减少内存使用的同时实现了高性能。
键值（KV）缓存的增长是自回归解码中的一个主要瓶颈，因为内存和带宽与上下文长度呈线性关系。现有的KV驱逐方法通常依赖静态启发式或代理分数，这些分数无法很好地跟踪未来令牌的效用，并导致随着相关性变化而出现脆弱的驱逐。为了解决这个问题，我们引入了KVpop，它通过直接监督保留或丢弃决策来学习固定预算的KV驱逐策略。评分器针对新颖的未来注意力目标进行训练，无需具体化密集注意力图即可高效计算。我们进一步引入了一种基于延迟内存的评分器，它在所有学习的驱逐方法中独一无二地将评分延迟固定步数，以利用近未来上下文。在AIME和HMMT数学推理任务上，KVpop在Qwen3-4B模型上，在KV缓存压缩75%时保留了98%的完全注意力性能，在压缩88%时保留了97%，始终优于已有的驱逐基线。Qwen3-8B显示出更强的结果，达到了接近完全教师的性能。这些结果表明，使用未来注意力信号监督驱逐可以降低内存成本，同时保持质量。
查看 arXiv 页面查看 PDF

评论

LukasLukas
论文作者
论文提交者

键值（KV）缓存的增长是自回归解码中的一个主要瓶颈，因为内存和带宽会随上下文长度线性扩展。现有的KV驱逐方法通常依赖静态启发式或代理分数，这些方法无法很好地跟踪未来令牌的效用，并在相关性发生变化时导致脆弱的驱逐。为解决此问题，我们引入了KVpop，它通过直接监督保留或丢弃决策来学习一种固定预算的KV驱逐策略。评分器针对一种新颖的未来注意力目标进行训练，无需具体化密集注意力图即可高效计算。我们进一步引入了一种延迟的基于记忆的评分器，它在所有学习到的驱逐方法中独树一帜，可将评分延迟固定步数，以利用近未来的上下文。在AIME和HMMT数学推理任务中，KVpop在Qwen3-4B模型上，以75%的KV缓存压缩率保持了98%的全注意力性能，以88%的压缩率保持了97%的性能，持续优于已建立的驱逐基线。Qwen3-8B显示出更强的结果，接近完全教师模型的性能。这些结果表明，利用未来注意力信号监督驱逐可以降低内存成本，同时保持质量。

Aamer MihaysiAamer Mihaysi

KV缓存的增长是长上下文LLM部署的隐形杀手。大多数驱逐策略只是基于静态分数进行猜测，一旦对话焦点转移，这种策略就会失效。KVpop通过监督保留或丢弃决策以针对未来注意力目标的方法是处理这个问题的正确方式——它将缓存管理视为一个预测问题而非启发式问题。我特别对用于稳定性的延迟机制感兴趣。如果这真的能在上下文切换时减少内存开销而不增加困惑度，那么对于延迟预算紧张的生产代理系统来说，这将是一个巨大的胜利。