提出新型缓存修剪方法，提升长文本推理效率与精度。
OBCache: Optimal Brain KV Cache Pruning for Efficient Long-Context LLM Inference
- 基于注意力输出扰动量化令牌重要性，实现精准缓存淘汰
- 在LLaMA和Qwen模型上显著提升长上下文任务准确率
- 适合需要高效长序列推理的AI系统开发者
大语言模型（LLMs）支持更长上下文窗口以实现强大应用，但需缓存所有键值（KV）状态，内存开销随序列长度和批量大小线性增长。现有缓存淘汰方法虽利用注意力稀疏性，但通常通过累积注意力权重启发式排序令牌，未考虑其对注意力输出的真实影响。本文提出最优大脑缓存（OBCache），将缓存淘汰建模为逐层结构化剪枝问题。基于最优大脑损伤（OBD）理论，通过测量剪除令牌引起的注意力输出扰动来量化令牌重要性，推导出孤立键、孤立值及联合键值对的闭式评分。该方法不仅考虑注意力权重，还融合值状态与注意力输出信息，增强现有策略的输出感知能力。在LLaMA与Qwen模型上的实验表明，用OBCache的输出感知评分替代原有启发式评分，能持续提升长上下文准确性。代码已开源：https://github.com/DreamSoul-AI/OBCache。
原文摘要 · Abstract (English)
Large language models (LLMs) with extended context windows enable powerful applications but impose significant memory overhead, as caching all key-value (KV) states scales linearly with sequence length and batch size. Existing cache eviction methods address this by exploiting attention sparsity, yet they typically rank tokens heuristically using accumulated attention weights without considering their true impact on attention outputs. We propose Optimal Brain Cache (OBCache), a principled framework that formulates cache eviction as a layer-wise structured pruning problem. Building upon the Optimal Brain Damage (OBD) theory, OBCache quantifies token saliency by measuring the perturbation in attention outputs induced by pruning tokens, with closed-form scores derived for isolated keys, isolated values, and joint key-value pairs. Our scores account not only for attention weights but also for information from value states and attention outputs, thereby enhancing existing eviction strategies with output-aware signals. Experiments on LLaMA and Qwen models demonstrate that replacing the heuristic scores in existing works, which estimate token saliency across different query positions, with OBCache's output-aware scores consistently improves long-context accuracy. Code is available at https://github.com/DreamSoul-AI/OBCache.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。