The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity

The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity Martin J. Wainwright mjwain@mit.edu Lab for Information and Decision Systems Statistics and Data Science Center EECS and Mathematics, Massachusetts Institute of Technology August 14, 2026

arXiv:2608.13520v1 [cs.LG] 13 Aug 2026

Abstract We study masking diffusion for discrete sampling and introduce a path-resolved measure of data geometry called the unmasking growth complexity (UGC). Its local increments directly control Kullback– Leibler (KL) discretization error, yielding a unified analysis of Bernoulli-subset and fixed-cardinality unmasking schemes. In log-reveal-odds coordinates, this structure yields optimized single-block and multi-block schedules, and quantifies the gains from adapting computational effort to data geometry. Crucially, we show how UGC increments can be estimated from samples via KL increments along coupled reveal trajectories. This leads to certified-optimal samplers that achieve a prescribed KL error with high probability and iteration complexity within a constant factor of the corresponding oracle procedure. Collapsing the UGC path yields the aggregate UGC mass, which connects to classical multivariate dependence measures, and complexity measures from previous analyses of discrete diffusion. In the finepartition limit, the squared integral of the square-root UGC density determines the sharp leading-order optimal Euler discretization error. Examples exhibit substantial dimension-dependent gains over coarse √ e d) improvements achievable with a constant number of adaptively placed blocks. schedules, including Ω(

1

Introduction

The problem of sampling from a high-dimensional distribution is fundamental in nature, and has a wide range of applications. Efficient sampling algorithms underpin the utility of Monte Carlo approximation [RC04, RK08]; support uncertainty quantification in Bayesian models [GCS+ 13, BGJM11]; and lie at the heart of generative AI [RBD+ 22, CHIS23, YZSea25, CMFW24]. Recent years have witnessed tremendous practical and theoretical advances in the use of diffusion1 sampling algorithms, which generate samples via a sequence of “denoising” operations [SWMG15, SE19, HJA20, SSDK+ 21, CCL+ 23, CLL23, BDBDD24]. Initial work in the area focused on sampling from continuous distributions in Rd , in which case the denoising step corresponds to estimating a signal embedded in Gaussian noise [HJA20, SSDK+ 21]. A more recent line of work, and the general focus of this paper, has focused on sampling from discrete distributions, using various kinds of denoising processes [HNJ+ 21, AJH+ 21, CBB+ 22, LME24, SHW+ 24, SAS+ 24, CCL25, DHW26]. In this paper, we study the problem of sampling a random vector Z ∈ (A)d , where A is a discrete alphabet, and d is the ambient dimension. While various samplers have been devised, we focus on unmasking samplers, which traverse a path from a fully unobserved vector, denoted by Z ∗ = (⋆, . . . , ⋆) with ⋆ meaning masked or unobserved, back to a sample Z ∼ PZ from the target distribution. The path is traversed via a sequence of unmasking operations, in which a subset of the masked coordinates are revealed. Unmasking algorithms differ in the form and sequence by which these unmasking steps take place. Unmasking and closely related 1 In full generality, the diffusion terminology is a misnomer, since not all such models involve a diffusion process in the formal probabilistic sense; this includes the unmasking models studied in this paper. Nonetheless, we adopt the conventional terminology.

1

mask-based generative samplers have proven effective across a range of applications, including machine translation [GLLZ19], image and video synthesis [CZJ+ 22, YCS+ 23], text-to-image generation [CZB+ 23], speech synthesis [WZL+ 25], language modeling [SAS+ 24, NZY+ 25], and protein design [WZY+ 24].

1.1

Overview

This paper is motivated by two broad questions associated with diffusion samplers, as articulated in our companion paper on Gaussian diffusion [Wai26], which we paraphrase here: Q1: Can the performance of masked diffusion sampling be explained and quantified, in some generality, by a measure tied to data geometry? Q2: Is it possible to exploit a masked-diffusion measure of data geometry to design, optimize and certify practical sampling schemes? Recent work has made substantial progress on both questions. Chen et al. [CCL25] analyze fixed-cardinality unmasking and derive an exact information-profile representation of its expected KL discretization error, with consequences for schedule design and bounds based on total and dual total correlation. Lavenant and Zanella [LZ25] derive an information-profile representation of the factorization error for random-order unmasking, and study optimal schedules in asymptotic scaling regimes. From a different direction, Dmitriev et al. [DHW26] analyze discrete diffusion via a continuous-time Markov chain (CTMC) formulation; for masking diffusion, their modified τ -leaping guarantees are governed by an effective total correlation that can be substantially smaller than the classical worst-case measures. Collectively, these results provide important answers to Q1 and partial answers to Q2, but fall short of finite-sample guarantees for learning and certifying geometry-adaptive schedules from data. Our answer to Q1 is based on the unmasking growth complexity or UGC for short. It is a path-resolved measure of data dependence along the reveal process. We show that its local increments directly control KL discretization error, thereby obtaining a common analysis of both fixed-cardinality unmasking and a natural Bernoulli-subset variant. When collapsed over the full reveal path, the UGC complexity reduces to a coarse measure that closely connected to classical multivariate dependence measures [Wat60, Han75, Han78, TSE94], as well as to the effective total correlation appearing in the CTMC analysis [DHW26]. Consequently, the resulting single-block guarantees sharpen existing fixed-cardinality bounds while giving Bernoulli unmasking guarantees comparable to those available for CTMC samplers. The main contributions of our paper are in the context of Q2, and in particular, our use of the full UGC complexity path to design and analyze sampling algorithms that are certified-optimal. More precisely, we show that UGC complexity has a natural additive structure along the reveal path, so that its local increments can be estimated from samples and used to allocate computational effort where the target distribution is most difficult to unmask. We exploit this structure to derive and optimize blockwise schedules, construct data-dependent samplers with high-probability KL certificates, and finally, to optimize the block boundaries themselves. In the fine-partition limit, the log-reveal-odds UGC density emerges as the intrinsic local geometry: its square-root integral governs the limiting partition complexity and the sharp leading-order optimal Euler discretization error.

1.2

Related work

So as to put our results in context, we now provide a broader discussion of related work. There are various types of discrete diffusion models, all of which replace the additive Gaussian noise for continuous-space diffusions by stochastic corruption kernels on discrete state spaces. Early work developed multinomial and more general structured discrete diffusion processes [HNJ+ 21, AJH+ 21], among them the absorbing-state corruption that underlies a masking diffusion process. Other work [CBB+ 22] formulated discrete diffusion models using the formalism of continuous-time Markov chains, with subsequent work developing discrete analogues of score-based modeling in which the reverse CTMC is parameterized through learned probability ratios [LME24]. Masked diffusion is a particularly important subclass, in which coordinates are corrupted

2

by replacement with a distinguished mask symbol and generation proceeds by progressively reconstructing masked coordinates [SHW+ 24, SAS+ 24]. As described above, recent work has made progress on characterizing the accuracy–parallelism trade-off in random-order unmasking with fixed-cardinality subsets. In particular, Li and Cai [LC25] established information-theoretic convergence guarantees for parallel masked-diffusion sampling. Chen et al. [CCL25] derived an exact characterization of the expected KL divergence in terms of a one-dimensional information profile; they also gave explicit sampling guarantees involving classical multivariate dependence measures [Wat60, Han75, Han78], namely the total correlation (TC) and dual total correlation (DTC). A portion of our analysis exploits their exact KL representation. Lavenant and Zanella [LZ25] derived an information-profile representation for the factorization error of random-order unmasking, and then studied optimal scheduling via a continuum scaling limit. We compare their asymptotic schedule analysis more closely with our fine-partition and Euler results following Theorem 3. The τ -leaping method for continuous-time Markov chains (CTMC) originates in stochastic chemical kinetics, where its consistency and approximation errors have been studied extensively [Gil01, RPCG05, Li07, AGK11]. For discrete diffusion, recent work has developed convergence guarantees for τ -leaping under both standard and absorbing corruption processes [RCRY25, LHL+ 25, LLLS25, DHW26], as well as higherorder variants of CTMC discretization [RCZ+ 25]. Among CTMC analyses, most related to our work is the paper of Dmitriev et al. [DHW26], who analyze a modified τ -leaping sampler for a general class of CTMC discrete diffusion models. For masking diffusion, they give guarantees in terms of a measure that they call effective total correlation, which turns out to be closely related to the coarse UGC complexity that we introduce. In contrast to CTMC-based analyses, we study Bernoulli and fixed-cardinality unmasking; our coarse UGC bounds yield matching guarantees for these direct schemes. The coarse UGC complexity is also connected to several multivariate dependence measures, including TC, DTC, and the finer-grained measure of Tononi et al. [TSE94]. We develop these connections in Section A. Finally, in our companion paper [Wai26] on Gaussian diffusion sampling, we introduced the denoising growth complexity (DGC), a pathwise complexity that controls KL discretization error, and whose log-scale DGC density specifies optimal sampling schedules. Despite the substantial differences between Gaussian and masking diffusion, the two theories exhibit a remarkable degree of parallelism, suggesting a common underlying principle linking denoising growth, KL discretization error, and optimal scheduling.

2

Overview: unmasking growth complexity and sampling

So as to orient the reader, we begin by providing a high-level overview of the main ideas and quantities underlying our analysis. We first describe the canonical unmasking process. Using it, we define the unmasking growth complexity (UGC), which measures how informational difficulty is distributed along the reveal path. We then pass to log-reveal-odds coordinates, where this complexity is represented by a density whose geometry determines how sampling effort should be allocated along the path. Our goal here is primarily conceptual: to explain the distinction between coarse and geometry-aware schedules, to illustrate the data geometry of the UGC density on several examples, and to preview how data-driven adaptive partitioning can exploit this geometry. Precise forms of our guarantees for samplers, estimation procedures, and supporting technical results are developed in the subsequent sections.

2.1

Unmasking reveal process

Given a discrete alphabet A, our goal is to sample a d-dimensional random vector Z ∈ Ad with distribution PZ . Underlying the sampling algorithms that we analyze is an unmasking stochastic process {Xt , t ∈ [0, 1]},
d where Xt ∈ A ∪ {⋆} and the new symbol ⋆ denotes a masked or unobserved entry. At time t = 0, we have X0 = (⋆, . . . , ⋆) almost surely, whereas at time t = 1, we have X1 ∼ PZ , where PZ is the target distribution from which we would like to sample. The evolution of Xt along the path is controlled by a sequence {Ui }di=1 of i.i.d. Unif[0, 1] random variables,

3

independent of Z. These variables define the reveal subset sequence, indexed by t ∈ [0, 1], via  Rt := i ∈ [d] | Ui ≤ t , and its complement Rtc := [d] \ Rt .

(1a)

The subset Rt corresponds to the subset of coordinates i for which the hidden value Zi has been revealed by time t. More formally, the unmasking process at time t is given by  
d ⋆ for t ∈ [0, 1], (1b) ∈ A ∪ {⋆} Xt := ZRt , ZR c t ⋆ where ZRt = (Zi , i ∈ Rt ) are the variables revealed by time t, and ZR c = (⋆, . . . , ⋆) is a sub-vector of t c missing values in positions indexed by Rt . By construction, the unmasking process (1b) defines a family of probability distributions {Pt , t ∈ [0, 1]}, where P0 denotes a degenerate distribution with all its mass on the masked sequence, whereas P1 ≡ PZ is the probability distribution of the target variable Z.

2.2

Unmasking growth complexity and its geometry

For a reveal time t ∈ [0, 1], we define the Bernoulli unmasking gain h(t) :=

d X

Info Zi ; Xt | i ∈ M(Xt ) ,

(2a)

i=1

where M(Xt ) ⊆ {1, . . . , d} is the subset of indices that are masked at reveal time t, and Info denotes the (conditional) mutual information. The function h has a denoising interpretation, since when i ∈ M(Xt ), all other coordinates in Xt are revealed independently with probability t; the conditional mutual information term for the ith coordinate measures how much these revealed variables reduce uncertainty about Zi . Moreover, the derivative h′ turns out to have a simple information-theoretic representation: in particular, it corresponds to the second derivative h′ (t) = −

d2 Info(Z; Xt ), dt2

(2b)

where Info(Z; Xt ) denotes the mutual information between the target vector Z ∈ Ad and the partially
d unmasked vector Xt ∈ A ∪ {⋆} at reveal time t ∈ (0, 1). See equation (51a) in the proof of Theorem 1 for the underlying details. In terms of this mutual information derivative, the unmasking growth complexity assigns a non-negative number to any sub-interval [p, q] of the unit interval [0, 1] via Z q H(p, q) := t(1 − t)h′ (t) dt for any 0 ≤ p < q ≤ 1. (2c) p

Based on the identity (2b), we see that H(p, q) is a weighted integral of the information curvature. The essential feature of this complexity measure is that it is a path-resolved quantity: rather than assigning a single complexity to the target distribution PZ , it assigns a complexity to every interval of the reveal path. Our analysis shows how these local increments directly control the associated KL discretization error. Moreover, for any triple 0 ≤ p < q < r ≤ 1, we have the additivity property H(p, r) = H(p, q) + H(q, r).

(3)

This additivity allows the global sampling problem to be decomposed into local pieces, whose complexities can be estimated and controlled separately, and is what ultimately enables geometry-adaptive refinement of the sampling schedule. By collapsing the path-resolved complexity to the full interval [0, 1], we obtain the aggregate UGC mass R1 given by H(0, 1) = 0 t(1−t)h′ (t)dt. Interestingly, this aggregate quantity has several connections to classical 4

measures of multivariate dependence. By a non-trivial argument (see the proof of Proposition 3), it turns out to be equivalent to a multivariate dependence measure, first introduced by Tononi, Sporns and Edelman [TSE94] in 1994, which spawned a rich line of work (e.g., [BZ12a, BZ12b, OBAJ08, BBB09, RMGJ19, VPFS23]). A separate argument shows its close relation to the effective total correlation introduced by Dmitriev et al. [DHW26] in their study of CTMC unmasking algorithms. Moreover, the value H(0, 1) can be upper bounded by classical measures of multivariate dependence [Wat60, Han75, Han78], known as total correlation and dual total correlation. We elaborate upon these and other connections in Section A. Role of the log-reveal-odds density: The natural coordinate for progress along the unmasking path is not reveal time t nor its logarithm, but rather the log-reveal-odds λ = φ(t) := log(t/(1 − t)). Our analysis gives one-step KL bounds governed by the multiplicative change in reveal odds; consequently, equal increments in λ correspond to equal multiplicative changes in reveal odds, and place the early and late stages of unmasking on a common scale. Expressing the UGC complexity in this coordinate yields a density q that localizes the informational difficulty of sampling along the path. As our results show, regions where q(λ) is large require finer resolution. More precisely, the log-reveal-odds UGC density function is given by q(λ) := r2 (1 − r)2 h′ (r)

λ

e where r = φ−1 (λ) := 1+e λ.

By construction, the UGC increment is given by H(p, q) =

R φ(q) φ(p)

(4)

q(λ) dλ.

The high-level conclusions  of this paper take a particularly simple form in terms of q. We focus on the canonical reveal interval d1 , 1 − d1 , which corresponds under the log-reveal-odds transformation to the symmetric interval [−ℓd , ℓd ] with ℓd := log(d − 1). • Single-block unmasking schemes use a single multiplier across the reveal path and hence ignore the location of the UGC mass. Their iteration complexity is governed by the coarse or aggregate UGC complexity Z ℓd CUGC := 2ℓd q(λ) dλ. (5a) −ℓd

The measure CUGC depends only on the total UGC mass over the relevant portion of the path, and not on how this mass is distributed. This coarse measure connects directly to past work: it enables us to show that a single-block Bernoulli unmasking sampler matches the guarantees given for CTMC unmasking [DHW26], while sharpening guarantees from past work on fixed-cardinality sampling [CCL25]. See Corollary 1 and Proposition 3 for details. • Geometry-aware schemes exploit the local UGC geometry, taking finer steps where q is large and coarser steps where it is small. Crucially, we show that UGC increments can be estimated from samples, and use these estimates to choose the partition adaptively from data (see Proposition 2). Under progressive refinement, the resulting complexity converges from above to the fine-partition complexity !2 Z ℓd p PUGC := q(λ)dλ , (5b) −ℓd

We establish non-asymptotic convergence rates and high-probability finite-sample guarantees for the resulting samplers; see Theorems 2 and 3 for details. A related square-root structure for an information profile appears in the asymptotic analysis of fixed-cardinality schedules by Lavenant and Zanella [LZ25]; we discuss this connection in more detail following Theorem 3. In summary, the potential gain from data-dependent optimization of the sampling schedule is governed by the ratio Rℓ 2ℓd −ℓd d q(λ) dλ (i) CUGC = R (6) Ratio(PZ ) := 2 ≥ 1, ℓd p PUGC q(λ)dλ −ℓd 5

where the lower bound (i) follows from the Cauchy–Schwarz inequality. Equality holds if and only if q is constant, whereas larger values of Ratio(PZ ) arise when the UGC mass is unevenly distributed or sharply concentrated along the path. 2.2.1

Geometric behavior of UGC density

Since the UGC density q captures the essential structure of the problem, it is useful to compute and plot it for three different ensembles, each chosen to illustrate a qualitative aspect of our theory. Noisy repeated bit: We begin with a very simple example. For any p ∈ [0, 1], we use V ∼ Ber(p) to denote a Bernoulli random variable with P(V = 1) = p and P(V = 0) = 1 − p. In this ensemble, we construct a binary random vector Z = (Z1 , . . . , Zd ) ∈ {0, 1}d by first drawing U ∼ Ber(1/2), and then setting Zi = U ⊕ W i

for each i = 1, . . . , d,

where each Wi is an independent Ber(η)-variable for some η ∈ [0, 1/2], and ⊕ denotes addition modulo two. By construction, we have P(Zi = U ) = 1 − η, hence our use of the term “noisy repeated bit”. Noisy bit ( = 0.01)

0.2 0.1 6

4

2

0

Log reveal odds

(a)

2

4

0.3 0.2 0.1 0.0

6

4

2

0

Log reveal odds

(b)

Noisy bit ( = 0.45)

0.4

UGC density q( )

0.3

0.0

Noisy bit ( = 0.30)

0.4

UGC density q( )

UGC density q( )

0.4

2

4

0.3 0.2 0.1 0.0

6

4

2

0

Log reveal odds

2

4

(c)

Figure 1. Plots of the log-reveal-odds UGC density q for the noisy repeated-bit ensemble for η ∈ {0.01, 0.30, 0.45}. The density has a sharp peak for η = 0.01, and then flattens out and shifts to the right as η → 0.5. Up to two digits of accuracy, we have Ratio(PZ ) ∈ {4.51, 2.16, 1.65} in panels (a), (b), and (c), respectively. For this example, we have Ratio(PZ ) ≍ log(d) as the dimension grows.

Figure 1 gives plots of the UGC density (4) in dimension d = 128, and flip probabilities η ∈ {0.01, 0.30, 0.45}. For η = 0.01, the density has a sharp peak at a low value of the log-reveal-odds parameter λ. As η increases, the dependence among the coordinates of the multivariate random vector Z ∈ {0, 1}d decreases; it takes a higher reveal level before information about the hidden bit U is revealed, as reflected by the rightward shift in the mode of the UGC density. In parallel with this rightward shift, the entire density collapses towards zero, so that H(0, 1) → 0 as η → 1/2. Discrete mixture models: Our second ensemble corresponds to the discrete analog of a mixture model. More precisely, given a fixed set of M cluster centers, say C 1 , . . . , C M ∈ {0, 1}d , we generate a binary random vector Z = (Z1 , . . . , Zd ) ∈ {0, 1}d with components Zi = CiJ ⊕ ξi for i = 1, . . . , d, where the random cluster index J is drawn uniformly at random from {1, . . . , M } and the ξi ∼ Ber(η) are drawn i.i.d., and independently of the cluster index. This procedure generates an M -component mixture distribution with uniform weight 1/M on each component. We studied this family for varying dimensions d divisible by four, number of mixture components M = 2d/4 , and cluster centers sampled independently and uniformly from the Boolean hypercube {0, 1}d . Figure 2 shows a sequence of UGC densities for this family, with the dimension d and the number M = 2d/4 of mixture

6

9

6 3 0

2 1 0 Log reveal odds (a)

1

Discrete mixture (d = 48)

9 UGC density q( )

Discrete mixture (d = 32)

UGC density q( )

UGC density q( )

9

6 3 0

2 1 0 Log reveal odds (b)

1

Discrete mixture (d = 64)

6 3 0

2 1 0 Log reveal odds

1

(c)

Figure 2. Plots of the log-reveal-odds UGC density q for the discrete mixture model. For each dimension d, the latent vector Z ∈ {0, 1}d is generated by selecting uniformly from M = 2d/4 cluster centers, and then flipping each coordinate independently with probability η = 0.02. Panels (a)–(c) correspond to d ∈ {32, 48, 64}, and hence M ∈ {256, 4096, 65536}, respectively. Increasing the dimension produces a progressively sharper and taller spike. Up to two digits of accuracy, we have Ratio(PZ ) ∈ {2.19, √ 3.13, 3.85} in panels (a), (b), and (c), e d) as the dimension grows. respectively. For this example, it can be shown that Ratio(PZ ) = Θ(

components increasing as we move from left to right. The density is unimodal in all cases; as the dimension increases, the peak √ location remains fixed but its height increases. It can be shown that the ratio (6) scales e as Ratio(PZ ) = Θ( d), so that there are significant gains from using geometry-aware stepsizes. Moreover, since the UGC density is concentrated around a single increasingly narrow transition window, a partition with only K = 3 blocks—one covering this transition and two covering its tails—achieves the fine-partition complexity up to logarithmic factors. See Section 2.2.2 and Figure 2 for further discussion of these factors. Hierarchical mixture model: The previous examples, being quite simple in nature, all led to unimodal UGC densities. We now consider a hierarchical model in which multiple peaks emerge, and the peaks have interesting meanings in terms of the underlying data geometry. For an integer L ≥ 2, we construct a model that places mass on exactly 2L prototypes in the Boolean hypercube {0, 1}d , say {V 1 , . . . , V M }, where M = 2L . Each prototype is defined by a path down a binary tree of depth L; see the left column of Figure 3 for some examples. The geometry of the prototypes is related to the tree structure in the following way: for any tree level ℓ ∈ {1, . . . , L}, a pair of prototypes V a and V b whose corresponding tree paths first diverge at level ℓ of the tree are separated PL by Hamming distance ∆ℓ , for some pre-specified set of distances {∆ℓ }L , and total dimension d = ℓ=1 ℓ=1 ∆ℓ . Panel (a) of Figure 3 shows a toy example with L = 2, (∆1 , ∆2 ) = (12, 2), and dimension d = 12 + 2 = 14. It has M = 2L = 4 prototype vectors in {0, 1}14 , assigned the index labels {00, 01, 10, 11}, plotted as rows in the corresponding black-white matrix. Consider the pair of prototypes indexed by 00 and 11 respectively; they diverge at level ℓ = 1 of the tree, and their Hamming distance is equal to ∆1 = 12. Similarly, the prototypes indexed by 00 and 01 diverge at level ℓ = 2 of the tree, and have Hamming distance ∆2 = 2. Panel (c) illustrates a larger construction with L = 4, still in a toy dimension d = 34 for illustrative purposes. Panels (b) and (d) show the UGC densities for the L = 2 and L = 4 models, using larger separations and dimensions: the L = 2 model has dimension d = 32776 and (∆1 , ∆2 ) = (32768, 8), whereas the L = 4 model has dimension d = 34952 and (∆1 , ∆2 , ∆3 , ∆4 ) = (32768, 2048, 128, 8). Note how the UGC density for the L = 2 model has two distinct peaks in log-reveal odds. The level ℓ = 1 ambiguity is resolved at an earlier reveal time, since the corresponding prototypes differ by Hamming distance ∆1 = 32768; the second peak corresponds to the later reveal time at which the level ℓ = 2 ambiguity, corresponding to Hamming separation ∆2 = 8, is resolved. Similar comments apply to the L = 4 peaks for the UGC density in panel (d).

7

0.6

Prototype hierarchy (L = 2)

root

1

2

00 01 10 11

UGC density q( )

0.5

1

24

0.4 0.3 0.2 0.1 0.0

Schematic Hamming scales: 1 = 12,

2=2

(a) 0.6 1

2

3

4

0000 0001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111

8

4

Log reveal odds

0

(b) Prototype hierarchy (L = 4)

0.5

UGC density q( )

root

12

1

40

Schematic Hamming scales: 1 = 20,

2 = 8,

3 = 4,

4=2

(c)

0.4 0.3 0.2 0.1 0.0

12

8

4

Log reveal odds

0

(d)

Figure 3. Hierarchical binary-prototype mixtures and their log-reveal-odds UGC densities. The rows correspond to hierarchy depths L = 2 and L = 4. Left column panels (a) and (c) show the leaf prototypes, while right column panels (b) and (d) show the corresponding UGC densities q.

2.2.2

From coarse to fine geometry

Let us describe how we exploit finer-grained UGC geometry for algorithmic purposes. A single-block scheme uses one geometric multiplier across the full log-reveal-odds interval, and its performance is governed by the coarse complexity CUGC . In order to exploit local UGC geometry, we instead partition the interval [−ℓd , ℓd ] into K blocks and allow a different geometric multiplier on each block. For a partition P = {[λk , λk+1 ]}K−1 k=0 , Rλ define the block lengths Sk = λk+1 − λk and the associated UGC increments Hk = λkk+1 q(λ) dλ. Optimizing the geometric multipliers over these blocks leads to the partition complexity C(P) :=

K−1 Xp

Sk Hk

!2 .

(7a)

k=0

Proposition 1 shows that this quantity governs the iteration complexity of the resulting K-block sampler. Although the optimal multipliers depend on the unknown increments Hk , these increments can be estimated from clean samples. We use these estimates to choose the partition and its multipliers from data, leading to the certified guarantees; see Proposition 2 and Theorem 2. Progressive refinement interpolates between the coarse single-block complexity and the fine-partition limit (5b). In particular, in Theorem 3, we prove that inf C(P) = P

 Z ℓd p

2 q(λ)dλ

where the infimum is over partitions of [−ℓd , ℓd ].

(7b)

−ℓd

Thus, increasing the number of blocks allows the sampler to exploit increasingly fine features of the local √ UGC geometry, and ultimately achieve the fine-grained complexity based on q.

8

3

q( ) Single block

Square-root UGC density

Square-root UGC density

Illustrative example: Figure 4 illustrates this coarse-to-fine interpolation for the discrete-mixture model underlying Figure 2. For d = 64, the single-block scheme has complexity CUGC ≈ 65.0. A three-block partition that allocates finer resolution around the narrow high-density region reduces the partition complexity to approximately 23.5, compared with the fine-partition value PUGC ≈ 16.9.

2 1 area2 = CUGC

0

4

2 0 2 Log reveal odds

4

3

q( ) K = 3 partition

2 1 area2 = C( )

0

4

(a)

2 0 2 Log reveal odds

4

(b)

Figure 4. Block geometry for the d = 64 random discrete mixture with M = 2d/4 and η = 0.02. Both √ panels show the square-root UGC density q versus the log-reveal-odds λ. Panel (a) uses a single geometric multiplier, yielding CUGC ≈ 65.0. Panel (b) uses the illustrated K = 3 block partition, yielding C(P) ≈ 23.5. The fine-partition complexity for this example is PUGC ≈ 16.9.

The three-block scheme devotes more iterations to the narrow high-UGC region and fewer to the flatter tails. Even this coarse adaptation captures most of the available improvement: the complexity drops from approximately 65.0 to 23.5, substantially closing the gap to the fine-partition value 16.9. This example illustrates the basic role of adaptive partitioning: computation is allocated according to where the UGC mass is concentrated, rather than uniformly across the reveal path.

3

Unmasking samplers and KL discretization error

We now turn to the description and analysis of some simple unmasking samplers. At a high level, they are based on discretized approximations to the unmasking reveal process {Xt , t ∈ [0, 1]} from equation (1b). Recalling that M(x) denotes the set of masked coordinates in the vector x, their updates involve the singlesite posterior distributions
µi (a, x) := P Zi = a | Zj = xj for all j ∈ / M(x) , defined for each i ∈ M(x), (8)
d where a ∈ A and x ∈ A ∪ {⋆} . We use µ to denote the full collection of these distributions, and often refer to it as the denoiser. It summarizes residual uncertainty in Z given the observation x. In practice, the denoiser µ is estimated based on samples from PZ as follows. Given a clean sample Z ∼ PZ , we can generate training examples by independently masking random subsets of coordinates. We then predict the original value of each masked coordinate from the resulting partially observed sample using cross-entropy loss. With this set-up, the population-optimal predictors are given by the posterior distributions µi (·, x), and the training procedure generates estimates µ bi for each i = 1, . . . , d. For simplicity, we describe the algorithms using the exact denoiser µ. Using a learned denoiser µ b contributes an additional term to the KL error bound in a standard way; see equation (18) following the statement of Theorem 1 for details.

3.1

Bernoulli and fixed-cardinality unmasking

We now describe the two standard unmasking samplers that we analyze in this paper. At a high level, they involve two basic operations: choosing a subset of indices and then, conditional on this subset, making a 9

random choice of values for the revealed variables using the denoiser µ. The samplers that we analyze differ only in how the random subset is chosen: the Bernoulli sampler uses coin flips, thereby generating a set with a random cardinality, whereas the fixed-cardinality sampler chooses a random subset with fixed cardinality. Both samplers are used in practice, and our theory gives a unified analysis of both in terms of UGC functions. Given a subset A ⊆ [d] of coordinates, both samplers make use of the conditional distribution  !   O  O O δXi ⊗ Pumask (X ′ | X, A) := µi (·, X) ⊗  δ⋆  , for A ⊆ M(X), (9a) i∈A

i∈M(X) /

i∈M(X)\A

which fixes all observed variables (i.e., indices i ∈ / M(X)); fixes all missing variables not in A; and makes a random update to variables in A. Moreover, for an iteration count N ≥ 1, both samplers are based on a reveal-time grid 0 < t0 < t1 < · · · < tN < 1.

(9b)

b 0 ∼ Qt , and then generate a sequence {X b j }N moving The samplers are initialized with a random choice X 0 j=0 forward in time across the grid (9b). 3.1.1

Bernoulli unmasking sampler

We begin by describing the Bernoulli unmasking sampler. It chooses subsets according to the binomial distribution PBer (A | X, β) := β card(A) (1 − β)card(M(X)\A)

with support A ⊆ M(X),

(10)

and then updates entries according to the conditional distribution (9a). More precisely, given the initializab 0 and reveal-time grid (9b), it generates the sequence tion X Draw random binomial subset: Random unmasking:

b j , βj , Aj ∼ PBer · | X

where βj :=

tj+1 − tj , 1 − tj

b j+1 ∼ Pumask · | X b j , Aj , X

(11a) (11b)

for iterations j = 0, 1, . . . , N − 1. 3.1.2

Fixed-cardinality unmasking sampler

The fixed-cardinality sampler is very similar, except that it replaces the binomial subset choice (10) with a black-box that draws a uniformly random subset of a specified size. More precisely, the fixed-cardinality sampler only admits iteration counts N < d − 1, and it requires that each reveal time be d-aligned, meaning b 0 with exactly d − s0 masked that tj = sj /d for integers 1 ≤ s0 < · · · < sN ≤ d − 1. It initializes X j b coordinates. Then, at iteration j = 0, . . . , N − 1, the current state X has exactly d − sj masked coordinates. Draw fixed-cardinality subset: Random unmasking:

n o b j ) card(A) = sj+1 − sj . Aj ∼ Unif A ⊆ M(X
b j+1 ∼ Pumask · | X b j , Aj , X

(12a) (12b)

For future reference, we note that underlying the fixed-cardinality sampler is a de-Poissonized version of the unmasking process {Xt , t ∈ [0, 1]}. It is defined only at d-aligned times t = s/d, and the associated variable Xt ∼ Pt

has exactly s revealed coordinates,

(13)

for a subset of cardinality s chosen uniformly at random. Note that we are overloading our notation here, but the specific version of Pt being used will be clear from context. 10

3.1.3

Initialization and completion steps

Both samplers involve initialization and completion steps. In our analysis, the initialization cost is measured by the KL divergence b 0 ∼ Qt . where Xt0 ∼ Pt0 and X 0

DKL (Pt0 ∥Qt0 )

(14a)

Here Xt0 is generated by the standard reveal process for the Bernoulli sampler (1b), and from its dePoissonized analogue (13) for the fixed-cardinality sampler. Moreover, both samplers involve a final completion step, based on a completion kernel CT (· | x) that maps each partially revealed terminal state
d x ∈ A ∪ {⋆} to a probability law on fully specified vectors in Ad . We denote the resulting completed output by Zb and its law by PZb . The defect of any completion kernel is measured by ΓC (T ) := EXT [DKL (L(Z | XT ) ∥ CT (· | XT ))] .

(14b)

While these terms appear in our bounds, they are not dominant terms. For example, the following completion kernel C has defect equal to zero, and uses only card(M(XT )) serial updates. Fix an ordering of the indices in M(XT ), and then sample from the conditional distributions µi in sequence to fill in the missing entries. With the endpoint T = d−1 d , we have card(M(XT )) = 1 for the exact cardinality sampler and E[card(M(XT ))] = 1 for the Bernoulli sampler, so that the computational cost of implementing this completion kernel is minimal. Similarly, for initialization, it is easy to initialize the fixed-cardinality sampler with exactly one unmasked entry, so that the initialization cost vanishes for t0 = 1/d. The bounds below are stated for a general interval [t0 , T ] and arbitrary initialization-completion choices, but we later specialize to the interval d1 , 1 − d1 , which we refer to as the canonical reveal time interval.

3.2

Unified analysis

We are now set up to state a single theorem that provides a unified analysis of both the Bernoulli and fixedcardinality unmasking samplers. Our original definition (2a) of the Bernoulli denoising gain h emphasizes the connection to the reveal process. Here, in order to make transparent the connections with fixed-cardinality sampling, we make use of the following equivalent representation: h(t) =

d X

  EAi,t Info(Zi ; ZAi,t ) ,

(15a)

i=1

where Ai,t ⊆ [d] \ {i} is a random subset obtained by including each coordinate of [d] \ {i} independently with probability t. See Section B for the proof of this alternative representation. The analogue for the fixed-cardinality sampler is given by the fixed-cardinality unmasking gain coefficients card

hj

:=

d X

  EBi,j Info(Zi ; ZBi,j )

where Bi,j is uniform over j-cardinality subsets of [d] \ {i}.

(15b)

i=1

:= 0. These coefficients are defined for j = 1, . . . , d − 1, with hcard 0 These two notions lead to the two versions of UGC growth complexity, given by Z q H(p, q) := t(1 − t)h′ (t) dt for pairs 0 ≤ p < q ≤ 1, and Bernoulli:

(16a)

p

Exact cardinality:

Hcard (a, b) :=

 b−1    X j j  card 1− hj − hcard j−1 , d d j=a+1

0 ≤ a < b ≤ d.

(16b)

Our main theorem bounds the KL error of both the Bernoulli and fixed-cardinality unmasking samplers in terms of increments of these respective UGC functions. It applies to any reveal-time grid (9b), and the t is natural. analysis shows why the reveal odds function ψ(t) := 1−t 11

Theorem 1. For any grid (9b), the Ber-unmasking sampler produces output Zb such that −1
NX DKL PZ PZb ≤ j=0



 ψ(tj+1 ) − 1 H(tj , tj+1 ) + DKL (Pt0 ∥Qt0 ) + ΓC (T ). ψ(tj )

(17a)

a Moreover, for any d-aligned grid tj = dj , the Card-unmasking sampler produces output Zbcard such that

DKL PZ PZbcard ≤

N −1 X j=0

(

) ψ(tj+1 )
− 1 Hcard (aj , aj+1 ) + DKL (Pt0 ∥Qt0 ) + ΓC (T ). ψ tj + d1

(17b)

We prove these two claims in Sections 6.1 and 6.2, respectively. The proofs are relatively short, and follow the same template of starting from an exact representation of the KL discretization error over the interval, and then bounding it in terms of H or Hcard respectively. For the Bernoulli sampler, we first derive the exact KL representation, from which the upper bound (17a) follows from an elementary argument. On the other hand, the fixed-cardinality guarantee (17b) makes use of the exact representation of the KL are rescaled versions discretization error, derived by Chen et al. [CCL25]; in particular, the coefficients hcard j of their information coefficients. In contrast to our upper bounds, neither form of the exact KL error has simple additive structure in terms of the log-reveal odds. This structure in our upper bounds turns out to be very useful: as we discuss in the sequel (see Corollary 1 and Section A), a simple stepsize choice in the bound (17b) gives an immediate sharpening of the fixed-cardinality sampling guarantees given in the paper [CCL25]. More generally, these upper bounds lend themselves naturally to developing and analyzing K-block sampling schemes (see Section 4). Estimated denoisers: If the sampler replaces each exact single-site posterior µi (·, x) by an estimate µ bi (·, x), while leaving the random-subset law and reveal probabilities unchanged, then the right-hand side of equation (17a) acquires the additional additive term Eden (b µ) :=

N −1 X j=0

h βj E

X

i DKL µi (·, Xtj ) µ bi (·, Xtj )

t

−t

j where βj := j+1 1−tj .

(18)

i∈M(Xtj )

Here the expectation is under the true reveal-process law of Xtj . A similar comment applies to the Cardunmasking sampler. We establish this fact as part of the proof in Section 6.1. Relation between UGC functions: By inspection, it is clear that the Ber-UGC and Card-UGC functions, as defined in equation (2c) and equation (16b) respectively, are closely related. This intuition can be formalized as follows. For any 0 ≤ p < q ≤ 1, define the multinomial triple (U0 , U1 , U2 ) ∼ Multinomial (d + 1; p, q − p, 1 − q). We then have the equivalence H(p, q) =

  d E Hcard (AU , BU ) , d+1

(19)

where AU := max{U0 − 1, 0}, and BU := min{U0 + U1 , d}, and we set Hcard (a, b) = 0 whenever a ≥ b. Thus, the Ber-UGC over [p, q] is the average of Card-UGC quantities over randomized cardinality endpoints. See Section C.1.2 for the proof of the identity (19).

3.3

A single-block guarantee and its consequences

−1 By inspection, the bounds (17a) and (17b) are very well-suited to unmasking-time grid points {tj }N j=0 for which the reveal odds ψ(tj ) = tj /(1 − tj ) evolve in a geometric way. In this section, we use Theorem 1 to derive guarantees for this particularly simple choice of stepsizes.

12

3.3.1

Guarantees for a single geometric block

For a parameter ρ > 0, consider the sequence n o ψ(tj+1 ) = min (1 + ρ)ψ(tj ), ψ(T ) .

Geo(ρ) :

(20)

When the unmasking algorithm is run with this geometric sequence, we refer to it as the Geo(ρ)-unmasking algorithm. Corollary guarantee for a reveal-odds geometric scheme). Given an iteration budget  1 (Single-block   ψ(T ) N ≥ log ψ(t , the Geo(ρ)-unmasking sampler with multiplier ρ := exp N1 log(ψ(T )/ψ(t0 )) − 1 0) produces completed output Zb such that DKL (PZ ∥PZb ) ≤

 ψ(T )  2H(t0 , T ) + DKL (Pt0 ∥Qt0 ) + ΓC (T ). log N ψ(t0 )

(21)

This result is an immediate consequence of Theorem 1, and we include the proof here. It makes essential use of the additivity property (3) of H. Proof. By construction, the schedule traverses the interval [t0 , T ] in exactly N steps, and moreover, we have ψ(tj+1 ) ψ(tj ) − 1 ≤ ρ at each round. Applying the bound (17a) from Theorem 1 yields DKL (PZ ∥PZb ) ≤ ρ

N −1 X

H(tj , tj+1 ) +DKL (Pt0 ∥Qt0 ) + ΓC (T )

j=0

|

(i)

{z

}

= ρH(t0 ,T )

(ii) 2H(t , T ) 0

≤

N

log

 ψ(T )  ψ(t0 )

+ DKL (Pt0 ∥Qt0 ) + ΓC (T ),

where step (i) uses the additivity property (3), and step (ii) follows since exp(u) − 1 ≤ 2u for u ∈ [0, 1],  applied with u = N1 log

ψ(T ) ψ(t0 )

.

Exact cardinality sampler: An analogous result holds for the fixed-cardinality uniform sampler, using a slightly different geometric schedule. Given a pair of integers a0 < A between 1 and d − 1, we define t0 = ad0 ψ(T )
and T = A 1 , starting from a0 , we recursively choose aj+1 as d . Given an iteration budget N ≥ log the largest integer in {aj + 1, . . . , A} such that   aj + 1 aj+1
ψ ≤ (1 + ρ)ψ d d

ψ(t0 + d )

where ρ := exp

n

1 N log

ψ(T )
1 ψ(t0 + d )

o

− 1.

With these choices, running the fixed-cardinality sampler yields completed output Zbcard such that  ψ(T ) 
2Hcard (a0 , A) DKL PZ PZbcard ≤ log + DKL (Pt0 ∥Qt0 ) + ΓC (T ). N ψ(t0 + d1 ) The proof uses the bound (17b) from Theorem 1, and then applies the same reasoning as the proof of Corollary 1.

13

3.3.2

Some consequences

At a high level, Corollary 1 and its analogue for fixed-cardinality sampling show that the iteration complexity when using a single geometric multiplier protocol is governed by the aggregate UGC mass functionals H(0, 1) and Hcard (0, d), respectively, for the Bernoulli and fixed-cardinality versions. These relations allow us to recover and sharpen several existing results on unmasking samplers. As we show in Section A, the aggregate UGC complexities associated with Bernoulli and fixed-cardinality unmasking are equivalent up to universal constants, and they are closely related to the complexity governing the CTMC guarantees of Dmitriev et al. [DHW26]. Consequently, the single-block guarantee in Corollary 1 yields comparable guarantees for these different unmasking mechanisms, while sharpening bounds based only on total correlation and dual total correlation. We defer these comparisons, together with further information-theoretic aspects of aggregate UGC, to Section A.

3.4

Why the UGC path matters

The aggregate UGC complexity is a coarse measure, relative to the full log-reveal-odds density that captures the full UGC path. It is useful to consider a stylized pair of distributions, based on some primitives analyzed in the paper [DHW26], that makes this discrepancy very explicit: while their aggregate complexities are identical, the UGC path structure is extremely different, and as we will see in the next section, the resulting optimal algorithms are very different. Recall that the log-reveal-odds UGC density is given by q(λ) = r2 (1 − r)2 h′ (r)

Repeated common bit (d = 64)

0.4 0.3 0.2 0.1 0.0

8

4

0

Log reveal odds

4

0.3 0.2 0.1 0.0

8

(a)

(23)

Global parity (d = 64)

0.4

UGC density q( )

UGC density q( )

λ

e where r = φ−1 (λ) = 1+e λ.

8

4

0

Log reveal odds

4

8

(b)

Figure 5. Log-reveal-odds UGC densities for the repeated-common-bit model in panel (a) and the globalparity model in panel (b), both with d = 64. The two densities are reflections of one another about λ = 0 in log-reveal-odds coordinates: their information is concentrated near opposite ends of the reveal path even though their full-path UGC values agree.

Repeated bit: First, suppose that we draw a hidden V ∼ Ber(1/2), and then generate the random vector Z ∈ {0, 1}d by setting Z1 = · · · = Zd = V . This construction is simply the noiseless (η = 0) instance of the noisy-repeated-bit example discussed previously in Section 2.2.1, and moreover, we can compute its UGC density function explicitly. From equation (23), it suffices to compute h and then its derivative. Conditional on coordinate i remaining masked, its mutual-information contribution equals Ent(V ) = log 2 once at least one of the other d − 1 coordinates is revealed; it is equal to zero otherwise. It follows that

14

 hrep (r) = d log 2 1 − (1 − r)d−1 . Taking derivatives and combining with the formula (23) yields  qrep (λ) = d(d − 1) log 2 r2 (1 − r)d .

(24a)

Global-parity model: In the global-parity model, draw Z1 , . . . , Zd−1 independently from Ber(1/2) and set Zd = Z1 ⊕ · · · ⊕ Zd−1 . Calculations similar to those in the previous example show that  qpar (λ) = d(d − 1) log 2 (1 − r)2 rd . (24b) We plot these two UGC densities in Figure 5. Comparing equation (24a) and equation (24b) reveals the relation2 qpar (λ) = qrep (−λ), which corresponds to reflection across λ = 0 in log-reveal-odds. This reflection can be seen by comparing the two panels. As can be verified via a straightforward calculation, both models have H(0, 1) = (d−1) (d+1) log 2, so that they have matched aggregate UGC mass. However, the full UGC path reveals the differences in their distributional structure: the repeated-bit dependence is resolved near the beginning of the path, whereas parity dependence persists until the end. Moreover, as we develop in the next section, this path difference has important consequences for designing optimal algorithms.

4

Certified-optimal bounds for the UGC path

Thus far, we have considered the single-block consequence of Theorem 1, based on one geometric multiplier over the full reveal interval. The additivity (3) of the UGC function allows a more refined approach. We first partition the reveal interval into K blocks, and then optimize the stepsizes within each block. Doing so leads to a UGC partition complexity that refines the coarse UGC bound. We then show how the required UGC increments can be estimated from samples, yielding certified data-dependent samplers and allowing the block boundaries themselves to be optimized. Finally, by refining the partition, we recover the fine-partition complexity PUGC in equation (5b), which characterizes the optimal discretization error up to a factor of two.

4.1

An optimal K-block guarantee

For some  0 < t0 < T < 1, we begin with a prescribed K-block partition of the interval [t0 , T ], say of the form P := [bk , bk+1 ] | k = 0, . . . , K − 1 , where b0 = t0 and bK = T . For each block, define its log-reveal-odds length and UGC mass by Sk := φ(bk+1 ) − φ(bk ),

and

Hk := H(bk , bk+1 ),

(25)

where we recall that φ(r) := log(r/(1 − r)). PK−1 Given a total iteration budget N , suppose that we allocate positive integers Nk with k=0 Nk = N . We then run the Geo(ρ)-sampler to traverse the k th -block in Nk iterations. Applying Theorem 1 blockwise yields DKL (PZ ∥PZb ) ≤

K−1 Xh

i  exp Sk /Nk − 1 Hk + DKL (Pt0 ∥Qt0 ) + ΓC (T ).

(26a)

k=0

Here we have again used the additivity property (3) to bound the error of each block by the scalar multiple of its UGC mass. Hence, for the prescribed partition, the best guarantee furnished by Theorem 1 is obtained from the integer program min

N0 ,...,NK−1 >0

K−1 Xh

i  exp Sk /Nk − 1 Hk

k=0

subject to

K−1 X

Nk = N.

k=0

It is easy to see that this integer program can be solved efficiently by dynamic programming. 2 This follows from the identity φ−1 (−λ) = 1 − φ−1 (λ), which interchanges r and 1 − r under reflection.

15

(26b)

4.2

Guarantees for an explicit procedure

Rather than solving the dynamic program directly, we seek an explicit allocation that exposes the structure of its solution. Doing so leads to the UGC-based partition complexity !2 K−1 Xp Sk Hk , (27) CUGC (P) := k=0

which interpolates between the single-block complexity (5a) and the fine-partition limit (5b). Proposition 1 (Near-optimal geometric schedules partition). Given any K-block parn for a K-block  o tition P of [t0 , T ] and an iteration budget N ≥ 2 K + 2 log p r ) CUGC (P) Sk , ρk := min 1, 4 N Hk

