🤖 machine learning

BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

该论文介绍了 BudgetDraft，这是一种多视图训练方法，通过为稀疏 KV 草拟器（sparse-KV drafter）配备接受度感知学习，使其在不增加推理时显存成本的情况下，在中长文本上下文中保持高 Token 接受率并实现显著的端到端加速（高达 6.55 倍）。

原作者： Liang He, Jingbo Wen, Qishi Zhan, Yixiong Chen, Kangning Cui, Qizhen Lan, Xilu Wang

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Liang He, Jingbo Wen, Qishi Zhan, Yixiong Chen, Kangning Cui, Qizhen Lan, Xilu Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图写一个长篇故事，但你有一个严格的规则：你只能通过观察一本不断缩小的、记录着前文句子的微型笔记本来决定下一步写什么。这有点像现代 AI 模型在处理长对话或文档时，在内存有限的计算机上运行的方式。

这篇论文介绍了一种名为 BudgetDraft 的新方法，旨在帮助 AI 在不因内存限制而产生混乱的情况下更快地写作。以下是其工作原理的拆解，通过简单的概念进行说明：

问题所在：“内存错位”

把 AI 尝试写故事的过程想象成一个两人团队：

起草者（速度型选手）： 一个小巧、快速的助手，负责快速猜测接下来的几个词。为了节省空间，这个助手只保留故事至今为止的一张微小的“便签”（即稀疏内存）。
校验者（老板）： 一个庞大、聪明的经理，负责检查助手的猜测是否正确。老板的大脑中保留着整个故事（即全量内存）以确保质量。

故障点：
当故事较短时，助手的微型便签还足够使用，他们大多数时候都能猜对。但随着故事变长（4,000 到 16,000 个单词），助手的微型便签就变得毫无用处了。他们开始胡乱猜测，因为他们忘记了故事的开头。
论文将此称为 “16K 崩溃”。助手的猜测变得非常糟糕，以至于老板几乎拒绝了所有内容，导致整个过程变得极其缓慢——有时甚至比不进行猜测、逐字书写还要慢。

解决方案：BudgetDraft

作者意识到，以往的方法都是针对一种特定尺寸的便签来训练助手的。如果计算机的内存发生轻微变化，助手就会失败。

BudgetDraft 是一种新的训练方法，它教会助手具备灵活性。

创意类比：“多视角”健身房

想象你在训练一名篮球运动员（助手）投篮。

旧方法： 你只练习从正好 10 英尺远的地方投篮。如果比赛中篮筐移动到了 12 英尺处，球员就会投丢。
BudgetDraft 方法： 在练习过程中，你随机将篮筐移动到 5 英尺、10 英尺、15 英尺和 20 英尺处。你告诉球员：“无论篮筐在哪里，你都必须瞄准教练（校验者）所指的篮板上的那个确切位置。”

通过在练习中使用许多不同的“预算”（内存大小），助手学会了一项通用的技能。他们不再依赖于特定的内存大小，而是学会了如何让自己的猜测与老板的预期保持一致，无论他们有多少可用内存。

实际运作方式

训练： 助手被展示同一个故事，但被迫使用不同量的内存（有时是 256 个单词，有时是 1024 个，等等）。
目标： 即使在自己的内存非常稀疏时，助手也必须匹配老板对下一个词的“首选选择”。
结果： 助手变得具有“预算鲁棒性”。无论计算机拥有充足的内存还是极少的内存，助手都能保持正确的猜测。

实验结果

论文在三种不同类型的长文本（书籍、会议记录和长篇故事）上测试了该方法。

速度： 在标准的高端计算机上，即使是对于非常长的文本（高达 16,000 个单词），BudgetDraft 也比旧的缓慢方法快 2 到 6 倍。
稳定性： 不同于以往的方法在内存限制改变时会崩溃，BudgetDraft 在所有不同的内存设置下都能平稳运行。
简洁性： 它不需要为计算机添加额外的复杂机制；它只是在训练期间让现有的“助手”变得更聪明。

总结

BudgetDraft 解决了 AI 在内存有限的计算机上编写长文本时变得缓慢且混乱的问题。它通过训练 AI 的“猜测助手”使其具备适应能力，教会它无论拥有微小内存还是大内存都能做出好的猜测，从而确保 AI 保持快速且准确。

技术摘要：BudgetDraft —— 面向稀疏 KV 推测解码的接受度感知多视图训练

1. 问题陈述

本文解决了在部署中长上下文推理（4K–16K token）时，推测解码（Speculative Decoding, SD）面临的关键瓶颈。虽然 SD 通过使用小型“草拟器”（drafter）模型提出 token，再由大型“验证器”（verifier）模型进行验证，从而加速自回归（AR）生成，但标准的部署方式面临着**“稀疏/全量不匹配”**的问题。

约束条件： 为了管理峰值 GPU 显存（VRAM）和延迟，草拟器通常被迫在固定的预算下以稀疏 KV 缓存运行，而验证器则维持全量 KV 缓存以确保输出质量。
失效模式： 随着上下文长度从 4K 增加到 16K，稀疏草拟器与全量验证器之间的对齐度迅速下降。论文识别出了一个**“16K 崩溃”**现象，即接受率骤降至接近于零，导致 SD 失效（加速比 $\le$ 1 $\times$ ）。
敏感性问题： 朴素的稀疏草拟对特定的 KV 预算高度敏感。为单一预算调优的草拟器在不同内存约束下部署时往往会失效，且预算大小与接受率之间的关系是非单调的（较小的预算有时比较大的预算产生更高的接受率）。

