技术摘要：KVpop — 基于预测性在线剪枝的键值（KV）缓存压缩

1. 问题陈述

基于 Transformer 的大语言模型（LLM）中的自回归解码依赖于键值（KV）缓存来避免重新计算过去的 Token 表示。然而，随着上下文长度的增加，KV 缓存呈线性增长，这在长上下文推理过程中造成了显著的内存和带宽瓶颈。

现有的降低成本的解决方案通常分为两类：

• 静态/启发式驱逐（Static/Heuristic Eviction）： 诸如滑动窗口（Sliding Window）或 Sink-Token 保留（例如 StreamingLLM）等方法依赖于固定规则，而基于评分的方法则使用在线代理指标（例如累积注意力）。这些方法往往无法追踪未来 Token 的效用，导致在相关性发生偏移时出现脆弱的驱逐行为。
• 学习型驱埃（Learned Eviction）： 诸如动态内存压缩（DMC）或动态内存稀疏化（DMS）等方法通过训练策略来合并或丢弃 Token。然而，其中许多方法在 Token 插入时就做出驱逐决策，或者依赖于无法直接基于未来效用监督特定“保留或丢弃”决策的可微松弛方法。此外，带有保护窗口的方法虽然延迟了 Token 的移除，但往往在 Token 进入时就决定了其命运，忽略了在受保护窗口内积累的证据。

核心挑战在于预测哪些 Token 对未来查询是有用的，因为局部显著的 Token 可能会变得无关紧要，而当前注意力较低的 Token 在稍后可能会变得至关重要（例如在推理轨迹中）。

2. 方法论：KVpop

KVpop（具有预测性在线剪枝功能的 KV 压缩）是一种稀疏注意力后置组件（retrofit），旨在不改变基础模型架构的情况下，强制执行固定的每头 KV 预算。它通过两个主要机制解决了先前工作的局限性：未来注意力监督（Future-Attention Supervision）和延迟状态化评分（Delayed Stateful Scoring）。

2.1 架构与预算

KVpop 维护一个由 $B = s + w + k$ 定义的每 KV 头固定预算：

• $s$ Sink Tokens： 始终保留前几个 Token。
• $w$ 保护窗口（Protected Window）： 最近的 $w$ 个 Token 免于被驱逐。
• $k$ 长程 Top-K（Long-Range Top-K）： 较旧的 Token 根据学习到的重要性评分竞争固定的 $k$ 个插槽。

2.2 未来注意力目标监督

与基于过去注意力或局部统计量进行评分的方法不同，KVpop 使用 Token 在离开保护窗口后的**未来注意力质量（future attention mass）**来监督评分器。

• 目标定义： 对于一个 Token $t$，其目标是来自未来查询（位置 $d \ge t + w$）所接收到的聚合注意力质量的对数。
• 高效计算： 为了避免在训练期间物化密集注意力图带来的 $O(S^2)$ 开销，KVpop 使用了转置注意力传递（transposed-attention pass）。通过交换查询（queries）和键（keys）的角色，并复用稀疏注意力算子返回的 log-sum-exp (LSE) 归一化因子，系统可以高效地计算未来注意力目标，而无需物化密集矩阵。
• 损失函数： 训练目标是一个边界感知的配对逻辑损失（boundary-aware pairwise logistic loss）。它将新符合条件的 Token（$t_{new}$）的预测得分与教师模型的截断 Token（$t_{bnd}$）的得分进行比较。该损失函数推动学生模型正确地相对于边界对 $t_{new}$ 进行排序，确保 top-$k$ 选择与教师模型的未来效用排序相匹配。

2.3 带延迟决策的状态化评分

KVpop 引入了一个延迟内存评分器（通过 mLSTM 实现），该评分器将评分决策推迟到 Token 到达驱逐边界时。

• 机制： 评分器不是在 Token 插入时立即评分，而是通过其内部状态更新 Token 的表示及其后续上下文。当 Token 退出保护窗口（$q - w$）时，评分器读取累积的记忆以生成重要性评分。
• 优势： 这使得评分器能够整合“近未来”上下文（即在 Token 处于保护窗口期间积累的证据），从而在做出保留或丢弃决策之前进行评估，解决了即时评分的局限性。

2.4 推理效率

