A Probabilistic Interpretation of KV Cache Eviction
Renato Geh, Alex Chen, Daniel Israel, Aditya Grover, Guy Van den Broeck
cs.CL, cs.AI
2026-08-28
UCLA证明精确KV驱逐是NP完全问题，把top-k写成零方差无界偏估计，采样后用SNIS在解码时纠偏，跨任务更稳。
KV 缓存驱逐的卖点很直白：把一部分 key/value 从缓存里丢掉，换吞吐和显存，质量最好别掉。H2O、SnapKV、TOVA、StreamingLLM 这类方法已经能在经验上做到这一点，套路几乎一样：给每个 token 打分，留 top-k。
缺的是定义。真实目标是驱逐之后、后续解码步的注意力还对；文献里通常只在驱逐当下算完分数就切掉，解码时当那些条目从来没存在过。UCLA 这篇把问题写成判定题，并指出现有启发式其实是一类偏差不受控的估计器。
判定版 KVEviction：给定 K、Q、V、误差 ε 和保留比例 r，问有没有大小为 ⌊n·r⌋ 的子集，让子集上的 softmax 注意力跟全集差不超过 ε。从 Partition 归约可证这是 NP-complete。更麻烦的是，这还只是驱逐当下的近似；真实部署要保的是未来步，那时被丢掉的条目已经算不回来。
注意力本身是期望。softmax(qKᵀ) 在 n 个位置上给出分类分布 p，输出是 V 在 p 下的期望。驱逐等于改这个分布的支撑，再在残存条目上重新归一化。少掉的概率质量会把留下的条目系统性放大或缩小，误差还会顺着层和后续 token 往下传。
现有 top-k 方法对应零方差估计：同一组分数永远留下同一批条目。方差为零，偏差可以任意大。概率驱逐把那些分数看成未归一化的提案分布 π，有放回采样 m 次，没被抽到的条目才丢掉。解码时用自归一化重要性采样（self-normalized importance sampling，用 π 和采样计数去校正目标期望）把驱逐前、后两段期望拆开再加权：驱逐之后新生成的 token 可以精确算，驱逐之前的那一段用 SNIS 估，两边的配重来自对归一化常数的同一套估计。偏差和均方误差都有 O(1/m) 上界。
重要性权重还可以加温度 τ。τ=1 是标准校正；τ→∞ 等于关掉校正、方差下降、偏差上升。分组查询注意力里，一组头共享同一份 K/V，提案改成组内混合分布，保证每个头的支撑都被盖住。
评测模型是 Llama3.2-3B 和 Qwen3-4B，基线来自 KVPress 里的 StreamingLLM、SnapKV、TOVA、H2O 和 K-norm。LongBench 只用了上下文短于 3000 token 的 HotpotQA、QASPER、TriviaQA；RULER 抽了 130 条。压缩比在实验里按驱逐比例扫，r 接近 1 表示 prompt 几乎被砍光。
平均分上，带校正的概率驱逐在中低压缩段能打平或略过现有方法。按任务两两比的 win score，带谐波先验的最小方差提案 πmin-h 在两个模型和多组任务上排第一。确定性启发式会在个别任务上崩：Llama3 上 StreamingLLM 在问答分割第一、在 MultiKey-NIAH 垫底；Qwen3 上 K-norm 在 MultiKey-NIAH 常进前三、在 CommonWords 垫底。概率方法把偏差压下来之后，这种任务错位会少一些。
压缩预算是全局软约束，每个头按自己的分布自适应花钱。实测是底层少砍、高层多砍，跟 PyramidKV「下层多留预算」的观察一致。m 很大时再加温度会把偏差抬上去、误差变差；只有 m 极小、方差炸开时，升温才有用。论文没有给出可直接引用的分任务 F1 表，也没有端到端吞吐数字。
给 KV 驱逐补了一套统计语言。H2O 这类打分不必丢掉，可以改写成提案，再在解码时校正。生产里如果只砍 top-k 且不做校正，等于接受无界偏差，换来的是零方差的确定性。
对已经在用驱逐的人，这篇更接近一次改视角，不是立刻可插的新算子。校正要记采样计数，变长头还需要稀疏实现。
作者自己写明：SNIS 既有偏差也有方差，理想的零偏差零方差做不到，因为那等于用更少的内存完美复原被丢掉的缓存。额外内存是 O(k·h·b)，相对 KV 本身的 O(d·k·h·b) 渐近可忽略，但仍要改注意力核。变长压缩还依赖高效稀疏张量，文中指向 Ada-KV 一类实现，自己没有交出生产级 kernel。
评测面偏窄：3B/4B、LongBench 截断到 3000 token、RULER 子集。硬度证明保的是驱逐当下的注意力，不是解码未来。文中也没有报告校正后的 tokens/s 或显存墙钟。