RASS: How Rank-Aware Speculative Sampling Speeds Up Diffusion Models by 20%

A Fresh Approach to Accelerating Diffusion Models: Starting with Speculative Sampling

Diffusion models have achieved remarkable results in image generation, but their iterative denoising process demands a large number of steps, keeping inference costs stubbornly high. Accelerating this process without sacrificing generation quality is one of the central research challenges today.

Speculative sampling offers an appealing solution. The idea is to use a cheap "draft" model to generate multiple candidate states in parallel, then have the target model verify those candidates in parallel. With the right verification mechanism, this can significantly reduce expensive target model calls while preserving the target distribution exactly. This mirrors speculative decoding in large language models — both use a "guess and verify" strategy to amortize computation.

A new paper on arXiv, Rank-Aware Speculative Sampling for Diffusion Draft Trees (RASS), proposes improvements in this direction and demonstrates measurable gains across multiple benchmarks.

Speculative sampling was first systematically introduced for LLM inference: a small "draft model" generates several tokens sequentially, and a large "target model" verifies all candidates in parallel in a single pass. A carefully designed acceptance-rejection scheme guarantees that the final output distribution is identical to what the target model would produce independently. The key insight is that the target model's parallel verification of a batch of candidate tokens costs about the same as verifying a single token — so as long as the draft model's "hit rate" is high enough, overall throughput improves substantially. Adapting this framework to diffusion models introduces a new challenge: diffusion models operate over continuous high-dimensional distributions rather than discrete vocabularies, so the acceptance-rejection rules must be redesigned to handle comparisons and corrections between continuous probability densities.

From Single-Chain Drafts to Draft Trees: The Limitations of D-GRS

Early speculative sampling methods used a single-chain draft structure — generating candidates along one path and verifying them sequentially. More recent tree-based methods allocate the parallel compute budget more effectively: rather than betting on a single chain, they branch out at each node into multiple candidates.

The representative prior work cited in the paper is Diffusion Greedy Rejection Sampling (D-GRS). D-GRS generates $K$ conditionally independent candidates at each node and tests them one by one in the order they were generated. This design leverages the parallel compute budget and is more efficient than single-chain drafts.

However, the authors identify a neglected source of information: the $K$ sampled candidates implicitly carry a valuable "ranking" relationship, and extracting this ranking requires no additional target model calls. D-GRS simply tests candidates in generation order, effectively discarding this free information.

The draft tree approach stems from a reallocation of parallel compute resources: given a fixed budget, it is better to branch out horizontally at a shallow node than to unroll a single path deeply. This way, even if the acceptance rate on any individual path is low, the probability that at least one path in the tree is accepted increases substantially. D-GRS operates in this framework — it independently samples $K$ candidates at each node and tests them with the acceptance probability of the target distribution, keeping the first one that passes and discarding the rest. But "testing in generation order" means the first candidate tested is not necessarily the highest-quality one; the ordering itself introduces wasted randomness.

The Core Mechanism of RASS: Rank-Aware List Coupling

To address this, the paper proposes Rank-Aware Speculative Sampling (RASS) — a verification rule for speculative draft trees based on "rank-aware list coupling." Its core logic can be broken down into three steps:

1. Sort Candidates by Mean Displacement

RASS first sorts draft candidates by their "proposal-target mean displacement." This ranking captures how far each candidate deviates from the target distribution in terms of direction — the concrete realization of the "free information" mentioned above.

2. Sample a Rank Position Using Optimized Weights

Next, RASS samples a rank position using a set of optimized weights. The optimization objective is explicit: minimize the total variation (TV) distance between the selected proposal distribution and the target distribution. In other words, it proactively selects candidates more likely to be accepted by the target model, boosting the verification pass rate.

3. Maximal Coupling with Residual Correction

Finally, the selected candidate is maximally coupled with the target distribution, with a residual correction applied. This step is critical: it guarantees that regardless of how the rank weights are chosen, the entire process constitutes exact sampling — the final samples strictly follow the target distribution with no bias introduced.

The elegance of this design lies in the decoupling of efficiency and correctness: the rank weights are free parameters that can be optimized for throughput, while correctness is guaranteed by the coupling and residual correction — the two concerns are independent.

Experimental Validation: Consistent Speedups Across Settings

The paper evaluates RASS on four types of tasks, spanning the full spectrum from toy problems to real large-scale models:

• Gaussian mixture target distributions: validating theoretical behavior on controllable distributions;
• Unconditional pixel-space generation on FFHQ: a face image generation benchmark;
• Conditional generation on CIFAR-10: class-conditional image generation;
• Latent-space diffusion with Stable Diffusion 3.5: using COCO2014 prompts, directly applied to a mainstream large-scale text-to-image model.

The evaluation metric is the ratio of target model calls required by standard sampling versus speculative sampling — i.e., how many expensive target model evaluations are saved to produce one sample. RASS outperforms D-GRS across all evaluated settings, with speedups of approximately 20% on CIFAR-10 under matched compute budgets.

Validation on a production-scale latent diffusion model like Stable Diffusion 3.5 demonstrates that the method is not merely a theoretical toy — it has real potential for deployment in practical generation pipelines.

Latent diffusion models are the backbone of today's leading text-to-image systems: rather than running the diffusion process directly in pixel space, they use an encoder to compress images into a low-dimensional latent space, perform denoising iterations there, then decode back to pixels. This design reduces the computational cost of diffusion by orders of magnitude, making high-resolution image generation feasible. The Stable Diffusion series is the canonical example of this architecture, and version 3.5 further introduces a Transformer backbone (DiT structure) with notable gains in both parameter scale and generation quality. Validating RASS on such a model means handling more complex conditioning (text prompts), longer sampling chains, and continuous-distribution coupling in latent space — all of which speak to the method's practical relevance.

Significance and Outlook

RASS derives its value from a nearly free insight — the ranking information already implicit in the draft candidates — and converts it into meaningful inference speedup without sacrificing sampling accuracy. This class of acceleration technique, which modifies only the verification rule rather than the model itself, is especially friendly to large pre-trained diffusion models: no retraining required.

Of course, the 20% speedup is achieved under specific conditions (CIFAR-10, matched budget); gains vary across other settings, and the paper does not report a single unified number across all configurations. The overhead of rank weight optimization, as well as performance at higher resolutions and longer sampling chains, remains to be clarified in future work. But as a refined improvement to the speculative draft tree paradigm, RASS demonstrates the potential of "squeezing every bit of existing information" — a principle worth building on.

Background: Maximal Coupling

Maximal coupling is a classical construction in probability theory: given two probability distributions $P$ and $Q$, a maximal coupling is a joint distribution that maximizes the probability of the two random variables being equal, achieving exactly $1 - d_{TV}(P, Q)$ — the complement of the total variation distance, which is the theoretical upper bound on this "overlap probability." In speculative sampling, maximal coupling works as follows: when a draft candidate falls in a high-probability region of the target distribution, it is directly "accepted" as a target sample; otherwise, a new sample is drawn from the residual distribution (the normalized difference between the target and proposal distributions). This three-step accept-reject-correct procedure mathematically guarantees that the marginal distribution of the output is identical to direct sampling from the target distribution, introducing no systematic bias — hence the term "exact sampling."