💬 NLP

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

KVzap 是一种快速、具有输入自适应性的 KV 缓存剪枝方法，它在各种大语言模型和任务中均实现了具有最先进水平的压缩率且几乎没有精度损失，解决了由上下文长度增长引起的关键推理瓶颈问题。

原作者： Simon Jegou, Maximilian Jeblick

发布于 2026-02-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Simon Jegou, Maximilian Jeblick

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在读一本非常长的书，每当你翻一页时，你都必须在脑海中记录下之前读过的每一个单词，才能理解故事的内容。随着书变得越来越长，你的大脑（或者在这种情况下是计算机的内存）会因为试图记住每一个单词而感到不堪重负。这就是现代 AI 模型面临的问题：当它们阅读越来越长的文本时，它们会耗尽“精神空间”（内存）来追踪重要的部分。

这篇论文介绍了一个名为 KVzap 的新工具，它就像是这些 AI 模型的高效图书管理员。以下是它的工作原理，通过简单的概念进行了解：

问题：“精神杂乱”

当 AI 阅读一个句子时，它会为每个单词创建一个“键值对”（Key-Value，简称 KV 对）。你可以把这些想象成写着单词的便利贴，存放在一大堆里。

瓶颈： 如果 AI 阅读一份 10 万字的文档，它就会有 10 万张便利贴。保存所有这些内容会占用大量的内存，从而降低 AI 的速度并增加运行成本。
旧方法： 以前的方法尝试根据简单的规则丢弃便利贴（比如“只保留最后 100 个单词”或“保留出现频率最高的单词”）。但这些规则往往过于粗糙，会不小心丢掉重要信息，导致 AI 出错。

解决方案：KVzap（聪明的图书管理员）

KVzap 是一种新方法，它决定哪些便利贴该保留，哪些该丢弃，但它做得快速、聪明且忠实（不会丢失故事内容）。

以下是 KVzap 工作原理的类比：

1. “二次检查”问题（旧方法）
之前一种顶尖的方法叫做 KVzip，它的准确度很高。它的工作方式是这样的：为了决定一个单词是否重要，AI 会假装把整本书再读一遍，看看它关注了哪些单词。

缺陷： 这就像是要求一名学生读完一个章节后，为了决定哪些句子重要，又得把这个章节再读第二遍。这太慢了，而且非常耗能。此外，它无法在 AI 写作（解码）的过程中工作，只能在阅读时使用。

2. KVzap 的捷径
KVzap 通过训练一个微小的、超快速的“助手”（一个小型的神经网络）来解决这个问题，让它在无需重新阅读文本的情况下就能猜出哪些单词是重要的。

类比： 想象一位阅览过数千本书的资深图书管理员。她不需要重新阅读一本新书来寻找重点，而是只需扫一眼页面，就能立刻知道：“这一段很关键，留着；那一段只是填充内容，扔掉。”
它是如何学习的： 研究团队通过向这个“助手”展示 AI 的内部思维（隐藏状态），并要求它预测那个缓慢的“二次检查”方法会做出怎样的判断，从而训练它。这个助手学会了完美地模仿那位专家，但速度快了许多倍。

3. “滑动窗口”安全网
为了确保 AI 不会忘记即时的上下文（比如它刚刚写的最后几个句子），KVzap 保留了一个包含最近 128 个单词的“滑动窗口”。

类比： 即使你在总结一部长篇小说，你也总是把最后几页拿在手里，这样才不会迷失方向。KVzap 会保护这些最近的单词，绝不丢弃它们。

为什么这很重要？

论文声称 KVzap 是一个游戏规则改变者，主要基于三个原因：

它很快： 它几乎不会增加额外的处理时间。这就像是有一位图书管理员在你走进图书馆的同时就在帮你整理书籍。
它具有适应性： 它不使用固定规则（比如“保留 50% 的笔记”）。相反，它会观察文本。如果文本复杂且密集，它会保留更多笔记；如果文本具有重复性，它会丢弃更多内容。它会自动进行调整。
它很准确： 在长文本阅读任务（如回答来自 4,000 字文档的问题）和推理任务（如解决数学题）的测试中，KVzap 成功地将内存占用缩小了 2 到 4 倍，同时保持了与保留“所有”笔记时几乎完全相同的 AI 准确度。

结果

研究人员在流行的 AI 模型（如 Qwen 和 Llama）上进行了测试。他们发现：

KVzap 打败了 15 种现有的方法。
它在长上下文任务的排行榜上取得了最佳成绩。
它既适用于阅读长文档（预填充/prefilling），也适用于编写长篇故事或逐步解决问题（解码/decoding）。

总结

KVzap 就像是给 AI 装了一个“智能过滤器”，它能瞬间识别出一段长对话或文档中哪些部分是本质的，哪些只是噪音。它让 AI 能够处理更长的文本，而不会耗尽内存或产生混乱，且不会减慢处理速度。作者认为，这使得它已准备好用于主要的现实世界 AI 系统中。

技术摘要：KVzap —— 高速、自适应且忠实的 KV 缓存剪枝

1. 问题陈述

随着基于 Transformer 的语言模型上下文长度扩展到数万甚至数十万个 token，键值（KV）缓存已成为高效大语言模型（LLM）推理的主要瓶颈。KV 缓存存储了中间注意力状态，消耗了大量的 GPU 显存（例如，在 128k 序列长度下，65B 模型需消耗 335 GB），并增加了首字延迟（time-to-first-token），同时降低了解码吞吐量。

虽然架构修改（如 Grouped Query Attention、Multi-head Latent Attention）和稀疏注意力机制已经在头（ $H$ ）、维度（ $D$ ）和层（ $L$ ）轴上解决了压缩问题，但序列长度轴（ $T$ ）在广泛采用的方法中仍未得到充分解决。现有的权宜之计式（ad-hoc）剪枝方法往往无法满足生产环境部署的四个关键标准：