ψ(T ) ψ(t0 )

, define the geometric multipliers

(

k = 0, 1, . . . , K − 1.

(28a)

Then the K-block Ber-unmasking procedure based on P uses at most N score evaluations and yields an b satisfying output Z DKL (PZ ∥PZb ) ≤

4CUGC (P) + DKL (Pt0 ∥Qt0 ) + ΓC (T ). N

(28b)

⋆ ) denotes the optimal solution of the dynamic program (26b) with iteration Moreover, if (N0⋆ , . . . , NK−1 budget N , then K−1 X k=0

ρ(Nk⋆ )Hk ≥

CUGC (P) N

where ρ(Nk⋆ ) := exp(Sk /Nk⋆ ) − 1.

(28c)

The proof is an unmasking analogue of the argument used for Gaussian diffusion in our companion paper [Wai26]. The principal changes are the log-reveal-odds clock, the corresponding UGC block masses, and the boundary terms; the remainder follows mutatis mutandis. This parallelism reflects a common geometric structure underlying Gaussian diffusion and discrete unmasking, despite their different noise mechanisms and natural path coordinates. The upper and lower bounds in Proposition 1 show that, apart from the initial and terminal cost, the explicit allocation is within a factor of 4 of the optimal value of the dynamic program (26b). It also always improves upon the single-block complexity. Indeed, by Cauchy–Schwarz and the additivity (3) of the UGC complexity, we have !2 ! K−1 ! K−1 K−1 X X Xp  CUGC (P) = Sk Hk ≤ Sk Hk = φ(T ) − φ(t0 ) H(t0 , T ) = CUGC ([t0 , T ]). (29) k=0

k=0

k=0

Equality holds if and only if there is a scalar c > 0 such that Sk = cHk for every block. Thus, partitioning provides a strict improvement whenever the UGC mass is distributed non-uniformly relative to log-reveal-odds length.

4.3

Sandwich-UGC estimators from KL increments

While the bounds in Proposition 1 are attractive, the geometric multipliers ρk depend on knowledge of the UGC increments Hk for each block k. In this section, we describe how it is possible to estimate these increments, assuming that we have available a collection of i.i.d. samples {Z (ℓ) }m ℓ=1 from the original distribution PZ . This estimator forms the basis for the certified-optimal sampling schemes described in Section 4.4. 16

4.3.1

Kullback–Leibler unmasking increments

From the definition (2a), the Bernoulli unmasking gain h is defined in terms of the original reveal process Xp with additional conditioning on the event {i ∈ M(Xp )}. Consequently, it is useful to introduce the forcedmask reveal-time process that explicitly encodes this conditioning. Recalling the uniform random variables that underlie the masking process (1a), we define    (30a) Mt|i := {i} ∪ ℓ ∈ [d] \ {i} | Uℓ > t , and Xt|i := Mt|i , Z(Mt|i )c . Due to the shared set of uniform random variables, the process satisfies a natural coupling: for q > p, the variable Xq|i can be obtained from Xp|i by revealing each variable in M(Xp|i ) \ {i} independently with q−p probability 1−p . Moreover, by construction, the random variable Xt|i has the law of the original reveal process Xt conditional on i ∈ M(Xt ). For 0 ≤ p < q ≤ 1, we define the KL unmasking increment D(p, q) := (q − p)

d X

E [DKL (πi,q ∥ πi,p )]

where πi,t (·) ≡ µi ·, Xt|i for each t ∈ [0, 1],

(30b)

i=1

and the expectation is taken with respect to the joint forced-mask coupling (Xp|i , Xq|i ) for each i = 1, . . . , d. With this set-up, we have the following guarantee: Lemma 1. For every 0 ≤ p < q < 1, we have the exact relation D(p, q) = h(q) − h(p). q−p

(31a)

For every 0 < p < q < 1, we also have the sandwich relation 1 1 + cp,q D(p, q) ≤ H(p, q) ≤ D(p, q) cp,q cp,q

where cp,q :=

ψ(q) − 1 > 0. ψ(p)

(31b)

See Section 6.3.1 for the proof. Observe that the sandwich relation (31b) is particularly well-suited to intervals [p, q] for which the ratio of the reveal odds ψ(r) = r/(1 − r) is well-controlled. Concretely, for a dyadic interval with ψ(q)/ψ(p) = 2, the relation provides a factor two sandwich. We now exploit this fact to design an estimator for a partition of arbitrary length. 4.3.2

A tail-robust estimate of the UGC increment

We now specify and analyze a tail-robust estimate of the UGC increment. For an arbitrary pair 0 < p < q < 1, consider the reveal-odds-dyadic partition {vj }Jj=0 given by l  m ψ(q) ψ(vj ) := min{2j ψ(p), ψ(q)} where ψ(r) := r/(1 − r) and J := log2 ψ(p) . (32a) For each interval [vj , vj+1 ], Lemma 1 guarantees that D(vj , vj+1 ) ≤ H(vj , vj+1 ) ≤ 2D(vj , vj+1 ), and hence PJ−1 PJ−1 that j=0 D(vj , vj+1 ) ≤ H(p, q) ≤ 2 j=0 D(vj , vj+1 ), where we have used the additivity property (3) of H. (ℓ) Given a clean sample Z (ℓ) , we can simulate the forced-mask reveal process Xt|i over the interval t ∈ [0, 1], for each coordinate i = 1, . . . , d, and using these trajectories, we can construct the trajectory statistic   d J−1   X X (ℓ) (ℓ) Q(ℓ) := (vj+1 − vj )DKL µi (·, Xvj+1 |i )∥µi (·, Xvj |i ) . (32b)   i=1

j=0

17

PJ−1 By construction, we have E[Q(ℓ) ] = j=0 D(vj , vj+1 ), where the expectation is taken over the sample Z (ℓ) , along with the forced-mask trajectory randomness. Combined with the factor two sandwich property, we have constructed a statistic such that E[Q(ℓ) ] ≤ H(p, q) ≤ 2E[Q(ℓ) ].

(32c)

A tail-robust Given the statistic Q(ℓ) from equation (32b), the standard Monte Carlo estimate Pestimator: m 1 (ℓ) is given by m ℓ=1 Q . It is unbiased and can be analyzed, but can be overly sensitive to the tail behavior of the underlying KL divergences in the estimate. To introduce tail robustness, we instead analyze a truncated version of this estimator. It requires a moment bound, but provides exponential Bernstein-type tail control. More precisely, suppose that, for some moment order α ≥ 4, the changes in the one-coordinate denoisers satisfy the aggregate moment bound  ( ) 2/α d    α/2     X (ℓ) (ℓ)  µi ·, Xvj |i ≤ Bα for some known Bα > 0. max E DKL µi ·, Xvj+1 |i j=0,...,J−1   i=1

(33a) Under this condition, we define the τ -truncated estimator b m (p, q) := 2 H m

m X

min{Q(ℓ) , τ }

where τ := (q − p)Bα

n

3(m−1) 7 log(4/η)

o2/α

.

(33b)

ℓ=1

Here η ∈ (0, 1) is a user-chosen failure probability and {Q(ℓ) }m ℓ=1 are i.i.d. copies of the trajectory statistic Q(ℓ) from equation (32b). We also define the (unbiased) empirical variance of our truncated estimator as Vb :=

m 2 1 X b m (p, q) . 2 min{Q(ℓ) , τ } − H m−1

(33c)

ℓ=1

b m (p, q), the UGC Proposition 2 (Factor-two data-dependent sandwich on UGC). Given the estimate H increment H(p, q) can be sandwiched as (A) (B)  b m (p, q) + rbm (η) ≤ 2 H(p, q) + rbm (η) H(p, q) ≤ H

with probability at least 1 − η,

where the upper confidence correction takes the form s  1−2/α 2Vb log(4/η) 7 log(4/η) rbm (η) := + 4(q − p)Bα . m 3(m − 1)

(34a)

(34b)

See Section 6.3.2 for the proof. Interpretation of the tail condition: Condition (33a) controls the tails of the aggregate change in the one-coordinate denoisers over each reveal-odds-dyadic interval. This choice is important: it measures the same posterior changes that enter the trajectory statistic (32b), rather than the absolute uncertainty of the individual denoisers. A condition on the true-symbol log loss, for instance, could be large even when revealing additional coordinates produces no change in any posterior. As a basic example, suppose that the coordinates of Z are independent. For the exact Bayes denoisers considered here, the forced-mask observation Xt|i reveals only masking variables and a subset of the coordinates Z−i . Consequently, Zi is independent of Xt|i , and hence
µi ·, Xt|i = L(Zi ) for every t ∈ [0, 1]. 18

All of the KL increments in equation (33a) therefore vanish, so that the condition holds with any positive Bα , in agreement with the fact that the corresponding UGC increment is zero.

4.4

Sampling schemes that are certified-optimal

In this section, we turn to the key problem of designing samplers that are certified-optimal, meaning that their iteration complexity meets the optimum specified by the partition complexity functional, up to datadependent residual terms, and that for any target accuracy ε > 0 and failure probability η ∈ (0, 1), there are data-dependent parameter choices for the sampler guaranteeing that its output Zb satisfies the accuracy bound
with probability at least 1 − η. (35) DKL PZ ∥PẐ ≤ ε In Section 4.4.1, we specify and analyze a certified procedure for a given K-block partition, whereas in Section 4.4.2, we describe certified procedures based on dynamic programming for choosing the block boundaries. 4.4.1

Certified-optimal sampling using K-blocks

Suppose that we are given a K-block partition of the interval [t0 , T ], based on blocks of the form [bk , bk+1 ]. Our goal is to design a sampler that achieves the optimal K-bound from Proposition 1, but using a choice of geometric multipliers that can be estimated based on samples {Z (ℓ) }m ℓ=1 . Here we describe a simple way of doing so, one which exploits the guarantees from Proposition 1 as well as the UGC increment estimator analyzed in Proposition 2. Suppose that, for each block [bk , bk+1 ], we apply the tail-robust estimator from Proposition 2 with failure probability η/K. We then obtain a tail bound with the upper confidence limit s  1−2/α 2Vbk log(4K/η) 7 log(4K/η) + 4(bk+1 − bk )Bα rbk,m (η/K) := , (36a) m 3(m − 1) | {z } ≡b rk

where Vbk is the empirical variance estimate for the block. Finally, we define the surrogate partition complexity  K−1 2 Xp q b k,m + rbk . b := Sk H (36b) CUGC (P) k=0

With this set-up, we have the following certified analogue for the Ber-unmasking sampler with geometric multipliers, or the Geo(ρ)-scheme for short. Theorem 2 (Data-certified guarantees for multi-block scheme). Given a failureprobability  η ∈ (0, 1), and a K-block partition of the interval [t0 , T ], consider any iteration budget N ≥ 2 K + 2 φ(T ) − φ(t0 ) Then running the K-block Geo(ρ)-unmasking scheme with the geometric multipliers q   s  4 C  b UGC (P) Sk ρbk := min 1, for k = 0, 1, . . . , K − 1. (37a) b k,m + rbk   N H b such that yields a completed output Z DKL (PZ ∥PZb ) ≤

b UGC (P) 4C + DKL (Pt0 ∥Qt0 ) + ΓC (T ) N

and does so with at most N unmasking rounds. See Section 6.4 for the proof. 19

with probability at least 1 − η,

(37b)

.

Certified-optimal: Let us now clarify why the guarantee of Theorem 2 is both certified and optimal (up to constant factors). Suppose that our goal is to design a sampler that achieves a prescribed KL error ε > 0. As discussed in Section 3.1.3, it is straightforward to ensure that the initialization and terminal contribution satisfies DKL (Pt0 ∥Qt0 ) + ΓC (T ) ≤ ε/2. From Theorem 2, it can be seen that choosing an iteration number b UGC (P)/ε guarantees that the resulting sampler has output Z b such that N ≥ 8C DKL (PZ ∥PZb ) ≤ ε

with probability at least 1 − η.

(38)

Thus, the estimated partition complexity provides both a data-dependent schedule and an end-to-end certificate of its sampling accuracy. For this reason, we refer to it as a certified guarantee. It is also near-optimal in a precise sense. In particular, recall that the procedure in Proposition 1 achieves the DP optimum up to a constant factor of 4. Under the same initialization conditions, doing so requires an iteration number N ≥ 8CUGC (P)/ε. As shown in the proof of Theorem 2, on the simultaneous high-probability event from Proposition 2, we have ( ) q K−1 Xp √ p b UGC (P) ≤ 2 CUGC (P) + Sk rbk . C k=0

Thus, the increase in computational complexity incurred by replacing the oracle UGC increments by√their data-dependent estimates is captured explicitly by the confidence radii rbk , together with the factor of 2 in the sandwich bound. 4.4.2

Certified-optimal boundary selection

Our analysis thus far has focused on optimal schemes for a fixed set of K blocks. In this section, we formulate the problem of optimizing the split points that define a (near)-optimal set of K blocks via a simple dynamic program. More precisely, for an integer J > K that defines the resolution, suppose that we fix a fine deterministic grid of candidate reveal times τ0 = t0 < τ1 < · · · < τJ = T . Our goal is to select indices {ik }K−1 k=1 from the set {1, . . . , J − 1}, and use the corresponding reveal times tk = τik , together with the endpoints t0 = t0 and tK = T , to define a K-block partition. We seek to minimize the square root of the partition complexity K−1 Xp p CUGC (P) = Sk Hk ,

where Sk := φ(tk+1 ) − φ(tk ), and Hk := H(tk , tk+1 ).

k=0

Given the naturally sequential structure, this optimization problem can be solved by dynamic programming, as we now formalize. For candidate intervals 0 ≤ i < j ≤ J, define the edge cost q where S(a, b) := φ(b) − φ(a). (39a) e(i, j) := S(τi , τj )H(τi , τj ), With this notation, our ultimate goal is to compute VK (J) :=

min

0=j0 <j1 <···<jK =J

K−1 X

e(jk , jk+1 ).

(39b)

k=0

This quantity is the minimum cost of a K-block partition of the interval [τ0 , τJ ] = [t0 , T ]. It is easy to see that it can be computed via dynamic programming, where we introduce intermediate quantities Vk (j), defined for each k ∈ {0, . . . , K} and j ∈ {0, . . . , J} as the minimum cost of a k-block partition of the subinterval [τ0 , τj ], with value +∞ when no such partition exists. Finally, we observe that the edge costs in equation (39a) depend on the UGC increments H(τi , τj ). As in the previous section, we can obtain confidence bounds for increments of this type using the estimator from Proposition 2, thereby obtaining a data-dependent version of the dynamic program. 20

5

Fine partition limit and optimal Euler discretization

In the preceding section, we developed and analyzed sampling methods based on finite K-block partitions, and showed how their performance depends on the local UGC mass of the chosen partition. We now characterize the limit under arbitrarily fine refinement and connect it to the optimal KL error of Euler discretization. This analysis highlights the central role of the log-reveal-odds UGC density q defined in equation (4).

5.1

Fine partition limit and its Euler optimality

We specialize to the canonical reveal interval Ifull := [1/d, 1 − 1/d]. Its image under the log-reveal-odds map is [−ℓd , ℓd ], where ℓd := log(d − 1). For a partition P = {[bk , bk+1 ]}K−1 k=0 of Ifull , define the corresponding logreveal-odds blocks Ik := [φ(bk ), φ(bk+1 )] and their lengths Sk := Len(Ik ). Our previously defined partition complexity can then be written directly in terms of q as CUGC (P) :=

K−1 X

s

!2

Z Sk

.

q(λ) dλ

(40a)

Ik

k=0

Observe that refinement can only decrease this complexity: whenever P ′ is a refinement of P, we have CUGC (P ′ ) ≤ CUGC (P). This fact motivates the fine-partition limit PUGC (Ifull ) := inf CUGC (P),

(40b)

P

where the infimum is over all finite partitions of Ifull . We now characterize this infimum explicitly in terms of √ q, and show that the same quantity governs the sharp leading-order KL error of optimal Euler discretization. Theorem 3 (Fine-partition limit and Euler optimality). Suppose that q is continuous and strictly positive on [−ℓd , ℓd ]. Then the fine-partition limit (40b) is given by PUGC (Ifull ) =

Z ℓd p

!2 q(λ) dλ

.

(41a)

−ℓd

Moreover, the optimal N -step KL discretization error satisfies inf

d−1 1 d =t0 <···<tN = d

N −1 X

Γumask (tj , tj+1 ) = | {z } j=0 KL discretization error

PUGC (Ifull ) + o(N −1 ), 2N

(41b)

where the infimum is over all N -step discretizations of the canonical interval [ d1 , d−1 d ]. See Section 6.5 for the proof. Taken together, the two claims in Theorem 3 show that the fine-partition complexity is precisely twice the leading-order constant governing the optimal Euler KL error, and that it can be approached arbitrarily closely by the K-block schemes analyzed in Proposition 1 and Theorem 2. Related work: A related square-root principle appears in the analysis of Lavenant and Zanella [LZ25], who study random-order fixed-cardinality unmasking and express its factorization error in terms of a cardinalityindexed information profile. They then consider a high-dimensional continuum limit in which, using our notation, both the ambient dimension d and the number of sampling rounds N diverge. In the regime d/N → +∞, optimization of the resulting limiting functional yields a square-root information-profile rule. Our results differ in both formulation and scope. The UGC density q is defined directly from the finitedimensional reveal path, with log-reveal odds providing the natural path coordinate, and our finite-step 21

bounds apply at fixed dimension and for arbitrary finite partitions. Moreover, rather than optimizing only a limiting variational problem, we optimize blockwise schedules at finite dimension, and use estimated UGC increments to construct data-dependent partitions with certified guarantees. Finally, equation (41b) gives the sharp leading-order optimal Euler error at fixed dimension d as the number of discretization steps N grows.

5.2

Gains and convergence to the fine-partition limit

We first ask how much can be gained by exploiting the local UGC geometry, and then how many blocks are needed to realize this gain. The maximal improvement over a single-block scheme is measured by Ratio(PZ ) :=

CUGC ([Ifull ]) ≥ 1. PUGC (Ifull )

(42)

For the repeated-bit and parity models (see Figure 5), we have Ratio(PZ ) ≍ log(d), yielding a growing but relatively modest separation. Much larger gains are possible when the UGC density becomes increasingly con√ e d) centrated with dimension. For example, the discrete mixture model in Figure 2 satisfies Ratio(PZ ) = Ω( with high probability over the cluster centers. We next exhibit the same phenomenon in a simpler random XORSAT model, for which the scaling can be analyzed via an easier argument. 5.2.1

√ e d) gains An example with Ω(

√ e d) separation, consider a dense random XORSAT ensemble. To obtain a tractable example with the same Ω( For ambient dimension d and latent dimension k, draw A ∈ {0, 1}d×k with i.i.d. Ber(1/2) entries, draw U ∈ {0, 1}k with i.i.d. Ber(1/2) entries, and set Z = AU modulo two. Thus, Z is supported on a random Boolean subspace. We parameterize the ensemble by the ratio α := k/d. Figure 6 plots the UGC density for d = 64 and α ∈ {0.2, 0.5, 0.8}. The density is unimodal, with its peak moving from right to left as α increases. To quantify the resulting gain, we focus on the symmetric case α = 0.5.

10

5

0

2

0

Log reveal odds

(a)

2

pre-transition

transition

post-transition

10

5

0

15

UGC density q( )

15

UGC density q( )

UGC density q( )

15

2

0

Log reveal odds

(b)

2

10

5

0

2

0

Log reveal odds

2

(c)

Figure 6. Plots of the log-odds UGC density q for the random XORSAT ensemble with d = 64 and parameter α ∈ {0.2, 0.5, 0.8} in panels (a)–(c), respectively. Panel (b) marks the pre-transition, transition, and posttransition regions used in the three-block construction.

22

Lemma 2. For the random XORSAT model with α = 0.5, for each even dimension d above a universal constant, there is a distribution PZ such that p for a universal constant c > 0, (43) Ratio(PZ ) ≥ c d log d and this gain is achieved, up to constants, by a K = 3 block scheme. See Section D for the proof. Our analysis uses the probabilistic method [AS16] to show that, for each sufficiently large dimension, there exists a realization of A, inducing a “good” distribution PZ , that satisfies the following properties. First, its coarse complexity satisfies CUGC ([I √ full ]) = Θ(d log d). Aspa consequence, it suffices to construct a three-block partition with complexity O( d log d). Set δd = c1 log d/d for a sufficiently large constant c1 > 0, and define the reveal time intervals       Ipre := 1/d, 1/2 − δd , Itrans := 1/2 − δd , 1/2 + δd , I := 1/2 + δd , 1 − 1/d . post

Our proof shows that the “good” PZ has UGC masses H(Ipre ) = O(d−10 ),

H(Itrans ) = Θ(d),

and

H(I ) = O(d−10 ). post

(44)

Thus, consistent with the behavior shown in Figure p 6, essentially all of the UGC mass lies in the transition interval I log d/d). This combination leads to the claimed scaling trans , which has log-reveal-odds width O( √ O( d log d) for the 3-block partition. 5.2.2

Uniform partitions: a generic convergence bound

The XORSAT example shows that a small number of well-placed blocks can already be near-optimal. For comparison, let UK denote the uniform K-block partition of [−ℓd , ℓd ], which ignores local geometry when choosing the block boundaries. Here we quantify its convergence to the fine-partition limit under a simple global regularity condition. √ In particular, suppose that f := q is continuous on [−ℓd , ℓd ], and define its modulus of continuity by ωf (h) := sup |f(λ) − f(λ′ )| λ,λ′

where the sup ranges over pairs λ, λ′ ∈ [−ℓd , ℓd ] with |λ − λ′ | ≤ h.

(45)

Lemma 3 (Accuracy of uniform K-block approximations). For the uniform partition UK , we have   p p 2ℓd . (46) CUGC (UK ) ≤ PUGC (Ifull ) + ℓd ωf K The proof is the same modulus-of-continuity argument used for Gaussian diffusion in our companion paper (Lemma 4, [Wai26]). As a particular consequence, if f is Lipschitz with constant Lip(f), then p p 2 Lip(f)ℓ2d CUGC (UK ) ≤ PUGC (Ifull ) + . K

(47)

This generic bound can be conservative because it depends on a global regularity constant and fixes the block boundaries in advance. For the XORSAT example, the Lipschitz constant Lip(f) grows with dimension even though K = 3 adaptively placed blocks are near-optimal. When samples are available, the data-dependent DP procedure in Section 4.4.2 instead optimizes the block boundaries directly.

23

6

Proofs

In this section, we collect the proofs of a subset of our results, including the proof of Theorem 1, with the Bernoulli and fixed-cardinality versions given in Section 6.1 and Section 6.2, respectively; the proof of our data-certification procedures, including Lemma 1 and Proposition 2 in Section 6.3; the proof of Theorem 2 in Section 6.4; and finally, the proof of Theorem 3 in Section 6.5. We defer the proofs of other results in the paper to the appendices.

6.1

Proof of Theorem 1: Bernoulli version

We begin by proving the bound (17a) for the Bernoulli unmasking sampler. 6.1.1

Main argument

Recall from equation (1b) the reveal process {Xt , t ∈ [0, 1]}. For any pair 0 < p < q < 1, we let Kp,q denote the transition kernel of the reveal process in moving from Xp to Xq . On the other hand, the unmasking b p,q denote the associated sampler is based on the updates in equation (11a) and equation (11b), and we let K b b transition kernel that moves from Xp to Xq . We introduce the shorthand notation  i h b p,q (· | Xp ) , (48) Γumask (p, q) := EXp DKL Kp,q (· | Xp ) K corresponding to the averaged Kullback–Leibler (KL) discrepancy between the two transition kernels over the interval [p, q]. Here the expectation is taken over the marginal distribution of the reveal process variable Xp . The core technical result at the heart of Theorem 1 is the following bound: Lemma 4 (Multiplicative control of the one-step unmasking defect). For every 0 < p < q < 1, we have   ψ(q) r Γumask (p, q) ≤ − 1 H(p, q) where ψ(r) := . (49) ψ(p) 1−r See Section 6.1.2 for the proof. Given this lemma, the proof of Theorem 1 is very simple. Recalling the definition (14b) of the completion defect ΓC (T ), we have DKL PZ

N −1 X
(i) PZb ≤ DKL (Pt0 ∥Qt0 ) + Γumask (tj , tj+1 ) + ΓC (T ) j=0 (ii)

≤ DKL (Pt0 ∥Qt0 ) +

N −1  X j=0

 ψ(tj+1 ) − 1 H(tj , tj+1 ) + ΓC (T ), ψ(tj )

where step (i) follows from the KL chain rule, combined with the data processing inequality, whereas step (ii) follows by applying the bound (49) repeatedly to each of the KL increments. Extension to estimated denoisers: Fix a state x at step j = 0, . . . , N − 1. In the existing kernel b t ,t (· | x) is the frozen product transition formed notation, Ktj ,tj+1 (· | x) is the exact block transition, K j j+1 b from the exact one-site posteriors, and Ktj ,tj+1 [b µ](· | x) is the frozen product transition formed from the t −tj estimated posteriors. The two frozen kernels have the same reveal probability βj := j+1 1−tj . Since the exact

24

b t ,t (· | x) have the same one-coordinate marginals, splitting the log-likelihood ratio block transition and K j j+1 through the product of these marginals gives     X b t ,t [b b t ,t (· | x) + βj µ ](· | x) = D K K (· | x) DKL (µi (·, x) ∥ µ bi (·, x)) . DKL Ktj ,tj+1 (· | x) K t ,t KL j j+1 j j+1 j j+1 i∈M(x)

Indeed, for each i ∈ M(x), both one-coordinate kernels leave the coordinate masked with probability 1 − βj . Conditional on revealing it, they use µi (·, x) and µ bi (·, x), respectively. Their one-coordinate KL divergence is bi (·, x)). Averaging the first KL term over x = Xtj under the exact reveal-process given by βj DKL (µi (·, x) ∥ µ law gives exactly Γumask (tj , tj+1 ). Averaging the decomposition and summing over j = 0, . . . , N − 1 in the same KL chain-rule argument adds the term displayed after equation (17a), which completes the proof. 6.1.2

Proof of Lemma 4

Our proof is based on the exact one-step representation Z q Γumask (p, q) = (q − u)h′ (u) du,

(50a)

p

which we return to prove below. Taking this claim as given, let us prove the bound (49). Introduce the shorthand f (u) = u (1 − u), and observe that f (u) ≥ 0 for u ∈ [0, 1]. Note that for every u ∈ [p, q], we have (q − u) ≤ (q − p), and f (u) ≥ p(1 − q), so that     q−p q(1 − p) ψ(q) (i) (ii) q−u ≤ f (u) = − 1 f (u) = − 1 f (u) (50b) p(1 − q) p(1 − q) ψ(p) where step (i) follows from simple algebra; and step (ii) uses the definition ψ(r) = r/(1 − r). Thus, we can write Z q    Z q (iv) ψ(q) ψ(q) (iii) (v) ′ ′ −1 − 1 H(p, q), Γumask (p, q) = f (u) h (u)du = (q − u)h (u) du ≤ ψ(p) ψ(p) p p where step (iii) follows from the exact representation (50a); step (iv) follows from the inequality (50b) along with the fact that h′ (t) ≥ 0 for all t ∈ [0, 1]; and step (v) follows from the definition (2c) of the UGC function. Thus, we have established the claim (49). Proof of the exact representation (50a):

We first claim that h satisfies the relation

d Info (Z; Xu ) , (51a) du Pd where we have introduced the shorthand S(Z) := To verify the i=1 Ent(Zi ) for the entropy sum. claim (51a), fix some δ > 0 sufficiently small. Conditional on Xu , each coordinate i ∈ M(Xu ) is revealed by time u + δ with probability δ/(1 − u). Revealing two or more coordinates has probability O(δ 2 ), while revealing only coordinate i contributes its current posterior entropy. Thus, dividing the first-order expansion of the information increment by δ and letting δ ↓ 0 gives h(u) = S(Z) −

d d X  d 1 X   Info (Z; Xu ) = E 1 i ∈ M(Xu ) Ent (Zi | Xu ) = Ent (Zi | Xu , i ∈ M(Xu )) . du 1 − u i=1 i=1

Since the event i ∈ M(Xu ) is independent of Zi , we have the equivalence Ent (Zi | Xu , i ∈ M(Xu )) = Ent(Zi ) − Info (Zi ; Xu | i ∈ M(Xu )) .

25

Summing over i and using the definitions of S(Z) and h(u) proves the claim (51a). We now prove the exact representation. Conditional on Xp = x, let A ⊆ M(x) be the set of coordinates revealed between p and q. The exact and frozen-product kernels induce the same law for the random subset A. Conditional on A, they unmask coordinates using, respectively, the joint posterior of ZA and the product of its one-site posterior marginals. Consequently, we can decompose the KL divergence as " #   X b Ent (Zi | Xp = x) Xp = x − E [Ent (ZA | Xp = x, A) | Xp = x] . DKL Kp,q (· | x) Kp,q (· | x) = E {z } | i∈A T2 | {z } T1

q−p Note that each coordinate still masked at time p is included in A with probability βp,q := 1−p . Averaging over Xp and using the definition (2a) of h, we find that

E[T1 ] = βp,q (1 − p) {S(Z) − h(p)} = (q − p) {S(Z) − h(p)} . Moreover, conditional on Xp , the new part of Xq is (A, ZA ). Since A is conditionally independent of Z and the random sub-vector ZA is a function of (Z, A), we have E[T2 ] = Info (Z; A, ZA | Xp ) = Info (Z; Xq | Xp ) . Combining these two identities yields Γumask (p, q) = (q − p) {S(Z) − h(p)} − Info (Z; Xq | Xp ) .

(51b)

It remains to analyze the trailing mutual information term. We have Z q Z q d (i) (ii) (iii) Info (Z; Xq | Xp ) = Info (Z; Xq ) − Info (Z; Xp ) = Info (Z; Xu ) du = {S(Z) − h(u)} du, p du p (51c) where step (i) follows because Z −→ Xq −→ Xp is a Markov chain; step (ii) follows from the fundamental theorem of calculus; and step (iii) follows from the identity (51a). Finally, we have Z q Z qZ u Z q (iv) (v) (vi) ′ Γumask (p, q) = {h(u) − h(p)} du = h (v) dv du = (q − v)h′ (v) dv, p

p

p

p

where equality (iv) follows by substituting the representation (51c) into equation (51b) and simplifying terms; step (v) follows from the fundamental theorem of calculus; and step (vi) follows from Fubini’s theorem. This proves the exact representation (50a).

6.2

Proof of Theorem 1: Exact cardinality sampling

For integers 0 ≤ a < b ≤ d − 1, let A be a uniformly random a-subset of [d], independent of Z, and define   n o ⋆ X [a] := ZA , ZA where A ∼ Unif S ⊆ [d] card(S) = a , c where A is a subset of cardinality n a, chosen uniformly atorandom, and independently of Z. Conditional on A, we choose (B | A) ∼ Unif S ⊆ Ac card(S) = b − a , corresponding to the uniform-subset rule (12a) that defines the exact cardinality sampler. [a] b card for the transition generated by Write Kcard to cardinality b, and K a,b for the exact transition from X a,b the Card-unmasking sampler. In analogy with the Bernoulli defect (48), define the one-step KL defect h  i [a] [a] b card . (52a) Γcard (a, b) := EX [a] DKL Kcard a,b (· | X ) Ka,b (· | X )

26

[a] The kernels use the same conditional law N for B. Conditional on (X , B), the exact kernel samples from L(ZB | ZA ), whereas the sampler uses i∈B L(Zi | ZA ), corresponding to the product update (9a). We now apply the one-block specialization of the exact error representation of Chen et al. [Theorem 3.3 CCL25]. After translation3 to our normalization and notation, it guarantees that

Γcard (a, b) =

b−1  1 X (b − j) hcard − hcard j j−1 . d j=a+1

(52b)

We claim that this exact representation admits the following upper bound: Lemma 5 (Fixed-cardinality one-step unmasking defect). For all integers 0 ≤ a < b ≤ d − 1, we have ( ) ψ(b/d)
− 1 Hcard (a, b). Γcard (a, b) ≤ (53) ψ (a + 1)/d Proof. If b = a + 1, both sums are empty. Otherwise, we introduce the shorthand ρa,b :=

ψ(b/d) d(b − a − 1)
−1 = , (a + 1)(d − b) ψ (a + 1)/d

d(b−j) and make note of the upper bound j(d−j) ≤ ρa,b , valid for every a + 1 ≤ j ≤ b − 1. Since hcard is nonj card card decreasing, the increments hj − hj−1 are nonnegative. Putting together the pieces, we arrive at the upper bound

Γcard (a, b) =

b−1  1 X (b − j) hcard − hcard j j−1 d j=a+1

b−1  1 X j(d − j) hcard − hcard ≤ ρa,b 2 j j−1 d j=a+1

= ρa,b Hcard (a, b),

where the last equality follows from the definition (16b) of Hcard . This completes the proof of the upper bound (53). Note that Lemma 5 is the fixed-cardinality analogue of the Bernoulli sampler bound given in Lemma 4. Consequently, the KL bound (17b) for the fixed-cardinality sampler follows from the same argument, instead using Lemma 5 to bound the one-step errors.

6.3

Proof of data-certification results

In this section, we provide proofs of our data-based certification results, including that of Lemma 1 in Section 6.3.1; as well as that of Proposition 2 in Section 6.3.2. 6.3.1

Proof of Lemma 1

We break our proof into parts, one for each claim in the statement. 3 To be clear, their information curve at index j + 1 is hcard /d, and their previous cardinality and block size are a and b − a. j Since B is recoverable from the input and output mask patterns, retaining the subset partition does not change this transition KL divergence.

27

Proof of the claim (31a): We first prove the exact expression (31a). For each coordinate i, let Xp|i and Xq|i denote the coupled observations at times p and q, with coordinate i kept masked. Note that the observation Xp|i is obtained from Xq|i by a random remasking operation independent of Zi conditional on Xq|i . Consequently, the triple Zi −→ Xq|i −→ Xp|i forms a Markov chain, so that we can write

Info(Zi ; Xq|i | Xp|i ) = Info Zi ; Xq|i − Info Zi ; Xp|i . (54a) At the same time, by the conditional KL representation of conditional mutual information, we have
  Info(Zi ; Xq|i | Xp|i ) = E(Xp|i ,Xq|i ) DKL L(Zi | Xq|i , Xp|i ) L(Zi | Xp|i )
  (54b) = E(Xp|i ,Xq|i ) DKL L(Zi | Xq|i ) L(Zi | Xp|i ) , where the second equality follows from the Markov relation Zi −→ Xq|i −→ Xp|i , which implies L(Zi | Xq|i , Xp|i ) = L(Zi | Xq|i ). Combining equations (54a) and (54b) yields the identity 


E DKL L(Zi | Xq|i ) L(Zi | Xp|i ) = Info Zi ; Xq|i − Info Zi ; Xp|i . Summing both sides of this identity over the indices i = 1, . . . , d, and using the definitions of D and h in equation (30b) and equation (2a), respectively, we find that D(p, q) = h(q) − h(p), q−p

(54c)

as claimed in equation (31a). Proof of the sandwich relation (31b):

Turning to the claim (31b), for any pair 0 < p < q < 1, we have

p(1 − q) ≤ t(1 − t) ≤ q(1 − p)

for all t ∈ [p, q].

Multiplying both sides by h′ (t) ≥ 0 and integrating yields Z q   p(1 − q) h(q) − h(p) ≤ t(1 − t)h′ (t) dt ≤ q(1 − p) h(q) − h(p) . p

Using the definition (2c) of H(p, q) and the exact identity (54c), we have proved that p(1 − q) q(1 − p) D(p, q) ≤ H(p, q) ≤ D(p, q). q−p q−p

(54d)

ψ(q) q−p Finally, an elementary calculation gives cp,q := ψ(p) − 1 = p(1−q) , so that we have

p(1 − q) 1 = , q−p cp,q

and

q(1 − p) 1 + cp,q = . q−p cp,q

Combining with equation (54d) proves the claim (31b). 6.3.2

Proof of Proposition 2

Introduce the shorthand µ(p, q) := E[Q(ℓ) ] = have

PJ−1

j=0 D(vj , vj+1 ) for the expected value of Q

µ(p, q) ≤ H(p, q) ≤ 2µ(p, q), 28

(ℓ)

. Note that we (55a)

as previously stated in equation (32c). We first derive the required moment bound for the trajectory statistic from the denoiser condition (33a). Set r := α/2, so that r ≥ 2. Introducing the shorthand notation Kj :=

d X

     (ℓ) (ℓ) µi ·, Xvj |i , DKL µi ·, Xvj+1 |i

for j = 0, . . . , J − 1,

i=1

PJ−1 our statistic Q(ℓ) can be written as Q(ℓ) = j=0 (vj+1 − vj )Kj . We then apply Minkowski’s inequality and make use of the moment condition (33a), thereby obtaining the moment bound J−1 io1/r J−1 n h X X    1/r ≤ E (Q(ℓ) )r (vj+1 − vj ) E Kjr ≤ (vj+1 − vj )Bα = (q − p)Bα . | {z } j=0

j=0

≡Mα

Pm 1

(ℓ) b m (p, q) = . We first bound In terms of the shorthand Y (ℓ) := 2 min{Q(ℓ) , τ }, we can write H ℓ=1 Y m the truncation bias. Since r > 1, we have   h i 2E (Q(ℓ) )r 2Mαr 0 ≤ 2µ(p, q) − E[Y (ℓ) ] = 2E (Q(ℓ) − τ )+ ≤ ≤ r−1 . (55b) r−1 τ τ

We are now set up to apply the empirical Bernstein inequality of Maurer and Pontil [Theorem 4 MP09]. We apply it to both the i.i.d. variables Y (ℓ) /(2τ ) ∈ [0, 1], and then separately to the i.i.d. variables 1 − Y (ℓ) /(2τ ). Observe that the two collections have the same empirical variance. We apply a failure probability of η/2 to each application, so that a union bound over both tail bounds has failure probability at most η. Via this argument, we are guaranteed to have the two-sided tail bound s b b m (p, q) ≤ 2V log(4/η) + 14τ log(4/η) E[Y (ℓ) ] − H (55c) m 3(m − 1) with probability at least 1 − η. Combining the truncation bias (55b) with the tail bound (55c), we find that s r b b m (p, q) ≤ 2V log(4/η) + 14τ log(4/η) + 2Mα 2µ(p, q) − H r−1 m 3(m − 1) τ with probability at least 1 − η. We claim that the right-hand side of this bound is precisely rbm (η). In −1/r log(4/η) particular, we can express the truncation level (33b) as τ = Mα cη , where cη := 73(m−1) . In terms of 1−1/r

log(4/η) this shorthand, we have 14τ3(m−1) = 2Mα cη equivalence follows, and we conclude that

2M r

1−1/r

and τ r−1α = 2Mα cη

. Since 1/r = 2/α, the claimed

b m (p, q) ≤ rbm (η) 2µ(p, q) − H

(55d)

with probability at least 1 − η. It remains to prove the sandwich (34a) claimed in the statement. Conditioned on the event (55d), the upper inequality in equation (55a) gives b m (p, q) + rbm (η). H(p, q) ≤ 2µ(p, q) ≤ H For the other side, the same concentration event and the lower inequality in equation (55a) give  b m (p, q) + rbm (η) ≤ 2µ(p, q) + 2b H rm (η) ≤ 2 H(p, q) + rbm (η) . Combining the last two displays proves the claim (34a). 29

6.4

Proof of Theorem 2

We need to verify two separate claims: the stated bound (37b) on the KL divergence, and the fact that the procedure terminates in at most N unmasking rounds. Throughout the proof, we make use of the shorthand q s b UGC (P) 4 C Sk ak := so that ρbk = min{1, ak }. (56) b N Hk,m + rbk 6.4.1

Verifying the KL bound (37b)

Our first step is to apply Proposition 2 with failure probability η/K ∈ (0, 1) to each block [bk , bk+1 ]. The b k,m ≥ 0 along with the upper confidence correction estimator in this proposition gives us an estimate H rbk > 0. By the union bound over all K blocks, with probability at least 1 − η, we are guaranteed to have b k,m + rbk Hk ≤ H

for all k = 0, . . . , K − 1,

(57)

where we have used inequality (A) in equation (34a). We condition on this “good event” throughout the remainder of the proof. m l Sk bk := bk specified in the For each block, we define the iteration count N log(1+b ρk ) , using the stepsizes ρ theorem statement. Observe that from our choice of ρbk and the definition (20) of the geometric schedule, it bk iterations to traverse the block, and every step inside this block has reveal-odds multiplier takes at most N at most 1 + ρbk . We now apply Theorem 1 to this concatenated multi-block schedule. By the additivity (3) of H, the contribution from block k takes the form ρbk Hk . Summing up these terms, along with the boundary correction, yields DKL (PZ ∥PZb ) ≤

K−1 X

ρbk Hk + DKL (Pt0 ∥Qt0 ) + ΓC (T )

k=0

X (i) K−1 ≤

b k,m + rbk ) + DKL (Pt ∥Qt ) + ΓC (T ) ρbk (H 0 0

k=0

X (ii) K−1 ≤

k=0

√ 4

bUGC (P) C N

|

q

Sk b k,m + rbk ) + DKL (Pt ∥Qt ) + ΓC (T ) (H 0 0 b k,m +b H rk

{z

}

=ak

q b UGC (P) K−1 Xq 4 C b k,m + rbk ) + DKL (Pt ∥Qt ) + ΓC (T ) Sk (H = 0 0 N k=0

b UGC (P) (iii) 4C =

N

+ DKL (Pt0 ∥Qt0 ) + ΓC (T ),

where step (i) uses the good event bound (57); step (ii) uses the fact that ρbk ≤ ak , along with the definition b UGC (P). of ak ; and equality (iii) follows from the definition of C 6.4.2

Verifying the score-evaluation budget

It remains to verify that the procedure terminates in at most N unmasking rounds, so that it satisfies the prescribed score evaluation budget. Since ρbk ∈ [0, 1] and log(1 + u) ≥ u/2 for u ∈ [0, 1], we have bk ≤ 1 + N

Sk 2Sk ≤ 1+ = 1 + 2Sk max{1, a1k } log(1 + ρbk ) ρbk q N b k,m + rbk ). ≤ 1 + 2Sk + q Sk (H b 2 CUGC (P) 30

b UGC (P). Doing so yields We now sum over blocks, using the definition (25) of Sk , and the definition of C K−1 X k=0

K−1 Xq N b k,m + rbk ) Sk (H Sk + q b k=0 2 CUGC (P) k=0  N = K + 2 φ(T ) − φ(t0 ) + 2 ≤ N,

bk ≤ K + 2 N

K−1 X

 where the final inequality follows from the assumed lower bound N ≥ 2K + 4 φ(T ) − φ(t0 ) .

6.5

Proof of Theorem 3

We split our proof into two parts, corresponding to the two claims in the theorem. 6.5.1

Proof of the claim (41a) p Rℓ Introducing the shorthand f(λ) := q(λ) and F⋆ := −ℓd d f(λ) dλ, we need to prove that inf P CUGC (P) = F⋆2 . By the definition (40b), for a partition of [−ℓd , ℓd ] into intervals Lk with length |Lk |, applying the Cauchy– Schwarz inequality on each interval gives s Z X XZ p CUGC (P) = |Lk | q(λ) dλ ≥ f(λ) dλ = F⋆ . (58) k

Lk

k

Lk

Conversely, let UK be aRuniform K-block partition with blocks {Lk }. By the integral form of the mean-value theorem, we can write Lk q(λ)dλ = |Lk |q(ξk ) for some ξk ∈ Lk , and consequently, we obtain s q

X p inf CUGC (P) ≤ CUGC (UK ) = P

Z |Lk |

q(λ) dλ = Lk

k

X

f(ξk )|Lk |.

k

The quantity on the right-hand side is a Riemann sum for f; taking the limit as K → +∞ shrinks the intervals to zero, and so it converges to F⋆ . We have thus shown that inf P CUGC (P) ≤ F⋆2 . Combined with the lower bound (58), the proof is complete. 6.5.2

Proof of the claim (41b) Rℓ √ Recall our shorthand f = q and F⋆ = −ℓd d f(λ) dλ. For a pair −ℓd ≤ x < y ≤ ℓd , introduce the shorthand
eλ γ(x, y) := Γumask R(x), R(y) where R(λ) = 1+e λ . By changing variables in equation (50a), and using the derivative R′ (t) = R(t){1 − R(t)}, we obtain Z y R(y) − R(t) γ(x, y) = q(t) dt. (59a) R′ (t) x Define the constants m := min[−ℓd ,ℓd ] R′ > 0, M := maxλ∈[−ℓd ,ℓd ] R′ (λ) < ∞, and qmin := minλ∈[−ℓd ,ℓd ] q(λ) > 0. We then have R(y) − R(t) ≥ m(y − t) and R′ (t) ≤ M , from which equation (59a) implies that m where c0 := qmin 2M .

γ(x, y) ≥ c0 (y − x)2

(59b)

In addition, we claim that γ(x, x + δ) =

1 q(x)δ 2 + o(δ 2 ) 2 31

uniformly in x.

(59c)

Writing t = x + s, uniform continuity of R′ on [−ℓd , ℓd ], together with the fact that it is bounded away from zero over this interval, imply that R(x + δ) − R(x + s) = (δ − s){1 + o(1)} R′ (x + s) uniformly in x and 0 ≤ s ≤ δ. Similarly, the continuity of q gives q(x + s) = q(x) + o(1) uniformly. Substituting these two estimates into equation (59a) and integrating over s ∈ [0, δ] yields the local expansion bound (59c). Upper bound: For the upper bound, we choose the grid {λk,N }N k=0 to have equal f-mass, so that R λk+1,N F⋆ f(λ) dλ = N . Since f is continuous and bounded away from zero, this mesh has width O(N −1 ), λk,N F2

and hence we have q(λk,N )(λk+1,N − λk,N )2 = N⋆2 + o(N −2 ) uniformly in k. Equation (59c) therefore gives N −1 X

γ(λk,N , λk+1,N ) =

k=0

F⋆2 + o(N −1 ). 2N

(59d)

−2 Lower bound: For the lower bound, we take a grid {µk,N }N of the infimum, k=0 whose cost is within N and write δk,N = µk+1,N − µk,N . The preceding upper bound shows that the infimum is O(N −1 ), so the near-optimal grid {µk,N } also has total cost O(N −1 ). Consequently, the quadratic lower bound (59b) gives

c0

N −1 X

2 δk,N ≤

k=0

N −1 X

γ(µk,N , µk+1,N ) = O(N −1 ),

k=0

P 2 2 −1 ), the mesh tends to zero, so the and hence = O(N ). Since maxk δk,N ≤ k δk,N = O(N P 2
−1 uniform local expansion may be summed over the grid with total remainder o ). Doing k δk,N = o(N so yields (N −1 )2 N −1 N −1 X X (i) 1 1 X 2 −1 q(µk,N )δk,N + o(N ) ≥ f(µk,N )δk,N γ(µk,N , µk+1,N ) = + o(N −1 ) 2 2N 2 k δk,N

P

k=0

−1

k=0

k=0

2 (ii) F⋆

=

2N

+ o(N −1 ),

2 where step (i) follows P from the Cauchy–Schwarz inequality and q = f ; and the last step follows from the Riemann-sum limit k f(µk,N )δk,N = F⋆ + o(1).

Finally, combining this lower bound with the upper bound (59d) yields the claim.

7

Discussion

In this paper, we have shown that the unmasking growth complexity (UGC) controls the performance of random subset unmasking schemes, and reveals the structure required to optimize their performance. It is a path-based measure with additive local increments that bound the corresponding KL discretization errors. The reveal odds determine the natural multiplicative scale for a finite unmasking step, making log-revealodds the appropriate path coordinate. The log-reveal-odds UGC density gives a local description of sampling difficulty: regions carrying little UGC mass can be traversed rapidly with large stepsizes, whereas regions of concentrated mass require smaller stepsizes. We formalized this resource-allocation principle via a notion of partition complexity. At one extreme, it reduces to a (coarse) single-block complexity; under progressive refinement, it converges to a functional of the square root of the UGC density, which also determines the sharp leading-order complexity of optimally scheduled Euler discretization. 32

Notably, this paper moves beyond treating the underlying geometry as an oracle quantity. We show how the UGC increments can be reliably estimated using KL increments along a coupled reveal trajectory. By doing so, we obtain sampling methods that are certified-optimal, meaning that they are certified (with high probability) to achieve a pre-specified KL error, and their iteration complexity is within a constant factor of the optimal iteration complexity. Thus, our analysis connects three core tasks in diffusion sampling: characterizing the geometry of the target distribution, designing a sampling schedule, and certifying the accuracy of the resulting sampler. Let us make note of some important qualifications to this conclusion. Our strongest data-dependent guarantees rely on access to Bayes denoisers and on moment control for their KL increments. With learned denoisers, approximation error contributes an additional term to the KL guarantee, and it needs to be estimated for a fully certified guarantee. Moreover, our notion of optimality is relative to frozen-poster Euler discretizations. It is possible that higher-order discretizations and other more sophisticated procedures could admit different local geometries, and hence potentially different optimal clocks for stepsize allocation. This paper has focused exclusively on the utility of UGC complexity for designing and optimizing unmasking samplers. However, thinking beyond the context of sampling, the view of data geometry afforded by the UGC density is of independent interest for statistical analysis. As a notable example, Figure 3 shows the UGC density q for a discrete mixture model with a hierarchical structure, where the q-modes correspond to reveal times at which successive levels of the hierarchy are resolved. Our methods allow us to reliably estimate increments of the UGC density, and it would be interesting to do so for real-world discrete datasets. In our companion paper [Wai26] on Gaussian diffusion sampling, we introduced the denoising growth complexity (DGC), and provided analogous results on sampling performance and stepsize optimization, again culminating in a characterization via the DGC density. Despite the substantial differences between Gaussian diffusion and discrete masking, the two theories exhibit a striking parallel. The natural path coordinates differ (log-heat-time versus log-reveal-odds), but both the UGC and DGC densities capture a form of information curvature along the noising path, and their shape dictates how computational effort should be allocated. Acknowledgements This work was partially supported by a Guggenheim Fellowship, an NSF grant (DMS-2311072), and the Ford Professorship at MIT. We thank Yuting Wei for her inspiring talk during the MIT Statistics and Data Science conference in spring 2026.

References [AGK11]

D. F. Anderson, A. Ganguly, and T. G. Kurtz. Error analysis of tau-leap simulation methods. The Annals of Applied Probability, 21(6):2226–2262, 2011.

[AJH+ 21]

J. Austin, D. D. Johnson, J. Ho, D. Tarlow, and R. van den Berg. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, volume 34, 2021.

[AS16]

N. Alon and J. H. Spencer. The Probabilistic Method. John Wiley & Sons, 4 edition, 2016.

[BBB09]

L. Barnett, C. L. Buckley, and S. Bullock. Neural complexity and structural connectivity. Physical Review E, 79(5):051914, 2009.

[BDBDD24] J. Benton, V. De Bortoli, A. Doucet, and G. Deligiannidis. Nearly d-linear convergence bounds for diffusion models via stochastic localization. In International Conference on Learning Representations (ICLR) 2024, 2024. [BGJM11]

S. Brooks, A. Gelman, G. L. Jones, and X. L. Meng. Handbook of Markov Chain Monte Carlo. CRC Press, 2011.

[BZ12a]

J. Buzzi and L. Zambotti. Approximate maximizers of intricacy functionals. Probability Theory and Related Fields, 153(3–4):421–440, 2012.

[BZ12b]

J. Buzzi and L. Zambotti. Mean mutual information and symmetry breaking for finite random fields. Annales de l’Institut Henri Poincaré, Probabilités et Statistiques, 48(2):343–367, 2012.

33

[CBB+ 22]

A. Campbell, J. Benton, V. D. Bortoli, T. Rainforth, G. Deligiannidis, and A. Doucet. A continuous time framework for discrete denoising models. In Advances in Neural Information Processing Systems, volume 35, 2022.

[CCL+ 23]

S. Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. R. Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations, 2023.

[CCL25]

S. Chen, K. Cong, and J. Li. Optimal inference schedules for masked diffusion models. arXiv preprint arXiv:2511.04647, 2025.

[CHIS23]

F.-A. Croitoru, V. Hondru, R. T. Ionescu, and M. Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(9):10850–10869, 2023.

[CLL23]

H. Chen, H. Lee, and J. Lu. Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions. In International Conference on Machine Learning, 2023.

[CMFW24] M. Chen, S. Mei, J. Fan, and M. Wang. An overview of diffusion models: Applications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771, 2024. [CZB+ 23]

H. Chang, H. Zhang, J. Barber, A. Maschinot, J. Lezama, L. Jiang, M.-H. Yang, K. P. Murphy, W. T. Freeman, M. Rubinstein, Y. Li, and D. Krishnan. Muse: Text-to-image generation via masked generative transformers. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 4055–4075. PMLR, 2023.

[CZJ+ 22]

H. Chang, H. Zhang, L. Jiang, C. Liu, and W. T. Freeman. MaskGIT: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11315–11325, 2022.

[DHW26]

D. Dmitriev, Z. Huang, and Y. Wei. Efficient sampling with discrete diffusion models: Sharp and adaptive guarantees. arXiv preprint arXiv:2602.15008, 2026.

[GCS+ 13]

A. Gelman, J. Carlin, H. S. Stern, D. B. Dunson, A. Vehtari, and D. K. Salomon. Bayesian data analysis. CRC Press, 2013.

[Gil01]

D. T. Gillespie. Approximate accelerated stochastic simulation of chemically reacting systems. The Journal of Chemical Physics, 115(4):1716–1733, 2001.

[GLLZ19]

M. Ghazvininejad, O. Levy, Y. Liu, and L. Zettlemoyer. Mask-Predict: Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 6112–6121. Association for Computational Linguistics, 2019.

[Han75]

T. S. Han. Linear dependence structure of the entropy space. Information and Control, 29(4):337–368, 1975.

[Han78]

T. S. Han. Nonnegative entropy measures of multivariate symmetric correlations. Information and Control, 36(2):133–156, 1978.

[HJA20]

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems (NeurIPS), 33:6840–6851, 2020.

[HNJ+ 21]

E. Hoogeboom, D. Nielsen, P. Jaini, P. Forré, and M. Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances in Neural Information Processing Systems, volume 34, 2021.

[LC25]

G. Li and C. Cai. Breaking AR’s sampling bottleneck: Provable acceleration via diffusion language models. In Advances in Neural Information Processing Systems, volume 38, 2025.

[LHL+ 25]

Y. Liang, R. Huang, L. Lai, N. Shroff, and Y. Liang. Absorb and converge: Provable convergence guarantee for absorbing discrete diffusion models. In Advances in Neural Information Processing Systems, volume 38, 2025.

[Li07]

T. Li. Analysis of explicit tau-leaping schemes for simulating chemically reacting systems. Multiscale Modeling & Simulation, 6(2):417–436, 2007.

[LLLS25]

Y. Liang, Y. Liang, L. Lai, and N. Shroff. Discrete diffusion models: Novel analysis and new sampler guarantees. In Advances in Neural Information Processing Systems, volume 38, 2025.

34

[LME24]

A. Lou, C. Meng, and S. Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 32819–32848. PMLR, 2024.

[Lor86]

G. G. Lorentz. Bernstein Polynomials. Chelsea Publishing Company, New York, second edition, 1986.

[LZ25]

H. Lavenant and G. Zanella. Error bounds and optimal schedules for masked diffusions with factorized approximations. arXiv preprint arXiv:2510.25544, 2025.

[MP09]

A. Maurer and M. Pontil. Empirical Bernstein bounds and sample variance penalization. In Proceedings of the 22nd Conference on Learning Theory, 2009.

[NZY+ 25]

S. Nie, F. Zhu, Z. You, X. Zhang, J. Ou, J. Hu, J. Zhou, Y. Lin, J.-R. Wen, and C. Li. Large language diffusion models. In Advances in Neural Information Processing Systems, volume 38, 2025.

[OBAJ08]

E. Olbrich, N. Bertschinger, N. Ay, and J. Jost. How should complexity scale with system size? The European Physical Journal B, 63(3):407–415, 2008.

[Phi03]

G. M. Phillips. Interpolation and Approximation by Polynomials. Springer-Verlag, New York, 2003.

+

[RBD 22]

R. Rombach, E. Blattmann, S. L. Dhariwal, A. M. D. M. L., and P. E. S. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10684–10694, 2022.

[RC04]

C. P. Robert and G. Casella. Monte Carlo Statistical Methods. Springer Texts in Statistics. Springer, New York, 2 edition, 2004.

[RCRY25]

Y. Ren, H. Chen, G. M. Rotskoff, and L. Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via a stochastic integral framework. In International Conference on Learning Representations, 2025.

[RCZ+ 25]

Y. Ren, H. Chen, Y. Zhu, W. Guo, Y. Chen, G. M. Rotskoff, M. Tao, and L. Ying. Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. In Advances in Neural Information Processing Systems, volume 38, 2025.

[RK08]

R. Y. Rubinstein and D. P. Kroese. Simulation and the Monte Carlo Method. John Wiley and Sons, Hoboken, NJ, 2nd edition, 2008.

[RMGJ19]

F. E. Rosas, P. A. M. Mediano, M. Gastpar, and H. J. Jensen. Quantifying high-order interdependencies via multivariate extensions of the mutual information. Physical Review E, 100(3):032305, 2019.

[RPCG05]

M. Rathinam, L. R. Petzold, Y. Cao, and D. T. Gillespie. Consistency and stability of tau-leaping schemes for chemical reaction systems. Multiscale Modeling & Simulation, 4(3):867–895, 2005.

[SAS+ 24]

S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin, J. T. Chiu, A. Rush, and V. Kuleshov. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems, volume 37, 2024.

[SE19]

Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (NeurIPS), pages 11895–11907, 2019.

[SHW+ 24]

J. Shi, K. Han, Z. Wang, A. Doucet, and M. K. Titsias. Simplified and generalized masked diffusion for discrete data. In Advances in Neural Information Processing Systems, volume 37, 2024.

[SSDK+ 21] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. [SWMG15] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256–2265. PMLR, 2015. [TSE94]

G. Tononi, O. Sporns, and G. M. Edelman. A measure for brain complexity: Relating functional segregation and integration in the nervous system. Proceedings of the National Academy of Sciences of the United States of America, 91(11):5033–5037, 1994.

[VPFS23]

T. F. Varley, M. Pope, J. Faskowitz, and O. Sporns. Multivariate information theory uncovers synergistic subsystems of the human cerebral cortex. Communications Biology, 6:451, 2023.

[Wai26]

M. J. Wainwright. Denoising growth complexity: Data geometry and certified schedules for diffusion sampling. Technical Report arXiv:2607.26285, Massachusetts Institute of Technology, 2026.

35

[Wat60]

S. Watanabe. Information theoretical analysis of multivariate correlation. IBM Journal of Research and Development, 4(1):66–82, 1960.

[WZL+ 25]

Y. Wang, H. Zhan, L. Liu, R. Zeng, H. Guo, J. Zheng, Q. Zhang, X. Zhang, S. Zhang, and Z. Wu. MaskGCT: Zero-shot text-to-speech with masked generative codec transformer. In International Conference on Learning Representations, 2025.

[WZY+ 24]

X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu. Diffusion language models are versatile protein learners. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 52309–52333. PMLR, 2024.

[YCS+ 23]

L. Yu, Y. Cheng, K. Sohn, J. Lezama, H. Zhang, H. Chang, A. G. Hauptmann, M.-H. Yang, Y. Hao, I. Essa, and L. Jiang. MAGVIT: Masked generative video transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10459–10469, 2023.

[YZSea25]

L. Yang, Z. Zhang, Y. Song, and et al. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys (to appear), 2025. arXiv:2209.00796.

A

Consequences of single-block theory

In this appendix, we discuss some consequences of Corollary 1, which gives single-block stepsize schedules with iteration complexity depending on the aggregate UGC masses H(0, 1) and Hcard (0, d), for the Bernoulli and fixed-cardinality samplers, respectively. For this reason, it is interesting to relate the aggregate UGC mass to other complexity measures, and thereby to relate these complexity measures to the sampling guarantees in the papers [CCL25, DHW26]. We begin by observing that the aggregate UGC mass H(0, 1) turns out to be equivalent to a measure of multivariate dependence first introduced by Tononi, Sporns, and Edelman [TSE94] for modeling brain function. This complexity measure, hereafter referred to as TSE, has since been studied by various authors [OBAJ08, BBB09, RMGJ19, VPFS23], and also extended to a more general notion of intricacy [BZ12a, BZ12b]. The TSE complexity is given by TSE(PZ ) :=

d  X k=1

k ek − ed d



  where ek := E Ent(ZAk ) ,

(60)

with Ak being a uniformly random subset of [d] with k elements. We establish its equivalence to the aggregate UGC mass as part of Proposition 3 below. Turning to sampling connections, the paper [CCL25] analyzed the fixed-cardinality sampler, and gave guarantees in terms of the minimum of two classical measures of multivariate dependence [Wat60, Han75, Han78], known as the total correlation and dual total correlation, defined by TC(PZ ) :=

d X

Ent(Zi ) − Ent(Z)

and

DTC(PZ ) := Ent(Z) −

i=1

d X

Ent(Zi | Z−i ).

(61a)

i=1

Here Ent denotes the Shannon entropy, whereas Z−i denotes the (d − 1)-dimensional subvector of all coordinates except the ith . The single-block sampling guarantee in terms of Hcard (0, d), or equivalently H(0, 1), sharpens these guarantees for the fixed-cardinality sampler. As we show in Section B.2, in particular via a Bernstein polynomial representation, the total and dual total correlation measures have the simple representations Z 1 Z 1 ′ TC(PZ ) = (1 − t)h (t)dt and DTC(PZ ) = th′ (t)dt. (61b) 0

Since H(0, 1) =

R1 0

0

t (1 − t)h′ (t)dt, we see immediately that H(0, 1) ≤ min{TC(PZ ), DTC(PZ )}. 36

(61c)

In their work on continuous-time Markov chain (CTMC) sampling, Dmitriev et al. [DHW26] gave a CTMC unmasking sampler with iteration complexity controlled by a functional DHW(PZ ), which they refer to as effective total correlation, that is also upper bounded by the minimum of the TC and DTC complexities. We also show as part of Proposition 3 that this functional is lower bounded by H(0, 1), and can be upper bounded by a constant multiple of it. Thus, both the Bernoulli and fixed-cardinality samplers, in their singleblock instantiation, inherit the guarantees of the DHW analysis. In particular, their paper studies various interesting examples, including stochastic block models and quantized versions of low-dimensional structure, for which DHW(PZ ), and hence H(0, 1) and Hcard (0, d), are relatively small. Notably, they construct an instance with min{TC(PZ ), DTC(PZ )} = Θ(d), whereas their complexity has constant scaling. Our analysis shows that the guarantees on CTMC-based unmasking [DHW26] have analogues for both fixed-cardinality unmasking [CCL25, LZ25] and Bernoulli unmasking. With this context, we now formally state the complexity relations: Proposition 3 (Connections between aggregate complexity measures). For the full interval [0, 1], the Bernoulli aggregate-UGC complexity is related to TSE complexity via the relation H(0, 1) =

2 TSE(PZ ). d+1

(62a)

Moreover, the fixed-cardinality complexity Hcard (0, d) is sandwiched by the Bernoulli version H(0, 1) ≤ Hcard (0, d) ≤ 2H(0, 1).

(62b)

Finally, the DHW complexity satisfies the sandwich relation (iii)

(iv)

H(0, 1) ≤ DHW(PZ ) ≤

e H(0, 1). e−1

(62c)

We prove these claims in Section C.1; let us highlight a few interesting features here. The proofs of both claims (62a) and (62b) make use of a Bernstein polynomial representation of h, a result of independent interest. In particular, we show in Section B.2 that h(t) =

 d−1  X d−1 j=0

j

tj (1 − t)d−1−j hcard j ,

(63)

so that h is a polynomial of degree at most d − 1. In proving the sandwich (62c), we show that the DHW complexity functional has the explicit represenR1 tation DHW(PZ ) = 0 t min{1, log(1/t)}h′ (t)dt, from which the estimates (62c) follow from simple algebra. The analysis leading to this representation reveals an interesting algorithmic connection to Bernoulli unmasking: the DHW unmasking algorithm, while described and analyzed as a form of τ -leaping [Gil01] in a continuous-time Markov chain, is closely related to the Bernoulli unmasking sampler. See Section C.2 for details of this relation.

B

Alternative representations of h

In this section, we develop and explore some alternative representations of h and its derivative h′ that provide useful insight for subsequent proofs.

37

B.1

Proof of the identity (15a)

In this appendix, we prove the random subset identity (15a), namely that we can write h(t) =

d X

  EAi,t Info(Zi ; ZAi,t ) ,

i=1

where Ai,t ⊆ [d] \ {i} is a random subset obtained by including each coordinate of [d] \ {i} independently with probability t. To prove this representation, for each i ∈ [d], let Ei denote the event that coordinate i remains masked at reveal time t, and let A ≡ Ai,t ⊆ [d] \ {i} denote the set of revealed coordinates. Conditional on Ei , the observation Xt is equivalent to the
pair (A, ZA ). Consequently, we have the mutual information identity Info(Zi ; Xt | Ei ) = Info Zi ; A, ZA . By the chain rule for mutual information, we have

Info Zi ; A, ZA = Info(Zi ; A) + Info Zi ; ZA | A , | {z } =0

using the fact that Zi and A ≡ Ai,t are independent. Conditioning on A = a does not change the joint law of (Zi , ZA ), so that we have shown that X   Info(Zi ; Xt | Ei ) = Info(Zi ; ZA | A) = P(A = a) Info(Zi ; Za ) = EA Info(Zi ; ZA ) for each i = 1, . . . , d. a

Summing this identity over the coordinate index i yields the claim (15a).

B.2

Representation via Bernstein polynomials

In this section, we develop representations of both h and h′ as Bernstein polynomials in terms of the coefficients hcard j . These representations have some useful immediate consequences, and are used in the proofs in Section C. Recall the definition of the coefficients := hcard j

d X

  EBi,j Info(Zi ; ZBi,j )

where Bi,j is uniform over j-cardinality subsets of [d] \ {i}.

(64)

i=1

Letting Ai,t be a Bernoulli random subset of [d] \ {i} with inclusion probability t, we can write hcard = j

d h i X E Info(Zi ; ZAi,t ) | card(Ai,t ) = j . i=1

j d−1−j Since P[card(Ai,t ) = j] = d−1 , it follows by combining the tower property with the represenj t (1 − t) tation (15a) that h(t) =

 d−1  X d−1 j=0

j

tj (1 − t)d−1−j hcard j .

(65a)

Moreover, by differentiating equation (65a), we obtain h′ (t) = (d − 1)

 d−2  X  d−2 j card t (1 − t)d−2−j hcard . j+1 − hj j j=0

(65b)

This calculation uses standard properties of derivatives of Bernstein polynomials [Lor86, Phi03]. (In particular, differentiation transforms coefficients into their first differences.) As one immediate consequence, from equation (65a), we see that h is a polynomial of degree at most d − 1. It follows immediately that the derivative h′ exists. Let us develop two additional consequences. 38

B.2.1

Proof of h-monotonicity

We now prove that h is monotone by showing that h′ (t) ≥ 0. The Bernstein polynomial representation (65b) card of h′ shows that it suffices to prove that hcard j+1 ≥ hj . For 0 ≤ j ≤ d − 2 and a fixed coordinate i, let Aj be a random subset, chosen uniformly from all the j-element subsets of [d] \ {i}, and, conditionally on Aj , let J be uniform over its complement inside [d] \ {i}. This defines a coupling between the subsets. Then the random set Aj ∪ {J} is uniform over the (j + 1)element subsets, while the chain rule for mutual information gives the decomposition Info(Zi ; ZAj ∪{J} ) − Info(Zi ; ZAj ) = Info(Zi ; ZJ | ZAj ) ≥ 0. Averaging first over the coupling and then over i shows that card ≥ 0, as claimed. hcard j+1 − hj B.2.2

Proof of the total-correlation representations (61b)

Now let us prove the representations of the total correlation TC and dual total correlation DTC given in equation (61b). From the random subset representation (15a) of h, we have h(1) =

d X

Info(Zi ; Z−i ) =

i=1

d X 

Ent(Zi ) − Ent(Zi | Z−i ) ,

i=1

where the second equality expands the definition of mutual information. ByR comparing with the defini1 tions (61a), we see immediately that h(1) = TC(PZ ) + DTC(PZ ). Since h(1) = 0 h′ (t)dt, using the fact that R1 R1 h(0) = 0, if we can prove the identity TC(PZ ) = 0 (1 − t)h′ (t)dt, then it follows that DTC(PZ ) = 0 th′ (t)dt. Accordingly, the remainder of our analysis focuses on the TC representation. For j = 1, . . . , d, let ej be the entropy coefficient defined in equation (60), along with e0 = 0. Note that Pd we have e1 = d1 i=1 Ent(Zi ), and ed = Ent(Z). For any i ∈ [d] and A ⊆ [d] \ {i}, we have the mutual information identity Info(Zi ; ZA ) = Ent(Zi ) + Ent(ZA ) − Ent(ZA∪{i} ). It can be averaged over i and over coefficients, we find that uniform j-cardinality subsets A. Doing so and using the definition (15b) of the hcard j
hcard = d e1 + ej − ej+1 . j

(66)

Substituting this identity into the Bernstein representation (65a) and integrating over the interval [0, 1] yields  Z d−1  d−1 X d − 1 card  1 j (i) 1 X card (ii) (iii) t (1 − t)d−1−j dt} = h(t) dt = hj = de1 − ed = TC(PZ ), hj j d 0 0 j=0 j=0

