KVpop -- 基于预测性在线剪枝的键值缓存压缩

TL;DR KVpop 通过直接监督未来注意力来学习 token 淘汰策略，以固定预算压缩 KV 缓存，在数学推理上以 75% 压缩率保留 98% 性能，为长上下文推理提供高效内存方案。

问题

问题背景

在基于 Transformer 的大语言模型（LLM）中，KV 缓存是支撑自回归解码高效的关键：生成每个新 token 时，模型重用已计算的键、值张量，避免重计算历史表示。但随着上下文长度增加，KV 缓存的内存占用和带宽需求线性增长，成为长文本推理（如对话、代码生成、文档理解）的核心瓶颈，限制批量大小与并发能力。

现有方法的局限

当前 KV 缓存压缩方法主要分为几类：

• 静态策略：如滑窗 + "sink token" 保留初始 token，仅豁免最近窗口，无法感知内容重要性变化。
• 基于分数的启发式方法：如 H2O、StreamingLLM，在线计算注意力权重或局部相似度作为重要性代理，但这些分数基于当前查询，难以反映 token 在未来时间步的实际效用。
• 学习型方法：如 DMC 通过可学习的合并策略压缩历史，但合并操作破坏了原始 token 的细粒度访问；DMS 训练二值淘汰门，但仅在 token 离开保护窗口后才决策，缺少对全局时序动态的建模。

整体上，现有淘汰策略缺乏明确的 未来注意力监督信号，导致在推理任务（如数学题解、多步逻辑）等需要长程依赖的场景中，压缩率稍高即出现性能崩塌。

为什么该问题既困难又重要

预测 token 的未来效用本质上是 时序决策问题：局部高注意力 token 可能很快失效，而暂时沉寂的 token 可能在数十步后成为关键证据。当压缩率升高（如超过 75%）时，任何误淘汰都可能切断长程依赖链，严重损害生成质量。从工程角度，KV 缓存占用了推理内存的主要部分，有效压缩直接决定大模型在有限硬件上的部署成本与吞吐。工业界对于长上下文服务（例如检索增强生成 RAG、代码库级分析）的需求日益迫切，亟需能保持输出质量的缓存淘汰技术。

行业类比

类似 CDN 的内容驱逐算法，需预测哪些文件未来会被再次请求；KV 缓存淘汰需预测哪些 token 会在后续自回归生成中被注意力机制重新关注。

核心洞察

• KVpop 用未来注意力信号直接监督 token 保留决策，而非传统的静态启发式或历史注意力代理分数。这使得驱逐策略能精准预判 token 在后续生成中的实际效用，从根本上解决了“当前不重要但未来关键”的 token 被误删的问题，在数学推理等需长程依赖的任务中明显优于 StreamLLM、H2O 等基于局部或历史分数的基线方法。
• 提出的延迟记忆评分器（delayed memory-based scorer）在每层引入固定步数的近未来上下文窗口，让状态化驱逐决策能利用该窗口内的注意力模式稳定更新，这是现有学习型 KV 驱逐方法中首次显式利用未来局部信息优化在线决策，有效缓解了逐层独立打分带来的 token 保留抖动，在高压缩率下尤其显著。

方法

输入与问题定义

在自回归生成中，Transformer 的 KV 缓存 存储所有历史 token 的 key 和 value，内存随序列长度线性增长。KVpop 的目标是固定预算下在线决定每个 token 是保留还是驱逐，使未来注意力质量损失最小。

核心模块：基于未来注意力监督的驱逐策略

1. 重要性得分器训练

KVpop 训练一个二元驱逐策略，直接预测每个 token 在未来是否会被用到。训练信号不是当前注意力分布，而是未来注意力目标：对于某一步生成的查询，汇总后续多个查询对该 token 的注意力权重，形成“未来效用”分数。然后以此作为监督，让得分器学习判断 token 是否值得保留。

2. 高效目标计算

为避免显式构建未来密集注意力矩阵，KVpop 利用转置注意力和稀疏归一化近似，在不显著增加计算开销下获得未来注意力目标。具体实现中，通过只计算部分查询位置、采样和 Fenwick 树维护 top-k 掩码等技巧，使训练可行且高效（见附录 C、D）。

3. 延迟记忆得分器（Delayed Memory Scorer）

KVpop 引入一种独特的在线得分器：不是每步立即评分，而是延迟固定步数。在此期间，得分器利用mLSTM 等状态化结构积累近期上下文信息，观察 token 在近未来的实际被使用情况，再做出保留/驱逐决策。这使模型能利用短期依赖提升准确性，避免过早丢弃稍后才重要的 token。

输出与推理

推理时，得分器为缓存中每个 token 实时输出重要性分数，保留 top-k，其余驱逐。通过 Fenwick 树高效维护排序和阈值，实现动态剪枝。

差异点

与基于静态启发式（滑动窗口、sink token）或代理分数（如局部注意力权重）的驱逐方法不同，KVpop 直接以未来注意力为监督目标，并通过延迟评分机制使学习到的策略更贴合 token 真实生命周期，这是同类学习型驱逐方法中独有的设计。

实验

实验设计

实验在 Qwen3-4B 和 Qwen3-8B 两个规模上验证 KVpop 的 KV 缓存压缩能力，采用 AIME 与 HMMT 数学推理基准。不同压缩率（如 75%、88%）下，以全注意力模型性能为教师目标，评估 KVpop 保持原始精度的程度，并与多种 KV 淘汰基线对比，包括滑动窗口、StreamingLLM、H2O 等启发式方法，以及 DMC、DMS 等学习型压缩策略。

