PaperlayerView exact revision on arXiv

Source-linked AI summary

Rank-Aware Speculative Sampling for Diffusion Draft Trees

Marcello Bullo, Yanxiao Liu, Öykü Sıla Güner, Arpan Mukherjee, Deniz Gündüz

arXiv:2610.02251v1cs.LGstat.CO

TL;DR

Speculative diffusion sampling seeks to reduce the sequential target-model evaluations required for generation while preserving the target law, but D-GRS tests exchangeable draft candidates in generation order. RASS ranks candidates and maximally couples the selected proposal with the target; across evaluated settings it improves target-call speedup over D-GRS, with gains reaching 20.9% on CIFAR-10 at matched budgets.

  • Problem

    D-GRS tests exchangeable, conditionally i.i.d. draft candidates in generation order, imposing a sequence constraint that can be suboptimal relative to list coupling.

    Source: §1 · p. 1

  • Method

    RASS ranks draft proposals by their target-to-proposal likelihood ratios, selects a rank with discrepancy-dependent weights, and maximally couples the selected proposal with the target.

    Source: §1 · p. 2, §3 · p. 5

  • Results

    RASS consistently improves mean target-call speedup over D-GRS across the evaluated settings, reaching 20.9% improvement on CIFAR-10 at ε = 0.6 and B = 63.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    RASS preserves the target law for every feasible rank-weight choice while improving mean target-call speedup over D-GRS at matched parallel budgets.

    Source: §1 · p. 2, §7 · p. 9

    Target-call speedups do not by themselves establish wall-clock acceleration, which depends on hardware capable of efficiently executing larger parallel batches, drafting, and verification.

    Source: §6 · p. 9

Abstract

from arXiv · show

Speculative sampling accelerates diffusion generation by verifying inexpensive draft states in parallel while preserving the target law. Recent tree-based methods allocate the parallel compute budget more effectively than single-chain drafts, as demonstrated by Diffusion Greedy Rejection Sampling (D-GRS). D-GRS generates $K$ conditionally independent candidates per node, and sequentially tests them in their generation order. Yet the sampled candidates admit an informative ranking without additional target-model evaluations. To exploit this, we introduce Rank-Aware Speculative Sampling (RASS), a verification rule for speculative draft trees based on rank-aware list coupling. RASS orders draft candidates along the proposal-target mean displacement and samples a rank with weights optimized to minimize total variation between the selected-proposal and target laws. Finally, the selected candidate is maximally coupled with the target, with residual correction ensuring exact sampling for any choice of rank weights. We evaluate RASS on a Gaussian-mixture target, unconditional pixel-space generation on FFHQ, conditional generation on CIFAR-10, and latent diffusion with Stable Diffusion 3.5 using COCO2014 prompts. Measured by the ratio of standard to speculative sampling's target-model evaluation counts, RASS improves on D-GRS across the evaluated settings, with gains reaching approximately 20% on CIFAR-10 at matched compute budgets.

1 Introduction

Speculative diffusion sampling reduces sequential target evaluations through parallel verification, but D-GRS imposes a drafting-order sequence on exchangeable candidates. RASS replaces that constraint with rank-aware verification, preserving exactness while improving target-call speedup across the reported settings.

1 Introduction

  • Speculative diffusion sampling uses cheaper drafted proposals, corrected to exact target samples, to reduce expensive target-model evaluations and inference delay.Its speed-up depends on how effectively parallel compute is allocated.Source: §1 · p. 1
  • D-GRS tests conditionally i.i.d. draft-tree children in generation order, even though their exchangeability leaves that ordering unnecessary and potentially suboptimal.The sequence constraint can discard exploitable symmetry relative to list coupling.Source: §1 · p. 1
  • RASS ranks i.i.d. proposals by target-to-proposal likelihood ratios and uses the ranking to construct a proposal distribution for verification.The rank distribution interpolates between uniform selection for close proposal and target distributions and greater preference for the highest-likelihood proposal as discrepancy grows.Source: §1 · p. 2
  • The paper derives a universal upper bound on acceptance for list coupling when drafts are i.i.d., tightening bounds that allow arbitrary dependence.This bound provides a benchmark for exact verifiers under the independent drafting constraint.Source: §1 · p. 2
  • RASS preserves the target law for every feasible rank-weight choice and improves mean target-call speedup across the reported Gaussian-mixture, CIFAR-10, FFHQ, and Stable Diffusion 3.5 experiments.Gains reach 20.9% at matched draft trees and parallel budgets.Source: §1 · p. 2

