论文阅读-CriticalKV
论文
[2502.03805] CriticalKV: Optimizing KV Cache Eviction from an Output Perturbation Perspective
背景
问题:长序列推理中庞大的 KV(键值)缓存面临着高存储和运行时间成本的巨大挑战
传统做法:基于注意力权重剪枝较不重要的条目来减小 KV 缓存大小
局限性:在删减缓存时,普遍默认“注意力权重(Attention Weights)大 = 条目重要” 。
核心工作:提出了一个通过分析注意力输出扰动(Output Perturbation)来识别关键 KV 缓存条目,除了注意力权重之外,KV 条目中的 Value(值)状态以及预训练参数矩阵也是至关重要的 。
技术方案:提出了一种扰动约束选择算法,该算法通过优化最坏情况下的输出扰动来识别关键条目 。我们证明了我们的算法是一种通用的、即插即用的增强模块,且仅带来微不足道的计算开销 。
实验定量结果:覆盖了 3 种基线方法(SnapKV、AdaKV、HeadKV) 、3 种不同架构和规模的模型(Llama-3.1-8B、Mistral-7B、Qwen-2.5-32B),在相同的缓存留存率下(例如 40% 的缓存大小),由剪枝带来的模型精度降级(Compression Loss)平均减少了 50% 以上
核心发现
决定剪枝引发的输出误差(Output Error)的,远不止注意力权重本身。至关重要的是,包含在 KV 条目内部的 Value(值)状态的模长,以及模型固有的预训练输出投影矩阵(Output Projection Matrix,即 WO),同样对最终的输出扰动起着决定性的约束作用。
- 我不去猜哪个条目重要,我直接去计算一个纯粹的数学问题——“如果我把条目 i 删掉,注意力层的最终输出结果会偏离多少(即扰动值 ΔO 有多大)?”
- 通过严格的数学推导,作者抓住了前人的盲区:自注意力机制的最终输出是 O = Softmax(QKT)V ⋅ WO。前人只盯着前面的 Softmax(QKT)(注意力权重)看,却完全忽略了后面乘上的 V(Value 状态)和 WO 矩阵。如果一个 Token 虽然注意力权重很高,但它的 V 向量模长极小,或者它经过 WO 投影后几乎被坍缩掉,那删掉它其实根本不影响大局!
基于上述理论发现,我们提出了一种“扰动约束选择算法(Perturbation-Constrained Selection Algorithm)”。该算法的核心在于,它在数学上显式地去优化和控制“最坏情况下的输出扰动上界”。我们进一步证明了,该算法是一个通用的(Universal)、即插即用的(Plug-and-play)增强模块,能够无缝嵌入现有的各种淘汰策略中,且仅引入微不足道的线性计算开销。在 Ruler 和 LongBench 等长文本基准上的广泛实验表明,我们的方法能够将压缩带来的精度损失平均降低一半以上。
我的理解
也就是说这篇工作,
1.还是在kvcache淘汰方法的基础上,作为一个即插即用(Plug-and-play)的增强模块,嵌套在它们的底层代码里
2.主要聚焦在如何判断哪些kv条目是有用的,即挑选权重时,不能光看attention权重
3.使用扰动算法挑选出:谁被扔掉后引起的输出扰动上界最小,就优先淘汰谁
大语言模型缓存淘汰(Cache Eviction)
核心思想是在推理过程中,动态地剪枝(Prune)掉较不重要的历史 KV 条目。然而,现有的工作绝大多数依赖自注意力权重(Attention Weights)作为主要的启发式指标(Heuristic Metric)来评估条目的重要性。这些方法普遍遵循一个经验性假设,即注意力权重较小的条目对最终输出的贡献微乎其微。遗憾的是,这种纯经验驱动的方法缺乏严格的数学形式化证明与理论根基。
代表论文[2306.14048] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
核心发现
惊人的注意力稀疏性(Sparsity):即使 LLM 是在全连接(Dense)状态下训练出来的,但在推理阶段,注意力矩阵有 95% 以上 的元素其实都是高度稀疏、接近于零的 。这意味着,只有大约 5% 的 KV Cache 对当前 Token 的预测真正起作用 。
重击者 Token(Heavy Hitters, 简称 H2):模型在解码时,注意力得分(Attention Scores)的累积呈现出明显的幂律分布 。也就是说,有极少数的核心 Token(如句首 Token、高频共现词等)在整个生成过程中持续贡献着绝大部分的注意力权重 。
功能坍塌风险:实验表明,如果强行把这些 H2 Token 从缓存中移除,模型的准确率会断崖式下跌,发生严重的“功能崩塌” 。
H2O 的动态缓存淘汰机制
基于上述发现,理想的策略应当是将宝贵的显存留给那些最关键的 Token 。H2O 提出了一种均分显存预算的动态淘汰策略,将 KV Cache 预算(例如只保留总长度的 20%)平分给两类 Token :
最近的 Token(Local Tokens):由于自然语言的局部相关性,刚生成不久的临近词通常拥有很强的关联。
重击者 Token(H2 Tokens):全文本中累积注意力得分最高、影响力最大的核心 Token 。
动态维护流程
在每个解码(Decoding)步骤中:
- 新生成的 Token 产生的 KV 键值对会被正常放入缓存中 。
- 当缓存中的 Token 数量超过了预设的显存空间预算 k 时,淘汰机制就会触发 。
- 算法会利用局部统计数据(Local
Statistics)——即只累加当前步骤及之前所有步骤产生的注意力得分
,算出一个累积总分 。
- 算法进行贪心选择:在原有的 H2 集合与新加入的 Token 中,找出累积注意力得分最低的那个 Token,将其从缓存中无情剔除,确保缓存容量始终恒定 。
其他对比方法
1️⃣ SnapKV (Li et al., 2024)
- 核心原理:SnapKV 发现,在 Prefill(提示词前缀编码)阶段结束时,各个注意力头(Attention Head)重点关注哪些 Token 的格局已经基本固化了。因此,它在 Prefill 结束时,对最后的观察窗口(Observation Window)内的注意力权重算一次均值/最大池化(Pooling),选出得分最高的特征簇(Clustered Features),然后一锤定音锁死需要保留的历史 KV 位置,后续 Decode 阶段不再动态更改。
- 本文如何增强它:原本 SnapKV 在选出高分簇后直接按注意力权重大小做 Top-K 留下 KV;本文将其改为“既看权重,又看 Value 模长与矩阵投影带来的扰动上界”的动态筛选。
2️⃣ AdaKV (Feng et al., 2024) / HeadKV (Fu et al., 2024)
- 核心原理:这两篇工作聚焦于注意力头之间的动态预算分配(Head-wise Dynamic Budget Allocation)。传统的淘汰方法(如 H2O)对每个注意力头都是一视同仁的(例如每个头都强制只留 128 个词的缓存)。但 AdaKV 和 HeadKV 发现,有的头属于“钝感头”(注意力很分散),有的头属于“敏感头”(注意力极度聚焦)。它们通过计算每个头的注意力熵(Entropy)或实时统计量,动态地把宝贵的显存预算挪给更需要的头。
- 本文如何增强它:本文的算法与这类分配策略是完全正交(Orthogonal)的。AdaKV 负责在全球视角下决定“给这个注意力头分配多少个缓存位置(Budget b)”,而本文的算法则在拿到这个预算 b 后,在注意力头内部以最低扰动代价精准挑出这 b 个条目。
Critical KV Cache Entry Selection(关键 KV 缓存条目选择)
KV 缓存淘汰问题
给定严格的硬件资源限制,KV 缓存的最大容量受到用户定义预算 b(其中 b ≪ t)的约束。KV 缓存淘汰算法的目标是识别一个关键索引子集 St ⊂ {1, 2, …, t},且满足 |St| = b。被淘汰的补集表示为 Stc。剪枝后的缓存矩阵 K̂t, V̂t ∈ ℝb × d 仅由 St 索引的条目组成。因此,近似的注意力权重和受到扰动的层输出公式化为:
$$\hat{A}_t = \text{Softmax}\left(\frac{q_t \hat{K}_t^T}{\sqrt{d}}\right), \quad \hat{o}_t = \hat{A}_t \hat{V}_t W_O$$
最坏情况输出扰动
假设我们的显存预算(Budget)只允许我们保留一个大小为 b 的子集 S(即 |S| = b),而剩下的条目集合 Sc(大小为 N − b)必须被无情淘汰(Evict)。
如果直接丢弃 Sc 中的条目,那么剪枝后的近似输出 Ô 就变成了:
Ô = ∑i ∈ SAiViWO
这个公式的含义是利用残缺的历史 Key 矩阵计算出重构后的注意力权重,乘上对应的 Value 向量,最后通过当前头的输出投影矩阵
这样计算累加后的结果,就是多头注意力要吐出的一维向量
论文的核心目标是:衡量剪枝前后的输出绝对误差,即输出扰动(Output Perturbation) ΔO。作者在这里严谨地引入了 L1 范数(Norm)来定量刻画:
$$\Delta O = \| O - \hat{O} \|_1 = \left\| \sum_{i=1}^N A_i V_i W_O - \sum_{i \in S} A_i V_i W_O \right\|_1 = \left\| \sum_{i \in S^c} A_i V_i W_O \right\|_1$$
这是一个多元向量求和的范数。由于我们在推理阶段根本无法预测下一个未知的 Query 会让这些向量如何交织,因此为了提供最稳固的理论保证,必须去推导它的最坏情况上界(Worst-case Upper Bound)。
利用矩阵与向量分析中最经典的三角不等式(Triangle Inequality),作者完成了惊艳的去耦合(Decoupling)推导:
ΔO = ∥∑i ∈ ScAiViWO∥1 ≤ ∑i ∈ Sc∥AiViWO∥1
因为自注意力权重 Ai 经过 Softmax 标量化后天然满足 Ai ≥ 0,所以标量 Ai 可以直接从范数内部提出来。于是公式完美收敛为:
ΔO ≤ ∑i ∈ ScAi∥ViWO∥1
关键度得分(The Criticality Score)
上面的不等式右侧,就是作者千方百计寻找的“误差天花板”(Perturbation Upper Bound)。
为了让大模型在剪枝后不“变傻”,我们的终极优化目标是让真正的扰动 ΔO 尽可能小。既然直接控制 ΔO 极难,那最严谨、最鲁棒的工程解法就是直接把这个上界死死压住。
为了让留下的子集 S 使得未留下的子集 Sc 的误差之和最小:
minS∑i ∈ ScAi∥ViWO∥1 ⇔ maxS∑i ∈ SAi∥ViWO∥1
这级联导出了这篇论文最核心的算法判定准则。每一个 KV 缓存条目 i 的关键度得分(Criticality Score) ℳi 被严谨地定义为:
ℳi = Ai⋅∥ViWO∥1
与传统方法对比
我们可以通过下面这个简单的表格,直观对比传统方法与本文方法在筛选机制上的根本演进:
| 维度 | 传统方法(H2O / SnapKV) | 本文方法(CriticalKV) |
|---|---|---|
| 评估指标 | ℳilegacy = Ai | ℳiours = Ai⋅∥ViWO∥ |
| 数学假设 | 隐式假设 ∥ViWO∥ 为常数常数噪声 | 承认 ∥ViWO∥ 具有高方差和支配级影响力 |
| 物理意义 | 仅关注 Query 喜不喜欢它(选择概率) | 同时关注选择概率与它本身携带并投射的信息量级 |
实验
Setup
基座模型:Llama-3.1、Mistral 和 Qwen-2.5
评估基准:采用了来自两个严格基准的 29 个数据集:Ruler 和 LongBench
Ruler 是目前公认对长文本剪枝极其严苛的基准(比传统的 LongBench 更难欺骗),它包含了大量长距离依赖和噪声干扰任务。
集成基线:本文的算法是作为增强插件(Plugin),直接替换掉了 SnapKV、AdaKV 和 HeadKV 原本纯看 Attention 权重的 Top-K 筛选核心代码,并保持它们其余的窗格和跨头预算分配逻辑完全不动
Ruler Benchmark
RULER 基准通过提升标准检索任务的难度来评估大语言模型的有效上下文长度。它包含四大核心任务类别:单值检索、多值检索、变量追踪和通用聚合。
固定物理缓存预算 b,然后让输入的总文本长度 N 从 4K 一路疯狂拉长到 128K。
在 Llama-3.1-8B 的 SnapKV 框架下,当序列长度超过 32K 时,原始的纯注意力权重策略经历了严重的性能下滑。相比之下,通过集成我们的准则 ℳi,增强后的模型在高达 64K 甚至 128K 的长度下依然维持了高准确率,将淘汰缓存后的有效上下文长度有效地扩展了 2 到 4 倍。
CWE (Common Words Extraction) / FWE (Frequent Words Extraction):常见词/高频词提取。测试模型能否在海量长文本中准确统计和提取出现频率最高的词汇,考察全局信息聚合能力。
NIAH (Needle In A Haystack ):
- 这是目前测试长文本检索最核心的指标,要求模型从几万字的文档中找出隐藏的一句话。
- Multikey 1/2/3:多键检索(文档里藏了多组不同的钥匙和宝箱,看模型能不能全部对应找出来)。
- Multiquery:多查询检索(一次性问模型多个捞针问题)。
- Multivalue:多值检索(一个键对应了多个不同的值,看模型能不能找全)。
- Single 1/2/3:单针检索(最基础的捞针,放在文章的不同位置)。
QA1 / QA2:
- 长文本问答(Quality Assurance)。通常基于 HotpotQA 等数据集改装,需要模型结合长文本中的多处碎片信息进行深度推理才能回答。
VT (Variable Tracking):
- 变量跟踪。模拟代码运行,在极长的上下文中追踪多个变量(如
a=1,b=a+2,c=b...)的状态改变,测试模型的精确长程依赖能力。
Cache Size(缓存大小):代表了 KV 缓存的保留比例。
具体子任务: 包括
cwe(单词提取)、fwe、niah(大海捞针测试的各种变体,如单针 single、多键 multikey、多查询 multiquery、多值 multivalue)、qa(长文本问答)以及vt(变量追踪) 。
结论一:随着横轴的 Cache Size 从 100% 缩减到 40% 甚至 20%,所有的虚线(SnapKV, AdaKV, HeadKV 基础版本)都呈现出严重的急速下滑趋势 。这说明在极高压缩率下,仅仅依靠“注意力权重(Attention Weights)”来决定删除哪些缓存,会错误地删掉很多关键信息,导致模型在长文本任务中“失忆” 。
LongBench Evaluation
知识补全
KV 缓存条目(entry)
当模型处理或生成一个词(Token)时,它会为这个词计算出对应的 Key 向量(键状态)和 Value 向量(值状态) 。这两个向量作为一个整体,被存入显存中,这就构成了 KV Cache 中的一个 Entry(条目/记录) 。
WO如何理解
多头注意力机制(Multi-Head Attention)中有一个全局的输出投影矩阵,写作 WOglobal,用于让所有头不同维度上的特征进行加权组合。而这篇论文里写在每个头公式里的 WO,正是全局权重矩阵中属于当前头的那一个“切片”(Block/Slice)。
为了让多头的信息彻底混合,并且映射回模型的残差连接空间,Transformer 在多头注意力层的最后,放置了一个全连接线性层。这个线性层的权重矩阵,就是 WO(维度是 dmodel × dmodel)。
什么时候进行kvcache的淘汰机制
触发条件:当前缓存里的 Token 数量已经达到了规定的硬件上限(例如 2048 个),此时第 2049 个新生成的 Token 带着它刚算出来的 K, V 向量准备挤进显存。
在给kvcache显存满的时候,开始进行token的末位淘汰,这个论文就是计算每个每个关键度得分 ℳi = Ai⋅∥ViWO∥1,找出得分最低的那 1 个 Token,直接从显存里扔掉