Title: Cost-Aware Diffusion Draft Trees for Speculative Decoding

Abstract:

Speculative decoding enhances inference speed by employing a lightweight drafter to propose tokens, which are then verified in parallel by a larger target language model. Methods like DFlash, a block diffusion drafter, produce per-position marginals by generating an entire draft block in a single pass. DDTree leverages these marginals to construct a candidate tree designed to maximize the expected acceptance length, constrained by a fixed node budget. However, we find that acceptance length is non-decreasing with respect to the budget; consequently, larger trees are always preferred regardless of the associated verification costs, leaving no principled framework for selecting an optimal budget.

To address this, we propose \textbf{CaDDTree} (Cost-aware Diffusion Draft Tree), a novel approach that directly optimizes token throughput—defined as the expected number of tokens generated per unit of time—by simultaneously determining the tree structure and the node budget. By explicitly modeling both draft and verification latencies, we demonstrate that the throughput objective simplifies into a one-dimensional search over the budget for each round. Furthermore, we prove that when verification costs are convex, the throughput function is \emph{unimodal}, which allows for an efficient greedy stopping rule. Unlike previous methods, CaDDTree eliminates the need for offline budget searches, dynamically adjusting the budget each round based on current per-position distributions and verification costs. Evaluations on Qwen3-4B and Qwen3-8B across eight benchmarks covering reasoning, coding, and instruction-following tasks reveal that CaDDTree matches or exceeds the performance of DDTree with oracle budget selection on almost all tasks.

Source: arXiv Generated at: 2026-06-02 00:00:00 UTC