Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

Daniil Dmitriev; Zhihan Huang; Yuting Wei

Source-linked AI summary

Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

Daniil Dmitriev, Zhihan Huang, Yuting Wei

arXiv:2602.15008v2Feb 16, 2026cs.LGcs.ITmath.STstat.ML

TL;DR

Existing analyses can scale poorly with both vocabulary size and dimension, motivating sharper sampling-efficiency guarantees for discrete diffusion. The paper analyzes τ-leaping samplers under CTMC dynamics and obtains dimension-adaptive results for uniform and masking noising processes, alongside scoped lower bounds.

Problem
Existing theoretical analyses suggest convergence rates scale at least linearly with vocabulary size S and ambient dimension d, motivating sharper guarantees for discrete-diffusion sampling.
Method
The paper develops sharp convergence analyses for τ-leaping and modified τ-leaping strategies under uniform and masking noising processes in a CTMC formulation.
Results
For uniform diffusion, τ-leaping requires tilde O(d/ε) iterations for ε KL accuracy, while masking diffusion admits bounds governed by target-distribution structure that can scale logarithmically with d.
Takeaways & Limitations
The results establish an algorithmic lower bound for τ-leaping and show that masking-based sampling can adapt to intrinsic data structure without boundedness assumptions on the score estimator.
Takeaways & Limitations
The lower bound is specific to τ-leaping and existing path-measure-based analysis techniques; whether modified algorithms improve uniform diffusion remains open.

Abstract

from arXiv · show

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.

1 Introduction

Discrete diffusion models extend diffusion-based generation to categorical, language, and graph-structured data, but their sampling efficiency remains theoretically challenging. This work develops sharp CTMC-based guarantees for τ-leaping samplers, including dimension-optimal uniform diffusion rates and structure-adaptive masking diffusion rates.

Background: Discrete diffusion models use noisy forward and reverse processes, with CTMC formulations pairing rate matrices and learned score functions.Uniform and absorbing rate matrices are prominent choices, and sampler performance depends sensitively on the rate matrix.
Motivation: Existing analyses typically scale linearly with vocabulary size S and dimension d, making sampling costly for high-dimensional applications.GPT-2-based tasks illustrate the issue with S = 50,257 and d = 102–10^3.
Contributions: The paper establishes sharp convergence guarantees for uniform and masking discrete diffusions, addressing whether sampling can be efficient or automatically adapt to structure.The analysis centers on τ-leaping strategies under the CTMC formulation.
Uniform diffusion: Uniform diffusion requires only Õ(d/ε) discretization steps for ε-error in KL divergence, improving the prior Õ(d^2S/ε) bound by removing vocabulary dependence and one factor of d.The paper also establishes a matching algorithmic lower bound showing linear d dependence is essentially unavoidable for τ-leaping.
Masking diffusion: Masking diffusion uses a refined τ-leaping sampler governed by effective total correlation D, which is bounded by d log S but can be sublinear or O(1) for structured data.The method allows automatic adaptation to low-dimensional target distributions without prior structural knowledge.
Analysis: The guarantees require only control of score entropy loss rather than boundedness or smoothness assumptions on the score estimator.The analysis uses a Girsanov change of measure and martingale properties to separate approximation and discretization errors.

2 Preliminaries of discrete diffusion models

Discrete diffusion models are formulated with forward and reverse CTMCs, whose reverse dynamics are approximated using learned score functions and discretized samplers. The section defines uniform and masking noising processes and introduces τ-bridging strategies, which preserve parallel updates while enabling CTMC-based analysis.

2.1 A continuous-time Markov chain formulation: A discrete diffusion model represents d-dimensional data with a forward CTMC that progressively transforms the data distribution toward an easy-to-sample distribution.The process is defined through rate matrices and one-step transition probabilities in the continuous-time limit.
2.1 A continuous-time Markov chain formulation: Uniform and masking noising processes differ in their limiting distributions: uniform diffusion approaches the uniform measure, whereas masking diffusion approaches an all-MASK Dirac measure.The initial data distribution is supported on non-masked states for the masking process.
2.1 A continuous-time Markov chain formulation: The reverse CTMC is determined by the forward rates and marginal distributions, with density ratios serving as discrete score functions.The reverse rates are related to the forward rates through the time-reversed marginals, while the score generalizes ∇x log q_t(x) from continuous diffusion.
2.2 Score estimation: Because exact marginals and scores are generally intractable, practical samplers learn score estimates using score entropy loss and assume control of aggregated estimation error.The estimates are available only at discrete time points, so continuous dynamics must also be discretized.
2.3 Score-based sampling algorithms: τ-leaping updates coordinates independently and in parallel, while τ-bridging strategies restrict updates to CTMC-based transitions for analytical tractability.The paper’s τ-bridging class retains parallel coordinate updates and evolves d independent one-dimensional CTMCs over each discretization interval.
2.3 Score-based sampling algorithms: The modified scheme enforces at most one transition per coordinate per discretization interval and achieves sublinear complexity for masking diffusion under mild distributional assumptions.The paper presents this as the first such guarantee for the masking noising process.

