BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference
作者: Janghyeon Kim, Minsoo Kim, Kyuhong Shim, Jungwook Choi
分类: cs.LG, cs.CL
发布日期: 2026-09-04
备注: ICML 2026. Code: https://github.com/aiha-lab/BeaconKV
💡 一句话要点
提出BeaconKV以解决大规模推理模型推理中的KV缓存压缩问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大规模推理模型 KV缓存压缩 思维链 信标查询 内存优化 推理效率 人工智能
📋 核心要点
- 现有的KV缓存压缩方法在长时间推理中面临严重的内存瓶颈,无法有效处理大规模推理模型的需求。
- 论文提出BeaconKV,通过维护信标查询来压缩KV缓存,避免存储整个查询历史,从而提高推理效率。
- 实验结果显示,BeaconKV在内存使用上减少了5.8倍,同时保持了缓存的准确性,并提升了4.3倍的推理吞吐量。
📝 摘要(中文)
大型推理模型(LRMs)通过扩展的思维链(CoT)生成实现了卓越的问题解决能力,但随之而来的键值(KV)缓存随着序列长度线性增长,造成严重的内存瓶颈,常常超出GPU的容量。现有的KV缓存压缩方法依赖于最近的查询来估计未来的token重要性,隐含假设这些查询可以作为未来注意力模式的可靠代理。然而,我们发现这一假设在长时间推理中并不成立:某些解码步骤会生成重新关注远程上下文的思维重访token(TRT)。基于这一发现,我们提出了BeaconKV,一种无训练的KV缓存压缩方法,通过维护信标查询,作为每个全局查询簇的紧凑代表,来预测哪些KV对将被重新访问,而无需存储整个查询历史。实验表明,BeaconKV在四个开源LRMs和多种推理基准上普遍优于现有压缩方法,实现了高达5.8倍的内存减少,同时几乎保持了完整的缓存准确性,并提高了4.3倍的吞吐量。
🔬 方法详解
问题定义:论文要解决的问题是大型推理模型在推理过程中KV缓存的线性增长导致的内存瓶颈。现有方法依赖于最近的查询来预测未来的token重要性,但在长时间推理中,这种方法的假设并不成立。
核心思路:论文的核心解决思路是提出BeaconKV,通过维护信标查询作为每个全局查询簇的代表,来预测哪些KV对将被重新访问,而无需存储整个查询历史。这种方法能够有效减少内存使用,同时保持推理的准确性。
技术框架:BeaconKV的整体架构包括信标查询的维护和KV对的压缩。首先,通过对查询的系统分析,识别出思维重访token(TRT)并将其聚类为相似性组。然后,利用这些信标查询来指导KV缓存的压缩过程。
关键创新:最重要的技术创新点在于引入信标查询的概念,作为全局查询簇的紧凑代表。这一方法与现有的依赖于最近查询的压缩方法本质上不同,能够更准确地预测未来的注意力模式。
关键设计:在设计中,BeaconKV采用了无训练的方式,避免了复杂的训练过程。关键参数设置包括信标查询的数量和聚类算法的选择,以确保在压缩过程中尽可能保留重要的KV对。
🖼️ 关键图片
📊 实验亮点
实验结果表明,BeaconKV在四个开源LRMs上表现优异,内存使用减少了高达5.8倍,同时几乎保持了完整的缓存准确性,并在推理吞吐量上提升了超过4.3倍,显示出其在实际应用中的巨大潜力。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、图像识别和其他需要大规模推理的人工智能任务。通过有效的KV缓存压缩,BeaconKV能够显著提高推理速度和内存利用率,具有广泛的实际价值和未来影响。
📄 摘要(原文)
Large Reasoning Models (LRMs) achieve superior problem-solving through extended Chain-of-Thought (CoT) generation, but the resulting key-value (KV) cache grows linearly with sequence length and creates severe memory bottlenecks, often exceeding GPU capacity for long reasoning traces. Existing KV cache compression methods rely on recent queries to estimate future token importance, implicitly assuming these serve as reliable proxies for future attention patterns. We demonstrate that this assumption fails in long-horizon reasoning: certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context, such as task-solving plans formulated early in the trace. Through systematic analysis, we discover that queries corresponding to the TRT cluster into a small number of similarity groups in the embedding space. Based on this insight, we propose BeaconKV, a training-free KV cache compression method that maintains beacon queries, compact representatives for each global query cluster, to anticipate which KV pairs will be revisited without storing the entire query history. Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms existing compression methods, achieving up to $5.8\times$ memory reduction while nearly preserving full cache accuracy and improving throughput by over $4.3\times$.