Cost-Aware Diffusion Draft Trees for Speculative Decoding

Technical Summary: Cost-Aware Diffusion Draft Trees for Speculative Decoding

Problem Statement
Speculative decoding accelerates autoregressive language model inference by using a lightweight drafter to propose candidate tokens, which are then verified in parallel by the target model. While tree-based verification allows for the simultaneous verification of multiple candidate continuations, existing methods for constructing these draft trees face a fundamental misalignment in their optimization objectives.

Current approaches, such as DDTree, construct draft trees to maximize the expected acceptance length under a fixed node budget. However, the authors observe that acceptance length is non-decreasing with respect to the budget; larger trees always yield longer acceptance lengths, offering no principled basis for selecting an optimal budget size. Furthermore, the optimal budget is not static; it varies significantly across decoding rounds depending on the drafter's confidence and the current context length. Existing methods that adapt tree shape per step still rely on a fixed, offline-selected total node budget, failing to respond to per-round variability in the trade-off between coverage and verification cost.

Methodology: CaDDTree
The paper introduces CaDDTree (Cost-aware Diffusion Draft Tree), a framework that directly optimizes token throughput (expected tokens generated per unit time) rather than acceptance length. The method operates on block diffusion drafters (e.g., DFlash), which generate per-position marginal distributions in a single forward pass.

1. Throughput Objective: The authors model the round cost as $C(T; \ell) = C_d + C_v(|T|; \ell)$, where $C_d$ is the drafting cost and $C_v$ is the verification cost dependent on the tree size $|T|$ and context length $\ell$. The objective is to maximize the ratio of expected accepted tokens ($1 + \Phi(T)$) to the total cost.
2. Decomposition: The authors prove that maximizing throughput over all valid trees reduces to a one-dimensional search over the node budget $n$. For any fixed $n$, the optimal tree structure is the "greedy best-first" tree (selecting nodes with the highest prefix probabilities), which is already the construction used by DDTree.
3. Unimodality and Greedy Stopping: A key theoretical contribution is the proof that under a mild convexity condition on the verification cost function, the surrogate throughput objective is unimodal with respect to the tree size. This property allows for an efficient greedy algorithm: the algorithm incrementally adds nodes to the tree and stops as soon as the marginal gain in throughput becomes non-positive.
4. Adaptive Budget: Unlike prior work, CaDDTree does not require an offline grid search for a fixed budget. Instead, it adapts the budget $n^*$ for every decoding round based on the current per-position distributions and the estimated verification cost (which varies with context length).
5. Cost Profiling: The verification cost function $C_v(n; \ell)$ is estimated once offline via convex regression on measured latencies across various context lengths and node counts. This ensures the cost model remains convex and monotonic, satisfying the unimodality guarantee.

Key Contributions

• Objective Shift: Proposes direct throughput optimization as the primary criterion for draft tree construction, replacing the misaligned acceptance-length maximization.
• Theoretical Guarantee: Proves that the throughput objective is unimodal under convex verification costs, enabling a provably optimal greedy stopping rule without additional training.
• Adaptive Mechanism: Introduces a method that dynamically selects the optimal tree size per round, eliminating the need for manual or offline budget tuning.
• Implementation: Instantiates the framework using block diffusion drafters, leveraging their single-pass marginal distributions to build trees efficiently.

Experimental Results
The method was evaluated on Qwen3-4B and Qwen3-8B across eight benchmarks covering reasoning (MATH-500, GSM8K, AIME 2025), coding (HumanEval, MBPP, LiveCodeBench), and instruction following (MT-Bench, Alpaca).

• Performance: CaDDTree matches or surpasses DDTree with oracle budget selection (the best possible fixed budget found via grid search) on nearly all tasks.
• Efficiency: The method achieves lower per-token generation times (higher throughput) compared to autoregressive baselines and fixed-budget DDTree. For instance, on Qwen3-8B with MATH-500, CaDDTree achieves a per-token time of 4.60 ms compared to 5.83 ms for the autoregressive baseline.
• Adaptivity: Analysis shows that the optimal budget varies significantly by round. On reasoning tasks like MATH-500, the selected budget fluctuates widely (standard deviation $\approx 42$) depending on drafter confidence, often deviating significantly from the fixed oracle budget. On instruction-following tasks, the budget is more stable but still distinct from fixed grid points.
• Robustness: Ablation studies indicate the method is robust to perturbations in cost parameter estimates, suggesting precise offline calibration is not strictly necessary.

Significance
The paper claims that CaDDTree provides a principled, cost-aware approach to speculative decoding that resolves the tension between tree coverage and verification latency. By proving the unimodality of the throughput objective, the authors enable an efficient, training-free algorithm that automatically adapts to the specific difficulty of each decoding step. This eliminates the need for heuristic budget selection and demonstrates that dynamic, cost-aware tree construction can consistently outperform static, oracle-tuned baselines across diverse tasks and model sizes.