Provably adaptive sampling with uniform and remasking discrete diffusion models

Provably adaptive sampling with uniform and remasking discrete diffusion models Daniil Dmitriev∗,

Zhihan Huang˚ ,

Yuting Wei˚

arXiv:2608.23554v1 [cs.LG] 24 Aug 2026

August 25, 2026

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, ˆ ˙ DTCpX0 q N “O ε discretization steps suffice to achieve sampling error Opεscore ` εq, where ε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 DTCpX0 q, rather than directly by the ambient dimension d. In particular, this shows that the unfavorable dimension dependence of τ -leaping for uniform diffusion is a consequence of the sampler rather than the forward process itself. 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 DTCpX0 q. Numerical experiments on structured synthetic distributions illustrate the predicted dimension-adaptive behavior.

Contents 1 Introduction 1.1 Sampling efficiency and algorithm design . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2 An information-theoretic perspective on adaptive sampling . . . . . . . . . . . . . . . . . . . 1.3 Our main contributions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.4 Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.5 Other related works . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

2 3 3 4 5 5

2 Problem setup 2.1 Continuous-time Markov chain . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.2 Learning the reverse process . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

6 6 7

∗ Department of Statistics and Data Science, {daniild,zhihanh,ytwei}@wharton.upenn.edu

the

Wharton

1

School,

University

of

Pennsylvania;

email:

3 Discrete diffusion sampling with leave-one-out sampler 3.1 Leave-one-out sampler . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3.2 Bayes-optimal sampler . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

9 9 10

4 Main results

12

5 Numerical examples

14

6 Discussion

16

A Details on the remasking process

17

B Technical preparations

19

C Proof of our main results 20 C.1 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 C.2 Proof of Proposition 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 C.3 Proof of Proposition 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 C.4 Discretization error control . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 C.5 Approximation error control . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 D Proofs of results in Appendix B 31 D.1 Proof of Proposition 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 D.2 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 D.3 Proof of Lemma 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 D.4 Proof of Lemma 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 E Connection to the effective total correlation

1

35

Introduction

Diffusion models have achieved remarkable success across a wide range of generative modeling tasks and have recently emerged as a compelling alternative to autoregressive (AR) modeling for discrete sequence generation, such as natural language and protein sequences (Austin et al. (2021); Campbell et al. (2022); Lou et al. (2024)). Unlike AR models, which generate tokens sequentially according to a rigid left-to-right factorization, discrete diffusion models enable parallel generation and iterative refinement, making them an increasingly important component of the modern generative modeling toolbox. Among existing approaches, masking diffusion has been particularly successful: by progressively replacing data tokens with a dedicated mask token and learning to recover them, masking diffusion has demonstrated strong empirical performance at scale (Ou et al. (2025); Sahoo et al. (2024); Shi et al. (2024)). Nevertheless, the vanilla masking diffusion sampling process has an inherent limitation: unmasking is typically monotone, so once a position is assigned a token, it cannot be revisited in subsequent denoising steps. Therefore, incorrect early predictions may persist and become part of the context used to generate other tokens, potentially leading to error propagation (Huang et al. (2026); Kim et al. (2025); Xu et al. (2025)). Recent remasking and iterative-refinement methods try to overcome this limitation by explicitly allowing uncertain predictions to be remasked and regenerated, suggesting that the ability to revise intermediate decisions is an important property of discrete generative models (Wang et al. (2026); Zhao et al. (2026)). Uniform diffusion naturally provides this capability. Since corruption and denoising operate entirely within the original vocabulary, a token can transition between different valid states throughout the sampling trajectory, supporting error correction without relying on a special mask state. Importantly, uniform diffusion is not merely of theoretical interest; recent large-scale systems, including Google DeepMind’s DiffusionGemma (Google DeepMind (2026)) and the 7B Sumi model (Ye et al. (2026)), adopt uniform-state diffusion, while other recent frameworks, such as GIDD (Von Rütte et al. (2025)) and XDLM (Liu et al. (2026)), have explored hybrid uniform-masking processes. These developments motivate a systematic reconsideration of discrete diffusion

2

beyond masking diffusion as a practical and complementary foundation for large-scale discrete generative modeling.

1.1

Sampling efficiency and algorithm design

A central question for discrete diffusion models is their sampling efficiency, namely, how many discretization steps or model evaluations are needed to generate a sample to a prescribed accuracy. This question is closely tied to the particular choice of the sampler, as different samplers may yield different results. Many samplers were proposed recently, including the τ -leaping sampler (Campbell et al. (2022)) and its variants (Liang et al. (2025c)), uniformization sampler (Chen and Ying (2025)), DMPM sampler (Pham et al. (2025)), and others. Importantly, certain samplers (e.g., uniformization) only allow one transition at a time. Although theoretically interesting, such a restriction diminishes the main advantage of discrete diffusion models, parallel generation. In what follows, we consider the class of samplers that by design may perform several updates per discretization step, which includes the τ -leaping sampler. Recent work by Dmitriev et al. (2026) reveals a striking adaptivity property of masking diffusion: the sampling complexity can automatically adapt to the intrinsic dimensionality or structural complexity of the target distribution, leading to substantially better efficiency than worst-case guarantees when the target distribution exhibits favorable structure. This echoes a growing literature on continuous diffusion models that establishes analogous forms of adaptation to low-dimensional or structured target distributions; see, e.g., Huang et al. (2024); Li et al. (2025); Li and Yan (2024); Liang et al. (2025a). For uniform discrete diffusion, however, the picture appears less optimistic. Beyond a few special cases, Dmitriev et al. (2026) show that the sampling complexity of the widely adopted τ -leaping algorithm scales linearly with the ambient dimension d of the target distribution. Such linear dependence on the ambient dimension essentially precludes fast parallel generation, as AR models also require a number of model evaluations linear in d. This contrast between the adaptive guarantee for masking diffusion and the lower bound for uniform diffusion raises a fundamental question: Is this unfavorable dimension dependence an intrinsic limitation of uniform discrete diffusion, or merely a consequence of the τ -leaping sampler? In this work, we show that the latter is true and propose a sampler that yields adaptive guarantees for the uniform discrete diffusion. Resolving this question is important both for understanding the statistical and computational limits of uniform diffusion and for guiding the design of more efficient sampling algorithms. We focus on first-order sampling methods, where each denoising step requires only a single evaluation of the learned model, and investigate the fundamental limits of their sampling efficiency, in contrast to higher-order samplers that use multiple model evaluations per step (e.g., Ren et al. (2026)). A related open problem concerns remasking discrete diffusion models, which have become widely adopted mechanisms for revising previously generated tokens but currently lack a comparable theoretical understanding of their sampling efficiency. Our goal is therefore to characterize when uniform and remasking diffusion can exploit low-dimensional or structured target distributions, and to determine how their sampling complexity can adapt to intrinsic structure in a manner analogous to masking diffusion.

1.2

An information-theoretic perspective on adaptive sampling

To understand the adaptive sampling behavior described above, we turn to information-theoretic quantities that capture the intrinsic structure of discrete distributions. As a simple example, consider a d-dimensional binary distribution that is uniform over the two strings 0d and 1d . While the ambient dimension can be arbitrarily large, the distribution contains only one bit of uncertainty. Moreover, this global structure remains detectable after corruption. Under masking diffusion, observing a single unmasked coordinate determines the original string, whereas under uniform diffusion, the imbalance between the numbers of zeros and ones in the corrupted sequence provides information about the initial state. Thus, the effective difficulty of recovering the underlying sample can remain small even in high dimensions. This intuition can be formalized using information-theoretic measures of dependence. One such quantity, which has appeared in recent analyses of adaptive sampling (Chen et al. (2025); Dmitriev et al. (2026); Zhao

3

Estimated KL divergence

Masking Remasking Uniform

0.8 0.6 0.4 0.2 0.0

0

100

200 300 Dimension d

400

500

Figure 1: N “ 20 discretization steps. Target distribution is a Markov chain, see Section 5. Remasking process uses pM “ 0.5. We use early stopping δ “ 1e´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 tk`1 ´tk “ pT ´δq{N ; the remasking and uniform processes use geometric discretization tk`1 ´ tk “ κ minp1, T ´ tk`1 q with κ « 1.3. Results are averaged over 7 runs. and Cai (2026)), is the dual total correlation: DTCpXq “ HpXq ´

d ÿ

HpX i | X ´i q,

(1)

i“1

where X “ pX 1 , . . . , X d q. The DTC measures the dependence among the coordinates that remains after conditioning each coordinate on all the others. Importantly, DTCpXq ď HpXq ď d log S, and there exist high-dimensional distributions whose DTC remains bounded independently of d. Hence, guarantees expressed in terms of DTC can be substantially sharper than worst-case bounds that scale directly with the ambient dimension. More broadly, our analysis is based on an information-theoretic characterization of the sampling error rather than on DTC alone. We show that the discretization error can be expressed through mutual information between different coordinates of the forward process at different times. For the uniform and remasking processes, we shall see that this general characterization can then be controlled in terms of the DTC of the target distribution.

1.3

Our main contributions

Our main contributions are as follows: • An efficient leave-one-out sampler for uniform and remasking diffusion. We study a firstorder sampler based on leave-one-out conditional probabilities. On each discretization interval, the resulting approximate reverse process decomposes into independent one-dimensional CTMCs and can be executed for all coordinates in parallel. For the uniform process, our sampler recovers the recently proposed leave-one-out bridge plug-in and cavity ancestral samplers (Gourevitch et al. (2026); Noguerales et al. (2026)), while our CTMC formulation provides a unified construction that also applies to the remasking process, with the masking process as a special case. • Adaptive sampling guarantees beyond ambient dimension. Our main result, Theorem 1, establishes for both the uniform and remasking processes the bound KLpqT ´tN }poutput q À KLpqT }qnoise q ` εscore ` κDTCpX0 q.

4