2 Background

Speculative diffusion sampling replaces a fully sequential target trajectory with inexpensive proposal drafting, batched target evaluation, and exact verification. Draft trees allocate parallel compute across candidate children, while verification follows an accepted branch and corrects at the first rejection.

2 Background

  • A standard N-step diffusion trajectory requires N sequential neural-network evaluations, motivating methods that reduce sequential target dependence.The reverse process is simulated with Euler–Maruyama using the learned target drift.Source: §2 · p. 3
  • A speculative round drafts candidate future states with a cheaper proposal, evaluates target drifts in parallel, then accepts or corrects candidates while preserving the target Markov-chain law.Figure 2 depicts drafting, batched verification evaluations, and correction at the first rejection.Source: §2 · p. 3, Figure 2 · p. 3
  • Drafting can be organized as a rooted tree whose non-root nodes are candidate states, with each child sampled from a proposal transition conditional on its parent.A chain is the branching-factor-one special case, and same-parent children are conditionally i.i.d. proposals.Source: §2 · p. 4
  • Full-tree verification uses B = |V| and evaluates leaves in advance, allowing their target drifts to be reused after a fully rejected tree; the internal-node alternative uses B = |I|.The main experiments use the full-tree budget; the trade-off between these strategies is problem dependent.Source: §2 · p. 4
  • At a visited node, verification couples i.i.d. Gaussian proposal children to the Gaussian target, accepting a child to continue or returning a residual correction to end the round.For isotropic Gaussians with shared covariance, verification reduces to the scalar projection along the proposal-to-target mean displacement.Source: §2 · p. 4

3 Rank-Aware Speculative Sampling

RASS selects a draft by rank along the proposal-to-target mean displacement, then maximally couples that selected proposal with the target. It chooses rank weights to improve acceptance while residual correction preserves exact sampling.

3 Rank-Aware Speculative Sampling

  • Candidate ranks follow their projections along the proposal-to-target mean displacement, with rank K denoting the child having the largest projection toward the target mean.The rank-weight vector lies in the probability simplex and depends on displacement δ and draft count K.Source: §3 · p. 5
  • Uniform rank weights recover the original proposal, whereas always selecting the top rank can overshoot the target and reduce overlap when displacement is small.RASS therefore adaptively interpolates between uniform and top-rank selection.Source: §3 · p. 5
  • RASS samples a rank using weights ω, selects the corresponding ranked child, and maximally couples its induced proposal distribution with the target.The selected rank distribution is optimized to make the selected-child law close to the target, maximizing maximal-coupling acceptance.Source: §3 · pp. 5–6
  • On rejection, RASS draws from the normalized positive residual between the target density and selected-proposal density, preserving exactness.The selected rank weights determine the selected-proposal density and its acceptance probability.Source: §3 · pp. 5–6

4 Theoretical Analysis

RASS is exact for every feasible rank distribution, and the analysis gives an upper bound on acceptance for any verifier using independent Gaussian drafts. Comparisons show higher RASS acceptance than D-GRS across tested positive displacements, with the largest gains at intermediate displacement.

4 Theoretical Analysis

  • RASS is exact for every K ≥ 1, δ ≥ 0, and feasible rank-weight choice.This allows acceptance optimization without changing the target law.Source: §4 · p. 6
  • The optimal list-acceptance bound applies to independent Gaussian drafts and is tighter than a bound allowing arbitrary proposal dependence.A verifier can return a draft above a threshold only if the draft list contains one; any shortfall must be supplied through rejection and correction.Source: §4 · p. 6
  • RASS has higher acceptance than D-GRS across tested positive mean displacements, with the largest gains at intermediate δ.The figure compares acceptance against the optimal i.i.d.-list bound as well as D-GRS.Source: §4 · pp. 6–7, Figure 3 · p. 7

5 Experiments

RASS improves mean target-call speedup over D-GRS across the evaluated Gaussian-mixture, pixel-space, and latent-diffusion settings. The experiments also compare sample quality and show that gains depend on the task and budget.