Z 1

j R1 d−1−j where step (i) uses the beta-integral identity 0 d−1 dt = d1 ; step (ii) uses the identity (66) j t (1 − t) along with some algebra; and step (iii) follows from the definition (61a) of the total correlation, combined with the definition (60) of the entropy coefficients. R1 R1 Finally, integrating by parts shows that 0 h(t)dt = 0 (1 − t)h′ (t)dt, where the boundary term vanishes since h(0) = 0. This completes the proof of the TC representation in equation (61b).

C

Proofs of connecting results

In this section, we prove various results that connect the UGC aggregate complexity to other complexity notions (Section C.1), and compare the DHW τ -leaping unmasking algorithm with the Bernoulli unmasking procedure (Section C.2).

C.1

Proof of Proposition 3

The proposition consists of three claims, and we prove each of them in turn.

39

C.1.1

Proof of TSE identity (62a)

R1 Beginning with the definition (2c), integration by parts yields H(0, 1) = 0 (2t−1)h(t) dt, where the boundary term vanishes since t(1 − t) = 0 at the endpoints. Substituting the Bernstein representation (65a) of h into this equation yields   Z 1  d−1 d−1 X X 1 card d − 1 j d−1−j (2t − 1)t (1 − t) (2j − d + 1) hcard H(0, 1) = dt = (67) hj j , d(d + 1) j 0 j=0 j=0 where the second step follows from a standard beta-integral identity. To complete the proof, we need to relate the entropy coefficients ej = E[Ent(ZAj )], where Aj is a coefficients. For each j = 0, . . . , d − 1, we can write the uniformly random subset of cardinality j, to the hcard j mutual information as Info(Zi ; ZA ) = Ent(Zi ) + Ent(ZA ) − Ent(ZA∪{i} ). Averaging this identity over i and
the j-element subsets A ⊆ [d] \ {i} yields the relation hcard = d e + e − e . Substituting this identity 1 j j+1 j into the expression (67) yields   d−1 d−1 d−1  X 1  X 1 n X e1 (2j − d + 1) + (2j − d + 1)(ej − ej+1 ) = 2 ej − (d − 1)ed H(0, 1) =  d+1 d + 1  j=1 j=0 j=0 | {z } =0

 d  2 X j 2 = ej − ed = TSE(PZ ), d + 1 j=1 d d+1 which establishes the claim (62a). C.1.2

