面向推理模型的值感知随机 KV 缓存驱逐 Value-Aware Stochastic KV Cache Eviction for Reasoning Models
保护大值状态并用随机采样提升多样性的训练免 KV 缓存驱逐方法。
前置知识
KV 缓存 (KV Cache)
Transformer 自回归生成时,为每个历史 token 保存的键(key)和值(value)向量,避免重复计算注意力。缓存随序列长度线性增长,是推理阶段的显存与带宽瓶颈。
本文所有方法都围绕如何压缩 KV 缓存展开,不理解缓存的线性增长就无法理解动机。
推理模型 / 思维链 (Chain of Thought)
通过生成长思维链提升准确率的大语言模型(如 Qwen3)。问题虽短输出却极长,AIME25 平均生成 17790 token,远超 prompt 的 198 token,加剧 KV 缓存压力。
本文专门针对推理模型解码阶段的超长输出,这是 KV 缓存爆炸的主要来源。
稀疏注意力 (Sparse Attention)
注意力只关注部分历史 token 的方法,分两类:选择类保留完整缓存但激活稀疏子集(内存随序列线性增长),驱逐类永久丢弃低重要性 KV 对(实现静态内存但可能丢信息)。
选择与驱逐的根本区别是本文的核心对比,理解二者权衡才能看懂 VASE 的定位。
注意力汇聚点 (Attention Sink)
序列开头的少数 token 会吸收不成比例的注意力权重,常与低幅度值状态共现,对驱逐方法至关重要。本文指出与之互补的大幅度值状态同样关键。
区分sink(低幅度)与本文重点(大幅度值)是理解其创新点与已有工作的差异。
KV 缓存量化 (KV Cache Quantization)
把 KV 缓存从高精度压缩到低比特(如2/4 bit)以省显存。逐 token 量化时 Range 大的值状态会产生不成比例的重建误差(outlier 问题),与本文发现的值幅度强相关。
本文建立了值幅度与量化误差的关联,是连接稀疏注意力与量化两条压缩路线的关键。
研究动机
推理模型(如 Qwen3)通过生成超长思维链提升准确率,却带来严重内存与计算瓶颈。Chen 等人发现模型对简单算术题“2 + 3 = ?”竟用超 900 个 token 回答;而竞赛题更夸张,AIME25 上 Qwen3-4B 平均生成 17790 个 token,prompt 却只有 198 token。自回归模型需为每个历史 token 保存键值表示,序列增长导致 KV 缓存线性膨胀,给显存带宽与容量巨大压力。稀疏注意力分两类:选择类(如 SeerAttention-R)保留完整缓存只激活稀疏子集,内存随序列线性增长;驱逐类(如 SnapKV、R-KV、CurDKV)在缓存达预算后永久丢弃低重要性 KV 对,能实现静态内存但精度下降明显。具体地,4× 压缩下 SnapKV 在 Qwen3-4B 平均准确率仅 49.15%,远低于完整模型 65.04%,存在明显的效率-精度权衡缺口,驱逐方法在推理任务上的表现始终弱于选择类。
本文的目标是本文目标是设计训练免(training-free)的 KV 缓存驱逐方法,在保持静态内存占用的前提下逼近甚至超越基于选择的方法 SeerAttention-R 的准确率,弥合驱逐与选择之间的精度差距。具体包括:在约 25% 有效注意力稀疏度(即 4× 压缩)下,让 Qwen3-4B 与 Qwen3-14B 在六项推理任务(AIME25、AIME26、HMMT25、GPQA-Diamond、MATH、LiveCodeBench-v6)上平均准确率超过 R-KV、CurDKV 等驱逐基线并匹敌 SeerAttention-R;同时在实际吞吐与峰值显存上优于现有驱逐方法,并支持 FlashAttention2 内核以利工程落地。作者还希望揭示值状态(value states)在长链推理中的关键作用,为 KV 缓存量化等压缩方向提供理论启示,最终把稀疏注意力与量化两条独立压缩路线打通。
与已有工作不同的是,已有驱逐方法几乎都只围绕“键(key)”做文章:SnapKV 用注意力分数对键打分但完全忽略值;R-KV 在 SnapKV 基础上加基于键余弦相似度的冗余惩罚以提升多样性,同样不看值;CurDKV 虽通过 CUR 矩阵分解的杠杆分数同时考虑键和值,却用固定随机投影矩阵且采用确定性 topk 选择,无法保证保留集合多样性。本文独特切入角度是:作者通过分析 Qwen3-4B 完整 KV 缓存,发现值状态存在极端倾斜的幅度分布——少数 token 拥有异常大的 Range 值,而这些大值状态对维持推理进程至关重要(驱逐它们使模型陷入重复循环,准确率暴跌至 14.3%)。这是首次系统地把“值感知”与“多样性/随机性”两条独立轴线结合:用值幅度为关键 token 预留专门槽位,用随机采样替代确定性 topk 以提升保留集合代表性。
核心方法
VASE 整体思路是在周期性驱逐框架下(持久预算 $K$ + 缓冲区 $B$,每 $B$ 步触发一次驱逐),同时利用两种互补信号决定保留哪些 KV 对——一是值状态幅度信息(保护大 Range 值不被丢弃),二是随机性(提升保留集合多样性与代表性覆盖)。直觉上,注意力输出 $o_t = \sum_i \alpha_i^{(t)} v_i$ 是对值向量的加权和,因此幅度大的值向量对输出有不成比例影响;同时确定性 topk 会让某些 token 每步都被永久淘汰,而随机采样能让每个 token 每步都有非零保留概率。基于此,作者提出两个变体:VASE-AttnV 建立在 SnapKV 之上,VASE-DKV 建立在 CurDKV 之上,分别用不同方式注入值感知与随机性。方法完全训练免,可直接用于已有 Qwen3 模型无需微调。
核心创新在于把 KV 缓存驱逐分解为三个正交维度——基于键的打分、基于值的打分、以及多样性,并首次将三者结合(见 Table 1)。VASE-AttnV 本质区别是不再用单一注意力分数做 topk,而是先按 $\mathrm{Range}(v_i)=\max_j v_j-\min_j v_j$ 选出 $N_v$ 个最大幅度值状态放入保留集 $R_V$,再从剩余候选以正比于 $\bar{\alpha}_i$ 的权重加权采样填充剩余 $K-|R|$ 槽位。这样每步保留概率 $\pi_i^{(t)}\geq 1-(1-\bar{\alpha}_i^{(t)}/Z_t)^{K-|R|}>0$,经 $T$ 步后总保留概率始终非零,避免 topk 下低于截断线 token 被确定性淘汰。VASE-DKV 保留 CurDKV 乘积分数 $\ell^{(K)}_i\cdot\ell^{(V)}_i$($\ell(G)=\|G^\top v_i\|_2^2$)天然感知值幅度,但关键区别是每个驱逐步 $t$ 重新采样高斯投影 $G_t$ 而非沿用固定 $G$,避免某些 token 被单一投影永久判低分。
方法步骤详情
VASE-AttnV 步骤:(1)维护持久预算 $K$ 与缓冲区 $B$,缓存满 $N=K+B$ 时触发驱逐;(2)对非缓冲区 $K$ 个候选按 $\mathrm{Range}(v_i)$ 排序,选最大幅度 $N_v$ 个(默认 $N_v=K/4$)加入保留集 $R_V$,连同缓冲区 $R_B$ 构成 $R$;(3)对剩余候选以权重 $\bar{\alpha}_i^{(t)}$ 加权随机采样填充剩余 $K-|R|$ 槽位;(4)丢弃未被选中的 $B$ 个对并清空缓冲区。VASE-DKV 步骤:每个驱逐步 $t$ 独立重采样高斯投影 $G_t\in\mathbb{R}^{d\times r}$($r=20$),计算键杠杆分数 $\ell^{(K)}_{i,t}=\|G_t^\top k_i\|_2^2$ 与值杠杆分数 $\ell^{(V)}_{i,t}=\|G_t^\top v_i\|_2^2$,用乘积分数排序保留最高 $K$ 个。两变体都保护缓冲区内最近 $B$ 个 token。所有超参数均在 GSM8K 上调好后直接迁移。
技术新颖性
技术新颖性有三点。首先,作者首次系统刻画值状态幅度的极端倾斜分布,并通过对照实验(故意驱逐最大 Range 值,使 GSM8K 准确率从 88.0% 暴跌到 14.3%)证明这些大值状态对推理进程不可或缺——与以往强调注意力汇聚点(对应低幅度值)的工作互补,首次指出“大幅度值同样重要且应被上权重”。其次,VASE-AttnV 提出用值幅度预留专门槽位而非乘以注意力分数的新机制,确保大值状态总是被保留,这与 Guo 等人将注意力分数乘以值范数的方法本质不同。第三,作者揭示随机性作为独立轴线的价值:把 topk 换加权采样、把固定投影换逐步重采样,分别给 SnapKV 带 6.6%、CurDKV 带 9.0% 提升。此外作者建立了值幅度与逐 token KV 量化误差的强相关性(多数层 Pearson > 0.8),把稀疏注意力与量化两压缩方向打通,启示了混合精度方案。
实验结果
核心发现分四部分。(1)消融(GSM8K,$K=512$):值打分($N_v$ 16→256)使准确率从 73.3% 升到 80.5%(+16.2%);采样给 SnapKV、CurDKV 分别带 +6.6%、+9.0%;两者结合接近完整模型 88.4%,证明互补。(2)主结果(Table 2):Qwen3-4B 上 VASE-AttnV 达 59.09% 超 SeerAttention-R(58.81%)与 R-KV(54.69%,+4.4%),VASE-DKV 达 57.48% 比 CurDKV(49.78%)高 7.7%;Qwen3-14B 上 VASE-AttnV(65.81%)匹敌 SeerAttention-R(65.37%)超 R-KV 4.9%。(3)代码生成(Figure 5):VASE-AttnV 63.5% 与 R-KV 62.6% 持平,SeerAttention-R 仅 45.3%,训练免方法域偏移更鲁棒。(4)系统:VASE-DKV 吞吐 411 tokens/s 是 Full 的 3.1×,显存 VASE-DKV<VASE-AttnV<R-KV<Full。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 六任务平均 (Qwen3-4B) | pass@1 准确率 (4×压缩) | 59.09% (VASE-AttnV) | R-KV 54.69% | +4.4 个百分点 |
| 六任务平均 (Qwen3-14B) | pass@1 准确率 (4×压缩) | 65.81% (VASE-AttnV) | R-KV 60.90% | +4.9 个百分点,匹敌选择类 SeerAttention-R 65.37% |
| GSM8K (Qwen3-4B) | pass@1 准确率 (K=512) | 87.6% (VASE-DKV) | CurDKV 78.6% | +9.0 个百分点(随机性贡献) |
| LiveCodeBench-v6-Medium (Qwen3-4B) | pass@1 准确率 (2048预算) | 63.5% (VASE-AttnV) | SeerAttention-R 45.3% | +18.2 个百分点(域偏移鲁棒性) |
| 解码吞吐 (Qwen3-14B, A100-80G) | tokens/s (16K输出, 2048预算) | 411 (VASE-DKV) | Full 133 | 3.1× 加速 |
| 值打分消融 (GSM8K, Qwen3-4B) | pass@1 准确率 | 80.5% (Nv=256) | SnapKV 64.3% | +16.2 个百分点(仅值感知) |
局限与改进
作者在附录 E 坦承若干局限。首先评估仅聚焦推理模型与解码阶段,未覆盖预填充阶段的长上下文压缩基准,尽管作者认为核心方法可扩展到预填充。其次仅在 Qwen3 模型上实现,因选择基线 SeerAttention-R 只发布 Qwen 检查点,但作者强调值打分仅依赖值缓存的 min/max 统计本质上是模型无关的。第三虽建立了值幅度与量化的关联,但量化方法本身的开发超出范围,未提出整合方案。从我观察看还有几点:所有超参数仅在 GSM8K(小学数学)上调优后迁移到竞赛级任务,GSM8K 分布可能不代表所有任务泛化性需更广验证;VASE-AttnV 的值预留预算 $N_v=K/4$ 固定未随任务/层自适应;Table 5 显示小测试集(AIME 仅 30 题)标准误较大置信区间重叠单点比较统计显著性偏弱;随机采样引入的额外方差可能影响需要确定性输出场景的可复现性。
独立分析的弱点
独立分析主要弱点如下。(1)超参数敏感性:$N_v=K/4$、$\lambda=0.5$、$r=20$ 均在 GSM8K 上选定后固定迁移,缺对各任务自适应机制。改进方向是引入层/任务自适应的值预留比例或学习轻量门控。(2)通用性验证不足:仅 Qwen3 系列未在 Llama、Mistral 等架构验证值幅度倾斜是否普适,也缺长上下文(如 RULER、128k)场景。改进方向是跨模型族扩展并报告长上下文基准。(3)随机采样方差控制:加权采样引入额外生成随机性,虽报告标准误但需确定性输出(代码、形式化验证)场景可能不利。改进方向是提供可控随机种子或温度调节并量化方差对端到端可靠性影响。(4)仅解码阶段:预填充阶段同样存在内存瓶颈未探索预填充压缩。改进方向是把值感知+随机性推广到预填充。(5)与量化割裂:虽发现 Range 与量化误差强相关但未实现混合精度整合。改进方向是设计 outlier 感知混合精度 KV 量化,把大 Range 值放高精度保留缓存其余低比特量化。
未来方向
作者明确提出的方向:(1)混合精度方案——把大 Range 值状态放入高精度保留缓存,其余 KV 缓存量化为低比特,一旦高精度预算用尽就把最早值状态量化后“驱逐”到低精度缓存,实现压缩比极限突破;(2)深入分析大值状态携带什么信息、如何影响长链推理,因移除它们会让模型陷入重复循环,作者猜测这些值帮助在潜在推理步间过渡防止模型塌缩到自我强化路径。基于成果可延伸的方向:(3)把 VASE 的值感知+随机性两轴推广到预填充阶段用于长 prompt 压缩;(4)探索值幅度作为通用信号来指导 KV 缓存跨层共享、跨请求合并等更激进优化;(5)结合推测解码与 KV 驱逐研究草稿/验证两阶段如何联合管理缓存;(6)把随机性思路扩展到选择类方法(如 SeerAttention-R)检验能否进一步改善其在域偏移下鲁棒性。
复现评估
复现评估总体较好。代码已在 https://github.com/terarachang/VaSE 开源(基线多为 CC-BY-4.0)。模型(Qwen3-4B/14B,Apache 2.0)和数据集(GSM8K、MATH、GPQA-Diamond、LiveCodeBench-v6、AIME、HMMT)均有公开来源。作者复用 SeerAttention-R 官方评测脚本,超参数($K\in\{1024,2048,4096\}$、$B=64$、$N_v=K/4$、$r=20$、$\lambda=0.5$)与压缩比设定(按各任务平均 token 数 $N_{avg}/4$)都有明确交代。算力需求较高:实验在 NVIDIA A100/H100 上,单卡 80GB 显存即可跑通,但每任务耗时 10–72 GPU 小时不等,全量复现六任务、两模型、多次随机种子(pass@1 平均 8–16 次)规模可观。难度中等:方法本身训练免实现清晰,但需熟悉 vLLM/FlashAttention2 内核与周期性驱逐框架工程实现,随机采样正确注入点(每层每头独立)需仔细处理,对一般研究者工程门槛偏高。
论文图表
左图为不同层值状态 $\mathrm{Range}(v)$ 的小提琴分布,可见存在极端幅度离群点;右图为对比实验:驱逐最大幅度值会使 GSM8K 准确率从 88.0% 暴跌到 14.3%,远低于随机驱逐基线(53.2%)。
这是论文核心动机的实证依据,直接证明大幅度值状态对精度至关重要,引出值感知设计的必要性。
两个 GSM8K 示例:驱逐大 Range 值后模型从合理推理陷入重复循环(例1 不断重复“Wait, no...”,例2 反复自我反思“Alternatively, maybe...”),无法收敛到最终答案。
直观展示大值状态被移除后的灾难性失败模式,是理解值感知动机的最有力证据。