ResKV 不再把 KV 缓存压缩等同于“删词”:给被淘汰的信息留一份残差账本

长上下文推理有个绕不开的成本:模型每生成一个新 token,都要回看历史 token 的 Key 和 Value。上下文越长,KV cache 占用的显存和内存带宽就越高。多数压缩方法因此先给历史 token 打分,只保留一小部分,其余直接淘汰;另一些方法把被淘汰的状态合并进保留项。前者省得干脆,却把被删除 token 对注意力的贡献一起抹掉;后者保住部分信息,却可能改动本应精确保留的 Key 和 Value。

7 月 31 日提交到 arXiv 的论文 ResKV,换了一个角度:被淘汰的信息不必逐个保存,但它们对 softmax 注意力分子和分母的总体贡献,可以被压缩成“残差统计量”。于是,在固定的 KV 槽位预算内,ResKV 把缓存拆成两部分:

  • 主缓存保存优先级高的 token,Key 和 Value 保持原样;
  • 残差缓存用少量条目概括被淘汰 token 的聚合贡献。

关键不是在注意力算完后补一个修正项。主缓存和残差条目进入同一次 softmax 归一化,残差由此同时补回分子中的 Value 加权贡献和分母中的概率质量。它既没有把遗漏信息塞进保留 token,也没有增加总 KV 槽位数。论文原文

两道控制,避免残差喧宾夺主

残差信息并非每层、每个 KV head 都同样有用。ResKV 在缓存构建阶段,用一组拟合查询和验证查询尝试多个残差预算比例;只有当残差缓存能降低相对于完整缓存的注意力输出重建误差时,才给该层和 KV head 分配残差槽位,否则全部预算仍交给主缓存。

解码阶段还有一道动态门控。论文用主缓存中最大的注意力权重衡量分布是否尖锐:如果某个查询已经高度集中在主缓存的少数条目上,就压低残差权重,避免模糊明确的检索峰值;如果注意力较分散,就允许残差承担更多概率质量,补回散落在上下文中的证据。

固定预算下,提升集中在紧预算与未知查询

实验使用 LLaMA-3.1-8B-Instruct 和 Qwen-2.5-7B-Instruct,在 LongBench 与 RULER 上测试,并覆盖保留 10%、20%、30%、40% KV 的预算。ResKV 被接到 AdaKV 和 SnapKV 上,与原基线使用相同的总槽位数。

论文报告,LongBench 展示的 32 组配置全部提升,两个骨干模型的平均增益为 1.02 分;在只保留 10% 和 20% KV 时,平均增益分别为 1.43 和 1.17 分。RULER 展示的 64 组配置中有 63 组提升,平均增益为 3.38 分。更值得注意的是,不知道未来查询内容就先压缩缓存的 query-agnostic 设置,平均提升 4.54 分,高于 query-aware 设置的 2.22 分。这更接近真实服务:系统通常不能等下一条用户问题出现后,再重新决定旧上下文该留什么。

组件消融也给出了清晰边界。在 10% KV、query-agnostic 设置下,去掉验证代理,会让两个代表性 RULER 任务分别下降 4.00 和 7.04 分;去掉动态门控,分别下降 4.80 和 4.48 分;把共享 softmax 改成主缓存与残差各自归一化,RepoBench-P 下降 3.00 分。

这不是免费的加速,但方向更像“信息压缩”

效率图显示,ResKV 的峰值显存与 SnapKV 基本重合,额外内存开销可忽略,并能在 128K 上下文维持稳定解码;代价是每一步都要计算残差分支和门控,因此吞吐低于 SnapKV。论文没有把它包装成零成本胜利:它用适度计算开销,换取固定显存预算下更完整的注意力信息。

ResKV 真正重要的判断是:**KV cache 压缩不该只有“留或删”这一个旋钮。**当长上下文里的证据分散时,单个 token 看似不重要,它们合在一起却可能决定答案。下一步值得关注的,不只是更聪明地挑 token,而是如何把被删掉的信息压成仍能参与计算的结构。