Proof of the bound (62b)

For any 0 ≤ p < q ≤ 1, define the multinomial triple (U0 , U1 , U2 ) ∼ Multinomial (d + 1; p, q − p, 1 − q). As claimed previously, we have the identity H(p, q) =

  d E Hcard (AU , BU ) , d+1

(68)

where AU := max{U0 − 1, 0}, and BU := min{U0 + U1 , d}. The claims (62b) follow as a special case. For (p, q) = (0, 1), we have (AU , BU ) = (0, d) almost surely, d Hcard (0, d), from which it follows that H(0, 1) ≤ Hcard (0, d) ≤ 2H(0, 1), as claimed. and hence H(0, 1) = d+1
j j We now prove the claim (68). Introduce the shorthand notation ∆j := hcard − hcard j j−1 and cj := d 1 − d , and observe that b−1 X

Hcard (a, b) =

cj ∆j .

(69a)

j=a+1

Moreover, using the representation (65b) of h′ as a Bernstein polynomial, we can write   d−1 X d j t(1 − t)h′ (t) = cj ∆j d t (1 − t)d−j . j j=1

(69b)

Generate d + 1 i.i.d. Unif[0, 1] random variables Yi , and let Y(j+1) denote the (j + 1)-st order statistic. Its
density is given by (d + 1) dj tj (1 − t)d−j . Consequently, by integrating the Bernstein expansion (69b) over [p, q], we obtain d−1
d X H(p, q) = cj ∆j P Y(j+1) ∈ (p, q) . d + 1 j=1

