Provably adaptive sampling with uniform and remasking discrete diffusion models

Daniil Dmitriev; Zhihan Huang; Yuting Wei

cs.LGAug 24, 2026

Provably adaptive sampling with uniform and remasking discrete diffusion models

Authors: Daniil Dmitriev, Zhihan Huang, Yuting Wei

Organizations: Department of Statistics and Data Science, the Wharton School, University of Pennsylvania

Abstract

Discrete diffusion models offer a promising alternative to autoregressive generation by enabling parallel updates, but their sampling efficiency can depend strongly on the choice of the forward process and the sampler. For the uniform forward process, existing lower bounds for the standard $τ$ -leaping sampler scale linearly with the ambient dimension $d$ , raising the question of whether this dependence is intrinsic to the forward process. We answer this question in the negative. We consider a first-order sampler based on the leave-one-out denoiser for uniform and remasking processes whose coordinate updates can be performed in parallel. In both cases, the sampler can correct denoising mistakes during the sampling process, which becomes necessary when many coordinates are updated together. Our main result establishes an adaptive sampling guarantee: up to logarithmic factors, $N = O(\mathrm{DTC}(X_0) / \varepsilon)$ discretization steps suffice to achieve sampling error $O(\varepsilon_{\mathrm{score}}+\varepsilon)$ , where $\varepsilon_{\mathrm{score}}$ is the error in score estimation. Thus, the sampling complexity is governed by the intrinsic dependence structure of the target distribution, as measured by its dual total correlation $\mathrm{DTC}(X_0)$ , rather than directly by the ambient dimension $d$ . Our analysis proceeds through a Bayes-optimal auxiliary sampler that separates discretization error from score-estimation error. We also derive an exact information-theoretic representation of the discretization error in terms of the mutual information between different coordinates of the forward process at different times. This representation applies to general forward processes and, in the uniform and remasking cases, can be controlled by $\mathrm{DTC}(X_0)$ . Numerical experiments on structured synthetic distributions illustrate the predicted dimension-adaptive behavior.

Figures & tables

Figure 1:

N=20

discretization steps. Target distribution is a Markov chain, see Section 5 . Remasking process uses

p_{M}=0.5

. We use early stopping

\delta=1\mathrm{e}{-5}

and time horizon

T=8

. Different types of discretization schedules are adopted for different noising processes to optimize their performance: the masking process uses constant discretization

t_{k+1}-t_{k}=(T-\delta)/N

; the remasking and uniform processes use geometric discretization

t_{k+1}-t_{k}=\kappa\min(1,T-t_{k+1})

with

\kappa\approx 1.3

. Results are averaged over 7 runs.

	True reverse process	Bayes-optimal sampler	Our approximation
Score function	$\displaystyle\frac{\Pr(X_{t}^{i}=b\mid X_{t}^{-i}=x_{t}^{-i})}{\Pr(X_{t}^{i}=x_{t}^{i}\mid X_{t}^{-i}=x_{t}^{-i})}$	$\displaystyle\frac{\Pr(X_{t}^{i}=b\mid X_{\ell}^{-i}=x_{\ell}^{-i})}{\Pr(X_{t}^{i}=x_{t}^{i}\mid X_{\ell}^{-i}=x_{\ell}^{-i})}$	$\displaystyle\frac{\widehat{\Pr}(X_{t}^{i}=b\mid X_{\ell}^{-i}=x_{\ell}^{-i})}{\widehat{\Pr}(X_{t}^{i}=x_{t}^{i}\mid X_{\ell}^{-i}=x_{\ell}^{-i})}$
LOO denoiser	$\Pr(X_{0}^{i}=b\mid X_{t}^{-i}=x_{t}^{-i})$	$\Pr(X_{0}^{i}=b\mid X_{\ell}^{-i}=x_{\ell}^{-i})$	$\widehat{\Pr}(X_{0}^{i}=b\mid X_{\ell}^{-i}=x_{\ell}^{-i})$

Table 1: Here,

\ell=T-t_{k}

is the end point of the

k

-th discretization interval and

T-t_{k+1}\leq t\leq\ell

. The score function row compares the true value

s_{t}(x_{t}\odot_{i}b,x_{t})

, the Bayes-optimal

\widetilde{s}_{t}(x_{t}\odot_{i}b,x_{t})

, and ours

\widehat{s}_{t}(x_{t}\odot_{i}b,x_{t})

. The reverse process conditions on the correct context at time

t

. The Bayes-optimal estimator conditions on all information available at the start of the discretization step, i.e., at time

\ell

. Our approximation is based on the Bayes optimal one, but instead of the true probabilities

\Pr(\cdot)

it uses the estimate

\widehat{\Pr}(\cdot)

, similarly for the leave-one-out denoiser. To compute the probability of

X_{t}^{i}=b

used in the score function, all three approaches use the same formula (see Equations 7 and 9 ) but with different leave-one-out denoisers.

Figure 2:

N=30

discretization steps. Target distribution is the Markov chain with

dp_{\mathrm{flip}}=2

. The left plot shows the performance on the geometric grid

t_{k+1}-t_{k}=\kappa\min(1,T-t_{k+1})

and the right plot shows the result for the constant grid

t_{k+1}-t_{k}=(T-\delta)/N

. We use early stopping at