Experiments

  • Across the Gaussian-mixture, pixel-space, and SD3.5 experiments, RASS consistently improved mean target-call speedup over D-GRS at matched parallel budgets.RASS and D-GRS used identical tree configurations; all three methods preserved the discretized target law, with comparable FID and CLIP scores reported.Source: §5 · pp. 7–8
  • The experiments measured mean per-trajectory target-call speedup, using 100 trajectories per Gaussian-mixture and SD3.5 configuration and 500 for CIFAR-10 and FFHQ.The full tree, including leaves, defined the parallel evaluation budget.Source: §5 · p. 7, Figure 4 · p. 8, Figure 5 · p. 8

Gaussian-mixture target

  • Across the tested Gaussian-mixture budgets, RASS achieved higher mean speedup than D-GRS, with a 6.4% higher maximum.Within the tested grid, RASS exceeded D-GRS’s highest observed mean speedup using 28% of its budget.Source: §5 · p. 7

Pixel-space diffusion

  • 20.9% relative improvement over D-GRS on CIFAR-10 at ε = 0.6 and B = 63; RASS improved speedup in every tested pixel-space configuration.On FFHQ, improvements were 1.4% to 4.6%.Source: §5 · p. 8, Figure 4 · p. 8

High-dimensional Latent-space Text-to-image

  • For SD3.5, RASS reached 1.152× mean speedup versus 1.142× for D-GRS and 1.105× for RMC.At matched configurations, its relative improvement over D-GRS reached 1.93% at (K, L) = (5, 3), or B = 156.Source: §5 · p. 8, Figure 5 · p. 8

6 Limitations

Reported speedups measure reductions in sequential target-model evaluations at matched parallel budgets, not guaranteed wall-clock acceleration. Direct applicability is also constrained by the Gaussian draft and covariance assumptions.

6 Limitations

  • Target-call savings do not guarantee wall-clock gains, which depend on hardware handling parallel batches, drafting, and verification; RASS directly applies to stochastic samplers with conditionally independent Gaussian drafts sharing the target’s isotropic covariance.The paper identifies full-pipeline optimization and extensions to more general proposal and covariance families as future work.Source: §6 · p. 9

7 Conclusion and Future Work

RASS is a rank-aware verification rule for speculative diffusion draft trees, improving target-call speedup while preserving the target law. The conclusion identifies limits of rank mixtures and systems integration, while the related work and algorithm sections place RASS within existing sampling approaches and describe its pluggable verifier setting.

7 Conclusion and Future Work

  • Across Gaussian-mixture, pixel-space, and latent text-to-image experiments, RASS consistently improves mean target-call speedup over D-GRS and RMC at matched parallel budgets.RASS preserves the target law for every feasible rank-weight choice, while optimized rank mixtures remain separated from the optimal list-coupling bound.Source: §7 · p. 9

A Related Work

  • RASS is a verification rule for draft-tree speculative sampling; unlike approaches that change the sampling distribution or require training, it is orthogonal to the proposal used to draft.In continuous, high-dimensional Gaussian diffusion, rank-1 structure reduces list coupling to a scalar, and rank selection with optimized weights becomes a K-dimensional LP independent of d.Source: A Related Work · p. 13

B Speculative Sampling Algorithm

  • The draft-tree sampler supports pluggable VERIFY rules, and exactness at each node lets verification continue recursively after an accepted child.The evaluation set determines the parallel budget; the main experiments evaluate the full tree, while an internal-node alternative is also described.Source: B Speculative Sampling Algorithm · p. 13

C.1 Proof of Proposition 1

The proposition characterizes the selected child’s projection and orthogonal component under rank selection. Theorem 2 establishes that RASS preserves the target law for any feasible rank weights.

C.1 Proof of Proposition 1

  • The selected child’s orthogonal component has law N(0, Id − ee⊤) conditional on the projections and rank draw.The selected projection’s density is described using the density of a Beta(r, K + 1 − r) random variable.Source: §C.1 · p. 14

C.2 Proof of Theorem 2

  • For every K ≥ 1, δ ≥ 0, and feasible rank weights ω ∈ WK, the RASS verifier is exact.At each visited node, the next committed state has the target transition law; induction over committed steps gives the target Markov-chain law.Source: §C.2 · p. 15