40

(69c)

Now recalling our i.i.d. uniform samples {Yi }d+1 i=1 , define the multinomial count vector U = (U0 , U1 , U2 ) via U0 := card{ℓ | Yℓ ∈ [0, p]},

U1 := card{ℓ | Yℓ ∈ (p, q)}

and

U2 := card{ℓ | Yℓ ∈ [q, 1]}.

By construction, we have Y(j+1) ∈ (p, q) if and only if AU < j < BU . Thus, we can rewrite the expression (69c) as     BX d−1 U −1 X d d E E H(p, q) = cj ∆j 1{AU < j < BU } = cj ∆ j  d+1 d + 1 j=1 j=AU +1

=

  d E Hcard (AU , BU ) , d+1

where the last step uses the representation (69a). This completes the proof of equation (68). C.1.3

Proof of the DHW sandwich (62c)

For u ≥ 0, we define a reveal time r = e−u , so that u = log(1/r) is the log-inverse-reveal time. The analysis of Dmitriev et al. [DHW26] involves a forward masked vector Yu ; it is related to our reveal process Xr via law

the relation Yu = Xe−u . In the process Yu , each coordinate Yui equals Zi with probability r = e−u and equals ⋆ otherwise, and they define the conditional-information density X I(u) := Info(Yui ; Yuj | Yu−(i,j) ), i̸=j −(i,j)

