SIRS: Single-Query Efficient Rejection Sampling for Masked Diffusion Models

Masked Diffusion Models (MDMs) accelerate generation by unmasking multiple tokens in parallel, but this independent sampling approach ignores correlations between tokens, causing the output distribution to deviate. Anari et al. (2026) used rejection sampling to correct this bias by decomposing the log acceptance probability into a sum of token-level terms—yet this introduced a new cost: every unmasked token requires a score query. At high levels of parallelism, this overhead accumulates rapidly.

To address this tension, recent research proposes the Single-Increment Rejection Sampler (SIRS), which aims to estimate the log acceptance probability with a single query, thereby dramatically reducing query cost while maintaining unbiasedness.

Where Does the Bias in Parallel Masked Diffusion Come From?

The generation process of a masked diffusion model can be understood as follows: starting from a fully masked sequence, each round unmasks a batch of tokens in parallel until all positions are filled. This "intra-block parallelism" strategy significantly reduces the number of forward passes, but the cost is that sampling assumes tokens within a block are mutually independent.

In real data distributions, tokens within the same block often have strong dependencies—such as the syntactic and semantic constraints between adjacent words in natural language. Ignoring these dependencies is equivalent to approximating the joint distribution with the product of marginal distributions, which inevitably introduces bias in the output. Anari et al.'s work theoretically characterizes this bias and provides a correction scheme based on rejection sampling: decomposing the joint acceptance probability into a sum of per-token contributions, then evaluating them one by one.

Masked Diffusion Models (MDMs) are a variant of diffusion models on discrete sequences. Unlike continuous-space diffusion, which adds and removes noise gradually, MDM's "noising" process progressively replaces tokens in a sequence with a special [MASK] token. The training objective is to teach the model to predict the original tokens at masked positions given the currently partially visible sequence. At generation time, the process starts from a fully masked sequence and repeatedly calls the model to predict conditional distributions at each position and sample fills—this process is called "unmasking." Since the conditional distribution at each position can be computed simultaneously, MDMs naturally support parallel decoding, in stark contrast to autoregressive models that generate tokens one at a time.

Score query in the MDM context refers to calling the model's forward pass to obtain the conditional distribution or log probability at each masked position given the currently unmasked context. Rejection sampling needs to know the ratio between the proposal distribution and the target distribution, so each evaluation of the acceptance probability requires a model query. Anari et al.'s scheme decomposes the block-level acceptance probability into a sum of token-level terms, each depending on the model's output at that position—so the number of tokens to be unmasked in a block determines the number of queries required. When blocks are large, the query count grows linearly, becoming the primary bottleneck.

SIRS's Core Idea: Estimating the Entire Increment with One Query

The key insight of SIRS is: rather than querying the score separately for each token in a block, randomly sample a single increment and use the result of that one query to estimate the log acceptance probability for the entire block. Specifically, each forward pass of the model returns the conditional distributions of all remaining tokens in the block. SIRS uses this information to perform Rao-Blackwellization on the increment—averaging over tokens in the block without requiring additional queries, thereby reducing the estimation variance.

Variance reduction directly yields a practical benefit: less clipping. In rejection sampling, if the acceptance probability estimate fluctuates too much, more aggressive clipping is needed to ensure numerical stability—and clipping introduces additional distribution bias. SIRS keeps clipping within a smaller range through more stable estimates.

SIRS then leverages the Bernoulli factory proposed by Chen et al. (2026) to convert a Poisson-distributed clipped estimate into per-token accept/reject decisions. This conversion guarantees the statistical correctness of the sampling process, rather than being a simple heuristic approximation.

Rao-Blackwellization is a classic variance reduction technique named after statisticians C. R. Rao and David Blackwell. The basic idea is: if an estimator depends on an auxiliary random variable over which one can take conditional expectations, then replacing the original estimator with the conditional expectation of that auxiliary variable given a sufficient statistic can reduce variance without introducing bias. Formally, for any random variables X and Y, Var(E[X|Y]) ≤ Var(X)—taking the conditional expectation over Y can only make the variance smaller or unchanged.

In SIRS, the "increment" is the randomly sampled auxiliary variable, and the conditional distributions of tokens in the block returned by the model's forward pass provide exactly the information needed to take the conditional expectation over the increment. Thus SIRS can analytically average over the increment without querying each token individually, obtaining a lower-variance acceptance probability estimate. This explains why a single query suffices for a sufficiently stable estimate: variance reduction comes not from increasing sample size, but from analytically eliminating part of the randomness.

