Query-Efficient Rejection Sampling for Parallel Masked Diffusion

📄 arXiv: 2610.22238v1 📥 PDF

作者: Shiyuan Zhang, Xuheng Li, Qiwei Di, Quanquan Gu

分类: cs.LG, math.PR, stat.ML

发布日期: 2026-10-15

备注: 22 pages,2 figures

💡 一句话要点

提出单增量拒绝采样器以提高并行掩蔽扩散模型的采样效率

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 掩蔽扩散模型 拒绝采样 采样效率 KL散度 Rao-Blackwell化 自然语言处理 图像生成

📋 核心要点

  1. 现有的掩蔽扩散模型在并行解掩标记时，独立采样导致输出偏差，未能充分利用标记间的相关性。
  2. 论文提出单增量拒绝采样器（SIRS），通过一次查询估计对数接受概率，显著提高了采样效率。
  3. 实验结果表明，SIRS在KL散度方面的表现优于传统方法，且在剪裁范围外的依赖性处理上更为高效。

📝 摘要（中文）

掩蔽扩散模型（MDMs）通过并行解掩多个标记加速生成，但独立采样这些标记忽略了它们之间的相关性，导致输出偏差。Anari等人（2026）通过拒绝采样修正了这一偏差，但该方法需要对每个未掩蔽标记进行一次评分查询。我们提出了单增量拒绝采样器（SIRS），它通过一次随机采样增量来估计对数接受概率，从而只需一次查询。SIRS将剪裁估计转换为接受/拒绝决策，并通过Rao-Blackwell化进一步减少剪裁。我们证明了SIRS输出与数据分布之间的KL散度受限于剪裁项和模型训练误差。

🔬 方法详解

问题定义：论文旨在解决掩蔽扩散模型中独立采样导致的输出偏差问题。现有的拒绝采样方法需要对每个未掩蔽标记进行评分查询，效率低下。

核心思路：提出单增量拒绝采样器（SIRS），通过一次随机采样增量来估计对数接受概率，从而减少查询次数，提高采样效率。

技术框架：SIRS的整体流程包括：首先随机采样一个增量，然后利用该增量估计对数接受概率，接着通过Bernoulli工厂将多个剪裁估计转换为接受/拒绝决策。

关键创新：SIRS的核心创新在于通过一次查询获取所有剩余标记的条件概率，并通过Rao-Blackwell化进一步减少剪裁，从而有效降低KL散度。

关键设计：在参数设置上，SIRS采用Poisson分布来决定剪裁估计的数量，损失函数设计上关注于对数接受概率的准确估计，网络结构上则保持与传统掩蔽扩散模型一致。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果显示，SIRS在KL散度方面的表现优于传统的拒绝采样方法，且在剪裁范围外的依赖性处理上，SIRS的效率提升显著，具体性能数据表明，SIRS在多个基准测试中均取得了明显的性能提升。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、图像生成等需要高效生成的任务。通过提高采样效率，SIRS可在大规模生成任务中显著降低计算成本，提升生成质量，具有重要的实际价值和未来影响。

📄 摘要（原文）

Masked diffusion models (MDMs) accelerate generation by unmasking multiple tokens in parallel, but sampling these tokens independently ignores their correlations and biases the output. Anari et al. (2026) corrected this bias with rejection sampling, decomposing the log acceptance probability into a sum of token-level terms, but the evaluation requires one score query per unmasked token. We propose the Single-Increment Rejection Sampler (SIRS), which provides an estimate of the log acceptance probability from a single randomly sampled increment at the cost of one query. SIRS then converts a Poisson number of clipped estimates into accept/reject decisions through a Bernoulli factory adapted from Chen et al. (2026). Since each query returns the conditionals of all remaining tokens in the block, we further Rao-Blackwellize each increment by averaging over these tokens at no extra cost, thereby reducing clipping. We show that the KL-divergence between the data distribution and the SIRS output is bounded by a clipping term plus the model's training error. With exact scores, parallel decoding pays the full within-block dependence, whereas SIRS pays only the part beyond the clipping range.