denotes the collection where Yu R ∞of unmasked coordinates other than i and j. They define the effective total correlation by DHW(PZ ) := 0 min{1, u}I(u) du. In order to prove the sandwich relation, it suffices to prove the identity I(u) = e−2u h′ (e−u )

for all u ≥ 0.

Indeed, when this identity holds, the change of variables t = e−u guarantees that Z 1 DHW(PZ ) = w(t)h′ (t) dt, where w(t) := t min{1, − log t}.

(70a)

(70b)

0

In terms of the shorthand v(t) := t (1 − t), it is straightforward to verify that v(t) ≤ w(t) ≤

e v(t) e−1

for all t ∈ [0, 1],

from which the sandwich (62c) follows. Proof of the identity (70a): We begin by relating their definitions to equivalent objects in our notation. Fix an ordered pair of indices i ̸= j and let Ai,j r be the random subset of [d] \ {i, j} obtained by retaining each coordinate independently with probability r. If either i or j is masked, one of the first two arguments in the conditional mutual information is deterministic, and its contribution is zero. The probability that both coordinates are retained is r2 . Conditional on this event and on Ai,j r = A, the three arguments reduce to Zi , Zj , and ZA , respectively. Since the masking pattern is independent of Z, the law of total expectation gives   Info(Yui ; Yuj | Yu−(i,j) ) = r2 E Info(Zi ; Zj | ZAi,j ) . r