3 Main results

The paper develops sharp KL-divergence guarantees for τ-bridging samplers under uniform and masking discrete diffusion. Uniform diffusion achieves dimension-linear complexity with a matching τ-leaping lower bound, while masking diffusion adapts to intrinsic target structure.

Overview: τ-bridging strategies preserve parallel coordinate updates while enabling tractable analysis for uniform and masking noising processes.Each update evolves d independent one-dimensional CTMCs over one discretization interval.
Uniform diffusion: The uniform-diffusion KL bound combines score-estimation, initialization, and discretization errors, with discretization scaling linearly in d and logarithmically in S.Initialization error decays exponentially with the diffusion horizon, while score error accumulates across discretization steps.
Uniform diffusion: d/ε steps suffice for uniform diffusion up to logarithmic factors, improving the previous d^2S/ε complexity.The bound is for ε-accurate sampling in KL divergence.
Uniform diffusion: The uniform τ-leaping complexity cannot generally improve below linear dependence on d for informative target distributions.This lower bound applies uniformly to early-stopping and non-early-stopping schedules.
Uniform diffusion: The lower bound is algorithm-dependent, while analogous improvements for uniform diffusion using modified samplers remain open.It does not constitute an information-theoretic or minimax lower bound.
Masking diffusion: For masking diffusion, the convergence bound depends on effective total correlation, which is bounded by d log S and can remain small for structured distributions.The resulting sampler adapts to target structure without prior knowledge or algorithmic modification.

15 end

The masking sampler corrects a score-induced rate-matrix mismatch and obtains complexity controlled by intrinsic dependence measures. These measures yield sublinear guarantees for several structured distributions, including hidden Markov models, quantized low-dimensional data, random graphs, and latent parity models.

Adaptive masking sampler: Algorithm 1 rescales the score estimate to correct the mismatch between surrogate and true rate matrices in masking diffusion.This rescaling produces a sharper upper bound than truncated τ-leaping.
Adaptive masking sampler: The refined sampler implicitly exploits information-theoretic structure without additional hyperparameters.The paper connects this automatic adaptation to the fundamental role of total-correlation quantities.
Complexity guarantees: Under a constant step size, iteration complexity is controlled by dual total correlation; under an exponential-then-constant schedule, effective total correlation governs it.Both schedules achieve KL error at most ε_score + ε under the stated horizon and conditions.
Examples: For hidden Markov models, dual total correlation grows sublinearly with sequence length when topics evolve slowly.This gives a sublinear sampling-complexity guarantee in that structured setting.
Examples: For quantized continuous data with intrinsic dimension k, dual total correlation scales linearly with k rather than ambient dimension d.This example is motivated by image generation.
Examples: For sparse random regular graphs and stochastic block models, dual total correlation grows at most linearly in n up to logarithmic factors rather than quadratically.The graph distributions are over binary edge variables.
Examples: In the latent parity model, total correlation and dual total correlation are both order d, while effective total correlation remains constant-order.This supports constant-order complexity for a distribution with highly structured latent signal.

4 Discussion

The paper establishes sharp guarantees for uniform and masking discrete diffusion samplers. It also identifies adaptive masking samplers and leaves noising-process selection and score learning as open questions.

Uniform τ-leaping requires rO(d/ε) iterations for ε-accurate KL sampling.This improves prior bounds, while a matching lower bound applies to a large class of distributions.
Masking diffusion admits an upper bound governed by the target distribution’s intrinsic complexity.The bound can scale logarithmically with ambient dimension through a modified τ-leaping strategy.
Both analyses require only control of score entropy loss, without score-estimator boundedness or early stopping.
The modified masking sampler permits parallel coordinate updates while preserving CTMC dynamics.
Choosing suitable noising mechanisms and learning accurate discrete score functions remain open research questions.

A.1 Details and formal results

The formal examples show how effective total correlation yields sublinear sampling complexity for structured distributions. HMMs, images, and random graphs can therefore have complexity below their ambient dimension.

