Crystal-KV: 基于答案优先原则的思维链大语言模型高效KV缓存管理
计算与语言
2026-01-27 v1 人工智能
机器学习
摘要
大语言模型(LLM)中的思维链(CoT)推理显著提升了复杂任务的准确性，但由于存储在键值(KV)缓存中的长思考阶段序列，导致了过多的内存开销。与所有令牌同等重要的传统生成任务不同，CoT强调最终答案，这使得传统的KV压缩策略失效。在本文中，我们提出了Crystal-KV，一个专为CoT推理定制的高效KV缓存管理框架。我们的关键见解是答案优先原则。通过将答案偏好映射到思考阶段的注意力图，我们区分了主要维持推理流程但偶尔引入误导性上下文的SlipKV，以及真正有助于最终答案正确性的CrystalKV。接着，我们提出了一种基于注意力的最近最少使用频率算法。该算法精确识别SlipKV条目的效用何时过期并将其驱逐，保留CrystalKV而不中断推理流程。最后，我们引入了一种自适应缓存预算分配算法。基于CrystalKV的动态比例，它估计每个层/头的重要性，并在推理过程中调整KV缓存预算，放大关键组件以提高预算利用率。结果表明，Crystal-KV实现了最先进的KV缓存压缩，显著提高了吞吐量，并实现了更快的响应时间，同时保持甚至提高了CoT推理的答案准确性。
引用
@article{arxiv.2601.16986,
title = {Crystal-KV: Efficient KV Cache Management for Chain-of-Thought LLMs via Answer-First Principle},
author = {Zihan Wang and Cheng Tang and Lei Gong and Cheng Li and Chao Wang and teng wang and Wenqi Lou and Xuehai Zhou},
journal= {arXiv preprint arXiv:2601.16986},
year = {2026}
}