关键发现

• 高压缩下性能依然稳健：Qwen3-4B 在 75% 压缩时仅损失 2% 全注意力性能，88% 压缩时仍保持 97%。Qwen3-8B 表现更强，几乎无损逼近教师模型。
• 延迟记忆评分器（delayed memory-based scorer） 是关键：通过延迟评分并利用近未来上下文，显著提升状态感知淘汰效果，尤其在高压缩率下缩小与全注意力差距。
• 未来注意力监督信号 驱动淘汰策略更精准：直接对“保留/丢弃”决策进行二元监督，避免传统基于局部注意力分数的误判，使得不重要但近期活跃的 token 能被正确移除。

与基线对比深度解读

KVpop 与现有方法的根本差异在于 预测性在线淘汰：

• 对比 静态启发式（如滑动窗口、sink token）：后者仅按时间远近或固定初始窗口保留 token，无法动态适应注意力偏移，而 KVpop 通过学习未来注意力目标，能在上下文需求变化时实时调整保留集合。
• 对比 代理分数方法（如 H2O 基于累积注意力分数）：这些分数是注意力历史的滞后反映，对即将到来的需求变化不敏感。KVpop 直接训练对未来 token 有用的保留策略，更具前瞻性。
• 对比 延迟淘汰方法（如 DMS）：DMS 将近期 token 完全保护在窗口内，延迟淘汰决策，而 KVpop 的延迟评分机制在保留近期上下文的同时，仍允许对窗口内 token 进行评估和选择性丢弃，平衡了近期与远期的实用性。

工程启示：KVpop 在维持推理质量的前提下，可大幅缩减 KV 缓存内存占用，对长上下文场景（如多轮对话、长文档理解）的显存优化部署有直接价值。

行业影响

落地场景

KVpop 面向所有依赖 Transformer 自回归解码的在线推理服务，尤其长上下文场景受益最大。典型产品包括：

• 智能客服与对话系统：需要维护长对话历史，KVpop 可减小显存占用，支持更长会话。
• 代码补全与编程助手：上下文窗口常包含整个代码库，KV cache 压缩能大幅降低推理延迟。
• 文档分析 / 问答平台：处理长 PDF、法律合同或学术论文时，KVpop 维持高召回率，同时控制内存成本。
• 实时流式生成：如语音转文字后的大模型总结、直播弹幕分析等，压缩策略可保障吞吐与低延迟。

商业价值

• 推理成本下降：在保持 97-98% 原始性能的前提下，实现 75-88% 的 KV cache 压缩，直接减少 GPU 显存需求。相同硬件可支撑更高并发，或在同批请求下使用更便宜的 GPU 实例，降低云服务开销。
• 用户体验提升：允许服务端支持更长上下文（如 32k → 128k），避免因截断历史导致的上下文丢失，提升回答连贯性。
• 差异化竞争力：对于 API 或企业级部署，更高吞吐与更低成本意味着更优的利润率或定价空间。

与现有产品 / 工作流的集成

KVpop 作为一种 轻量级在线 eviction 策略，可无缝嵌入主流推理框架：

• 直接替换现有策略：在 vLLM、TensorRT-LLM、HuggingFace TGI 等框架中，可接管 KV cache 管理模块，替代静态滑动窗口或启发式淘汰方法。
• 无需重新训练基座模型：只需附加一个 小型 scorer 网络（延迟记忆机制），通过离线微调或在线适配即可生效。scorer 可采用与加速器兼容的 FlexAttention 稀疏核，降低额外开销。
• 与量化 / 卸载技术协同：可与 KV cache 量化（如 INT8）或层级卸载结合，进一步压缩总内存占用，形成多级优化管线。

用例示例：

1. 电商平台智能导购：用户与 AI 助手多轮对话中，历史包含复杂商品对比和偏好。KVpop 压缩使模型在单卡上维持 128 轮对话，不丢失关键意图，同时保持低延迟。
2. 医疗影像报告生成：根据长序列的影像描述和既往病史，模型需关注远距离的异常描述。KVpop 压缩保留高价值 token，在 2 倍上下文中仍生成准确报告，节省高端 GPU 租赁费用。

局限

• **评估范围受限**：实验主要在 AIME 和 HMMT 数学推理数据集上进行，且仅使用 Qwen3-4B/8B 模型。论文未展示在长文档理解、多轮对话或代码生成等更广泛任务上的泛化性能，也未在更大规模模型（如 13B+）或其他架构（如非 Qwen 系列）上验证。对于极长上下文（如 128K token）下 scorer 的训练稳定性和驱逐效果，缺乏实证分析。
• **训练与计算开销**：虽然提出的 **future-attention target** 计算避免了密集注意力图的物化，但仍需在训练时维护额外的注意力统计量，且延迟记忆 scorer（delayed memory-based scorer）引入了额外的状态管理开销。在极长上下文场景下，高效 Top-k 稀疏注意力的 Fenwick 树操作可能成为瓶颈，实际推理时的延迟和吞吐影响缺乏详细的消融实验。
• **方法复杂度与对比基线**：相较于静态启发式（如 sliding window）或无需训练的在线估计方法，KVpop 需要专门的 scorer 训练阶段，涉及多目标损失和超参数（如延迟步数、预算比例、记忆维度）的调优。在与 **DMS** 等可学习的驱逐方法对比时，未充分讨论训练数据需求、收敛速度以及在不同预训练模型上的适配代价，可能限制其在资源受限环境下的直接应用。