Query-Efficient Rejection Sampling for Parallel Masked Diffusion
Abstract
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.
Related
- Serial Depth of Adaptive Parallel Sampling: Lower Bounds and Schedules for Masked Diffusion Models
- Lower Bounds for Parallel Diffusion Sampling
- Structured Masked Diffusion for Shared-Codebook Unsourced Random Access Decoding
- A Uniformly Efficient Rejection Sampler for the Multivariate Expectile-Based Distribution
- Representation-based Masked Diffusion Model
BibTeX
@misc{zhang2026queryefficient,
title = {{Query-Efficient Rejection Sampling for Parallel Masked Diffusion}},
author = {Shiyuan Zhang and Xuheng Li and Qiwei Di and Quanquan Gu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.22238},
eprint = {2610.22238},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.22238},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Query-Efficient Rejection Sampling for Parallel Masked Diffusion
AU - Shiyuan Zhang
AU - Xuheng Li
AU - Qiwei Di
AU - Quanquan Gu
PY - 2026
DA - 2026/10/15
PB - arXiv
DO - 10.48550/arxiv.2610.22238
UR - https://arxiv.org/abs/2610.22238
AB - 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.
ER -