41

card(A) Moreover, P(Ai,j (1 − r)d−2−card(A) . Expanding the expectation and summing over the r = A) = r ordered pairs therefore yields X X I(u) := r2 rcard(A) (1 − r)d−2−card(A) Info(Zi ; Zj | ZA ) with r = e−u . (71a) i̸=j A⊆[d]\{i,j}

In order to complete the proof, it suffices to show that h′ satisfies the identity X X rcard(A) (1 − r)d−2−card(A) Info(Zi ; Zj | ZA ). h′ (r) =

(71b)

i̸=j A⊆[d]\{i,j}

Our claim (70a) then follows by comparing the two representations (71b) and (71a). To prove the identity (71b), fix an integer k ∈ {0, . . . , d − 2}. Beginning with the definition (15b) of hcard k and averaging over the choice of the additional coordinate, we find that   n o X X d − 2  card (d − 1) hk+1 − hcard = Info(Zi ; ZA∪{j} ) − Info(Zi ; ZA ) , k k {z } i̸=j A⊆[d]\{i,j} | Info(Zi ;Zj |ZA )

card(A)=k

where the underbrace relation follows from the chain rule for mutual information. Substituting this identity into our Bernstein representation (65b) for h′ , and then summing over k yields equation (71b).

C.2

Relation with the DHW τ -leaping masking sampler

