论文

RASS:扩散草稿树的秩感知推测采样

Rank-Aware Speculative Sampling for Diffusion Draft Trees

模型推理投机采样

摘要

推测采样通过并行验证廉价的草稿状态在保持目标分布的同时加速扩散生成。近期树形方法比单链草稿更有效地分配并行算力,以扩散贪心拒绝采样(D-GRS)为代表:每个节点生成K个条件独立候选并按生成顺序依次检验。但这些候选本身携带无需额外目标模型评估即可获得的排序信息。为此我们提出秩感知推测采样(RASS),一种基于秩感知列表耦合的草稿树验证规则:按提议-目标平均位移对候选排序,再以优化选出的权重采样秩,使被选提议与目标分布间的总变差最小;被选候选与目标最大耦合,残差校正保证对任意秩权重都精确采样。我们在高斯混合目标、FFHQ无条件像素空间生成、CIFAR-10条件生成及使用COCO2014提示的Stable Diffusion 3.5潜扩散上评估。以标准采样与推测采样的目标模型评估次数之比衡量,RASS在所有评估设置上优于D-GRS,匹配算力下CIFAR-10增益约20%。