Bernoulli factory is an algorithmic framework for converting random samples from one probability distribution into samples from another. In rejection sampling, we need to generate a Bernoulli decision that accepts with probability p and rejects with probability 1-p based on an estimated acceptance probability p—but p is often an estimate rather than an exact value, and may appear in some transformed form. A Bernoulli factory can construct a Bernoulli distribution exactly parameterized by p using only "black-box" queries about p (such as determining whether p exceeds a threshold), without needing to know p's analytical form.

In SIRS, the clipped acceptance probability estimate follows a Poisson distribution and cannot be used directly as an acceptance probability. The Bernoulli factory of Chen et al. (2026) happens to handle exactly this type of input, converting it precisely into per-token accept/reject decisions. This step is key to sampling unbiasedness: if the estimate were simply used directly as the acceptance probability, the resulting distribution would have additional bias; after conversion through the Bernoulli factory, as long as the estimator itself is unbiased, the final acceptance decisions are statistically correct.

Theoretical Guarantee: Bounded KL Divergence

SIRS provides an explicit error bound: the KL divergence between the data distribution and SIRS's output distribution is bounded by a clipping term plus the model's training error.

The significance of this bound is that it decouples sampling error from the model's own quality—the extra bias introduced by the sampling algorithm depends only on the clipping range, and does not grow uncontrollably with block size or parallelism. The paper further compares two cases:

• With exact scores, parallel decoding must pay the full cost of intra-block dependencies;
• SIRS only pays the dependency cost that exceeds the clipping range.

In other words, as long as the clipping range is set appropriately, SIRS's actual cost can be far lower than naive parallel rejection sampling.

KL divergence (Kullback-Leibler divergence) is an asymmetric measure of the difference between two probability distributions, defined as D_KL(P||Q) = Σ_x P(x) log(P(x)/Q(x)). It is always non-negative and equals zero if and only if the two distributions are almost everywhere equal. In error analysis of generative models, KL divergence is often used to characterize the gap between the sampling algorithm's output distribution and the true data distribution: the smaller the divergence, the closer the sampling result is to the target distribution in the information-theoretic sense.

The bound provided by SIRS decomposes the KL divergence into two parts: one arising from the bias introduced by the clipping operation, and the other from the training error of the model's score estimate itself. The value of this decomposition lies in showing that the extra error introduced by the sampling algorithm is independent of hyperparameters like block size and parallelism, and is controlled only by the clipping threshold. Therefore, one can push the sampling error arbitrarily small by tightening the clipping range without sacrificing parallel efficiency—at the cost of more rejections and resampling, forming a controllable precision-efficiency trade-off.

Implications for Accelerating Masked Diffusion

Masked diffusion models have attracted attention in recent years for tasks such as text generation and code completion, with their parallel decoding capability being the main efficiency advantage over autoregressive models. But the higher the parallelism, the more severely intra-block dependencies are ignored, and the sharper the quality-speed trade-off becomes.

SIRS's value lies in not avoiding this trade-off, but rather using more refined statistical tools to push the cost close to the theoretical lower bound. The single-query estimation approach decouples query cost from block size, while Rao-Blackwellization further compresses variance at no additional cost. For application scenarios requiring large-batch sampling during inference, this improvement in query efficiency may be more practically valuable than simply increasing parallelism.

Of course, SIRS's effectiveness depends on the quality of score estimates and the choice of clipping threshold. How to tune these in real-world deployment remains an engineering question to be answered.

One-line Summary

SIRS achieves efficient unbiased rejection sampling for masked diffusion models with a single query plus Rao-Blackwellization.

Paragraph Summary

Masked diffusion models assume intra-block independence when unmasking tokens in parallel, causing the output distribution to deviate from the true joint distribution. Anari et al.'s rejection sampling correction eliminates this bias but requires one score query per token—a prohibitive overhead at high parallelism. SIRS proposes a single random increment query to estimate the entire block's log acceptance probability, using the conditional distributions returned by the model's forward pass for Rao-Blackwellization—reducing variance and clipping at zero extra query cost. It then uses a Bernoulli factory to convert the Poisson-distributed clipped estimate into per-token accept/reject decisions, guaranteeing statistical correctness. The theory proves that the KL divergence between SIRS's output and the data distribution is bounded by the clipping term plus training error, and that with exact scores, only the dependency cost beyond the clipping range needs to be paid.