让大模型自动删掉不重要的记忆,省内存还提速。
Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility

- 用轻量预测器判断每个记忆条目的未来价值,只保留高价值的。
- 实测能将缓存大小压缩3到10倍,长文本压缩效果更明显。
- 适合需要长时间推理的智能体和超长文本生成场景。
在现代测试时计算和智能体范式下,语言模型需处理越来越长的序列。基于Transformer的高效文本生成正受到键值缓存内存占用和带宽的制约。为此,我们提出自剪枝键值注意力(SP-KV),通过预测未来键值对的利用价值来减少长期键值缓存的大小。该机制以细粒度运行:一个轻量级利用价值预测器为每个键值对打分,近期键值始终通过局部窗口可用,而较旧的键值仅在预测价值超过阈值时才写入缓存并用于全局注意力。语言模型与利用价值预测器通过纯下一词预测损失端到端联合训练,并从预训练检查点迁移而来。不同于固定压缩率,SP-KV实现动态稀疏化:根据输入自适应调整,通常使键值缓存大小缩小3至10倍,长序列更具可压缩性。这显著提升了内存使用效率和解码速度,验证损失及下游任务性能几乎无损。此外,该方法揭示了结构化的层与头特异性稀疏模式,可用于指导混合局部-全局注意力架构的设计。
原文摘要 · Abstract (English)
Under modern test-time compute and agentic paradigms, language models process ever-longer sequences. Efficient text generation with transformer architectures is increasingly constrained by the Key-Value cache memory footprint and bandwidth. To address this limitation, we introduce Self-Pruned Key-Value Attention (SP-KV), a mechanism designed to predict future KV utility in order to reduce the size of the long-term KV cache. This strategy operates at a fine granularity: a lightweight utility predictor scores each key-value pair, and while recent KVs are always available via a local window, older pairs are written in the cache and used in global attention only if their predicted utility surpasses a given threshold. The LLM and the utility predictor are trained jointly end-to-end exclusively through next-token prediction loss, and are adapted from pretrained LLM checkpoints. Rather than enforcing a fixed compression ratio, SP-KV performs dynamic sparsification: the mechanism adapts to the input and typically reduces the KV cache size by a factor of $3$ to $10\times$, longer sequences often being more compressible. This leads to vast improvements in memory usage and decoding speed, with little to no degradation of validation loss nor performance on a broad set of downstream tasks. Beyond serving as an effective KV-cache reduction mechanism, our method reveals structured layer- and head-specific sparsity patterns that we can use to guide the design of hybrid local-global attention architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。