Cost-Aware Diffusion Draft Trees for Speculative Decoding

📄 arXiv: 2606.01813v1 📥 PDF

作者: Shuai Zhang, Huachuan Qiu, Hongliang He, Yong Dai

分类: cs.CL

发布日期: 2026-06-01

💡 一句话要点

提出CaDDTree，通过优化token吞吐量实现更高效的推测解码。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 扩散模型 语言模型加速 token吞吐量 成本优化

📋 核心要点

  1. 现有DDTree方法在推测解码中，预算选择缺乏原则性，忽略了验证成本，导致效率降低。
  2. CaDDTree通过联合优化树结构和节点预算，直接最大化token吞吐量，考虑了draft和验证的延迟。
  3. 实验表明，CaDDTree在多个任务上与具有oracle预算的DDTree性能相当或更优，无需离线预算搜索。

📝 摘要（中文）

推测解码通过轻量级的draft模型生成token草案，并由目标语言模型并行验证，从而加速推理过程。DFlash等块扩散draft模型一次性生成整个草案块，产生每位置的边际分布；DDTree利用这些分布构建候选树，在固定节点预算下最大化预期接受长度。然而，接受长度随预算非递减，总是倾向于更大的树，忽略了验证成本，缺乏预算选择的原则性依据。我们提出了CaDDTree（Cost-aware Diffusion Draft Tree），一种通过联合选择树结构和节点预算，直接优化token吞吐量（单位时间内生成的预期token数）的方法。我们显式地建模了draft和验证延迟，证明了吞吐量目标可以分解为每个回合的预算一维搜索，并且在凸验证成本下，吞吐量函数是单峰的，从而实现了一种高效的贪婪停止规则。CaDDTree不需要离线预算搜索，而是根据当前的每位置分布和验证成本，在每个回合调整预算。在Qwen3-4B和Qwen3-8B上，对涵盖推理、编码和指令遵循任务的八个基准测试的实验表明，CaDDTree在几乎所有任务上都匹配或超过了具有oracle预算选择的DDTree。

🔬 方法详解

问题定义：推测解码旨在加速大型语言模型的推理速度。现有的DDTree方法虽然利用扩散模型生成草案，并通过构建树结构来选择验证token，但其预算选择策略缺乏对验证成本的考虑。DDTree倾向于选择更大的树，即使验证这些token的成本很高，导致整体吞吐量下降。因此，如何根据draft和验证的成本，动态地选择最优的树结构和预算，是需要解决的关键问题。

核心思路：CaDDTree的核心思路是直接优化token吞吐量，即单位时间内生成的预期token数量。它显式地建模了draft和验证的延迟，并将预算选择问题转化为一个优化问题，目标是在给定当前每位置分布和验证成本的情况下，找到能够最大化吞吐量的树结构和节点预算。通过这种方式，CaDDTree能够根据实际的draft和验证情况，动态地调整预算，从而提高整体的推理效率。

技术框架：CaDDTree的整体框架包括以下几个主要步骤：1) 使用块扩散draft模型生成每位置的边际分布；2) 根据这些分布构建候选树；3) 显式建模draft和验证延迟；4) 将吞吐量优化问题分解为每个回合的预算一维搜索；5) 利用单峰性质，使用贪婪停止规则选择最优预算。该框架的关键在于将吞吐量作为优化目标，并根据实际的draft和验证成本动态调整预算。

关键创新：CaDDTree最关键的创新在于其成本意识（Cost-aware）的设计。它不像DDTree那样只关注接受长度，而是将验证成本纳入考虑，直接优化token吞吐量。此外，CaDDTree证明了在凸验证成本下，吞吐量函数是单峰的，这使得可以使用高效的贪婪算法来找到最优预算，避免了耗时的离线预算搜索。

关键设计：CaDDTree的关键设计包括：1) 显式建模draft和验证延迟，这使得能够准确地计算吞吐量；2) 将吞吐量优化问题分解为每个回合的预算一维搜索，降低了计算复杂度；3) 利用单峰性质，使用贪婪停止规则选择最优预算，提高了效率。此外，CaDDTree不需要任何特定的损失函数或网络结构，可以与现有的块扩散draft模型无缝集成。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CaDDTree在Qwen3-4B和Qwen3-8B模型上，对涵盖推理、编码和指令遵循任务的八个基准测试中，在几乎所有任务上都匹配或超过了具有oracle预算选择的DDTree。这意味着CaDDTree能够在不需要离线预算搜索的情况下，达到甚至超过最优的性能。

🎯 应用场景

CaDDTree可应用于各种需要加速大型语言模型推理的场景，如在线对话系统、机器翻译、文本生成等。通过提高token吞吐量，CaDDTree能够降低推理延迟，提升用户体验，并降低计算成本。该方法尤其适用于对延迟敏感的应用，例如实时语音助手和交互式AI应用。

📄 摘要（原文）

Speculative decoding accelerates inference by having a lightweight drafter propose tokens verified in parallel by the target language model. Block diffusion drafters such as DFlash generate an entire draft block in one pass, yielding per-position marginals; DDTree uses these to build a candidate tree that maximizes expected acceptance length under a fixed node budget. We observe, however, that acceptance length is non-decreasing in budget: it always favors larger trees regardless of verification cost, offering no principled basis for budget selection. We introduce \textbf{CaDDTree} (Cost-aware Diffusion Draft Tree), a method that directly optimizes token throughput (expected tokens generated per unit time) by jointly selecting the tree structure and node budget. We model draft and verification latencies explicitly, show that the throughput objective decomposes into a per-round one-dimensional search over the budget, and prove that under a convex verification cost the throughput function is \emph{unimodal}, enabling an efficient greedy stopping rule. CaDDTree requires no offline budget search, adapting the budget each round from the current per-position distributions and verification cost. Experiments on Qwen3-4B and Qwen3-8B across eight benchmarks spanning reasoning, coding, and instruction-following tasks show that \caDDTree{} matches or surpasses DDTree with oracle budget selection on nearly all tasks.