CompressKV：语义检索头在生成前了解哪些 token 不重要
计算与语言
2025-08-05 v1 人工智能
摘要
最近的大型语言模型(LLM)技术显著提升了长上下文处理能力，但不断增长的 key-value(KV)缓存规模对内存和执行效率构成严峻挑战。大多数KV缓存压缩方法依赖于基于分组查询注意力(GQA)的LLM中所有注意力头的启发式token驱逐策略。该方法忽略了注意力头的不同功能，导致关键token被驱逐，从而降低LLM性能。为解决上述问题，我们首先识别出每层中既能检索提示的初始和最终token，又能检索文本中重要token并注意其周围语义上下文的注意力头。随后，我们利用这些头来确定重要token并保留其对应的KV缓存对。进一步，我们对每层的缓存驱逐误差进行分析，引入基于层的自适应KV缓存分配策略。实验结果表明，所提出的CompressKV在LongBench和Needle-in-a-Haystack基准测试的 various 内存预算下 consistently 优于最先进的方法。我们的代码已公开：https://github.com/TUDa-HWAI/CompressKV.git。
引用
@article{arxiv.2508.02401,
title = {CompressKV: Semantic Retrieval Heads Know What Tokens are Not Important Before Generation},
author = {Xiaolin Lin and Jingcun Wang and Olga Kondrateva and Yiyu Shi and Bing Li and Grace Li Zhang},
journal= {arXiv preprint arXiv:2508.02401},
year = {2025}
}