Accelerating Diffusion Sampling via Speculative Draft Trees

1 Introduction

• Diffusion inference repeatedly evaluates a large neural network over tens to hundreds of denoising iterations, making image generation slow.Parallel simulation reduces sequential depth but typically needs multiple refinement rounds and substantial parallel resources.Source: §1 · pp. 1–2
• Speculative diffusion sampling drafts inexpensive states for target verification, allowing accepted drafts to skip denoising iterations while preserving the target distribution exactly.Unlike approximate acceleration techniques, it incurs no loss in image quality.Source: §1 · p. 2
• Existing speculative samplers form linear proposal chains, a budget allocation that can underuse parallelism, particularly when acceptance is low.The broader allocation problem is how to spend a fixed amount of parallel computation to maximize draft acceptance.Source: §1 · pp. 2–3
• Up to 7.4% speedups over existing methods are reported for practical batch sizes ≤256, while preserving exact sampling.Figure 1 separately reports up to a 14% reduction in target model evaluations relative to RMC.Source: Figure 1 · p. 2, §1 · p. 3
• The paper connects speculative sampling to relative entropy coding and introduces draft trees with greedy rejection sampling for richer compute allocation.The authors report consistent improvement over existing methods while preserving exact sampling.Source: §1 · p. 3

2 Problem Setting

• The reverse-time process includes a churn parameter ε that controls stochasticity, with ε = 1 recovering the standard reverse-time SDE and ε = 0 yielding the deterministic probability-flow ODE.Euler–Maruyama discretization on a uniform time grid produces the sampled reverse process.Source: §2 · p. 3
• Speculative sampling uses an inexpensive proposal model to draft multiple states before verifying them in parallel with the target reverse drift.The method assumes drafting is significantly cheaper than verification; one batched target evaluation can verify all drafts within compute budget B.Source: §2 · p. 4
• The proposal can use delayed reverse drift, replacing the current target transition mean with one based on an earlier state, and this paper adopts it throughout.The cited prior work reports higher-quality proposals from delayed reverse drift than from lightweight draft models.Source: §2 · p. 4

3 Speculative Algorithms

• Speculative diffusion sampling drafts candidates according to a topology, verifies them in parallel with the target model, then examines verified states to accept a proposal or sample the next target state.The algorithms differ primarily in draft topology and acceptance criterion.Source: §3 · pp. 4–5

3.1 A General Template For Speculative Diffusion

• The general template varies only the draft topology and verification rule, and assumes isotropic Gaussian draft and target transitions with a shared variance schedule and different means.The template is agnostic to how proposals are produced.Source: §3.1 · p. 5

3.1 Draft Topology.

• A draft topology is a rooted tree whose nodes represent candidate states at process steps, with each node conditioned on exactly one parent.The framework covers chains, depth-dependent widths, and irregular or pruned trees; verification may evaluate all nodes or only internal nodes.Source: §3.1 · pp. 5–6

3.1 Round Structure.

• Each round expands the tree by depth, batches target evaluations for internal nodes, then traverses downward applying verification until the first rejection.If no child qualifies, the verification rule returns an exact target sample from the normalized residual distribution.Source: §3.1 · p. 7

3.1 Rank-1 Projection.

• For isotropic Gaussian proposal and target transitions sharing covariance, exact sampling reduces to transforming the scalar projection along the mean displacement while retaining the orthogonal component.The projection is standard normal under the proposal and shifted by the normalized mean displacement under the target.Source: §3.1 · pp. 7–8

3.2 Reflection Maximal Coupling (RMC)

• RMC uses a chain and reflection maximal coupling; early rejection makes downstream drafts unusable, wasting the fixed parallel compute budget.On rejection, it reflects the scalar proposal about the midpoint between proposal and target means to sample from the residual distribution.Source: §3.2 · pp. 8–9, Figure 2 · p. 9

3.3 Diffusion Greedy Rejection Sampling (D-GRS)

• D-GRS replaces the chain with a K-ary draft tree and adapts greedy rejection sampling to the resulting list-coupling problem.At each node, it examines children in order, greedily assigns remaining target mass, and samples from the residual if no child is accepted.Source: §3.3 · pp. 9–12
• List coupling is computationally costly for continuous diffusion transitions, so D-GRS uses an order-preserving sequence coupling motivated by relative entropy coding.The method processes proposals in sampled order and selects the first accepted proposal, preserving exactness through residual sampling when none is accepted.Source: §3.3 · pp. 11–12

4 Performance Analysis

