跳转至

KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction

会议: NeurIPS 2025 (Oral)

arXiv: 2505.23416

代码: 有

领域: 模型压缩 / LLM推理优化

关键词: KV Cache压缩, 查询无关, 上下文重建, 长上下文推理, 缓存复用

一句话总结

提出 KVzip，一种查询无关的 KV Cache 驱逐方法，通过利用 LLM 自身从缓存的 KV 对中重建原始上下文来量化每个 KV 对的重要性，实现 3-4× 的 KV Cache 压缩率和约 2× 的 FlashAttention 解码延迟降低，同时在多查询场景下显著优于现有查询感知方法。

研究背景与动机

随着 LLM 处理的上下文长度不断增长，KV Cache 的内存开销和注意力计算延迟成为部署瓶颈：

内存挑战：长上下文场景（如 170K token）下 KV Cache 占据大量 GPU 内存

延迟挑战：注意力计算的延迟随 KV Cache 长度线性增长

现有 KV Cache 压缩方法主要是查询感知的（query-aware），即根据当前查询的注意力分数来决定保留哪些 KV 对。但这种方法有根本缺陷：

  • 缓存无法复用：每次查询需要重新压缩，无法在多查询场景中复用同一压缩缓存
  • 性能退化：即使在 90% 的缓存预算比下，查询感知方法在多查询场景也会出现性能下降
  • 需要查询预知：必须在处理查询时才能决定哪些 KV 对重要

KVzip 的核心洞察：KV 对的重要性应该由其对上下文重建的贡献来衡量，而非由特定查询的注意力分数决定。

方法详解

整体框架

KVzip 工作流程：

  1. 预填充阶段：LLM 正常处理上下文，生成完整的 KV Cache
  2. 重要性评估：使用 LLM 自身评估每个 KV 对对重建原始上下文的重要性
  3. KV 驱逐：移除重要性低的 KV 对，保留压缩后的 KV Cache
  4. 多查询推理：压缩后的 KV Cache 可被多个不同查询复用

关键设计

1. 基于上下文重建的重要性评估

KVzip 利用 LLM 本身作为评估工具：给定一组候选保留的 KV 对，测量 LLM 能否从这些 KV 对中重建原始输入上下文。重建能力越强，说明保留的 KV 对包含了更多关键信息。

具体而言，重要性分数通过以下方式计算： - 计算每个 KV 对被移除后上下文重建损失的增量 - 增量越大，说明该 KV 对越重要，应该保留

2. 查询无关性

重要性评估完全基于上下文本身，不依赖任何未来查询。这意味着： - 压缩只需执行一次 - 压缩后的 KV Cache 可为任意后续查询服务 - 特别适合 prefill-once, decode-many 的部署场景（如文档QA、多轮对话）

3. 分层驱逐策略

不同注意力层和不同注意力头对 KV 对的重要性可能不同，KVzip 允许在每一层独立估计重要性并执行驱逐，实现更精细的压缩。

损失函数 / 训练策略

KVzip 是一种无需额外训练的方法，直接利用现有 LLM 的能力进行重要性评估。上下文重建损失即为标准的交叉熵损失，衡量从压缩 KV Cache 中重建原始 token 序列的能力。

实验关键数据

主实验

表1: 多查询场景下的性能（保留 25% KV Cache）

方法 类型 QA准确率 检索准确率 推理准确率 代码理解准确率 平均
Full Cache (100%) - 基准 基准 基准 基准 100%
H2O 查询感知 显著下降 显著下降 显著下降 显著下降 <80%
SnapKV 查询感知 下降 下降 下降 下降 <85%
KVzip 查询无关 接近基准 接近基准 接近基准 接近基准 ~97%

表2: 不同模型上的压缩效果

模型 KV Cache压缩率 FlashAttention延迟降低 性能损失
LLaMA3.1-8B 3-4× ~2× 可忽略
Qwen2.5-14B 3-4× ~2× 可忽略
Gemma3 3-4× ~2× 可忽略

表3: 不同缓存预算比下的对比（多查询场景）

缓存预算比 H2O (查询感知) SnapKV (查询感知) KVzip (查询无关)
90% 性能退化 性能退化 接近满缓存
50% 严重退化 明显退化 轻微影响
25% 不可用 大幅下降 仍可接受

消融实验

上下文重建 vs 其他重要性度量

重要性度量方法 25% 预算下平均性能
随机驱逐 很差
注意力分数 (首个query) 中等但不稳定
累积注意力分数 较好但退化
上下文重建 (KVzip) 最佳且稳定

长上下文扩展测试

上下文长度 KVzip 压缩率 性能保持
8K tokens 3× >98%
32K tokens 3× >97%
128K tokens 4× >95%
170K tokens 4× >94%

关键发现

  1. 查询感知方法在多查询场景失效：即使 90% 预算比，H2O 等方法也因为针对特定查询优化而在其他查询上退化
  2. 上下文重建是更好的重要性信号：比注意力分数更能反映 KV 对的全局重要性
  3. 跨模型一致性：在 LLaMA3.1、Qwen2.5、Gemma3 上均有效
  4. 超长上下文可行性：在 170K token 上下文长度下仍能维持有效压缩
  5. 实际加速：FlashAttention 解码延迟降低约 2×，内存减少 3-4×

亮点与洞察

  1. 查询无关范式：重新定义了 KV Cache 压缩问题——重要性不应由查询决定，而应由上下文信息密度决定
  2. 自监督信号：巧妙利用 LLM 自身作为"评判者"评估 KV 对重要性，无需额外标注或训练
  3. Prefill-once 部署：完美契合长文档QA、代理推理等需要对同一上下文多次查询的实际场景
  4. NeurIPS 2025 Oral：说明审稿人对这个简洁而有效的方法给予了高度评价
  5. 即插即用：不修改模型架构，可直接应用于任何 Transformer LLM

局限与展望

  1. 重要性评估的开销：需要额外的前向计算来评估 KV 对重要性（一次性成本）
  2. 静态压缩：一旦压缩完成，无法根据新查询动态调整保留的 KV 对
  3. 可与量化结合：KV Cache 驱逐与量化是正交方法，结合可获得更大压缩率
  4. 注意力头异质性：不同注意力头可能有不同的最优压缩率，自适应策略值得探索
  5. 流式场景：当上下文持续增长时（如对话），需要增量式重要性更新

相关工作与启发

  • H2O (Zhang et al., 2024)：基于注意力分数的 KV Cache 驱逐
  • SnapKV (Li et al., 2024)：通过观察窗口注意力模式进行压缩
  • StreamingLLM (Xiao et al., 2023)：保留 attention sink token 的流式推理
  • GQA / MQA：通过共享 KV 头减少 KV Cache 大小
  • 启发：KV Cache 压缩的关键不是找到"对当前查询重要的 token"，而是找到"包含上下文信息的 token"

评分

维度 分数 (1-5) 说明
创新性 5 查询无关压缩范式，上下文重建的新颖角度
技术质量 4 方法简洁有效，NeurIPS Oral
实验充分度 5 多模型、多任务、多长度、多对比方法
实用性 5 即插即用，解决实际部署痛点
写作质量 4 问题定义清晰，结果展示充分

相关论文