注意力质量只是代理,可能排错淘汰项
KV 缓存保存自回归解码中每个已处理 token 的键值状态,因此其内存和注意力开销随序列长度线性增长。免训练压缩只保留缓存条目的一部分,淘汰规则由推理时统计计算,而不是从任务轨迹中学习,也不通过额外的缓存策略训练来优化。
现有免训练方法用间接信号决定保留什么。基于注意力的规则保留在当前上下文中获得高注意力的条目,或预期在未来解码中获得高注意力的条目;另一些方法偏好有助于从压缩缓存重建原始上下文的条目。论文认为这些代理目标未必与淘汰的预测效果一致,因为注意力可能高估冗余条目,而重建效用可能偏向恢复上下文,而非保持下一词分布。
下一词 KL 是行为损失中可观测的部分
论文以淘汰前的同一模型为参照定义行为保持。对于长度为 tau 的未来续写,完整缓存与压缩缓存条件分布之间的差异,按自回归 KL 的链式法则可分解为当前前缀处的下一词 KL,加上对未来 token 取期望的剩余散度。
压缩决策作出时只有第一项可观测,因为后续项依赖尚未出现的未来 token 和查询状态。因此该方法把可用查询位置上的下一词分布作为压缩时的行为目标,并按淘汰项对输出分布的估计扰动来评分。
图 1 支持这一选择。作者使用 LLaMA3.1-8B,按每种方法自身评分的百分位对采样淘汰掩码分组,并测量同一参考续写上后续位置的平均 KL。保持行为的评分与未来位置 KL 的关联最清晰(LongBench 上 Spearman rs = 0.90,MATH500 上 0.86,AIME25 上 0.89),而基于代理的评分关联较弱(例如 SnapKV 分别为 0.22、0.45 和 0.32)。
用淘汰前统计估计 logit 偏移
理想目标是在保留掩码上最小化完整缓存预测到压缩缓存预测的 KL,但为大量候选掩码获得精确 logit 扰动需要针对每个候选运行掩码缓存前向传播。该方法改为从完整缓存正向统计估计扰动。
对每一层、每个头和每个评分位置,它计算一个删除反事实:固定淘汰前的查询、键和值,移除被掩码的条目,并在保留条目上重新归一化注意力。把得到的头输出扰动在各头上求和,就得到该层的残差流扰动。
最终状态扰动由直接残差路径贡献近似,即对各层扰动求和,并把后续注意力、MLP 和归一化层的响应视为近似误差。通过最终归一化和 LM 头的一阶展开,将其转换为估计的 logit 偏移,再代入 KL 目标。对基于 RMSNorm 的模型,最终归一化的雅可比向量积有闭式解;对没有最终归一化的模型,它就是恒等映射。
Top-K KL 让评分足够便宜可用
全词表 KL 评估开销很大,因此实现的评分在完整缓存分布的 top-K 支撑集上评估 KL,使用重新归一化的 top-K 概率和数值稳定的 log-sum-exp。当 K 为全词表时,该评分退化为估计 logit 的 KL。
直接搜索选出 B 个淘汰条目不可行,因为淘汰代价不可加:移除多个条目会同时改变重新归一化的注意力输出和由此得到的预测分布。该方法采用两阶段贪心过程。它先由单条淘汰评分最低的合格条目构建候选池,再通过反复加入使掩码级估计 logit KL 增幅最小的候选来细化选择,并在每一步后更新缓存的掩码统计。
细化过程中,只有受影响层的移除质量和移除贡献统计发生变化,其他层的统计可复用。一条稳定性规则会拒绝任何使保留注意力质量低于阈值 epsilon 的候选,避免不稳定的注意力重新归一化。
预算越紧,相对注意力启发式的增益越大
该方法在两种推理设置下评估。在预填充时压缩中,提示缓存于长上下文编码后压缩一次,评分位置取自最后 m_pre 个提示位置。在生成时压缩中,提示条目始终受保护,每生成 I_comp 个 token 触发一次压缩,只有最近保护窗口之外的生成 token 条目合格。
在不同架构和两种设置下,该方法在匹配的保留 KV 预算下,下游任务质量相对轻量注意力启发式有显著提升,压缩越激进增益越大。预算按最终保留的层-KV 头-token 预算匹配,受保护条目计为保留条目,因此没有方法获得未计入的受保护缓存。
保持行为的评分会带来额外的压缩时计算,但在所评估的设置中,压缩后的推理仍比完整缓存推理更快。默认超参数在每个机制内固定,不按基准调优。
行为保证只覆盖压缩时的查询
由于淘汰决策必须在读取缓存的未来解码查询已知之前作出,该方法只对可用的压缩时查询保持行为。它不保证任意未见未来续写的行为,后续查询可能放大或削弱淘汰的影响。
残差路径近似是第二条边界:它保留可由淘汰前统计计算的部分,并把后续注意力、MLP 和归一化层的响应视为近似误差。论文在附录中验证这一排序信号,而不是声称与掩码缓存前向传播精确等价。
预算紧张时优先用保持行为的评分
这篇论文中与决策相关的分界是预算压力。在匹配的保留 KV 预算下,保持行为的评分优于轻量注意力启发式,且压缩越激进差距越大,因此缓存预算越紧,越有理由付出额外的压缩时计算。
第二个分界是压缩机制。预填充时压缩在编码后对提示末尾位置评分一次,而生成时压缩保护提示条目并按间隔对近期生成位置评分;两种机制使用各自的实现超参数,因为其合格候选和调度不同。读者比较方法时应匹配保留的层-KV 头-token 预算并计入受保护条目,因为这些选择决定比较是否公平。
本文引用的文献
Behavior-Preserving KV Cache Compression
一个免训练框架按压缩缓存与完整缓存下一词分布之间的估计 KL 给候选 KV 淘汰项评分,该评分由淘汰前正向统计计算,并在匹配的保留 KV 预算下优于轻量注意力启发式,压缩越激进增益越大。