In this section, we formalize the exact relationship between Algorithm 1 of Dmitriev et al. [DHW26] and the Bernoulli unmasking sampler (11). Any DHW grid of the form 0 = u0 < · · · < uN = U can be used to define a sequence of reveal times tj := e−(U −uj ) , and the corresponding inverse transformation U − uj = log(1/tj ). In terms of the shorthand βk := (tk+1 − tk )/(1 − tk ), for the Bernoulli unmasking update with estimated denoiser µ bi (·, x), the corresponding transition kernel is given by KBer k,i (⋆ | x) = 1 − βk ,

and

KBer bi (a, x) k,i (a | x) = βk µ

for each a ∈ A.

(72a)

On the other hand, for a current state x and masked coordinate i ∈ M(x), write the DHW estimated score DHW as qk,i (a, x) := bslog(1/tk ) (xi→a , x). Assuming its total mass is positive, define its normalized shape and relative mass by πk,i (a, x) := P

DHW (a, x) qk,i DHW b∈A qk,i (b, x)

,

and

ck,i (x) :=

1 − tk X DHW qk,i (b, x). tk

(72b)

b∈A

These ingredients along with the DHW stepsize yield the kernel ck,i (x) KDHW , k,i (⋆ | x) = (1 − βk )

and

n o ck,i (x) KDHW (a | x) = 1 − (1 − β ) πk,i (a, x). k k,i

(72c)

Taking the differences between the expressions (72a) and (72c), we obtain the kernel differences Ber ck,i (x) KDHW − (1 − βk ), k,i (⋆ | x) − Kk,i (⋆ | x) = (1 − βk ) n o Ber ck,i (x) KDHW πk,i (a, x) − βk µ bi (a, x). k,i (a | x) − Kk,i (a | x) = 1 − (1 − βk )

(72d) (72e)

Thus, the two kernels are very closely related, but not identical in general. DHW ⋆ We claim that when there is no score error, so that qk,i = qk,i and µ bi = µi , then the two kernels coincide. First of all, we claim that Proposition 6 of Dmitriev et al. [DHW26], when translated into our notation, shows that X tk tk (ii) (i) ⋆ ⋆ µi (a, x), and qk,i (a, x) = . (72f) qk,i (a, x) = 1 − tk 1 − tk a∈A

42

Equation (ii) follows from equation (i), since the denoiser µi is a conditional probability distribution. To establish equation (i), if x has coordinate i masked and y is obtained by filling that coordinate with a, then the marginal convention in Proposition 6 gives q0 (y)/q0 (x) = P(Zi = a | x) = µi (a, x). Moreover, their CTMC time u is related to our reveal probability by the transformation tk = e−u , so that (eu − 1)−1 = tk /(1 − tk ). Using equation (72f) we see that the exact score has total mass tk /(1 − tk ), so that ck,i (x) = 1 in ⋆ (a, x) by this total mass cancels the common factor tk /(1 − tk ), equation (72b). Moreover, normalizing qk,i so that πk,i (a, x) = µi (a, x). Thus, the right-hand side of equation (72d) is (1 − βk ) − (1 − βk ) = 0, while the right-hand side of equation (72e) is βk µi (a, x) − βk µi (a, x) = 0. For learned scores, the mask-probability gap is caused by the non-unit score mass ck,i (x). The important difference is that rescaling the DHW score changes its kernel, whereas the Bernoulli reveal probability remains fixed by the grid.

D

Proof of Lemma 2: XORSAT scaling

Let us first complete the proof, taking the bounds p (44) as given. Note that the outer blocks have log-revealodds length O(log d), whereas Len(Itrans ) = O( log d/d). Hence, we can compute 2 p p Len(Ipre )H(Ipre ) + Len(Itrans )H(Itrans ) + Len(I )H(I ) "r #2   1/4 √ log d log d + = O d  d10 d p
= O d log d ,

CUGC (P) =

p

post

post

where we have used the claimed bound (44). Combining this bound with the scaling CUGC ([Ifull ]) = Θ(d log d) establishes the claim (43).  We now return to prove the bounds (44). With the shorthand Itrans = [sleft , s ] and c′ = min sleft (1 − sleft ), s (1 − s )}, it follows from the definition (2c) of H that right

right

right

h(1) − h(0) , as well as 4 h(1) − h(s ) and H(I ) ≤ . 4

c′ {h(s ) − h(sleft )} ≤ H(Itrans ) ≤ right

H(Ipre ) ≤

h(sleft ) − h(0) 4

right

post

(73a) (73b)

Consequently, we can control the UGC increments by controlling the differences in h. From the definition Pd h(1) = i=1 Info(Zi ; Z−i ) and the binary nature of each Zi , it follows that h(1) ≤ d log 2, and hence that h(1) − h(0) = O(d). For controlling the other increments, we study averages of h over the random choice of A ∈ {0, 1}d×k . The increments of the averaged h take a simple form. For a positive integer m ≤ k, let Bm ∈ {0, 1}m×k be a Boolean random matrix with m independent uniform rows in {0, 1}k , and let b be a fresh independent uniform row. Defining the function g(m) := P{b ∈ / rowspan(Bm )}

(74a)

we claim that   EA [h(s) − h(t)] = E g(Mt ) − g(Ms ) , d log 2

where Mr ∼ Bin(d − 1, r) is a binomial RV,

(74b)

and also E[g(Msleft )] = 1 − O(d−11 )

and

E[g(Ms

right

)] = O(d−11 ).

Taking these two claims as given for the moment, let us complete the proof of the bounds (44). 43

(74c)

  Transition interval: Equation (74c) gives E g(Msleft ) − g(Ms ) = 1 − O(d−11 ) = Θ(1). Consequently, equation (74b) gives EA [h(s )−h(sleft )] = Θ(d), and combining with inequality (73a), this yields E[H(Itrans )] = Θ(d). right

right

Pre- and post-transition intervals: For the pre-transition interval, M0 = 0 deterministically, so that  equation (74b) gives EA [h(sleft )] − EA [h(0)] = d(log 2) g(0) − E[g(Msleft )] . Since g is non-increasing and 0 ≤ g ≤ 1, it follows that 0 ≤ g(0) − E[g(Msleft )] ≤ 1 − E[g(Msleft )] = O(d−11 ), and hence EA [h(sleft )] − EA [h(0)] = O(d−10 ). Similarly, we have M1 = d − 1 deterministically, and equation (74b) gives EA [h(1)] − EA [h(s )] =  d(log 2) E[g(Ms )] − g(d − 1) . Using monotonicity and non-negativity, right

right

0 ≤ E[g(Ms

right

)] − g(d − 1) ≤ E[g(Ms

right

)] = O(d−11 ),

and hence EA [h(1)] − EA [h(s )] = O(d−10 ). Combining these two estimates with inequalities (73b), we conclude that E[H(Ipre )] = O(d−10 ) and E[H(I )] = O(d−10 ). right

post

To pass from these averaged bounds to a single realization of A, first note that the two outer H terms are nonnegative. Hence, by Markov’s inequality and the preceding expectation bounds, with probability −10 ). For the middle interval, bounded below by a positive constant, both H(Ipre ) and H(I ) are O(d  the endpoint estimates (74c) and identity (74b) give EA h(s ) − h(sleft ) = d(log 2) 1 − O(d−11 ) . Since 0 ≤ h(s ) − h(sleft ) ≤ d log 2, Markov’s inequality applied to the nonnegative deficit d log 2 − {h(s ) − h(sleft )} shows that   d log 2 PA h(s ) − h(sleft ) < = O(d−11 ). 2 post

right

right

right

right

Thus, by a union bound, for all sufficiently large d there exists a realization of A for which the two outer bounds hold and h(s ) − h(sleft ) = Ω(d) simultaneously. Inequality (73a) then gives H(Itrans ) = Θ(d) for this same realization. right

Proof of the identity (74b): Fix a coordinate i, and condition on coordinate i being masked at reveal probability r. Let R ⊆ [d] \ {i} denote the set of revealed coordinates among the remaining coordinates. Then |R| = Mr ∼ Bin(d − 1, r). Coordinate i is associated with the ith row ai of A, and we use AR to denote the matrix formed by the rows indexed by R. All of the following linear-algebraic statements are over the Boolean field F2 . Conditional on A and the revealed values ZR , there are two possibilities. If ai ∈ rowspan(AR ), then Zi is a linear combination of the revealed coordinates ZR , and hence is determined by them. On the other hand, if ai ∈ / rowspan(AR ), then the linear functional defined by ai is nonzero on the nullspace of AR . Consequently, among the latent assignments compatible with ZR , exactly half give Zi = 0 and half give Zi = 1. By this reasoning, the conditional entropy takes the simple form  Ent(Zi | ZR , A) = (log 2) 1 ai ∈ / rowspan(AR ) . (75) Now condition on the event |R| = m. By independence and permutation invariance of the rows, the submatrix AR has the same law as Bm , while ai is a fresh independent uniform row. From the definition (74a)  of g, it follows that EA Ent(Z  i | ZR , A) |R| =  m = (log 2)g(m).   Averaging this identity over the binomial random variable yields E Ent(Z | Z , A) = (log 2)E g(M ) . The same argument with no revealed A i R r   coordinates gives EA Ent(Zi | A) = (log 2)g(0). Combining the pieces, the ith term in the sum (2a) that defines h has expectation    EA Ent(Zi | A) − Ent(Zi | ZR , A) = (log 2) g(0) − E[g(Mr )] .  Summing over the d exchangeable coordinates yields the identity EA [h(r)] = d(log 2) g(0) − E[g(Mr )] . Taking the difference between r = s and r = t in this identity yields the claim (74b). 44

Proof of equation (74c):

We first claim that g satisfies the bounds

(a)

1 − g(m) ≤ 2m−k

(b)

for m ≤ k,

and

g(m) ≤ 2k−m

for m > k.

(76)

To prove the first bound (76)(a), we observe that conditional on Bm , its row span contains 2rank(Bm ) vectors out of the 2k possible rows in {0, 1}k . Hence a fresh uniform row lies in this span with probability at most 2rank(Bm )−k ≤ 2m−k . Combining with the definition of g yields the claim. As for the bound (76)(b), when m > k, the event that the fresh row lies outside the row span requires  rank(Bm ) < k. Union-bounding over nonzero right-null vectors therefore gives g(m) ≤ P rank(Bm ) < k ≤ (2k − 1)2−m < 2k−m , as claimed. We now show that E[g(Msleft )] = 1 − O(d−11 ). Recalling that Msleft ∼ Bin(d − 1, sleft ), its mean is r p log d k := µ− E[Msleft ] = (d − 1)sleft = k − − c1 d log d + c1 . d d √ For the threshold T− := k − (c1 /4) d log d, we therefore have r 3c1 p k log d c1 p T− − µ− = d log d + − c1 ≥ d log d 4 d d 2 for all sufficiently large d. Defining the bad event E = {Msleft > T− }, Hoeffding’s inequality gives   2 2(T− − µ− )2 ≤ d−c1 /2 , P(E) ≤ exp − d−1 so that P(E) = O(d−11 ) for c1 sufficiently large. On the complement E c , the bound (76)(a) guarantees that √ Msleft −k −(c1 /4) d log d 1 − g(Msleft ) ≤ 2 ≤2 . Since 0 ≤ g(m) ≤ 1, it follows that √

1 − E[g(Msleft )] ≤ P(E) + 2−(c1 /4) d log d = O(d−11 ), which proves the claim.  A similar argument  at the right endpoint s . The same Hoeffding bound guarantees that √ applies P Ms < k + (c1 /4) d log d = O(d−11 ), while on the complementary event, we have the bound right

right

g(Ms

right

Taking expectations gives E[g(Ms

right

) ≤ 2k−Ms

√

right

≤ 2−(c1 /4) d log d .

)] = O(d−11 ).

Combining the two pieces completes the proof of the claimed bounds (74c).

45

Related documents

• Recursive Scaling in Masked Diffusion Models#1129307
• Local Diffusion Models and Phases of Data Distributions#158764
• Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration#924125
• Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration#919896
• Diffusion contrast induced dynamics in a time-periodic competition-diffusion system with equal total resources#620354