Rank-Aware Speculative Sampling for Diffusion Draft Trees

2026-09-30

AI Summary

This study addresses the inefficient allocation of parallel computation budgets in draft trees for speculative sampling in diffusion models. To this end, it proposes a verification mechanism based on rank-aware list coupling, termed RASS. This method leverages candidate ranking information to optimize weight selection, achieving exact coupling with the target distribution without incurring additional evaluation overhead. By integrating tree-structured draft generation, total variation minimization, and maximal coupling theory, RASS significantly outperforms D-GRS across multiple benchmarks. Notably, under an equivalent computational budget on CIFAR-10, it improves the acceleration ratio by approximately 20%, effectively enhancing the inference efficiency of diffusion models.

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.

Problem

• Speculative Sampling
• Diffusion Models
• Draft Trees
• Candidate Ranking

Innovation

• Speculative Sampling
• Diffusion Models
• Draft Trees
• Rank-Aware Coupling
• Exact Sampling