针对推理过程的内存瓶颈,提出答案优先的缓存管理方法
Crystal-KV: Efficient KV Cache Management for Chain-of-Thought LLMs via Answer-First Principle
- 根据答案重要性区分有效与无效缓存,动态识别可淘汰内容
- 在不降低准确率前提下,压缩缓存达60%以上,吞吐提升2.1倍
- 适合需要长推理链的大模型应用,如复杂问答与数学推演
大语言模型的思维链(CoT)推理虽能显著提升复杂任务准确性,但因长期推理序列存储于键值(KV)缓存中,导致内存开销过大。传统生成任务中所有词元同等重要,而CoT强调最终答案,使常规压缩策略失效。本文提出Crystal-KV,一种面向CoT推理的高效缓存管理框架。核心思想为答案优先原则:通过将答案偏好映射至思考阶段注意力图,区分出仅维持推理流程但可能引入误导信息的滑动缓存(SlipKV)与真正贡献于答案正确性的晶体缓存(CrystalKV)。进一步提出基于注意力的最近最少频繁使用算法,精准判断SlipKV失效时机并剔除,保留CrystalKV而不干扰推理流程。最后设计自适应缓存预算分配机制,依据各层/头中CrystalKV比例动态调整缓存分配,强化关键组件以提升资源利用效率。实验表明,Crystal-KV实现业界领先的缓存压缩效果,显著提升吞吐量,加快响应速度,同时保持甚至提高答案准确性。
原文摘要 · Abstract (English)
Chain-of-Thought (CoT) reasoning in large language models (LLMs) significantly improves accuracy on complex tasks, yet incurs excessive memory overhead due to the long think-stage sequences stored in the Key-Value (KV) cache. Unlike traditional generation tasks where all tokens are uniformly important, CoT emphasizes the final answer, rendering conventional KV compression strategies ineffective. In this paper, we present Crystal-KV, an efficient KV cache management framework tailored for CoT reasoning. Our key insight is the answer-first principle. By mapping answer preferences into think-stage attention map, we distinguish between SlipKV, which mainly maintains the reasoning flow but may occasionally introduce misleading context, and CrystalKV, which truly contributes to the correctness of the final answer. Next, we propose an attention-based Least Recently Frequently Used algorithm. It precisely identifies when a SlipKV entry's utility expires and evicts it, retaining CrystalKV without disrupting reasoning flow. Finally, we introduce an adaptive cache budget allocation algorithm. Based on the dynamic proportion of CrystalKV, it estimates the importance of each layer/head and adjusts the KV cache budget during inference, amplifying critical components to improve budget utilization. Results show that Crystal-KV achieves state-of-the-art KV cache compression, significantly improves throughput, and enables faster response time, while maintaining, or even improving, answer accuracy for CoT reasoning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。