Rank-Aware Speculative Sampling for Diffusion Draft Trees
Abstract
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.
Related
- Accelerating Diffusion Sampling via Speculative Draft Trees
- Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
- TAPS: Target-Aware Permanent Sampling for Graph Diffusion
- Multiple Importance Sampling for Guided Speculative Inference
- Curvature Diffusion of Inverse-weight Lin--Lu--Yau Ricci Flow on Finite Trees
BibTeX
@misc{bullo2026rankaware,
title = {{Rank-Aware Speculative Sampling for Diffusion Draft Trees}},
author = {Marcello Bullo and Yanxiao Liu and Öykü Sıla Güner and Arpan Mukherjee and Deniz Gündüz},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.02251},
eprint = {2610.02251},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.02251},
}
RIS
TY - UNPB
TI - Rank-Aware Speculative Sampling for Diffusion Draft Trees
AU - Marcello Bullo
AU - Yanxiao Liu
AU - Öykü Sıla Güner
AU - Arpan Mukherjee
AU - Deniz Gündüz
PY - 2026
DA - 2026/10/01
PB - arXiv
DO - 10.48550/arxiv.2610.02251
UR - https://arxiv.org/abs/2610.02251
AB - 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.
ER -