AhaKV：基于自适应整体注意力驱动的 KV 缓存驱逐方法，用于高效推理大语言模型
计算与语言
2025-06-05 v1 人工智能
摘要
大语言模型（LLM）显著推动了人工智能领域的发展。然而，由于模型参数数量庞大以及（Key-Value）KV 缓存在推理过程中消耗大量内存，其部署具有资源密集型特征。虽然已有多个研究提出通过驱逐不必要的标记来减少 KV 缓存，但这些方法依赖注意力得分的累积值作为驱逐得分来量化标记的重要性。我们指出，累积注意力得分存在偏差，且随着标记在数学期望中的位置而减小。结果是，保留下来的标记集中在初始位置，限制了模型获取全局上下文信息的能力。为解决此问题，我们提出了自适应整体注意力 KV（AhaKV），通过根据注意力得分信息熵的期望自适应调节 softmax 的比例来消除累积注意力得分的偏差。为了利用自注意力机制中整体注意力信息，AhaKV 利用价值向量的信息，这在以往工作中被忽略。我们在理论上证明了该方法适合偏差消除。在固定缓存预算下，我们在不同模型上部署 AhaKV。实验表明，AhaKV 成功消除了偏差，跨全局上下文保留关键标记，并在多个基准任务上相较于其他相关工作实现了业界最佳效果。
引用
@article{arxiv.2506.03762,
title = {AhaKV: Adaptive Holistic Attention-Driven KV Cache Eviction for Efficient Inference of Large Language Models},
author = {Yifeng Gu and Zicong Jiang and Jianxiu Jin and Kailing Guo and Ziyang Zhang and Xiangmin Xu},
journal= {arXiv preprint arXiv:2506.03762},
year = {2025}
}
备注
14 pages, 8 figures