C.3 Proof of Proposition 3

Proposition 3 characterizes the highest acceptance probability possible when coupling a target with a list of independent drafts. In the Gaussian setting, the proof reduces the problem to one dimension and shows the bound is attainable.

C.3 Proof of Proposition 3

  • For independent uniform drafts and target law ν, the maximum target-preserving acceptance probability is 1 − d(ν), and the bound is attainable.The proof constructs finite couplings of draft and target bins, then takes a weak limit to obtain a coupling of the original variables.Source: §C.3 · pp. 16–17
  • For equal-covariance Gaussian proposal and target laws, the acceptance problem reduces to scalar projections along the direction between their means.The proof applies the probability integral transform to express the scalar draft coordinates as uniform variables and the target coordinate through its CDF.Source: §C.3 · p. 16
  • A weak-limit construction attains the resulting Gaussian acceptance bound, completing the proof.The limiting coupling preserves both the independent uniform draft-list marginal and the target marginal.Source: §C.3 · pp. 17–18
  • Because the transformed Gaussian target density is increasing, an upper interval maximizes target mass among sets of any fixed Lebesgue measure.The proof then optimizes over the interval threshold to obtain the maximal achievable acceptance probability.Source: §C.3 · p. 18

C.4 Proof of Theorem 4

Theorem 4 expresses acceptance for a rank-weighted proposal through the overlap of its scalar density with the target density. The proof derives the residual region and compares uniform and top-rank choices; the implementation section introduces their computation.

C.4 Proof of Theorem 4

  • For fixed rank weights, the acceptance region consists of points where the target scalar density exceeds the selected-proposal density, and it is a union of disjoint open intervals.The acceptance contribution is summed over those intervals using their residual mass.Source: §C.4 · pp. 18–20
  • With uniform rank weights, the equally weighted order-statistic mixture is uniform on (0, 1), yielding the standard Gaussian proposal density.The proof uses this density to derive the corresponding rejection probability.Source: §C.4 · p. 20
  • With top-rank weights and K > 1, the proposal and target densities cross at exactly two points, enclosing the interval where the proposal density is larger.The proof establishes the two crossings by showing the log density ratio has one maximum and is negative in the tails.Source: §C.4 · pp. 20–21

D Implementing RASS

  • The implementation section describes computing rank weights and sampling the residual distribution using the scalar densities qδ and pω.These are the routines used in lines 5 and 11 of Algorithm 1.Source: D Implementing RASS · p. 21

D.1 Rank Weights

Rank weights are optimized by maximizing proposal-target overlap with a quadrature-based linear program. Residual correction uses interval masses and inverse-CDF sampling, while numerical crossing detection is not guaranteed to find every interval.

D.1 Rank Weights

  • The rank-weight objective is concave, and its quadrature approximation can be optimized as a linear program over the probability simplex.The program has K + M variables and 2M upper-bound inequalities; its size is independent of state dimension d.Source: §D.1 · pp. 21–22
  • 2.1 × 10−5 was the largest observed quadrature loss, versus a 3.7 × 10−2 gap to optimal list-coupling acceptance in the tested settings.The evaluation covered K ∈ {2, 3, 4, 6, 8} and 14 values of δ from 0.01 to 6.Source: §D.1 · p. 22
  • As δ approaches zero, optimized weights approach uniformity; increasing δ shifts weight toward higher ranks, nearly all to the highest rank beyond the reported thresholds.The near-top-rank thresholds are δ ≳ 2 for K = 4 and δ ≳ 2.5 for K = 8; irregular small-δ changes can reflect nonunique optimal solutions.Source: §D.1 · p. 23

D.2 Residual Sampling

  • The inverse-CDF residual sampler selects an interval and solves for its within-interval draw, using one uniform variable and one scalar root solve without another target-model evaluation.The residual interval masses can be computed without numerical integration.Source: §D.2 · pp. 23–24
  • Residual support can include multiple intervals: for K = 2, ω = (0, 1), and δ = 1, the left interval contains about 0.8% of residual mass.The implementation searches 513 grid points for crossings and refines detected sign changes, but a finite grid alone cannot guarantee every crossing is found.Source: §D.2 · p. 24
  • The corrected state can reuse the first drafted child’s orthogonal component because it remains Gaussian and independent of scalar verification and residual sampling.An independent draw from the orthogonal Gaussian component gives the same output law.Source: §D.2 · p. 25

