DMS：动态内存稀疏化——1000步训练实现8倍KV Cache压缩

ArXiv ID: 2506.05345

作者: Adrian Lancucki等

机构: NVIDIA, University of Edinburgh

发布日期: 2025年6月

摘要

推理时扩展（Inference-time Scaling）通过生成更长或更多的推理路径来提升LLM的推理能力，但其效率受限于KV Cache的内存瓶颈。本文提出推理时超缩放（Inference-Time Hyper-Scaling）的概念：通过压缩KV Cache，在相同计算预算内生成更多token，从而进一步提升推理精度。为此，作者提出DMS（Dynamic Memory Sparsification），一种仅需1000步训练即可实现8倍KV Cache压缩的稀疏化方法，同时精度优于无训练的稀疏注意力方案。DMS的关键设计包括延迟淘汰策略和可学习的逐头逐层淘汰策略。在Qwen-R1 32B上，DMS在AIME 24提升12.0分、GPQA提升8.6分、LiveCodeBench提升9.7分。

核心概念：推理时超缩放

传统推理时扩展的瓶颈

传统 Inference-Time Scaling:
更多推理token → 更大KV Cache → 更多内存 → 受限于GPU显存

Hyper-Scaling (DMS):
更多推理token → KV Cache压缩 → 内存不变 → 可生成更多token
                                          → 推理精度更高

核心洞察：推理时计算的瓶颈不是FLOPs而是KV Cache内存。如果能压缩KV Cache，就能在相同硬件上”免费”获得更多推理计算预算。

DMS方法

1. 延迟淘汰

与传统方法（H2O、TOVA等）立即淘汰不重要token不同，DMS将token标记为”即将淘汰”后，仍保留一个滑动窗口期，让模型有机会吸收其信息：

def delayed_eviction(kv_cache, eviction_decisions, delay_window=64):
    """
    标记为淘汰的token不立即删除，
    而是在delay_window步后才实际移除
    """
    for token_idx in eviction_decisions:
        # 标记token的淘汰时间
        kv_cache[token_idx].evict_at = current_step + delay_window

    # 实际淘汰
    for token in kv_cache:
        if current_step >= token.evict_at:
            kv_cache.remove(token)

延迟淘汰的优势：

• 模型在淘汰前有时间将关键信息编码到后续token中
• 类似于”隐式合并”——信息不是丢弃而是被转移
• 显著减少信息损失，尤其在高压缩率下

2. 可学习的淘汰策略

DMS为每个注意力头、每一层学习独立的淘汰策略：

def learnable_eviction_policy(Q, K, head_idx, layer_idx):
    """
    使用Gumbel-sigmoid实现可微分的淘汰决策
    """
    # 每个头/层有独立的淘汰参数
    eviction_logits = eviction_head[head_idx][layer_idx](K)

    # Gumbel-sigmoid：训练时可微分，推理时离散
    if training:
        eviction_probs = gumbel_sigmoid(eviction_logits, temperature=0.5)
    else:
        eviction_probs = (eviction_logits > 0).float()

    return eviction_probs

3. 极低训练成本

DMS仅需约1000步微调即可retrofit到现有模型：

training_config:
  num_steps: 1000
  learning_rate: 1e-4
  trainable_params: eviction_heads_only  # 仅训练淘汰决策头
  frozen_params: all_model_weights       # 原模型参数冻结
  compression_ratio: 8x
  estimated_time: ~30分钟 (单GPU)

实验结果

推理精度提升（Qwen-R1 32B，等计算预算）

基准	原模型	DMS (8x压缩)	提升
AIME 24	基线	+12.0	显著
GPQA	基线	+8.6	显著
LiveCodeBench	基线	+9.7	显著

解读：DMS允许模型在相同内存下生成更多推理token，而更多的推理token直接转化为更高的推理精度。

与其他压缩方法对比

方法	训练需求	压缩率	精度保持
DMS	1K步	8x	优秀
TOVA	无训练	4x	中等
H2O	无训练	4x	中等
DMC	100K步	8x	良好

DMS在训练效率和压缩效果之间取得了最佳平衡。

与生态系统的关系

DMS与推理时扩展（Test-Time Scaling）方法高度互补：

Chain-of-Thought / Self-Consistency
          ↓
     更长推理序列
          ↓
    KV Cache膨胀 → DMS压缩 → 内存可控
          ↓
  可生成更多推理路径 → 精度提升

部署建议

最佳场景：

• 推理型模型（如R1、o1类模型）的serving优化
• 需要长推理链的数学和代码任务
• GPU显存受限但需要强推理能力的部署

实施步骤：

1. 使用预训练好的推理模型
2. 运行1000步DMS retrofit训练（约30分钟）
3. 部署时启用8x KV Cache压缩
4. 将节省的内存分配给更长的推理序列

个人评价

DMS提出的”推理时超缩放”概念非常有洞察力——它将KV Cache压缩从”节省内存”的工具升级为”提升推理精度”的手段。1000步训练成本的极低门槛使其具备极强的实用性。来自NVIDIA的背景也意味着后续可能有更好的框架集成支持。在推理型模型日益流行的当下，DMS解决的正是最紧迫的瓶颈之一。

---

评分: 4.40/5.0

论文: https://arxiv.org/abs/2506.05345