Hidden Markov models: For HMMs with rare latent-state transitions, sampling complexity scales with intrinsic topic dimension rather than document length.The latent space is modeled as [S]^k with k much smaller than d.
Low-dimensional structures: Low-dimensional image structure yields effective total correlation O(k), which is o(d) when latent dimension k is much smaller than observation dimension d.The construction uses a continuous latent variable, Lipschitz decoding, additive Gaussian noise, and entrywise quantization.
Low-dimensional structures: The resulting image distributions can be sampled efficiently despite high-dimensional observation spaces.
Random graph models: Generalized random graph models remain subquadratic when latent variables have sublinear complexity relative to n^2.Discrete latent variables with o(n) entropy yield o(n^2) dual total correlation, while continuous latent dimension d_z=o(n) also gives sublinear rates.
Effective total correlation: Effective total correlation can remain small even when total correlation and dual total correlation are individually large.A mixed distribution demonstrates D(p_ex) much smaller than both B(p_ex) and C(p_ex).

A.2 Proofs of results in Section A.1

The proofs bound effective total correlation using latent-variable structure and conditional independence. For HMMs, this reduces the complexity to the entropy of the latent Markov chain.

Conditional independence of coordinates given W implies B(X)≤I(X;W), and additionally B(X)≤H(W) when W is discrete.
The proof uses entropy inequalities and conditional independence rather than direct control of the full observed distribution.
For the HMM, conditional independence of observations given latent states bounds B(q_data) by the latent-chain entropy H(z).
The latent-chain entropy is controlled by the initial-state entropy and transition probability under the stated Markov assumptions.

A.2.2 Proof of Proposition 2

The proof of Proposition 2 bounds image-distribution complexity through mutual information between continuous latent variables and noisy decoded observations. Quantization preserves the relevant conditional-independence structure.

Entrywise quantization preserves coordinate conditional independence given the latent variable z.This lets Lemma 1 transfer a bound on I(x;z) to the effective total correlation B(q_data).
The proof bounds B(q_data) by I(x_con;z) using data processing through the quantization map.
Partitioning the latent space into hypercubes of size σ/L controls decoder variation using Lipschitzness.The argument combines data processing, conditional mutual information, Gaussian-channel bounds, and hypercube diameter estimates.
The final bound follows by controlling within-cell decoder variation and the number of latent-space cells.
The proof combines the intermediate estimates to obtain the stated effective-complexity bound for quantized image data.

A.2.3 Proof of Proposition 3

The proposition bounds information quantities for random regular graphs and block stochastic block models. Random regular graphs have an entropy-based upper bound of order kn log n, while block models have subquadratic dependence on n.

The set of k-regular graphs is considered under the parity condition that nk is even.
O(kn log n) bounds the entropy of a uniformly sampled k-regular graph through B(G) ≤ H(G) ≤ log(|G_n,k|).
Conditional on the latent block assignments, graph edges are sampled independently, supporting the entropy bound for the stochastic block model.
n log n bounds the block-model quantity because B(G) ≤ H(z_1,...,z_n) ≤ n log r ≤ n log n.
The block-model setting is also a special case of a generalized random geometric graph model under analogous assumptions.

A.2.5 Proof of Proposition 5

The proof analyzes a structured binary distribution by decomposing its entropy and conditional entropy, then bounds the resulting information quantities and score-related integrals. Under a balanced coordinate partition, these quantities scale as stated below.

The balance assumption min{r,1-r}=Θ(1) keeps both coordinate groups proportional to d.
log(2)|I_1| lower bounds the entropy H(x), while log(2)(|I_1|-1) appears in the corresponding auxiliary entropy bound.
d log 2 equals H(x_i)+H(i*), since the distinguished coordinate is Bernoulli(1/2) and the remaining coordinates contribute binary entropy.
0 is the conditional entropy H(x_i|x_-i), because each coordinate can be recovered from the others in the structured distribution.
Ω(d) lower bounds H(x_i)-H(x), yielding a linear-in-d information quantity under the partition structure.
For uniform and masking noising, Proposition 6 gives equivalent score-function formulations based on their respective transition structures.

C.1 Proof of Theorem 1

The proof decomposes sampling KL error into initialization, discretization, and score-estimation contributions. It controls these terms using CTMC reversibility identities, entropy inequalities, and bounds on score magnitudes and temporal variation.

The output KL divergence is decomposed into an initialization term and accumulated conditional transition discrepancies along the discretized reverse process.
e^-T d log(S) bounds the initialization contribution for uniform noising, reflecting convergence of the forward process toward the uniform distribution.
The score entropy loss enters directly when bounding one discretization-error term in the KL decomposition.
The logarithmic score bound |log s_t(y,x)| ≲ log(S)+max{log(t^-1),0} controls score magnitudes for neighboring states.
The function φ(t) is non-increasing, allowing temporal error terms to be bounded through its monotonicity and a maximum step size Δ.
Choosing T=log(d log(S)/ε) and applying the theorem yields KL(q_data||q_output) ≲ ε_score+e^-T d log(S)+Td.
Information-theoretic identities express φ_i,c(t) as KL divergence between a forward marginal and its pushforward under a coordinate-update bijection.

