Rank-Aware Speculative Sampling for Diffusion Draft Trees
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.
摘要： 推测采样（Speculative sampling）通过并行验证低成本的草稿状态，在保持目标分布的同时加速了扩散模型的生成过程。正如扩散贪婪拒绝采样（D-GRS）所展示的那样，近期基于树的方法比单链草稿能更有效地分配并行计算预算。D-GRS 在每个节点生成 $K$ 个条件独立的候选样本，并按生成顺序依次进行测试。然而，这些采样出的候选样本本身就包含有价值的排序信息，且无需额外的目标模型评估。
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.
为了利用这一点，我们引入了秩感知推测采样（Rank-Aware Speculative Sampling, RASS），这是一种基于秩感知列表耦合（rank-aware list coupling）的推测草稿树验证规则。RASS 按照提议分布与目标分布之间的均值偏移对草稿候选样本进行排序，并根据优化后的权重对秩进行采样，以最小化所选提议分布与目标分布之间的总变差（Total Variation）。最后，将所选候选样本与目标进行最大耦合，并通过残差校正确保在任何秩权重选择下都能实现精确采样。
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.
我们在高斯混合目标、FFHQ 上的无条件像素空间生成、CIFAR-10 上的条件生成，以及使用 COCO2014 提示词的 Stable Diffusion 3.5 潜在扩散模型上对 RASS 进行了评估。通过比较标准采样与推测采样的目标模型评估次数之比，结果显示 RASS 在所有评估设置中均优于 D-GRS，在计算预算相同的情况下，其在 CIFAR-10 上的性能提升达到了约 20%。