在推理期间，KVpop 强制执行固定预算 $B$。Top-$k$ 长程 Token 的选择通过对静态 Token 优先级和 Fenwick 树进行运行阈值处理进行优化，允许在注意力算子（例如 FlexAttention）内部即时生成稀疏掩码，而无需存储密集的 $S \times S$ 掩码。

3. 核心贡献

1. 监督驱逐策略： 作者将固定预算的 KV 驱逐形式化为一个在驱逐边界处的监督预测问题，利用未来注意力目标来训练长程 top-$k$ 决策。
2. 高效的目标计算： 他们引入了一种仅用于训练的转置注意力传递，用于计算未来注意力目标，而无需物化密集注意力图，确保在推理时没有额外开销。
3. 延迟评分： 他们证明了状态化评分器可以通过将评分决策推迟到 Token 变得可驱逐时，来利用近未来上下文，从而优于即时评分方法。
4. 经验验证： 该方法在数学推理（AIME, HMMT）和领域外任务（GPQA-D, LiveCodeBench）上进行了评估，表明学习型驱逐策略可以在激进压缩下保持高性能。

4. 实验结果

作者在 Qwen3-4B 和 Qwen3-8B 模型上对该方法进行了评估，并在 Nemotron-Math v2 数据集上进行了长度为 16,384 的序列训练。

4.1 数学推理 (AIME & HMMT)

• 压缩率 75% (预算 2048)：
  • Qwen3-4B： KVpop 保留了 98% 的全注意力教师模型性能。
  • Qwen3-8B： KVpop 保留了 100%（接近全量教师性能）。
  • 对比： KVpop 显著优于无训练基线（StreamingLLM, TOVA）和学习型基线 DMS。例如，在 75% 压缩率下，StreamingLLM 在 Qwen3-4B 上的相对性能降至 ~76%，而 KVpop 保持在 98%。
• 压缩率 88% (预算 4096)：
  • Qwen3-4B： KVpop 保留了 97% 的教师性能。
  • Qwen3-8B： KVpop 保留了 100% 的教师性能。
  • 对比： 无训练方法出现了显著崩溃（例如，StreamingLLM 降至 ~47% 的相对性能），而 KVpop 依然保持稳健。

4.2 领域外泛化

在 GPQA-Diamond（科学推理）和 LiveCodeBench（代码生成）上进行评估：

• 仅在数学推理数据上训练的 KVpop 模型在这些领域表现出了良好的泛化能力。
• 在 75% 压缩率下，Qwen3-4B 上的 KVpop 在 GPQA-D 上达到了与密集教师模型相当的准确度（0.59 vs 0.59），并在 LCB 上取得了具有竞争力的结果。

4.3 推理效率

• 内存： 与密集注意力（其 VRAM 使用量随生成长度线性扩展，在 131k token 时达到 36GB）不同，KVpop 和 DMS 保持了近乎恒定的内存占用（19GB），且不受序列长度影响。
• 延迟： 在长生成长度（高达 131k tokens）下，KVpop 的端到端解码延迟低于 DMS 和密集注意力。作者将其归功于 KVpop 的固定每头预算，这避免了 DMS 中复杂的“参差不齐（ragged）”缓存结构，从而简化了编译和执行。

4.4 消融实验

• 延迟评分： 移除延迟（在插入时立即评分）会导致 Token 准确率下降 0.2 个点，证实了整合近未来上下文的益处。
• 无状态 vs 有状态： 有状态变体（KVpop）通常与无状态变体（KVpopmlp）持平或略优，特别是在较高的压缩率下。

5. 重要性与主张

论文声称，使用未来注意力信号进行监督驱逐是比静态启发式或代理评分更优的 KV 缓存压缩策略。通过直接训练一个预测哪些 Token 对未来查询重要的策略，KVpop 实现了：

1. 高保真度： 即使在高压缩率（高达 88%）的情况下，它也能保持接近全量教师模型的性能，而在此阶段其他方法往往会失效。
2. 鲁棒性： 学习到的策略能有效地迁移到领域外任务，表明模型学习到了通用的 Token 效用原则，而非仅仅是对特定训练分布的过拟合。
3. 高效性： 它强制执行硬性内存限制，从而在实现恒定内存使用的同时，支持稳定的长上下文解码，并提供了比 DMS 等动态门控方法更高的吞吐量。

作者总结道，KVpop 表明，当通过未来效用进行适当监督并允许利用近未来上下文时，学习型驱逐策略可以有效地解决自回归生成中的内存瓶颈，且不会牺牲模型质量。