跳转至

Inference-Time Hyper-Scaling with KV Cache Compression

会议: NeurIPS 2025
arXiv: 2506.05345
代码: HuggingFace Models（NVIDIA开源）
领域: 模型压缩/推理加速
关键词: KV Cache压缩, 推理时缩放, 稀疏注意力, Dynamic Memory Sparsification, 长序列推理

一句话总结

提出"推理时超缩放"（Inference-Time Hyper-Scaling）范式：通过高效压缩KV缓存，在相同计算/内存预算下生成更长或更多并行推理序列，显著提升推理模型在数学、代码、科学推理等任务上的准确率。

研究背景与动机

推理时缩放（Inference-Time Scaling）是当前提升LLM推理能力的主要手段，通过生成更长的推理链（如Chain-of-Thought）或多条并行路径（如Best-of-N），以计算换准确率。然而，在Transformer LLM中，生成成本的真正瓶颈并非生成的token数量，而是KV缓存的大小：

内存瓶颈：KV缓存随序列长度线性增长，对于32B参数模型，即使是8K上下文也需要大量GPU内存

延迟瓶颈：自回归生成是内存带宽受限任务，每一步都需要读取全部KV缓存，缓存越大延迟越高

缩放上限：在固定内存预算下，KV缓存大小限制了可并行生成的序列数或单序列最大长度

作者观察到一个关键洞察：如果能在高压缩比下压缩KV缓存同时保持准确率，就可以在相同计算预算下生成更多token，实现推理准确率的进一步提升——这就是"超缩放"的核心思想。

方法详解

整体框架

Inference-Time Hyper-Scaling的核心流程：

  1. 使用KV缓存压缩方法将缓存压缩至原来的 \(1/C\)（如 \(C=8\)）
  2. 在相同内存预算下，可以运行 \(C\) 倍多的并行推理路径（宽度缩放 \(W\)），或生成 \(C\) 倍长的序列（深度缩放 \(D\)）
  3. 最终通过多数投票（Majority Voting）聚合并行结果，或直接取最长序列的答案

关键前提是压缩方法必须在高压缩比下仍能保持生成质量。

关键设计：Dynamic Memory Sparsification (DMS)

DMS是本文提出的KV缓存稀疏化方法，核心创新有三点：

1. 延迟驱逐（Delayed Eviction）

  • 与传统方法（如H2O、StreamingLLM）在写入时就决定驱逐不同，DMS保留完整缓存直到容量上限
  • 达到上限后才批量执行一次稀疏化，选择保留最重要的 \(B\) 个token
  • 优势：在驱逐前，所有token已通过多层注意力隐式"合并"了信息，被保留的token包含了被驱逐token的语义

2. 可学习的稀疏化阈值

  • 为每层每头学习一个评分函数，基于注意力分数决定哪些KV对应保留
  • 采用Gumbel-Sigmoid实现可微分的token选择
  • 仅需标准语言建模损失（next-token prediction）训练

3. 极简训练

  • 仅需 1K 步训练（约数小时），使用少量文本数据即可为预训练模型添加DMS能力
  • 训练好的DMS模块可以直接应用于不同长度和不同任务的推理

损失函数 / 训练策略

  • 损失函数：标准的next-token prediction交叉熵损失
  • 训练时在随机位置触发稀疏化，使模型学会在不同压缩时机下保持生成质量
  • 仅训练稀疏化相关的少量参数（评分头），原模型参数冻结

实验关键数据

主实验

在多个推理模型上验证Hyper-Scaling的效果（使用DMS实现8倍压缩）：

模型 任务 基线（Full Cache） DMS 8× + Hyper-Scale 提升
Qwen-R1 32B AIME 24 72.0 84.0 +12.0
Qwen-R1 32B GPQA 61.6 70.2 +8.6
Qwen-R1 32B LiveCodeBench 57.3 67.0 +9.7
Qwen3-8B AIME 24 62.7 73.3 +10.6
Qwen3-8B GPQA 54.8 62.1 +7.3
Qwen3-8B LiveCodeBench 48.5 58.6 +10.1

核心发现：在等效计算预算（相同内存读取次数）下，DMS压缩+更多并行的组合一致优于不压缩的原始模型。

DMS与其他压缩方法的对比

方法 类型 训练需求 8×压缩后AIME 24 8×压缩后GPQA
Full Cache — — 72.0 61.6
StreamingLLM Training-free 无 30.2 38.1
H2O Training-free 无 41.5 42.8
SnapKV Training-free 无 55.3 50.7
DMS (本文) Learned 1K步 68.7 59.4

DMS在8倍压缩比下仅损失3-4分，远优于所有training-free方法。

消融实验

延迟驱逐的重要性：

策略 AIME 24 (8×) GPQA (8×)
即时驱逐（Eager） 52.1 47.3
延迟驱逐（Delayed） 68.7 59.4

延迟驱逐带来16+分的提升，验证了"隐式信息合并"的关键作用。

缩放策略对比（等效计算预算）：

策略 宽度W 深度D AIME 24
原始 1 1× 72.0
纯宽度缩放 8 1× 80.3
纯深度缩放 1 8× 76.5
混合缩放 4 2× 84.0

混合缩放（同时增加并行数和序列长度）效果最佳。

关键发现

  1. 超缩放有效：KV缓存压缩→更多token→更高准确率，这一链条在多个模型和任务上一致成立
  2. 压缩质量是关键：只有高质量的压缩方法（如DMS）才能让超缩放带来正收益；低质量压缩反而有害
  3. DMS训练极其高效：1K步训练即可达到接近Full Cache的质量，且训练后的模块跨任务泛化良好
  4. 宽度+深度混合缩放最优：单纯增加并行数或序列长度不如两者结合

亮点与洞察

  1. 范式创新：将KV缓存压缩从"效率优化"工具重新定位为"推理能力增强"工具，开辟了全新的方向
  2. 隐式信息合并：延迟驱逐机制利用注意力机制本身的信息聚合能力，被保留的token自然吸收了邻近被删token的信息
  3. 实用性极强：1K步训练、8倍压缩、显著提升——可直接用于现有推理模型的部署优化
  4. 多数投票的可靠性：在并行推理场景下，多数投票是比简单选择更鲁棒的答案聚合方式

局限与展望

  1. DMS仍需少量训练，对于频繁更新的模型有额外成本
  2. 超缩放在简单任务上收益递减，对于不需要长推理链的任务可能效率不高
  3. 评分函数仅基于注意力分数，未考虑token内容的语义重要性
  4. 多数投票作为聚合策略较为简单，更复杂的验证机制可能进一步提升效果

相关工作与启发

  • StreamingLLM / H2O / SnapKV：training-free的KV缓存压缩方法，本文证明它们在高压缩比下质量不足
  • Best-of-N / Majority Voting：推理时缩放的基础策略，DMS为其提供了更大的缩放空间
  • KVQuant / KIVI：量化方式压缩KV缓存，与稀疏化的DMS是正交的，可以组合使用
  • 启发：模型效率优化不仅可以降成本，还可以提能力——这一思路可推广到其他模块的压缩

评分

  • 新颖性：⭐⭐⭐⭐⭐（范式级创新）
  • 技术深度：⭐⭐⭐⭐
  • 实验充分度：⭐⭐⭐⭐⭐
  • 实用性：⭐⭐⭐⭐⭐
  • 写作质量：⭐⭐⭐⭐

相关论文