r pDTCpX0 q{εq , discretization steps to achieve sampling error Consequently, it suffices to take N “ O Opεscore ` εq. Thus, the sampling complexity adapts to the intrinsic dependence structure of the target distribution. For the uniform process, this shows that the unfavorable dimension dependence of the standard τ -leaping sampler is not due to the forward process itself, circumventing the lower bound of Dmitriev et al. (2026). Our result also provides an adaptive sampling guarantee for the remasking process, for which theoretical guarantees were previously lacking. • A Bayes-optimal decomposition of sampling error. To disentangle errors arising from discretization and score estimation, we introduce a Bayes-optimal auxiliary sampler that uses the exact leave-one-out conditional probabilities available at each discretization point. The discretization error measures the discrepancy between the true reverse process and the Bayes-optimal sampler, whereas the approximation error measures the additional discrepancy introduced by the learned sampler. The former depends only on the target distribution, forward process, and time discretization, whereas the latter is controlled by the standard score entropy loss and can also be characterized directly through errors in the leave-one-out denoiser. Moreover, in Theorem 2, we give an exact information-theoretic representation of the discretization error in terms of mutual information between different coordinates of the forward process at different times. This characterization contains no explicit dependence on ambient parameters such as the dimension or vocabulary size, and may be of independent interest.

1.4

Notation

For a positive integer n, we denote rns “ t1, . . . , nu. We use d, S, and T to denote the ambient dimension, the vocabulary size, and the time horizon, respectively. Let V “ rSs Y Aux, where Aux is the set of auxiliary states. We have Aux “ H or Aux “ tMASK, REMASKu. Let qdata denote a distribution on rSsd . For x “ px1 , . . . , xd q P V d and i P rds, we denote x´i :“ px1 , . . . , xi´1 , xi`1 , . . . , xd q P V d´1 . For x P V d , i P rds, and b P V, we denote x di b P V d as follows: px di bqj “ xj for j ‰ i, and px di bqi “ b. We use KL, H, and I to denote the KL divergence, the entropy, and the mutual information, respectively. We adopt r hides logarithmic factors in d, S, and standard asymptotic notation: Op¨q, op¨q, Ωp¨q, and À; notation Op¨q 1{ε. We let Dϕ px, yq “ ϕpxq ´ ϕpyq ´ px ´ yqJ ∇ϕpyq be the Bregman divergence for ϕ : Rn Ñ R, and denote Dpa, bq “ ab ´ 1 ´ log ab as the Bregman divergence for the scalar function ϕpxq “ ´ log x.

1.5

Other related works

Sampling guarantees for discrete diffusion. A central question in the theory of discrete diffusion models is how many sampling steps are required to generate an accurate sample. Early work by Chen and Ying (2025) studies an exact uniformization-based sampler and establishes guarantees in both KL divergence and total variation, which eliminates discretization error but requires a number of sampling steps that scales linearly with d. Subsequent works analyze a broader range of discretized samplers, including τ -leaping, Euler, and Tweedie τ -leaping (Liang et al., 2025c), as well as higher-order schemes (Ren et al., 2026). Related convergence guarantees have also been developed for discrete Markov probabilistic models (Pham et al., 2025), absorbing or masking processes (Liang et al., 2025b, 2026), and more general masked and random-walk dynamics (Conforti et al., 2025). For the uniform process, Dmitriev et al. (2026) establish a r sharp Opd{ϵq complexity for the standard τ -leaping sampler, together with a matching algorithmic lower bound. Adaptive guarantees. More recent work has sought sampling guarantees that depend on the intrinsic structure of the target distribution rather than directly on the ambient dimension. Li and Cai (2025) establish a sampling complexity bound in terms of mutual information, while Chen et al. (2025); Zhao and Cai (2026) sharpen this dependence to information-theoretic quantities such as total correlation and dual total correlation. Cai and Li (2026) further show that confidence-based unmasking schedules for diffusion language models can achieve sublinear sampling complexity. An interesting recent work by Wainwright (2026) develops an information-theoretic measure of data geometry for masking diffusion, yielding datadependent sampling guarantees and optimized sampling schedules. Using a CTMC framework, Dmitriev et al. (2026) establish a guarantee in terms of effective total correlation, which is upper bounded by both

5

total correlation and dual total correlation. These adaptive guarantees, however, are specific to masking diffusion, either in the diffusion-language-model formulation or in the CTMC framework. In contrast, our results establish adaptive guarantees for the uniform and remasking processes. Leave-one-out denoisers. Score estimation is a central component of discrete diffusion models, particularly following the score entropy objective introduced by Lou et al. (2024). More recently, Gourevitch et al. (2026) and Noguerales et al. (2026) observe that training the leave-one-out denoiser, or a cavity estimator, is equivalent to score estimation and yields better empirical performance. Our work also adopts the leave-one-out denoiser as a central object, but from a complementary perspective: we show that it arises naturally from the CTMC formulation of the reverse process and use it to construct the sampler and analyze its approximation and discretization errors.

2

Problem setup

2.1

Continuous-time Markov chain

Let us begin by introducing the framework of discrete diffusion models, which are used to approximate a target distribution qdata on a discrete domain rSsd . Analogous to their continuous counterparts, discrete diffusion models consist of a forward process and a reverse process evolving over the discrete state space. Both processes can be formulated in terms of continuous-time Markov chains (CTMCs), which we introduce next, following Campbell et al. (2022). Definition 1. A stochastic process pXt qtPr0,T s on V d with rate matrices pQt qtPr0,T s and initial distribution q0 is a continuous-time Markov chain (CTMC) if X0 „ q0 and 1. pXt qtPr0,T s satisfies the Markov property: for any 0 ď u ă v ď T , Xv is conditionally independent of pXt qtău given Xu ; 2. As ∆t Ó 0, for any x, y P V d , PrpXt`∆t “ y | Xt “ xq “ Ity “ xu ` Qt px, yq∆t ` op∆tq. d

d

Here, for any t P r0, T s, the rate matrix Qt P RV ˆV satisfies: 1. Qt px, yq ě 0 for any x ‰ y P V d , ř 2. Qt px, xq “ ´ y‰x Qt px, yq. For any fixed initial distribution q0 , the marginals pqt qtPr0,T s of Xt are the solutions to the Kolmogorov forward equation: for all x P V d ,

ÿ d PrpXt “ xq “ Qt py, xq PrpXt “ yq. dt d yPV

We refer the reader to Feinberg et al. (2014); Feller (1940) for a rigorous treatment of CTMCs. Forward process. We define a forward process as a CTMC pXt qtPr0,T s corresponding to the particular choice of rate matrices pQt qtPr0,T s and we assume that there exist pQt qtPr0,T s , with each Qt P RVˆV , such that 1. Qt px, yq “ 0 if dH px, yq ě 2, 2. Qt px, yq “ Qt pxi , y i q, if dH px, yq “ 1 and xi ‰ y i . This requirement states that the forward process acts on each coordinate independently. In this work, we focus on the following time-homogeneous (Qt ” Q) forward processes: 1. Uniform process: when V “ rSs and Qpa, bq “ 1{S for all a, b P rSs,

6

2. Remasking process: when V “ rSs Y tMASK, REMASKu, and for a P rSs, b P V, and 0 ă pM ď 1, we have Qpa, bq “ pM Itb “ MASKu ` p1 ´ pM qItb “ REMASKu

and

QpREMASK, aq “ 1{S.

Here, MASK is an absorbing state and the masking process corresponds to pM “ 1. For the remasking process, choosing pM ă 1 allows the reverse process to correct unmasked elements, as the coordinates in the forward process will transition through the REMASK state one or several times with positive probability. As discussed, this is crucial to mitigate a well-known disadvantage of the masking process, where once a coordinate is unmasked, it cannot be changed later during the sampling process; see Wang et al. (2026); Zhao et al. (2026). Self-loops. A convenient way to model how the uniform process acts on a single coordinate is as follows: the process makes a transition from state a P rSs with rate 1 to a new state b „ UnifprSsq, with a KK b. For technical convenience, we define diagonal elements in the uniform case to be Qpa, aq “ 1{S for a P rSs, which makes Q the transition matrix of a discrete-time Markov chain, rather ř than the rate matrix of a CTMC. As the only difference is in the diagonal elements, we explicitly write ´ b‰a Qpa, bq whenever the diagonal entry of the rate matrix is used in the analysis. Similarly, for the remasking process, we set QpMASK, MASKq “ 1. Reverse process.

For the forward process pXt qtPr0,T s with its marginal distribution pqt qtPr0,T s , there exists Ð

Ð

a time-reversed CTMC pX t qtPr0,T s with an initial distribution qT and rate matrices pQt qtPr0,T s : Ð

Qt px, yq “ QT ´t py, xq

PrpXT ´t “ yq , PrpXT ´t “ xq

for all x ‰ y P V d , d Ð

such that its marginals coincide with the forward process: Xt “ X T ´t for all t P r0, T s. We refer to this CTMC as the reverse process (Campbell et al. (2022)). Analogous to the continuous case, where the Stein score function ∇ log pt pxq determines the reverse process, we define the (concrete) score function as the ratio of the probability mass functions st py, xq :“

PrpXt “ yq . PrpXt “ xq

(2)

As by our construction, Qt py, xq ą 0 only if dH px, yq “ 1, we often denote y “ x di b for some i P rds and b P V. To construct the reverse process, it is therefore sufficient to know the score functions (2), for every t P r0, T s, x P V d and y “ x di b.

2.2

Learning the reverse process Ð

Computing rate matrices of the reverse process pQt qtPr0,T s requires access to score functions at every t P r0, T s, which is not tractable in general. Instead, they are replaced in practice by a data-driven estimator spt py, xq at discrete time points 0 ă T ´tN ´1 ă . . . ă T ´t0 “ T such that spt py, xq « st py, xq “ PrpXt “ yq{ PrpXt “ xq. p t px, yq “ QT ´t py, xqp sT ´t py, xq. The estimated rate matrices are defined as Q To measure the accuracy of the estimated score, we use the score entropy loss introduced by Lou et al. (2024), which has become a standard objective for training discrete diffusion models. This loss quantifies the discrepancy between the estimated score spt py, xq and the true score st py, xq associated with the forward process: « ff ÿ spy, xt q, spy, xt qq . LSE pt, sp, sq :“ Ext „qt Qt py, xt qspy, xt qDpp (3) y‰xt

Here t ě 0 and sp, s P V d ˆ V d Ñ Rě0 . To disentangle the effect of score estimation from the discretization or optimization errors that govern the efficiency of the sampling algorithms being considered, we isolate the score estimation error and make the following assumption. 7

Assumption 1. Let 0 “ t0 ă t1 ă . . . ă tN ď T be the time discretization. We assume that Nÿ ´1

(4)

ptk`1 ´ tk qLSE pT ´ tk , spT ´tk , sT ´tk q ď εscore .

k“0

This assumption on the score entropy loss is a standard way to control the approximation error of discrete diffusion models and has appeared in several prior works, including Conforti et al. (2025); Dmitriev et al. (2026); Liang et al. (2025c); Lou et al. (2024). A leave-one-out formulation. We state a leave-one-out denoiser formulation that plays an important role in our analysis and has appeared previously in the context of both remasking (Zhao et al., 2026) and uniform diffusion (Gourevitch et al., 2026). The key idea is that the estimator must approximate the leaveone-out conditional probabilities PrpX0i “ b | Xt´i “ x´i q instead of PrpX0i “ b | Xt “ xq. This can be achieved either by imposing the leave-one-out structure directly through the design of the estimator or by modifying the cross-entropy objective used to train the denoiser. We refer to Gourevitch et al. (2026) for a more detailed discussion. In our setting, the leave-one-out formulation leads to a cleaner characterization of the reverse dynamics and, more importantly, facilitates the control of the resulting sampling errors. This role is different from its use in prior work, where the emphasis is more closely tied to the practical construction and training of the denoiser. Formally, recall the definition of the score function st px di b, xq “ PrpXt “ x di bq { PrpXt “ xq. Dividing numerator and denominator by PrpXt´i “ x´i q, we obtain PrpXti “ b | Xt´i “ x´i q PrpXt “ x di bq . (5) “ PrpXt “ xq PrpXti “ xi | Xt´i “ x´i q ␣ ( There exists a simple bijection between the sets tst px di b, xqubPV and PrpXti “ b | Xt´i “ x´i q bPV , as ř ´i i ÿ “ x´i q 1 i PrpXt “ b | Xt st px di b, xq “ b‰x i “ ´ 1. ´i i i ´i i PrpXt “ x | Xt “ x q PrpXt “ x | Xt´i “ x´i q b‰xi st px di b, xq “

Indeed, we can use this equality together with the convention st px, xq “ 1 to compute for all b P V, st px di b, xq PrpXti “ b | Xt´i “ x´i q “ ř . c st px di c, xq Therefore, computing the score is equivalent to computing the leave-one-out denoiser probabilities. In the following, in light of this bijection, we refer to both quantities interchangeably. For both uniform and remasking processes, we find it instrumental to introduce the notation νpt, bq :“ PrpXti “ b | D a jump at the i-th coordinate on r0, tsq.

(6)

For both uniform and remasking processes, PrpXti “ b | Xt´i “ x´i q can be expressed in terms of νpt, bq and PrpX0i “ b | Xt´i “ x´i q: PrpXti “ b | Xt´i “ x´i q “ PrpX0i “ b | Xt´i “ x´i qe´t ` νpt, bqp1 ´ e´t q.

(7)

Indeed, in both cases, for any fixed i P rds and time t ě 0, either there was a jump at the i-th coordinate on r0, ts, and thus Xti K K Xt´i , or there was no jump and Xti “ X0i , which happens with probability e´t . • For the uniform process, νpt, bq “ 1{S and we obtain PrpXti “ b | Xt´i “ x´i q “ PrpX0i “ b | Xt´i “ x´i qe´t `

1 p1 ´ e´t q. S

(8)

• For the remasking process, Lemma 2 gives the explicit expression for νpt, bq. Observe that the only unknown data-dependent quantities are tPrpX0i “ b | Xt´i “ x´i qubPrSs . As st px di b, xqu is equivalent to assuming access to a! consequence, assuming) access to score estimators tp ´i i ´i x PrpX0 “ b | Xt “ x q at the discretization points t P tT ´ t0 , . . . T ´ tN ´1 u. bPrSs

8

3

Discrete diffusion sampling with leave-one-out sampler

With the score estimator spt that satisfies Assumption 1, the remaining task is to construct a tractable approximation to the reverse-time process. Since spt is only evaluated at the time discretization 0 ă T ´ tN ´1 ă . . . ă T ´ t0 “ T , we extend these estimates over r0, T s and use the resulting approximation to simulate the reverse dynamics. The standard choice, both in theoretical analyses and in practice, is the τ -leaping sampler (Campbell et al. (2022)) and its variants (Liang et al. (2025c)). However, recent work by Dmitriev et al. (2026) shows that, for the uniform process, the τ -leaping sampler can lead to suboptimal sampling complexity. In this section, we study the following leave-one-out sampler which leads to improved sampling efficiency and overcomes the theoretical barrier of τ -leaping.

3.1

Leave-one-out sampler

Consider the k-th discretization interval and let u “ T ´ tk`1 and ℓ “ T ´ tk . As discussed above, we x i “ b | X ´i “ x´i qu for b P rSs and i P rds, and assume that the sampling algorithm has access to tPrpX 0 ℓ ℓ x i “ b | X ´i “ x´i qubPrSs,iPrds . Following Eqn. (7), we define our goal is to sample xu given xℓ and tPrpX 0 ℓ ℓ x i “ b | X ´i “ x´i q as follows: PrpX t ℓ ℓ x i “ b | X ´i “ x´i q “ PrpX x i “ b | X ´i “ x´i qe´t ` νpt, bqp1 ´ e´t q. PrpX t 0 ℓ ℓ ℓ ℓ

(9)

Together with Eqn. (2), we define our score estimator spt for t P ru, ℓs: spt px di b, xq :“

x i “ b | X ´i “ x´i qe´t ` νpt, bqp1 ´ e´t q x i “ b | X ´i “ x´i q PrpX PrpX t 0 ℓ ℓ ℓ ℓ “ , x ti “ xi | X ´i “ x´i q x i “ xi | X ´i “ x´i qe´t ` νpt, xi qp1 ´ e´t q PrpX PrpX ℓ

0

ℓ

ℓ

(10)

ℓ

p t px, yq “ QT ´t py, xqp p t px, xdi and the corresponding rate matrix Q sT ´t py, xq. Note that both spt pxdi b, xq and Q ´i bq do not depend on x . This means that effectively on the discretization interval ru, ℓs, the approximate CTMC can be decomposed into d independent one-dimensional CTMCs, allowing us to simulate the dynamics of the CTMC in parallel for each fixed discretization interval, as is done in Algorithm 1, see Proposition 1. In view of the Bayes formula, we can write PrpXui “ b | Xℓ “ xℓ q “

PrpXℓi “ xiℓ | Xui “ bq PrpXui “ b | Xℓ´i “ x´i ℓ q ´i ´i i i PrpXℓ “ xℓ | Xℓ “ xℓ q

9 PrpXℓi “ xiℓ | Xui “ bq PrpXui “ b | Xℓ´i “ x´i ℓ q. Here, PrpXℓi “ xiℓ | Xui “ bq does not depend on qdata and can be computed from the properties of Q, and ´u PrpXui “ b | Xℓ´i “ x´i PrpX0i “ b | Xℓ´i “ x´i ` looooooooomooooooooon νpu, bqp1 ´ e´u q, ℓ q “ looooooooooooooooomooooooooooooooooon ℓ qe no jump on r0,us

jump on r0,us

where νpu, bq is defined in Eqn. (6). Therefore, the sampling algorithm applies the following in parallel over i P rds: Given xℓ , sample xiu from µ pi , ´ ¯ x i “ b | X ´i “ x´i qe´u ` νpu, bqp1 ´ e´u q . µ pi pbq 9 PrpXℓi “ xiℓ | Xui “ bq PrpX (11) 0 ℓ ℓ Finally, we remark that as PrpXℓi “ a | Xui “ bq depends only on the matrix Q, it can be explicitly computed. Thus, we set x i “ a | X i “ bq “ PrpX i “ a | X i “ bq, PrpX ℓ u ℓ u for all 0 ď u ď ℓ and a, b P V. Putting things together and instantiating the above to the processes that we consider gives: • For the uniform process, ¯´ ¯ ´ ´pℓ´uq x i “ b | X ´i “ x´i qe´u ` 1´e´u PrpX Itb “ xiℓ ue´pℓ´uq ` 1´e S 0 ℓ ℓ S pi pbq “ . µ ´i ´i ´ℓ 1´e´ℓ i i x PrpX0 “ xℓ | Xℓ “ xℓ qe ` S 9

(12)

Algorithm 1: Our sampling algorithm Input: Initial distribution: p0 , Discretization steps: 0 “ t0 ă t1 ă . . . ă tN ď T , x i “ b | X ´i “ ¨ q for t P tt0 , . . . , tN ´1 u, b P rSs, and i P rds. Leave-one-out denoiser: PrpX 0 T ´t d Output: Sample x pPV . 1 Sample xT from p0 2 for k “ 0, . . . , N ´ 1 do 3 u Ð T ´ tk`1 4 ℓ Ð T ´ tk 5 for i P rds in parallel do 6 Let µ pi be a probability distribution ´ over V with ¯ i i i x i “ b | X ´i “ x´i qe´u ` νpu, bqp1 ´ e´u q pi pbq 9 PrpX “ x | X “ bq PrpX µ ℓ

ℓ

u

0

ℓ

ℓ

pi Sample xiu from µ 8 end 9 end 10 return xT ´tN 7

• For the masking process, for i P rds with xiℓ “ MASK, pi pbq “ µ

e´u ´ e´ℓ x i PrpX0 “ b | Xℓ´i “ x´i ℓ q 1 ´ e´ℓ

for b P rSs,

and

µ pi pMASKq “

1 ´ e´u , 1 ´ e´ℓ

recovering Algorithm 1 from Dmitriev et al. (2026). • For the general remasking process, Lemma 4 gives explicit expressions. The next proposition provides a simple algorithm to simulate this CTMC for both uniform and remasking processes. The proof of this result is deferred to Section C.2. Proposition 1. Fix k P t0, . . . , N ´ 1u and let u “ T ´ tk`1 , ℓ “ T ´ tk . Let xℓ and xu be as in Algorithm 1. p t defined in Eqn. (10). Then, Let yu be the distribution of a CTMC initialized at xℓ with the rate matrices Q d

xu “ yu . Consequently, Algorithm 1 simulates the full dynamics of the CTMC with initial distribution p0 and rate p t qtPr0,t s . matrices pQ N For the uniform process, our sampler coincides with prior work, e.g., the leave-one-out bridge plug-in sampler of Gourevitch et al. (2026) after identifying their noise schedule with αt “ e´t . The two approaches, however, arise from different perspectives. Gourevitch et al. (2026) derive the leave-one-out predictor as the optimal target for the bridge plug-in parameterization and study its implications for training and inference. In contrast, we derive the same sampling transition directly from the CTMC formulation: on each discretization interval, we condition the leave-one-out denoiser on the state available at the beginning of the interval and construct an approximate reverse-time CTMC, which decomposes into independent onedimensional processes that can be simulated in parallel. This viewpoint provides a natural Bayes-optimal intermediate process for separating approximation and discretization errors and, moreover, extends within a unified framework to both the uniform and remasking processes.

3.2

Bayes-optimal sampler

To separate the error due to score estimation from the error due to time discretization, we introduce an oracle counterpart of our sampler. On each discretization interval ru, ℓs, this oracle has access to the exact leave-one-out conditional probabilities at the beginning of the interval, but, like our practical sampler, it 10

Score function LOO denoiser

True reverse process

Bayes-optimal sampler

Our approximation

PrpXti “ b | Xt´i “ x´i t q PrpXti “ xit | Xt´i “ x´i t q

PrpXti “ b | Xℓ´i “ x´i ℓ q PrpXti “ xit | Xℓ´i “ x´i ℓ q

x i “ b | X ´i “ x´i q PrpX t ℓ ℓ ´i i i x PrpXt “ xt | X “ x´i q

PrpX0i “ b | Xt´i “ x´i t q

PrpX0i “ b | Xℓ´i “ x´i ℓ q

x i “ b | X ´i “ x´i q PrpX 0 ℓ ℓ

ℓ

ℓ

Table 1: Here, ℓ “ T ´ tk is the end point of the k-th discretization interval and T ´ tk`1 ď t ď ℓ. The score function row compares the true value st pxt di b, xt q, the Bayes-optimal srt pxt di b, xt q, and ours spt pxt di b, xt q. 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 ℓ. Our approximation is based x on the Bayes optimal one, but instead of the true probabilities Prp¨q it uses the estimate Prp¨q, similarly i for the leave-one-out denoiser. To compute the probability of Xt “ b used in the score function, all three approaches use the same formula (see Eqns. (7) and (9)) but with different leave-one-out denoisers. does not observe the evolving context Xt´i for t P ru, ℓs. Thus, it provides a natural intermediate process between the true reverse process and our approximation. More precisely, we define the following. Definition 2. Fix k P t0, . . . , N ´ 1u and let u “ T ´ tk`1 , ℓ “ T ´ tk . For t P ru, ℓs, define srt px di b, xq “

PrpXti “ b | Xℓ´i “ x´i ℓ q , PrpXti “ xi | Xℓ´i “ x´i ℓ q

(13)

r t px, yq “ QT ´t py, xqr and the corresponding rate matrix by: Q sT ´t py, xq. Equivalently, srt is obtained from our score estimator spt by replacing the estimated leave-one-out probabilities at the discretization point with their population counterparts. The following lemma justifies that this term is indeed Bayes optimal: among estimators restricted to the r t is the conditional expectation of the information available at the beginning of the discretization interval, Q true reverse rate. d´1 Lemma 1. Fix k P t0, . . . , N ´ 1u. Let t ď ℓ :“ T ´ tk . For any x´i and xit , b P V, it holds that ℓ PV

”Ð r T ´t px di xi , x di bq “ E ´i ´i QT ´t pxt , xt di bq Q t x „q t

t

ı i i Xℓ´i “ x´i ℓ , Xt “ xt .

We refer to Table 1 for the comparison of the true score function st , Bayes-optimal srt , and the approximation spt that we use. We emphasize that, for the time t ď ℓ “ T ´ tk and coordinate i P rds, the true reverse process evaluates the score function using the full current context Xt´i , while the Bayes-optimal sampler and our sampler can only access the context Xℓ´i available at the beginning of the interval. Our practical sampler introduces one additional approximation by replacing the exact conditional probabilities with their learned estimates compared to the Bayes-optimal sampler. In contrast, the score function used by the standard τ -leaping sampler (Campbell et al. (2022)) is given by: x i “ xi ` pb ´ xi q | X ´i “ x´i q PrpX t ℓ ℓ ℓ ℓ . sτt pxt di b, xt q “ x i “ xi | X ´i “ x´i q PrpX ℓ

ℓ

ℓ

ℓ

To evaluate the likelihood of transition xit Ñ b, the τ -leaping sampler uses transition xiℓ Ñ xiℓ ` pb ´ xit q at the beginning of the discretization interval. This construction implicitly relies on an ordinal structure of the state space, an issue already noted in Campbell et al. (2022), where truncation was proposed for non-ordered vocabularies. Moreover, τ -leaping ignores the time index t, and computes probabilities with respect to the beginning of the discretization interval ℓ. Both our sampler and the Bayes-optimal sampler mitigate these drawbacks.

11

4

Main results

We now present our main sampling guarantees for the uniform and remasking diffusion models. Our result shows that, under suitable choices of discretization, the sampling complexity is governed by the intrinsic dependence structure of the target distribution, as measured by its dual total correlation. All proofs for the results in this section are given in Appendix C. Theorem 1. Let 0 “ t0 ă t1 ă . . . ă tN ă T and suppose that for some κ P p0, 1q, tk`1 ´ tk ď κ minp1, T ´ tk`1 q for all k P t0, . . . , N ´ 1u. Consider either of the following two processes: 1. Uniform: pqt qtPr0,T s are the marginals of the uniform process, 2. Remasking: pqt qtPr0,T s are the marginals of the remasking process. Under Assumption 1, Algorithm 1 initialized from p0 “ qnoise outputs a sample xT ´tN „ poutput such that KLpqT ´tN }poutput q À KLpqT }qnoise q ` εscore ` κDTCpX0 q. ` ˘ In particular, for T “ O logpε´1 d log Sq , and corresponding choice of qnoise 1 , it suffices to take ˙ ˆ DTCpX0 q r , N “O ε

(14)

discretization steps to guarantee KLpqT ´tN }poutput q À εscore ` ε. This result separates the sampling error into three sources: initialization error, score-estimation error, and an intrinsic discretization error controlled by the dependence structure of the target distribution. In particular, the resulting step complexity scales with DTCpX0 q rather than explicitly with the ambient dimension d. For structured high-dimensional distributions with DTCpX0 q ! d, this can yield a substantially sharper guarantee than dimension-dependent worst-case bounds; see concrete examples in Dmitriev et al. (2026). For the uniform process, this result shows that the unfavorable dimension dependence previously established (Dmitriev et al. (2026)) for the standard τ -leaping sampler is not intrinsic to the forward process itself, but can instead arise from the choice of sampling algorithm. The same analysis also yields an adaptive sampling guarantee for the remasking process, for which, to the best of our knowledge, no comparable theoretical guarantee was previously available. Proof sketch. Using the Bayes-optimal sampler introduced in Section 3, we decompose the KL divergence between two path measures into approximation and discretization errors, see Proposition 2. To upper bound the discretization error, we express it as the integral of the second partial derivative of the mutual information (Theorem 2) and then upper bound it for both considered forward processes (Proposition 3). The proof of the latter result is based on Grönwall’s inequality. The following proposition formalizes the specific error decomposition underlying this argument. It is developed for general forward processes by comparing them with auxiliary CTMCs, and therefore separates the general information-theoretic part of our analysis from the process-specific bounds developed later. The proof of this result is included in Section C.3. Proposition 2. Let 0 “ t0 ă t1 ă . . . ă tN ď T be the time discretization. Recall st , spt from Eqns. (5) and (10), and srt from Definition 2. Then, KLpqT ´tN }poutput q ď KLpqT }qnoise q `

Nÿ ´1 k“0

Lpkq approx `

Nÿ ´1

pkq

Ldiscr ,

k“0

“ UniftrSsd u for the uniform process and qnoise “ µbd with µpbq “ Itb P rSsue´T {S ` νpT, bqp1 ´ e´T q for the remasking process. 1q noise

12

where, for k P t0, . . . , N ´ 1u, u “ T ´ tk`1 , and ℓ “ T ´ tk , « ff żℓ ÿ pkq Lapprox :“ Ext ,xℓ „qt,ℓ Qt py, xt qr st py, xt qDpp st py, xt q, srt py, xt qq dt u

and

«

żℓ pkq Ldiscr :“

(15)

y‰xt

ff ÿ

Ext ,xℓ „qt,ℓ u

Qt py, xt qst py, xt qDpr st py, xt q, st py, xt qq dt.

(16)

y‰xt

Proposition 2 provides a general and interpretable decomposition of the sampling error. A key feature of this decomposition is that the discretization error is independent of the particular sampler: it measures the discrepancy between the true reverse process and the Bayes-optimal sampler, and therefore depends only on the target distribution, the forward process, and the chosen time discretization. In contrast, the approximation error measures the discrepancy between a particular sampler and its Bayes-optimal counterpart, thereby isolating the error arising from approximating the reverse dynamics. This separation allows the discretization error to be studied independently of sampler-specific approximations, and the resulting analysis applies beyond Algorithm 1, including the τ -leaping sampler and its variants. Discretization error. We first characterize the discretization error independently of the particular forward process. The following theorem provides an exact information-theoretic representation in terms of the mutual information between one coordinate and the remaining coordinates at different times. The proof is deferred to Appendix C.4. Theorem 2. Fix k P t0, . . . , N ´ 1u, let u “ T ´ tk`1 and ℓ “ T ´ tk , and recall Ldiscr from Eqn. (16). Then, ÿ ż ℓ ż ℓ B2 pkq IpXti ; Xv´i qdvdt. (17) Ldiscr “ BtBv u t pkq

iPrds

We emphasize that Theorem 2 holds for any forward process and shows that the discretization error does not explicitly depend on the problem parameters, such as ambient dimension d and vocabulary size S, but instead on the information-theoretic properties of both the data distribution and the forward process. We next specialize this characterization to the uniform and remasking processes, and leave potential applications for other cases, e.g., discrete Gaussian, or semantic-dependent (see, e.g., Austin et al. (2021)), for future work. Proposition 3. Fix k P t0, . . . , N ´ 1u and let pXt qtPr0,T s be the uniform or remasking process. For u “ T ´ tk`1 and ℓ “ T ´ tk with ℓ ´ u ď κ minp1, uq, we have pkq

Ldiscr À κ pDTCpXu q ´ DTCpXℓ qq . řN ´1 pkq Using a telescoping sum, Proposition 3 immediately leads to k“0 Ldiscr ď κDTCpX0 q. We remark that while the statement requires ℓ ´ u ď κ minp1, uq, and thus imposes a geometric discretization grid tk`1 ´ tk ď κ minp1, T ´ tk`1 q, the discretization error on the interval r0, δs can also be made small, for δ small enough. Therefore, the early stopping requirement in Theorem 1 arises from the control of the approximation error rather than the discretization error, which we shall discuss next. Approximation error. We next study the approximation error Lapprox defined in Eqn. (15), which measures the discrepancy between a practical sampler and its Bayes-optimal counterpart. We first show that it is directly controlled by the score entropy loss appearing in Assumption 1. The proofs for the following two results are deferred to Appendix C.5. Proposition 4. Consider the k-th interval of the time discretization 0 “ t0 ă t1 ă . . . ă tN ă T and let u “ T ´ tk`1 , ℓ “ T ´ tk . Recall the true score st from Eqn. (2) and our score estimator spt from Eqn. (10). If ℓ ´ u ď κ minp1, uq with 0 ă κ ă 1, then Lpkq pℓ , sℓ q. approx À pℓ ´ uqLSE pℓ, s 13

(18)

Geometric grid Estimated KL divergence

Constant grid

Masking Remasking Uniform

2.0 1.5 1.0 0.5 0.0 0

100

200

300

Dimension d

400

500

0

100

200

300

Dimension d

400

500

Figure 2: N “ 30 discretization steps. Target distribution is the Markov chain with dpflip “ 2. The left plot shows the performance on the geometric grid tk`1 ´ tk “ κ minp1, T ´ tk`1 q and the right plot shows the result for the constant grid tk`1 ´tk “ pT ´δq{N . We use early stopping at δ “ 1e´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 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. Let us next give an alternative characterization directly in terms of the leave-one-out denoiser. This formulation makes explicit how errors in estimating the conditional distribution of X0i from the context Xt´i contribute to the sampling error. Importantly, the coordinate Xti itself is excluded from the conditioning context. Theorem 3. Let pXt qtPr0,T s be the CTMC corresponding to the uniform or remasking process. If tk`1 ´ tk ď κ minp1, T ´ tk`1 q for k P t0, . . . , N ´ 1u, then Nÿ ´1

Lpkq approx À κ

k“0

Nÿ ´1 ÿ

ˇ ´ ¯ › ˇ pX0i ˇ XT´i´tk . e´pT ´tk`1 q KL µX0i › µ

k“0 iPrds

ˇ ´ ¯ › ˇ pX0i ˇ XT´i´tk is the conditional KL divergence between PrpX0i “ ¨ | XT´i´tk q Here, for fixed i P rds, KL µX0i › µ x i “ ¨ | X ´i q. and PrpX 0

T ´tk

Theorem 3 gives a quantitative bound for the approximation error with respect to the cross-entropy loss (Austin et al. (2021); Sahoo et al. (2024)). The cross-entropy loss is widely used in practice for training discrete diffusion models (e.g., DiffusionGemma Team et al. (2026)), although not in the leave-one-out formulation. This result may be of future interest, e.g., for proving bounds on the sample complexity of score estimation.

5

Numerical examples

In this section, we conduct experiments on two synthetic target distributions to illustrate the adaptive sampling behavior predicted by our theory. In particular, we examine how the sampling error depends on the ambient dimension d, the choice of time discretization, and the structural complexity of the target distribution. For both target distributions considered, a binary Markov chain and a mixture of binary strings, the leaveone-out probabilities and thus the Bayes-optimal score function can be computed efficiently. We therefore use the exact values for our sampler, so that the approximation error vanishes and the remaining sampling error is solely due to discretization. 14

Estimated KL divergence

0.7 0.6 0.5 0.4 0.3 0.2 0.1 0.0

Tau-leaping Truncated tau-leaping Our sampler

20

30

40

50

60 70 Dimension d

80

90

100

Figure 3: N “ 40 discretization steps. Target distribution is the binary Markov chain with dpflip “ 2. We use early stopping at δ “ 1e´5 and T “ 8. Results are averaged over 7 runs. We compare three different samplers for the uniform process: τ -leaping, truncated τ -leaping, and ours. A binary Markov chain. We consider a Markov chain tak udk“1 on t0, 1u of length d, where a1 „ Bernp1{2q and for k “ 1, . . . , d ´ 1, ak`1 “ ak with probability 1 ´ p, and ak`1 “ 1 ´ ak otherwise. We choose p “ 2d´1 , so that a typical sample consists of several (three on average) consecutive blocks of the same digit. Let qdata be the distribution of this Markov chain. A simple computation shows DTCpqdata q ď Hpqdata q “ Oplog dq. Figure 1 compares three processes for the discrete diffusion models: (i) masking process, (ii) remasking process with pM “ 1{2, and (iii) uniform process. While all three processes benignly depend on the dimension (given only N “ 20 of discretization steps), uniform and remasking processes consistently outperform the standard masking process. This can be attributed to the fact that both uniform and remasking processes allow the sampler to correct early mistakes, thus can perform well under extremely coarse discretization grids. We leave the theoretical justifications of this observation for future work. Figure 2 shows two choices of discretization scheme for all three processes. We observe that the masking process performs much better on the constant grid, while uniform, in contrast, takes advantage of the geometric grid. Remasking also performs better on the geometric grid, but also performs well on the constant grid. Figure 3 compares three different samplers for the uniform process: τ -leaping, its truncated version, where only one jump per coordinate is allowed per discretization step, and our sampler. We observe that our sampler shows the best performance out of the three samplers. The additional errors of the τ -leaping and truncated τ -leaping samplers arise from the non-zero approximation error of these samplers, as they inaccurately follow the Bayes-optimal sampler. In these binary Markov chain experiments, we estimate the KL divergence by fitting an autoregressive model to the output samples. We use 1500 samples to fit the model parameters and 2500 samples for the KL estimation. Exact computations in small dimensions show good empirical agreement with this estimation. Mixture of binary strings. We independently sample k binary strings s1 , . . . , sk uniformly from t0, 1ud and define the empirical distribution k 1 ÿ qrdata “ δs , k i“1 i where δx denotes the Dirac delta distribution at point x P t0, 1ud . To ensure full support on t0, 1ud , we consider the smoothed distribution qdata “ p1 ´ εqr qdata ` ε Unifpt0, 1ud q, 15

ε “ 10´10 .

Estimated KL divergence

1.0 0.8

Masking Uniform Remasking

0.6 0.4 0.2 0.0

102

103 Mixture size k

Figure 4: N “ 20 discretization steps. Dimension d “ 2000. Target distribution is a sparse mixture with k components. We use early stopping at δ “ 1e´5 and T “ 8. Results are averaged over 7 runs. Masking uses constant grid; uniform and remasking use geometric grid. In the regime of log k ! d, we have that DTCpqdata q “ Oplog kq. Figure 4 shows how the estimated KL divergence scales with increasing k from k “ 80 to k “ 5000. As the X-axis is plotted in logarithmic scale, we see that, for a fixed number of discretization steps N “ 20, the KL divergence grows logarithmically with k, consistent with Theorem 1. For this setting, we also observe that the uniform and remasking processes incur smaller sampling errors than the widely used masking process. Providing a rigorous explanation for this phenomenon is an interesting direction for future work. In the mixture experiments, to estimate the KL divergence, we collect all the generated samples that do not match any of the k binary strings into a single bin and compute the KL divergence between this restricted distribution over k ` 1 elements. While this only provides a lower bound on the true KL divergence, we find that this approximation is accurate in low dimensions and can scale to higher dimensions.

6

Discussion

This paper establishes adaptive sampling guarantees for the uniform and remasking processes. Our results show that the linear dependence on the ambient dimension d exhibited by existing results for uniform discrete diffusion is not intrinsic to the forward process. Instead, the leave-one-out sampler studied here admits guarantees controlled by an information-theoretic measure of the target distribution. Thus, sampling efficiency depends not only on the choice of forward process, but also critically on how the reverse dynamics are approximated and discretized. To the best of our knowledge, this is the first work to establish such adaptive guarantees for both the uniform and remasking processes. Our analysis also reveals a common structure underlying the two processes and suggests that the techniques developed here extend naturally to other discrete diffusion models used in practice. A key property required by our arguments is that the forward process is unstructured : after a coordinate undergoes a jump, its new value is independent of its initial value. Our work suggests several directions for future investigation. • In our numerical experiments, the uniform and remasking processes consistently outperform the widely used masking process. Providing a theoretical explanation for this empirical observation is an interesting direction for future work. • Extending the present techniques to other forward processes, such as discrete Gaussian or semantically dependent processes, may help clarify the connections between discrete and continuous diffusion models. 16

• It would also be interesting to determine whether higher-order samplers can further improve the dependence on accuracy or intrinsic complexity. • Finally, characterizing the sample complexity required for accurate score estimation would be an important step toward a more unified theory of sampling with discrete diffusion models.

Acknowledgements This work is supported in part by Wharton Dean’s Research Fund, the NSF grants CCF-2106778, CCF2418156 and CAREER award DMS-2143215. This work is also supported by the NSF under Cooperative Agreement No. 2433450.

A

Details on the remasking process

In this section, we provide explicit expressions for the probabilities used in the construction of the remasking process. We emphasize that the exact form of these expressions is not used in the proofs (with the exception of νpt, cq, for c P rSs Y tREMASKu, used in the proofs of Lemma 6 and Proposition 4), and are given here for completeness. ? Lemma 2. Consider the remasking process with parameter pM and let ρ “ 1 ´ pM . Recall νpt, cq from Eqn. (6). We have $ ´t e ’ for b P rSs, & S pcoshptρq ´ 1q, 1 ´t νpt, bq “ ˆ e ρ sinhptρq, for b “ REMASK, 1 ´ e´t ’ % ´t 1 ´ e pcoshptρq ` ρ sinhptρqq , for b “ MASK. Proof. We assume 0 ă pM ă 1 with the case pM “ 1 interpreted by continuity from pM Ñ 1. To compute νpt, bq, observe that it is enough to study a CTMC where all rSs states are represented by a single state, and the rate matrix is as follows: ¨ ˛ ´1 1 ´ pM pM ´1 0 ‚. R“˝ 1 0 0 0 ¯ ´ ´1 1 ´ pM be the top-left submatrix of The order of the three states is rSs, REMASK, MASK. Let B “ 1 ´1 R and observe that the corresponding part of Rk equals B k . Next, we have ˙k 8 k ˆ ÿ t 0 1 ´ pM t e expptBq “ 0 k! 1 k“0 8 ÿ

ˆ ˙ 8 ÿ t2m`1 t2m m m 0 1 ´ pM p1 ´ pM q I ` p1 ´ pM q “ 1 0 p2mq! p2m ` 1q! m“0 m“0 ˆ ˙ a a 1 0 1 ´ pM sinhpt 1 ´ pM q “ coshpt 1 ´ pM qI ` ? , 1 0 1 ´ pM ? and we compute for b P rSs (recall that ρ “ 1 ´ pM ), 1 PrpXti P rSs | jump on r0, tsq S ` ˘ 1 “ PrpXti P rSsq ´ PrpXti P rSs and no jump on r0, ts ´t p1 ´ e qS ` ˘ 1 pexp ptBqq11 ´ e´t “ ´t p1 ´ e qS e´t “ pcoshptρq ´ 1q . p1 ´ e´t qS

νpt, bq “

17

(19)

Similarly, νpt, REMASKq “ PrpXti “ REMASK | jump on r0, tsq “

e´t ρ sinhptρq expptBq12 “ . 1 ´ e´t 1 ´ e´t

Finally, νpt, MASKq follows from νpt, MASKq ` νpt, REMASKq `

ÿ

νpt, bq “ 1.

bPrSs

? Lemma 3. Consider the remasking process with parameter pM and let ρ “ 1 ´ pM . Let 0 ď u ă ℓ with ∆ “ ℓ ´ u ą 0 and a, b P V. Then, $ ´∆ ’ e´∆ Ita “ bu ` e S pcoshp∆ρq ´ 1q if a, b P rSs ’ ’ ’ ’ ’ e´∆ ρ sinhp∆ρq if a P rSs, b “ REMASK ’ ’ ’ ´∆ ’ ’ 1 ´ e pcoshp∆ρq ` ρ sinhp∆ρqq if a P rSs, b “ MASK, ’ ’ ’ & e´∆ sinhp∆ρq if a “ REMASK, b P rSs, S ρ PrpXℓi “ b | Xui “ aq “ ´∆ e coshp∆ρq if a “ b “ REMASK, ’ ’ ¯ ´ ’ ’ ’1 ´ e´∆ coshp∆ρq ` sinhp∆ρq ’ if a “ REMASK, b “ MASK ’ ρ ’ ’ ’ ’ 1 if a “ b “ MASK ’ ’ ’ %0 otherwise. Proof. The proof follows from Eqn. (19). ? pi pbq defined Lemma 4. Consider the remasking process with parameter pM and let ρ “ 1 ´ pM . Recall µ in Eqn. (11): ´ ¯ x i “ b | X ´i “ x´i qe´u ` νpu, bqp1 ´ e´u q . µ pi pbq 9 PrpXℓi “ xiℓ | Xui “ bq PrpX 0 ℓ ℓ Then, for xiℓ P rSs, $´ ¯´ ¯ ’ x i “ b | X ´i “ x´i q ` coshpuρq´1 Itb “ xi u ` coshp∆ρq´1 , ’ PrpX 0 ℓ & ℓ ℓ S S 1 µ pi pbq 9 sinhpuρq sinhp∆ρq, S ’ ’ %0,

if b P rSs, if b “ REMASK, if b “ MASK.

For xiℓ “ REMASK,

µ pi pbq 9

$´ ¯ ’ x i “ b | X ´i “ x´i q ` coshpuρq´1 ρ sinhp∆ρq, ’ PrpX 0 & ℓ ℓ S

if b P rSs,

ρ sinhpuρq coshp∆ρq, ’ ’ %0,

if b “ REMASK, if b “ MASK.

For xiℓ “ MASK, ¯ ´ $ ˘ coshpuρq´1 ` ´i ´i ´u x i ´∆ ’ pcoshp∆ρq ` ρ sinhp∆ρqq e PrpX “ b | X “ x q ` 1 ´ e , ’ 0 ℓ ℓ S & ´ ´ ¯¯ pi pbq 9 e´u ρ sinhpuρq 1 ´ e´∆ coshp∆ρq ` sinhp∆ρq , µ ρ ’ ’ % 1 ´ e´u pcoshpuρq ` ρ sinhpuρqq ,

if b P rSs, if b “ REMASK, if b “ MASK.

Normalization factor should be computed separately for each of the three cases: xiℓ P rSs, xiℓ “ REMASK, and xiℓ “ MASK. Proof. The proof follows from Lemmas 2 and 3. 18

B

Technical preparations

This section contains results that are used in the proofs. Importantly, all results here concern only the forward process, and not the reverse process. Recall the definitions of the total correlation and the dual total correlation: for a random vector X “ pX 1 , . . . , X d q, ÿ ÿ HpX i q ´ HpXq and DTCpXq :“ HpXq ´ HpX i | X ´i q. TCpXq :“ iPrds

iPrds

The following proposition is the basis of our main results, as it provides explicit expressions for the first and second partial derivatives of the mutual information. Proposition 5. Let pXt qtPr0,T s be a CTMC with rate matrices pQt qtPr0,T s as in Definition 1. Then, (i) DTCpXt q ` TCpXt q “

ÿ

IpXti ; Xt´i q;

iPrds

(ii) for i P rds, it satisfies “` ´i ˘ ‰ B Qv log PrpXti “ xit | ¨ q px´i IpXti ; Xv´i q “ Exi ,x´i v q ; t v Bv (iii) for i P rds, we have “` i ˘ i ‰ B B Qt log Prp ¨ | Xv´i “ x´i IpXti ; Xv´i q “ HpXti q ` Exi ,x´i v q pxt q ; v t Bt Bt (iv) ÿ B d “ IpXti ; Xv´i q DTCpXv q, Bv dv t“v

and

iPrds

ÿ B d “ TCpXt q; IpXti ; Xv´i q Bt dt v“t

iPrds

(v) for fixed x´i v , j, and b, let us define rpaq “

PrpXti “ a | Xv´i “ x´i v q . ´i ´i i PrpXt “ a | Xv “ xv dj bq

It then obeys ȷ „ ÿ ÿ B2 rpaq i ´i j i ´i IpXt ; Xv q “ Ex´i Qv pxv , bqQt pa, cq PrpXt “ a | xv dj bq ˆ rpaq log ` rpcq ´ rpaq . v BtBv rpcq j‰i a,b,c We use the following lemma to upper bound Bregman (Itakura-Saito) divergence by the KL divergence. Lemma 5. Let p, q P rα, 1s for some α ą 0. Then, p

´q q

´ 1 ´ log pq

p p p q log q ´ q ` 1

¯ď

1 1 ď . minpp, qq α

The case p “ q is interpreted by continuity. The next definition introduces F pt, bq, which informally quantifies the probability mass that is moved to the state b at time t, given that a jump to the state b appeared. We recall that, in the uniform process case, this includes a possible self-loop jump b Ñ b.

19

Definition 3. Let t ě 0 and b P V. Define F : Rě0 ˆ V Ñ r0, 1s as follows: ÿ d Qpa, bq PrpXti “ aq. PrpXti “ bq ` PrpXti “ bq “ dt aPV

Fpt, bq :“

Lemma 6. Let 0 ă t ď T and b P rSs. Recall νpt, bq :“ PrpXti “ b | D a jump at the i-th coordinate on r0, tsq. Then, 1. for the uniform process, F pt, bq “ 1{S , 2. for the remasking process, Fpt, bq “

a e´t a 1 ´ pM sinhpt 1 ´ pM q. S

Consequently, for both processes, νpt, bq Á Fpt, bq. Proof. For the uniform process, the result is immediate, as Qpa, bq “ 1{S for all a, b P rSs which implies F pt, bq “ 1{S for all b P rSs. For the remasking process, for b P rSs, the only non-zero element is QpREMASK, bq “ 1{S , which gives Fpt, bq “

PrpXti “ REMASKq . S

(20)

Using Lemma 2, we compute PrpXti “ REMASKq “ νpt, REMASKqp1 ´ e´t q “ e´t and thus

a a 1 ´ pM sinhpt 1 ´ pM q,

? νpt, bq coshpt 1 ´ pM q ´ 1 ? “? “: f ptq. Fpt, bq 1 ´ pM p1 ´ e´t q sinhpt 1 ´ pM q

Observe that f ptq is a strictly increasing function and limtÑ0 f ptq “ 1{2. This shows that νpt, bq{Fpt, bq ě 1{2 and concludes the proof. Importantly, for both uniform and remasking processes, F pt, bq quantifies not only the marginal probability mass (averaged over the initial state X0i ) moved to the state b, but also when conditioning on a specific value X0i “ c or on the context Xu´i “ x´i at any time u, as the following lemma shows. Lemma 7. Consider the uniform or remasking process. For i P rds, b P V, c in the support of X0i , and t ě 0, we have ÿ Qpa, bq PrpXti “ a | X0i “ cq “ F pt, bq. (21) aPV

Consequently, for i P rds, b P V, x P V d , and t, u ě 0, we have ÿ Qpa, bq PrpXti “ a | Xu´i “ x´i q “ Fpt, bq,

(22)

aPV

x : as long as PrpXu´i “ x´i q ą 0. Furthermore, the same holds for Prp¨q ÿ ÿ x i “ a | X i “ cq “ x i “ a | X ´i “ x´i q. Fpt, bq “ Qpa, bqPrpX Qpa, bqPrpX t 0 t u aPV

C

Proof of our main results

C.1

Proof of Theorem 1

aPV

Proposition 1 shows that a sample obtained from Algorithm 1 has the same distribution as the one obtained p t defined by Eqn. (10). In the following we analyze this CTMC. Using Proposition 2, we using a CTMC Q have Nÿ ´1 Nÿ ´1 pkq KLpqT ´tN }poutput q ď KLpqT }qnoise q ` Lpkq Ldiscr . (23) approx ` k“0

20

k“0

In view of Proposition 3, the discretization error satisfies Nÿ ´1

pkq

Ldiscr À

k“0

Nÿ ´1

tk`1 ´ tk pDTCpXT ´tk`1 q ´ DTCpXT ´tk qq minp1, T ´ tk`1 q k“0

ďκ

Nÿ ´1

pDTCpXT ´tk`1 q ´ DTCpXT ´tk qq

(24)

k“0

ď κDTCpXT ´tN q ď κDTCpX0 q. Next, Proposition 4 gives Nÿ ´1

Lpkq approx ď

Nÿ ´1

ptk`1 ´ tk qLSE pT ´ tk , spT ´tk , sT ´tk q ď εscore ,

(25)

k“0

k“0

where the last inequality follows from Assumption 1. Collecting Eqns. (23), (24) and (25) concludes the proof of Eqn. (14). ´ ¯ δ ´1 Next, observe that under our condition on the step size, we can pick κ “ O T `log , where δ “ T ´tN N is the early stopping parameter. Let qnoise “ µbd , where: µpbq “ Itb P rSsu We have

e´T ` νpT, bqp1 ´ e´T q. S

˜

¸

KLpqT }qnoise q “ KL qT

â

qTi

iPrds

ÿ `

KLpqTi }µq,

where qTi is the i-th marginal of qT . Let X, Y „ q0,T . By the convexity of the KL divergence, ` ˘ KLpqTi }µq ď e´T KL q0i }UnifprSsq “ e´T plog S ´ HpXi qq . We also have

˜ KL qT

(26)

iPrds

(27)

¸ â

qTi

“ TCpY q

iPrds

ÿ “

HpYi q ´ HpY q (28)

iPrds

ÿ ď

HpYi q ´ HpY | Xq

iPrds

ÿ “

pHpYi q ´ HpYi | Xi qq ,

iPrds

where the last line follows as the coordinates of Y are conditionally independent given X. For fixed i P rds, HpYi q ´ HpYi | Xi q “ IpXi ; Yi q ď IpXi ; Yi , Bi q,

(29)

for Bi “ Itjump at the i-th coordinate on [0, T]u with Bi „ Bernp1 ´ e´T q. As Bi KK Xi , we obtain IpXi ; Yi , Bi q “ IpXi ; Yi | Bi q “ e´T HpXi q. Collecting Eqns. (26) to (30) gives KLpqT }qnoise q ď e´T

ÿ

HpXi q ` e´T

iPrds

ÿ iPrds

˘ Choosing T “ O logpε´1 d log Sq concludes the proof. `

21

plog S ´ HpXi qq “ e´T d log S.

(30)

C.2

Proof of Proposition 1

Fix a discretization interval ru, ℓs and condition on the current state xℓ . By Eqn. (10), for a transition that changes only coordinate i, the rate p t px, x di bq Q ´i depends on the frozen context x´i ℓ , but not on the evolving coordinates x . Consequently, on the interval ru, ℓs, the approximate CTMC decomposes into d independent one-dimensional CTMCs. It therefore suffices to verify that, for each i P rds, the update in Algorithm 1 coincides with the transition of the corresponding one-dimensional CTMC. Fix i P rds and x´i P V d´1 . To make the argument explicit, introduce an auxiliary one-dimensional ℓ i forward process pZt qtPr0,ℓs with rate matrix Q and initial distribution

` ˘ x X i “ b | X ´i “ x´i . PrpZ0i “ bq “ Pr 0 ℓ ℓ By construction and Eqn. (9), its marginal at time t is ` ˘ x X i “ b | X ´i “ x´i . PrpZti “ bq “ Pr t ℓ ℓ Hence, the score in Eqn. (10) is precisely the score of this auxiliary one-dimensional process: spt px di b, xq “

PrpZti “ bq . PrpZti “ xi q

p t to coordinate i is It follows from the standard time-reversal formula for CTMCs that the restriction of Q exactly the reverse generator of pZti q. We initialize this reverse process at the observed endpoint Zℓi “ xiℓ . Therefore, its distribution at time u is PrpZui “ b | Zℓi “ xiℓ q. By Bayes’ rule, PrpZui “ b | Zℓi “ xiℓ q 9 PrpZℓi “ xiℓ | Zui “ bq PrpZui “ bq ” ` ı ˘ x X i “ b | X ´i “ x´i e´u ` νpu, bqp1 ´ e´u q , “ PrpXℓi “ xiℓ | Xui “ bq Pr 0 ℓ ℓ where the last equality follows from Eqn. (9). The right-hand side is exactly µi pbq in Eqn. (11). Thus, the update of coordinate i in Algorithm 1 has the same law as the corresponding coordinate of the CTMC pt . generated by Q Since the coordinate processes are independent on each discretization interval conditional on xℓ , Algorithm 1, which samples all coordinates independently in parallel, has the same transition kernel from xℓ to p t . Applying this argument successively over all discretization intervals xu as the CTMC with rate matrices Q proves the claim.

C.3

Proof of Proposition 2

For the divergence Dpα, γq “ αγ ´ 1 ´ log αγ , a straightforward calculation shows, for any α, β, γ ą 0: γDpα, γq “ βDpα, βq ` γDpβ, γq ` pβ ´ γq log

α . β

p T ´t pxt , yq, Fix xt P V d and let y “ x di b for i P rds, b P V, such that Qt py, xt q ą 0. We pick α “ Q Ð r T ´t pxt , yq, and γ “ QT ´t pxt , yq and note that both Q p T ´t pxt , yq and Q r T ´t pxt , yq are functions of xℓ β“Q

and xit but not of x´i t . The law of total expectation gives « ff ´ ¯ p T ´t pxt , yq Ð Q r T ´t pxt , yq ´ QT ´t pxt , yq log Exℓ ,xt Q r T ´t pxt , yq Q « ff ´ ”Ð ı¯ p T ´t pxt , yq Q r T ´t pxt , yq ´ E ´i QT ´t pxt , yq | Xℓ “ xℓ , X i “ xi log “ Exℓ ,xit Q “ 0, t t xt r T ´t pxt , yq Q 22

ı ”Ð r T ´t pxt , yq “ E ´i QT ´t pxt , yq | Xℓ “ xℓ , X i “ xi 2 . as by Lemma 1, Q t t xt Putting these together, we have obtained the following decomposition: ´ ¯ı ”Ð Ð p T ´t pxt , yq, QT ´t pxt , yq Exℓ ,xt QT ´t pxt , yqD Q ´ ¯ı ” ´ ¯ı ”Ð Ð r T ´t pxt , yq, QT ´t pxt , yq . r T ´t pxt , yqD Q p T ´t pxt , yq, Q r T ´t pxt , yq ` Ex ,x QT ´t pxt , yqD Q “ Exℓ ,xt Q ℓ t Using Girsanov’s change-of-measure theorem (Campbell et al. (2022)), we arrive at żT ´ ¯ ÿ Ð Ð p T ´t pxt , yq, QT ´t pxt , yq dt E QT ´t pxt , yqD Q KLpqT ´tN }poutput q ď KLpqT }qnoise q ` tN

“ KLpqT }qnoise q `

y‰xt

Nÿ ´1 ż T ´tk k“0

T ´tk`1

y‰xt

Nÿ ´1 ż T ´tk

´ ¯ Ð Ð r T ´t pxt , yq, QT ´t pxt , yq dt QT ´t pxt , yqD Q

ÿ

`

E k“0

“ KLpqT }qnoise q `

´ ¯ r T ´t pxt , yqD Q p T ´t pxt , yq, Q r T ´t pxt , yq dt Q

ÿ E

Nÿ ´1

T ´tk`1

Lpkq approx `

k“0

y‰xt Nÿ ´1

pkq

Ldiscr ,

k“0

which concludes the proof.

C.4

Discretization error control

C.4.1

Proof of Theorem 2

Fix t P ru, ℓs. Using Lemma 1, we arrive at ´ ¯ ÿ Ð Ð r T ´t pxt , yq, QT ´t pxt , yq Ext ,xℓ QT ´t pxt , yqD Q y‰xt

»

¨ Ð

ÿ

Ð

r T ´t pxt , yq ´ QT ´t pxt , yq ´ QT ´t pxt , yq log –Q

“ Ext ,xℓ

˛fi r ˝ QT ´t pxt , yq ‚fl Ð

QT ´t pxt , yq

y‰xt Ð

ÿ Ð Q pxt , yq “ Ext ,xℓ QT ´t pxt , yq log T ´t r QT ´t pxt , yq y‰xt ÿ st py, xt q “ Ext ,xℓ Qt py, xt qst py, xt q log rt py, xt q s y‰x t

ÿ ÿ “ Ext ,xℓ

Qt pb, xit qst pxt di b, xt q log

iPrds bPV

st pxt di b, xt q . srt pxt di b, xt q

In the last line, we add the term corresponding to b “ xit , as both st pxt , xt q and srt pxt , xt q equal to 1, and therefore this term equals to 0. Furthermore, note that the overall expression does not depend on xiℓ , as srt pxt di b, xt q “

PrpXti “ b | Xℓ´i “ x´i ℓ q . PrpXti “ xit | Xℓ´i “ x´i ℓ q

We obtain ÿ ÿ ÿ ÿ st pxt di b, xt q Prpxt di bq st pxt di b, xt q Ext ,xℓ Qt pb, xit qst pxt di b, xt q log “ Ext ,x´i Qt pb, xit q log , ℓ srt pxt di b, xt q Prpxt q srt pxt di b, xt q bPV bPV iPrds

iPrds

2 As X ´i K K Xℓi | Xti from the Markovian property, conditioning on Xℓ “ xℓ and Xti “ xit is equivalent to conditioning on t ´i ´i Xℓ “ xℓ and Xti “ xit .

23

which, after relabeling xit Ø b and using that Xℓ´i KK Xti | Xt´i , gives ÿ ÿ

Ext ,x´i Qt pb, xit q ℓ

iPrds bPV

ÿ ÿ st pxt di b, xt q Prpxt di bq st pxt , xt di bq log “ Ext ,x´i Qt pxit , bq log . ℓ srt pxt di b, xt q srt pxt , xt di bq Prpxt q bPV iPrds

We continue as follows: st pxt , xt di bq rt pxt , xt di bq s iPrds bPV „ ȷ ÿ ÿ PrpXti “ xit | Xℓ´i “ x´i PrpXti “ xit | Xt´i “ x´i t q ℓ q “ Ext ,x´i Qt pxit , bq log ´ log ℓ PrpXti “ b | Xt´i “ x´i PrpXti “ b | Xℓ´i “ x´i t q ℓ q iPrds bPV ˆˆ ˆ ˙ ˙ ˙ ÿ B B “ ´ IpXti ; Xv´i q IpXti ; Xv´i q Bt Bt v“ℓ v“t iPrds ż ℓ 2 ÿ B “ IpXti ; Xv´i qdv, t BvBt ÿ ÿ

Ext ,x´i ℓ

Qt pxit , bq log

iPrds

where we used Proposition 5 (iii) in the third line. This proves Eqn. (17), as żℓ pkq Ldiscr “

¯ ´ ÿ ż ℓ ż ℓ B2 ÿ Ð Ð r IpXti ; Xv´i qdvdt. Exℓ ,xt QT ´t pxt , yqD QT ´t pxt , yq, QT ´t pxt , yq dt “ u u t BtBv y‰x iPrds

t

C.4.2

Proof of Proposition 3

Theorem 2 together with Fubini’s theorem imply: ÿ ż ℓ ż ℓ B2 IpXti ; Xv´i qdvdt BtBv iPrds u t ÿ ż ℓ ż v B2 “ IpXti ; Xv´i qdtdv BtBv u u iPrds ÿ żℓ B B ´ dv. IpXti ; Xv´i q IpXti ; Xv´i q “ Bv u Bv t“v t“u

pkq Ldiscr “

iPrds

Next, Proposition 6 gives ´

B B IpXti ; Xv´i q “ ´ p1 ` Opκqq IpXti ; Xv´i q , Bv Bv t“u t“v

which implies using Eqn. (31), pkq

Ldiscr À ´κ

żℓ ÿ u iPrds

B From Proposition 5 (iv) we have Bv IpXti ; Xv´i q

t“v

B IpXti ; Xv´i q dv. Bv t“v d DTCpXv q, therefore “ dv

żℓ pkq

Ldiscr À ´κ

d DTCpXv qdv “ κpDTCpXu q ´ DTCpXℓ qq. u dv

This concludes the proof.

24

(31)

C.4.3

Statement and proof of Proposition 6

Proposition 6. Let pXt qtPr0,T s be the uniform or remasking process. Then, for i P rds and t, v ą 0, ˆ ˙ 1 B B2 i ´i IpXt ; Xv q À ´ 1 ` IpXti ; Xv´i q. ´t BtBv 1´e Bv

(32)

Furthermore, if v ě u with v ´ u ď κ minp1, uq, ´

B B IpXti ; Xv´i q “ ´ p1 ` Opκqq IpXti ; Xv´i q . Bv Bv t“u t“v

Proof. Recall Proposition 5 (ii) and (v): ÿ ÿ B IpXti ; Xv´i q “ ´Ex´i Qpxjv , bq PrpXti “ a | x´i v dj bq rrpaq log rpaq ´ rpaq ` 1s , v Bv j‰i a,bPV ȷ „ ÿ ÿ B2 rpaq i ´i j i ´i IpXt ; Xv q “ Ex´i ` rpcq ´ rpaq , Qpxv , bqQpa, cq PrpXt “ a | xv dj bq rpaq log v BtBv rpcq j‰i a,b,cPV ´i i ´i i where, for fixed x´i v , j, and b, we recall rpaq “ PrpXt “ a | Xv “ xv q { PrpXt “ a | Xv “ xv dj bq. We decompose

rpaq log

´ ¯ ´ ¯ rpaq ` rpcq ´ rpaq “ rpaq log rpaq ´ rpaq ` 1 ` rpcq ´ 1 ´ rpaq log rpcq , rpcq

which, together with Lemma 7, immediately shows (relabeling a Ø c in the second line) B B2 IpXti ; Xv´i q ` IpXti ; Xv´i q BtBv Bv ÿ ÿ “ Ex´i Qpxjv , bqQpa, cq PrpXti “ a | x´i v dj bq rrpcq ´ 1 ´ rpaq log rpcqs v j‰i a,b,cPV

ÿ ÿ “ Ex´i v

Qpxjv , bqFpt, aq rrpaq ´ 1 ´ log rpaqs .

j‰i a,bPV

Next, we apply Lemmas 5 and 6, and obtain ÿ ÿ B2 B Qpxjv , bqFpt, aq rrpaq ´ 1 ´ log rpaqs IpXti ; Xv´i q “ ´ IpXti ; Xv´i q ` Ex´i v BtBv Bv j‰i a,bPV B IpXti ; Xv´i q Bv ÿ ÿ 1 ´i ` E Qpxjv , bq PrpXti “ a | x´i v dj bq rrpaq log rpaq ´ rpaq ` 1s 1 ´ e´t xv j‰i a,bPV ˆ ˙ 1 B “´ 1` IpXti ; Xv´i q. ´t 1´e Bv À´

This finishes the proof of Eqn. (32). Next, we use this bound in the following variation of Grönwall’s inequality: ˆ ˙ ˙ żvˆ B 1 B IpXti ; Xv´i q dt IpXti ; Xv´i q ď exp ´C 1` , ´t Bv 1 ´ e Bv u t“v t“u for some universal constant C ą 0. This gives B ´ IpXti ; Xv´i q ď ´eCpv´uq Bv t“u

ˆ v ˙C e ´1 B . IpXti ; Xv´i q eu ´ 1 Bv t“v

25

Under the condition that v ´ u ď κ minp1, uq, we have that e

Cpv´uq

ˆ v ˆ ˙C ˆ ˙˙C ˆ ˆ ˙˙ e ´1 ev´u ´ 1 κueκ v´u v´u ¨ “ e e ` u “ p1 ` Opκqq 1 ` O κ ` “ 1 ` Opκq, eu ´ 1 e ´1 u

which leads to ´

B B IpXti ; Xv´i q “ ´ p1 ` Opκqq IpXti ; Xv´i q . Bv Bv t“u t“v

This concludes the proof.

C.5

Approximation error control

C.5.1

Proof of Theorem 3

To facilitate the proof of Theorem 3, let us introduce the following proposition. Proposition 7. Fix k P t0, . . . , N ´ 1u in the time discretization 0 “ t0 ă t1 ă . . . ă tN ă T . Consider the uniform or remasking process and recall Fpt, bq from Definition 3. Letting u “ T ´ tk`1 and ℓ “ T ´ tk , we have « ÿ ÿżℓ PrpXti “ b | x´i pkq ℓ q Lapprox “ Ex´i PrpXti “ b | x´i q log ℓ i ℓ x t “ b | x´i q PrpX ℓ iPrds bPV u ˜ ¸ff ´i i PrpXt “ b | xℓ q PrpXti “ b | x´i ℓ q ` F pt, bq ´ 1 ´ log dt. x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX ℓ

ℓ

Proposition 7 leads to the following control of the approximation error. The proofs of these two results are deferred to later in the section. Corollary 1. Fix k P t0, . . . , N ´ 1u in the time discretization 0 “ t0 ă t1 ă . . . ă tN ă T . Consider the uniform or remasking process and let u “ T ´ tk`1 , ℓ “ T ´ tk . We have Lpkq approx À

ˇ ¯ › ℓ ´ u ÿ ´u ´ ˇ pX0i ˇ Xℓ´i . e KL µX0i › µ minp1, uq iPrds

In view of Corollary 1, we arrive at Nÿ ´1

Lpkq approx À

k“0

Nÿ ´1

ˇ ´ ¯ ÿ › tk`1 ´ tk ˇ pX0i ˇ XT´i´tk e´pT ´tk`1 q KL µX0i › µ minp1, T ´ tk`1 q k“0 iPrds

ďκ

Nÿ ´1 ÿ

ˇ ´ ¯ › ˇ pX0i ˇ XT´i´tk , e´pT ´tk`1 q KL µX0i › µ

k“0 iPrds

which concludes the proof of Theorem 3. Proof of Proposition 7. Recall that for t P ru, ℓs, we have srt px di b, xq “

PrpXti “ b | Xℓ´i “ x´i q PrpXti “ xi | Xℓ´i “ x´i q

and

spt px di b, xq “

x i “ b | X ´i “ x´i q PrpX t ℓ . x ti “ xi | X ´i “ x´i q PrpX ℓ

Here and below, for clarity of the exposition, we omit stochastic process notation Xti “ ¨ , Xvi “ ¨ , etc., when i i time and dimension indices are clear from the context. Thus, e.g., Prpxiℓ | x´i ℓ q stands for PrpXℓ “ xℓ |

26

Xℓ´i “ x´i ℓ q. Using Eqn. (15), we write żℓ Lpkq approx “

ÿ Ext ,xℓ

u

y‰xt

żℓ “

´ ¯ r T ´t pxt , yqD Q p T ´t pxt , yq, Q r T ´t pxt , yq dt Q

ÿ Ext ,xℓ

u

st py, xt qD pp st py, xt q, srt py, xt qq dt Qt py, xt qr

y‰xt

żℓ ÿ

(33) ÿ

“ u iPrds

i Prpxiℓ | x´i ℓ q Prpxt | xℓ q

Ex´i ℓ

xiℓ ,xit PV

ÿ ˆ

Qpb, xit qr st pxt di b, xt qD pp st pxt di b, xt q, srt pxt di b, xt qq dt.

bPV

The terms corresponding to b “ xit are added as their contribution isř0 by Dpp st pxt , xt q, srt pxt , xt qq “ Dp1, 1q “ 0. We also recall that we define diagonal elements of Q such that bPV Qpa, bq “ 1, for all a P V. Using the equality Prpxit | xℓ q Prpxit , xℓ q Prpx´i Prpxℓ | xit q 1 Prpxiℓ | xit q ℓ q “ ¨ “ ¨ “ , ´i i i Prpxℓ q Prpxit | x´i Prpxit , x´i Prpx´i Prpxiℓ | x´i ℓ q ℓ q ℓ | xt q Prpxℓ | xℓ q ℓ q we can rewrite i i st pxt di b, xt q “ Prpxiℓ | x´i Prpxiℓ | x´i ℓ q Prpxt | xℓ qr ℓ q Prpxt | xℓ q

PrpXti “ b | x´i i i ℓ q “ PrpXti “ b | x´i ℓ q Prpxℓ | xt q. ´i i Prpxt | xℓ q

Plugging this into Eqn. (33), we continue: Lpkq approx “

żℓ ÿ

ÿ

u iPrds

ℓ

xiℓ ,xit ,bPV

żℓ ÿ

ÿ

“ u iPrds

i i Qpb, xit q Prpb | x´i st pxt di b, xt q, srt pxt di b, xt qq dt ℓ q Prpxℓ | xt qD pp

Ex´i

Ex´i ℓ

st pxt di b, xt q, srt pxt di b, xt qq dt Qpb, xit q Prpb | x´i ℓ qD pp

xit ,bPV

Next, we use: x i | x´i q x i “ b | x´i q x i “ b | x´i q Prpxi | x´i q Prpx PrpX PrpX t t t t ℓ ℓ ℓ ℓ `log . ´1´log i “ b | x´i q i | x´i q x it | x´i q PrpXti “ b | x´i q PrpX Prpx Prpx t t ℓ ℓ ℓ ℓ (34) şℓ ř pkq Splitting Eqn. (34) into four parts, we write Lapprox “ u iPrds rI1 ` I2 ` I3 ` I4 s dt, where, for fixed t P ru, ℓs and i P rds:

st pxt di b, xt q, srt pxt di b, xt qq “ D pp

ÿ

I1 “ Ex´i ℓ

Qpb, xit q PrpXti “ b | x´i ℓ q

xit ,bPV

x i “ b | x´i q Prpxi | x´i q PrpX t t ℓ ℓ ´i ´i i i x Prpxt | x q PrpXt “ b | x q ℓ

ℓ

x i “ b | x´i q Prpxi | x´i q PrpX t t ℓ ℓ ´i i ℓ x Prpxt | xℓ q xit ,bPV ˜ ¸ ÿ Prpxi | x´i q ÿ t ℓ x i “ b | x´i q , “ Ex´i Qpb, xit qPrpX t ℓ ℓ x it | x´i q Prpx i bPV ℓ xt PV ÿ I2 “ ´Ex´i Qpb, xit q PrpXti “ b | x´i ℓ q, ÿ

Qpb, xit q

“ Ex´i

ℓ

xit ,bPV

¨ I3 “ Ex´i

˛ ÿ

ÿ ˝

ℓ

bPV

Qpb, xit q‚PrpXti “ b | x´i ℓ q log

xit PV

PrpXti “ b | x´i ℓ q , i x PrpXt “ b | x´i q ℓ

27

and

I4 “ ´Ex´i ℓ

ÿ xit PV

Prpxit | x´i ℓ q log i x Prpxt | x´i q

˜

¸ ÿ

Qpb, xit q PrpXti “ b | x´i ℓ q

.

bPV

ℓ

Collecting I1 ` I2 ` I4 gives: «˜

¸ PrpXti “ b | x´i ℓ q I1 ` I2 ` I4 “ Ex´i ℓ x ti “ b | x´i q PrpX aPV bPV ℓ ˜ ¸˜ ¸ff ÿ PrpXti “ b | x´i ´i i ℓ q ´ Qpa, bq PrpXt “ a | xℓ q 1 ` log x ti “ b | x´i q PrpX aPV ℓ ¸ ˜ ´i ´i i i ÿ PrpXt “ b | xℓ q PrpXt “ b | xℓ q ´ 1 ´ log , “ Ex´i Fpt, bq ´i i ℓ x x ti “ b | x´i q PrpXt “ b | x q PrpX ÿ

¸˜

ÿ

x i “ a | x´i q Qpa, bqPrpX t ℓ

bPV

ℓ

ℓ

where we used Lemma 7 in the last line. This proves that ˜ ¸ff żℓ ÿ« ÿ PrpXti “ b | x´i PrpXti “ b | x´i pkq ℓ q ℓ q I3 ` Ex´i Fpt, bq Lapprox “ ´ 1 ´ log dt. ℓ x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX u bPV

iPrds

ℓ

Observe that by our definition of the diagonal elements of Q, we have

I3 “ Ex´i

ÿ ˝

ℓ

bPV

i xit Qpb, xt q “ 1, and therefore,

˛

¨ ÿ

ℓ

ř

Qpb, xit q‚PrpXti “ b | x´i ℓ q log

xit PV

PrpXti “ b | x´i ℓ q x ti “ b | x´i q PrpX ℓ

ÿ PrpXti “ b | x´i ℓ q PrpXti “ b | x´i “ Ex´i q log . ℓ ´i i ℓ x t “b|x q PrpX bPV ℓ We obtain Lpkq approx “

ÿ ÿżℓ iPrds bPV

u

« Ex´i PrpXti “ b | x´i ℓ q log ℓ

PrpXti “ b | x´i ℓ q i x PrpXt “ b | x´i q ℓ

˜ ` F pt, bq

PrpXti “ b | x´i PrpXti “ b | x´i ℓ q ℓ q ´ 1 ´ log ´i i i x x PrpXt “ b | x q PrpXt “ b | x´i q ℓ

¸ff dt,

ℓ

which concludes the proof. Proof of Corollary 1. Proposition 7 shows: « ÿ ÿżℓ PrpXti “ b | x´i pkq ℓ q Lapprox “ Ex´i PrpXti “ b | x´i q log ℓ i ℓ x t “ b | x´i q PrpX ℓ iPrds bPV u ˜ ¸ff ´i i PrpXt “ b | xℓ q PrpXti “ b | x´i ℓ q ` F pt, bq ´ 1 ´ log dt. x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX ℓ

ℓ

Without loss of generality, we assume b P rSs. Indeed, in the uniform process, V “ rSs, and in the remasking process, as the target distribution is not supported on tMASK, REMASKu, we have that PrpXti “ MASK | ´i x i x´i ℓ q “ νpt, MASKq “ PrpXt “ MASK | xℓ q. The same holds for b “ REMASK. Thus, terms in the pkq sum corresponding to b P tMASK, REMASKu do not contribute to Lapprox and we may assume b P rSs. Using Eqns. (7) and (9) together with Lemma 6 we have ´ ¯ ´i ´t ´t x i min PrpXti “ b | x´i ℓ q, PrpXt “ b | xℓ q ě νpt, bqp1 ´ e q Á F pt, bqp1 ´ e q, and proceed with Lemma 5 as follows: 28

˜

¸ PrpXti “ b | x´i PrpXti “ b | x´i ℓ q ℓ q Fpt, bq ´ 1 ´ log x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX bPV ℓ ℓ ˜ ¸ ´i i ÿ PrpX “ b | x q 1 t ´i ℓ x i À PrpXti “ b | x´i ´ PrpXti “ b | x´i ℓ q log x ℓ q ` PrpXt “ b | xℓ q 1 ´ e´t PrpXti “ b | x´i q bPV ℓ ˜ ¸ ´i i ÿ PrpX “ b | x 1 t ´i i ℓ q PrpXt “ b | xℓ q log . “ x ti “ b | x´i q 1 ´ e´t PrpX ÿ

bPV

ℓ

This implies that Lpkq approx À

żℓ ÿ u iPrds

PrpXti “ b | x´i 1 ´i i ℓ q PrpX “ b | x q log . t ℓ ´t i x 1´e PrpXt “ b | x´i q bPV ÿ

Ex´i ℓ

ℓ

Using again Eqns. (7) and (9) and the convexity of the KL divergence gives ÿ

PrpXti “ b | x´i ℓ q log

bPV

ÿ PrpXti “ b | x´i PrpX0i “ b | x´i ´i ´t i ℓ q ℓ q ď e , PrpX “ b | x q log 0 ℓ ´i ´i i i x x PrpXt “ b | x q PrpX “ b | x q bPV

ℓ

0

ℓ

which shows ¸ ˇ ˇ ¯ ´ ¯ ´ ÿ › › ℓ´u ÿ 1 ˇ ´i ˇ ´i › › i i i i µ p µ p KL µ X dt ď KL µ X . ˇ ˇ X0 X0 X0 X0 ℓ ℓ t eu ´ 1 u e ´1

˜ż ℓ Lpkq approx À

iPrds

iPrds

This concludes the proof. C.5.2

Proof of Proposition 4

Proposition 7 gives: Lpkq approx “ Exℓ „qℓ

ÿ żℓ ÿ iPrds u bPV

« PrpXti “ b | x´i ℓ q log

PrpXti “ b | x´i ℓ q i x PrpXt “ b | x´i q ℓ

˜ ` F pt, bq

PrpXti “ b | x´i PrpXti “ b | x´i ℓ q ℓ q ´ 1 ´ log ´i i i x t “b|x q x t “ b | x´i q PrpX PrpX ℓ ℓ

¸ff dt.

1 Uniform process. We first prove the result for the uniform process, ř where Fpt, bq “ {S . The first term is the KL divergence, which is the Bregman ř divergence for ϕpµq “ bPV µpbq log µpbq, and the second term is the Bregman divergence for ϕpµq “ ´ bPV log µpbq. Here, µ is a probability distribution over V. Let t P ru, ℓs. For a general Bregman divergence Dϕ with ϕ : Rn Ñ R, using Hadamard’s lemma, we have

ˆż 1 J

˙ 2

J

Dϕ px, yq :“ ϕpxq ´ ϕpyq ´ px ´ yq ∇ϕpyq “ px ´ yq

p1 ´ rq∇ ϕpy ` rpx ´ yqqdr px ´ yq. 0

For fixed t P ru, ℓs and i P rds this gives, using px log xq2 “ 1{x, f ptq :“

ÿ

PrpXti “ b | x´i ℓ q log

bPV

PrpXti “ b | x´i ℓ q x ti “ b | x´i q PrpX ℓ

¯2 ż 1 ÿ´ x i “ b | x´i q “ PrpXti “ b | x´i q ´ PrpX t ℓ ℓ

1´r

´i i x ti “ b | x´i q 0 r PrpXt “ b | xℓ q ` p1 ´ rqPrpX ℓ

bPV

29

dr,

and using p´ log xq2 “ 1{x2 : ˜ ¸ ´i i ÿ PrpX i “ b | x´i q PrpX “ b | x q t t ℓ ℓ ´ 1 ´ log gptq :“ x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX bPV ℓ ℓ ¯2 ż 1 ÿ´ 1´r ´i x i PrpXti “ b | x´i “ ´ ¯2 dr. ℓ q ´ PrpXt “ b | xℓ q 0 x ti “ b | x´i q bPV r PrpXti “ b | x´i q ` p1 ´ rq PrpX ℓ ℓ To upper bound f ptq with f pℓq, we need to upper bound, for fixed r P r0, 1s, 1 ´i x i r PrpXti “ b | x´i ℓ q ` p1 ´ rqPrpXt “ b | xℓ q

with

1 ´i x i r PrpXℓi “ b | x´i ℓ q ` p1 ´ rqPrpXℓ “ b | xℓ q

.

1 ´pℓ´tq ´pℓ´tq x i “ b | x´i q. Recall that PrpXℓi “ b | x´i PrpXti “ b | x´i q, similarly for PrpX ℓ ℓ q “ e ℓ q ` S p1 ´ e ℓ ´i ´i x i “ b | x q, we obtain: Letting α “ PrpXti “ b | xℓ q and β “ PrpX t ℓ ´pℓ´tq ´i x i q{S r PrpXℓi “ b | x´i e´pℓ´tq prα ` p1 ´ rqβq ` p1´e ℓ q ` p1 ´ rqPrpXℓ “ b | xℓ q “ ´i ´i x ti “ b | x q rα ` p1 ´ rqβ r PrpXti “ b | x q ` p1 ´ rqPrpX

ℓ

ℓ

“ e´pℓ´tq ` ´pℓ´tq

ďe “ where the inequality follows as α, β ě p1´e

´t

`

p1´e´pℓ´tq q{S

rα ` p1 ´ rqβ

(35)

p1´e´pℓ´tq q{S p1´e´t q{S

1 ´ e´ℓ , 1 ´ e´t

q{S . Since

´ ¯2 ´ ¯2 x i “ b | x´i q “ e2pℓ´tq PrpX i “ b | x´i q ´ PrpX x i “ b | x´i q , PrpXti “ b | x´i q ´ PrpX t ℓ ℓ ℓ ℓ ℓ ℓ we obtain ÿ bPV

PrpXti “ b | x´i ℓ q log

ℓ ÿ PrpXℓi “ b | x´i PrpXti “ b | x´i ´i i ℓ´t e ´ 1 ℓ q ℓ q PrpX “ b | x q log ď e . ℓ ℓ t´1 ´i i i x x e PrpXt “ b | x q PrpX “ b | x´i q bPV

ℓ

ℓ

ℓ

Squaring Eqn. (35), we can also upper bound gptq with gpℓq: ˜ ¸ ÿ PrpX i “ b | x´i q PrpXti “ b | x´i t ℓ ℓ q ´ 1 ´ log x ti “ b | x´i q x ti “ b | x´i q PrpX PrpX bPV ℓ ℓ ˜ ¸ ˆ ℓ ˙2 ÿ ´i i PrpXℓ “ b | xℓ q PrpXℓi “ b | x´i q e ´1 ℓ ď ´ 1 ´ log . x i “ b | x´i q x i “ b | x´i q et ´ 1 bPV PrpX PrpX ℓ ℓ ℓ ℓ Together with the bound ℓ ´ u ď κ minp1, uq this proves that for the uniform process, LSE pt, spt , srt q À LSE pℓ, spℓ , srℓ q, and as srℓ “ sℓ , gives żℓ Lpkq approx :“ Exℓ „qℓ

LSE pt, spt , srt qdt À pℓ ´ uqLSE pℓ, spℓ , sℓ q. u

Remasking process. Now, we focus on the remasking process, where by Lemma 2 for b P rSs, νpt, bq “

´ ¯ a e´t coshpt 1 ´ p q ´ 1 . M p1 ´ e´t qS

30

´v ´v Using that for all v ą 0 we have PrpXvi “ b | x´i PrpX0i “ b | x´i qνpv, bq, we express ℓ q“e ℓ q ` p1 ´ e ´pℓ´tq ´ℓ ´t ´pℓ´tq PrpXℓi “ b | x´i PrpXti “ b | x´i νpt, bq, ℓ q“e ℓ q ` p1 ´ e qνpℓ, bq ´ p1 ´ e qe

x i “ b | x´i q. Repeating the steps from Eqn. (35), we get (recall the notation and analogously for PrpX ℓ ℓ ´i i x i “ b | x´i q) α “ PrpXt “ b | xℓ q and β “ PrpX t ℓ ´i x i r PrpXℓi “ b | x´i ℓ q ` p1 ´ rqPrpXℓ “ b | xℓ q x ti “ b | x´i q r PrpXti “ b | x´i q ` p1 ´ rqPrpX ℓ

ℓ

e´pℓ´tq prα ` p1 ´ rqβq ` p1 ´ e´ℓ qνpℓ, bq ´ p1 ´ e´t qe´pℓ´tq νpt, bq “ rα ` p1 ´ rqβ “ e´pℓ´tq `

(36)

p1 ´ e´ℓ qνpℓ, bq ´ p1 ´ e´t qe´pℓ´tq νpt, bq . rα ` p1 ´ rqβ

If pM “ 1, we have that νpt, bq “ 0 for b P rSs and all t ě 0, thus the latter expression equals e´pℓ´tq ď 1. When pM ă 1, we continue Eqn. (36) as follows: ´i x i r PrpXℓi “ b | x´i ℓ q ` p1 ´ rqPrpXℓ “ b | xℓ q x ti “ b | x´i q r PrpXti “ b | x´i q ` p1 ´ rqPrpX ℓ

ℓ

p1 ´ e´ℓ qνpℓ, bq ´ p1 ´ e´t qe´pℓ´tq νpt, bq ď e´pℓ´tq ` p1 ´ e´t qνpt, bq ´ℓ p1 ´ e qνpℓ, bq “ , p1 ´ e´t qνpt, bq

(37)

Using the expression for νpt, bq for the remasking process, we obtain ? p1 ´ e´ℓ qνpℓ, bq e´ℓ pcoshpℓ 1 ´ pM q ´ 1q ? “ . p1 ´ e´t qνpt, bq e´t pcoshpt 1 ´ pM q ´ 1q As the ratio is strictly decreasing with respect to pM , we consider the case pM “ 0, which gives p1 ´ e´ℓ qνpℓ, bq “ p1 ´ e´t qνpt, bq

ˆ

1 ´ e´ℓ 1 ´ e´t

˙2 ď p1 ` κq2 À 1.

The rest of the proof follows closely the argument for the uniform process. This concludes the proof.

D

Proofs of results in Appendix B

D.1

Proof of Proposition 5

The first property follows immediately from the definition of the mutual information, the total correlation, and the dual total correlation. The second and third properties follow as: “` ´i ˘ ‰ B HpXti | Xv´i q “ ´Exi ,x´i Qv log Prpxit | ¨ q px´i v q , t v Bv “` i ˘ i ‰ B HpXti | Xv´i q “ ´Exi ,x´i Qt log Prp ¨ | x´i and v q pxt q , t v Bt IpXti ; Xv´i q “ HpXti q ´ HpXti | Xv´i q. For the fourth property, we proceed as follows:

31

(38)

ˆ ˙ d d 1 HpXv q “ Exv log dv dv PrpXv “ xv q ˙ ÿ ˆB 1 “ PrpXv “ xv q log Bv PrpXv “ xv q d xv PV

ÿ ÿ “ Exv

Qv pxiv , bq log

PrpXv “ xv q PrpXv “ xv di bq

Qv pxiv , bq log

PrpXvi “ xiv | Xv´i “ x´i v q . PrpXvi “ b | Xv´i “ x´i v q

iPrds bPV

ÿ ÿ “ Exv iPrds bPV

(39)

Similarly, it can be easily calculated that d HpXvi | Xv´i q dv ÿ ÿ “ Exv Qv pxjv , bq log jPrds bPV

PrpXvi “ xiv | Xv´i “ x´i v q PrpXvi “ pxv dj bqi | Xv´i “ pxv dj bq´i q

ÿÿ PrpXvi “ xiv | Xv´i “ x´i PrpXvi “ xiv | Xv´i “ x´i v q v q j ` E Q px , bq log . x v v v ´i ´i ´i ´i i i i PrpXv “ b | Xv “ xv q PrpXv “ xv | Xv “ xv dj bq j‰i bPV bPV (40) Taking these collectively yields ¨ ˛ ÿ d d ˝ DTCpXv q “ HpXv q ´ HpXvi | Xv´i q‚ dv dv ÿ

“ Exv

Qv pxiv , bq log

iPrds

PrpXvi “ xiv | Xv´i “ x´i v q i “ xi | X ´i “ x´i d bq PrpX v v j v v iPrds j‰i bPV ÿ “` ˘ ´i ‰ ´i i “ Exv Qv log Prpxv | ¨ q pxv q ÿ ÿÿ

“ ´Exv

Qv pxjv , bq log

iPrds

ÿ B , “ IpXti ; Xv´i q Bv t“v iPrds

which, together with (i), proves (iv). For the last property, we continue from (ii): B2 IpXti ; Xv´i q BtBv ˜ ¸ ÿÿ Prpxit | x´i B v dj bq j “ Ex´i Exi |x´i Qv pxv , bq log v t v Bt Prpxit | x´i v q j‰i bPV ˆ ˙ ÿ ÿ PrpXti “ c | x´i Prpxit | x´i v dj bq v dj bq j i “ Exi ,x´i Q px , bqQ px , cq log ´ log v v t t t v PrpXti “ c | x´i Prpxit | x´i v q v q j‰i b,cPV ˆ ˙ ÿ ÿ PrpXti “ c | x´i PrpXti “ c | x´i v dj bq v q j i ` Exi ,x´i Qv pxv , bqQt pc, xt q ´ . t v Prpxit | x´i Prpxit | x´i v dj bq v q j‰i b,cPV

(41)

PrpX i “a|X “x´i q

v t v For fixed x´i v , j, b, recall rpaq “ PrpX i “a|X “x´i d bq . Then, we have t

v

v

j

„ ȷ ÿ Prpxit | x´i PrpXti “ c | x´i rpxit q v dj bq v dj bq i ´i i ´i Exi ,x´i log ´ log “ E Prpx | x d bqrpx q log , j t v t xv t v rpcq PrpXti “ c | x´i Prpxit | x´i v q v q ´i xt PV

32

(42)

and

˙ PrpXti “ c | x´i PrpXti “ c | x´i v dj bq v q Exi ,x´i ´ t v Prpxit | x´i Prpxit | x´i v dj bq v q cPV ˆ ˙ i ´i ÿ PrpXt “ c | xv dj bq PrpXti “ c | x´i v q i i ´i Qt pc, xt q Prpxt | xv q ´ “ Ex´i v Prpxit | x´i Prpxit | x´i v dj bq v q xit ,cPV ˆ ˙ ÿ Prpxit | x´i Prpxit | x´i v dj bq v q i i ´i Q px , cq PrpX “ c | x q “ Ex´i ´ t t t v v PrpXti “ c | x´i PrpXti “ c | x´i v dj bq v q xit ,cPV ˆ ˙ ÿ PrpXti “ c | x´i Prpxit | x´i v q v q i i ´i Qt pxt , cq Prpxt | xv dj bq “ Ex´i ´ v PrpXti “ c | x´i Prpxit | x´i v dj bq v dj bq xit ,cPV ÿ ` ˘ i Qt pxit , cq Prpxit | x´i “ Ex´i v dj bq rpcq ´ rpxt q , v ˆ

ÿ

Qt pc, xit q

(43)

xit ,cPV

where we relabeled xit Ø c in the third line. Plugging Eqns. (42) and (43) into Eqn. (41) and setting a “ xit , we obtain „ ȷ ÿ ÿ B2 rpaq rpaq rpaq j i ´i IpXti ; Xv´i q “ Ex´i Q px , bqQ pa, cq PrpX “ a | x d bqrpcq log ` 1 ´ ě 0, v v t j t v v BtBv rpcq rpcq rpcq j‰i a,b,c (44)

as x log x ` 1 ´ x ě 0 for all x ě 0.

D.2

Proof of Lemma 1 Ð

i i r T ´t px di xi , x di bq “ Qt pb, xi qr As Q t t st px di b, x di xt q and QT ´t pxt , xt di bq “ Qt pb, xt qst pxt di b, xt q, we proceed by showing ” ı i i srt px di b, x di xit q “ Ex´i st pxt di b, xt q Xℓ´i “ x´i , X “ x t t . ℓ t

Here and below we omit stochastic process notation, such as Xt “ ¨, Xℓ´i “ ¨, etc., when the time and dimension indices are clear from the context. We have ȷ „ ” ı Prpxt di bq ´i ´i ´i i i i | xℓ , xt Ex´i st pxt di b, xt q Xℓ “ xℓ , Xt “ xt “ E t Prpxt q ÿ ˇ ´i i ˘ ` Prpxt di bq ˇ x , xt “ Pr x´i t ℓ Prpx q t ´i xt PV d´1 ˇ ´i ˘ ` ´i ˇ i ˘ ` ˇ xt Pr xt ˇ xt ÿ Prpxt di bq Pr x´i ℓ ˇ ˘ ` “ ¨ ˇ xit Prpxt q Pr x´i ´i ℓ d´1 xt PV

 ´i ´i i Prpx´i Prpx Prpxt di bq  t | xt q ℓ | xt q ¨ .   ´i i ´i  i i Prpxℓ | xt q Prpxt q Prpx  t | xt q d´1

ÿ “ x´i t PV

ř ´i ´i i i Observe that in the denominator we have Prpxit q Prpx´i | x´i d´1 Prpxt q Prpxℓ t q. ℓ | xt q “ Prpxℓ , xt q “ x´i t PV We continue: ř ´i ” ı Prpxt di bq Prpx´i ℓ | xt q x´i ´i ´i i i t Ex´i st pxt di b, xt q Xℓ “ xℓ , Xt “ xt “ ř ´i t Prpxt q Prpx´i ℓ | xt q x´i t ř ´i ´i Prpxt di bq Prpx´i ℓ | xt q{ Prpxℓ q x´i t “ ř ´i ´i ´i x´i Prpxt q Prpxℓ | xt q{ Prpxℓ q t

33

Prpxt di b | x´i ℓ q x´i t ř ´i Prpxt | xℓ q x´i t

ř “ “

PrpXti “ b | x´i ℓ q . i i PrpXt “ xt | x´i ℓ q

´i i i Recalling that by Definition 2, srt px di b, x di xit q “ PrpXti “ b | x´i ℓ q { PrpXt “ xt | xℓ q finishes the proof.

D.3

Proof of Lemma 5

First, assume that q ă p. In this case, as x ´ 1 ´ log x ď x log x ´ x ` 1 we have

p

´q q

´ 1 ´ log pq

p p p q log q ´ q ` 1

for x ě 1,

1 1 1 “ ď . q minpp, qq α

¯ď

The case p ă q follows similarly as xpx ´ 1 ´ log xq ď x log x ´ x ` 1

for x ď 1.

x Finally, as limxÑ1 xx´1´log log x´x`1 “ 1, the case p “ q follows by continuity.

D.4

Proof of Lemma 7

Recall the definition of νpt, aq from Eqn. (6). We have PrpXti “ aq “ PrpX0i “ aqe´t ` νpt, aqp1 ´ e´t q, and PrpXti “ a | X0i “ cq “ Ita “ cue´t ` νpt, aqp1 ´ e´t q. Taking the difference and summing over all a P V, we get ˜ ¸ ÿ ÿ ÿ ` ˘ i i i i Qpa, bq PrpXt “ a | X0 “ cq ´ PrpXt “ aq “ Qpa, bqIta “ cu ´ Qpa, bq PrpX0 “ aq e´t aPV

aPV

aPV

˜

¸

Qpc, bq ´

“

ÿ

Qpa, bq PrpX0i “ aq

e´t .

aPV

For the uniform process, the term in the bracket equals zero as all Qpa, bq “ 1{S . For the remasking process, as c P rSs we consider three cases: (i) b P rSs, (ii) b “ REMASK, and (iii) b “ MASK. In the first case, we have Qpc, bq “ 0 and the only a P V ř such that Qpa, bq ‰ 0 is a “ REMASK, for which PrpX0i “ aq “ 0. In the second case, Qpc, bq “ 1 ´ pM “ aPV Qpa, bq PrpX0i “ aq. The third case follows equivalently. This proves Eqn. (21). Next, we have ÿ Qpa, bq PrpXti “ a | Xu´i “ x´i q aPV

˜ ÿ “

PrpX0i “ xi0 | Xu´i “ x´i q

xi0 PV

¸ ÿ

Qpa, bq PrpXti “ a | X0i “ xi0 q

aPV

ÿ

“ Fpt, bq

PrpX0i “ xi0 | Xu´i “ x´i q

xi0 PV

“ Fpt, bq, x | ¨q in Eqns. (21) and (22), which proves Eqn. (22). Finally, using Eqn. (12), we can replace Prp¨ | ¨q with Prp¨ which concludes the proof. 34

E

Connection to the effective total correlation

In Dmitriev et al. (2026), for the masking noising process, the effective total correlation is studied: ż8 ÿ Dpqdata q :“ minp1, tq IpXti ; Xtj | Xt´i,j qdt. 0

i‰j

The next proposition connects this quantity with the presented results. Proposition 8. For the masking noising process, ż1 DTCpXt qdt “ Dpqdata q. 0

Proof. We show that both expressions are equal to ż8 d ´ minp1, tq DTCpXt qdt. dt 0 Indeed,

ż8 minp1, tq

´ 0

ż1 ż8 d d d DTCpXt qdt “ ´ t DTCpXt qdt ´ DTCpXt qdt dt 0 dt 1 dt ż1 t“0 t“1 ` “ tDTCpXt q DTCpXt qdt ´ DTCpXt q 0

t“1

t“8

ż1 DTCpXt qdt.

“ 0

Furthermore, observe that as IpXti ; Xv´i q “ IpXti ; Xv

´pi,jq

q ` IpXti ; Xvj | Xv´i,j q for any j ‰ i, we have

ÿ B ÿ B B IpXti ; Xv´i q “ IpXti ; Xv´i q “ IpXti ; Xvj | Xv´i,j q, Bv Bvj Bvj j‰i j‰i where BvB j denotes taking the derivative only with respect to the j-th coordinate of Xv´i . Next, for the masking noising process, B IpXti ; Xvj | Xv´i,j q “ IpXti ; MASK | Xv´i,j q ´ IpXti ; Xvj | Xv´i,j q “ ´IpXti ; Xvj | Xv´i,j q. Bvj Using that

B i ´i i Bv IpXt ; Xv q

d “ dv DTCpXv q, as shown in Proposition 5 (iv), we obtain

ř

t“v

ż8 minp1, tq 0

ÿ

IpXti ; Xtj | Xt´i,j qdt “ ´

ż8 minp1, tq 0

i‰j

d DTCpXt qdt, dt

which finishes the proof. For the masking process, as Fpt, bq “ 0 for b P rSs, one can show a stronger version of Proposition 6: B2 B IpXti ; Xv´i q “ ´ IpXti ; Xv´i q, BtBv Bv

(45)

which implies that the discretization error can be written as Nÿ ´1 k“0

pkq

Ldiscr À ´κ

ż8 minp1, tq 0

d DTCpXt qdt “ κDpqdata q. dt

(46)

Therefore, our results recover exactly the best previously known bound for the masking noising process and obtain the first adaptive bound for the uniform and remasking noising processes. 35

References Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981–17993. Cai, C. and Li, G. (2026). Confidence-based decoding is provably efficient for diffusion language models. arXiv preprint arXiv:2603.22248. Campbell, A., Benton, J., De Bortoli, V., Rainforth, T., Deligiannidis, G., and Doucet, A. (2022). A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279. Chen, H. and Ying, L. (2025). Convergence analysis of discrete diffusion model: exact implementation through uniformization. Journal of Machine Learning, 4(2):108–127. Chen, S., Cong, K., and Li, J. (2025). Optimal inference schedules for masked diffusion models. arXiv preprint arXiv:2511.04647. Conforti, G., Durmus, A., and Pham, L.-T.-N. (2025). Non-asymptotic convergence of discrete diffusion models: Masked and random walk dynamics. arXiv preprint arXiv:2512.00580. DiffusionGemma Team, Taïga, A. A., Assiene, J., Calandriello, D., Chaabouni, R., Gante, J., von Glehn, T., Keating, N., Knutsen, C., Kukla, M., et al. (2026). Diffusiongemma technical report. arXiv preprint arXiv:2608.00146. Dmitriev, D., Huang, Z., and Wei, Y. (2026). Efficient sampling with discrete diffusion models: Sharp and adaptive guarantees. In Hanneke, S. and Lattimore, T., editors, Proceedings of Thirty Ninth Conference on Learning Theory, volume 336 of Proceedings of Machine Learning Research, pages 2038–2104. PMLR. Feinberg, E. A., Mandava, M., and Shiryaev, A. N. (2014). On solutions of kolmogorov’s equations for nonhomogeneous jump markov processes. Journal of Mathematical Analysis and Applications, 411(1):261– 270. Feller, W. (1940). On the integro-differential equations of purely discontinuous markoff processes. Transactions of the American Mathematical Society, 48(3):488–515. Google DeepMind (2026). Diffusiongemma. Google DeepMind. An experimental open text diffusion model based on Gemma 4. Gourevitch, S., Janati, Y., Shariatian, D., Simsekli, U., Moulines, E., Xing, E. P., and Durmus, A. (2026). Uniform diffusion models revisited: Leave-one-out denoiser and absorbing state reformulation. arXiv preprint arXiv:2605.22765. Huang, Z., Wang, Y., Chen, Z., and Qi, G.-J. (2026). Don’t settle too early: Self-reflective remasking for diffusion language models. In International Conference on Learning Representations, volume 2026, pages 56998–57026. Huang, Z., Wei, Y., and Chen, Y. (2024). Denoising diffusion probabilistic models are optimally adaptive to unknown low dimensionality. arXiv preprint arXiv:2410.18784. Kim, J., Kim, S., Lee, T., Pan, D. Z., Kim, H., Kakade, S., and Chen, S. (2025). Fine-tuning masked diffusion for provable self-correction. arXiv preprint arXiv:2510.01384. Li, G. and Cai, C. (2025). Breaking AR’s sampling bottleneck: Provable acceleration via diffusion language models. Advances in Neural Information Processing Systems, 38. Li, G., Cai, C., and Wei, Y. (2025). Dimension-free convergence of diffusion models for approximate gaussian mixtures. arXiv preprint arXiv:2504.05300. Li, G. and Yan, Y. (2024). Adapting to unknown low-dimensional structures in score-based diffusion models. Advances in Neural Information Processing Systems, 37:126297–126331. 36

Liang, J., Huang, Z., and Chen, Y. (2025a). Low-dimensional adaptation of diffusion models: Convergence in total variation. arXiv preprint arXiv:2501.12982. Liang, Y., Huang, R., Lai, L., Shroff, N., and Liang, Y. (2025b). Absorb and converge: Provable convergence guarantee for absorbing discrete diffusion models. Advances in Neural Information Processing Systems, 39. Liang, Y., Liang, Y., Lai, L., and Shroff, N. (2025c). Discrete diffusion models: Novel analysis and new sampler guarantees. Advances in Neural Information Processing Systems, 39. Liang, Y., Tan, Z., Shroff, N., and Liang, Y. (2026). Sharp convergence rates for masked diffusion models. arXiv preprint arXiv:2602.22505. Liu, Y., Zhao, Y., Xie, Z., Ye, Q., Jiao, J., Hu, Y., Cao, S., and Liu, Y. (2026). Balancing understanding and generation in discrete diffusion models. In Proceedings of the 43rd International Conference on Machine Learning. Spotlight. Lou, A., Meng, C., and Ermon, S. (2024). Discrete diffusion modeling by estimating the ratios of the data distribution. In International Conference on Machine Learning, pages 4735–4763. PMLR. Noguerales, R. C., Schölkopf, B., Hofmann, T., and Raoufi, A. (2026). What does a discrete diffusion model learn? arXiv preprint arXiv:2607.05381. Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., and Li, C. (2025). Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In The Thirteenth International Conference on Learning Representations. Pham, L.-T.-N., Shariatian, D., Ocello, A., Conforti, G., and Durmus, A. O. (2025). Discrete markov probabilistic models: An improved discrete score-based framework with sharp convergence bounds under minimal assumptions. In International Conference on Machine Learning. Ren, Y., Chen, H., Zhu, Y., Guo, W., Chen, Y., Rotskoff, G., Tao, M., and Ying, L. (2026). Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. Advances in Neural Information Processing Systems, 38:167228–167282. Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J., Rush, A., and Kuleshov, V. (2024). Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:130136–130184. Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. (2024). Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131–103167. Von Rütte, D., Fluri, J., Ding, Y., Orvieto, A., Schölkopf, B., and Hofmann, T. (2025). Generalized interpolating discrete diffusion. arXiv preprint arXiv:2503.04482. Wainwright, M. J. (2026). The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity. arXiv preprint arXiv:2608.13520. Wang, G., Schiff, Y., Sahoo, S., and Kuleshov, V. (2026). Remasking discrete diffusion models with inferencetime scaling. Advances in Neural Information Processing Systems, 38:147282–147339. Xu, M., Geffner, T., Kreis, K., Nie, W., Xu, Y., Leskovec, J., Ermon, S., and Vahdat, A. (2025). Energy-based diffusion language models for text generation. In International Conference on Learning Representations, volume 2025, pages 33769–33789. Ye, M., Kudo, K., Ikeda, W., Matsuda, R., Sakaguchi, K., and Suzuki, J. (2026). Sumi: Open uniform diffusion language model from scratch. Zhao, Y. and Cai, C. (2026). Adaptation to intrinsic dependence in diffusion language models. arXiv preprint arXiv:2602.20126. Zhao, Y., Shi, J., Chen, F., Druckmann, S., Mackey, L., and Linderman, S. (2026). Informed correctors for discrete diffusion models. Advances in Neural Information Processing Systems, 38:125510–125538. 37

Related documents

• Monte Carlo Approximations of Time-Nonlocal Diffusions in Bounded Domains#1033902
• Accurate Sampling from Diffusion Models#1077973
• Backward Kolmogorov Transport: Sampling Invariant Laws of Reversible Diffusions from Trajectory Data#1116995
• OnlineCache: Learning Dynamic Caching Policies with Error Correction for Efficient Diffusion Inference#422263
• Optimize Your Sampling: Tuned Diffusion Sampling with Bayesian Optimization#463206