快速且轻量化： 开销可以忽略不计。
阶段无关性： 既适用于预填充（长上下文）也适用于解码（推理）。
优化友好： 与 FlashAttention2 和 PagedAttention 等算子兼容。
忠实度高： 极小化精度损失。

目前的先进方法如 KVzip 虽然实现了高压缩率，但需要对长提示词进行计算昂贵的双向前向传播，这使其不适用于在线推理和解码任务。

2. 方法论

作者引入了 KVzap，这是一种针对 KVzip 的快速、输入自适应近似方案，可在预填充和解码阶段运行。该方法由三个核心组件组成：

2.1. 增强型评分策略 (KVzip+)

作者首先改进了 KVzip 的评分机制。KVzip 根据“上下文重建”对 token 进行评分，即衡量当提示词重复时，一个 token 所能获得的最大注意力权重。作者通过引入受 Devoto 等人 (2025) 启发的一个归一化项来增强这一过程，该项考虑了 token 的值向量（value vector）对残差流（residual stream）的贡献。
KVzip+ 评分 ( $s^+_i$ ) 定义为：
$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$
其中 $a_{ji}$ 是注意力权重， $W_O$ 是输出投影矩阵， $v_i$ 是值向量， $h_j$ 是隐藏状态。

2.2. 代理模型 (Surrogate Modeling)

为了避免 KVzip+ 昂贵且耗时的双向前向传播成本，KVzap 训练了轻量级的代理模型（线性层或两层 MLP），直接从模型的输入隐藏状态 ( $h$ ) 预测对数评分 ( $\log(s^+)$ )。

输入： 隐藏状态 $h_t \in \mathbb{R}^{D_h}$ 。
输出： 每个 KV 头对应的评分 $\in \mathbb{R}^H$ 。
训练： 这些模型在从 Nemotron-Pretraining-Dataset 中采样的 120 万对（隐藏状态，对数评分）上进行训练。MLP 变体使用宽度为 $D_h/8$ 且带有 GELU 激活函数的结构。
性能： 代理模型与 KVzip+ 预言机（oracle）之间的平均 $R^2$ 相关性达到 0.60–0.80，证明了可以通过隐藏状态准确近似这些昂贵的评分。

2.3. 自适应剪枝与驱逐

与强制执行固定预算（例如，始终保留 50% 的 token）的 KVzip 不同，KVzap 采用了基于阈值的剪枝：

阈值化： 预测评分低于固定阈值 $\tau$ 的 token 会被丢弃。这使得压缩率能够根据输入提示词的信息密度进行动态调整。
滑动窗口： 为了保留局部上下文，遵循 StreamingLLM 的方法，始终保留最近的 $w=128$ 个 token。
实现： 剪枝发生在注意力操作之后。该过程计算效率极高，增加的开销微乎其微（MLP 模型为 1.1%，线性模型为 0.02%）。

3. 核心贡献

KVzip+： 一种改进的 KV 缓存重要性评分指标，结合了隐藏状态归一化。
代理近似： 证明了 KVzip+ 评分可以通过在隐藏状态上训练的轻量级模型进行有效近似，从而消除了双向前向传播的需求。
KVzap： 一种全新的剪枝技术，通过应用这些代理模型配合基于阈值的驱逐和滑动窗口，在实现最先进性能的同时，满足了所有四个实际标准（速度、阶段无关性、优化友好性和忠实度）。

4. 实验结果

实验在 Qwen3-8B、Llama-3.1-8B-Instruct 和 Qwen3-32B 上针对三个基准测试进行了评估：RULER（长上下文检索/聚合）、LongBench（多任务长上下文）以及 AIME25（数学推理）。

压缩 vs. 精度： KVzap 在几乎不损失精度的情况下实现了 2–4 倍的 KV 缓存压缩。
- 在 RULER 4k 上，KVzap 的表现达到或超过了 KVzip+ 预言机，并优于其他 15 种剪枝方法（包括 Expected Attention 和 Duo Attention）。
- 在 LongBench 上，KVzap 在 2–3 倍压缩下仍能保持近乎完美的精度。
- 在 AIME25 上，即使丢弃超过 50% 的 KV 缓存，KVzap-MLP 仍能保持推理精度。
自适应性： 基于阈值的方法能根据任务难度自动调整压缩率。例如，与密度更高的真实世界 LongBench 数据相比，它在合成的 RULER 数据上实现了更高的压缩率。
效率： 该方法增加了微不足道的计算和内存开销。在解码过程中，增加的 FLOPs 利用的是原本会因内存带宽限制而停顿的 GPU 周期。

5. 意义与主张

本文声称 KVzap 弥合了学术界剪枝研究与现实世界推理引擎之间的鸿沟。其意义在于：

实用性： 它是首个在不牺牲速度或不需要架构修改的情况下实现最先进压缩的方法，能够与主流推理引擎（vLLM, SGLang, TRT-LLM）兼容。
忠实性： 它证明了 LLM 并没有充分利用 KV 缓存，并且可以通过隐藏状态识别出未使用的键值对，而无需对基础模型进行端到端的重新训练。
可扩展性： 该方法在 8B 到 32B 参数规模的模型上得到了验证，并设计为与未来的稀疏注意力架构兼容。

作者承认存在局限性，指出 KVzap 是一个事后（post-hoc）添加的组件而非端到端训练的目标，并且需要进一步的工程化工作，通过优化的算子（如 PagedAttention）将压缩率转化为实际的运行时间加速。然而，他们认为 KVzap 结合了简单性、高压缩率和鲁棒性能的特点，使其成为生产部署的理想候选方案。