• D-GRS outputs Y distributed according to the target Q.The analysis identifies this as the exactness result for Algorithm 3.Source: §4 · p. 12
• Proposal and normalized residual branches together recover Q exactly.The proof calculates target mass from the first K proposal branches and the contribution from a sample drawn from the normalized residual distribution.Source: §4 · p. 12
• The probability of accepting at least one of K proposals is given as 1 − G_{K+1}.The passage introduces this as Algorithm 3's probability of accepting at least one drafted child.Source: §4 · p. 12
• RMC achieves the optimal acceptance probability for maximal coupling, and Figure 6 plots per-step acceptance probabilities for different K values.The analysis compares D-GRS with the RMC baseline.Source: §4 · pp. 12–13

5 Experiments

• Across experiments, D-GRS and RMC are compared at equal target-call cost, with each D-GRS configuration matched to an RMC chain of length B.The experiments cover a Gaussian-mixture target, pixel-space CIFAR-10 and FFHQ generation, and Stable Diffusion 3 latent-space sampling.Source: §5 · p. 13

5.1 Gaussian Mixture Target

• The Gaussian-mixture evaluation uses a five-component isotropic Gaussian target in 500 dimensions, 30 generation steps, and ε = 0.1.It sweeps L ∈ {2, …, 7} and K ∈ {1, …, 7}, with speedups measured against the standard sampler at (K, L) = (1, 1).Source: §5.1 · p. 13
• At B = 155, D-GRS achieves 2.14× versus RMC’s 1.96× speedup, reducing target calls by 8–9% without compromising exactness.For L ≥ 3, D-GRS strictly dominates; at (K, L) = (6, 6), it reaches 2.33× versus RMC’s 1.99×.Source: §5.1 · p. 14

5.2 Pixel-Space Diffusion

• On CIFAR-10, RMC has a marginal advantage at B = 14, but for B ≥ 62 its speedup plateaus while D-GRS speedups continue increasing.The pixel-space comparison covers CIFAR-10 and FFHQ across stochasticity levels ε ∈ {0.1, 0.3, 0.6}.Source: §5.2 · p. 15, Figure 4 · p. 15
• On FFHQ, D-GRS strictly dominates RMC at every tested budget and churn level, despite lower absolute speedups than on CIFAR-10.FFHQ speedups are approximately 2.0–2.2×, versus 2.2–2.6× on CIFAR-10; higher churn degrades both methods across the pixel-space tasks.Source: §5.2 · p. 15

5.3 High-dimensional Latent Space Diffusion

• For SD3 at L = 2, D-GRS reaches nearly 1.18× speedup at K = 6 (B = 42), a 6.7% increment relative to RMC.Across the parameter sweep, RMC plateaus between 1.10× and 1.12×, while D-GRS speedup increases monotonically and peaks near 1.18×.Source: Figure 5 · p. 16, §5.3 · p. 17

6 Conclusions and Future Work

• Up to 8.3%: D-GRS improves NFE speed-up over RMC in practical settings by adapting greedy rejection sampling to draft-tree list coupling.The conclusion identifies RMC’s allocation of its speculative budget to a single draft chain as the topological limitation addressed.Source: §6 · p. 17
• Exact-coupling efficiency rapidly degrades with generation dimension, and adapting draft topology online to proposal quality and compute budget remains an open direction.The authors describe this work as a first step toward more efficient speculative diffusion sampling.Source: §6 · p. 17

A.1 Diffusion Models

• Diffusion formulations including DDPMs, score networks, SDEs, and flow matching are mathematically equivalent under transformations of output parameterization and noise schedule.The passage notes that training objectives do not uniquely determine the sampling algorithm: denoising-trained models can use probability-flow ODE sampling, while flow-matching models can use stochastic schemes.Source: §A.1 · p. 26
• Diffusion acceleration includes training-free solvers and trained few-step models, including SDE methods that learn marginal transition kernels or pathwise solution maps.The passage distinguishes training supervision and matching objectives among training-based methods.Source: §A.1 · pp. 26–27

A.2 Speculative Sampling for Language Models

• Speculative decoding uses a smaller draft model to propose future tokens for parallel target verification, with rejection correction preserving the target distribution exactly.The passage also describes tree-based inference and verification, multi-draft architectures, and list coupling as developments in this paradigm.Source: §A.2 · p. 27

A.3 Relative Entropy Coding

• Relative entropy coding is an exact remote-sampling framework in which a sender with target Q communicates an index to a receiver sharing proposal P; its communication cost is governed by D_KL(Q∥P).The passage names this family channel simulation, channel synthesis, and reverse channel coding in related literature.Source: §A.3 · p. 27
• The paper employs greedy rejection sampling; Poisson-process schemes are another prominent REC class, while inexact channel-simulation approaches fall outside standard REC.The passage says those approximate approaches fall short of the theoretical relative-entropy limit.Source: §A.3 · pp. 27–28