C.4 Efficient sampling for high-entropy distributions

The high-entropy regime permits an adaptive exponential-then-constant discretization schedule whose error depends on the initial divergence from the uniform distribution. The proof constructs an equivalent CTMC on each interval and controls its KL error term by term.

Theorem 4: Theorem 4 uses an exponential-then-constant schedule with step sizes bounded by κ min(1,T-t_k+1), for 0<κ<0.9.
Theorem 4: The resulting error upper bound depends only on KL(q_0||Unif(X)), which may be small when the data distribution is close to uniform.
Theorem 4: Sublinear dependence on d occurs when KL(q_0||Unif(X))=o(d).
Proof strategy: The initialization error is bounded by e^-T d(1+log S+T) ≲ e^-T d log S.
Proof strategy: On each interval [t_k,t_k+1], the sampler has the same endpoint distribution as a CTMC with a constructed rate matrix.
Proof strategy: The KL analysis separates initialization, score-entropy-loss, martingale, and conditional-mutual-information contributions.
Proof strategy: Both constant and exponential-then-constant schedules achieve KL(q_0||p_T) ≲ ε_score+ε under their respective step-size conditions.

D.3 τ-leaping for masking discrete diffusion

The masking-noise analysis uses truncated τ-leaping because multiple jumps on one coordinate are ill-defined, and derives a KL guarantee under a step-size schedule with early stopping. Its bound tracks score entropy loss plus additional discretization terms arising from score-time approximation.

Sampler construction: Truncated τ-leaping is analyzed because multiple jumps on one coordinate are ill-defined under masking noise.The sampler uses the truncated version rather than classical τ-leaping for this reason.
Guarantee: Theorem 5 imposes a step-size condition and early stopping at T − δ, producing an exponential-then-constant schedule.The theorem specifies bounded increments relative to min(1, T − t_k), while the analysis requires δ > 0.
Guarantee: The output distribution p_T−δ is controlled under the truncated sampler, with KL error bounded by score entropy loss and ε.The stated guarantee has the form KL(q_δ || p_T−δ) ≲ ε_score + ε.
Assumptions: Under one-sided score boundedness, the additional constant can be upper bounded using Cauchy–Schwarz.The corollary assumes a lower bound on the relevant score probabilities.
Error decomposition: The masking analysis adds two error terms because discrete-time score estimates require both time interpolation and score-function approximation.Unlike the cleaner uniform-noise analysis, the truncated sampler needs an additional approximation step and early stopping.
Information quantity: For masking data, the effective quantity is bounded by the minimum of total correlation and dual total correlation.Lemma 16 states D(q_data) ≤ min(B(q_data), C(q_data)).

E.6 Proof of Lemma 12

The proof establishes decay of coordinate-wise conditional KL terms under the masking forward semigroup, then differentiates and integrates these bounds to control the target information quantity.

Semigroup decomposition: The forward generator decomposes into commuting coordinate and remaining-coordinate operators.This decomposition enables separate analysis of one-coordinate evolution and the complementary coordinates.
Contraction: The induced one-coordinate channel contracts conditional KL divergence by e^−u after time u.The proof derives this from a strong data processing inequality for the S-ary symmetric channel.
Contraction: The contraction controls the evolution of φ_i,c(t), a coordinate-wise conditional KL quantity, throughout the semigroup.The inequality holds for every u ≥ 0 and supplies the derivative bound used later.
Integration: The proof then combines the derivative estimate with entropy and mutual-information identities to obtain the lemma’s integrated bound.The final steps use differential-entropy identities and Gaussian-noise calculations where applicable.

F.3 Proof of Lemma 13

The proof verifies that the CTMC transition over each interval matches the masking τ-leaping update, then expands the resulting score terms to establish the required identity.

Transition matching: The masking CTMC decomposes across coordinates, allowing the transition law to be checked coordinate by coordinate.Coordinates that are not masked remain unchanged under both constructions.
Transition matching: For masked coordinates, the proof computes probabilities of remaining masked or becoming symbol c and matches them to Algorithm 1.The resulting CTMC distribution equals the algorithmic update distribution.
Score identity: The score expansion introduces f_i(y) as a log-ratio involving q_0 and the coordinate value y_i.This function is independent of the i-th coordinate, enabling Dynkin’s formula on the remaining coordinates.
Conclusion: After applying Dynkin’s formula and rearranging masked-coordinate terms, the proof combines the resulting equations to conclude the desired identity.The final rearrangement uses the masking transition relation and expectation cancellations.