D.3 RASS Computational Overhead

With cached rank weights and rejection sampling, RASS adds little scalar verification cost relative to D-GRS, while numerical implementation remains subject to error. Delayed-drift proposals reuse cached model evaluations during drafting and select a candidate evaluation for the next round.

D.3 RASS Computational Overhead

  • Rank-weight solves take 5.8–8.0 ms but are cached by (eδ, K); rejection-sampling corrections are needed only with probability ηω.At small δ, an individual correction can require many trials, but its rarity offsets this in expectation; trial time still varies with (δ, K).Source: §D.3 · p. 25
  • With cached weights, RASS rejection-sampling verification takes 30–55 µs and stays within 1.5% of D-GRS for every tested branching factor.D-GRS takes 30–55 µs; RASS inverse-CDF sampling takes 117–168 µs. These are scalar verification timings averaged over tested displacements.Source: §D.3 · pp. 25–26, Table 3 · p. 26

D.4 Exactness Under Numerical Approximation

  • The exactness guarantee applies to the ideal verifier; floating-point arithmetic, residual-interval grid search, and root-solver tolerances can introduce implementation errors.Quadrature and rounding in RANKWEIGHTS do not invalidate the theorem, which covers every feasible rank-weight choice.Source: §D.4 · p. 26

E.1 Delayed-drift proposals and prefetching

  • Delayed-drift drafting holds a cached model field fixed while using each node’s state and schedule; prefetching retains an available evaluation for the next round.With no initial cached evaluation, the sampler obtains one target evaluation before drafting. Internal-node-only evaluation can leave no evaluated candidate at the next root after certain endings.Source: §E.1 · p. 27

E.2 Results with internal-node evaluation

With internal-node or full-tree budgets, RASS has higher mean speedup than D-GRS in every evaluated configuration. Internal-node evaluation reduces batch size but can make the next proposal less accurate; qualitative samples are reported as indistinguishable in quality and content from target samples.

E.2 Results with internal-node evaluation

  • Internal-node evaluation removes leaf evaluations from the batch and prefetch cache, which can reduce next-proposal accuracy after a leaf acceptance or last-level rejection.The trade-off is a smaller batch per round versus a more accurate cache with full-tree evaluation; which budget works better depends on model, lookahead, and hardware batch cost.Source: §E.2 · pp. 28, 30
  • On the Gaussian mixture, internal-node RASS is 0.6%–6.7% faster than D-GRS in all 144 configurations with K ≥2.At fixed budgets B ≤100, its best mean speedup is 1.1%–2.1% higher than with full-tree evaluation; across the full grid, the highest speedups differ by at most 0.6%.Source: §E.2 · pp. 28–29
  • RASS exceeds D-GRS in every evaluated configuration under both budget definitions across the reported settings.The figures cover Gaussian-mixture, pixel-space, and SD3.5 results; Figure 2 in this appendix reports the pixel-space comparison under both budgets.Source: Figure 7 · p. 29, §E.2 · p. 30, Figure 12 · p. 33
  • On CIFAR-10, internal-node RASS at B = 31 is within 2.6% of full-tree RASS at B = 63 and exceeds D-GRS’s highest full-tree speedup at B = 341.Across internal-node configurations, RASS gains 3.9%–18.7% over D-GRS on CIFAR-10 and 1.4%–2.4% on FFHQ.Source: §E.2 · p. 29
  • On SD3.5, the best mean RASS speedup within B ≤43 is 1.152× with internal-node evaluation and 1.151× with full-tree evaluation.Internal-node RASS exceeds D-GRS by 0.2%–1.2% in all ten configurations, but all ten confidence intervals overlap.Source: §E.2 · p. 29, Figure 14 · p. 34

E.3 Qualitatively Assessment of Exactness

  • Across CIFAR-10, FFHQ, and SD3.5, speculative samples are reported as indistinguishable in quality and content from target samples; text-conditioned samples match their prompts.The displayed comparisons use four randomly selected samples per run, with shared class labels on CIFAR-10 and shared prompt and initial noise on SD3.5.Source: §E.3 · p. 31

Related papers

Loading 2610.02251v1…