• The appendix restates the Section 4 theoretical guarantees before giving their proofs.Source: B On Performance Analysis · p. 28
• Throughout the proofs, proposal P and target Q are Gaussian distributions on R^d with shared covariance σ²I_d, and Y is the proposal selected by Algorithm 3 at index J.Source: B On Performance Analysis · p. 28
• Φ denotes the standard Gaussian survival function, and (x)+ denotes max{x, 0}.Source: B On Performance Analysis · p. 28

• Algorithm 3 outputs Y distributed exactly as the Gaussian target Q.Source: §B.1 · p. 28
• The proof invokes the residual-measure characterization of greedy rejection sampling to track target mass remaining after the first j proposals.Source: §B.1 · p. 28
• Combining the intermediate identities establishes that Y ∼ Q exactly.Source: §B.1 · p. 28

• With Δ = (μq − μp)/σ and δ = ||Δ||, the proof standardizes the Gaussian laws of s to compute the superlevel-set masses; setting j = K completes the proof.Source: §B.2 · pp. 29–30
• P(∃k ∈ {1, …, K}: Y = Ŷ_k) = 1 − G_{K+1}; Figure 6 plots per-step D-GRS acceptance for different K.Source: §B.2 · p. 29, Figure 6 · p. 31
• The proof evaluates G_{K+1} through the probability of reaching proposal j and the P- and Q-masses of superlevel sets of the density ratio.The integrand vanishes outside the relevant superlevel set.Source: §B.2 · p. 29
• For equal-covariance Gaussian P and Q, the log target-to-proposal ratio is a monotone function of scalar s, and its superlevel sets are half-spaces {s ≥ τ_j}, with τ_j = ln λ_j.Source: §B.2 · pp. 29–30
• When K = 1, the method recovers RMC and its acceptance probability matches De Bortoli et al. (2025).Source: §B.2 · pp. 30–31
• If δ = 0, P = Q and the first proposal is accepted almost surely; the argument also extends to any common positive-definite covariance Σ.Source: §B.2 · p. 31

• The section provides additional details about the experimental settings.Source: C Extended Experimental Results · p. 31

C.1 Implementation Details

• The delayed-drift proposal uses the target drift evaluated at the round’s root as the proposal mean at every lookahead depth.Source: §C.1 · p. 31
• Figure 7 plots Gaussian-mixture NFE speedup against tree width K, with separate D-GRS lines for lookahead depths L = 2 to 7.Points average per-replicate speedups over 100 independent trials; shaded bands show ±1 standard error.Source: Figure 7 · p. 32
• Root-drift prefetching reuses an already computed drift instead of making an extra target evaluation during drafting.The delayed-drift proposal otherwise requires one extra target call per round to evaluate the target drift at the root.Source: §C.1 · pp. 31–32

C.2 The Value of K and L

• The chain baseline saturates near 2× because geometric acceptance decay, rather than draft budget, limits its committed steps per target evaluation.With measured α ≈ 0.52, the predicted ceiling is 2.09×, close to the observed 2×.Source: §C.2 · pp. 32–33
• 2.34× is the best speedup at (K, L) = (7, 6), with widening from K = 5 to K = 7 at L = 6 adding 0.09×.D-GRS surpasses RMC from (K, L) = (2, 3) onward; at L = 2 its best configuration reaches only 1.876×.Source: §C.2 · p. 33
• The relative value of depth L and width K depends on the per-step acceptance probability, controlled by the ratio δ.Depth is limiting below a problem-dependent crossover; above it, depth’s marginal benefit is exhausted around L ≈ 5, while added width can still help.Source: §C.2 · p. 33

C.3 SD3 Conversion

• A conversion without retraining expresses SD3 as a reverse-time SDE whose Gaussian transition has state-independent variance, the structure required by the rank-1 reduction.This conversion lets both RMC and D-GRS apply to SD3 without modification.Source: §C.3 · pp. 33, 35
• At ε = 0, SD3’s transitions degenerate to point masses with total variation distance one, so every proposal is rejected and speculation reduces to standard sampling.Source: §C.3 · p. 35

C.4 Qualitative Assessment of Correctness

• The qualitative assessment displays generated CIFAR-10 and FFHQ samples and SD3 samples conditioned on COCO 2014 prompts.Figures 8 and 9 show CIFAR-10 and FFHQ samples; Figure 10 shows SD3 samples under the reported setting, with prompts listed in Table 2.Source: §C.4 · p. 35, Table 2 · p. 36, Figure 8 · p. 37