Parallelism, critical windows, and separations among diffusion language models Liye Wang †
Sitan Chen *
arXiv:2609.20539v1 [cs.LG] 17 Sep 2026
September 13, 2026
Abstract A popular selling point of diffusion large language models (dLLMs) is their capacity for parallelism: the ability to generate sequences of text far more efficiently than autoregressive models, which require one forward pass per token. Yet among the many competing paradigms for dLLMs, from masked to uniform to Gaussian diffusion, principled understanding of how these different proposals compare in parallelism remains limited. In this work, we initiate a fine-grained comparison of the capacity for parallelism among these three leading approaches and prove the following: • Uniform and Gaussian diffusion can sample in a number of forward passes which scales with the dual total correlation of the underlying distribution, a measure of intrinsic complexity which can be much smaller than the context length. Previously, it was only known how to achieve this using masked diffusion [CCL26, LZ25]. √ e d) forward passes are necessary • For a certain family of random empirical measures, we show that Θ( and sufficient to sample using uniform or Gaussian diffusion, yet there exist approximate score oracles e for which Ω(d) forward passes are needed for masked diffusion. This establishes the first provable separation in parallelism between the three prevailing dLLM paradigms. Contrary to popular intuition that masked diffusions are harder to parallelize because they must commit to token values, the latter separation instead comes from the fact that the critical windows in masked diffusion sampling are asymptotically narrower than those in uniform and Gaussian diffusion sampling.
* SEAS, Harvard University. Email: [email protected]. This work was supported by the Harvard Dean’s Competitive Fund
for Promising Scholarship and was completed in part during a visit to the Simons Institute for the Theory of Computing. † Tsinghua University. Email: [email protected].
1
Contents 1
Introduction 1.1 Result 1: Scaling with intrinsic complexity . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2 Result 2: Separating masked diffusion from other paradigms . . . . . . . . . . . . . . . . . 1.3 Related work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
4 5 6 9
2
Preliminaries 2.1 Uniform diffusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.2 Gaussian diffusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.3 Masked diffusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.4 Information-theoretic quantities . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
10 10 11 12 14
3
Uniform diffusion can scale with dual total correlation 3.1 Conditional TC versus DTC decrement . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3.2 Score error and approximate reverse kernel . . . . . . . . . . . . . . . . . . . . . . . . . . 3.3 Sampler and telescoping argument . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
15 16 18 21
4
Gaussian diffusion can scale with dual total correlation 4.1 Conditional TC versus DTC decrement . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.2 Reverse DPI for Gaussian diffusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.3 Score error and approximate reverse kernel . . . . . . . . . . . . . . . . . . . . . . . . . . 4.4 Sampler and telescoping argument . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
23 24 25 26 28
5
Construction of random empirical measure √ e d) upper bound for uniform diffusion O( 6.1 Quantifying the critical window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6.2 Identifying the critical window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6.3 Sampling after locating the window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . √ e d) upper bound for Gaussian diffusion O( 7.1 Quantifying the critical window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7.2 Identifying the critical window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7.3 Sampling after locating the window . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
31
6
7
32 33 35 39 43 44 47 53
8
e Ω(d) lower bound for masked diffusion 58 8.1 Adversarial oracle construction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 59 8.2 Distribution indistinguishability: the lower bound proof . . . . . . . . . . . . . . . . . . . . 60
9
62
Outlook √ e d) lower bound for uniform diffusion A Ω( A.1 Adversarial oracle construction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.2 Distribution indistinguishability: the lower bound proof . . . . . . . . . . . . . . . . . . . .
2
68 70 73
B Deferred proofs B.1 Proofs from Section 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.1.1 Proof of Lemma 2.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.1.2 Proof of Lemma 2.8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Proofs from Section 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2.1 Proof of Proposition 3.6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2.2 Proof of Lemma 3.10 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2.3 Proof of Lemma 3.11 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3 Proofs from Section 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3.1 Proof of Lemma 4.6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3.2 Proof of Lemma 4.7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4 Proofs from Section 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.1 Proof of Lemma 6.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.2 Proof of Lemma 6.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.3 Proof of Lemma 6.6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.4 Proof of Lemma 6.7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.5 Proof of Lemma 6.8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.6 Proof of Proposition 6.10 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4.7 Proof of Lemma 6.11 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5 Proofs from Section 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.1 Proof of Lemma 7.5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.2 Proof of Lemma 7.6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.3 Proof of Lemma 7.8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.4 Proof of Lemma 7.9 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.5 Proof of Proposition 7.11 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.6 Proof of Lemma 7.12 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.7 Proof of Lemma 7.13 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
75 75 75 76 76 76 76 77 78 78 78 79 79 80 80 80 81 81 82 82 83 83 84 84 86 86 88
C Details of simulation for Figure 4
88
3
1
Introduction
Diffusion large language models (dLLMs) have recently emerged as a powerful alternative to autoregressive large language models (LLMs) for generative modeling over discrete domains [AJH+ 21, LME23, SAS+ 24, NZY+ 25]. Despite the moniker, diffusion language modeling is not one, but a multitude of competing frameworks (Figure 1) all built around the same guiding principle of denoising: given a corruption process that converts data into noise, learn how to undo it in order to transform fresh noise into fresh samples. LLMs are themselves a special case of this paradigm, where the corruption process is right-to-left erasure, and a central question in the theory and practice of dLLMs is whether there is a better choice of corruption process. One of the most widely touted selling points of moving beyond right-to-left erasure is that it unlocks the ability to perform few-step generation: Whereas LLMs must decode one token at a time, dLLMs based on alternative corruption processes can in principle generate a sequence of length d in o(d) forward passes. The cost of doing so, however, remains poorly understood. It is well known that there is some statistical price to pay, as such decoding strategies fundamentally incur discretization error coming from approximating certain posterior distributions using only their marginals. But the extent to which this price can be rendered negligible depends heavily on how the sampling algorithm at inference time is tuned. Various empirical works have suggested that this price also depends heavily on the choice of corruption process itself, arguing that few-step generation is easier to achieve for some frameworks than for others [SDG+ 25, LYA+ 26, HQL+ 26]. For example, it is commonly claimed that corruption processes involving independent erasures, which give rise to the popular paradigm of masked diffusion models, are less conducive to few-step generation. The informal reasoning is that by default, masked diffusion models have to commit to token values in discrete steps of the generation process, meaning that mistakes arising from committing multiple tokens in parallel cannot subsequently be corrected [SSP+ 25, SDG+ 25, SLY+ 26]. Other popular choices of corruption process, e.g., involving random re-assignments of token values (uniform diffusion models [LME23, SDG+ 25, Dif26]) or Brownian motion in a latent space (Gaussian diffusion models [LTG+ 22, DSR+ 22, RDE+ 26, LYA+ 26, HQL+ 26, Die26]), are claimed not to suffer from this as they can progressively revise their output over the course of sampling (see Figure 1 for a depiction of these different dLLM paradigms). Yet these intuitions are largely heuristic, and the foundations of few-step generation under any of these paradigms remain nascent. A flurry of recent theoretical works [LC26, CCL26, DHW26a, Wai26, LZ25, ZC26] have suggested that with the right inference schedule, masked diffusion models can achieve fewstep generation, where the number of steps scales with an intrinsic complexity measure associated to the data distribution. In contrast, no such results were known for other frameworks like uniform and Gaussian diffusion. In this work, we thus ask: How does the choice of corruption process affect the few-step generation capabilities of dLLMs? Framework: score oracle queries. When comparing different dLLM frameworks, one basic figure of merit is the number of forward passes of the model needed to generate an accurate sample. More formally, a single forward pass of the model provides us with the following information. Given a clean sample X0 , the corruption process degrades it into a noisy sample Xt , and a single forward pass provides an approximation to the coordinatewise posterior marginals law((X0 )i | Xt = z) simultaneously for all coordinates i, for any noise level t and conditioning z of one’s choice. As these marginals are equivalent, up to affine transformation, to the so-called annealed score functions of the data distribution, in this work we refer to a single forward pass of the model as a query to the (approximate) score oracle. We define these oracles more formally in Sections 2.1–2.3. When proving positive results about few-step generation, we will analyze
4
(a) C ORRUPTION PROCESS
(b) O NE INFERENCE STEP Xt
resample each token w.p. 1 − e−t score oracle output: posterior marginals of X0
UNIFORM
X0
Xt
XT i.i.d. uniform tokens
Xs , s < t every token resampled Xt
one-hot
e−t X0 + σt G score oracle output: posterior mean E[X0 | Xt ]
G AUSSIAN A ∈ [S]d
X0 = (eAi )i ∈ RSd
XT ≈ N (0, I)
Xt
Xs , s < t every coordinate updated ⋆
⋆⋆
Xt
mask each token w.p. 1 − e−t ⋆
MASKED
X0
⋆⋆
Xt
score oracle output: marginals of masked tokens
⋆⋆⋆⋆⋆⋆⋆
XT all tokens masked
⋆
Xs , s < t unmasked tokens committed
noise level t
Figure 1: The three dLLM paradigms considered in this work. specific algorithms that leverage these oracles; when proving no-go results, we will rule out all algorithms that leverage these oracles.
1.1
Result 1: Scaling with intrinsic complexity
First, we prove that both uniform and Gaussian diffusion can match the scaling established in prior work for masked diffusion. More specifically, we prove the following. Given a categorical distribution q over sequences X = (X1 , . . . , Xd ) in a product space Σd , its dual total correlation DTC(q) is, roughly speaking, the total entropy in the distribution which is not explained by the local entropies H(Xi | X−i ). Prior work showed that, up to logarithmic factors, masked diffusion models can sample from q in a number of score oracle queries scaling linearly in DTC(q), which can in general be much smaller than the dimension d. We begin by showing that such a result is also possible for uniform diffusion. Theorem 1.1 (Informal, see Theorem 3.1). Let q be an arbitrary distribution over Σd . There is an ale e hides logarithmic factors in |Σ|, d, and ε, to an gorithm that makes O(DTC(q)/ε) queries, where O(·) accurate estimate of the uniform diffusion score for q and outputs a sample from a distribution qb for which KL(q ∥ qb) ⩽ ε. By a similar mechanism albeit with a more involved argument, we also show: Theorem 1.2 (Informal, see Theorem 4.1). There is a sampler with the same guarantee as in Theorem 1.1 e which uses O(DTC(q)/ε) queries to an accurate estimate of the Gaussian diffusion score for q. The proof strategy for these results is a marked departure from the one for masked diffusion in [CCL26]. Crucial to that work was an exact characterization for the error incurred by the standard algorithm for masked
5
diffusion sampling; with this characterization in hand, it was enough to substitute an appropriate choice of e step size schedule into that expression to obtain the desired rate of O(DTC(q)/ε). In contrast, for Theorems 1.1 and 1.2, in lieu of an exact characterization, we prove that the error incurred in a single step of the sampler from ignoring conditional correlations can be controlled in terms of the increase in dual total correlation when going from one noise level to a lower noise level. This is proven using a certain reverse data processing inequality (DPI) which shows that the distance between two distributions cannot contract too quickly if they are only slightly noised. With this estimate, we can then telescope the DTC increases together and conclude a query complexity that scales linearly in DTC(q). This argument is illustrated in Figure 2. Xtj+1 Xtj
DTC(q)
DTC(qt )
discarded dependence = TC(Xtj | Xtj+1 )
KL(q ∥ qb)
tj
tj+1
tj+1 ≈ (1 + a) tj
⩽
P
P
chain rule
j TC(Xtj | Xtj+1 ) ≲ reverse DPI
a
noise level t
j DTC(qtj ) − DTC(qtj+1 )
⩽
telescoping
a · DTC(q)
=
a ≍ ε/DTC(q)
ε
Figure 2: Illustration of the telescoping argument behind Theorems 1.1 and 1.2. The sampler’s KL error is a sum over conditional total correlations corresponding to dependencies discarded at each sampling step (inset). The key step is a reverse data-processing inequality which bounds each summand, roughly, by step size times the drop in DTC(qt ) in that step.
1.2
Result 2: Separating masked diffusion from other paradigms
Theorems 1.1 and 1.2 establish that uniform and Gaussian diffusion can adapt to a certain measure of the intrinsic complexity of the data distribution at least as well as masked diffusion. But they leave open whether or not there is a genuine separation among these paradigms. In our next set of results, we exhibit such a separation by constructing a simple family of random empirical measures for which the query complexity of sampling with a masked diffusion score oracle is strictly higher than with a uniform or Gaussian diffusion oracle. More precisely, distributions in this family are given by sampling eκd many points uniformly at random from the Boolean hypercube, where 0 < κ < log 2 is an (unknown) absolute constant (see Section 5 for a formal definition). Recently, √ [XP26] showed e d) queries are that given an approximate Gaussian diffusion score oracle for such a distribution, at least Ω( necessary even to determine κ to sufficient precision, let alone generate samples close to the distribution. In Appendix A, we prove that this lower bound also applies to uniform diffusion, using similar ideas. √Our main result in this second part of the paper is to show that for these random empirical measures, e O( d) query complexity is actually tight for both uniform and Gaussian diffusion: Θ(d) random Theorem 1.3 (Informal, see Theorems 6.1 and 7.1). For empirical measures √ 2 q supported on 2 e d/ε ) queries to any approximate points on the Boolean hypercube, there is an algorithm that uses O( uniform diffusion score oracle for q that has score error 1/poly(d) and, with high probability over the randomness of the support of q, succeeds in sampling from q to total variation error ε. The same result also holds given any approximate Gaussian diffusion score oracle that has score error 1/poly(d).
6
Notably, for this family of distributions, the dual total correlation is Θ(d), so the scaling achieved in Theorem 1.3 provably goes beyond the dual total correlation scaling implied by Theorems 1.1 and 1.2. In contrast, e we show that for this same problem instance, Ω(d) queries are needed in the case of masked diffusion. Theorem 1.4 (Informal, see Theorem 8.1). Let q be an empirical measure supported on 2Θ(d) random points on the Boolean hypercube. There is an approximate masked diffusion oracle such that any algorithm for sampling making oe(d) queries to this oracle produces samples from a distribution qb such that with high probability over the randomness of q, TV(q, qb) ≥ 0.99. Altogether, these results establish the first provable separation in few-step generation ability among the three prevailing paradigms for diffusion language modeling. (a) U NIFORM AND G AUSSIAN DIFFUSION e −1/2 ) critical window√Θ(d e hidden among Θ( d ) positions Xt
recovery error
√ e d) O(
Xt
posterior collapsed
indistinguishable from noise
queries
noise level
(b) M ASKED DIFFUSION e −1 ) critical window Θ(d hidden among Θ(d) positions Xt
⋆
Xt ⋆ ⋆
⋆
⋆⋆ ⋆⋆⋆
e Ω(d) queries
# masked tokens
Figure 3: Critical window for masked diffusion is narrower than for uniform and Gaussian diffusion, so more queries are needed to locate it before one can generate a sample. We highlight that this separation manifests through a mechanism fundamentally different from the “commitment” issue mentioned previously. Instead, it comes from a difference in the sharpness of a certain phase transition arising in the sampling dynamics √ for masked diffusion versus uniform and Gaussian diffusion. e e To see this, we note that in both the Ω( d) lower bound for uniform and Gaussian diffusion and the Ω(d) lower bound for masked diffusion, the key idea is that for random empirical measures, unless one queries the score oracle near the right noise level, the approximation error in the score oracle can be designed to suppress any useful signal about the underlying distribution q. Crucially, however, the range of acceptable noise levels for which this is the case is far narrower for masked diffusion than for uniform and Gaussian diffusion: whereas the informative range for masked diffusion occupies an O(log(d)/d) fraction of the full noise spectrum (Proposition 8.6), the informative range for uniform and Gaussian diffusion occupies a far p e wider O( log(d)/d) fraction (Propositions 6.4 and 7.4). As a result Ω(d) queries are needed in the former √ e d) suffice in the latter case. Figure 4 illustrates this discrepancy numerically. case, whereas only O( These critical windows have previously been studied in the context of speciation transitions in diffusion sampling [BBDBM24, LC24, LKC25, SFW25, RA23], but to our knowledge this is the first time critical windows have been used to derive new upper bounds on the query complexity of diffusion sam√ e pling. Roughly speaking, after identifying the location of the window using O( d) queries to a uniform 7
1
(a) masked diffusion
d = 64 d = 256 d = 1024 d = 4096
=
0 1
(b) uniform diffusion
0.5 0 1
(c) Gaussian diffusion
d 1/2
10 2
10 3
0.5 0
10 1
width of the critical window (nats)
recovery probability Pr[X0 = y
Xt ]
0.5
masked diffusion uniform diffusion Gaussian diffusion
(d) 0.10
0.15
0.20
0.25
information per coordinate of the noisy sample (nats)
0.30
64
128
256
d 1
512 1024 2048 4096
dimension d
Figure 4: Simulated critical windows for the random empirical measure with M = eκd codewords, κ = 0.2, extending a simulation of [XP26] from Gaussian diffusion to masked and uniform diffusion. (a)– (c) Probability that the posterior at a given noise level recovers the planted codeword, plotted against the information per coordinate of the noisy sample. The information per coordinate I is the mutual information between a uniformly random bit and its corrupted version: (m/d) log 2 for masked diffusion with m revealed coordinates, and Iunif (t), Igauss (t) for uniform and Gaussian diffusion. By Eq. (22), every process has its critical level at I = κ, placing the three transitions on a √ common scale. (d) Width of the transition versus d: it decays like 1/d for masked diffusion but only like 1/ d for uniform and Gaussian diffusion. or Gaussian diffusion score oracle, we then argue that at noise levels above this critical window, the score is well-approximated by that of the uniform distribution, and at noise levels below the critical window, the posterior distribution over the clean sample is so pronounced that one can easily “round” to the correct point in the support. Thus, it suffices to design an algorithm that can efficiently locate and “cross” the window, the details of which are given in Sections 6.2 and 7.2. The idea behind this argument is illustrated in Figure 3. Concurrent work. Dmitriev, Huang, and Wei [DHW26b] also independently established a DTC-adaptive bound for uniform diffusion. Their proof strategy bears similarities to the proof of our Theorem 1.1, in particular bounding the sampling error by a telescoping sum of differences between dual total correlation at successive noise levels. Interestingly, their work gives a characterization of the discretization error for any forward process in terms of the change in the mutual information between one coordinate and the remaining coordinates at other times. The remaining results in our works are disjoint: [DHW26b] additionally prove a DTC-adaptive bound for remasking diffusion [WSSK26], whereas we prove one for Gaussian diffusion and study separations among the various frameworks.
8
1.3
Related work
Theory for discrete diffusion sampling. Campbell et al. [CBDB+ 22] formulated continuous-time discrete denoising through continuous-time Markov chains. Chen and Ying [CY24] used uniformization to implement the reverse chain associated with an estimated score and obtained total-variation and KL guarantees on the hypercube, notably obtaining zero discretization error in O(d) uniform diffusion score oracle queries. Li and Cai [LC26] initiated the theoretical study of the query complexity of masked diffusion as a function of information-theoretic properties of the data distribution. Chen, Cong, and Li [CCL26] and Lavenant and Zanella [LZ25] derived an exact characterization for the sampling error in terms of the data distribution’s information curve, which can be used to obtain query complexity bounds scaling linearly in the total correlation and dual total correlation. Closely related to our lower bound for masked diffusion, [CCL26] also gave an Ω(d) lower bound with a different construction based on MDS codes. Dmitriev et al. [DHW26a] and Zhao and Cai [ZC26] proved sharp convergence guarantees for standard uniform and masking τ -leaping schemes, showing that the latter can adapt to the total and dual total correlation of the distribution. Separately, Ren et al. [RCRY25] developed a stochastic-integral framework for KL analysis of timediscretized reverse processes. Building on this framework, Ren et al. [RCZ+ 25] developed and analyzed higher-order solvers for discrete diffusion. More broadly, there is a large body of orthogonal work that analyzes the convergence of CTMC-based discrete diffusion samplers, with particular attention to the dependence on the alphabet size [ZCG25, PSO+ 25, LHL+ 25, LLLS25, PCRD26, KLZ+ 26]. Additionally, there is an extensive theoretical literature on Gaussian diffusion models — see the recent notes of [Lu26] for a comprehensive overview. Most relevant to our work is the lower bound of Xun and √ e Price [XP26] showing that Ω( d) queries are necessary for Gaussian diffusion. The second half of this paper is based on one of the constructions studied in that work. In addition, there have been a number of works on algorithmic stochastic localization for sampling from distributions over the Boolean hypercube [EAMS22, MW23, EAMS23], essentially by regarding the domain as a subset of Euclidean space, performing Gaussian diffusion, and rounding at the end. Critical windows and speciation. Numerous works [GVS+ 23, RA23, SFLW24, SFW25, BBDBM24, LC24, LKC25, ABBM26] have studied an intriguing property of real-world diffusion models whereby key aspects of the final generation emerge over a narrow range of times in the reverse process. These works characterize this behavior as a phase transition, interchangeably termed a critical window or a speciation transition, and provide various probabilistic models of data under which this arises. This phenomenon is the same one that manifests in the context of the random empirical measures considered in [XP26] and that we study in the second part of this work, which is in turn closely related to concepts from the statistics and information theory literature like the all-or-nothing phenomenon in high-dimensional inference [NWZ20] and channel resolvability and soft covering [HV93].
Organization The DTC-adaptive bounds for uniform and Gaussian diffusion are proved in Sections 3 and 4. The random empirical measure is introduced in Section 5. The uniform and Gaussian upper bounds for this family are proved in Sections 6 and 7, and the masked lower bound is proved in Section 8. The uniform diffusion query lower bound appears in Appendix A. The deferred proofs appear in Appendix B.
9
2
Preliminaries
2.1
Uniform diffusion
Given alphabet size S ≥ 2, define Σ = [S] and let X = Σd . Let q be a distribution over X . Let ∆(Σ) denote the probability simplex over Σ. For every t ≥ 0 define 1 − αt αt = e−t , βt = , S and, for z ∈ Σ, define the one-coordinate transition kernel KtU (· | z) = αt 1[z = ·] + βt . These kernels (KtU ) form a semigroup. Given X0 ∼ q, let (Xt )t≥0 denote the Markov process with transition kernel (KtU )⊗d , and let qt := law(Xt ). Given y ∈ X , i ∈ [d], and a ∈ Σ, let y i←a be the point obtained by replacing the i-th coordinate of y [d]×Σ with a. The (uniform diffusion) score at time t is the function st : X → R≥0 given by the likelihood ratio qt (y i←a ) . qt (y)
st (y)[i, a] =
The score naturally admits an interpretation both in terms of posterior marginals and in terms of Gibbs sampling marginals: Lemma 2.1. Given a vector v ∈ RS>0 and given h > 0 and b ∈ Σ, define probability distributions gibbs ppost over Σ by h,b , ph X αh 1[a = b] + βh ′ := [v](a) ppost v − β v . a h a h,b αh ′ a ∈Σ
pgibbs [v](a) := h
va
1 αh
P
a′ ∈Σ va′
− βh .
Let h = t − s. Given y ∈ X , P[Xsi = a | Xt = y] = ppost h,yi [st (y)[i, ·]](a) , and P[Xsi = a | (Xt )−i = y−i ] = pgibbs [st (y)[i, ·]](a) . h Here, vectors with the subscript −i denote the vectors obtained by deleting coordinate i. We defer the proof to Appendix B. We will consider estimated scores (b st ) and quantify their error as follows. First, without loss of generality, we will assume that the estimated scores are normalized, i.e., for all t, y ∈ X , and i ∈ [d], sbt (y)[i, yi ] = 1. Given c, s > 0, define the (entropic) Bregman divergence by ψ(c, s) := s − c + c log(c/s) . Definition 2.2 (Score error – uniform diffusion). Given a normalized score estimate sbt at time t > 0, its
10
score error is εunif (t) := EY ∼qt
d X X
ψ st (Y )[i, a], sbt (Y )[i, a] .
i=1 a∈Σ
We will assume that εunif (t) < ∞ for all t > 0, so that sbt (y)[i, a] > 0 for all i, a, t, y. This is equivalent to the standard score-entropy loss used in training uniform diffusion models [LME23]. In Section 3.2, we discuss a robust analogue of Lemma 2.1 for converting an approximate score into a valid estimate for the posterior marginals. For lower-bound arguments, it will be convenient to make the score estimate and target distribution explicit, using the oracle language. An approximate uniform diffusion oracle O, queried at (t, y), returns a score estimate sbO t (y). We write εunif (O; q, t) := EY ∼qt
d X X
ψ st (Y )[i, a], sbO t (Y )[i, a]
i=1 a∈Σ
for the corresponding score error. When q and O are clear from context, we abbreviate εunif (O; q, t) by εunif (t) and write sbt for sbO t . Finally, for the case of binary alphabet Σ = {−1, +1}, the forward likelihood ratio has a convenient overlap form which will be useful in Section 6. Lemma 2.3 (Binary likelihood ratio). For every x, z ∈ {−1, +1}d , 2d (KtU )⊗d (x | z) = exp
d X
! log(1 + e−t xi zi )
= (1 + e−t )d tanh(t/2)dH (x,z) ,
i=1
where dH (x, z) denotes the Hamming distance between x and z. Proof. Coordinatewise, 2KtU (xi | zi ) = 1 + e−t xi zi . Multiplying these identities gives the first equality. A matching coordinate contributes 1 + e−t and a mismatching coordinate contributes 1 − e−t = (1 + e−t ) tanh(t/2), which gives the second.
2.2
Gaussian diffusion
We begin by describing Gaussian diffusion for general distributions over Rd , before specializing to embeddings of discrete distributions. Let q be a distribution on Rd . Define the Ornstein–Uhlenbeck process (Xt )t≥0 by √ dXt = −Xt dt + 2 dBt , X0 ∼ q , where (Bt )t≥0 is a standard Brownian motion, and let qt := law(Xt ). Defining p σt = 1 − e−2t , we see that the transition kernel for this process is given by 1 ∥x − e−t z∥2 G Kt (x | z) = exp − . 2σt2 (2πσt2 )d/2 These kernels (KtG ) form a semigroup. 11
The (Gaussian diffusion) score at time t is the function st : Rd → Rd given by st (x) = ∇ log qt (x) , The score naturally admits an interpretation in terms of posterior expectations: Lemma 2.4 (Tweedie’s formula). For any x ∈ Rd and t > 0, E[X0 | Xt = x] = et x + et σt2 st (x) . We will consider estimated scores (b st ) and quantify their error as follows. Definition 2.5 (Score error – Gaussian diffusion). Given score estimate sbt at time t, its L2 score error is defined by 1/2 εgauss (t) := σt2 EXt ∼qt ∥st (Xt ) − sbt (Xt )∥22 . b t (x) for the denoiser m b t (x) := et x + By Lemma 2.4, σt2 st (x) − sbt (x) = e−t E[X0 | Xt = x] − m 2 −2t 2 t 2 b t (Xt ) − E[X0 | Xt ]∥2 is, up to the factor e−2t , e σt sbt (x) associated to sbt , so εgauss (t) = e EXt ∼qt ∥m the excess risk of m b t for the denoising objective E∥m b t (Xt ) − X0 ∥22 used to train Gaussian diffusion mod+ els [SDWMG15, Vin11, SE19, HJA20, SSDK 21], whose minimizer is the posterior mean. In this work, we will study Gaussian diffusion in the context of sampling from discrete distributions. For general finite alphabet Σ, let qpre be any distribution over X = Σd . We consider the following standard embedding. Given a sample A = (A1 , . . . , Ad ) from qpre over X = [S]d , apply the one-hot encoding to (1) (d) (i) obtain X0 = (X0 , . . . , X0 ) ∈ RSd , where X0 = eAi and ea ∈ RS denotes the a-th standard basis vector. We then take q to denote the pushforward of the data distribution under this one-hot encoding map. We will also consider the special case of binary alphabet Σ = {−1, 1}, in which case we will more directly regard qpre as a distribution over Rd in the natural way by regarding X = {−1, 1}d as a subset of Rd .
2.3
Masked diffusion
Given alphabet size S ≥ 2, unlike in the previous sections we distinguish between [S] and Σ, defining Σ = [S] ∪ {⋆}, where ⋆ ∈ / [S] is a special mask symbol. Let X = [S]d ,
X = Σd .
Let q be a distribution over X . As in the case of uniform diffusion, let αt = e−t . For z ∈ [S], define the one-coordinate transition kernel KtM (· | z) = αt 1[z = ·] + (1 − αt ) 1[⋆ = ·] . These kernels form a semigroup. Given X0 ∼ q, let (Xt )t≥0 denote the Markov process with transition ⊗d kernel KtM , and let qt := law(Xt ). For y ∈ X , define the set of revealed coordinates by I(y) := {i ∈ [d] : yi ̸= ⋆} . We identify y with the partial assignment (I(y), yI(y) ); conversely, given I ⊆ [d] and x ∈ X , we write (I(Xt ))
x(I) ∈ X for the string obtained from x by masking the coordinates outside I, so that Xt = X0 12
and
observing Xt is equivalent to observing (I(Xt ), X0,I(Xt ) ). We call a partial assignment (I, xI ) consistent if P[X0,I = xI ] > 0. Given a consistent partial assignment and an unrevealed coordinate i ∈ / I, define the posterior marginal qi|I (a | xI ) := P[X0,i = a | X0,I = xI ] , a ∈ [S] . Given y ∈ X , i ∈ / I(y), and a ∈ [S], let y i←a be given by replacing the i-th coordinate of y with a. The [d]×[S] given by (masked diffusion) score at time t is the function st : X → R≥0 st (y)[i, a] :=
qt (y i←a ) , qt (y)
i∈ / I(y) ,
a ∈ [S] ,
and st (y)[i, a] is undefined for i ∈ I(y). The masked diffusion score has a direct interpretation in terms of posterior marginals. / I and a ∈ [S], Lemma 2.6. Let y ∈ X be consistent and let I = I(y) ⊊ [d]. For any i ∈ st (y)[i, a] =
αt q (a | yI ) . 1 − αt i|I
P[X0,i =a,X0,I =yI ] |I|+1 , while qt (y i←a ) = αt (1 − αt )d−|I|−1 P[X0,i = P[X0,I =yI ] |I| a, X0,I = yI ] and qt (y) = αt (1 − αt )d−|I| P[X0,I = yI ]. Taking the ratio gives the desired result.
Proof. Observe that qi|I (a | yI ) =
By Lemma 2.6, the masked diffusion score is equivalent to posterior marginals, up to a deterministic time-dependent factor. We will consider estimated scores (b st ) and quantify their error as follows. First, we will assume that P t the estimated scores are consistent, i.e., that for all t, y ∈ X , and i ∈ / I(y), a∈[S] 1−α bt (y)[i, a] = 1, αt s 1−αt and define the estimated posterior marginals by qbi|I (a | yI ) := αt sbt (y)[i, a], which we assume do not depend on t, as is the case for the true posterior marginals. By Lemma 2.6, querying such a score estimate at (t, y) is then the same as querying an approximate masked diffusion oracle O with the partial assignment O (· | x ) over [S] for all i ∈ (I, xI ) = (I(y), yI(y) ) and receiving the distributions qbi|I / I; we use both I O descriptions interchangeably, writing qbi|I for qbi|I when O is clear from context. On inconsistent partial assignments the posterior marginals are undefined, and the oracle may return arbitrary distributions. Definition 2.7 (Score error – masked diffusion). Given consistent score estimate sbt at time t, its score error is defined as X X 1 − αt 1 εmask (t) := EY ∼qt ψ(st (Y )[i, a], sbt (Y )[i, a]) , αt d − |I(Y )| i∈[d]\I(Y ) a∈[S]
with the convention that the bracket is 0 when I(Y ) = [d]. For the associated oracle O and 0 ⩽ m < d, the score error of O at level m is, for any t > 0, the same quantity conditioned on |I(Y )| = m (it does not depend on t, see below): X X 1 − αt (I) (I) 1 εmask (O; q, m) := E ψ st X0 [i, a], sbt X0 [i, a] . I∼Unif ([d] ) αt d − m m i∈I / a∈[S]
X0 ∼q
Since I(Xt ) is a uniformly random subset of size |I(Xt )| ∼ Bin(d, αt ) independent of X0 , we have εmask (t) = Em∼Bin(d,αt ) [εmask (O; q, m) 1[m < d]], and εmask (O; q, m) does not depend on t (see Lemma 2.8). 13
The level m is the natural notion of noise level for masked diffusion: unlike for uniform and Gaussian diffusion, it is observable from Xt , and it is the quantity that samplers control in practice. The following lemma reinterprets the masked diffusion score error in terms of the estimated posterior marginals. Lemma 2.8. Let (b st ) be consistent, with associated oracle O. Then for every 0 ⩽ m < d, " # 1 X O εmask (O; q, m) = E KL(qi|I (· | X0,I ) ∥ qbi|I (· | X0,I )) . I∼Unif ([d] d−m m) X0 ∼q
(1)
i∈I /
In particular, εmask (O; q, m) depends on O only through the marginals it returns on consistent partial assignments with m revealed coordinates. We defer the proof to Appendix B. The right-hand side of Eq. (1) is, up to an additive term independent of O P and the weighting over levels m, the cross-entropy loss i∈I bi|I (X0,i | X0,I ) on masked positions, / − log q commonly used to train masked diffusion models [SAS+ 24, SHW+ 24].
2.4
Information-theoretic quantities
Here we formally define the relevant information-theoretic quantities that we use to quantify the intrinsic complexity of a data distribution. Definition 2.9. Given a random vector X = (X1 , . . . , Xd ), the total correlation of X is given by X TC(X) = KL(law(X) ∥ ⊗i law(Xi )) = I(Xi ; Xi+1:d ) . i
The dual total correlation of X is given by DTC(X) := H(X) −
X
H(Xi | X−i ) .
i
More generally, their conditional versions are X X TC(X | Y ) := H(Xi | Y ) − H(X | Y ) = I(Xi ; Xi+1:d | Y ) , i
i
and DTC(X | Y ) := H(X | Y ) −
X
H(Xi | X−i , Y ) .
i
The following lemma bounds the extent to which a coordinatewise channel decreases dual total correlation in terms of a sum of conditional mutual information terms. Lemma 2.10. Let A = (A1 , . . . , Ad ) and let B = (B1 , . . . , Bd ) be obtained by applying independent coordinate channels to A. Then DTC(A) − DTC(B) = DTC(A | B) +
d X i=1
14
I(Bi ; A−i | B−i ) .
In particular, d X
I(Bi ; A−i | B−i ) ⩽ DTC(A) − DTC(B) .
(2)
i=1
Proof. By definition, DTC(A) − DTC(B) − DTC(A | B) = H(A) − H(B) − H(A | B) −
X
H(Ai | A−i ) − H(Bi | B−i ) − H(Ai | A−i , B) . (3)
i
Note that H(A) − H(B) − H(A | B) = −H(B | A) = −
X
H(Bi | Ai ) ,
i
where the second equality follows because B is obtained through a product channel applied to A. For the same reason, H(Ai | A−i , B) = H(Ai | A−i , Bi ) , and H(Ai | A−i ) − H(Ai | A−i , Bi ) + H(Bi | Ai ) = H(Ai | A−i ) − H(Ai | A−i , Bi ) + H(Bi | Ai , A−i ) = H(Bi | A−i ) . Substituting these into Eq. (3), DTC(A) − DTC(B) − DTC(A | B) =
X
H(Bi | B−i ) − H(Bi | A−i )
i
Finally, note that H(Bi | A−i ) = H(Bi | A−i , B−i ), so the above can be written as X = I(Bi ; A−i | B−i ) , i
as claimed.
3
Uniform diffusion can scale with dual total correlation
In this section we prove our first main result that uniform diffusion can achieve query complexity scaling with the dual total correlation of the distribution, in analogy to what was previously shown for masked diffusion [CCL26]. Theorem 3.1. Let q be any distribution on [S]d , and suppose that score estimates (b st ) satisfying εunif (t) ⩽ εunif for all t, along with a number DTC ≥ DTC(q), are available. For every 0 < ε < 1, there is a sampler whose output law qb satisfies KL(q ∥ qb) ⩽ ε + O(εunif · log(dS/ε)) , and which uses DTC edS O 1+ log ε ε score oracle queries. 15
In Section 3.1, we prove the key estimate bounding the conditional TC in terms of the amount by which the DTC decreases when going from a lower noise level to a higher noise level, enabling a telescoping argument. In Section 3.2, we control the effect of score error. Finally, in Section 3.3, we describe the sampling algorithm and complete the proof of Theorem 3.1.
3.1
Conditional TC versus DTC decrement
The main result of this section is an upper bound on the conditional TC, TC(Xs | Xt ), in terms of the amount by which the dual total correlation decreases when going from noise level s to a higher noise level t: Lemma 3.2. For 0 < s < t, let h = t − s. Then TC(Xs | Xt ) ⩽
e2h − 1 (DTC(qs ) − DTC(qt )) . 1 − e−s
We first make a simple but key observation about the Gibbs sampling marginals which will drive the proof of the lemma above: Proposition 3.3. For 0 < s ⩽ t and for all i ∈ [d], a ∈ Σ, y ∈ X , P[(Xs )i = a | (Xt )−i = y−i ] ≥ βs . Proof. If in addition to (Xt )−i = y−i , one also conditions on all of X0 , then (Xs )i and (Xt )−i become conditionally independent. But conditioned on X0 , the probability that (Xs )i = a is, by definition of the transition kernel KsU , lower bounded by βs . Averaging this bound over the conditional law of X0 given (Xt )−i = y−i , we prove the claim. We next establish the following helper lemma, motivated by the observation above. It shows that if two distributions over Σ place some nonzero amount of mass on every element, then applying a small amount of noise does not decrease the KL divergence between them by too much, yielding a reverse data-processing inequality (DPI) for uniform diffusion on a single token. Lemma 3.4 (Reverse DPI for uniform diffusion). Let p, p′ be distributions over Σ such that mina∈Σ min{pa , p′a } ≥ c/S for some c > 0. Then e−2t c . 1 − e−t + e−t c R 1 1−θ Proof. Because for any x, y > 0, x log(x/y) − x + y = (x − y)2 0 y+θ(x−y) dθ, we have KL(pKtU ∥ p′ KtU ) ≥ KL(p ∥ p′ ) ·
Z 1 X ′ 2 KL(p ∥ p ) = (pa − pa ) ′
0
a
1−θ p′a + θ(pa − p′a )
dθ .
By assumption, for all a ∈ Σ, p′a + θ(pa − p′a ) ≥ c/S for all θ ∈ [0, 1], and thus −t
e
1−e (p′a + θ(pa − p′a )) + S
−t
⩽ e
16
−t
1 − e−t ′ + (pa + θ(pa − p′a )) . c
(4)
(1)
(2)
noise
Xt ⩽
Xs 1
i
e2h − 1 · 1 − e−s
Xs 1
d
I (Xs )i+1:d ; (Xs )i | Xt
Xt
i
d
I (Xs )i+1:d ; (Xt )i | (Xt )−i
Figure 5: The two conditional mutual informations compared in Corollary 3.5, for a sequence of d tokens at noise levels s < t. Gray cells are conditioned on, white cells are marginalized out, and the arc joins the two arguments of the mutual information. In (2), the coordinate (Xs )i is replaced by (Xt )i and only the remaining coordinates (Xt )−i are conditioned on. Additionally, ((pKtU )a − (p′ KtU )a )2 = e−2t (pa − p′a )2 . Applying the same identity that gives rise to Eq. (4), we conclude that e−2t KL(pKtU ∥ p′ KtU ) ≥ −t KL(p ∥ p′ ) , e + (1 − e−t )/c as desired. We can use the above estimate to compare two quantities: (1) the mutual information between a coordinate i of Xs and some other coordinates of Xs , conditioned on the sequence Xt at a higher noise level t, and (2) the same quantity except (Xs )i is replaced with (Xt )i and only the remaining coordinates of Xt are conditioned upon. See Figure 5 for an illustration. Corollary 3.5. For 0 < s < t, let h = t − s. Then for every i ∈ [d − 1], I((Xs )i+1:d ; (Xs )i | Xt ) ⩽ I((Xs )i+1:d ; (Xt )i | (Xt )−i ) ·
e2h − 1 . 1 − e−s
(5)
Proof. Conditioned on (Xt )−i , we have a Markov chain (Xs )i+1:d − (Xs )i − (Xt )i , as (Xt )i is generated from (Xs )i by applying KhU with randomness independent of ((Xs )i+1:d , (Xt )−i ). Note that I(A; B | C) = EA,C KL(law(B | A, C) ∥ law(B | C)), so I((Xs )i+1:d ; (Xt )i | (Xt )−i ) = E KL(law((Xt )i | (Xs )i+1:d , (Xt )−i ) ∥ law((Xt )i | (Xt )−i )),
(6)
I((Xs )i+1:d ; (Xs )i | (Xt )−i ) = E KL(law((Xs )i | (Xs )i+1:d , (Xt )−i ) ∥ law((Xs )i | (Xt )−i )) .
(7)
Conditioned on any realization of (Xs )i+1:d and (Xt )−i , the conditional distribution on (Xs )i is given by the mixture law((Xs )i | (Xs )i+1:d , (Xt )−i ) = e−s law((X0 )i | (Xs )i+1:d , (Xt )−i ) +
1 − e−s , S
so if we take p and p′ in Lemma 3.4 to be law((Xs )i | (Xs )i+1:d , (Xt )−i ) and law((Xs )i | (Xt )−i ), and if we take t therein to be h = t − s, then by the display above and Proposition 3.3, the condition in the lemma is satisfied with c = 1 − e−s , and we conclude from this and Eqs. (6) and (7) after averaging over all
17
conditionings that I((Xs )i+1:d ; (Xt )i | (Xt )−i ) ≥ I((Xs )i+1:d ; (Xs )i | (Xt )−i ) ·
e−2h (1 − e−s ) . 1 − e−h + e−h (1 − e−s )
(8)
Finally, by the Markov chain property and then by chain rule, I((Xs )i+1:d ; (Xs )i | (Xt )−i ) = I((Xs )i+1:d ; (Xs )i , (Xt )i | (Xt )−i ) = I((Xs )i+1:d ; (Xs )i | Xt ) + I((Xs )i+1:d ; (Xt )i | (Xt )−i ) .
(9)
Note that the second term on the right-hand side is the same as the left-hand side of Eq. (8). Denoting the factor on the right-hand side of Eq. (8) by C, if we substitute Eq. (9) into Eq. (8) and rearrange, we conclude that I((Xs )i+1:d ; (Xs )i | Xt ) ⩽ (1/C − 1) · I((Xs )i+1:d ; (Xt )i | (Xt )−i ) . Finally, we have 1/C − 1 =
(1 − e−h )(1 + e−h (1 − e−s )) e2h − 1 ⩽ , 1 − e−s e−2h (1 − e−s )
(10)
as claimed. This corollary allows us to establish our key result relating the conditional total correlation TC(Xs | Xt ) (in which the terms on the left-hand side of Eq. (5) appear) to the dual total correlation decrement in Eq. (2) (in which the terms on the right-hand side of Eq. (5) appear). Proof of Lemma 3.2. By the definition of TC and Corollary 3.5, TC(Xs | Xt ) =
d−1 X i=1
d−1
e2h − 1 X I((Xs )i ; (Xs )i+1:d | Xt ) ⩽ I((Xt )i ; (Xs )i+1:d | (Xt )−i ) . 1 − e−s i=1
Note that I((Xt )i ; (Xs )i+1:d | (Xt )−i ) ⩽ I((Xt )i ; (Xs )−i | (Xt )−i ) , and the sum of these over i = 1, . . . , d − 1 is, by Lemma 2.10, at most DTC(qs ) − DTC(qt ), as desired. Henceforth, for convenience, we denote the factor in Lemma 3.2 by cs (h) :=
e2h − 1 , 1 − e−s
so that TC(Xs | Xt ) ⩽ cs (h)(DTC(qs ) − DTC(qt )) .
3.2
Score error and approximate reverse kernel
In this section, we show that the approximate score sb can be converted into approximate posterior marginals, by giving a “robust” version of Lemma 2.1. We then bound the distance between these approximate posterior marginals and the true posterior marginals in terms of the score entropy loss εunif .
18
Given b ∈ Σ, define n va St,b := v ∈ RS>0 : vb = 1 and P
a′ ∈Σ va′
o ≥ βt ∀ a ∈ Σ .
As the S + 1 constraints defining this set are linear, this is a convex set. Furthermore, by the constraint corresponding to a = b, the set is compact. Proposition 3.6. For every t > 0, y ∈ X , and i ∈ [d], the true score (st (y)[i, a])a∈Σ is an element of St,yi . Furthermore, for 0 < s < t, h = t − s, and any v ∈ St,b , the distribution pgibbs [v] over Σ defined h gibbs in Lemma 2.1 is a valid probability distribution and satisfies ph [v](a) ≥ βs for all a ∈ Σ. Likewise, U P gibbs ppost [v](a) is a valid probability distribution. a′ ∈Σ va′ Kh (b | a) ph h,b [v](a) = We defer the proof to Appendix B. Let projt,b : RS>0 → RS>0 denote the Bregman projection X
projt,b (v) := arg min ′
v ∈St,b
ψ(va′ , va ) .
a∈Σ
This is uniquely defined because St,b is convex and compact and ψ is strictly convex in its first argument. By the Pythagorean theorem for Bregman divergences, for any v ∗ ∈ St,b and any v ∈ RS>0 , X
ψ(va∗ , projt,b (v)a ) ⩽
a∈Σ
X
ψ(va∗ , va ) .
(11)
a∈Σ
Given score estimates (b st ), define their projections (b s♯t ) by st (y)[i, ·]) . sb♯t (y)[i, ·] := projt,yi (b For 0 < s < t and h = t − s, we will consider the approximate reverse kernel unif bt→s R (y, x) :=
d Y
i bt→s s♯t (y)[i, ·]] . R (y, ·) := ppost h,yi [b
i bt→s R (y, xi ) ,
i=1
We will compare this against the true product reverse kernel: d Y unif i := Rt→s (y, x) Rt→s (y, xi ) , i=1
i Rt→s (y, ·) := ppost h,yi [st (y)[i, ·]] = P[(Xs )i = · | Xt = y] .
(12)
Lemma 3.7. For 0 < s < t and h = t − s, unif unif bt→s (Y, ·)) ⩽ cs (h)εunif (t) . EY ∼qt KL(Rt→s (Y, ·) ∥ R
Proof. Fix any y ∈ X and i ∈ [d] and condition on (Xt )−i = y−i . For convenience, denote by p the conditional law of (Xs )i , and also define the corresponding approximation pb = pgibbs [b s♯t (y)[i, ·]] given by h the score estimate. Additionally, given any b ∈ Σ, denote by p|b the conditional law of (Xs )i upon further conditioning on (Xt )i = b. Note that i p|yi = Rt→s (y, ·) , so we would like to bound KL(p|yi ∥ pb|yi ) for a suitable approximation pb|yi defined below. Finally, let 19
r = pKhU and rb = pbKhU , noting that r is the conditional law of (Xt )i given (Xt )−i = y−i , and rb is an approximation thereof, so that st (y)[i, a] =
r(a) r(yi )
and
sb♯t (y)[i, a] =
rb(a) . rb(yi )
By chain rule applied to the joint conditional law of (Xs )i , (Xt )i and its approximation, KL(p ∥ pb) = KL(r ∥ rb) + Eb∼r KL(p|b ∥ pb|b ) ,
(13)
where pb|b (·) := pb(·)KhU (b | ·)/b r(b) is the posterior law of a sample from pb given that its image under P U the channel Kh equals b. Writing v = sb♯t (y)[i, ·], we have rb(a) = αh pb(a) + βh = va / a′ ∈Σ va′ , and substituting this into the definition of pb|yi and using vyi = 1 gives bi pb|yi = ppost h,yi [v] = Rt→s (y, ·) . We will lower bound the expectation on the right-hand side of Eq. (13) by the contribution from b = yi to get KL(p ∥ pb) ≥ KL(r ∥ rb) + r(yi ) · KL(p|yi ∥ pb|yi ) . By the reverse DPI in Lemma 3.4, whose hypothesis holds with c = 1 − e−s by Propositions 3.3 and 3.6, KL(r ∥ rb) ≥ KL(p ∥ pb) ·
e−2h (1 − e−s ) . 1 − e−h + e−h (1 − e−s )
Recall that this factor is the same one appearing in Eq. (8), which was denoted in the proof of Corollary 3.5 by C and which satisfies 1/C − 1 ⩽ cs (h) by Eq. (10). Rearranging, we have KL(p|yi ∥ pb|yi ) ⩽
cs (h) KL(r ∥ rb) . r(yi )
We have X
r(a) rb(y ) X r(a) 1 1 i − + log · rb(yi ) r(yi ) r(yi ) rb(a) r(yi ) a∈Σ 1 1 1 =ψ , + KL(r ∥ rb) r(yi ) rb(yi ) r(yi ) 1 ≥ KL(r ∥ rb) , r(yi )
ψ(st (y)[i, a], sb♯t (y)[i, a]) =
a
so KL(p|yi ∥ pb|yi ) ⩽ cs (h)
X
ψ(st (y)[i, a], sb♯t (y)[i, a]) ⩽ cs (h)
a
X
ψ(st (y)[i, a], sbt (y)[i, a]) ,
a
where in the last step we used the Pythagorean theorem in Eq. (11). Summing over coordinates i, averaging over y ∼ qt , and recalling Definition 2.2, we conclude the claimed bound.
20
3.3
Sampler and telescoping argument
We consider a sampler of the following form. Fix a grid of times 0 < t0 < t1 < · · · < tM . 1. Initialize at Unif(X ) at time tM . 2. For each j = M − 1, . . . , 0: bunif • Apply R tj+1 →tj to get the next iterate. 3. Given the iterate Y at time t0 , output a draw from (KtU0 )⊗d (· | Y ). The following lemma provides a straightforward chain rule calculation that utilizes the main estimates from the previous subsections (Lemma 3.2 and Lemma 3.7). Lemma 3.8. Let qb denote the output law of the above sampler. Then KL(q ∥ qb) ⩽ KL(qtM ∥ Unif(X ))+H(qt0 )−H(q)+
M −1 X
h i ctj (tj+1 −tj )· DTC(qtj )−DTC(qtj+1 )+εunif (tj+1 ) .
j=0
Proof. Let P be the joint law of (XtM , XtM −1 , . . . , Xt0 , X0 ) under the forward process, written in reverse btunif →t order. Let Pb be the comparison path law that initializes its first coordinate from Unif(X ), uses R j+1 j from tj+1 to tj , and uses (KtU0 )⊗d (· | Xt0 ) for the last transition. For 0 ≤ j < M , define unif
Ejdisc := EXtj+1 KL(law(Xtj | Xtj+1 ) ∥ Rtj+1 →tj (Xtj+1 , ·)), unif
bunif Ejsc := EXtj+1 KL(Rtj+1 →tj (Xtj+1 , ·) ∥ R tj+1 →tj (Xtj+1 , ·)), E0end := EXt0 KL(law(X0 | Xt0 ) ∥ (KtU0 )⊗d (· | Xt0 )) . unif
Because the marginals of the product reverse kernel Rtj+1 →tj are by definition the marginals of law(Xtj | Xtj+1 = y), btunif →t (Xt , ·)) = Ejdisc + Ejsc . EXtj+1 KL(law(Xtj | Xtj+1 ) ∥ R j+1 j+1 j The X0 marginal of P is q, while the last-coordinate marginal of Pb is qb. Data processing and chain rule for KL therefore give KL(q ∥ qb) ≤ KL(P ∥ Pb) = KL(qtM ∥ Unif(X )) +
M −1 X
(Ejdisc + Ejsc ) + E0end .
j=0
By Lemma 3.9 below, the last term is equal to H(qt0 ) − H(q). Furthermore, by definition Ejdisc = TC(Xtj | Xtj+1 ), so by Lemma 3.2, Ejdisc ⩽ ctj (tj+1 − tj ) DTC(qtj ) − DTC(qtj+1 ) . Finally, by Lemma 3.7, Ejsc ⩽ ctj (tj+1 − tj )εunif (tj+1 ) . 21
Substituting all of these bounds proves the claim. The lemma below was used above to control the error incurred by the rounding step at the end of the sampler. Lemma 3.9. For any t > 0, we have EXt ∼qt KL(law(X0 | Xt ) ∥ (KtU )⊗d (· | Xt )) = H(qt ) − H(q) . Proof. The joint law of (X0 , Xt ) is q(x)(KtU )⊗d (y | x), while the comparison joint law is qt (y)(KtU )⊗d (x | y). The kernel is symmetric, so their log-likelihood ratio is log q(x) − log qt (y). Averaging gives −H(q) + H(qt ). We next control the two endpoint terms in Lemma 3.8 by taking t0 sufficiently small and tM sufficiently large. Lemma 3.10. For any T > 0, KL(qT ∥ Unif(X )) ⩽ e−T d log S. We defer the proof to Appendix B. Lemma 3.11. For 0 < t ⩽ 1, H(qt ) − H(q) ⩽ dt log(eS/t). We defer the proof to Appendix B. We are now ready to prove our main result, a DTC-adaptive query complexity for uniform diffusion sampling: Proof of Theorem 3.1. It remains to set the step sizes so that the sum in Lemma 3.8 telescopes (see Figure 6). Define ε , TU = log(4d log(S)/ε) , LU = log(16edS/ε), δU = 16dLU and aU =
ε , 12 max{DTC, ε}
u0 = eδU − 1,
U = eTU − 1 ,
noting that aU ⩽ 1/12. Recursively set uj+1 = min((1 + aU )uj , U ) , and let NU be the first index for which uj = U . Set tj = log(1 + uj ), so that δU = t0 < · · · < tNU = TU . Because uj /(1 + uj ) = 1 − e−tj , we have etj+1 −tj =
1 + uj+1 1 + (1 + aU )uj ⩽ = 1 + aU (1 − e−tj ) , 1 + uj 1 + uj
and hence, writing xj := aU (1 − e−tj ) ⩽ 1, ctj (tj+1 − tj ) =
(1 + xj )2 − 1 (2 + xj )xj e2(tj+1 −tj ) − 1 ⩽ = ⩽ 3aU . −t −t 1−e j 1−e j 1 − e−tj
Therefore, NX U −1
ctj (tj+1 − tj ) DTC(qtj ) − DTC(qtj+1 ) + εunif (tj+1 ) ≤ 3aU DTC(q) + 3aU NU εunif
j=0
⩽ ε/4 + 3aU NU εunif , 22
t = log(1 + u)
×(1 + aU )
u0
uj
U
(log scale)
tj TU
δU
tj+1 ≈ (1 + aU ) tj
tj+1 − tj ≈ log(1 + aU )
Figure 6: Step schedule for uniform diffusion sampler since DTC(qt ) is nonincreasing in t by Lemma 2.10, so the increments are nonnegative and telescope. )) ≤ ε/4. By Lemma 3.11 and the choice of By Lemma 3.10 and the choice of TU , KL(qTU ∥ Unif(X ε δU , H(qδU ) − H(q) ⩽ dδU log(eS/δU ) = 16 1 + logLULU ⩽ ε/8. Finally, uj increases by a factor 1 + aU until the last step, so NU ≤ 1 +
log(U/u0 ) . log(1 + aU )
bunif Here log(U/u0 ) ≲ LU and 1/ log(1 + aU ) ≤ 2/aU ≲ 1 + DTC/ε. Each application of R tj+1 →tj requires [d]×Σ only the score matrix sbtj+1 (y) ∈ R at the current iterate y, i.e., one score oracle query, so the sampler makes NU queries in total, which proves the query bound. Substituting all of the above bounds into Lemma 3.8 gives the claimed KL bound.
4
Gaussian diffusion can scale with dual total correlation
We now prove the Gaussian diffusion analogue of Theorem 3.1. Recall from Section 2.2 that we consider the pushforward q of a discrete distribution qpre on X = Σd to (RS )d via the one-hot encoding. When we refer to the dual total correlation of q, we treat each block of S = |Σ| entries as a single coordinate in the definition of DTC(q), so that DTC(q) = DTC(qpre ). Given X0 ∼ q, let (Xt ) be the Ornstein–Uhlenbeck process of Section 2.2, with qt := law(Xt ) and score st = ∇ log qt . The error of estimated scores (b st ) is quantified by the L2 score error εgauss (t) of Definition 2.5. Theorem 4.1. Let q be any one-hot encoding of a distribution qpre on X , and suppose that score estimates (b st ) satisfying εgauss (t) ⩽ εgauss for all t, along with a number DTC ≥ DTC(q), are available. For every 0 < ε < 1, there is a sampler whose output law qb on [S]d satisfies KL(qpre ∥ qb) ⩽ ε + O(ε2gauss · log(dS/ε)) , and which uses DTC edS O 1+ log ε ε score oracle queries. The general outline of the argument remains the same as the one for uniform diffusion in Section 3. An important difference is that there is no analogue of Proposition 3.3, so the reverse DPI used to relate conditional 23
TC to the decrease in DTC must be proven using a different route (see Section 4.2). Change of variable. Throughout this section, in place of the OU parametrization it will be more convenient to work with the additive noise parametrization. As such, we consider the change of variable u = e2t − 1 ,
t(u) = 12 log(1 + u) ,
Zu := et(u) Xt(u) ,
so that Zu = X0 + Bu
(14)
for a standard Brownian motion (Bu )u≥0 in RSd independent of X0 . In lieu of qt and st , we use πu and gu = ∇ log πu to denote the density of Zu and its score, noting that gu (z) := √
z 1 st(u) √ . 1+u 1+u
Similarly, we will consider the score estimate gbu (z) := √
z 1 sbt(u) √ , 1+u 1+u
which has score error 1+u 1 EXt(u) ∼qt(u) ∥b st(u) (Xt(u) ) − st(u) (Xt(u) )∥22 = εgauss (t(u))2 , 1+u u2 (15) 2 2 where in the last step we re-expressed the σt(u) prefactor in the definition of εgauss via σt(u) = u/(1 + u). EZu ∼πu ∥b gu (Zu ) − gu (Zu )∥22 =
4.1
Conditional TC versus DTC decrement
The main result of this subsection is the following analogue of Lemma 3.2: an upper bound on the conditional TC, TC(Zu | Zv ), in terms of the amount by which the dual total correlation decreases when going from noise level u to higher noise level v. Here and below, for laws with densities DTC is defined by the same formula with differential entropies, and Lemma 2.10 applies verbatim (also when A is discrete and B continuous). Lemma 4.2. For 0 < u < v, TC(Zu | Zv ) ⩽
v 3/u−3/v e − 1 DTC(πu ) − DTC(πv ) . u
As in the uniform diffusion case, the key ingredient in the proof is a reverse DPI for a single token. In the uniform diffusion case, however, the proof of the reverse DPI crucially relied on a lower bound on the mass that the relevant conditional distributions placed on each possible token. In the Gaussian setting we cannot rely on such a property. For a distribution α = (α1 , . . . , αS ) over Σ and u > 0, if we regard it as a distribution over standard basis vectors in RS and apply Gaussian noise at level u, the resulting distribution is a mixture of Gaussians with density X fα,u (z) = αa ϕu (z − ea ) , a∈Σ
where ϕu denotes the density of N (0, u · IdS ). 24
The main estimate we show is the following reverse DPI for such mixture distributions. Lemma 4.3 (Reverse DPI for Gaussian diffusion). For every α, β ∈ ∆(Σ) and 0 < u < v, u e3/v−3/u . KL(fα,v ∥ fβ,v ) ≥ KL(fα,u ∥ fβ,u ) · v Lemma 4.3 is a consequence of standard identities for the additive Gaussian channel, and we defer its proof to Section 4.2 and first derive Lemma 4.2 from it. In exact analogy with Corollary 3.5 in the uniform diffusion case, we use this bound to compare two quantities: (1) the mutual information between a block i of Zu and the later blocks of Zu , conditioned on (i) (i) the sequence Zv at a higher noise level v, and (2) the same quantity except Zu is replaced with Zv and only the remaining blocks of Zv are conditioned upon (recall Figure 5). Corollary 4.4. For 0 < u < v and i ∈ [d − 1], I(Zu(i+1:d) ; Zu(i) | Zv ) ⩽ (i+1:d)
v 3/u−3/v e − 1 · I(Zu(i+1:d) ; Zv(i) | Zv(−i) ) . u
(−i)
(i)
(i)
Proof. Conditionally on (Zu , Zv ), the block Zu has law fα,u and Zv has law fα,v with α = (i+1:d) (−i) law(Ai | Zu , Zv ), because the noise in block i is independent of the conditioning; likewise with (−i) (−i) β = law(Ai | Zv ) when conditioning on Zv alone. Lemma 4.3 therefore applies to these conditional laws, and the proof follows verbatim from the argument for Corollary 3.5 in the uniform diffusion case, with the parameter C defined therein taken to be C = uv e3/v−3/u in light of Lemma 4.3. 2h
e −1 Proof of Lemma 4.2. The proof also follows verbatim from the argument for Lemma 3.2, but where 1−e −s v 3/u−3/v therein is replaced with u e − 1.
Henceforth, for convenience denote the factor in Lemma 4.2 by cu (v) :=
v 3/u−3/v e − 1, u
so that TC(Zu | Zv ) ⩽ cu (v)(DTC(πu ) − DTC(πv )) .
4.2
Reverse DPI for Gaussian diffusion
We will use the following standard identities for the additive Gaussian channel, see, e.g., [Wib25, Lemmas 1 and 2]: Lemma 4.5. Let (µu )u≥0 and (νu )u≥0 be measures evolving according to the standard Gaussian channel, that is, which satisfy ∂u pu = 12 ∆pu for p = µ, ν. Denote the log-density ratio between them by ℓu := log(µu /νu ), and denote the relative Fisher information between them by FI(µu ∥ νu ) := Eµu ∥∇ℓu ∥22 . Then • ∂u KL(µu ∥ νu ) = − 21 FI(µu ∥ νu ) • ∂u FI(µu ∥ νu ) = −Eµu ∥∇2 ℓu ∥2F + 2Eµu (∇ℓu )⊺ (∇2 log νu )(∇ℓu ) 25
Specializing µu and νu above to fα,u and fβ,u respectively, we obtain the following: Lemma 4.6. Let α, β ∈ ∆(Σ) and u > 0. Given λ ∈ ∆(Σ), define mλ,u (z) :=
λa eza /u P zb /u b λb e
Cλ,u (z) := diag(mλ,u (z)) − mλ,u (z)⊗2 .
and a∈Σ
Then • ∂u KL(fα,u ∥ fβ,u ) = − 12 FI(fα,u ∥ fβ,u ) • ∂u FI(fα,u ∥ fβ,u ) ≥ − u2 + u32 FI(fα,u ∥ fβ,u ) We defer the proof to Appendix B. We are now ready to complete the proof of Lemma 4.3. Proof of Lemma 4.3. For convenience, define K(u) := KL(fα,u ∥ fβ,u ). By Lemma 4.6, we have 1 −K ′ (u) = FI(fα,u ∥ fβ,u ) ≥ 0 2
and
2 1 3 K ′′ (u) = − ∂u FI(fα,u ∥ fβ,u ) ⩽ − + 2 K ′ (u) . 2 u u
Let ξ = v/u > 1. By Grönwall’s inequality applied to −K ′ , for all s ≥ u we have Z ξs 2 3 1 1 + 2 dr = −K ′ (s) 2 e3/(ξs)−3/s ≥ −K ′ (s) 2 e3/v−3/u . −K (ξs) ≥ −K (s) exp − r r ξ ξ s ′
′
Note that K(∞) = 0, so Z ∞ K(u) = −
′
K (s) ds ⩽ −e
3/u−3/v
Z ∞
ξ 2 K ′ (ξs) ds =
u
u
v 3/u−3/v e K(v) , u
as claimed.
4.3
Score error and approximate reverse kernel
Note that from the true score we can read off information about the posterior per-token marginals. Indeed, by Tweedie’s formula (Lemma 2.4), mv (z) := E[X0 | Zv = z] = z + vgv (z) , and for every block i ∈ [d], m(i) v (z) = P[Ai = · | Zv = z] . Conditioned on (X0 , Zv = z), the bridge law of the forward process in Eq. (14) is u u u X0 , u 1 − IdSd . Zu ∼ N z + 1 − v v v Therefore, the reverse kernel’s marginals are given by a mixture of Gaussians: (i) Rv→u (y, ·) := law Zu(i) | Zv = y u X u u ea , u 1 − IdS . y (i) + 1 − = m(i) v,a (y) · N v v v a∈Σ
26
(16)
(i)
(i)
u (i) u To p sampleu from Rv→u (y, ·), one draws ξi = ea with probability mv,a (y) and outputs v y + (1 − v )ξi + u(1 − v ) Gi , where Gi ∼ N (0, IdS ) is drawn independently across blocks. This naturally suggests the following product kernel in analogy with Eq. (12):
G
Rv→u (y, ·) :=
d Y
(i) Rv→u (y, ·) .
i=1
In this subsection, we show that the approximate score gb can likewise give rise to an approximation of this product kernel. We then bound the distance between the approximation and the true product kernel in terms of the score error εgauss . In lieu of the Bregman projection from Section 3.2, here we simply use a Euclidean projection Π∆(Σ)d : given the rescaled estimate gbv , define m b v (z) := Π∆(Σ)d z + vb gv (z) .
(17)
Since mv (z) ∈ ∆(Σ)d by Eq. (16) and Euclidean projection onto the closed convex set ∆(Σ)d is nonexpansive, we have ∥m b v (z) − mv (z)∥2 ≤ v∥b gv (z) − gv (z)∥2 , (18) in analogy with the Pythagorean theorem in Eq. (11). We can then construct the approximate reverse kernel G bv→u R (y, ·) :=
d Y
(i) bv→u R (y, ·) ,
i=1
by replacing mv above with m b v to obtain b(i) (y, ·) := R v→u
X
m b (i) v,a (y) · N
a∈Σ
u
u u y (i) + 1 − ea , u 1 − IdS . v v v
We now proceed to bound the distance between the product kernel and its approximation. To do so, we would like to quantify the distance between two mixtures of Gaussians at a given noise level in terms of the squared Euclidean distance between their mixing coefficients: Lemma 4.7. For α, β ∈ ∆(Σ) and ρ > 0, KL(fα,ρ ∥ fβ,ρ ) ≤ e2/ρ − 1 ∥α − β∥22 . We defer the proof to Appendix B. From this, we can readily control the distance between the product kernel and its approximation in terms of the score estimation error: Lemma 4.8. For 0 < u < v, let 2(v − u) := wu (v) (1 + v) exp −1 . uv Then G G bv→u Ey∼πv KL(Rv→u (y, ·) ∥ R (y, ·)) ⩽ wu (v) εgauss (t(v))2 .
27
Moreover, if cu (v) ⩽ 1, then wu (v) ≤ 8 log(1 + cu (v)) .
(19)
Proof. As both kernels are products over blocks, it suffices to bound the KL for each block. Over a fixed block i ∈ [d], the kernels are given by mixtures of Gaussians. After translating both by (u/v)y (i) and from Lemma 4.7, with scaling by (1 − u/v)−1 , these mixtures become the mixtures fm(i) (y),ρ and fm (i) b v (y),ρ v ρ = uv/(v − u). By Lemma 4.7 and Eq. (18), d X
(i) (i) bv→u KL(Rv→u (y, ·) ∥ R (y, ·)) ≤ e2(v−u)/(uv) − 1 v 2 ∥b gv (y) − gv (y)∥22 .
i=1
Averaging over y ∼ πv and applying Eq. (15) gives the first part of the claim. For Eq. (19), write θ := (v − u)/(uv). The hypothesis implies v/u ≤ 2 and 3θ ≤ log 2, hence e2θ − 1 ≤ 4θ and also (v − u)/u ≤ 2 log(v/u). It follows that 4(v − u) 1 1 v wu (v) ≤ 4θ + ≤4 − + 8 log ≤ 8 log(1 + cu (v)) . u u v u
4.4
Sampler and telescoping argument
We consider a sampler of the following form. Fix a grid of noise levels 0 < u0 < u1 < · · · < uM . 1. Initialize at N (0, uM IdSd ) at noise level uM . 2. For each j = M − 1, . . . , 0: buG →u to get the next iterate. • Apply R j+1 j 3. Given the iterate Z at the smallest noise level u0 , round it to a string in X . The rounding in Step 3 proceeds as follows. In the uniform diffusion case, recall that we simply applied the forward kernel coordinatewise. In the Gaussian diffusion case, we will instead round the final iterate to a string in X by drawing from the product of block-wise posterior marginals (under uniform prior) at u = u0 : Qu (a1 , . . . , ad | z) :=
d Y
(i)
(i) Q(i) u (ai | z ) ,
i=1
(i) Q(i) u (a | z ) := P
exp(za /u) (i)
.
b∈Σ exp(zb /u)
We have the following analogue of the chain rule calculation from Lemma 3.8 which combines the main estimates from the previous subsections (Lemma 4.2 and Lemma 4.8). Lemma 4.9. Let qb denote the output law of the above sampler. Then KL(qpre ∥ qb) ⩽ KL(πuM ∥ N (0, uM IdSd )) + EZu0 ∼πu0 KL(law(A | Zu0 ) ∥ Qu0 (· | Zu0 )) +
M −1h X
i cuj (uj+1 ) DTC(πuj ) − DTC(πuj+1 ) + wuj (uj+1 ) εgauss (t(uj+1 ))2 .
j=0
28
Proof. The proof follows verbatim from the argument for Lemma 3.8, with the reversed forward path (XtM , . . . , Xt0 , X0 ) replaced by (ZuM , . . . , Zu0 , A), the initialization Unif(X ) by N (0, uM IdSd ), the kernel (KtU0 )⊗d by Qu0 , and Lemmas 3.2 and 3.7 by Lemmas 4.2 and 4.8. The error from the rounding step EZu0 ∼πu0 KL(law(A | Zu0 ) ∥ Qu0 (· | Zu0 )) will be analyzed in a subsequent lemma. We next control the initialization and rounding errors in Lemma 4.9 by taking uM sufficiently large and u0 sufficiently small, respectively; the following two results play the roles of Lemmas 3.10 and 3.11. Lemma 4.10. For every U > 0, KL(πU ∥ N (0, U IdSd )) ≤
d . 2U
Proof. By convexity of KL in its first argument, it suffices to compare N (x, U IdSd ) with N (0, U IdSd ) for a fixed x whose blocks are standard basis vectors. As the two covariances agree, the KL is simply ∥x∥22 /(2U ) = d/(2U ). Lemma 4.11. For u > 0, let p∗ (u) :=
1 S−1 erfc √ . 2 2 u
If p∗ (u) ≤ 1/2, then EZu ∼πu KL(law(A | Zu ) ∥ Qu (· | Zu )) ≤ d h2 (p∗ (u)) + p∗ (u) log(S − 1) . Proof. The left-hand side is at most E[− log Qu (A | Zu )], by nonnegativity of entropy. As Qu is a product P (i) (i) kernel, this decomposes into a sum di=1 E[− log Qu (Ai | Zu )]. Fix any block i ∈ [d] and a ∈ Σ, and consider the contribution of this block conditioned on Ai = a, namely (i) E[− log Q(i) (20) u (Ai | Zu ) | Ai = a] . This has the following interpretation. Nature samples A′ uniformly at random from Σ and γ ∼ N (0, uIdS ), (i) and one observes the vector z = eA′ + γ. Then Qu (a | z) is the observer’s posterior probability that (i) A′ = a. Conditional on A′ = a, we have z = ea + γ, which has the same distribution as Zu conditional on Ai = a, so in Eq. (20) the expectation is over the randomness of γ. As a result, Eq. (20) is independent of a. So we may freely replace the expectation over the marginal distribution of Ai under A ∼ qpre in (i) (i) E[− log Qu (Ai | Zu )] with an expectation over any distribution, in particular, over Ai ∼ Unif(Σ). Thus, (i) (i) E[− log Qu (Ai | Zu )] is nothing more than the conditional entropy of a uniformly random element A′ of Σ conditioned on observing eA′ + γ. Consider any decoder F : RS → Σ that tries to predict the former given the latter, and let perr := P[F (eA′ + γ) ̸= A′ ] denote the decoding error. By Fano’s inequality, the conditional entropy is at most h2 (perr ) + perr log(S − 1) , where h2 denotes binary entropy. Take the decoder F to simply choose the largest coordinate of eA′ + γ. For any realization A′ = a, the decoding error is the probability that γb − γa > 1 for some b ̸= a, which by √ S−1 a union bound is at most 2 erfc 1/(2 u) . Since perr ⩽ p∗ (u) ⩽ 1/2 and p 7→ h2 (p) + p log(S − 1) is increasing on [0, 1/2], the claimed bound then follows by summing over i ∈ [d]. We are now ready to prove our main result, a DTC-adaptive query complexity for Gaussian diffusion:
29
FG (u) = log u − u3
+ log(1 + aG ) FG (u0 )
FG (U )
FG (uj )
uj (log scale) u 0
U 3 3 − uj+1 uj
≈ log(1 + aG )
uj+1 ≈ (1 + aG ) uj
Figure 7: Step schedule for Gaussian diffusion sampler Proof of Theorem 4.1. It remains to set the noise levels so that the sum in Lemma 4.9 telescopes (see Figure 7). Define 4d 1 , , U = max 1, LG = log(64edS/ε), u0 = 8LG ε and aG =
ε , 4 max{DTC, ε}
FG (u) = log u −
3 , u
noting that aG ⩽ 1/4. Recursively set uj+1 by FG (uj+1 ) − FG (uj ) = min {log(1 + aG ), FG (U ) − FG (uj )} , and let NG be the first index for which uj = U ; this is well defined because FG is increasing. In the proof of Theorem 3.1, the schedule was geometric in the parameter u = et − 1; here, the potential FG replaces log u as the clock because FG (v) − FG (u) = log(v/u) + 3/u − 3/v = log(1 + cu (v)). In particular, every step satisfies uj+1 3/uj −3/uj+1 cuj (uj+1 ) = e − 1 ⩽ aG , uj and therefore NX G −1
cuj (uj+1 ) DTC(πuj ) − DTC(πuj+1 ) ⩽ aG DTC(πu0 ) ⩽ aG DTC(q) ⩽ ε/4 ,
j=0
since DTC(πu ) is nonincreasing in u by Lemma 2.10. Next, we handle the contribution from score estimation error. We have log(1 + cuj (uj+1 )) = log(uj+1 /uj ) + 3/uj − 3/uj+1 ⩽ log(1 + aG ) , so cuj (uj+1 ) ⩽ 1 and thus by Eq. (19), wuj (uj+1 ) ⩽ 8 log(1 + cuj (uj+1 )) . Furthermore, NX G −1
log(1 + cuj (uj+1 )) = FG (U ) − FG (u0 ) ⩽ log(U/u0 ) + 3/u0 ⩽ log(8LG U ) + 24LG ≲ LG . (21)
j=0
30
Therefore, NX G −1
wuj (uj+1 ) εgauss (t(uj+1 ))2 ≲ ε2gauss LG .
j=0
Next, we handle the initialization error. By Lemma 4.10 and the choice of U , KL(πU ∥ N (0, U IdSd )) ⩽
d ε ⩽ . 2U 8
Next, for the rounding error, we wish to apply Lemma 4.11. By standard Gaussian tail bounds, p∗ (u0 ) ⩽ Se−2LG ⩽ 1/2. Furthermore, h2 (p∗ (u)) ⩽ p∗ (u) log(e/p∗ (u)). Hence, EZu0 ∼πu0 KL(law(A | Zu0 ) ∥ Qu0 (· | Zu0 )) ⩽ dSe−2LG (1 + 2LG ) ⩽ dSe−LG ⩽
ε . 4
Finally, we bound the total number of score oracle queries. By design, we have NG ≤ 1 +
FG (U ) − FG (u0 ) . log(1 + aG )
Recall from Eq. (21) that FG (U ) − FG (u0 ) ≲ LG . Additionally, 1/ log(1 + aG ) ≤ 2/aG ≲ 1 + DTC/ε. buG →u requires only the score vector gbu (y) ∈ RSd at the current iterate y, from Each application of R j+1 j+1 j bG are formed. So the sampler makes NG queries in which Eq. (17) and the approximate reverse kernel R total, which proves the query bound. Substituting all of the above bounds into Lemma 4.9 gives the claimed KL bound.
5
Construction of random empirical measure
In the second part of this paper, we turn to proving upper and lower bounds to separate the three paradigms of diffusion language modeling. Here, we first briefly present the family of distributions that we consider, which are based on a construction from recent work [XP26] proving a query complexity lower bound for Gaussian diffusion. Fix constants 0 < κ− < κ+ < log 2. In the subsequent sections, we will specialize to the case of binary alphabet, so that 1 − e−t βt = Σ = {±1} , X = {±1}d , . 2 Let C ⊆ X denote the subset of the hypercube defining the codebook, and let M = |C| denote the size of the codebook. Throughout, we will take M = ⌈eκd ⌉ for constant κ ∈ [κ− , κ+ ]. The compact interval [κ− , κ+ ] ensures that all the critical values below remain in a fixed compact set. Definition 5.1. Given a subset C ⊆ {±1}d , denote the uniform distribution over C by q C := Unif(C) . The random empirical measure we consider is q C for C a uniformly random M -element subset of {±1}d . We will often refer to the elements of C as codewords. 31
In all of our bounds in the sequel, we assume that the sampler: • Knows the interval [κ− , κ+ ] but does not know κ • Knows that it has approximate score oracle access to a random empirical measure in the sense of Definition 5.1, but does not know the specific realization of C Notation for subsequent sections. In our analysis, constants denoted by c, C may depend on the known parameters κ− , κ+ , but on no other parameter; dependence on additional parameters is indicated by a subscript. For the uniform and Gaussian forward processes, we will overload the notation qtC to denote the law of Xt under either forward process, with X0 ∼ q C . All three main results below as well as the uniform lower bound in Appendix A use Definition 5.1. The upper bounds we prove for uniform and Gaussian diffusion hold for any κ ∈ [κ− , κ+ ], whereas the lower bounds for masked and uniform diffusion hold even if κ is additionally known to be restricted to the lattice {(k/d) log 2 : k ∈ Z}, in which case M = eκd = 2k . Define the scalar mutual information for the uniform and Gaussian channels by 1 + e−t 1 − e−t log(1 + e−t ) + log(1 − e−t ), 2 2 p Igauss (t) := I V ; e−t V + 1 − e−2t G , V ∼ Unif({±1}), Iunif (t) :=
G ∼ N (0, 1) .
Here and below, h2 (p) := −p log p − (1 − p) log(1 − p) denotes the binary entropy function (in nats); note that Iunif (t) = log 2 − h2 (βt ). Define critical times t⋆unif and t⋆gauss to be the unique solutions of d Iunif (t⋆unif ) = log M
and
d Igauss (t⋆gauss ) = log M .
(22)
Additionally, let 0 < τ− < τ+ be (known) constants such that t⋆unif , t⋆gauss ∈ [τ− , τ+ ] for every κ ∈ [κ− , κ+ ] and all sufficiently large d, and define I0 = [τ− /2, τ+ + 1] . (23) In the proofs of our upper bounds, all times queried by our algorithms will lie in I0 defined in Eq. (23), though we place no such constraint in the proofs of our lower bounds.
√ e d) upper bound for uniform diffusion 6 O( In this √ section, we give an algorithm for approximately sampling from random empirical measures using e O( d) queries to any approximate uniform diffusion score oracle for such a distribution. This query complexity is tight: in Appendix A we prove a matching lower bound. Theorem 6.1 (Uniform diffusion with approximate scores). There are constants B, C, c0 , d0 > 0, depending only on [κ− , κ+ ], such that the following holds. Let d ≥ d0 , let ε, δ ∈ (0, 1/4), and set Λ = log Suppose Λ ⩽ c0 d.
32
Cd . εδ
(24)
Given query access to any approximate uniform diffusion score oracle (b st ) for q C with score error satisfying εδ B −BΛ = sup εunif (t) ⩽ e , (25) Cd t∈I0 √ there is an algorithm using at most O( dΛ/ε2 ) score oracle queries such that with probability at least 1 − δ over C which is sampled uniformly from all M -element subsets of X , the output law of the algorithm qbC satisfies TV(b q C , q C ) ⩽ ε. In particular, for constant ε and any polynomially small δ = d−Θ(1) , polynomially small score accuracy √ suffices, and the algorithm uses O( d log d) score queries. For the rest of the section, Λ has the value in Eq. (24). We may decrease c0 and increase d0 later when necessary so that log d ⩽ Λ ⩽ c0 d. In Section 6.1, we quantify the critical window: above the window, qtC is close to the uniform distribution, while below the window the original sample can be recovered from its noisy observation. In Section 6.2, we show how to locate this window using approximate score oracle queries at points sampled from Unif(X ). Finally, in Section 6.3, we describe how to simulate the reverse process across the window and then decode the codeword, completing the proof of Theorem 6.1.
6.1
Quantifying the critical window
In this section, we characterize the critical window of noise levels such that, for higher noise levels t, the score of qtC is close to that of the uniform distribution, and for lower noise levels t, given a noisy sample from qtC , it is possible to decode the original sample from q C that gave rise to it with high accuracy. Throughout, we will denote the likelihood ratio between qtC and the uniform distribution over X by qtC (·) . 2−d
Lt (·) :=
First, we bound the expected distance between qtC and the uniform distribution for large t. Lemma 6.2. For any integer radius 0 ⩽ r ⩽ d, 1 EC TV(qtC , Unif(X )) ⩽ √ (1 + e−t )d/2 · tanh(t/2)(r+1)/2 + P[Bin(d, βt ) ⩽ r] . 2 M We defer the proof to Appendix B. Next, we show that given a noisy sample from qtC for small t, it is possible to recover the original sample from q C with high accuracy. Lemma 6.3. Given C ⊆ {±1}d , define the Bayes-optimal recovery error errU C (t) := inf
P
zb(·) z∼q C ,x∼(KtU )⊗d (·|z)
[b z (x) ̸= z] .
For any integer radius 0 ⩽ r ⩽ d, EC errU C (t) ⩽ M · P[Bin(d, 1/2) ⩽ r] + P[Bin(d, βt ) > r] . 33
We defer the proof to Appendix B. Setting parameters appropriately and using the two estimates above, we can thus characterize the critical window for sampling from the random empirical measure using uniform diffusion: above the critical window, the forward process has sufficiently mixed so that the law of Xt is close to uniform, while below the critical window, the Bayes-optimal decoder can recover the original codeword with high probability. Proposition 6.4 (Critical window). For every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are cK,A , dK,A > 0, depending only on A, K, and [κ− , κ+ ], such that the following holds whenever d ≥ dK,A and log d ⩽ Λ ⩽ cK,A d. Set r w=A
Λ . d
(26)
Then we have EC TV(qtC , Unif(X )) ⩽ e−KΛ ,
t ≥ t⋆unif + w,
−KΛ EC errU , C (t) ⩽ e
0 ⩽ t ⩽ t⋆unif − w .
C Proof. We first observe that errU C (t) is nondecreasing in t and, by data processing inequality, TV(qt , Unif(X )) is nonincreasing in t. Therefore, it suffices to consider the endpoints t = t⋆unif + w and t = t⋆unif − w. By decreasing cK,A we may assume w ⩽ min{τ− /2, 1}; then these points are within I0 , so βt is bounded away from 0 and 1/2 and
cw ⩽ |βt − βt⋆unif | ⩽ Cw,
for t ∈ {t⋆unif − w, t⋆unif + w} ,
(27)
d because dt βt = e−t /2 is bounded above and below by positive constants on I0 . Here c, C > 0 depend only on [κ− , κ+ ].
Above the critical window. Set t = t⋆unif +w, ∆ = βt −βt⋆unif , and truncation radius r = ⌊d(βt +βt⋆unif )/2⌋. The radius is at least d∆/2 below the mean dβt of Bin(d, βt ). Hoeffding’s inequality therefore gives 2
P[Bin(d, βt ) ⩽ r] ⩽ e−d∆ /2 .
(28)
It remains to bound the first term in Lemma 6.2. Eq. (22) gives log M = d(log 2 − h2 (βt⋆unif )). Rewrite 1 + e−t = 2(1 − βt ) and tanh(t/2) = βt /(1 − βt ). Since r + 1 ≥ dβt⋆unif + d∆/2, log
(1 + e−t )d tanh(t/2)r+1 M
∆ βt ⩽ d log 2(1 − βt ) + d βt⋆unif + log − log M 2 1 − βt d∆ 1 − βt = −d KL(Ber(βt⋆unif ) ∥ Ber(βt )) − log 2 βt ⩽ −cd∆ .
Lemma 6.2 and Eqs. (28)–(29) now give 2
EC TV(qtC , Unif(X )) ⩽ 12 e−cd∆ + e−d∆ /2 .
34
(29)
Below the critical window. Set t = t⋆unif − w, ∆ = βt⋆unif − βt , and again r = ⌊d(βt + βt⋆unif )/2⌋. Since r + 1 > dβt + d∆/2, Hoeffding’s inequality gives 2
P[Bin(d, βt ) > r] ⩽ e−d∆ /2 .
(30)
To control the chance of another point entering the ball, a standard binomial tail bound gives ⋆ M P[Bin(d, 1/2) ⩽ r] ⩽ exp log M − d log 2 − h2 ((βt + βtunif )/2) = exp −d h2 (βt⋆unif ) − h2 ((βt + βt⋆unif )/2) ⩽ e−cd∆ .
(31)
In the last inequality, we used the fact that βt⋆unif − (βt + βt⋆unif )/2 = ∆/2, and that h′2 (p) = log((1 − p)/p) is bounded away from 0 for p between (βt +βt⋆unif )/2 and βt⋆unif . Combining Eqs. (30) and (31) with Lemma 6.3 yields 2 −cd∆ EC errU + e−d∆ /2 . C (t) ⩽ e Finally, Eq. (27) implies d∆2 ≥ cA2 Λ. Since w ⩽ 1, d∆ ≥ cdw ≥ cdw2 = cA2 Λ. Choosing A large 2 enough in terms of K makes e−cd∆ + e−d∆ /2 at most e−KΛ . The monotonicity of TV(qtC , Unif(X )) and errU C (t) in t then gives the claimed bounds. p For the rest of the section, let w = A Λ/d as in Eq. (26), where A > 0 is a constant to be fixed below. We may decrease c0 and increase d0 later when necessary.
6.2
Identifying the critical window
Having characterized the critical window, we now give an algorithm for detecting its location given approximate uniform diffusion score access to q C . Given x ∈ X , define the denoiser mt (x) := E[X0 | Xt = x]. We will consider the following statistic: d
Tt (x) := ⟨x, mt (x)⟩ − de
−t
1 − e−2t X = (1 − st (x)[i, −xi ]) . 2e−t i=1
Intuitively, ⟨x, mt (x)⟩ corresponds to the overlap between x and its denoising with respect to the true data distribution q C , and de−t is the same but under denoising with respect to the null distribution Unif(X ). As t goes to ∞, the difference between these tends to zero, but as t goes to 0, if x is a uniformly random point from X , then we expect the denoiser under q C to achieve worse overlap than the denoiser under Unif(X ). This is formalized in the following: Proposition 6.5. There are constants a0 , b0 > 0, depending only on [κ− , κ+ ], with b0 > 4a0 , such that for every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are cK,A , dK,A > 0 such that the following holds whenever d ≥ dK,A and log d ⩽ Λ ⩽ cK,A d: P
[Tt (x) < −a0 dw] ⩽ e−KΛ ,
t ≥ t⋆unif + 2w ,
P
[Tt (x) > −b0 dw] ⩽ e−KΛ ,
t ⩽ t⋆unif − 2w ,
C,x∼Unif(X )
C,x∼Unif(X )
uniformly over t ∈ I0 .
35
Proof. The overlap ⟨x, mt (x)⟩ can be expressed succinctly as A′x (θ(t)) where Ax (θ) := log Ez∼Unif(C) [eθ⟨x,z⟩ ] , and θ(t) := atanh(e−t ). Also define B(θ) := d log cosh(θ) , observing that B ′ (θ(t)) = de−t . We have log Lt (x) = Ax (θ(t)) − B(θ(t))
and
Tt (x) = A′x (θ(t)) − B ′ (θ(t)) .
High-noise regime (t ≥ t⋆unif + 2w). Choose neighboring noise levels t and t+ = t + ηw and define θ = θ(t) and θ+ = θ(t+ ) so that 0 ⩽ θ − θ+ ≍ ηw. The function Ax (θ) is convex in θ, so A′x (θ) ≥ and thus Tt (x) ≥
Ax (θ) − Ax (θ+ ) , θ − θ+
log Lt (x) − log Lt+ (x) B(θ) − B(θ+ ) ′ + − B (θ) . θ − θ+ θ − θ+
Because B ′′ (u) = d sech2 (u) ⩽ d, by Taylor’s theorem the second term above is ≥ −C ′ dηw for some constant C ′ > 0. On the other hand, since EX∼Unif(X ) |Lt (X) − 1| = 2TV(qtC , Unif(X )), Proposition 6.4 and Markov’s inequality show that Lt (X), Lt+ (X) ∈ [1/2, 3/2] except with probability at most e−(K+1)Λ . On this event, the first term is ≥ −C ′′ /(ηw) for some constant C ′′ > 0. With constant η, 1/(ηw) = o(dw) since dw2 = A2 Λ = Ω(log d). Choosing constant η small enough, we can ensure −C ′′ /(ηw) − C ′ dηw ≥ −a0 dw for any prescribed constant a0 > 0. Low-noise regime (t ⩽ t⋆unif − 2w). First note that A′x (θ(t)) is an average of overlaps of the form ⟨x, z⟩ for z ∈ C. We bound the max over these overlaps by a Chernoff and union bound calculation. By Chernoff, for any u > 0, 1+u P [⟨x, z⟩ ≥ du] ⩽ e−d KL(Ber( 2 ) ∥ Ber(1/2)) . z∼Unif(X )
−t⋆unif
So for u = e + O(Λ/d), noting that KL(Ber( 1+u 2 ) ∥ Ber(1/2)) = Iunif (log(1/u)) and recalling that ⋆ Iunif (tunif ) = log(M )/d, we conclude by a union bound over C that with probability at least 1 − e−Θ(Λ) over the randomness of C, ⋆ max⟨x, z⟩ ⩽ de−tunif + O(Λ) . z∈C
Henceforth condition on this event. Then ⋆
Tt (x) ⩽ de−tunif − de−t + O(Λ) . ⋆
When t ⩽ t⋆unif − 2w, we have de−tunif − de−t ⩽ −b′ dw for some constant b′ > 0. On the other hand, since √ Λ ⩽ cK,A d gives Λ ⩽ ( cK,A /A) dw, the O(Λ) term is at most b′′ dw for some constant b′′ ∈ (0, b′ ) once cK,A is small enough. Taking b0 = b′ − b′′ , the second claimed bound follows. Take a0 small enough so that b0 > 4a0 . The quantity Tt (x) is therefore a good statistic for identifying the critical window. In practice, we will not have access to mt (x) or st (x), but we can estimate Tt (x) using estimated scores sbt (x), which we will show 36
are accurate enough to identify the critical window. We first define the empirical version of Tt (x). Let ℓt = βt /(1 − βt ). For a fixed initial point z, flipping one coordinate of x changes (KtU )⊗d (x | z) by a factor of either ℓt or ℓ−1 t . The score is a weighted average of these ratios, with weights given by the posterior distribution of z. So for every t > 0, x ∈ X , and i ∈ [d], ℓt ⩽ st (x)[i, −xi ] ⩽ ℓ−1 t .
(32)
In place of the Bregman projection from Section 3, here we simply clip the entries of the score vector to this range to get the clipped approximate score set (x)[i, −xi ] := min ℓ−1 bt (x)[i, −xi ]) . t , max(ℓt , s
(33)
Define the following approximation to Tt that the sampler can actually form using approximate score access: d −2t X
1−e Tbt (x) := 2e−t
1 − set (x)[i, −xi ] .
(34)
i=1
The statistic Tbt (x) can be computed with only one approximate score oracle query at x. Here we show that Tbt (x) is a good estimate of Tt (x). We will quantify this relative to the squared Euclidean error of the clipped score: d X 2 Ft (x) := set (x)[i, −xi ] − st (x)[i, −xi ] . (35) i=1
Lemma 6.6. Uniformly over t ∈ I0 and x ∈ X , we have p EX∼qtC Ft (X) ⩽ C ′ εunif (t) |Tbt (x) − Tt (x)| ⩽ C dFt (x) ,
(36)
for constants C, C ′ > 0 depending only on [κ− , κ+ ]. We defer the proof to Appendix B. We record the following standard estimates for Hamming-sphere sizes and binomial point masses, which will be used repeatedly below. Lemma 6.7. For 1 ⩽ r ⩽ d − 1 and v = r/d, edh2 (v) ⩽ d+1
d ⩽ edh2 (v) . r
(37)
Consequently, for any p ∈ (0, 1), 1 1 d(v − p)2 P[Bin(d, p) = r] ≥ exp (−d KL(Ber(v) ∥ Ber(p))) ≥ exp − . d+1 d+1 p(1 − p)
(38)
Moreover, for all v, p ∈ (0, 1), KL(Ber(v) ∥ Ber(p)) ⩽
(v − p)2 . p(1 − p)
We defer the proof to Appendix B. The following lemma allows us to control the likelihood ratio Lt (x) near the critical time t⋆unif . 37
(39)
Lemma 6.8. For every A, K, L > 0, there are constants CA,K,L , cA,K,L , dA,K,L > 0, depending only on A, K, L and [κ− , κ+ ], such that whenever d ≥ dA,K,L , log d ⩽ Λ ⩽ cA,K,L d, and |t − t⋆unif | ⩽ Lw, we have P[Lt (x) < e−CA,K,L Λ ] ⩽ e−KΛ , C
uniformly over x ∈ X . We defer the proof to Appendix B. We can now combine all of the ingredients above to prove that Tbt (x) can be used to identify the critical window. Proposition 6.9. There are constants 0 < a < b, depending only on [κ− , κ+ ], such that for every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are BK,A , cK,A , dK,A > 0 such that the following holds whenever d ≥ dK,A , log d ⩽ Λ ⩽ cK,A d, B ≥ BK,A , and the score oracle satisfies Eq. (25): P
[Tbt (x) < −adw] ⩽ e−KΛ ,
t ≥ t⋆unif + 2w ,
(40)
P
[Tbt (x) > −bdw] ⩽ e−KΛ ,
t⋆unif − 3w ⩽ t ⩽ t⋆unif − 2w .
(41)
C,x∼Unif(X )
C,x∼Unif(X )
uniformly over t ∈ I0 . Proof. Let a0 , b0 be the constants in Proposition 6.5, chosen so that b0 > 4a0 , and set a := 2a0
and
b := b0 − a0 .
Then 0 < a < b. By Lemma 6.6, the event |Tbt − Tt | > a0 dw implies Ft > cdw2 , and its probability under the true noised law qtC is at most Ce−BΛ /(dw2 ). High-noise regime (t ≥ t⋆unif +2w). Replacing qtC by the uniform law Unif(X ) adds at most TV(qtC , Unif(X )). Averaging over C and using Proposition 6.4 with exponent K +2 therefore gives an exponentially small error probability for a uniform query as well. Therefore, Markov’s inequality implies P
[|Tbt (x) − Tt (x)| > a0 dw] ⩽
C,x∼Unif(X )
Ce−BΛ + e−(K+2)Λ . dw2
(42)
Since dw2 = A2 Λ ≥ 1, choosing B sufficiently large makes the right-hand side at most e−(K+1)Λ for all sufficiently large d. On the other hand, Proposition 6.5 with exponent K + 2 gives P
[Tt (x) < −a0 dw] ⩽ e−(K+2)Λ .
C,x∼Unif(X )
If neither event in Eqs. (42) and (43) occurs, then Tbt (x) ≥ Tt (x) − a0 dw ≥ −2a0 dw = −adw .
38
(43)
Consequently, P
[Tbt (x) < −adw] ⩽
C,x∼Unif(X )
[Tt (x) < −a0 dw] +
P
[|Tbt (x) − Tt (x)| > a0 dw]
P
C,x∼Unif(X )
C,x∼Unif(X )
⩽ e−(K+2)Λ + e−(K+1)Λ ⩽ e−KΛ for all sufficiently large d, proving Eq. (40). Low-noise regime (t ∈ [t⋆unif − 3w, t⋆unif − 2w]). Lemma 6.8 (with L = 3) implies Lt (x) ≥ e−CΛ , where ′ C = CA,K ′ ,3 , with probability at least 1 − e−K Λ , where K ′ can be chosen arbitrarily large. On the event Lt (x) ≥ e−CΛ , we have the pointwise bound 2−d 1[Lt (x) ≥ e−CΛ ] ⩽ eCΛ qtC (x). Thus Markov’s inequality gives C ′ e−(B−C)Λ ′ P [|Tbt (x) − Tt (x)| > a0 dw] ⩽ e−K Λ + . (44) dw2 C,x∼Unif(X ) Choose K ′ > K + 2 and then with A fixed, choose B large enough so that the last display is at most e−(K+1)Λ . Proposition 6.5 with exponent K + 2 gives that throughout this interval [Tt (x) > −b0 dw] ⩽ e−(K+2)Λ ,
P
(45)
C,x∼Unif(X )
since every such t satisfies t ⩽ t⋆unif − 2w. If neither the event in Eq. (44) nor the event in Eq. (45) occurs, then Tbt (x) ⩽ Tt (x) + a0 dw ⩽ −(b0 − a0 )dw = −bdw . Therefore, P
[Tbt (x) > −bdw] ⩽
C,x∼Unif(X )
P
[Tt (x) > −b0 dw] +
C,x∼Unif(X )
P
[|Tbt (x) − Tt (x)| > a0 dw]
C,x∼Unif(X )
⩽ e−(K+2)Λ + e−(K+1)Λ ⩽ e−KΛ for all sufficiently large d. This proves Eq. (41). From now on, fix A ≥ A12 and choose B ≥ B12,A . We may decrease c0 and increase d0 so that c0 ⩽ c12,A , d0 ≥ d12,A , and w ⩽ min(τ− /16, 1/16). We further decrease c0 and increase d0 below whenever necessary. We can now conclude the proof of the main result of this subsection, namely an algorithm for estimating the critical noise level t⋆unif , essentially by scanning from high to low noise levels and stopping when the median of Tbt over enough random inputs becomes sufficiently negative (see Algorithm 1 for details). Proposition 6.10. Under the hypotheses of Theorem 6.1, Algorithm 1 uses O(Λ/w) score oracle queries. For at least a 1 − δ/4 fraction of C, it satisfies P[t⋆unif − 3w ⩽ b t ⩽ t⋆unif + 2w | C] ≥ 1 − ε/8 .
(46)
We defer the proof to Appendix B.
6.3
Sampling after locating the window
Having found the window, we initialize right at the upper edge of the window, using the uniform distribution, and approximately simulate the reverse process until we reach its lower edge. Our algorithm for simulation 39
Algorithm 1: Finding the uniform diffusion critical window Input: d, [κ− , κ+ ], w, Λ, and query access to the score oracle sb. Output: An estimate b t of the critical time t⋆unif . 1 Let T = {t1 > t2 > · · · > tm } be a decreasing grid covering [τ− − 3w, τ+ + 3w] with mesh w/20; 2 Set R ← 2⌈CΛ⌉ + 1; 3 for t ∈ T do 4 for j ← 1 to R do 5 Sample X (j) ∼ Unif(X ); 6 Query sbt (X (j) ) and compute Tbt (X (j) ) using Eqs. (33)–(34); 7 if median Tbt (X (1) ), . . . , Tbt (X (R) ) < − a+b 2 dw then 8 return t 9 return τ− ;
is a naive discretization; unlike the DTC-adaptive sampler in Section 3, we do not need to carefully tune the step sizes, and we use the clipped score in lieu of the Bregman projection. First, we specialize the first part of Lemma 2.1 to the binary alphabet setting that we consider here. For 0 < s < t, write h = t − s and define the affine function ph (r) =
βh (1 − βh )r − βh . αh
(47)
Lemma 2.1 implies that ph (st (y)[i, −yi ]) is the probability that (Xs )i ̸= yi conditional on Xt = y. As in Section 3, in a single step of our sampler it will flip coordinate i with probability pbit→s (y) = ph set (y)[i, −yi ] ,
(48)
independently over i. The only distinctions from Section 3 are that here we use the clipped score set , and because the alphabet is binary, we only have to specify a single probability for each coordinate. As before, bunif . we will denote the resulting approximate product reverse kernel by R t→s Next we bound the sampling error coming from using set instead of st in this product reverse kernel. Lemma 6.11. There are constants C, h0 > 0, depending only on [κ− , κ+ ], such that whenever s, t ∈ I0 and i (y, ·) = law(X i | X = y) 0 < h = t−s ⩽ h0 , the probabilities in Eq. (48) belong to (0, 1). Denoting Pt→s t s i (y, ·) the Bernoulli law defined by Eq. (48), we have and Pbt→s EY ∼qtC
d X
i i KL(Pt→s (Y ) ∥ Pbt→s (Y )) ⩽ Chεunif (t) .
(49)
i=1
We defer the proof to Appendix B. Finally, we control the overall error incurred by our sampler over the course of crossing the critical window. Here we re-use some of the basic chain rule calculations from Section 3, but because we take steps of size ≍ ε2 /d, we do not need to tune the step sizes so as to take advantage of telescoping. Proposition 6.12. Fix t− < t+ in I0 , let ∆ = t+ − t− , and let rj = t− + j∆/N for 0 ⩽ j ⩽ N . Suppose b be the composition of the approximate product reverse kernels from rN down to h = ∆/N ⩽ h0 , and let R 40
r0 . If εunif (rj ) ⩽ η for 1 ⩽ j ⩽ N , then any initial distribution ρ satisfies r b qtC ) ⩽ TV(ρ, qtC ) + C TV(ρR, − +
d∆ + ∆η . N
(50)
Proof. First initialize from qtC+ . Let P be the true reverse path law on (XrN , . . . , Xr0 ) and let Pb be the path law obtained from the approximate product kernels. Write Rj (y, ·) = law(Xrj | Xrj+1 = y) and bj (y, ·) = Q Pbri →r (y, ·). For every y, the identity R i j+1 j bj (y, ·)) = TC(Xr | Xr KL(Rj (y, ·) ∥ R j j+1 = y) +
d X
KL(Prij+1 →rj (y, ·) ∥ Pbrij+1 →rj (y, ·))
i=1
separates discretization error and score approximation error. The KL chain rule, Lemma 3.2 (with s = rj ≥ t− ), Lemma 6.11, and telescoping therefore give N −1
KL(P ∥ Pb) ⩽
N −1
X e2h − 1 X C C DTC(q ) − DTC(q ) + Chεunif (rj+1 ) r r j j+1 −t 1−e − j=0
⩽C
j=0
∆ ∆ DTC(qtC− ) − DTC(qtC+ ) + C N N
N −1 X
εunif (rj+1 )
j=0
d∆ ⩽C + C∆η , N where DTC(qtC− ) ⩽ H(qtC− ) ⩽ d log 2. Data processing inequality and Pinsker’s inequality imply Eq. (50) b qtC R) b ⩽ TV(ρ, qtC ). when the initial law is qtC+ . Replacing it by ρ adds TV(ρR, + + Below the window, the posterior on the codeword that would have generated the noisy iterate is typically concentrated on a single point of C. We can recover that point by taking the signs of the estimated posterior expectation. Use Lemma 2.1 once more, now for the transition from t to 0, and define m b t (y)i = yi 1 − 2pt set (y)[i, −yi ] ,
zbt (y) = sgn(m b t (y)1 ), . . . , sgn(m b t (y)d ) ,
(51)
with pt given by Eq. (47) with h = t. Lemma 6.13. For t ∈ I0 , X0 ∼ q C , and its forward observation Xt , the rule in Eq. (51) satisfies P[b zt (Xt ) ̸= X0 | C] ⩽ 5errU C (t) + Cεunif (t) .
(52)
In particular, the law of zbt (Xt ) is within this distance of q C in total variation. Proof. Define a Bayes-optimal decoder (deterministic function) z∗ : for each y, z∗ (y) is a point with largest posterior probability given Xt = y. Let π∗ (y) denote the corresponding posterior mass. The Bayes error is E[1 − π∗ (Xt ) | C] = errU C (t). Therefore by Markov’s inequality, P[π∗ (Xt ) < 3/4 | C] ⩽ 4errU C (t) . On the complementary event, every coordinate of the exact posterior expectation has the sign of z∗ (Xt ) and magnitude at least 1/2. 41
The map from a score to the corresponding posterior-mean coordinate in the first part of Eq. (51) is affine with bounded slope, since I0 is compact. Lemma 6.6 consequently gives EXt ∼qtC ∥m b t (Xt ) − mt (Xt )∥22 ⩽ Cεunif (t) . When π∗ (Xt ) ≥ 3/4, a wrong sign requires the error in at least one coordinate to have magnitude at least 1/2. Markov’s inequality bounds the probability of this event by Cεunif (t). Finally, z∗ (Xt ) ̸= X0 with probability errU C (t). These three possible failures give Eq. (52). The last assertion follows by coupling zbt (Xt ) with X0 . We are now ready to prove the main result of the section. Algorithm 2: Sampling from the random empirical measure with uniform diffusion Input: d, [κ− , κ+ ], target accuracy ε, failure probability δ, and a score oracle satisfying Eq. (25). Output: A sample from the approximate output law qbC . 1 Set Λ as in Eq. (24) and w as in Eq. (26); 2 Run Algorithm 1 to obtain b t; 3 Set t+ ← b t + 4w, t− ← b t − 4w, N ← Cdw , h ← 8w N ; ε2 4 Set rj ← t− + jh, 0 ⩽ j ⩽ N; 5 Sample YN ∼ Unif(X ) independently of Algorithm 1; 6 for j ← N to 1 do 7 Query sbrj (Yj ) and compute serj (Yj ) using Eq. (33); 8 for i ← 1 to d do 9 Set pbj,i ← ph serj (Yj )[i, −(Yj )i ] ; 10 Independently draw Bj,i ∼ Ber(b pj,i ) and set (Yj−1 )i ← (−1)Bj,i (Yj )i ; 11 Query s bt− (Y0 ) and compute zbt− (Y0 ) using Eqs. (33)–(51);
12 return z bt− (Y0 );
Proof of Theorem 6.1. We analyze Algorithm 2. Proposition 6.4, with exponent 12, and Markov’s inequality imply TV(qtC1 , Unif(X )) ⩽ ε/16, errU C (t2 ) ⩽ ε/80 , uniformly over t1 ≥ t⋆unif + w and t2 ⩽ t⋆unif − w, outside an event over C of probability at most δ/2. Proposition 6.10 removes a further event over C of probability at most δ/4 where the search of the critical window fails. Now fix a C outside these failure events. Conditional on this C, the window search succeeds with probability at least 1 − ε/8. Now suppose Algorithm 1 succeeds, and condition on its returned value b t. The successful search ensures that t+ ≥ t⋆unif + w, t− ⩽ t⋆unif − w, t+ − t− = 8w . Both t+ and t− belong to I0 by the choice of w. Now we proceed to analyze the law of Y0 generated by applying approximate reverse kernels to the 8ε2 uniform sample YN in Algorithm 2. Since N ≥ Cdw/ε2 and ε ⩽ 1/4, we have h = 8w N ⩽ Cd ⩽ 1 2Cd . Thus, by increasing d0 if necessary, we have h ⩽ h0 , as required by Proposition 6.12. Therefore,
42
Proposition 6.12 gives r TV(law(Y0 | C, b t), qtC− ) ⩽ ε/16 + C
8dw 3ε + 8we−BΛ ⩽ , N 16
(53)
after choosing B and the constant in N sufficiently large. Finally consider the final recovery step. If the input law were exactly qtC− , Lemma 6.13 and Eq. (52) would bound the distance of its output law from q C by 5ε/80+Ce−BΛ ⩽ ε/8. Combining this with Eq. (53) (after applying the same recovery rule) gives TV(law(b zt− (Y0 ) | C), q C ) ⩽
5ε , 16
uniformly over every possible successful value of b t. Algorithm 1 fails with probability at most ε/8. Thus TV(b q C , q C ) ⩽ 7ε/16 < ε. The total failure probability over C is at most 3δ/4 < δ. 2 Finally, the window search stage uses O(Λ/w) queries, the preverse simulation stage uses N = O(dw/ε + 1), and the final recovery stage uses one. Substituting w = A Λ/d gives the claimed query complexity.
√ e d) upper bound for Gaussian diffusion 7 O( In this √ section, we give an algorithm for approximately sampling from random empirical measures using e d) queries to any sufficiently accurate Gaussian diffusion score oracle for such a distribution. This O( query complexity is also tight: a matching lower bound was proved in [XP26]. Theorem 7.1 (Gaussian diffusion with approximate scores). There are constants B, C, c0 , d0 > 0, depending only on [κ− , κ+ ], such that the following holds. Let d ≥ d0 , let ε, δ ∈ (0, 1/4), and set Λ = log
Cd . εδ
(54)
Suppose Λ ⩽ c0 d. Given query access to any approximate Gaussian diffusion score oracle (b st ) for q C with score error satisfying εδ B −BΛ sup εgauss (t) ⩽ e = , (55) Cd t∈I0 √ there is an algorithm using at most O( dΛ/ε2 ) score oracle queries such that with probability at least 1 − δ over C sampled uniformly from all M -element subsets of X , the output law of the algorithm qbC satisfies TV(b q C , q C ) ⩽ ε. In particular, for constant ε and any polynomially small δ = d−Θ(1) , polynomially small score accuracy √ suffices, and the algorithm uses O( d log d) score queries. For the rest of the section, Λ has the value in Eq. (54). We may decrease c0 and increase d0 later when necessary so that log d ⩽ Λ ⩽ c0 d. Define the reference distribution νt := law(e−t U + σt G),
U ∼ Unif(X ),
43
G ∼ N (0, Id ) .
(56)
Thus νt is obtained by applying the same forward process to the uniform distribution on the whole cube. The algorithm can of course sample from νt without knowing C. In Section 7.1, we quantify the critical window: above the window, qtC is close to the reference distribution νt , while below the window the original sample can be recovered from its noisy observation. In Section 7.2, we show how to locate this window using approximate score oracle queries at points sampled from νt . Finally, in Section 7.3, we describe how to simulate the reverse process across the window and then recover the original point, completing the proof of Theorem 7.1.
7.1
Quantifying the critical window
We begin by comparing the density of an observation Xt started from one point Y ∈ X with the density obtained from a uniform initial point on the whole cube. Normalize the observation by its noise standard deviation: Z = Xt /σt = λt Y + G, λt = e−t /σt . (57) For a general parameter u > 0, write pu for the density of uU + G with U uniform on X . Write ϕd for the standard Gaussian density, so 2
pu (z) = EY ∼Unif(X ) ϕd (z − uY ) = ϕd (z)e−u d/2
d Y
cosh(uzi ) .
(58)
i=1
Consequently, define the log-likelihood ratio of a point y ∈ X given an observation z as d
X ϕd (z − λt y) = λt ⟨z, y⟩ − ıt (y; z) = log log cosh(λt zi ) . pλt (z) i=1
Intuitively, ıt (y; z) measures how y stands out from the uniform distribution given the observation z. Under the joint distribution of Y ∼ Unif(X ) and Z = λt Y + G, this is a sum of d independent terms λt zi yi − log cosh(λt zi ), each with mean Igauss (t). Below we will show that it concentrates around its mean, and that the mean is strictly decreasing with a controlled speed. Lemma 7.2. Uniformly over t ∈ I0 and v ≥ 0, P
Y ∼Unif(X ), Z=λt Y +G
[|ıt (Y ; Z) − dIgauss (t)| ≥ v] ⩽ 2 exp −cv 2 /d .
(59)
Moreover, there are constants 0 < cI < CI such that cI ⩽ −I′gauss (t) ⩽ CI ,
t ∈ I0 .
(60)
Proof. For the concentration bound, write ıt (Y ; Z) =
d X
ξt,i ,
ξt,i = λt Zi Yi − log cosh(λt Zi ) .
i=1
Let Hi = Yi Gi . Then H1 , . . . , Hd are independent standard Gaussians, and since Zi Yi = λt + Hi and cosh is even, ξt,i = ft (Hi ) , ft (h) = λt (λt + h) − log cosh λt (λt + h) . 44
Moreover, ft′ (h) = λt 1 − tanh λt (λt + h) . ′ Since λt is bounded √ for t ∈ I0 , |ft (h)| ⩽ C uniformly in t and h. Hence the function (h1 , . . . , hd ) 7−→ Pd i=1 ft (hi ) is C d-Lipschitz. Gaussian concentration gives Eq. (59). To check the derivative, consider one coordinate. Write V ∼ Unif({−1, 1}), W ∼ N (0, 1) independent, and u = λt . Then
Igauss (t) = I(V ; uV + W ) = u2 − EW log cosh(u2 + uW ) . Differentiate in u and use Gaussian integration by parts, d I(V ; uV + W ) = 2u − 2u E tanh H − u Esech2 H, du
H = u2 + uW .
Since E[V | uV + W ] = tanh(u(uV + W )), E tanh H = E[V E[V | uV + W ]] = E[E[V | uV + W ]2 ] = E tanh2 H . Substituting this identity into the derivative, we get d I(V ; uV + W ) = u EW sech2 (u2 + uW ) . du This derivative is continuous and strictly positive for u > 0. Both λt and |λ′t | are bounded above and away from zero on I0 , and λ′t < 0. The chain rule now gives Eq. (60). For a fixed C, denote the smallest probability of failing to recover Y ∼ q C from λt Y + G by errG C (t): b errG C (t) := inf P [Y (λt Y + G) ̸= Y | C] . Yb Y ∼q C
The following lemma quantifies our intuition for the log-likelihood ratio ıt (Y ; Z). Lemma 7.3. For every t > 0 and γ > 0, EC TV(qtC , νt ) ⩽ P[ıt (Y ; Z) > log M − γ] + 21 e−γ/2 ,
(61)
−γ EC errG , C (t) ⩽ P[ıt (Y ; Z) ⩽ log M + γ] + e
(62)
where both probabilities are taken with respect to Y ∼ Unif(X ) and Z = λt Y + G. Proof. First consider bounding EC TV(qtC , νt ). Fix t, and write Ly (z) = ϕd (z − λt y)/pλt (z). For every z, the average of Ly (z) over the whole cube is 1. We have 1 TV(qtC , νt ) =
2
Z pλt (z)
1 X Ly (z) − 1 dz , M y∈C
P and M −1 y∈C Ly (z) is the empirical average of Ly (z). To control its fluctuation, truncate the likelihood ratio at R = M e−γ and define e y (z) = Ly (z) 1[Ly (z) ⩽ R], L
e c (z) = Ly (z) 1[Ly (z) > R], L y 45
e Y (z) . µ(z) = EY ∼Unif(X ) L
We have the following decomposition: 2EC TV(qtC , νt ) ⩽ EC |
Z
1 Xe pλt (z) Ly (z) − µ(z) dz + EC M y∈C {z } |
Z
1 X e cy (z) dz pλt (z) L M y∈C {z } 2
1
Z + EY ∼Unif(X ) |
pλt (z)LY (z) 1[LY (z) > R]dz . {z } 3
Fix z now. For 1 , by Cauchy–Schwarz and the variance bound for sampling without replacement, 1 Xe EC Ly (z) − µ(z) ⩽ M y∈C
r
1 e Y (z)2 ⩽ E L M Y ∼Unif(X )
r
R = e−γ/2 . M
e Y and µ(z) ⩽ 1. Integrating over pλ gives the same bound. Finally, e 2 ⩽ RL Here we used L t Y Z 2 = 3 = EY ∼Unif(X ) pλt (z)LY (z) 1[LY (z) > R] dz = P[ıt (Y ; Z) > log M − γ] , and Eq. (61) follows. Next we bound the Bayes recovery error. Given C and z, return the unique point y ∈ C with Ly (z) > γ M e , if there is one. This rule succeeds if the initial point passes the test and all other points fail it. The probability that the initial point fails is the first term in Eq. (62). We will bound the probability that any other point passes the test. Fix the initial point y and the observation z, and consider the randomness of C. The remaining M − 1 points are a uniform subset of X \ {y}. Since EY ′ ∼Unif(X ) LY ′ (z) = 1, by Markov’s inequality we have P
Y ′ ∼Unif(X )
[LY ′ (z) > M eγ ] ⩽ M −1 e−γ .
Removing y from the sampling population increases this bound by at most a factor of 2d /(2d − 1) < M/(M − 1). A union bound over the remaining points gives the bound e−γ , and Eq. (62) follows. Setting the parameters appropriately in the two estimates above characterizes the Gaussian critical window: above it the noised empirical measure is close to the reference distribution, while below it the original point can be recovered with high probability. Proposition 7.4 (Critical window). For every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are cK,A , dK,A > 0, depending only on A, K, and [κ− , κ+ ], such that the following holds whenever d ≥ dK,A and log d ⩽ Λ ⩽ cK,A d. Set r w=A
46
Λ . d
(63)
Then EC TV(qtC , νt ) ⩽ e−KΛ ,
t ≥ t⋆gauss + w ,
−KΛ EC errG , C (t) ⩽ e
0 ⩽ t ⩽ t⋆gauss − w .
Proof. By decreasing cK,A if necessary, we may assume w ⩽ min{τ− /2, 1}, so that t⋆gauss ± w ∈ I0 . Fix C. Applying an additional forward step to both distributions cannot increase TV(qtC , νt ). Also, an earlier observation can be used to simulate a later one, so errG C (t) cannot decrease with t. It is therefore enough to prove the claims at t⋆gauss + w and t⋆gauss − w. Above the critical window. Set t = t⋆gauss +w. Since d Igauss (t⋆gauss ) = log M , Eq. (60) gives d Igauss (t) ⩽ log M − cI dw. Take γ = cI dw/2 in Eq. (61). The event ıt (Y ; Z) > log M − γ requires a deviation for 2 ıt (Y ; Z) of at least cI dw/2 above the mean d Igauss (t). Lemma 7.2 bounds its probability by 2e−cdw . The 2 remaining term is at most e−cdw ≲ e−cdw . Below the critical window. Set t = t⋆gauss − w. Now d Igauss (t) ≥ log M + cI dw. With the same choice γ = cI dw/2 in Eq. (62), the event ıt (Y ; Z) ⩽ log M + γ requires a deviation of at least cI dw/2 below the 2 2 mean d Igauss (t). Its probability is again at most 2e−cdw , and the other term is at most e−cdw ≲ e−cdw . Finally, since dw2 = A2 Λ, increasing A gives the desired exponent. The monotonicity noted at the beginning extends the bounds to the stated time ranges. p For the rest of the section, let w = A Λ/d as in Eq. (63), where A > 0 is a constant to be fixed below. We may further decrease c0 and increase d0 later when necessary.
7.2
Identifying the critical window
Having characterized the critical window, we now show how to identify its location using approximate Gaussian score oracle queries at samples from the reference distribution. Define the density ratio of qtC to νt as Lt (x) :=
qtC (x) . νt (x)
(64)
Given x ∈ Rd , define the posterior means under q C and under the uniform distribution on the whole cube by −t e xi 0 mt (x) := E[X0 | Xt = x], mt (x)i := tanh , i ∈ [d] . (65) σt2 For z = x/σt , consider the difference between the two expected overlaps ⟨z, mt (x)⟩ and ⟨z, m0t (x)⟩: Tt (x) = ⟨z, mt (x) − m0t (x)⟩ .
(66)
Above the window, both distributions qtC and νt assign nearly the same density to most observations drawn from νt , and this gap is small. Below the window, a sample from νt is unlikely to be close to any point of C since C is exponentially sparse, so the gap is large. These two behaviors can be used to identify the critical window. We prove these two assertions before considering estimated scores.
47
As in the uniform diffusion proof, it is useful to express both the density ratio and the overlap difference through two log averages. For fixed C, z, define r⟨z,Y ⟩
Az (r) = log EY ∼qC e
r⟨z,Y ⟩
,
Bz (r) = log EY ∼Unif(X ) e
=
d X
log cosh(rzi ) .
(67)
i=1
At u = λt and z = x/σt , Bayes’ rule gives Tt (x) = A′z (u) − Bz′ (u) .
log Lt (x) = Az (u) − Bz (u),
(68)
Thus Az − Bz compares the two densities, while its derivative is exactly the quantity used by our test. Let πr (· | z) be the conditional distribution of U ∼ Unif(X ) given rU + G = z. Then −d r⟨z,y⟩−Bz (r)
πr (y | z) = 2
e
=
d Y i=1
erzi yi . 2 cosh(rzi )
(69)
Define the posterior information Iz (r) := KL(πr (· | z) ∥ Unif(X )) = d log 2 − H(πr (· | z)) = rBz′ (r) − Bz (r) .
(70)
Indeed, Eq. (69) gives log(πr (y | z)/2−d ) = r⟨z, y⟩ − Bz (r), and taking the posterior expectation gives the last identity in Eq. (70). Thus Iz (r) is the information about a uniformly random cube point contained in the observation z. The exponential of the posterior entropy is exactly 2d e−Iz (r) . When r = λt and Z ∼ pλt , taking expectation of IZ (λt ) over Z gives the mutual information d Igauss (t). At the critical time t⋆gauss this equals log M , so the entropy scale equals 2d /M , the reciprocal of the fraction of the cube contained in C. The derivatives of these quantities also have direct meanings. Under πr (· | z), the mean and variance of the overlap ⟨z, U ⟩ are Bz′ (r) and Bz′′ (r), and hence Iz′ (r) = rBz′′ (r) = r
V
U ∼πr (·|z)
[⟨z, U ⟩] .
In particular, the posterior information increases with the signal strength r. We need two elementary facts about the reference law. The first ensures enough overlap between pu and pv if u and v are close. Lemma 7.5. For u, v > 0 and every measurable set A ⊆ Rd , p 2 pu (A) ⩽ ed(u−v) /2 pv (A) .
(71)
More generally, let C be a random subset independent of (U, G), with the same law under parameters u and v. For every measurable event E determined by (C, Z), r d(u−v)2 /2 P [E] ⩽ e P [E] . (72) C,Z∼pu
C,Z∼pv
We defer the proof to Appendix B. The second fact says that, for a typical sample from pλt , the posterior information is close to its mean and changes at a speed of order d. 48
Lemma 7.6. For every K > 0, there are constants CK , cK , dK > 0, depending only on [κ− , κ+ ] and K, such that the following holds whenever d ≥ dK , log d ⩽ Λ ⩽ cK d, and t ∈ I0 . With probability at least 1 − e−KΛ for Z ∼ pλt , ∥Z∥22 ⩽ Cd,
(73)
cd ⩽ BZ′′ (r) ⩽ Cd,
(74)
√
|IZ (λt ) − dIgauss (t)| ⩽ CK dΛ .
(75)
The bounds on BZ′′ hold uniformly for r ∈ [λτ+ +1 /2, 2λτ− /2 ], and c, C depend only on [κ− , κ+ ]. We defer the proof to Appendix B. We can now use the above ingredients to prove that the statistic Tt can be used to distinguish whether one is above or below the critical window: Proposition 7.7. There are constants a0 , b0 > 0, depending only on [κ− , κ+ ], with b0 > 4a0 , such that for every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are cK,A , dK,A > 0, depending only on A, K, and [κ− , κ+ ], such that the following holds whenever d ≥ dK,A and log d ⩽ Λ ⩽ cK,A d: P
[Tt (X) < −a0 dw] ⩽ e−KΛ ,
t ≥ t⋆gauss + 2w,
(76)
P
[Tt (X) > −b0 dw] ⩽ e−KΛ ,
t ⩽ t⋆gauss − 2w ,
(77)
C,X∼νt
C,X∼νt
uniformly over t ∈ I0 . Proof. Write u = λt and z = X/σt . High-noise regime (t ≥ t⋆gauss + 2w). Choose a small constant θ > 0, set t+ = t + θw, and write v = λt+ . Since −λ′t is bounded above and below on the relevant compact interval, u − v ≍ θw. In particular, t and t+ are both above t⋆gauss + w. By Eq. (68), define the two density ratios, as functions of the common normalized observation z, by Lu (z) = eAz (u)−Bz (u) , Lv (z) = eAz (v)−Bz (v) . Proposition 7.4 and Markov’s inequality give P
C,Z∼pu
[|Lu (Z) − 1| > 1/2] +
2
P
C,Z∼pv
[|Lv (Z) − 1| > 1/2] ⩽ Ce−cdw .
Lemma 7.5 transfers the bound forLv from pv to pu . Since d(u − v)2 ⩽ Cθ2 dw2 , we may choose θ small X enough that, under C ∼ Unif M and Z ∼ pu , both Lu (Z) and Lv (Z) belong to [1/2, 3/2] except with 2 −cdw probability Ce . Lemma 7.6 gives BZ′′ (r) ⩽ Cd uniformly for r ∈ [v, u] outside another event of probability e−KΛ , for sufficiently large d.
49
Fix C, z for which these three conclusions hold. Since Az is convex, Tt (x) = A′z (u) − Bz′ (u) Az (u) − Az (v) − Bz′ (u) u−v log Lu (z) − log Lv (z) Bz (u) − Bz (v) = + − Bz′ (u) u−v u−v C log 3 ≥− − Cdθw . θw ≥
The last step uses u − v ≍ θw, Eq. (74), and Taylor’s formula. For any prescribed a0 > 0, further choose θ small enough so that the last term is at least −a0 dw/2, and then choose A large enough so that the first term is at least −a0 dw/2, using dw2 = A2 Λ. Increasing A also makes the total exceptional probability at most e−KΛ . Low-noise regime (t ⩽ t⋆gauss − 2w). Apply Lemma 7.6 with exponent K + 1 to Z ∼ pu , and fix an observation z satisfying Eq. (75). Since Igauss (t) − Igauss (t⋆gauss ) ≥ 2cI w and dIgauss (t⋆gauss ) = log M , choosing A sufficiently large gives Iz (u) ≥ log M + cI dw . (78) We show directly that C is unlikely to contain even one point y whose overlap ⟨z, y⟩ reaches the reference posterior mean Bz′ (u) up to a small gap b0 dw. Choose b0 > 0 so that ub0 ⩽ cI /2 uniformly in u. By Markov’s inequality and Eq. (70), ′
P
Y ∼Unif(X )
[⟨z, Y ⟩ ≥ Bz′ (u) − b0 dw] ⩽ e−u(Bz (u)−b0 dw) EY ∼Unif(X ) eu⟨z,Y ⟩ = exp{−Iz (u) + ub0 dw} ⩽ M −1 e−cI dw/2 .
By a union bound over the M points in C, except with probability e−cI dw/2 , every y ∈ C has overlap less than Bz′ (u) − b0 dw. Their posterior-weighted average is A′z (u), so on this event Tt (x) = A′z (u) − Bz′ (u) ⩽ −b0 dw . Adding the probability e−(K+1)Λ that Eq. (78) fails, and choosing A large enough that e−cI dw/2 ⩽ e−(K+1)Λ , which is possible since dw ≥ dw2 = A2 Λ, we obtain Eq. (77). Finally, choose a0 in the first part small enough that b0 > 4a0 . We now replace the posterior mean by the estimate obtained from the score. Define its clipped version coordinatewise by m e t (x)i := min 1, max −1, et xi + et σt2 sbt (x)i
,
i ∈ [d] .
(79)
Using this vector, define the clipped approximate score and the statistic available to the algorithm by set (x) :=
e−t m e t (x) − x , σt2
Tbt (x) := ⟨x/σt , m e t (x) − m0t (x)⟩ .
Both quantities can be computed from one approximate score oracle query at x. 50
(80)
The following lemma shows that the estimated overlap difference Tbt is indeed a good approximation of the true overlap difference Tt . Lemma 7.8. For every fixed C and t ∈ I0 , EqtC ∥m e t − mt ∥22 ⩽ Cεgauss (t)2 ,
EqtC ∥e st − st ∥22 ⩽ Cεgauss (t)2 .
(81)
√ Moreover, whenever ∥x/σt ∥2 ⩽ C1 d, √ |Tbt (x) − Tt (x)| ⩽ C1 d ∥m e t (x) − mt (x)∥2 .
(82)
We defer the proof to Appendix B. The score assumption controls errors on observations from qtC , while our queries use νt . Above the window, total variation already compares the two distributions. Near the window, we need the following weaker assertion: a sample from νt is unlikely to have extremely small density under qtC . Lemma 7.9. For every A, L, K > 0, there are constants CA,L,K , cA,L,K , dA,L,K > 0, depending only on A, L, K and [κ− , κ+ ], such that whenever d ≥ dA,L,K and log d ⩽ Λ ⩽ cA,L,K d, P
C,X∼νt
[Lt (X) < e−CA,L,K Λ ] ⩽ e−KΛ ,
(83)
uniformly over |t − t⋆gauss | ⩽ Lw. We defer the proof to Appendix B. The above lemmas allow us to use the estimated overlap difference Tbt to identify the critical window. Proposition 7.10. There are constants 0 < a < b, depending only on [κ− , κ+ ], such that for every K > 0 there is AK > 0 with the following property. For every A ≥ AK , there are BK,A , cK,A , dK,A > 0 such that the following holds whenever d ≥ dK,A , log d ⩽ Λ ⩽ cK,A d, B ≥ BK,A , and the score oracle satisfies Eq. (55): P
[Tbt (X) < −adw] ⩽ e−KΛ ,
t ≥ t⋆gauss + 2w,
P
[Tbt (X) > −bdw] ⩽ e−KΛ ,
t⋆gauss − 3w ⩽ t ⩽ t⋆gauss − 2w ,
C,X∼νt
C,X∼νt
uniformly over t ∈ I0 . Proof. Take a0 , b0 from Proposition 7.7, and put a = 2a0 and b = b0 − a0 . These satisfy a < b. We will show that replacing Tt by Tbt changes it by at most a0 dw with sufficiently high probability in each of the two ranges. Fix C and t, and let c > 0 be a small constant to be fixed below. By Lemma 7.8 and Markov’s inequality, √ Ce−2BΛ P [∥m e t (X) − mt (X)∥2 > c d w] ⩽ . dw2 X∼qtC
(84)
Choose c small enough event does not occur √ that Eq. (82) bounds the overlap error by a0 dw whenever this −(K+3)Λ and ∥X/σt ∥2 ⩽ C d. The latter condition fails under νt with probability at most e , by Lemma 7.6.
51
High-noise regime (t ≥ t⋆gauss + 2w). Replacing qtC by νt in Eq. (84) adds at most TV(qtC , νt ). Average √ over C and apply Proposition 7.4 with exponent K + 3. Together with the ∥X/σt ∥2 ⩽ C d probability bound, this gives Ce−2BΛ P [|Tbt (X) − Tt (X)| > a0 dw] ⩽ + 2e−(K+3)Λ . C,X∼νt dw2 Increasing B makes this at most e−(K+1)Λ . Outside this event and the event in Eq. (76), we have Tbt ≥ −2a0 dw = −adw. Low-noise regime (t⋆gauss − 3w ⩽ t ⩽ t⋆gauss − 2w). We apply Lemma 7.9 with L = 3 and exponent K + 3, and denote its density exponent by C1 . On the event Lt (X) = qtC (X)/νt (X) ≥ e−C1 Λ , integration under νt is bounded by eC1 Λ times integration under qtC . Then following the same steps as in the high-noise regime, we obtain −(2B−C1 )Λ
P
C,X∼νt
Ce [|Tbt (X) − Tt (X)| > a0 dw] ⩽
dw2
+ 2e−(K+3)Λ .
Choose B after C1 so that this is at most e−(K+1)Λ . Outside this event and the event in Eq. (77), we have Tbt ⩽ −(b0 − a0 )dw = −bdw. Finally, applying Proposition 7.7 with exponent K + 2 and adding the error probabilities proves both assertions. Here AK is the larger of the thresholds required by Proposition 7.4 at exponent K + 3 and by Proposition 7.7 at exponent K + 2. For every fixed A ≥ AK , apply Lemma 7.9 with L = 3 and exponent K + 3, choose BK,A sufficiently large in terms of its density exponent, and then choose cK,A sufficiently small and dK,A sufficiently large so that all three results apply. From now on, fix A ≥ A12 and choose B ≥ B12,A . We may decrease c0 and increase d0 so that c0 ⩽ c12,A , d0 ≥ d12,A , and w ⩽ min τ− /20, 1/20. We further decrease c0 and increase d0 below whenever necessary. The two bounds tell us how to search. We test decreasing times using samples from the reference distribution at each time. The test should not stop above t⋆gauss + 2w, and it should stop by the time it reaches [t⋆gauss − 3w, t⋆gauss − 2w]. Its behavior elsewhere is irrelevant. Algorithm 3: Finding the Gaussian diffusion critical window Input: d, [κ− , κ+ ], w, Λ, and query access to the score oracle sb. Output: An estimate b t of the critical time t⋆gauss . 1 Let T = {t1 > t2 > · · · > tm } be a decreasing grid covering [τ− − 3w, τ+ + 3w] with mesh w/20; 2 Set R ← 2⌈CΛ⌉ + 1; 3 for t ∈ T do 4 for j ← 1 to R do 5 Draw independent U (j) ∼ Unif(X ) and G(j) ∼ N (0, Id ); 6 Set X (j) ← e−t U (j) + σt G(j) ; 7 Query sbt (X (j) ) and compute Tbt (X (j) ) using Eqs. (79)–(80); 8 if median Tbt (X (1) ), . . . , Tbt (X (R) ) < −(a + b)dw/2 then 9 return t; 10 return τ− ;
52
We can now conclude the main result of this subsection: scanning from high to low noise and stopping when the median statistic becomes sufficiently negative locates the critical window. Proposition 7.11. Under the hypotheses of Theorem 7.1, Algorithm 3 uses O(Λ/w) score oracle queries. For at least a 1 − δ/4 fraction of C, P[t⋆gauss − 3w ⩽ b t ⩽ t⋆gauss + 2w | C] ≥ 1 − ε/8 .
(85)
We defer the proof to Appendix B.
7.3
Sampling after locating the window
Once the search has returned b t, we will start reverse simulation at t+ = b t + 4w and stop at t− = b t − 4w. At each step, we hold the estimated score fixed and solve the remaining linear SDE exactly. If a step goes from time r to time r − h, its update is p G ∼ N (0, Id ) . (86) Ynext = eh Y + 2(eh − 1)e sr (Y ) + e2h − 1 G, We need to bound the error from holding the score fixed for time h. A short interval alone is not enough: near the critical time, the posterior can change rapidly. The useful fact is that the total increase of its variance is at most d. We will charge the error on each step to the increase during that step and then sum these increases. For a fixed C, write Σt (x) = Cov(X0 | Xt = x),
V (t) = EXt ∼qtC Tr Σt (Xt ) .
(87)
Thus V (t) is the mean squared error of the posterior mean. The following identity makes the preceding plan precise. Lemma 7.12. Fix C. For t > 0, V ′ (t) = 2e−2t σt−4 EqtC Tr(Σ2t ) ≥ 0,
0 ⩽ V (t) ⩽ d .
(88)
We defer the proof to Appendix B. We next bound the change of the score along one exact reverse step. Stating this separately keeps the time-discretization calculation apart from the score-estimation calculation. Lemma 7.13. Fix C and r − h, r ∈ I0 , where 0 < h ⩽ 1. Let (Yu )0⩽u⩽h follow the exact reverse process C . Then from qrC to qr−h Z h
E∥sr−u (Yu ) − sr (Y0 )∥22 du ⩽ C dh2 + h(V (r) − V (r − h)) .
(89)
0
We defer the proof to Appendix B. We now combine the preceding score-change estimate with Girsanov’s theorem on each grid interval, and then sum the errors over the grid. Proposition 7.14. Fix C and a grid t− = r0 < r1 < · · · < rN = t+ in I0 . Put hj = rj − rj−1 , and suppose maxj hj ⩽ h0 , for a sufficiently small constant h0 > 0 depending only on I0 . Initialize Eq. (86) from qtC+ ,
53
and let qbtC− be its endpoint law. Then KL(qtC− ∥ qbtC− ) ⩽ C
N X j=1
dh2j + hj (V (rj ) − V (rj−1 )) + hj EqrC ∥e srj − srj ∥22 j
.
In particular, on an equal grid with ∆ = t+ − t− , N 2 X d∆ ∆ ∆ KL(qtC− ∥ qbtC− ) ⩽ C + (V (t+ ) − V (t− )) + E C ∥e sr − srj ∥22 . N N N qrj j
(90)
(91)
j=1
Proof. Fix C throughout. We compare the conditional path laws on each interval before summing over the grid. One grid interval. Fix a grid interval [r − h, r] and condition on Y0 = y. Let Pyr,h denote the path law of the exact reverse process √ Y0 = y , dYu = Yu + 2sr−u (Yu ) du + 2 dBu , e y denote the path law of and let P r,h √ dYeu = Yeu + 2e sr (y) du + 2 dBu ,
Ye0 = y .
The time-h endpoint of the latter process is exactly the update in Eq. (86). By Girsanov’s theorem, we have ey ) = 1 E y KL(Pyr,h ∥ P r,h 4 Pr,h
Z h 0
∥2sr−u (Yu ) − 2e sr (y)∥22 du =
Z h h i E ∥sr−u (Yu ) − ser (y)∥22 Y0 = y du . 0
Averaging over Y0 ∼ qrC therefore gives 0 e Y0 ) = EY0 ∼qrC KL(PYr,h ∥P r,h
The change-of-measure condition.
Z h 0
E ∥sr−u (Yu ) − ser (Y0 )∥22 du .
(92)
We briefly verify Novikov’s condition Z h
E exp 0
∥sr−u (Yu ) − ser (Y0 )∥22 du
<∞
used above in Girsanov’s theorem. By Tweedie’s formula and the clipping in Eqs. (79)–(80), √ ∥st (x)∥2 + ∥e st (x)∥2 ⩽ C(∥x∥2 + d), t ∈ I0 . Hence, along the exact reverse process, ∥sr−u (Yu ) − ser (Y0 )∥22 ⩽ C ∥Yu ∥22 + ∥Y0 ∥22 + d .
(93)
C , and every q C , t ∈ I , is a mixture of Gaussians N (e−t y, σ 2 I ), y ∈ C. Since Under this process Yu ∼ qr−u 0 t t d √ their variances are uniformly bounded and ∥e−t y∥2 ⩽ d, there is a constant c > 0, depending only on I0 ,
54
such that 2
sup EX∼qtC ec∥X∥2 < ∞ .
(94)
C,t∈I0
Indeed, this follows from the quadratic exponential moment of a Gaussian, which is finite for c < 2 sup 1
2 t∈I0 σt
after decreasing c if necessary. Moreover, Jensen’s inequality gives Z h Z 1 h Ch∥Yu ∥22 2 e du . ∥Yu ∥2 du ⩽ exp C h 0 0 Combining this with Eq. (94) and Cauchy–Schwarz, we get Z h E exp C ∥Yu ∥22 + ∥Y0 ∥22 + d du < ∞ 0
whenever h ⩽ h0 for a sufficiently small constant h0 > 0 depending only on I0 . By Eq. (93), Novikov’s condition follows. Bounding the one-step error. We now return to Eq. (92). Adding and subtracting sr (Y0 ), we separate the error caused by freezing the exact score from the error in the score estimate: Z h 0
E ∥sr−u (Yu ) − ser (Y0 )∥22 du ⩽ 2
Z h 0
E ∥sr−u (Yu ) − sr (Y0 )∥22 du + 2h EqrC ∥sr − ser ∥22 .
By Lemma 7.13, Z h 0
E ∥sr−u (Yu ) − sr (Y0 )∥22 du ⩽ C dh2 + h(V (r) − V (r − h)) .
Therefore we get the following bound on the KL divergence between the exact and approximate path laws on a short interval: 0 e Y0 ) ⩽ C dh2 + h(V (r) − V (r − h)) + 2h E C ∥sr − ser ∥2 . EY0 ∼qrC KL(PYr,h ∥P qr 2 r,h e path denote the path laws obtained by concatenating, Summing over the whole grid. Let Ppath and P respectively, the exact and approximate reverse processes over the grid intervals. The chain rule implies e path ) = KL(Ppath ∥ P
N X
Yr j e Yr j ) . EPpath KL(Prj ,h ∥P rj ,hj j
j=1
Under the exact reverse law, Yrj ∼ qrCj . Applying the preceding one-step bound on each interval therefore yields path
KL(P
e path ) ⩽ C ∥P
N N X X 2 hj EqrC dhj + hj V (rj ) − V (rj−1 ) + 2 j=1
j=1
55
j
2
srj − serj 2 .
The endpoint is a measurable function of the full path, so the data processing inequality gives the same upper bound for KL(qtC− ∥ qbtC− ). This proves Eq. (90). For an equal grid, hj = ∆/N with ∆ = t+ − t− . Hence N X
dh2j =
j=1
d∆2 , N
N X j=1
∆ V (t+ ) − V (t− ) . hj V (rj ) − V (rj−1 ) = N
Substituting these identities into the preceding bound gives Eq. (91). At the lower endpoint, we return the coordinatewise signs of m e t (x). This works for a reason stronger than a coordinatewise error estimate: most approximate samples below the window place almost all posterior mass on a single point of C. Lemma 7.15. Fix C and define Dt (x) := sgn(m e t (x)). If X0 ∼ q C and Xt is obtained from X0 through the Gaussian diffusion process, then P[Dt (Xt ) ̸= X0 | C] ⩽ 5errG e t (X) − mt (X)∥2 > 1/4] . C (t) + P [∥m X∼qtC
Proof. Fix Xt = x, let y∗ (x) ∈ C maximize the posterior probability P[X0 = y | Xt = x], and call this probability p∗ (x). If p∗ (x) ≥ 3/4, every coordinate of mt (x) has the sign of y∗ (x) and magnitude at least 1/2. An error at most 1/4 in the Euclidean norm of the estimated posterior mean preserves all these signs. On these two events, Dt (x) = y∗ (x). Averaging over Xt , we get E[1 − p∗ (Xt ) | C] = errG C (t). Markov’s inequality bounds the probability of G p∗ (Xt ) < 3/4 by 4errC (t). Add the probability of an error greater than 1/4 in the posterior mean estimate to obtain the claim. Algorithm 4: Sampling from the random empirical measure with Gaussian diffusion Input: d, [κ− , κ+ ], target accuracy ε, failure probability δ, and a score oracle satisfying Eq. (55). Output: A sample from the approximate output law qbC . 1 Set Λ as in Eq. (54) and w as in Eq. (63); 2 Run Algorithm 3 to obtain b t; 3 Set t+ ← b t + 4w, t− ← b t − 4w, and ∆ ← 8w; 2 4 Set N ← ⌈C(d∆ + d∆ )/ε2 ⌉ and rj ← t− + j∆/N for 0 ⩽ j ⩽ N ; 5 Draw YN ∼ νt+ independently of Algorithm 3; 6 for j = N, N − 1, . . . , 1 do 7 Query sbrj (Yj ) and compute serj (Yj ) using Eqs. (79)–(80); 8 Draw Gj ∼ N (0, Id ) independently and update Yj−1 by Eq. (86) with h = ∆/N ; 9 Query s bt− (Y0 ) and compute m e t− (Y0 ) using Eq. (79);
10 return sgn(m e t− (Y0 ));
Proof of Theorem 7.1. We analyze Algorithm 4, allowing for errors in the search, the reverse steps, and the final rounding. We first apply Proposition 7.4 at the deterministic times t⋆gauss + w and t⋆gauss − w. Markov’s inequality, with exponent 12 in that proposition, removes at most a δ/2 fraction of subsets and gives that for the
56
remaining fraction of C, TV(qtC , νt ) ⩽ ε/32,
t ≥ t⋆gauss + w,
errG C (t) ⩽ ε/160,
t ⩽ t⋆gauss − w .
(95)
Monotonicity extends the two bounds from the fixed two times to the two full ranges. The constant ε/160 is chosen so that 5errG C (t) ⩽ ε/32 in Lemma 7.15 below. Proposition 7.11 further removes at most another δ/4 fraction of C where the window search might fail. Fix a C satisfying all these properties for the remainder of the proof. With conditional probability at least 1 − ε/8, the critical window search returns b t ∈ [t⋆gauss − 3w, t⋆gauss + ⋆ ⋆ 2w]. Now, condition on any such value b t. Then t+ ≥ tgauss + w, t− ⩽ tgauss − w, and t+ − t− = ∆ = 8w. The choices of c0 , d0 ensure that the endpoints lie in I0 , even if the search fails. All samples used after the search are fresh. First initialize the reverse steps from qtC+ . Lemma 7.12 gives V (t+ ) − V (t− ) ⩽ d, and Lemma 7.8 bounds the total score term by C∆e−2BΛ . Proposition 7.14 therefore gives 2 d∆ + d∆ C C −2BΛ KL(qt− ∥ qbt− ) ⩽ C + ∆e ⩽ ε2 /128 . N The last inequality follows by choosing the constant in N and then B sufficiently large; the same choice ensures ∆/N ⩽ h0 . Pinsker’s inequality bounds TV(qtC− , qbtC− ) by ε/16. The actual algorithm starts from νt+ rather than qtC+ . Since both initial distributions are propagated by the same approximate reverse kernels, we have TV qbtC− , law(Y0 | C) ⩽ TV(qtC+ , νt+ ) ⩽ ε/32 . Combining this with the preceding endpoint bound, with triangle inequality, gives TV qtC− , law(Y0 | C) ⩽ ε/16 + ε/32 = 3ε/32 . If the final observation had law qtC− , Lemma 7.8 would give P [∥m e t− − mt− ∥2 > 1/4] ⩽ Ce−2BΛ ⩽ ε/32 .
qtC
−
C By Lemma 7.15 and Eq. (95), its rounded output law would then be within ε/16 of q in total variation. Rounding the actual endpoint adds at most TV qtC− , law(Y0 | C) of total variation error. Thus, conditional on a successful window search, the total error is at most 3ε/32+ε/16 = 5ε/32. Averaging over the window search adds at most its failure probability ε/8, using TV ⩽ 1 on the failure event, so the output law is within 5ε/32+ε/8 = 9ε/32 ⩽ ε of q C in total variation. The algorithm fails with probability at most δ/2+δ/4 < δ over the choice of C, completing the proof of correctness. √ 2 + 1 = The search uses O(Λ/w) = O( dΛ) queries. The reverse steps use N = O dw+dw 2 ε √ dΛ+Λ O + 1 queries, and rounding uses one more query, completing the proof. ε2
57
e 8 Ω(d) lower bound for masked diffusion We now consider the same random empirical measure under masked diffusion and exhibit an approximate e score oracle for which any sampler requires Ω(d) queries, in contrast to uniform and Gaussian diffusion. The intuition for this lower bound, which we state formally in Theorem 8.1 below, is as follows. If the codebook defining the empirical measure has size |C| = 2k , and a total of m coordinates have been revealed, a fixed partial assignment agrees with about 2k−m points of C. When m is well below k, these points have nearly balanced signs in the remaining coordinates. When m is well above k, a partial assignment chosen without knowledge of C is unlikely to agree with any point at all. We will use these two facts to construct an accurate masked diffusion oracle that usually returns uniform marginals in both ranges. Only queries at values of m close to the unknown value k need reveal information about the target. Let KdM := {k ∈ Z : κ− ⩽ k log 2/d ⩽ κ+ }, and consider the prior under which k ∼ Unif(KdM ) and, given k, C is a uniformly random 2k -element subset of X . For a partial assignment (I, xI ), write matchI (xI ) := {y ∈ C : yI = xI } for the set of codewords matching it. Then (I, xI ) is consistent with q C if and only if matchI (xI ) ̸= ∅, in which case |{y ∈ matchI (xI ) : yi = a}| , |matchI (xI )|
C qi|I (a | xI ) =
i∈ / I.
We measure the accuracy of an oracle O by its level-m score error εmask (O; q C , m) from Definition 2.7, in the form given by Lemma 2.8. Recall that this constrains O only on consistent partial assignments; on inconsistent partial assignments, the oracle can return arbitrary values. Theorem 8.1 (Masked diffusion lower bound). Fix ρ ∈ (0, 1/4). There are constants c, C, d0 > 0, depending only on [κ− , κ+ ] and ρ, such that the following holds for d ≥ d0 . Given η ∈ (0, 1/10), let Cd ℓ = log2 2 . (96) ηρ There is a family of oracles {Ok,C }k∈KM ,|C|=2k satisfying d
max εmask (Ok,C ; q C , m) ⩽ η,
k ∈ KdM , |C| = 2k ,
0⩽m<d
(97)
such that every adaptive randomized algorithm A making at most Q ⩽ cρ2
d ℓ O
queries, and which does not have knowledge of k, C, has output law qbA k,C satisfying O
[TV(b qA k,C , q C ) ≥ 1 − ρ] ≥ 1 − ρ .
P
k∼Unif(KdM ) {±1}d
C∼Unif
( 2k )
In particular, for constant ρ and any polynomially small η = d−Θ(1) , sampling with at most constant error for at least a constant fraction of the codebooks requires Ω(d/ log d) queries. In Section 8.1 we describe the oracle construction. In Section 8.2, we prove the lower bound.
58
8.1
Adversarial oracle construction
Define the oracle O0 to return Unif({±1}) in every unrevealed coordinate. We first check its average error when few coordinates are revealed. Lemma 8.2. For every k ∈ KdM and 0 ⩽ m < d, EC εmask (O0 ; q C , m) ⩽ 2m−k . Proof. Fix I with |I| = m and i ∈ / I. For a fixed assignment xI , define NIC (xI ) = |matchI (xI )|. CondiC C tional on NI (xI ) > 0, these NI (xI ) points form a uniformly drawn subset of the subcube {y : yI = xI }. C (x ) := |{y ∈ match (x ) : y = +1}|/N C (x ) as Exactly half of this subcube has yi = +1. Define Pi|I i I I I I I C the fraction of points in matchI (xI ) with yi = +1, so when NI (xI ) > 0, 1 C EC [Pi|I (xI ) | NIC (xI )] = , 2
h EC
C (xI ) − Pi|I
i 1 2 1 | NIC (xI ) ⩽ . C 2 4NI (xI )
The variance bound follows since we are sampling matchI (xI ) without replacement. The inequality log u ⩽ u − 1 also gives KL(Ber(p) ∥ Ber(1/2)) ⩽ 4(p − 1/2)2 . Thus for any NIC (xI ) > 0, h i 1 2 C C EC NIC (xI ) KL(Ber(Pi|I (xI )) ∥ Ber(1/2)) | NIC (xI ) ⩽ 4NIC (xI ) EC Pi|I (xI ) − | NIC (xI ) ⩽ 1 . 2 C (x )) ∥ Ber(1/2)) = 0. When NIC (xI ) = 0, define NIC (xI ) KL(Ber(Pi|I I C Now fix C. A sample X0 ∼ q satisfies X0,I = xI with probability NIC (xI )/2k . Hence the expected KL error is C EX0 ∼qC KL(qi|I (· | X0,I ) ∥ Ber(1/2)) =
1 2k
X
C NIC (xI ) KL(Ber(Pi|I (xI )) ∥ Ber(1/2)) .
xI ∈{±1}I
Now average over C. Each of the 2m summands has expectation at most one by the preceding calculation, so the result is at most 2m−k . Averaging over I and i proves the lemma. Call C good if the answers provided by O0 satisfy the desired error bound at every level m ⩽ k − ℓ: εmask (O0 ; q C , m) ⩽ η
for every 0 ⩽ m ⩽ k − ℓ .
(98)
Lemma 8.2, Markov’s inequality, and a union bound over the levels 0 ⩽ m ⩽ k − ℓ imply k−ℓ
P[C is not good] ⩽ C
1 X m−k 21−ℓ 2 ⩽ . η η
(99)
m=0
Definition 8.3 (Adversarial masked diffusion oracle). Define the oracle Ok,C as follows: • For a good C, it returns uniform marginals at all levels m ⩽ k − ℓ. At levels m > k − ℓ, it returns the exact conditional marginals on consistent assignments and uniform marginals on inconsistent ones.
59
• For C that is not good, at all levels m, it returns the exact conditional marginals on every consistent assignment and uniform marginals on inconsistent ones. Proposition 8.4 (Validity of the oracle). For every k ∈ KdM and |C| = 2k , the oracle Ok,C satisfies max0⩽m<d εmask (Ok,C ; q C , m) ⩽ η, i.e., Eq. (97). Proof. By Lemma 2.8, εmask (Ok,C ; q C , m) depends only on the answers of Ok,C on consistent assignments with m revealed coordinates. These answers are exact, so the error at level m vanishes, except when C is good and m ⩽ k − ℓ, in which case Ok,C agrees with O0 and the error is at most η by Eq. (98).
8.2
Distribution indistinguishability: the lower bound proof
Let R denote all the internal randomness of an algorithm A. The following compares the algorithm’s interaction with the constructed oracle to its interaction with O0 . For the latter, every answer is uniform, so the entire query sequence is fully determined by the internal randomness of the algorithm R. Proposition 8.5 (Coupling argument). Couple an algorithm’s run against the oracle Ok,C with its run against the oracle O0 , using the same internal randomness R. For an algorithm making at most Q queries, let pdiff (k, C) be the conditional probability over R given k, C that the runs receive different oracle answers at some point. Then (2ℓ + 1)Q 21−ℓ Ek,C pdiff (k, C) ⩽ + + Q2−ℓ . (100) η |KdM | Proof. Let D denote the event that the two coupled runs receive different oracle answers at some point. By definition, Ek,C pdiff (k, C) = Ek,C P[D | k, C] = P [D] = ER P [D | R] . (101) R
k,C,R
k,C
Now condition on any fixed realization of the internal randomness R = r. We will directly bound Pk,C [D | R = r] uniformly over r. It is enough to bound the probability of a first different answer, because the queries agree up to that point. Run the algorithm against oracle O0 . This fixes queries (I1 , xI1 ), . . . , (IQ , xIQ ), padding with irrelevant queries if the algorithm stops early. The sequence is independent of both k and C. Consider the randomness of k. For any query (Ij , xIj ), mj := |Ij | is within ℓ of at most 2ℓ + 1 possible integers k. A union bound shows that there exists a query with |mj − k| ⩽ ℓ with probability at most (2ℓ + 1)Q/|KdM | over k. Fix a k outside all these intervals {[mj − ℓ, mj + ℓ]}Q j=1 . We also remove the event that C is not good, whose probability is bounded by Eq. (99). High-noise regime (m < k − ℓ). At any query (I, xI ) with m < k − ℓ, Definition 8.3 returns uniform marginals when C is good, regardless of whether the assignment is consistent. Its answer agrees with O0 . Low-noise regime (m > k + ℓ). Fix one query (I, xI ) with m > k + ℓ. The subcube {y : yI = xI } specified by this assignment has 2d−m points, so P[matchI (xI ) ̸= ∅] ⩽ EC |matchI (xI )| = 2k−m ⩽ 2−ℓ . C
This probability is under the original uniform choice of C (we do not condition on C being good). A union bound over the at most Q queries with mj > k + ℓ therefore shows that the probability that any such queried assignment is consistent is at most Q2−ℓ . 60
Thus, for the fixed realization of R = r, if no query size mj lies in [k − ℓ, k + ℓ], if C is good, and if every queried assignment with mj > k + ℓ is inconsistent with C, then every query in the run against Ok,C receives exactly the same uniform answer as in the run against O0 . Indeed, before a first different answer the two runs have identical queries and answers, and hence make the same next query. By induction over the queries, no first different answer can occur under these three conditions. Consequently, uniformly over R = r, P [D | R = r] ⩽
k,C
(2ℓ + 1)Q 21−ℓ + + Q2−ℓ . η |KdM |
(102)
Averaging Eq. (102) over R and using Eq. (101) gives Eq. (100). We can now complete the proof of the main result of this section: Proof of Theorem 8.1. By Proposition 8.4, the oracle Ok,C satisfies the accuracy requirement for every instance. Fix any algorithm with query budget Q, and let ν be its output law when every oracle answer is uniform. This law is independent of k, C. For each fixed instance (k, C), couple the two runs of the algorithm, one against Ok,C and the other against O0 , by using the same internal randomness R, as in Proposition 8.5. Let Yk,C (R) and Y0 (R) denote O the corresponding outputs. Their marginal laws over R are respectively qbA k,C and ν. Hence, by the coupling inequality, O
TV(b qA k,C , ν) ⩽ P[Yk,C (R) ̸= Y0 (R)] R
⩽ P[the two runs receive different oracle answers] R
= pdiff (k, C) . The second inequality holds because, for every fixed realization of R, identical oracle answers yield the same final output. Since q C is supported on C, it follows that O
O
O
qA k,C , ν) ≥ 1 − ν(C) − pdiff (k, C) . TV(b qA k,C , q C ) ≥ 1 − qbA k,C (C) ≥ 1 − ν(C) − TV(b
(103)
A fixed point belongs to a uniform 2k -element subset with probability 2k−d . Averaging according to any output law ν independent of C thus gives EC ν(C) = EX∼ν EC 1[X ∈ C] = 2k−d ⩽ e−(log 2−κ+ )d . We have |KdM | ≥ c1 d for some constant c1 > 0 depending only on [κ− , κ+ ], for all sufficiently large d. Choose the constant c in the query budget small enough that the first term of Eq. (100) is at most ρ2 /4. Also Q ⩽ d for this choice. Taking C in Eq. (96) sufficiently large bounds the second and third terms in Eq. (100) by ρ2 /4 in total. For sufficiently large d, the expected mass EC ν(C) is at most ρ2 /4. Hence, Ek,C [ν(C) + pdiff (k, C)] ⩽ ρ2 . Markov’s inequality shows that Pk,C [ν(C) + pdiff (k, C) ⩽ ρ] ≥ 1 − ρ. Eq. (103) proves the theorem. Finally, ℓ = O(log d) for polynomially small η and constant ρ, giving Ω(d/ log d) queries.
61
Note that implicit in the argument above p was that the critical window for masked diffusion has width which is only an O(log d/d) (rather than O( log d/d) as in Gaussian and uniform diffusion) fraction of the total range of noise levels. We make this precise with the following analogue of Propositions 6.4 and 7.4: above the window, the revealed coordinates of a random codeword are close to uniform, while below it, the original G point can be recovered with high probability. In analogy with errU C (t) and errC (t), define errM C (m) := inf
P
zb(·) z∼q C , I∼Unif
z (I, zI ) ̸= z] , [b ([d] m)
the smallest error of any decoder that observes m uniformly random coordinates of a sample from q C . Proposition 8.6 (Critical window). Let |C| = 2k and ℓ ≥ 1, and fix 0 ⩽ m < d. If m ⩽ k − ℓ, then for every I ⊆ [d] with |I| = m and X0 ∼ q C (with law(X0,I ) taken conditionally on C), EC TV law(X0,I ), Unif({±1}I ) ⩽ 2−ℓ/2 .
(104)
−ℓ EC errM . C (m) ⩽ 2
(105)
If m ≥ k + ℓ, then Proof. Fix I with |I| = m and write N (xI ) := |matchI (xI )|, so that law(X0,I ) = N (·)/2k . Since C is a uniformly random 2k -element subset of X , each N (xI ) is hypergeometric with mean 2k−m and variance at most 2k−m . Thus 1 EC TV law(X0,I ), Unif({±1}I ) = 2 ⩽
X
EC N (xI )2−k − 2−m
xI ∈{±1}I
1 m −k √ k−m 1 −(k−m)/2 ·2 ·2 2 = ·2 ⩽ 2−ℓ/2 2 2
for m ⩽ k − ℓ ,
which gives Eq. (104). For Eq. (105), consider the decoder which outputs the unique element of matchI (zI ) whenever this set is a singleton, and something arbitrary otherwise. It is incorrect only if some other codeword agrees with z on I. Conditionally on z, the remaining 2k − 1 codewords form a uniformly random subset of X \ {z}, and each agrees with z on I with probability at most 2−m , so by a union bound k−m ⩽ 2−ℓ for m ≥ k + ℓ. EC errM C (m) ⩽ 2
9
Outlook
In this work, we compared the parallel sampling capabilities of three leading paradigms for diffusion language modeling. We first showed that the query complexity of uniform and Gaussian diffusion can adapt to the intrinsic complexity of the data distribution as quantified by its dual total correlation, matching what was previously known for masked diffusion [CCL26, LZ25]. On top of establishing parity in this regard across all three paradigms, we then exhibited a simple family of distributions over which uniform and Gaussian diffusion provably outperform any sampling algorithm based on masked diffusion. In contrast with prior heuristic reasoning which suggested that masked diffusion is less amenable to parallelism because it cannot revise the tokens it commits to during sampling, we showed that the origin of this separation comes from a different mechanism. Instead, it is due to the fact that the critical window of noise levels for masked diffusion, over which the posterior transitions from being uninformative to sharply concentrating on a small subset of the support, is asymptotically narrower (by a 62
√ e d)) than the analogous windows for uniform and Gaussian diffusion. To our knowledge, this factor of Θ( is the first connection between critical windows and discretization error for diffusion sampling. That said, this work only touches upon one of many important performance axes for diffusion language modeling and should not be interpreted reductively as advocating for one paradigm over another. Indeed, there are various other advantages of masked diffusion, for instance its any-order generation capability, which remain poorly understood in theory but incredibly relevant in practice [YGG+ 25, KSK+ 25]. On the computational side, prior work of Ghio et al. [GDKZ24] and Bhatt et al. [BGKM26] compared the algorithmic complexity of denoising under different corruption processes in the context of spin glass-like distributions. Lastly, in the case of Gaussian diffusion, we have only considered the most basic form of embedding, leaving open the exploration of richer encoding maps. Rigorously understanding all of these aspects, and more generally the interplay between corruption process, parallelism, complexity of denoising, and flexibility of inference remains an important open direction, and we leave further exploration of this to future work.
Acknowledgments We thank Fan Chen, Sinho Chewi, Kevin Cong, Khashayar Gatmiry, Holden Lee, Yiwen Kou, Jerry Li, Jianfeng Lu, Raghu Meka, Adil Salim, and Yimeng Wang for many illuminating discussions on discretization of discrete and continuous diffusion models in recent years. We also thank Martin Wainwright and Yuting Wei for coordinating with us on the discussion of our concurrent works in forthcoming versions of our papers.
Statement on AI usage Most of the key ideas of this work were due to the authors from spring and early summer of this year, with assistance from frontier models in a few crucial parts. Most notably, the telescoping argument for the DTCadaptive rates was suggested by AI. The original proof of the reverse DPI for Gaussian diffusion that was generated by AI was extremely difficult to verify, but further human-AI interaction resulted in the final form of the argument. In an earlier version, we also overlooked the need for the extra Bregman projection step in the uniform diffusion case, which was subsequently suggested by AI. While we had an entirely humangenerated argument for our main separation results, AI assisted with simplifying some parts of the argument. It suggested a simpler test statistic for locating the critical window than the one we initially considered. We previously had an alternative argument for sampling below the critical noise level by proving a bound on the relaxed Lipschitz constant of the posterior and then appealing to existing discretization analyses. We believed that this step could be simplified but did not attempt to optimize this part as it did not affect the final result. AI however made the simple but helpful observation that below the critical window, a one-step decoding sufficed. A coding agent also implemented the numerical simulation in Figure 4, with no further guidance beyond the initial prompt seeding it with the experiment of [XP26]. √ e d) upper bound for In fact, we could use it to prove a considerably more general version of the O( Gaussian diffusion than the one that appears here, generalizing from uniform distributions over random subsets of the hypercube to random empirical distributions for any background measure over the unit sphere satisfying concentration of convex Lipschitz functions. We chose not to include this result as Theorem 1.3 already captures the bulk of the intuition. Frontier models were also used to create all of the figures in this paper, and to catch minor typos and errors in notation and setting of parameters. Beyond these aspects, everything in this paper was human written.
63
References [ABBM26]
Beatrice Achilli, Marco Benedetti, Giulio Biroli, and Marc Mézard. Theory of speciation transitions in diffusion models with general class structure. Journal of Statistical Mechanics: Theory and Experiment, 2026(4):043304, 2026.
[AJH+ 21]
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, 2021.
[BBDBM24] Giulio Biroli, Tony Bonnaire, Valentin De Bortoli, and Marc Mézard. Dynamical regimes of diffusion models. Nature Communications, 15(1):9957, 2024. [BGKM26]
Alankrita Bhatt, Mukur Gupta, Germain Kolossov, and Andrea Montanari. Generating from discrete distributions using diffusions: Insights from random constraint satisfaction problems. arXiv preprint arXiv:2603.20589, 2026.
[CBDB+ 22] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279, 2022. [CCL26]
Sitan Chen, Kevin Cong, and Jerry Li. Optimal inference schedules for masked diffusion models. In Steve Hanneke and Tor Lattimore, editors, Proceedings of Thirty Ninth Conference on Learning Theory, volume 336 of Proceedings of Machine Learning Research, pages 1279–1311. PMLR, 29 Jun–03 Jul 2026.
[CY24]
Hongrui Chen and Lexing Ying. Convergence analysis of discrete diffusion model: Exact implementation through uniformization. arXiv preprint arXiv:2402.08095, 2024.
[DHW26a]
Daniil Dmitriev, Zhihan Huang, and Yuting Wei. Efficient sampling with discrete diffusion models: Sharp and adaptive guarantees. arXiv preprint arXiv:2602.15008, 2026.
[DHW26b]
Daniil Dmitriev, Zhihan Huang, and Yuting Wei. Provably adaptive sampling with uniform and remasking discrete diffusion models. arXiv preprint arXiv:2608.23554, 2026.
[Die26]
Sander Dieleman. Continuous diffusion language models. Blog post, https://sander. ai/2026/08/24/continuous-dlms.html, 2026.
[Dif26]
DiffusionGemma Team. arXiv:2608.00146, 2026.
[DSR+ 22]
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categorical data. arXiv preprint arXiv:2211.15089, 2022.
[EAMS22]
Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from the Sherrington– Kirkpatrick Gibbs measure via algorithmic stochastic localization. In IEEE Symposium on Foundations of Computer Science (FOCS), 2022.
DiffusionGemma technical report.
64
arXiv preprint
[EAMS23]
Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from mean-field Gibbs measures via diffusion processes. arXiv preprint arXiv:2310.08912, 2023.
[GDKZ24]
Davide Ghio, Yatin Dandi, Florent Krzakala, and Lenka Zdeborová. Sampling with flows, diffusion, and autoregressive neural networks from a spin-glass perspective. Proceedings of the National Academy of Sciences, 121(27):e2311810121, 2024.
[GVS+ 23]
Kristian Georgiev, Joshua Vendrow, Hadi Salman, Sung Min Park, and Aleksander Madry. The journey, not the destination: How data guides diffusion models. arXiv preprint arXiv:2312.06205, 2023.
[HJA20]
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[HQL+ 26]
Keya Hu, Linlu Qiu, Yiyang Lu, Hanhong Zhao, Tianhong Li, Yoon Kim, Jacob Andreas, and Kaiming He. ELF: Embedded language flows. arXiv preprint arXiv:2605.10938, 2026.
[HV93]
Te Sun Han and Sergio Verdú. Approximation theory of output statistics. IEEE Transactions on Information Theory, 39(3):752–772, 1993.
[KLZ+ 26]
Kelvin Kan, Xingjian Li, Benjamin J Zhang, Tuhin Sahai, Stanley Osher, and Markos A Katsoulakis. Vocabulary-size-independent convergence of discrete diffusion models: Adjoint equations induce the right space. arXiv preprint arXiv:2605.17232, 2026.
[KSK+ 25]
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham M. Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu, editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 30749– 30768. PMLR, 13–19 Jul 2025.
[LC24]
Marvin Li and Sitan Chen. Critical windows: non-asymptotic theory for feature emergence in diffusion models. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. JMLR.org, 2024.
[LC26]
Gen Li and Changxiao Cai. Breaking ar’s sampling bottleneck: Provable acceleration via diffusion language models. Advances in Neural Information Processing Systems, 38:11700– 11725, 2026.
[LHL+ 25]
Yuchen Liang, Renxiang Huang, Lifeng Lai, Ness Shroff, and Yingbin Liang. Absorb and converge: Provable convergence guarantee for absorbing discrete diffusion models. arXiv preprint arXiv:2506.02318, 2025.
[LKC25]
Marvin Li, Aayush Karan, and Sitan Chen. Blink of an eye: a simple theory for feature localization in generative models. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon LacosteJulien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu, editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 35047–35080. PMLR, 13–19 Jul 2025.
65
[LLLS25]
Yuchen Liang, Yingbin Liang, Lifeng Lai, and Ness Shroff. Discrete diffusion models: Novel analysis and new sampler guarantees. In Advances in Neural Information Processing Systems, 2025.
[LME23]
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2023.
[LTG+ 22]
Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-LM improves controllable text generation. In Advances in Neural Information Processing Systems, 2022.
[Lu26]
Jianfeng Lu. A mathematical introduction to diffusion models. arXiv:2607.01693, 2026.
[LYA+ 26]
Chanhyuk Lee, Jaehoon Yoo, Manan Agarwal, Sheel Shah, Jerry Huang, Aditi Raghunathan, Seunghoon Hong, Nicholas M Boffi, and Jinwoo Kim. Flow map language models: One-step language modeling via continuous denoising. arXiv preprint arXiv:2602.16813, 2026.
[LZ25]
Hugo Lavenant and Giacomo Zanella. Error bounds and optimal schedules for masked diffusions with factorized approximations. arXiv preprint arXiv:2510.25544, 2025.
[MW23]
Andrea Montanari and Yuchen Wu. Posterior sampling in high dimension via diffusion processes. arXiv preprint arXiv:2304.11449, 2023.
[NWZ20]
Jonathan Niles-Weed and Ilias Zadik. The all-or-nothing phenomenon in sparse tensor pca. Advances in Neural Information Processing Systems, 33:17674–17684, 2020.
[NZY+ 25]
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025.
[PCRD26]
Le-Tuyet-Nhi Pham, Giovanni Conforti, Zhenjie Ren, and Alain Durmus. Discrete flow matching: Convergence guarantees under minimal assumptions. arXiv preprint arXiv:2605.08882, 2026.
[PSO+ 25]
Le-Tuyet-Nhi Pham, Dario Shariatian, Antonio Ocello, Giovanni Conforti, and Alain Durmus. Bit-level discrete diffusion with Markov probabilistic models: An improved framework with sharp convergence bounds under minimal assumptions. arXiv preprint arXiv:2502.07939, 2025.
[RA23]
Gabriel Raya and Luca Ambrogioni. Spontaneous symmetry breaking in generative diffusion models. Advances in Neural Information Processing Systems, 36:66377–66389, 2023.
[RCRY25]
Yinuo Ren, Haoxuan Chen, Grant Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via a stochastic integral framework. In International Conference on Learning Representations, volume 2025, pages 42904–42941, 2025.
[RCZ+ 25]
Yinuo Ren, Haoxuan Chen, Yuchen Zhu, Wei Guo, Yongxin Chen, Grant M Rotskoff, Molei Tao, and Lexing Ying. Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. In Advances in Neural Information Processing Systems, 2025. 66
arXiv preprint
[RDE+ 26]
Daan Roos, Oscar Davis, Floor Eijkelboom, Michael Bronstein, Max Welling, Ismail Ilkan Ceylan, Luca Ambrogioni, and Jan-Willem van de Meent. Categorical flow maps. arXiv preprint arXiv:2602.12233, 2026.
[SAS+ 24]
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems, 2024.
[SDG+ 25]
Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin Chiu, and Volodymyr Kuleshov. The diffusion duality. In International Conference on Machine Learning, 2025.
[SDWMG15] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, 2015. [SE19]
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, 2019.
[SFLW24]
Antonio Sclocchi, Alessandro Favero, Noam Itzhak Levi, and Matthieu Wyart. Probing the latent hierarchical structure of data via diffusion models. arXiv preprint arXiv:2410.13770, 2024.
[SFW25]
Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. Proceedings of the National Academy of Sciences, 122(1):e2408799121, 2025.
[SHW+ 24]
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and generalized masked diffusion for discrete data. In Advances in Neural Information Processing Systems, 2024.
[SLY+ 26]
Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux, Jingyu Liu, John Thickstun, and Ante Jukic. Scaling beyond masked diffusion language models. arXiv preprint arXiv:2602.15014, 2026.
[SSDK+ 21]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[SSP+ 25]
Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dallatorre, Bernardo P de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. In International Conference on Learning Representations, 2025.
[Vin11]
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23(7):1661–1674, 2011.
[Wai26]
Martin J Wainwright. The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity. arXiv preprint arXiv:2608.13520, 2026.
67
[Wib25]
Andre Wibisono. Mixing time of the proximal sampler in relative fisher information via strong data processing inequality (extended abstract). In Nika Haghtalab and Ankur Moitra, editors, Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 5716–5717. PMLR, 30 Jun–04 Jul 2025.
[WSSK26]
Guanghan Wang, Yair Schiff, Subham Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling. Advances in Neural Information Processing Systems, 38:147282–147339, 2026.
[XP26]
Zhiyang Xun and Eric Price. Query lower bounds for diffusion sampling. arXiv preprint arXiv:2604.10857, 2026.
[YGG+ 25]
Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning. In International Conference on Learning Representations, volume 2025, pages 77875–77898, 2025.
[ZC26]
Yunxiao Zhao and Changxiao Cai. Adaptation to intrinsic dependence in diffusion language models. arXiv preprint arXiv:2602.20126, 2026.
[ZCG25]
Zikun Zhang, Zixiang Chen, and Quanquan Gu. Convergence of score-based discrete diffusion models: A discrete-time analysis. In International Conference on Learning Representations, 2025.
√ e d) lower bound for uniform diffusion A Ω( The upper bound in Section 6 searches for the critical time t⋆unif defined in Eq. (22). We show √ that this search e cannot be avoided with approximate scores. The hard √ distribution family contains Ω( d) possible values e of t⋆unif , separated by more than the width w = Θ(1/ d) of the critical window defined in Eq. (107) below, and a query far from the realized value usually receives the score of the uniform distribution. This is the continuous-time counterpart of the masked lower bound in Section 8. There, the hidden transition occurs when the number of revealed coordinates passes the hidden integer k. Here it occurs when the query time passes t⋆unif . Above t⋆unif , the noised random empirical measure is close to uniform. Below t⋆unif , a genuine sample is close to a point of C, but a fixed query chosen without knowing C is unlikely to be close to any point of C. For every integer k satisfying κ− ⩽ k log 2/d ⩽ κ+ , let t⋆unif (k) be the unique solution of d Iunif (t⋆unif (k)) = k log 2 = log M . The derivative I′unif (t) = −
(106)
e−t 1 + e−t log 2 1 − e−t
is bounded above and below in magnitude on [τ− , τ+ ]. Given A, CΛ > 0, η ∈ (0, 1/10), and ρ ∈ (0, 1/4), define r
CΛ d Λ = log , ηρ2
w=A
68
Λ . d
(107)
Fix a sufficiently large constant D > 0. Starting with the smallest admissible integer k, retain every ⌈Ddw⌉-th one, and denote the resulting set by κ+ d/ log 2 − ⌈κ− d/ log 2⌉ κ− d U := + j⌈Ddw⌉ : j = 0, 1, . . . , . Kd log 2 ⌈Ddw⌉ If k < k ′ are consecutive elements of KdU , the mean value theorem applied to Eq. (106) gives, for some ξ ∈ (t⋆unif (k ′ ), t⋆unif (k)), (k ′ − k) log 2/d t⋆unif (k) − t⋆unif (k ′ ) = . −I′unif (ξ) Consequently, after choosing D large enough, t⋆unif (k) − t⋆unif (k ′ ) ≥ 10w,
|KdU | ≥
cd cU ≥ , ⌈Ddw⌉ w
(108)
for all sufficiently large d. In particular, cU |KdU | ≥ A
r
d . Λ
(109)
We choose k uniformly from KdU and, conditional on k, choose C uniformly among the 2k -element subsets of X . Once k is fixed, we write t⋆unif for t⋆unif (k). Throughout, sCt denotes the exact score of qtC , that is, sCt (x)[i, a] = qtC (xi←a )/qtC (x). We measure the accuracy of a score oracle O by εunif (O; q C , t) from Definition 2.2. Theorem A.1 (Uniform diffusion lower bound). Fix ρ ∈ (0, 1/4). There are constants A, CΛ , c, c0 , d0 > 0, depending only on [κ− , κ+ ] and ρ, such that the following holds for d ≥ d0 . Given η ∈ (0, 1/10), take Λ, w, and KdU as defined above, and suppose Λ ⩽ c0 d. There is a family of score oracles {Ok,C }k∈KU ,|C|=2k , each d returning normalized score vectors with strictly positive entries, satisfying sup εunif (Ok,C ; q C , t) ⩽ η, t>0
k ∈ KdU ,
|C| = 2k ,
such that every adaptive randomized algorithm A that is not given k or C and makes at most r d 2 Q ⩽ cρ Λ
(110)
(111)
O
score oracle queries has output law qbA k,C satisfying P
k∼Unif(KdU )
O
[TV(b qA k,C , q C ) ≥ 1 − ρ] ≥ 1 − ρ .
(112)
C∼Unif (2Xk ) In particular, for constant ρ and any polynomially small η p = d−Θ(1) , sampling with at most constant error for at least a constant fraction of the instances requires Ω( d/ log d) score oracle queries. For the rest of the appendix, Λ and w have the values in Eq. (107). We may increase A, CΛ , d0 and decrease c0 later when necessary so that log d ⩽ Λ ⩽ c0 d and w ⩽ min{τ− /8, 1/8}. Fix a constant γ > 0 sufficiently small in terms of [τ− , τ+ ].
69
In Section A.1, we construct the oracle family and prove its accuracy. In Section A.2, we prove the query lower bound.
A.1
Adversarial oracle construction
Let 1 ∈ R[d]×{−1,1} be the score vector defined by 1[i, a] = 1 for every i ∈ [d] and a ∈ {−1, 1}, and define the reference oracle by O0 (t, x) := 1 for every query (t, x). This is the exact score oracle of Unif(X ). Proposition 6.4 already shows that qtC is close to Unif(X ) above t⋆unif ; the following lemma records the corresponding statement for the score error. Lemma A.2. For every K > 0, there are constants AK , cK , dK > 0, depending only on [κ− , κ+ ] and K, such that, whenever A ≥ AK , d ≥ dK , log d ⩽ Λ ⩽ cK d, and t ≥ t⋆unif (k) + 4w, EC TV(qtC , Unif(X )) ⩽ e−KΛ ,
EC εunif (O0 ; q C , t) ⩽ e−KΛ .
Both inequalities hold uniformly over k ∈ KdU and t ≥ t⋆unif (k) + 4w. Proof. Since t ≥ t⋆unif + 4w, Proposition 6.4 with exponent K + 2, whose constants at that exponent we take as the AK , cK , dK of the present lemma, gives EC TV(qtC , Unif(X )) ⩽ e−(K+2)Λ ⩽ e−KΛ . Eq. (32) places the exact scores in a fixed compact subinterval of (0, ∞), on which ψ(v, 1) ⩽ C|v − 1|. For each C and each coordinate i, we then have X EX∼qtC ψ(sCt (X)[i, −Xi ], 1) ⩽ C qtC (xi←−xi ) − qtC (x) x∈X
⩽ 4C TV(qtC , Unif(X )) .
(113)
The second inequality follows by inserting 2−d and using the fact that a coordinate flip permutes X . Summing Eq. (113) over i and averaging over C gives EC εunif (O0 ; q C , t) ⩽ 4Cd e−(K+2)Λ ⩽ e−KΛ , where the last inequality holds after increasing dK . Well below t⋆unif , a sample Xt from qtC is usually close to the point X0 of C from which it was generated. For t > 0, define Rt = ⌈d(βt + γw)⌉ + 1, NC,t = x ∈ X : min dH (x, y) ⩽ Rt . y∈C
The oracle will return the exact score on the neighborhood NC,t and the all-one score outside it. The next lemma separates the two facts needed later: a genuine sample from qtC usually lies in this neighborhood, whereas a fixed query usually does not. Lemma A.3. There are constants c, C, c1 , d1 > 0, depending only on [κ− , κ+ ], such that, whenever d ≥ d1 , log d ⩽ dw2 ⩽ c1 d, and 0 < t ⩽ t⋆unif − 4w, the following statements hold uniformly over the allowed values of k and t.
70
(i) For every C, 2
2
c qtC (NC,t ) ⩽ e−2γ dw .
(114)
c , then If the score output of the oracle O equals sCt on NC,t and the all-one score on NC,t 2
εunif (O; q C , t) ⩽ Cd(1 + log d)e−cdw .
(115)
P[x ∈ NC,t ] ⩽ e−cdw .
(116)
(ii) For every fixed x ∈ X , C
Proof. Fix k and write t⋆unif = t⋆unif (k). Since βr′ = e−r /2 and 0 ⩽ t < t⋆unif ⩽ τ+ , Z t⋆
e−r e−τ+ ⋆ dr ≥ (tunif − t) ≥ 2e−τ+ w . 2 2 t p Choose γ < e−τ+ /2. Since dw2 ≥ log d, i.e., w ≥ log d/d, after increasing d1 we have βt⋆unif − βt =
unif
Rt 2 1 ⩽ βt + γw + ⩽ βt⋆unif − cw < . d d 2
(117)
We start by showing Eq. (114). Fix C, draw Y ∼ q C , and generate X ∼ (KtU )⊗d (· | Y ). Then dH (X, Y ) ∼ Bin(d, βt ) and minz∈C dH (X, z) ⩽ dH (X, Y ). By Hoeffding’s inequality, we have 2
2
c qtC (NC,t ) ⩽ P[Bin(d, βt ) ≥ Rt ] ⩽ e−2γ dw ,
proving Eq. (114). We next bound the score error outside NC,t . For every x ∈ X and i ∈ [d], the binary likelihood ratio in Lemma 2.3 gives βt 1 − βt ⩽ sCt (x)[i, −xi ] ⩽ . 1 − βt βt Hence qtC (x)ψ(sCt (x)[i, −xi ], 1) ⩽ qtC (x) + log
1 − βt C i←−xi ). qt (x βt
(118)
If x ∈ / NC,t , then every one-coordinate neighbor of x has distance at least Rt from C. Summing Eq. (118) over x ∈ / NC,t and i ∈ [d] gives 1 − βt C εunif (O; q , t) ⩽ d 1 + log P[Bin(d, βt ) ≥ Rt ] . (119) βt If βt ≥ d−2 , then the logarithmic factor in Eq. (119) is at most C(1 + log d), and Hoeffding’s inequality proves Eq. (115). If 0 < βt < d−2 , set m = ⌈γdw⌉. By the Chernoff bound, edβt m P[Bin(d, βt ) ≥ Rt ] ⩽ . m For m ≥ 2, the function u 7→ (1 + log(1/u))um is increasing on (0, d−2 ]. Therefore, e m 1 − βt 2 1 + log P[Bin(d, βt ) ≥ Rt ] ⩽ C(1 + log d) ⩽ C(1 + log d)e−cdw . βt dm 71
(120)
The last inequality uses m ≥ γdw, w ⩽ 1, and hence m log(dm/e) ≥ cdw log d ≥ cdw2 for sufficiently large d. Eqs. (119)–(120) prove Eq. (115) uniformly over t. Now fix x before drawing C. Since C is a uniform 2k -element subset of X , P[x ∈ NC,t ] ⩽ EC |C ∩ {y : dH (x, y) ⩽ Rt }| C
Rt d 2k X ⩽ exp (k log 2 − d log 2 + dh2 (Rt /d)) . = d j 2
(121)
j=0
By Eq. (106), k log 2 = d Iunif (t⋆unif ) = d log 2 − h2 (βt⋆unif ) .
(122)
Moreover, h′2 (u) = log((1 − u)/u) is bounded below by a positive constant for u ⩽ βτ+ < 1/2. Eq. (117) consequently gives Z βt⋆ 1−u unif h2 (βt⋆unif ) − h2 (Rt /d) = log du ≥ cw . (123) u Rt /d Substituting Eqs. (122) and (123) into Eq. (121) proves Eq. (116). The oracle now follows the three regimes suggested by the two lemmas above. It gives the exact score within distance 4w of t⋆unif (the window). Below this window it gives the exact score only on NC,t and all-one score otherwise. Above the window it gives the all-one score whenever that answer has error at most η, and otherwise gives the exact score. Definition A.4 (Adversarial uniform diffusion oracle). For each fixed instance (k, C), define Ok,C by sCt (x), |t − t⋆unif | < 4w, C ⋆ st (x), t ⩽ tunif − 4w and x ∈ NC,t , Ok,C (t, x) := 1, . / NC,t , t ⩽ t⋆unif − 4w and x ∈ C ⋆ 1, t ≥ tunif + 4w and εunif (O0 ; q , t) ⩽ η, sC (x), t ≥ t⋆ + 4w and ε (O ; q C , t) > η. t
unif
unif
0
Proposition A.5 (Validity and indistinguishability of the oracle). The constants in Theorem A.1 can be chosen so that the oracles in Definition A.4 satisfy Eq. (110) for every allowed instance. For these choices, whenever |t − t⋆unif | ≥ 4w and x ∈ X is fixed independently of C, P[Ok,C (t, x) ̸= O0 (t, x)] ⩽ e−4Λ .
(124)
C
Proof. We first choose A large enough that Lemma A.2 holds with K = 6 and that, by Lemma A.3 (whose hypothesis log d ⩽ dw2 = A2 Λ ⩽ c1 d holds once c0 ⩽ c1 /A2 ), uniformly over t ⩽ t⋆unif (k) − 4w, 2
2
εunif (Ok,C ; q C , t) ⩽ Cd(1 + log d)e−cdw = Cd(1 + log d)e−cA Λ ⩽ η ,
(125)
once A and then d0 are sufficiently large. We then further decrease c0 and increase d0 so that, by Lemma A.3 again, uniformly over t ⩽ t⋆unif (k) − 4w, EC 1[x ∈ NC,t ] ⩽ e−cdw ⩽ e−4Λ .
72
(126)
Fix an instance and a time. In the interval |t − t⋆unif | < 4w, Definition A.4 returns the exact score. If t ⩽ t⋆unif − 4w, Eq. (125) gives error at most η. If t ≥ t⋆unif + 4w, the oracle returns the all-one score only when its error is at most η and returns the exact score otherwise. This proves Eq. (110). It remains to prove Eq. (124). Below the window, the oracle can differ from O0 only when x ∈ NC,t , so Eq. (126) gives the desired bound. Above the window, the two answers can differ only when εunif (O0 ; q C , t) > η. Markov’s inequality and Lemma A.2 with K = 6 give P[εunif (O0 ; q C , t) > η] ⩽ C
e−6Λ ⩽ e−4Λ , η
where the last step uses η ≥ e−2Λ = (ηρ2 /(CΛ d))2 , which holds since η, ρ < 1 and we may take CΛ ≥ 1. This finishes the proof of the proposition.
A.2
Distribution indistinguishability: the lower bound proof
Let R denote all the internal randomness of an adaptive algorithm A. The following compares the algorithm’s interaction with the constructed oracle to its interaction with O0 . For the latter, every answer is the all-one score, so the entire query sequence is determined by R and is independent of k and C. Proposition A.6 (Coupling argument). Couple an algorithm’s run against Ok,C with its run against O0 , using the same internal randomness R. For an algorithm making at most Q queries, let pdiff (k, C) be the conditional probability over R given k, C that the runs receive different oracle answers at some point. Then 1 −4Λ Ek,C pdiff (k, C) ⩽ Q +e . (127) |KdU | Proof. Let D denote the event that the two coupled runs receive different oracle answers at some point. By definition, Ek,C pdiff (k, C) = Ek,C P[D | k, C] = P [D] = ER P [D | R] . (128) R
k,C,R
k,C
Now condition on any fixed realization of the internal randomness R = r. We will bound Pk,C [D | R = r] uniformly over r. It is enough to bound the probability of a first different answer, because the queries agree up to that point. Run the algorithm against O0 , padding with irrelevant queries if it stops before making Q queries. This produces a deterministic sequence (tj , xj )Q j=1 , independent of k and C. Consider the randomness of k. By ⋆ Eq. (108), the possible values of tunif (k) are separated by at least 10w. Therefore, for each query time tj , the interval |tj − t⋆unif (k)| < 4w contains the critical time of at most one k ∈ KdU , and hence P[|tj − t⋆unif (k)| < 4w] ⩽ k
1 . |KdU |
A union bound shows that at least one query time lies in the critical window with probability at most Q/|KdU | over k. Fix a k for which every query time lies outside the critical window, i.e., |tj − t⋆unif (k)| ≥ 4w for every j. Since each xj was fixed before C was drawn, Proposition A.5 gives, for every j, P[Ok,C (tj , xj ) ̸= O0 (tj , xj )] ⩽ e−4Λ . C
73
A union bound over the Q queries shows that some oracle answer differs outside the window with probability at most Qe−4Λ . Consequently, for the fixed realization R = r, the probability that a query enters the critical window or that an oracle answer differs outside the window is at most the right-hand side of Eq. (127). If neither event occurs, the two runs have identical queries and answers and hence make the same next query at every step, so D does not occur. Averaging over R in Eq. (128) proves the proposition. Proof of Theorem A.1. Proposition A.5 gives an oracle satisfying Eq. (110) for every instance. Fix an adaptive randomized algorithm A, and let ν be its output law when every answer comes from O0 . Since O0 gives the same answer on every instance, ν is independent of k and C. For each fixed instance (k, C), couple the run against Ok,C with the run against O0 using the same O internal randomness R. Let Yk,C (R) and Y0 (R) denote their outputs. Their laws over R are qbA k,C and ν, respectively, and hence O
TV(b qA k,C , ν) ⩽ P[Yk,C (R) ̸= Y0 (R)] R
⩽ pdiff (k, C) . Since q C is supported on C, O
O
TV(b qA k,C , q C ) ≥ 1 − qbA k,C (C) ≥ 1 − ν(C) − pdiff (k, C) .
(129)
Fix k. A fixed point of X belongs to a uniformly chosen 2k -element subset with probability 2k−d . Since ν is independent of C, EC|k ν(C) = EX∼ν P [X ∈ C] = 2k−d ⩽ e−(log 2−κ+ )d . C|k
(130)
We now choose the constant c in Eq. (111). By Eqs. (109) and (111), Q cA 2 ⩽ ρ . U cU |Kd | Choose c ⩽ cU /(4A), so this term is at most ρ2 /4. Since Λ ≥ log d, r d −4Λ ρ2 Qe−4Λ ⩽ cρ2 e ⩽ cρ2 e−7Λ/2 ⩽ Λ 4 after increasing d0 . Increase d0 once more so that the right-hand side of Eq. (130) is at most ρ2 /2. Proposition A.6 then gives Ek,C ν(C) + pdiff (k, C) ⩽ ρ2 . Markov’s inequality yields P [ν(C) + pdiff (k, C) ⩽ ρ] ≥ 1 − ρ .
k,C
Eq.p (129) proves Eq. (112). Finally, Λ = O(log d) for polynomially small η and constant ρ, giving the stated Ω( d/ log d) lower bound.
74
B
Deferred proofs
B.1
Proofs from Section 2
B.1.1
Proof of Lemma 2.1
Proof. Fix y ∈ X and i ∈ [d], and denote v := st (y)[i, ·] ,
V :=
X
vb .
b∈Σ
Note that vyi = 1. Let ga := P[Xsi = a, (Xt )−i = y−i ] ,
G :=
X
ga = P[(Xt )−i = y−i ] .
a∈Σ
Conditionally on Xs , the coordinates of Xt are independent, and (Xt )i ∼ KhU (· | Xsi ). As a result, for every b ∈ Σ, X qt (y i←b ) = (131) ga KhU (b | a) = αh gb + βh G . a∈Σ
Summing Eq. (131) over b ∈ Σ and using αh + Sβh = 1 gives identifies the normalizing constant as V = G/qt (y) .
P
b∈Σ qt (y
i←b ) = G, so dividing by q (y) t
The Gibbs sampling marginal. Dividing Eq. (131) by qt (y) and then by V = G/qt (y), vb gb = αh · + βh . V G
(132)
As gb /G = P[Xsi = b | (Xt )−i = y−i ], solving Eq. (132) for gb /G gives the second claim. The posterior marginal. Write π := pgibbs [v] for the distribution just computed. As (Xt )i is conditionally h independent of (Xt )−i given Xsi , by Bayes’ rule we have πa KhU (yi | a) . U ′ a′ ∈Σ πa′ Kh (yi | a )
P[Xsi = a | Xt = y] = P
The denominator equals αh πyi + βh , which by Eq. (132) is vyi /V = 1/V . Substituting this and αh V πa = va − βh V , both read off from Eq. (132), we obtain the first claim, αh 1[a = yi ] + βh P[Xsi = a | Xt = y] = V πa αh 1[a = yi ] + βh = va − βh V = ppost h,yi [v](a) . αh
75
B.1.2
Proof of Lemma 2.8 (I)
Proof. Fix I with |I| = m, X0 , and i ∈ / I, and write Y = X0 ; the partial assignment (I, X0,I ) is O , and the homogeneity ψ(λc, λs) = λψ(c, s) for λ > 0, consistent. By Lemma 2.6, the definition of qbi|I 1 − αt X ψ(st (Y )[i, a], sbt (Y )[i, a]) αt a∈[S]
=
qi|I (a | X0,I ) O (a | X0,I ) − qi|I (a | X0,I ) + qi|I (a | X0,I ) log O qbi|I qbi|I (a | X0,I )
X a∈[S]
!
O = KL(qi|I (· | X0,I ) ∥ qbi|I (· | X0,I )) ,
P P O (a | X ) = 1. Averaging over i ∈ since a∈[S] qi|I (a | X0,I ) = a∈[S] qbi|I / I and taking the expectation 0,I over I and X0 gives Eq. (1).
B.2
Proofs from Section 3
B.2.1
Proof of Proposition 3.6
Proof. For the first part, write va := st (y)[i, a] = qt (y i←a )/qt (y). Every entry of KtU (· | z) is at least βt > 0, so qt has full support and v ∈ RS>0 ; moreover vyi = 1 since y i←yi = y. For the constraint P va / a′ ∈Σ va′ ≥ βt , va P
a′ ∈Σ va′
qt (y i←a ) = P[(Xt )i = a | (Xt )−i = y−i ] . i←a′ ) a′ ∈Σ qt (y
=P
That this is lower bounded by βt follows by Proposition 3.3. P For the second part, write wa := va / a′ ∈Σ va′ , so that p := pgibbs [v] is given by p(a) = (wa − βh )/αh . h P −h −t It is immediate that a p(a) = 1. Moreover, wa ≥ βt by definition, and βt − βh = e S−e = αh βs , so P U [v](a) = p(a) ≥ (βt − βh )/αh = βs . Finally, the identity ppost a′ va′ Kh (b | a) p(a) follows from the h,b P definitions, and summing it over a gives a′ va′ (αh p(b) + βh ) = vb = 1. B.2.2
Proof of Lemma 3.10
Proof. For each coordinate i ∈ [d], let Bi ∼ Ber(e−T ) be independent, and let U = (U1 , . . . , Ud ) ∼ e by Unif(X ) be independent of X0 ∼ q and B. Define the coupled random vector X ei := X
( (X0 )i ,
Bi = 1,
Ui ,
Bi = 0.
.
e ∼ qT (this is exactly the coordinatewise kernel (K U )⊗d with αT = e−T and βT = (1 − e−T )/S). Then X T Conditioned on B = b ∈ {0, 1}d , the coordinates with bi = 0 are i.i.d. uniform, and the coordinates with bi = 1 have joint law equal to the corresponding marginal of q. Let q (b) denote the conditional law of e given B = b. By convexity of KL in its first argument, X KL(qT ∥ Unif(X )) = KL(EB [q (B) ] ∥ Unif(X )) ⩽ EB KL(q (B) ∥ Unif(X )) .
76
For any fixed b, write I(b) := {i : bi = 1} and k := |I(b)|. Since q (b) is uniform on coordinates outside I(b), KL(q (b) ∥ Unif(X )) = KL(law(X0,I(b) ) ∥ Unif(ΣI(b) )) ⩽ k log S . Taking expectation over B gives KL(qT ∥ Unif(X )) ⩽ EB |I(B)| log S = de−T log S , as claimed. B.2.3
Proof of Lemma 3.11
Proof. Fix 0 < t ⩽ 1 and use the same mask coupling as in the proof of Lemma 3.10, with Bi ∼ Ber(e−t ) e ∼ qt defined by and X ( ei = (X0 )i , Bi = 1, , X Ui , Bi = 0, where U ∼ Unif(Σd ) is independent. e ⩽ H(X, e B) = H(B) + H(X e | B), we obtain Using H(X) e − H(X0 ) ⩽ H(B) + H(X e | B) − H(X0 ) . H(qt ) − H(q) = H(X) Conditioned on B = b, the refreshed coordinates contribute (d − |I(b)|) log S nats of entropy and are independent of everything else, while the kept coordinates are a function of X0 . Hence e | B = b) = H(X0,I(b) ) + (d − |I(b)|) log S ⩽ H(X0 ) + (d − |I(b)|) log S . H(X Averaging over B yields e | B) − H(X0 ) ⩽ EB [d − |I(B)|] log S = d(1 − e−t ) log S ⩽ dt log S . H(X It remains to bound H(B). Since B has i.i.d. coordinates with P[Bi = 0] = 1 − e−t , we have H(B) = d h2 (1 − e−t ) . For 0 < ρ ⩽ 1, the standard bound h2 (ρ) ⩽ ρ log(e/ρ) gives H(B) ⩽ d(1 − e−t ) log
e . 1 − e−t
For 0 < t ⩽ 1 we have 1 − e−t ⩽ t, and also ρ 7→ ρ log(e/ρ) is increasing on (0, 1], so H(B) ⩽ dt log(e/t) . Combining the two parts, H(qt ) − H(q) ⩽ dt log(e/t) + dt log S = dt log(eS/t) , as desired.
77
B.3
Proofs from Section 4
B.3.1
Proof of Lemma 4.6
Proof. Apply Lemma 4.5 with µu = fα,u , νu = fβ,u , and ℓu = log(fα,u /fβ,u ); the first part is then the first identity of that lemma. For the second part, since ∇ log fλ,u (z) = (mλ,u (z) − z)/u for every λ ∈ ∆(Σ), a direct calculation yields ∇ℓu = ∇ log fα,u − ∇ log fβ,u =
1 (mα,u (z) − mβ,u (z)) , u
1 1 ∇2 log νu = ∇2 log fβ,u = − IdS + 2 Cβ,u (z) , u u 1 ∇2 ℓu = ∇2 log fα,u − ∇2 log fβ,u = 2 (Cα,u (z) − Cβ,u (z)) . u By Lemma 4.5 and the fact that Cλ,u (z) ⪰ 0 for all λ, 2 1 ∂u FI(fα,u ∥ fβ,u ) ≥ − 4 Efα,u ∥Cα,u (z) − Cβ,u (z)∥2F − 3 Efα,u ∥mα,u (z) − mβ,u (z)∥22 u u 2 3 ≥− + FI(fα,u ∥ fβ,u ) , u u2 where in the last step we used Proposition B.1 below. The above proof used the following elementary fact: Proposition B.1. Given p ∈ ∆(Σ), define C(p) = diag(p) − p⊗2 . Then for all p, p′ ∈ ∆(Σ), ∥C(p) − C(p′ )∥2F ⩽ 3∥p − p′ ∥22 . Proof. Let s := p + p′ and δ := p − p′ . Then C(p) − C(p′ ) = diag(δ) − 12 (sδ ⊺ + δs⊺ ), so ∥C(p) − C(p′ )∥2F = ∥δ∥22 − 2 as ∥s∥22 ⩽ 4 and ⟨s, δ⟩2 ⩽ 2 B.3.2
1 1 sa δa2 + ∥s∥22 ∥δ∥22 + ⟨s, δ⟩2 ⩽ 3∥δ∥22 , 2 2 a
X
2 a sa δa by Cauchy–Schwarz.
P
Proof of Lemma 4.7
Proof. Write φa for the density of N (ea , ρIdS ) and h := α − β. We have Z P ( a ha φa )2 2 P KL(fα,ρ ∥ fβ,ρ ) ⩽ χ (fα,ρ ∥fβ,ρ ) ≤ . a βa φa For fixed h, the right-hand side is convex in β, so it is maximized at some vertex β = eb , in which case the denominator is φb . A direct calculation gives Z (ea − eb )⊺ (ec − eb ) φa (z)φc (z) dz = exp . φb (z) ρ
78
Using
P
a ha = 0 and expanding the quadratic form, with x = e
Z
(
P
a ha φa )
2
φb
= (x − 1)h2b + (x2 − x)
1/ρ , we conclude that
X
h2a ≤ (x2 − 1)∥h∥22 ,
a̸=b
as claimed.
B.4
Proofs from Section 6
p Throughout this subsection we use the notation of Section 6: Λ and w = A Λ/d are as in Eqs. (24) and (26), I0 = [τ− /2, τ+ + 1] is the interval of Eq. (23), Lt = qtC /2−d is the likelihood ratio with respect to the uniform distribution, ℓt = βt /(1 − βt ), set is the clipped approximate score of Eq. (33), Tbt is the statistic of Eq. (34), and Ft is the squared clipped score error of Eq. (35). B.4.1
Proof of Lemma 6.2
Proof. Define the truncated likelihood ratios e c (x) = 2d (K U )⊗d (x | z) · 1[dH (x, z) > r] L z t
e z (x) = 2d (K U )⊗d (x | z) · 1[dH (x, z) ⩽ r] . L t
and
We can write the likelihood ratio as Lt (x) =
qtC (x) e c (x) + Ez∼Unif(C) L e z (x) . = Ez∼Unif(C) L z 2−d
e cz (x). Note that µ(x) ⩽ 1 because the expectation of the (un-truncated) Also define µ(x) = Ez∼Unif(X ) L likelihood ratio is 1. Then e c (x) − µ(x) + Ex∼Unif(X ) [1 − µ(x)] 2EC TV(qtC , Unif(X )) ⩽ EC Ex∼Unif(X ) Ez∼Unif(C) L z {z } | {z } | 2
1
e z (x) . + EC Ex∼Unif(X ) Ez∼Unif(C) L {z } | 3
We have 1
2
2 e c (x) − µ(x) ⩽ EC Ex∼Unif(X ) Ez∼Unif(C) L z 2 1 e cz (x) − µ(x) ⩽ Ex,z∼Unif(X ) L M 1 e c (x)2 ⩽ E L M x,z∼Unif(X ) z 1 ⩽ (1 + e−t )d · tanh(t/2)(r+1) , M
where in the second step we used that the variance of the empirical average over samples drawn without replacement is at most the variance of the empirical average over samples drawn with replacement, and in e c (x)2 ⩽ (1 + e−t )d tanh(t/2)r+1 · 2d (K U )⊗d (x | z) by Lemma 2.3, together the last step we used that L z t
79
with Ex,z∼Unif(X ) 2d (KtU )⊗d (x | z) = 1. Finally, 2
e z (x) = = 3 = Ex,z∼Unif(X ) L
P
z∼Unif(X ),x∼(KtU )⊗d (·|z)
[dH (x, z) ⩽ r] .
The claimed bound follows upon observing that the above quantity is P[Bin(d, βt ) ⩽ r]. B.4.2
Proof of Lemma 6.3
Proof. We will bound the Bayes-optimal recovery error by the error of the nearest-neighbor decoder, i.e., the one which given x outputs anything in arg minz ′ ∈C dH (x, z ′ ). This decoder succeeds when the true z satisfies dH (x, z) ⩽ r and for all z ′ ∈ C\{z}, dH (x, z ′ ) > r. The probability that dH (x, z) > r is exactly the second term on the right-hand side of the claimed bound, because the distribution of z given by sampling random C and then taking z ∼ q C is uniform over X , and conditioned on z the quantity dH (x, z) is distributed as Bin(d, βt ). The probability that some z ′ ∈ C \ {z} has dH (x, z ′ ) ⩽ r is at most the expected number of such ′ z . Conditioned on any x, z, the remaining elements of C are a random (M − 1)-sized subset of the cube, −1 so the expected number of such z ′ is at most M · |Br (x)|, where Br (x) is the closed Hamming ball of 2d −1 radius r around x. The size of Br (x) is independent of x and equal to 2d P[Bin(d, 1/2) ⩽ r], so the claim follows. B.4.3
Proof of Lemma 6.6
Proof. For the first part of Eq. (36), the definitions of Tt and Tbt (Eq. (34)) and the Cauchy–Schwarz inequality give d
|Tbt (x) − Tt (x)| ⩽
1 − e−2t p 1 − e−2t X ⩽ s e (x)[i, −x ] − s (x)[i, −x ] dFt (x) , t i t i 2e−t 2e−t i=1
−2t
and the prefactor 1−e is bounded on I0 by a constant depending only on [κ− , κ+ ]. 2e−t For the second part, fix any entry c := st (x)[i, −xi ] of the true score. The function u 7→ ψ(c, u) decreases on (0, c] and increases on [c, ∞), since its derivative is 1 − c/u. By Eq. (32), c ∈ [ℓt , ℓ−1 t ], so clipping the estimated entry sbt (x)[i, −xi ] to this interval cannot increase ψ(c, ·). On the compact time interval I0 , all the clipping intervals [ℓt , ℓ−1 t ] lie in a single interval [m0 , M0 ] ⊂ (0, ∞). For c, u ∈ [m0 , M0 ], the second derivative of u 7→ ψ(c, u), which is c/u2 , is at least m0 /M02 . Taylor’s theorem at u = c gives (u − c)2 ⩽
2M02 ψ(c, u) . m0
Applying this inequality to each clipped entry u = set (x)[i, −xi ], summing over i ∈ [d], and averaging under qtC gives the second part of Eq. (36) with C ′ = 2M02 /m0 , which depends only on I0 and hence only on [κ− , κ+ ]. B.4.4
Proof of Lemma 6.7
Proof. Consider a binomial random variable with d trials and success probability v = r/d. Its most likely d −dh2 (v) value is r, so the probability at r lies between 1/(d + 1) and 1. Since this probability is r e , Eq. (37) follows. 80
For a different success probability p, multiplying the lower bound in Eq. (37) by pr (1 − p)d−r gives the first inequality in Eq. (38). The second follows from Eq. (39), which holds for all v, p ∈ (0, 1) by applying log u ⩽ u − 1 to the two terms of the KL divergence: KL(Ber(v) ∥ Ber(p)) ⩽
v 2 (1 − v)2 (v − p)2 + −1= . p 1−p p(1 − p)
B.4.5
Proof of Lemma 6.8 l m Proof. Choose r = d max{βt , βt⋆unif } + DΛ , where D is a sufficiently large constant; decreasing cA,K,L if necessary, 1 ⩽ r ⩽ d − 1. Write v = r/d. Let Nr be the number of elements of C at Hamming distance exactly r from x. Its expectation is d µr = EC Nr = M 2−d . r By Eq. (37) of Lemma 6.7 and log M = d Iunif (t⋆unif ) = d(log 2 − h2 (βt⋆unif )) (Eq. (22)), log µr ≥ d h2 (v) − h2 (βt⋆unif ) − log(d + 1) ≥ cDΛ − log(d + 1) . The last step uses that v − βt⋆unif ≥ DΛ/d and that h′2 is lower bounded away from zero in the relevant interval. Since Λ ≥ log d, choosing D large enough gives µr ≥ 4eKΛ . Define the Hamming sphere Sr (x) = {z ∈ X | dH (x, z) = r}. For z ∈ Sr (x), denote the indicator that P z is included in C by Iz = 1[z ∈ C]. Then Nr = z∈Sr (x) Iz . For distinct z, z ′ ∈ Sr (x), the events {z ∈ C} P and {z ′ ∈ C} are negatively correlated, so VC [Nr ] ⩽ z∈Sr (x) VC [Iz ] ⩽ µr . By Chebyshev’s inequality, P[Nr < µr /2] ⩽ 4/µr ⩽ e−KΛ . C
(133)
Henceforth condition on the event Nr ≥ µr /2. Each point z ∈ Sr (x) ∩ C has the same transition probability (KtU )⊗d (x | z). Their contribution alone yields Lt (x) =
qtC (x) Nr d r 2 β (1 − βt )d−r ≥ −d M t 2 1 d r 1 ≥ βt (1 − βt )d−r = P[Bin(d, βt ) = r] . 2 r 2
(134)
The choice of r gives 0 ⩽ v − βt ⩽ |βt − βt⋆unif | + DΛ/d + 1/d. Since |t − t⋆unif | ⩽ Lw and βt′ is bounded, |βt − βt⋆unif | ⩽ Cw. Thus D2 Λ2 1 2 + ⩽ C ′Λ . d(v − βt ) ⩽ C dw + d d 2
1 Eq. (38) now lower bounds the right-hand side of Eq. (134) by e−CΛ , the factor 2(d+1) being absorbed using Λ ≥ log d. Together with Eq. (133), this proves the lemma.
B.4.6
Proof of Proposition 6.10
Proof. Recall Algorithm 1: the grid T covers [τ− − 3w, τ+ + 3w] with mesh w/20, each grid point t is tested with R = 2⌈CΛ⌉ + 1 independent samples X (j) ∼ Unif(X ), and the search stops at the first grid 81
point at which the median of Tbt (X (j) ) over these samples falls below −(a + b)dw/2, where a < b are the constants of Proposition 6.9. The search only needs two kinds of decisions to be reliable: it must continue at every grid point above ⋆ tunif +2w, and it must stop by the time it has passed through [t⋆unif −3w, t⋆unif −2w]. Proposition 6.9 controls a single query in exactly these two regions. For a relevant grid point t (t > t⋆unif + 2w or t⋆unif − 3w < t < t⋆unif − 2w), let bC,t be the probability, conditional on C, that one trial lies on the wrong side of the threshold −(a + b)dw/2. Proposition 6.9 with exponent 12 gives EC bC,t ⩽ e−12Λ and hence PC [bC,t > 1/16] ⩽ 16e−12Λ . The grid has O(1/w) ≲ eΛ points for sufficiently large d. A union bound shows that, outside an event over C of probability at most 16e−11Λ ⩽ δ/4, every relevant trial has conditional error probability at most 1/16. Fix such an C. For each grid point the R trials are independent. The probability that at least half are wrong is at most 2R (1/16)R/2 ⩽ 2−R . Taking a union bound over the grid and choosing the constant C in the expression of R sufficiently large makes the probability that there is a grid point where at least half of the R trials are wrong at most ε/8. We may draw the randomness for all grid points in advance, so that stopping early does not change this simultaneous guarantee. On the event that more than half of the R trials are correct at all relevant grid points, the search never stops above t⋆unif + 2w, and the grid mesh ensures that it tests a point in [t⋆unif − 3w, t⋆unif − 2w], where it must stop, if it has not already done so. This proves Eq. (46). There are O(1/w) grid points and O(Λ) queries at each, resulting in the query bound. B.4.7
Proof of Lemma 6.11
Proof. Write u = e−s and v = e−h , so that e−t = uv. Substitution of the two clipping endpoints into Eq. (47) gives (1 − v)(1 − u) (1 − v)(1 + u) ph (ℓt ) = , ph (ℓ−1 . t )= 2(1 + uv) 2(1 − uv) Because s, t ∈ I0 , the factors other than 1 − v are bounded above and below by positive constants. Also 1 − v ≍ h for h ⩽ h0 . Since ph is increasing, every true or clipped approximate score gives a probability in [ch, Ch]. Choose h0 small enough that the upper bound is at most 1/2. Let r and re be the true and clipped approximate score entries for one coordinate, and write p = ph (r) and pb = ph (e r). The slope of ph is at most Ch, so |p − pb| ⩽ Ch|r − re|. Eq. (39) now gives KL(Ber(p) ∥ Ber(b p)) ⩽
(p − pb)2 ⩽ Ch(r − re)2 . pb(1 − pb)
Summing over the coordinates and applying Eq. (36) proves Eq. (49).
B.5
Proofs from Section 7
p Throughout this subsection we use the notation of Section 7: Λ is as in Eq. (54) and w = A Λ/d as in Proposition 7.4; I0 is the interval of Eq. (23); νt is the reference law of Eq. (56) and Lt = qtC /νt the density ratio of Eq. (64); λt = e−t /σt , pu , and ϕd are as in Eqs. (57) and (58); cI , CI are the constants of Eq. (60); Az , Bz are the log averages of Eq. (67), and πr (· | z) and Iz are the cube posterior and posterior information of Eqs. (69) and (70); mt , m0t , and Tt are as in Eqs. (65) and (66), and m e t , set , Tbt are their clipped versions from Eqs. (79) and (80); finally Σt and V are the posterior covariance and its expected trace from Eq. (87).
82
B.5.1
Proof of Lemma 7.5
Proof. Rather than comparing the mixture laws pu and pv directly, we lift them to the joint laws of the latent variable and observation (U, Z). Let Pr denote the joint law U ∼ Unif(X ),
Z = rU + G ,
for r > 0. Conditional on U = y, the laws of Z under Pu and Pv are Gaussians with means uy and vy and common covariance Id . Hence pu (Z | U ) dPu ϕd (Z − uU ) (U, Z) = = . dPv pv (Z | U ) ϕd (Z − vU ) Since ∥y∥22 = d for every y ∈ X , the squared distance between these two means is d(u − v)2 . Therefore, the second moment of the density ratio under Pv is " # Z dPu 2 ϕd (z − uY )2 2 EPv = EY ∼Unif(X ) dz = ed(u−v) . dPv ϕd (z − vY ) Thus, by change of measure and Cauchy–Schwarz, pu (A) = EPv
dPu (U, Z) 1[Z ∈ A] ⩽ dPv
" EPv
dPu dPv
2 #!1/2
2
Pv (Z ∈ A)1/2 = ed(u−v) /2
p pv (A) ,
proving Eq. (71). Since C has the same independent distribution under both laws, adjoining it to the two joint laws does not change the likelihood ratio or its second moment. The same argument therefore proves Eq. (72). B.5.2
Proof of Lemma 7.6
Proof. The coordinates of Z are independent mixtures of Gaussians with bounded means. Thus Eq. (73) holds except with probability e−cd . Since BZ′′ (r) =
d X
Zi2 sech2 (rZi ) ,
i=1
this also gives the upper bound for every r. For the lower bound, count the coordinates for which 1 ⩽ |Zi | ⩽ 2. Each coordinate has a probability bounded away from zero of satisfying this condition, uniformly in t. A Chernoff bound shows that at least cd coordinates satisfy it except with probability e−cd . Each such coordinate contributes a fixed positive amount to BZ′′ (r) throughout the prescribed interval. By Eq. (70), IZ (λt ) is a sum of d independent one-coordinate λt Zi tanh(λt Zi ) − log cosh(λt Zi ) terms. Each lies in [0, log 2], and its mean is Igauss (t). Hoeffding’s inequality gives Eq. (75). Choose CK so that its failure probability is at most e−(K+2)Λ , and then choose cK small enough to absorb the two probabilities e−cd .
83
B.5.3
Proof of Lemma 7.8
Proof. Every coordinate of mt (x) lies in [−1, 1]. Tweedie’s formula and the fact that clipping cannot increase Euclidean distance to this cube therefore give ∥m e t (x) − mt (x)∥2 ⩽ et σt2 ∥b st (x) − st (x)∥2 . Multiply this inequality by e−t σt−2 to obtain the corresponding bound for set − st . Squaring, averaging under qtC , and using the definition of εgauss proves Eq. (81). Eq. (82) follows by Cauchy–Schwarz. B.5.4
Proof of Lemma 7.9
Proof. For u = λt and z = x/σt , Eq. (68) writes the density ratio as Lt (x) =
qtC (x) 1 X u⟨z,y⟩−Bz (u) = e . νt (x) M
(135)
y∈C
Define the summand Lu,z (y) = eu⟨z,y⟩−Bz (u) =
πu (y | z) ϕd (z − uy) = . pu (z) 2−d
Thus Lu,z (y) is the posterior-to-prior ratio of y under the whole-cube prior, and EY ∼Unif(X ) Lu,z (Y ) = 1 . Eq. (135) is therefore the empirical average of these likelihood ratios over the M points of C, sampled without replacement from the cube. We will show that enough moderate likelihood ratios remain in this average to prevent it from being too small. Fix t in the stated range, and write u = λt and Z = X/σt . Since |t − t⋆gauss | ⩽ Lw and t⋆gauss ∈ [τ− , τ+ ], we have t ∈ I0 as soon as Lw ⩽ τ− /2, which holds after decreasing cA,L,K ; we may therefore apply Lemma 7.6 with exponent K + 3. Since |t − t⋆gauss | ⩽ Lw, Eqs. (60) and (75) imply √ √ √ IZ (u) ⩽ d Igauss (t) + CK+3 dΛ ⩽ log M + (CI LA + CK+3 ) dΛ = log M + C0 dΛ , where C0 depends only on A, L, K and [κ− , κ+ ]. Fix an observation z for which this bound and the uniform bounds cd ⩽ Bz′′ (r) ⩽ Cd hold. Until the last step, all probabilities are conditional on this z. Now we migrate to a slightly larger noise level. Set v = u − D0
p Λ/d ,
where D0 will be chosen sufficiently large. Since Iz′ (r) = rBz′′ (r) ≥ cd throughout the relevant parameter √ interval, decreasing the signal strength from u to v decreases the posterior information Iz by at least cD0 dΛ. Thus, choosing D0 sufficiently large compared with C0 gives √ Iz (v) ⩽ log M − 4 dΛ . (136) Decreasing cA,L,K if necessary ensures that v remains in the same fixed parameter interval.
84
Draw Y ∼ πv (· | z). The logarithm of the summand Lu,z (Y ) has mean EY ∼πv (·|z) log Lu,z (Y ) = EY ∼πv (·|z) [u⟨z, Y ⟩ − Bz (u)] = Iz (v) − KL(πv (· | z) ∥ πu (· | z)) ⩽ Iz (v) , (137) and variance V [log Lu,z (Y )] = u2 Bz′′ (v) ⩽ Cd . πv
Thus, under the slightly noisier posterior, the log likelihood ratio is typically well below log M . Fix a constant D > 0, to be chosen later, and define G = y ∈ X : Lu,z (y) ⩽ M e−DΛ . √ Take c√ dΛ. By Eqs. (136) and (137), the cutoff log M − DΛ lies at A,L,K sufficiently small that DΛ ⩽ least 3 dΛ above the mean of log Lu,z (Y ). Chebyshev’s inequality therefore gives πv (G | z) ≥ 1 −
C 1 ≥ Λ 2
for sufficiently large d. We next show that G still has nonnegligible posterior mass at the original signal strength u. Taylor’s theorem and the bound Bz′′ (r) ⩽ Cd imply KL(πv (· | z) ∥ πu (· | z)) = Bz (u) − Bz (v) − (u − v)Bz′ (v) ⩽ Cd(u − v)2 ⩽ CD02 Λ .
(138)
On the other hand, the data processing inequality for KL divergence gives KL(πv (· | z) ∥ πu (· | z)) ≥ KL(Ber(πv (G | z)) ∥ Ber(πu (G | z))) 1 ≥ πv (G | z) log − log 2 πu (G | z) 1 1 ≥ log − log 2 . 2 πu (G | z) Combining this with Eq. (138) gives πu (G | z) ≥ e−C1 Λ , where C1 depends on D0 but not on D. Now, consider only the points of C belonging to G, and define e= 1 L M
X
Lu,z (y) .
y∈C∩G
Its expectation over C is e = EY ∼Unif(X ) [Lu,z (Y ) 1[Y ∈ G]] = πu (G | z) . EC L
85
(139)
Moreover, Lu,z (y) ⩽ M e−DΛ on G. The variance bound for sampling without replacement therefore gives 1 V [Lu,z (Y ) 1[Y ∈ G]] M Y ∼Unif(X ) 1 EY ∼Unif(X ) Lu,z (Y )2 1[Y ∈ G] ⩽ M ⩽ e−DΛ EY ∼Unif(X ) [Lu,z (Y ) 1[Y ∈ G]]
e ⩽ V[L] C
⩽ e−DΛ πu (G | z) . Hence, by Chebyshev’s inequality and Eq. (139), e < πu (G | z)/2] ⩽ P[L C
4e−DΛ ⩽ 4e−(D−C1 )Λ . πu (G | z)
Choose D > C1 + K + 4. Then, except with probability at most e−(K+2)Λ over C, e ≥ πu (G | z)/2 ≥ e−(C1 +1)Λ . Lt (x) ≥ L Finally, add the probability that the conditions imposed on Z fail and enlarge the constant in the statement. This proves Eq. (83). B.5.5
Proof of Proposition 7.11
Proof. Recall Algorithm 3: the grid T covers [τ− − 3w, τ+ + 3w] with mesh w/20, each grid point t is tested with R = 2⌈CΛ⌉ + 1 independent samples X (j) ∼ νt , and the search stops at the first grid point at which the median of Tbt (X (j) ) over these samples falls below −(a + b)dw/2, where a < b are the constants of Proposition 7.10. We first control the choice of C, and then the randomness of the search conditional on C. At a fixed grid point above t⋆gauss + 2w, a wrong decision for one sample means Tbt < −(a + b)dw/2. In [t⋆gauss − 3w, t⋆gauss − 2w], it means the reverse inequality. Proposition 7.10, with exponent 12, bounds each of these probabilities by e−12Λ after averaging over C. At any one relevant grid point, Markov’s inequality shows that the conditional probability of a wrong decision exceeds 1/16 for at most a 16e−12Λ fraction of C. There are O(1/w) ⩽ eΛ grid points. A union bound therefore removes at most a 16e−11Λ ⩽ δ/4 fraction of C, and for every remaining C the conditional error is at most 1/16 at every relevant point. Fix C satisfying the preceding simultaneous guarantee. At one grid point the R trials are independent. The probability that at least half are wrong is at most 2R (1/16)R/2 ⩽ 2−R . Taking R = 2⌈CΛ⌉ + 1 with a sufficiently large constant C makes the union of these events over the grid have probability at most ε/8. Now condition on all relevant medians being correct. The search cannot stop above t⋆gauss +2w. The grid contains a point in [t⋆gauss − 3w, t⋆gauss − 2w], at which it must stop if it has not already done so. This gives Eq. (85). We may draw all samples in advance, so early stopping does not alter this argument. Finally, there are O(1/w) grid points and O(Λ) score oracle queries at each, giving the query complexity bound. B.5.6
Proof of Lemma 7.12
Proof. Let Y ∼ q C and observe Rr = rY + Br ,
86
r ≥ 0,
where B is an independent Brownian motion. For a fixed y, the likelihood of the trajectory (Rv )0⩽v⩽r relative to standard Brownian motion is exp ⟨y, Rr ⟩ − 2r ∥y∥22 . Hence the conditional probability of Y = y given the observations up to time r is proportional to q C (y) exp ⟨y, Rr ⟩ − 2r ∥y∥22 . In particular, this conditional distribution depends on the observations only through Rr . Define mr := E[Y | Rr ],
Γr := Cov(Y | Rr ) .
√ √ Notice that Rr / r has the same law as rY + G, where G ∼ N (0, Id ) independently. Hence the posterior at time r is the same as the Gaussian-diffusion posterior at time t when r = λ2t = e−2t /σt2 . Subtract the conditional drift from the observation and set Z r mv dv . Wr = Rr − 0
In the observation filtration Fr = σ(Rv : 0 ⩽ v ⩽ r), this is a continuous martingale with quadratic variation rId , hence a Brownian motion. Applying Itô’s formula to P[Y = y | Rr ], we have d P[Y = y | Rr ] = P[Y = y | Rr ]⟨y − mr , dWr ⟩ . By definition of mr and Γr , dmr =
X
y d P[Y = y | Rr ] =
y∈C
X
P[Y = y | Rr ]y(y − mr )⊺ dWr = Γr dWr .
(140)
y∈C
Since ∥Y ∥22 = d,
E Tr Γr = d − E∥mr ∥22 .
On the other hand, applying Itô’s formula to ∥mr ∥22 with Eq. (140) gives d∥mr ∥22 = 2⟨mr , Γr dWr ⟩ + Tr(Γ2r ) dr . Taking expectations therefore yields d E∥mr ∥22 = E Tr(Γ2r ) . dr Combining the preceding two identities, d E Tr Γr = −E Tr(Γ2r ) . dr Recall that the Gaussian diffusion posterior corresponds to the preceding observation model with r = λ2t = e−2t /σt2 . Since
d 2 λ = −2e−2t σt−4 , dt t
by the chain rule we get the derivative in Eq. (88). Finally, V (t) = d − E∥mt (Xt )∥22 lies in [0, d].
87
B.5.7
Proof of Lemma 7.13
Proof. Fix C throughout. There are two quantities to estimate: the norm of the score itself and the norm of its derivative. Both have direct posterior interpretations. −2t From Xt = e−t X0 + σt G, we have Cov(G | Xt ) = eσ2 Σt (Xt ). From Tweedie’s formula, we have t
st (Xt ) = −σt−1 E[G | Xt ]. Since d = E∥G∥22 = E∥E[G | Xt ]∥22 + E Tr Cov(G | Xt ), we obtain EqtC ∥st ∥22 = dσt−2 − e−2t σt−4 V (t) ⩽ dσt−2 ⩽ Cd .
(141)
Differentiating the posterior mean gives ∇mt (x) = e−t σt−2 Σt (x), and hence, by Tweedie’s formula, ∇st (x) = −σt−2 Id + e−2t σt−4 Σt (x) . All coefficients are bounded on I0 . Squaring the two terms separately and using Lemma 7.12 yields EqtC ∥∇st ∥2F ⩽ Cd + CEqtC Tr(Σ2t ) ⩽ C(d + V ′ (t)) .
(142)
This is where the increase of posterior variance enters the discretization error. The density qtC solves ∂t qtC = ∆qtC + ∇ · (xqtC ). Differentiating its logarithm gives ∂t st = ∆st + 2(∇st )st + st + (∇st )x . We now apply Itô’s formula to sr−u (Yu ), where dYu = (Yu + 2sr−u (Yu )) du + containing ∆st , (∇st )st , and (∇st )x cancel, leaving √ dsr−u (Yu ) = −sr−u (Yu ) du + 2 ∇sr−u (Yu ) dBu .
√
2 dBu . The terms
Solving this linear equation gives −u
sr−u (Yu ) − sr (Y0 ) = (e
√ Z u −(u−v) − 1)sr (Y0 ) + 2 e ∇sr−v (Yv ) dBv . 0
The stochastic integral has conditional expectation zero given Y0 . Eqs. (141)–(142) and Itô’s isometry therefore imply Z u 2 2 E∥sr−u (Yu ) − sr (Y0 )∥2 ⩽ Cdu + C (d + V ′ (r − v)) dv 0
⩽ Cdu + C(V (r) − V (r − u)),
0 ⩽ u ⩽ h ⩽ 1.
Integrating in u and using the monotonicity of V proves Eq. (89).
C
Details of simulation for Figure 4
Here we describe the numerical simulation that generated Figure 4, which was largely inspired by an experiment carried out for Gaussian diffusion by Xun and Price [XP26]. For a fixed codebook C, we consider a planted codeword y ∗ ∈ C. Starting from X0 = y ∗ , we apply the forward process associated to one of the three dLLM paradigms to obtain noisy sample Xt . The posterior
88
distribution given by conditioning on Xt places some mass on y ∗ , i.e., P[X0 = y ∗ | Xt ] , and panels (a)-(c) plot empirical estimates for the expectation of this quantity, which we call the recovery probability, over the randomness of the forward process and the codebook. In the simulation, we then define the width of the critical window as the distance between the noise levels at which this recovery probability goes from 0.2 to 0.8. To ensure that the noise levels are normalized the same way, we parametrize by the mutually information between a uniformly random bit and its corrupted version. For uniform and Gaussian diffusion, this information is given by Iunif (t) and Igauss (t), and for masked diffusion with m revealed coordinates, it is I = (m/d) log 2. The theory then predicts that the critical noise level occurs at I = κ, where recall that κ = d1 log M . Because the size of the codebook is exponential in the dimension, it is not possible to directly simulate score access to the random empirical measure. Instead, we use the same Poissonization trick used in [XP26]: after planting an arbitrary bitstring y ∗ , we take the remaining “spurious codewords” in the random codebook to be sampled from a Poisson point process on the hypercube with intensity M 2−d per point. We elaborate on how to perform an empirical estimate of the recovery probability under this approximation for each of the three frameworks. Masked diffusion. In this case, for a fixed number of revealed coordinates m, the expected recovery probability is simply the expectation of N1 , where N is the number of codewords (including y ∗ ) which agree with y ∗ on the revealed subset of coordinates. Under Poissonization, N − 1 is a draw from Poi(M 2−m ). Note that the value of m at which the expectation of N1 crosses from 0.2 to 0.8 is clearly independent of d in this setup, so under the normalization by I = (m/d) log 2, the width of the window in the simulation is Θ(1/d) as the theory predicts. Uniform diffusion.
In this case, by Lemma 2.3, P[X0 = y ∗ | Xt ] =
1 1+
Pd
j−dH (Xt ,y ∗ ) j=0 Nj tanh(t/2)
,
where Nj is the number of spurious codewords at Hamming distance j from Xt . Note that dH (Xt , y ∗ ) is distributed as Bin(d, βt ), whereas conditioned on Xt , each quantity Nj is independently distributed as d −d Poi(M j 2 ). This enables one to empirically estimate the recovery probability, where each sample in the empirical estimate only requires d Poisson samples and one binomial sample. Gaussian diffusion.
In this case,
P[X0 = y ∗ | Xt ] =
1 1+Z
for
X
Z=
y∈C\{y ∗ }
KtG (Xt | y) . KtG (Xt | y ∗ )
For spurious codeword y which differs from planted y ∗ on coordinates T ⊆ [d], KtG (Xt | y) 2e−t X ai = exp − 1 − e−2t KtG (Xt | y ∗ ) i∈T 89
ai = yi∗ Xt,i .
Note that each ai is independently distributed according to N (e−t , 1−e−2t ), and conditional on these draws, P T is a uniformly random subset of [d]. Under Poissonization, the sum i∈T ai is thus a Poisson process P with intensity M pa (s) ds, where pa is the density of the random variable di=1 bi ai for b ∼ Unif({0, 1}d ). This density is estimated using the saddlepoint method for approximating the pdf of a sum of independent random variables, and the Poisson point process is approximated by discretizing the approximated density into bins.
90