\delta=1\mathrm{e}{-5}

and

T=8

. Results are averaged over

7

runs. Given the same number of discretization steps

N

, the masking process works better with the constant grid, which agrees with the finding in Dmitriev et al. (2026) for distributions with small

\mathrm{DTC}

. Uniform process, in contrast, works better with the geometric grid, consistent with the finding of this work. Remasking process performs robustly on both grids.

Figure 3:

N=40

discretization steps. Target distribution is the binary Markov chain with

dp_{\mathrm{flip}}=2

. We use early stopping at

\delta=1\mathrm{e}{-5}

and

T=8

. Results are averaged over

7

runs. We compare three different samplers for the uniform process:

\tau

-leaping, truncated

\tau

-leaping, and ours.

Figure 4:

N=20

discretization steps. Dimension

d=2000

. Target distribution is a sparse mixture with

k

components. We use early stopping at

\delta=1\mathrm{e}{-5}

and

T=8

. Results are averaged over

7

runs. Masking uses constant grid; uniform and remasking use geometric grid.

Diffusion Model Sampling

Diffusion model sampling studies how to generate high-quality samples efficiently by designing reverse processes, solvers, and inference-time controls.

See all 82 papers

Feb 16, 2026

Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

Diffusion models over discrete spaces have recently shown striking empirical success, yet their theoretical foundations remain incomplete. In this paper, we study the sampling efficiency of score-based discrete diffusion models under a continuous-time Markov chain (CTMC) formulation, with a focus on

τ

-leaping-based samplers. We establish sharp convergence guarantees for attaining

\varepsilon

accuracy in Kullback-Leibler (KL) divergence for both uniform and masking noising processes. For uniform discrete diffusion, we show that the

τ

-leaping algorithm achieves an iteration complexity of order

\tilde O(d/\varepsilon)

, with

d

the ambient dimension of the target distribution, eliminating linear dependence on the vocabulary size

S

and improving existing bounds by a factor of

d

; moreover, we establish a matching algorithmic lower bound showing that linear dependence on the ambient dimension is unavoidable in general. For masking discrete diffusion, we introduce a modified

τ

-leaping sampler whose convergence rate is governed by an intrinsic information-theoretic quantity, termed the effective total correlation, which is bounded by

d \log S

but can be sublinear or even constant for structured data. As a consequence, the sampler provably adapts to low-dimensional structure without prior knowledge or algorithmic modification, yielding sublinear convergence rates for various practical examples (such as hidden Markov models, image data, and random graphs). Our analysis requires no boundedness or smoothness assumptions on the score estimator beyond control of the score entropy loss.

Daniil Dmitriev, Zhihan Huang, Yuting Wei

Department of Statistics and Data Science, the Wharton School, University of Pennsylvania

Sample Complexity Diffusion Model Sampling Diffusion Sampling Diffusion Models

Ask PDF

Apr 12, 2026

Query Lower Bounds for Diffusion Sampling

Diffusion models generate samples by iteratively querying learned score estimates. A rapidly growing literature focuses on accelerating sampling by minimizing the number of score evaluations, yet the information-theoretic limits of such acceleration remain unclear. In this work, we establish the first score query lower bounds for diffusion sampling. We prove that for

d

-dimensional distributions, given access to score estimates with polynomial accuracy

\varepsilon=d^{-O(1)}

(in any

L^p

sense), any sampling algorithm requires

\widetildeΩ(\sqrt{d})

adaptive score queries. In particular, our proof shows that, within any polynomial total-query budget, successful sampling requires searching over

\widetildeΩ(\sqrt{d})

distinct noise levels, providing a formal explanation for why multiscale noise schedules are necessary in practice.

Zhiyang Xun, Eric Price

UT Austin

Query Complexity Diffusion Sampling Information-Theoretic Lower Bounds Diffusion Models

Ask PDF

Jan 22, 2025

Low-dimensional adaptation of diffusion models: Convergence in total variation

This paper investigates how diffusion generative models leverage (unknown) low-dimensional structure to accelerate sampling. Focusing on two mainstream samplers -- the denoising diffusion implicit model (DDIM) and the denoising diffusion probabilistic model (DDPM), we prove that their iteration complexities under exact score functions are at most the order of

k/\varepsilon

(up to log factor), where

\varepsilon

is the precision in total variation distance and

k

is some intrinsic dimension of the target distribution. We further extend these convergence guarantees to the setting in which the score functions are learned from data rather than known exactly, showing that the convergence performance degrades gracefully under suitable score estimation assumptions. We then show that these assumptions are attainable via kernel-based score estimators with finite-sample guarantees that also adapt to the low-dimensional structure. Our results apply to a broad family of target distributions without requiring smoothness or log-concavity. Our findings provide the first rigorous evidence for the adaptivity of the DDIM-type samplers to unknown low-dimensional structure, and improve over the state-of-the-art DDPM theory regarding total variation convergence.

Jiadong Liang, Zhihan Huang, Yuxin Chen

Department of Statistics and Data Science, the Wharton School, University of Pennsylvania

Diffusion Sampling Diffusion Models Diffusion Model Adaptation

Ask PDF

Abstract

Diffusion Model Sampling

Explore similar work

Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

Query Lower Bounds for Diffusion Sampling

Low-dimensional adaptation of diffusion models: Convergence in total variation