CompressKV：语义检索头在生成前已知晓哪些令牌不重要
CompressKV: Semantic Retrieval Heads Know What Tokens are Not Important Before Generation
AI总结：
CompressKV通过识别语义检索头来确定重要令牌并保留其KV缓存，采用层自适应分配策略，在长上下文基准上超越现有方法。
摘要：
大型语言模型（LLM）的最新进展显著提升了长上下文处理能力。然而，不断增长的键值（KV）缓存大小对内存和执行效率构成了严峻挑战。大多数KV缓存压缩方法依赖于使用基于分组查询注意力（GQA）的LLM中所有注意力头进行启发式令牌驱逐。这种方法忽略了注意力头的不同功能，导致关键令牌被驱逐，从而降低了LLM的性能。为了解决上述问题，与先前工作中使用GQA-based LLM中所有注意力头来确定重要令牌不同，我们首先识别每一层中不仅能够检索提示的初始和最终令牌，而且能够检索文本中的重要令牌并关注其周围语义上下文的注意力头。随后，我们利用这些头来确定重要令牌并保留其对应的KV缓存对。此外，我们逐层分析缓存驱逐误差，并引入一种层自适应KV缓存分配策略。实验结果表明，所提出的CompressKV在LongBench和Needle-in-a-Haystack基准测试中，在各种内存预算下均持续优于最先进的方法。我们的代码公开于：https://github.com/TUDa-HWAI/CompressKV.git。
英文摘要：
Recent advances in large language models (LLMs) have significantly boosted long-context processing. However, the increasing key-value (KV) cache size poses critical challenges to memory and execution efficiency. Most KV cache compression methods rely on heuristic token eviction using all attention heads in Grouped Query Attention (GQA)-based LLMs. This method ignores the different functionalities of attention heads, leading to the eviction of critical tokens and thus degrades the performance of LLMs. To address the issue above, instead of using all the attention heads in GQA-based LLMs to determine important tokens as in the previous work, we first identify the attention heads in each layer that are not only capable of retrieving the initial and final tokens of a prompt, but also capable of retrieving important tokens within the text and attending to their surrounding semantic context. Afterwards, we exploit such heads to determine the important tokens and retain their corresponding KV cache pairs. Furthermore, we analyze the cache eviction error of each layer individually and introduce a layer-adaptive KV cache allocation strategy. Experimental results demonstrate the proposed CompressKV consistently outperforms state-of-the-art approaches under various memory budgets on LongBench and Needle-in-a-Haystack benchmarks. Our code is publicly available at: https://github.com/TUDa-HWAI/CompressKV.git.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