2. 方法论：BudgetDraft

作者提出了 BudgetDraft，这是一种在训练阶段的解决方案，旨在创建一个能够实现预算鲁棒性的单一草拟器，使其能够在不同的 KV 预算下保持稳定的接受率，且无需在推理阶段引入额外的组件。

核心架构：
该方法采用了**多视图稀疏训练（Multi-View Sparse Training）**策略，使草拟器（学生）在多种稀疏条件下同时学习与冻结的验证器（教师）进行对齐。

教师目标（Teacher Targets）： 验证器使用其全量 KV 缓存生成贪婪教师目标 ( $x^*_t$ )。
双分支损失函数（Dual-Branch Loss Function）： 训练目标结合了两个互补的损失函数：
1. 接受度感知损失 ( $L_A$ )： 全量缓存分支，草拟器使用全量前缀缓存来预测教师的贪婪 token。这确保了与验证器接受准则（Top-1 匹配）的直接对齐。
2. 多视图稀疏损失 ( $L_C$ )： 稀疏缓存分支，草拟器处理相同的序列，但使用从集合中（例如 $\{256, 512, 1024, 2048\}$ ）随机采样的 KV 预算 ( $B$ )。稀疏缓存通过对前缀进行分块，并根据累积注意力权重保留顶部的块来构建。
训练目标： 总损失为 $L = L_A + \lambda L_C$ 。通过在训练期间让草拟器接触同一前缀的多种稀疏视图，它学会了跨越不同的稀疏水平进行泛化，而不是过拟合于单一预算。
推理： 在部署时，训练好的草拟器使用选定的预算 $B$ 的稀疏 KV 缓存来草拟 token。验证器则使用其全量缓存进行验证。推理流程不需要任何结构性更改或中间模型。

3. 核心贡献

刻画接受度崩溃现象： 本文通过实验证明，朴素的稀疏推测解码在上下文长度从 4K 移动到 16K 时会遭受严重的接受率崩溃，并将 16K 识别为一个实际的失效边界。研究还揭示了 KV 预算大小与接受率之间的非单调关系。
BudgetDraft 框架： 提出了结合接受度感知对齐与多预算稀疏训练的多视图训练方法。这产生了一个预算不变的草拟器，能够在所有稀疏水平下恢复接受率，且没有额外的推理时开销。
实验验证： 证明了 BudgetDraft 在保持内存效率的同时，在单张 NVIDIA A100 GPU 上实现了显著的端到端加速，特别是在具有挑战性的 8K–16K 场景中。

4. 实验结果

实验在三个数据集（PG-19, LongBench, LWM）上进行，使用一个 68M 参数的草拟器和一个 7B 参数的验证器，并在单张 NVIDIA A100 GPU 上运行。

加速性能：
- 4K 上下文： 实现高达 6.55 $\times$ 的 AR 加速。
- 8K 上下文： 实现高达 4.46 $\times$ 的 AR 加速。
- 16K 上下文： 实现高达 2.10 $\times$ 的 AR 加速（在此场景下，朴素 SD 通常完全失效）。
鲁棒性： 与基线 SD（稀疏/全量）不同，后者在上下文长度增加或预算改变时表现出剧烈的性能下降或接近于零的接受率，而 BudgetDraft 在不同的 KV 预算（ $B \in \{256, 512, 1024, 2048\}$ ）下保持了稳定的接受率。
对比： 在 8K–16K 场景中，BudgetDraft 的表现优于结构化缓解方法 TriForce 和先进的草拟方法 EAGLE-3。例如，在 LWM 的 16K 场景下，BudgetDraft 实现了约 1.94 $\times$ 的加速，而 TriForce 为 1.19 $\times$ ，EAGLE-3 为 1.36 $\times$ 。
消融实验： 移除多视图稀疏损失 ( $L_C$ ) 会导致在预算偏移下的性能不稳定，这证实了在训练期间进行多预算采样的必要性。

5. 重要性与主张

本文声称 BudgetDraft 解决了目前限制推测解码在长上下文场景中实际部署的“稀疏/全量不匹配”问题。通过将复杂度转移到训练阶段，它实现了一个简单、内存友好且对不同硬件约束（VRAM 可用性）和工作负载并发性具有鲁棒性的推理流水线。

作者强调，该方法不需要复杂的结构性更改（如检索缓存）或在推理时使用多个模型。相反，它提供了一个轻量级的单一草拟器，可以适应不同的部署预算，使得推测解码在以往难以实现的、适用于中长上下文应用（4K–16K）的场景中变得可行。

论文承认的局限性：

位置外推： 68M 的草拟器具有原生位置嵌入限制（2048），需要进行外推才能处理 8K–16K 上下文，这可能仍会限制接受率。
验证器特异性： 草拟器是针对特定验证器训练的；如果更换验证器模型，则需要重新训练。
贪婪解码： 结果是基于贪婪解码得出的；将扩展到基于采样的解码（如温度采样）是未来的工作方向。