别让我找不到：有效的长上下文语言模型需要多少 KV 缓存？
计算与语言
2025-06-23 v1
摘要
语言模型处理日益增长的上下文长度，用于诸如书籍摘要等任务，这导致 KV 缓存的存储成本不断上升。许多先前工作提出了丢弃内存中 KV 条目的方法，但其方法针对特定情景，掩盖了诸多缺点，如高峰内存和性能下降，且难以进行公平比较。在本文中，我们提出 KV 足迹作为统一度量标准，其考虑了存储的 KV 条目数量及其在内存中的生命周期。我们评估基于在保持长上下文理解和生成性能的同时实现最小足迹的方法，上下文长度可达 128K 标记。该度量揭示了先前 KV 驱逐方法的高峰内存。一类方法——后填充驱逐（post-fill eviction）——由于不支持填充期间的驱逐，导致其足迹较高。我们改进了这些方法，使其能够在填充期间驱逐 KV，从而实现显著降低的 KV 足迹。随后，我们转向基于时间驱逐的方法，其中我们提出 PruLong，一种用于学习哪些注意力头需要保留完整 KV 缓存以及哪些不需要的端到端优化方法。PruLong 在保持长上下文性能的同时节省内存，实现比先前方法 12% 更小的 KV 足迹，同时在具有挑战性的召回任务中保持性能。我们的论文阐明了复杂的长上下文推理方法，并为未来开发以最小化 KV 足迹铺平了道路。
引用
@article{arxiv.2506.17121,
title = {Cache Me If You Can: How Many KVs Do You Need for Effective Long-Context LMs?},
author = {Adithya Bhaskar and Alexander Wettig and Tianyu Gao and Yihe Dong and Danqi Chen},
journal= {arXiv preprint arXiv:2506.17121},
year = {2025}
}
备注
We release our code publicly at https://github.com/princeton-pli/PruLong