Value-Aware Stochastic KV Cache Eviction for Reasoning Models
原作者: Ting-Yun Chang, Harvey Yiyun Fu, Deqing Fu, Chenghao Yang, Jesse Thomason, Robin Jia
原作者: Ting-Yun Chang, Harvey Yiyun Fu, Deqing Fu, Chenghao Yang, Jesse Thomason, Robin Jia
原始论文采用 CC BY 4.0 许可(http://creativecommons.org/licenses/by/4.0/)。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性,请参阅原始论文。 阅读完整免责声明
技术摘要:面向推理模型的价值感知随机 KV 缓存驱逐 (VASE)
1. 问题陈述
推理模型(如 Qwen3、OpenAI 的 o1)通过在生成最终答案之前生成长链条思维(Chain of Thought)来实现高准确率。然而,这种能力在解码阶段造成了显著的内存和计算瓶颈。随着序列长度的增加,用于存储每个过去 Token 表示的键值(KV)缓存所产生的开销巨大。
现有的解决方案分为两类:
- 基于选择的方法(Selection-based methods): 这些方法保留完整的 KV 缓存,但在计算注意力时仅激活一部分稀疏的 Token 子集。虽然准确,但其内存占用随序列长度线性增长(O(T)),无法解决内存瓶颈问题。
- 基于驱逐的方法(Eviction-based methods): 这些方法在达到预设预算后永久丢弃低重要性的 KV 对,提供静态内存占用和更好的吞吐量。然而,目前的驱逐方法在处理推理任务时,与基于选择的方法相比,存在明显的准确率下降问题,常导致模型进入重复推理循环或生成无意义的输出。
本文指出,当前的驱逐策略未能考虑到两个关键因素:**大模长值状态(large-magnitude value states)的不成比例影响,以及保留 Token 时对随机多样性(stochastic diversity)**的需求。
2. 方法论:VASE
作者提出了 价值感知随机 KV 缓存驱逐(Value-Aware Stochastic KV Cache Eviction, VASE),这是一个旨在弥合效率与准确率之间差距的无需训练(training-free)的驱逐框架。VASE 在一个周期性驱逐框架内运行(使用持久预算 K 和近期缓冲区 B),并引入了两个核心机制:
A. 对大模长值状态的保护
作者观察到,推理模型中的值状态(value states)呈现出强烈的偏斜分布,其中一小部分 Token 拥有异常大的向量模长(通过范围 Range(v)=max(v)−min(v) 来衡量)。
- 发现: 驱逐这些高模长值会导致灾难性的准确率崩溃(例如,在 GSM8K 上从 ~88% 降至 14%),并诱发重复循环,使模型在无法得出结论的情况下无休止地重新检查上下文。
- 机制: VASE 为这些高模长值预留了特定的 Token 预算(Nv),无条件保留模长最大的 Nv 个 Token。这确保了最具影响力的值向量永远不会被丢弃。
B. 引入随机性
目前的驱逐方法通常使用确定性的 top-k 选择,这可能导致保留的缓存缺乏多样性。
- 发现: 引入随机性能通过确保对完整上下文更具代表性的覆盖来提高准确率。
- 机制: VASE 不再确定性地选择得分最高的 Token,而是采用加权随机采样。
- VASE-AttnV: 将价值感知预留与基于注意力分数(源自 SnapKV)的随机采样相结合。
- VASE-DKV: 改进了 CurDKV 方法(该方法使用 CUR 矩阵分解的杠杆得分),通过在每次驱逐步骤中重新采样高斯投影矩阵 G。这防止了具有特定表示的 Token 被持续分配低分并被永久驱逐。
3. 核心贡献
- 识别关键因素: 本文确立了:(1)大模长值状态对于维持推理进程和防止重复循环至关重要;(2)驱逐决策中的随机性通过增加缓存多样性显著提升了准确率。
- VASE 框架: 一种新颖的、无需训练的驱逐方案,集成了价值状态模长保护和随机采样。它是第一个结合了基于 Key 的评分、基于 Value 的评分以及多样性促进的驱逐方法。
- 与量化的联系: 作者证明了大范围值状态也是每 Token KV 缓存量化中重建误差的主要来源,这表明 VASE 的见解可以推广到其他压缩技术中。
4. 实验结果
作者在 Qwen3-4B 和 Qwen3-14B 上针对六个推理任务(AIME25/26, HMMT25, GPQA-Diamond, MATH, LiveCodeBench-v6)进行了评估,采用了 4 倍 KV 缓存压缩率。
- 准确率 vs. 选择方法: VASE-AttnV 在保持静态内存占用的同时,在两种模型规模上都实现了比最强的选择类方法(SeerAttention-R)更高的平均准确率。
- Qwen3-4B: VASE-AttnV (59.09%) 优于 SeerAttention-R (58.81%) 和最强的驱逐基准 R-KV (54.69%),提升了 4.4%。
- Qwen3-14B: VASE-AttnV (65.81%) 匹配了 SeerAttention-R (65.37%),并优于 R-KV (60.90%) 4.9%。
- 消融实验:
- 价值感知: 为大模长值预留插槽使 GSM8K 的准确率比基准方法提高了高达 16.2%。
- 随机性: 在 CurDKV 中加入随机采样使 Qwen3-14B 的准确率提升了 9.2%。
- 效率: VASE-DKV 实现了最高的吞吐量(在 16K tokens 时比全量模型基准快 3.1 倍)以及所有测试方法中最低的峰值内存使用量。
- 代码生成: 在 LiveCodeBench 上,VASE 方法显著优于面临领域偏移问题的选择类方法 SeerAttention-R。
5. 重要性与主张
本文声称 VASE 成功弥合了长期困扰 KV 缓存驱逐方法的“效率-准确率”鸿沟。通过优先考虑大模长值状态并引入随机性,VASE 使推理模型能够在不牺牲通常与全量缓存或选择类方法相关的准确率的前提下,实现静态内存占用下的高效运行。
作者强调,关于价值状态模长重要性的发现具有超越驱逐领域的更广泛意义,特别是对于 KV 缓存量化——其中大范围值被确定为主要的误差来源。他们建议,未来的内存高效推理方法应考虑保护这些关键高模长状态的混合精度方法。
最终,VASE 提供了一个简单、有效且无需训练的方案,以支持 FlashAttention2 并实现长链推理模型的可扩展推理。
您所在领域的论文太多了?
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要,使用您的语言。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。