Self-Pruned Key-Value Attention:让LLM只保留未来有用的记忆
KV缓存是长上下文推理的瓶颈,这篇论文用自学习剪枝解决了内存爆炸问题,做LLM推理优化或长文本应用的开发者可以直接参考其方法。
一篇新论文提出Self-Pruned Key-Value Attention方法,让大语言模型在长文本生成时只保留对后续token有用的历史键值对,从而大幅压缩KV缓存。
该方法通过一个小型预测器为每个键值对打分,只保留高分项,同时确保最近token始终保留。模型在训练时通过正常的next-token预测学习剪枝策略,无需手工规则。
实验表明,模型通常只保留10%到33.7%的旧键值对,性能接近全注意力,解码速度在长上下文场景下提升2.
1到4.6倍。
This paper teaches LLMs to save memory by keeping only past tokens likely to matter later.
The problem is that long text generation makes the key-value cache grow, and this cache is the model’s working memory of earlier tokens.
Instead of saving every old token, the paper adds a small predictor that scores each key-value pair by how useful it seems for future tokens.
Recent tokens are always kept, because nearby words usually matter, but older tokens enter the long-term cache only when their score is high enough.
The authors trained this system together with the LLM using only normal next-token prediction, so the model learns its own pruning behavior rather than following a fixed hand-made rule.
They tested it across model sizes, long-context settings, downstream tasks, and decoding speed, then compared it with full attention and several cache-pruning methods.
The main result is that the model usually keeps only about 10% to 33.7% of older key-value entries, while matching normal performance closely and reaching 2.1 to 4.6 times faster decoding in some long-context batches.
----
Paper Link – arxiv. org/abs/2605.14037
Paper Title: "Self-Pruned Key-Value Attention: Learning When to Write by Predicting Future Utility"
This report is categorized under AITOP's "Papers" desk. Key entities and technical topics under continuous tracking include KV缓存, 长上下文, 注意力机制, 模型压缩.
Source attribution: curated from rohanpaul_ai (@rohanpaul_ai).