arXiv:2506.03762cs.CLcs.AI2025-06被引 3

提出AhaKV,让大模型推理时更高效地管理缓存,避免丢掉重要信息。

AhaKV: Adaptive Holistic Attention-Driven KV Cache Eviction for Efficient Inference of Large Language Models

  • 根据注意力分布自适应调整缓存淘汰权重,减少位置偏差。
  • 在固定缓存预算下,显著提升全局上下文保留能力,性能领先。
  • 适合需要高效推理的大模型部署场景,如实时对话系统。

大语言模型(LLMs)虽推动了人工智能发展,但其部署资源消耗巨大,尤其因推理时键值(Key-Value)缓存占用大量内存。现有方法通过淘汰无用标记来压缩缓存,依赖累积注意力分数作为淘汰依据。我们发现该分数存在位置偏差:随标记位置后移,数学期望递减,导致保留的标记集中在开头,削弱对全局上下文的访问。为此,我们提出自适应整体注意力驱动的键值缓存淘汰策略(AhaKV),通过自适应调节softmax尺度以抵消注意力分数的偏差,并利用被忽略的值向量信息进一步优化评分。理论上证明该方法可有效降低偏差。在多种模型上固定缓存预算进行实验,结果表明AhaKV能有效缓解偏差,保留跨全局上下文的关键信息,在多个基准任务中达到当前最优性能。

原文摘要 · Abstract (English)

Large Language Models (LLMs) have significantly advanced the field of Artificial Intelligence. However, their deployment is resource-intensive, not only due to the large number of model parameters but also because the (Key-Value) KV cache consumes a lot of memory during inference. While several works propose reducing the KV cache by evicting the unnecessary tokens, these approaches rely on accumulated attention score as eviction score to quantify the importance of the token. We identify the accumulated attention score is biased and it decreases with the position of the tokens in the mathematical expectation. As a result, the retained tokens concentrate on the initial positions, limiting model's access to global contextual information. To address this issue, we propose Adaptive holistic attention KV (AhaKV), it addresses the bias of the accumulated attention score by adaptively tuning the scale of softmax according the expectation of information entropy of attention scores. To make use of the holistic attention information in self-attention mechanism, AhaKV utilize the information of value vectors, which is overlooked in previous works, to refine the adaptive score. We show theoretically that our method is well suited for bias reduction. We deployed AhaKV on different models with a fixed cache budget. Experiments show that AhaKV successfully mitigates bias and retains crucial tokens across global context and achieve state-of-the-art results against other related work on several benchmark tasks.

大模型推理缓存优化注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。