Query-Efficient Rejection Sampling for Parallel Masked Diffusion

Shiyuan Zhang, Xuheng Li, Qiwei Di and Quanquan Gu

arXiv: 2610.22238v1 · PDF · 15 October 2026
Categories: cs.LG, math.PR, stat.ML
Licence: CC BY 4.0 The site shows the text.

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.

Items 2

Notations

F7Y7

We use normal-face letters for scalars, boldface lower-case letters for vectors, boldface upper-case letters for matrices, and italic boldface upper-case letters for random vectors. For N\in\mathbb{N}_+, [N] denotes the set \{1, 2, \dots, N\}. The clipping function is denoted by \operatorname{clip}_B(x)=\min(B, \max(-B, x)). We use (z)_+ and (z)_- to denote \max(z, 0) and \min(z, 0), respectively. We use \operatorname{KL}(\cdot\|\cdot) to denote the KL-divergence, D_{\mathrm{H}}(\cdot, \cdot) to denote the Hellinger distance, and \chi^2(\cdot\|\cdot) to denote the \chi^2-distance.

Setup

I2BH

All samplers reveal blocks of a uniform size, taken as consecutive chunks of a uniformly random permutation of [L] drawn independently of the tokens. We run SIRS-RB with block size n and B\in\{\tfrac18, \tfrac14, \tfrac12, 1, 2\}, and compare it with parallel sampling with block size n' and with sequential decoding (n'=1), which is exact. For each sampler we compute the KL-divergence \operatorname{KL}(p^*\|\widehat p_B) from the data distribution to the output distribution, in nats per sequence and averaged over permutations, and the expected number of score queries per sequence. Parallel sampling uses \lceil L/n'\rceil queries; in SIRS-RB we flip the coins of the Poissonized Bernoulli factory one at a time and reject at the first failure, which leaves the output distribution unchanged and saves queries.

Source and licence

The text on this page is from arXiv:2610.22238v1 by Shiyuan Zhang, Xuheng Li, Qiwei Di and Quanquan Gu, under the licence CC BY 4.0. The corpus made these changes: it read the items out of the LaTeX source, it expanded the macros of the paper, it put each proof after its statement, and it converted the text to HTML. The formulas are the TeX of the paper, and KaTeX renders them. A formula or a reference can look different from the PDF. The page is published under CC BY 4.0.

The record of each item is in analysis/objects/2610/2610.22238.jsonl.