Schedule optimization for tau-leaping in masked discrete diffusion
arXiv:2609.21960v1 [math.ST] 18 Sep 2026
Cecilia Secchi∗ and Giacomo Zanella†
Abstract Masked discrete diffusion models are commonly accelerated using the so-called tau-leaping discretization method, which reveals several coordinates in parallel at each sampling step. The sampler replaces the joint conditional law of each revealed block by a product distribution, incurring a factorization error εfact present even with perfectly learned predictors. We analyze the standard sampler on N coordinates with K sampling steps, whose random block sizes depend on a denoising schedule. Our analysis uses an exact integral representation of εfact in terms of a distribution-dependent dependence density ρ, which records how conditional dependence evolves as the revealed fraction of coordinates grows. We develop estimators for this profile and quantify how estimation errors affect schedule selection. We derive recursive stationarity equations for the finite-K optimization problem and, under a monotonicity condition, characterize its unique optimizer. In the joint limit N, K → ∞, we obtain an explicit characterization of the optimal limiting smooth schedule and quantify the cost of random block sizes relative to a deterministic planner. When ρN converges uniformly to a strictly positive continuous profile, optimizing over fixed smooth schedules can improve the leading constant but not the N/K scaling of εfact . By contrast, if ρN degenerates, suitable schedules can improve the asymptotic order relative to the uniform schedule. Examples based on stationary processes and exchangeable mixtures illustrate these two regimes.
1
Introduction
Discrete diffusion models (Austin et al., 2021; Shi et al., 2024; Sahoo et al., 2024; Campbell et al., 2022) have recently emerged as a competitive framework for generative modeling on discrete domains, with applications to text, images, music and biological sequences. These models inherit their structure from their continuous-space counterparts (Ho et al., 2020; Song et al., 2020): they are built from two processes, a forward noising process that gradually destroys information, and a reverse process that reconstructs samples from noise. This paper studies masked discrete diffusion models. A data point is a sequence x = (x1 , . . . , xN ) of N tokens drawn from a target distribution π on a finite vocabulary. In the forward process, each coordinate is independently replaced by a special mask symbol m according to a decreasing noise schedule αt over the time interval [0, 1]. The reverse process starts from the fully masked sequence and progressively reveals masked coordinates according to the reverse schedule βt := α1−t . At each reverse step, the model is trained to receive the set M ⊆ [N ] := {1, . . . , N } of the currently revealed coordinates and to output, for every unrevealed coordinate i ∈ / M , a one-coordinate ∗ Bocconi University, Department of Decision Sciences, Milan, Italy. [email protected]
† Bocconi University, Department of Decision Sciences and BIDSA, Milan, Italy. [email protected] GZ acknowledges support from the European Research Council (ERC), through StG “PrSc-HDBayLe” grant ID 101076564.
1
conditional distribution piθ (· | xM ). The ideal values of the standard learned time-independent predictors (Ou et al., 2024; Zheng et al., 2024; Kim et al., 2025) satisfy piθ (· | xM ) = πi (· | xM ),
M ⊆ [N ], i ∈ / M,
where πi ( · | xM ) is the true conditional law of X i given the coordinates X M = xM (see Section 2). Exact reverse sampling can be implemented by simulating the underlying continuous-time Markov chain (CTMC) with the Gillespie algorithm (Gillespie, 1977). In the masked setting, however, exact simulation reveals one coordinate at a time, and therefore requires N sequential updates; for long sequences this is expensive, as each update requires a model call. In practice, one uses the tau-leaping discretization (Gillespie, 2001; Campbell et al., 2022): the reverse time interval is divided into sub-intervals, and within each sub-interval several coordinates are revealed in parallel, using conditionals computed at the start of the step. This produces a sample in a prescribed number K ≪ N of iterations. The price of this acceleration is a systematic discretization error. When several coordinates are revealed simultaneously, the exact reverse law requires the joint conditional distribution of the whole block, whereas tau-leaping replaces it by the product of the one-coordinate conditionals. Thus, even with a perfectly trained model, the sampler incurs an error whenever the coordinates revealed in the same block are conditionally dependent given the previously revealed ones. We call this intrinsic discretization error the factorization error, which we denote εfact . The error is governed by the reverse schedule β: large increments reveal many coordinates in parallel and reduce the number of model calls, but introduce substantial conditional-independence bias; small increments reduce the bias but require more sequential steps. The schedule therefore sets a tradeoff between sampling cost and approximation accuracy. Most existing analyses give upper bounds on the error of a prescribed schedule. In this work we study the converse question: choosing the schedule that minimizes εfact for a given budget of K reverse steps. Contribution We first express the CTMC tau-leaping sampler for masked diffusion (Algorithm 1) in an equivalent formulation, in which a sampling trajectory is represented as an ordered random partition of the coordinate set (Algorithm 2). For this sampler we work with a π-dependent function ρ, the dependence density. Up to the factor N − 1, its value ρ(u) is the average conditional mutual information between two distinct coordinates when each of the others is independently revealed with probability u, and it yields the exact representation (Theorem 1) εfact (β) = N
K Z βt X k k=1
(βtk − u)ρ(u) du.
βtk−1
The formula separates the effect of the schedule from the dependence structure of the target: the schedule enters only through the weights (βtk − u), while all the relevant information about the target law is contained in ρ. The main payoff is that the shape of ρ determines how strongly the schedule matters. When ρN converges to a nondegenerate continuous profile g, optimization over fixed smooth schedules can improve the leading constant of εfact but preserves its N/K scaling (Theorem 2); the linear schedule is therefore order-optimal. This regime includes locally dependent distributions arising from ergodic Markov chains or more generally sufficiently regular stationary processes (Propositions 6, 7). For exchangeable mixtures, by contrast, the dependence is global and mediated by a latent variable. As a result, ρN concentrates near the origin and its limiting profile degenerates (Proposition 8); in this regime, schedules with finer early steps can change the asymptotic order of the error. With KN ≈ log N , we show that the uniform schedule has error Ω(N/ log N ) whereas a geometric schedule has error O(log N ) (Proposition 9). The representation has three further consequences. First, it yields explicit stationarity equations satisfied by every finite-step optimizer; under a monotonicity condition, these equations characterize the unique optimizer (Corollary 1). Second, the large-N limit yields a variational
2
problem with a closed-form leading-order optimizer (Corollary 2). Third, we quantify the asymptotic penalty caused by the block-size randomness intrinsic to tau-leaping, relative to a deterministic planner following the same limiting schedule (Proposition 4). The total correlation TC and the dual total correlation DTC can be expressed through ρ, and this result lets us place several existing schedule bounds and heuristics in a common framework (Propositions 2, 3). Finally, since ρ depends on the unknown target, we estimate its conditional mutual information coefficients through an auxiliary information profile. We compare two estimators of this auxiliary profile and derive a stability bound for schedule optimization (Section 6). Related work In the CTMC formulation of discrete diffusion models, Conforti et al. (2025); Dmitriev et al. (2026) establish non-asymptotic convergence guarantees for discrete diffusion samplers, including masked dynamics. Their analyses upper bound the terminal KL divergence by controlling the distance between path measures, decomposed into initialization, score-approximation, and time-discretization terms. A second, masked-specific, line of theory integrates out the continuous-time path and works directly with the terminal ordered partition, controlling the discrepancy between the algorithmic and target distributions through the global information-theoretic quantities TC and DTC. Li and Cai (2025) bound the error of any deterministic schedule, giving a step budget of order K ≳ TC + DTC to reach a fixed accuracy. Using two different schedule constructions, Chen et al. (2025) and Zhao and Cai (2026) each further reduce the step count to K ≳ TC log N or K ≳ DTC log N . In the CTMC formulation above, Dmitriev et al. (2026) reach the sharper rate K ≳ min{TC, DTC} log N . Related to the present work, Lavenant and Zanella (2025) and Chen et al. (2025) derive the information-profile representation of εfact on which our analysis builds. Starting from this result, Chen et al. (2025) study the difficulty of learning the optimal schedule, while Lavenant and Zanella (2025) study the N → ∞ scaling limit of the schedule-optimization problem for deterministic block sizes. The practical importance of the masking schedule is well recognized. Cosine-type and other monotone increasing schedules are standard in masked diffusion implementations (Shi et al., 2024; Zhang and Syed, 2025). In adaptive or online scheduling, Ben-Hamu et al. (2025) and Foresti et al. (2026) build schedules on the principle that each parallel update should contribute a roughly constant amount of information, a heuristic counterpart to the error-optimal recursion of Corollary 1; Cai and Li (2026) provide theoretical guarantees for algorithms of this type. Organization Section 2 reviews masked diffusion models and the tau-leaping sampler. Section 3 introduces the conditional mutual information profile and the dependence density, states the exact integral representation of the factorization error, derives the recursive optimality condition, and establishes upper bounds. Section 4 studies the asymptotic variational limit and the resulting optimal smooth schedules. Section 5 computes the dependence density for several families of target distributions. Section 6 studies estimation of ι and ρ through the auxiliary profile f , and quantifies the effect of estimation errors on schedule optimization. Technical proofs are collected in the appendix. Concurrent work. Concurrent and independent work by Wainwright (2026), which appeared online during the final writing phase of the present manuscript, derives similar integral representation for the factorization error, data-dependent schedule bounds and a fine-partition analysis. Notation. Unless otherwise stated, N ≥ 2 and the finite vocabulary X has size |X | = L ≥ 2. Let X ∗ := X ∪ {m} be the augmented vocabulary, where m denotes the mask symbol. The elements of (X ∗ )N are written as x = (x1 , . . . , xN ) ∈ (X ∗ )N . For any subset M ⊆ [N ], we denote by xM = (xi )i∈M ∈ (X ∗ )|M | the vector consisting of the entries of x indexed by M . Let π be the target distribution on X N . For any subset M ⊆ [N ], let πM denote the marginal distribution of xM under π. For each i ∈ [N ] \ M , define the conditional distribution πi (· | xM ) = π(xi = · | xM ).
3
2
Masked diffusions, tau-leaping and ordered partitions
Adopting the continuous-time Markov chain (CTMC) perspective, we first review the construction of masked diffusion models and the standard sampling procedure based on the tau-leaping approximation (Shi et al., 2024; Sahoo et al., 2024; Austin et al., 2021; Kim et al., 2025). We then give an equivalent formulation of the generative process in which a sampling trajectory is represented by an ordered random partition of the coordinate set; this representation is the starting point for our analysis of the factorization-error. Forward process. Let α : [0, 1] → [0, 1] be a continuous, strictly decreasing noise schedule with α0 = 1 and α1 = 0. Starting from a data point x0 ∼ π, the forward process masks the coordinates independently according to α as follows. At time t ∈ [0, 1], coordinate i is left unchanged with probability αt and is replaced by m with probability 1 − αt . Therefore, the conditional distribution of xt given x0 factorizes as qt|0 (xt | x0 ) =
N Y
i (xit | xi0 ), qt|0
i (xit | xi0 ) = Cat αt exi0 + (1 − αt )em . qt|0
i=1 L+1
Here ej ∈ R denotes the one-hot vector associated with token j, and Cat(µ) the categorical distribution on X ∗ with probability vector µ. Since α1 = 0, the terminal state is almost surely fully masked. Reverse process The reverse process starts from the fully masked state x1 = (m, . . . , m) ∈ (X ∗ )N and generates a sample from π at time 0. Unlike the forward process, the reverse dynamics generally couple the coordinates, so exact simulation proceeds through sequential single-coordinate updates. To enable parallel sampling, for 0 ≤ s < t ≤ 1 the joint reverse transition qs|t (· | xt ) is replaced by the product of its one-coordinate marginals. For each fixed xt , this product is the optimal factorized approximation of the true joint conditional, in the sense that it minimizes p 7→ KL qs|t (· | xt ) ∥ p over all product distributions p (Lou et al., 2023, Thm. 4.2 and App. A). Given xt ∈ (X ∗ )N , let Ut := {j ∈ [N ] : xjt ̸= m} denote the set of unmasked coordinates, and t write xU t for the corresponding values. For 0 ≤ s < t ≤ 1, the one-coordinate marginals of the true reverse law qs|t (xs | xt ) are given by (Ou et al., 2024, App. D.2): ( i δxit , i xt ̸= m, P (1) qs|t ( · | xt ) = Ut αs −αt i s Cat 1−α x∈X πi (x | xt )ex , xt = m, 1−αt em + 1−αt Ut Ut i t where πi (x | xU t ) = π(x0 = x | x0 = xt ) is the conditional law, under π, of the clean token at t coordinate i given that the coordinates xU t are observed. In practice, the unknown conditional Ut t πi (· | xt ) is replaced by a learned, time-independent predictor piθ (· | xU t ), which is a distribution over clean tokens and assigns zero mass to m. The sampler thus evolves according to the factorized kernel ( N i Y δxit , θ,i i θ,i θ xt ̸= m, q̂s|t (xs | xt ) = q̂s|t (xs | xt ), q̂s|t ( · | xt ) = U 1−αs αs −αt P i i t Cat 1−αt em + 1−αt x∈X pθ (x | xt )ex , xt = m. i=1
(2) This parallel update rule is the Tweedie tau-leaping approximation of the exact reverse dynamics (Lou et al., 2023).
2.1
Tau-leaping as a random ordered partition
Given the predictors piθ , a time grid 0 = t0 < t1 < · · · < tK = 1 and a continuous, strictly increasing denoising schedule β : [0, 1] → [0, 1], with β0 = 0 and β1 = 1, define a sampling scheme. Equivalently, β is determined by the forward noise schedule above through the relation βt = α1−t . Set βk − βk−1 for k = 1, . . . , K. βk := βtk qk := 1 − βk−1 4
The standard reverse tau-leaping sampler in Algorithm 1 is obtained from (1) with αt = β(1 − t) by taking (s, t) = (1 − tk , 1 − tk−1 ) at the k-th step; we write yk for its state after k steps. At each step, every currently masked coordinate is independently selected for unmasking with probability qk ; if selected, its value is sampled from the one-coordinate predictor piθ (· | yk−1 ). For a state y with unmasked set U , we use piθ (· | y) as shorthand for piθ (· | y U ). Algorithm 1 Standard reverse tau-leaping sampler Require: schedule 0 = β0 < β1 < · · · < βK = 1, predictor pθ 1: Initialize y0 ← (m, . . . , m) 2: for k = 1 to K do 3: for i = 1 to N do i 4: if yk−1 = m then 5: Sample yki ∼ Cat (1 − qk ) em + qk piθ (· | yk−1 ) 6: else i 7: yki ← yk−1 8: end if 9: end for 10: end for 11: return yK For the analysis, it is convenient to turn to an equivalent formulation of the same sampler, as done in Ben-Hamu et al. (2025); Li and Cai (2025); Lavenant and Zanella (2025); Chen et al. (2025). Algorithm 2 makes explicit that tau-leaping samples a random ordered partition: the draws of sk and zk depend on the past only through z<k , and not on the generated token values. Hence the entire block sequence z may equivalently be sampled in advance, after which the token values are generated block by block. For an integer 0 ≤ s ≤ |A|, let Unif(A; s) denote the uniform distribution over the subsets of A of cardinality s. Algorithm 2 Reverse tau-leaping sampler, second form Require: schedule 0 = β0 < β1 < · · · < βK = 1, predictor pθ 1: Initialize y0 ← (m, . . . , m) and z<1 ← ∅ 2: for k = 1 to K do 3: Sample sk ∼ Bin N − |z<k |, qk 4: Sample zk ∼ Unif Q [N ] \ z<kz;<ksk 5: Sample ykzk ∼ i∈zk piθ (· | yk−1 ) 6: Set z<k+1 ← z<k ∪ zk z z<k 7: Set yk<k ← yk−1 and yki ← m for all i ∈ / z<k+1 8: end for 9: return yK Proposition 1 (Equivalence of tau-leaping implementations). For every schedule 0 = β0 < β1 < · · · < βK = 1 and every predictor pθ , Algorithms 1 and 2 induce the same law on the trajectory (y0 , . . . , yK ) and, in particular, on the terminal output yK . The proof is given in Appendix A. Let ν β denote the law of z defined in Algorithm 2. It factorizes as ν β (z) =
K Y
ν β (zk | z<k ),
|z |
ν β (zk | z<k ) = qk k (1 − qk ) N −|z<k |−|zk |
k=1
for every admissible zk ⊆ [N ] \ z<k , and the conditional probability is zero otherwise. We use the convention 00 = 1. Since βK = 1 we have qK = 1, so every coordinate still masked before the final step is revealed at step K; hence (z1 , . . . , zK ) is an ordered partition of [N ] with possibly empty SK blocks: the zk are pairwise disjoint and k=1 zk = [N ]. 5
Conditionally on the ordered partition z, the model generates the terminal sample x := yK ∈ z<k X N block by block. Because revealed coordinates are never modified, yk−1 = xz<k , and the conditional law of x given z is pθ (x; z) :=
K Y
Y pθ xzk | xz<k := piθ xi | xz<k ,
pθ xzk | xz<k ,
i∈zk
k=1
with the convention that conditioning on xz<1 = x∅ means conditioning on the fully masked input. The joint law of the output and of the ordered partition is therefore palg (x, z) = pθ (x; z) ν β (z), and the law of the sampler output is its x-marginal, X palg (x) = pθ (x; z) ν β (z). z
2.2
Factorization error
The denoising schedule β determines how coordinates are allocated across parallel update blocks. To study its effect on sampling accuracy, we consider an upper bound on the KL divergence between the target distribution and the marginal output distribution. Following the literature (see e.g. Ben-Hamu et al., 2025; Li and Cai, 2025; Lavenant and Zanella, 2025; Chen et al., 2025), we define the learning error and the factorization error, respectively, by εlearn := Eπ(x)ν β (z)
K X X
πi (xi | xz<k ) , piθ (xi | xz<k )
(3)
π(xzk | xz<k ) . i z<k ) i∈zk πi (x | x
(4)
log
k=1 i∈zk
εfact := Eπ(x)ν β (z)
K X
log Q
k=1
The chain rule for KL divergence, together with its monotonicity under marginalization, yields DKL (π(x)∥palg (x)) ≤ DKL (π(x, z)∥palg (x, z)) = εlearn + εfact ,
(5)
where π(x, z) := π(x)ν β (z) and palg (x, z) := pθ (x; z)ν β (z). Both terms are nonnegative: εfact is an average of conditional total correlations, while εlearn is an average of one-coordinate conditional KL divergences. The learning error therefore vanishes when all one-coordinate conditionals are learned exactly: piθ (· | xM ) = πi (· | xM ),
M ⊆ [N ],
i∈ / M.
By contrast, the factorization error generally persists even under perfect learning. Indeed it is a purely algorithmic error induced by parallel updates and quantifies the discrepancy introduced by approximating each joint block conditional distribution by the product of its one-coordinate conditional marginals, Y π(xzk | xz<k ) ≈ πi (xi | xz<k ). i∈zk
Under perfect learning, εfact equals the joint KL divergence in (5) and thus provides an upperbound on the terminal KL divergence. This motivates our objective: for a fixed budget of K steps, we seek a schedule that minimizes the factorization error, min
0=β0 ≤β1 ≤···≤βK =1
6
εfact (β).
(6)
3
Dependence densities and optimal schedules
We introduce the definitions needed to express the factorization error as an integral functional of the schedule. To quantify the conditional dependence underlying this error, let ι(i) denote the average dependence between two unrevealed coordinates after observing i uniformly selected coordinates: ι(i) := Eσ [Iπ (X σi+1 ; X σi+2 | X σ≤i )] ≥ 0
i = 0, . . . , N − 2,
(7)
where σ is an independent, uniformly random permutation of [N ], σ≤i := {σ1 , . . . , σi }, with σ≤0 = ∅. For fixed, distinct coordinate indices a, b and a set A disjoint from them, Iπ (X a ; X b | X A ) denotes the standard conditional mutual information under π. Definition 1 (Dependence density). For a distribution π on X N , define ρ(u) := (N − 1)
N −2 X
ι(i)BiN −2 (u),
u ∈ [0, 1],
(8)
i=0
where Bin (u) := by (7).
n i n−i is the i-th Bernstein basis polynomial of degree n, and ι is given i u (1 − u)
Equivalently, ρ(u) = (N − 1)E[ι(Bu )],
Bu ∼ Bin(N − 2, u).
(9)
Thus ρ(u)/(N − 1) is the average conditional mutual information between two uniformly selected distinct coordinates when each of the other N − 2 coordinates is independently revealed with probability u. It records how conditional dependence varies with the reveal probability and is not normalized to integrate to one. Since ι(i) ≥ 0, the function ρ is nonnegative. The next lemma, proved in Appendix B, characterizes when it vanishes. Lemma 1. If π is a product measure, then ρ ≡ 0; otherwise ρ(u) > 0 for every u ∈ (0, 1). The following theorem gives an exact integral representation of the factorization error for the tau-leaping sampler. Its proof, presented in Appendix B, builds on the implicit factorization-error representation developed independently by Lavenant and Zanella (2025) and Chen et al. (2025), which we recall there in terms of conditional mutual information. Theorem 1 (Representation of εfact ). The factorization error defined in Equation (4) for Algorithm 2 admits the representation εfact (β) = N
K Z βk X k=1
(βk − u)ρ(u)du.
(10)
βk−1
This representation yields a recursive formula for the optimal schedule β that minimizes the factorization error. Corollary 1. If π is not a product measure, every minimizer of εfact (β) over the simplex {(β1 , . . . , βK−1 ) : 0 ≤ β1 ≤ · · · ≤ βK−1 ≤ 1} lies in its interior. Moreover every minimizer satisfies the recursion βk+1 = βk + Ψ(βk , βk − βk−1 ),
1 Ψ(b, δ) := ρ(b)
Z b ρ(u) du
k = 1, . . . , K − 1.
(11)
b−δ
If b 7→ Ψ(b, δ) is nondecreasing for every fixed δ, then the minimizer is unique. If π is a product measure, then εfact (β) = 0 for every schedule, so every schedule is optimal.
7
Remark 1 (On uniqueness). The monotonicity condition of Corollary 1 holds, for example, whenever log ρ is concave. The recursion (11) determines the entire schedule from the initial value β1 and stationary schedules correspond to the values of β1 for which the recursion terminates exactly at βK = 1. Under the monotonicity assumption the resulting shooting map is monotone and the unique minimizer can be computed by bisection on β1 . Without this assumption, the shooting equation may have several roots. In that case, one must identify all relevant stationary schedules, or use a global finite-dimensional optimization procedure, and compare their objective values. For instance, when K = 2, the stationary condition reduces to Z x ρ(u) du = (1 − x)ρ(x), x = β1 . 0
For the profile ρ(u) = 2 + sin(30u), this becomes 2x +
1 − cos(30x) − (1 − x)(2 + sin(30x)) = 0, 30
which has three roots in (0, 1); the first-order conditions alone therefore do not identify the minimizer. This ρ is not of the polynomial form (8) and serves just as an illustration. Remark 2 (Target-agnostic schedules). The optimal schedule requires full knowledge of ρ, which is generally unavailable in applications. It is therefore useful to consider simple target-agnostic schedules adapted to qualitative shapes of ρ such as whether the mass of ρ concentrates near the origin. For example, the geometric schedule βk+1 = (1 + a)βk studied in Proposition 3, formally solves the local optimality recursion for the idealized profile ρ(u) ∝ u−2 : every stationarity relation holds except the one adjacent to the origin where the profile diverges (the proof is given in Appendix B). No dependence density realizes this profile exactly (by (8), ρ is a polynomial, hence bounded) but it captures the behavior of genuine targets such as the exchangeable mixtures of Section 5.
3.1
Bounds for target-agnostic schedules
Several previous works bound the KL divergence between the algorithm distribution and the data distribution in terms of dependence functionals of π (Chen et al., 2025; Lavenant and Zanella, 2025; Li and Cai, 2025; Dmitriev et al., 2026; Zhao and Cai, 2026). We show that these quantities admit simple representations in terms of the dependence profile ρ. We use the standard notation H( · ), H( · | · ), and I( · ; · ) for entropy, conditional entropy and mutual information. Throughout, X = (X 1 , . . . , X N ) ∼ π and X −i = (X j )j̸=i . Definition 2. The total correlation and dual total correlation of π are defined, respectively, as T C(π) :=
N X
H(X i ) − H(X 1 , . . . , X N )
i=1
DT C(π) := H(X 1 , . . . , X N ) −
N X
H(X i | X −i ) .
i=1
Also, define their normalized sum D(π) : =
N 1 X T C(π) + DT C(π) = I(X i ; X −i ). N N i=1
The following Lemma (proved in Appendix B.1) relates these quantities to the dependence density. Lemma 2. The quantities D(π), T C(π), and DT C(π) can be expressed in terms of ρ as Z 1 Z 1 Z 1 u ρ(u) du. D(π) = ρ(u) du, TC(π) = N (1 − u)ρ(u) du, DTC(π) = N 0
0
0
8
Equivalently, D(π) =
N −2 X
ι(i),
TC(π) =
i=0
N −2 X
(N − i − 1)ι(i),
DTC(π) =
i=0
N −2 X
(i + 1)ι(i);
i=0
in particular D(π) ≤ min{TC(π), DTC(π)}. As a consequence of Theorem 1 and the decomposition (5), we obtain a universal upper bound on the terminal KL divergence DKL (π∥palg ) in terms of D(π). Thanks to the framework of Algorithm 2, the following result can be compared not only with the works we cited above, but also with results obtained from a CTMC perspective, namely Conforti et al. (2025, Theorem 3.1.1) and Dmitriev et al. (2026, Theorem 3). In contrast to path-space CTMC analyses, this bound is stated directly at the level of the terminal distribution and separates the learning and factorization terms appearing in (5). Proposition 2. For any schedule β, writing ∆βmax = maxk ∆βk , DKL (π(x)∥palg (x)) ≤ N ∆βmax D(π) + εlearn .
(12)
For the constant step size ∆βk = 1/K this becomes DKL (π(x)∥palg (x)) ≤
N D(π) + εlearn . K
A bound independent of the target distribution follows from the universal estimate D(π) ≤ log |X |, which holds since I(X i ; X −i ) ≤ H(X i ) ≤ log |X | for every i. Remark 3. Since E[sk ] = N ∆βk , the factor in (12) is the maximum expected block size, N ∆βmax = maxk E[sk ]. The bounds of Lavenant and Zanella (2025, Theorem 7) and Li and Cai (2025, Theorem 1) instead carry the expected maximum block size, in the form (E[maxk sk ] − 1)D(π). Since by Jensen maxk E[sk ] ≤ E[maxk sk ], the leading factor in (12) is sharper when E[maxk sk ] ≥ maxk E[sk ] + 1, that is, whenever the block sizes fluctuate by at least one token on average. Whereas (12) holds for an arbitrary schedule, one can also derive guarantees tailored to specific prescribed schedules, in the spirit of Chen et al. (2025, Theorem 1.9) and Dmitriev et al. (2026, Corollary 2), by working directly with the decomposition of Equation (29) in the proof of Theorem 1. l m Proposition 3. Let a > 0 with a/N ∈ (0, 1), and let K := 1 + log(N/a) log(1+a) . Consider the geometric schedule β0 = 0, βk = min{1, (1 + a)k−1 a/N }, k = 1, . . . , K, for which βK = 1. Then εfact ≤ a DTC(π) + D(π) ≤ 2a DTC(π). Symmetrically, consider the reverse geometric schedule, whose residual distance to the endpoint decays geometrically, βK = 1,
1 − βk = min{1, (1 + a)K−1−k a/N },
k = 0, . . . , K − 1,
for which β0 = 0 and 1 − βK−1 = a/N . Then εfact ≤ a TC(π) + D(π) ≤ 2a TC(π). The proof is in Appendix B Remark 4. For the geometric schedule, εfact ≤ ε is guaranteed by the choice a = ε (2 DTC(π))−1 , assuming a ≤ 1. Since the number of blocks satisfies K ≤ 1 + ⌈log(N/a)/ log(1 + a)⌉, for a ≤ 1 we have K ≲ a−1 log(N/a), and therefore K≲
DTC(π) N DTC(π) log . ε ε 9
4
Asymptotic limit
In this section, we study sequences of distributions {πN }N ≥1 whose dependence densities ρN converge uniformly to a limiting profile g. We derive a leading-order expression for εfact , identify the schedule minimizing the resulting variational functional, and quantify the excess cost of tauleaping relative to a planner with deterministic block sizes following the same limiting schedule. The proofs can be found in Appendix C. Throughout, ρN and ιN denote the dependence density and the conditional mutual information profile of πN , respectively. Assumption 1. For a sequence {πN }N ≥2 there exists a continuous function g : [0, 1] → R+ such that ∥ρN − g∥∞ −→ 0. N →∞
Let D∞ :=
R1 0
g(u) du. Assumption 1 and Lemma 2 imply D(πN ) → D∞ .
Lemma 3. Suppose there exists a continuous function g : [0, 1] → R+ such that, for N ≥ 3, j −→ 0. (13) max (N − 1)ιN (j) − g 0≤j≤N −2 N − 2 N →∞ Then ∥ρN (u) − g(u)∥∞ → 0 as N → ∞. The proof combines the Bernstein representation (8) with Bernstein’s approximation theorem. Theorem 2. Let {πN }N ≥2 satisfy Assumption 1 and let β ∈ C 1 ([0, 1]) be strictly increasing with β(0) = 0 and β(1) = 1. For each N, K, consider the schedule βk = β(k/K), k = 0, . . . , K, and write εN,K fact (β) for the corresponding factorization error. Then, as N, K → ∞ jointly and with no relation required between them, Z 1 N N N,K ′ 2 εfact (β) = g(β(t)) β (t) dt + o . (14) 2K 0 K Equation (14) yields a variational problem over smooth increasing schedules, whose minimizer is explicit. Corollary 2. Under the hypotheses of Theorem 2, assume g(u) > 0 for u ∈ [0, 1], the schedule minimizing the leading term in (14) satisfies Z yp 1 β ′ (t) ∝ p g(u) du, , that is, β ∗ (t) = G−1 t G(1) , G(y) := g(β(t)) 0 with corresponding factorization error ∗ εN,K fact (β ) =
N 2K
Z 1
p
0
2 g(u) du
N +o K
.
The variational problem and its solution are those of Lavenant and Zanella (2025, Prop. 13); we include the statement for completeness, the novelty here being the identification of the limit (14) for the tau-leaping sampler. Remark 5. For the linear schedule β(t) = t, Theorem 2 gives Z 1 N N εlin = g(u)du + o . fact 2K 0 K N,K ∗ Under the hypotheses of Corollary 2, write εopt fact := εfact (β ) for the error of its optimal limiting smooth schedule. Then R1 g(u)du εlin fact −→ R 0p 2 ≥ 1, opt N,K→∞ 1 εfact g(u)du 0
where the inequality follows from Jensen. Thus, under these hypotheses, optimizing the limiting functional over fixed smooth schedules improves the leading constant but not the N/K scaling. The same conclusion does not necessarily apply outside this regime, for example when g degenerates as in the exchangeable mixtures case of Section 5. 10
4.1
The cost of random block sizes
Given a schedule β, tau-leaping reveals a random number of coordinates at each step: as shown in the proof of Theorem 1, the joint vector of block sizes follows a multinomial distribution, so that sk ∼ Bin(N, ∆βk ) marginally. We compare the asymptotics of Theorem 2 with those obtained by Lavenant and Zanella (2025) for a planner that uses a uniformly random coordinate ordering and deterministic block sizes sdet k = ⌈N β(k/K)⌉ − ⌈N β((k − 1)/K)⌉,
k = 1, . . . , K.
For the same limiting schedule β, this comparison quantifies the excess factorization cost caused by the random block sizes of tau-leaping. For the comparisons in this subsection, assume the stronger coefficient convergence (13) and D∞ > 0. These hypothesis ensure that the deterministic limits of Lavenant and Zanella (2025) apply. When N, K → ∞ with N/K → ∞, the two factorization costs coincide to first order by Theorem 2 and Lavenant and Zanella (2025, Theorem 12). This is due to the fact that the random block sizes concentrate around their deterministic counterparts sdet k ≈ N ∆βk : for smooth sk −→ 1 in probability. schedules β with inf β ′ > 0, whenever N ∆βk → ∞ we have N ∆β k The difference becomes visible once the expected block sizes stay bounded. When N/K → s̄ ∈ [1, ∞), Theorem 2 gives the tau-leaping cost εN,K fact =
s̄ 2
Z 1
g(β(t)) β ′ (t)2 dt + o(1),
0
′
while, if β has finitely many local extrema, the deterministic limit of Lavenant and Zanella (2025, Theorem 15) is Z 1 1 g(β(t)) s̄ hs̄ (β ′ (t)) − β ′ (t) dt + o(1). εdet,N,K (β) = fact 2 0 Here hs̄ is the piecewise-linear interpolation of u 7→ u2 on the grid {0, 1/s̄, 2/s̄, . . . }, which has the closed form hs̄ (u) = u2 + s̄−2 {s̄u} 1 − {s̄u} , {s̄u} := s̄u − ⌊s̄u⌋ ∈ [0, 1). The next proposition shows that, under these assumptions, the limiting tau-leaping factorization cost is at least that of the deterministic planner. The limiting gap approaches 12 D∞ as s̄ subsequently grows. For the linear schedule with K = N , the effect of random block sizes is particularly clear: the deterministic planner reveals exactly one coordinate per step and has zero factorization error, whereas the tau-leaping factorization error converges to 12 D∞ . Proposition 4. Assume (13), D∞ > 0 and β as in Theorem 2, with β ′ having finitely many local extrema. Then, as N, K → ∞ with N/K → s̄ ∈ [1, ∞), the gap ∆s̄ := limN,K→∞ εN,K fact (β) − det,N,K εfact (β) exists and equals 1 ∆s̄ = 2
Z 1
"
θ(t) 1 − θ(t) g(β(t)) β (t) − s̄ 0
#
′
dt,
θ(t) := {s̄ β ′ (t)}.
Moreover 0 ≤ ∆s̄ ≤ 12 D∞ and ∆s̄ − 12 D∞
5
≤
1 8s̄
Z 1 g(β(t)) dt ≤ 0
∥g∥∞ −−−→ 0. s̄→∞ 8s̄
Examples
In this section we compute the dependence density ρ for various classes of target distributions π ∈ P(X N ). These examples show how the shape of ρ reflects the way dependence is distributed across the coordinates, and how it affects the factorization error. We include numerical simulations verifying the theoretical results for the stationary Markov chains and for exchangeable models. All proofs are deferred to Appendix D. 11
The first example is that of product measures. Since the coordinates are independent, conditioning on the revealed coordinates does not change the conditional law of unrevealed ones, so parallelizing the updates induces no factorization error. Proposition 5 (Product measures). Let π = π1 ⊗ · · · ⊗ πN . Then ρ ≡ 0 on [0, 1] and consequently εfact (β) = 0 for every schedule β. This is the degenerate case of Lemma 1; every schedule is optimal, and schedule design is vacuous. Next we consider distributions with local dependence, starting with a stationary Markov chain. Here ρ can be expressed through the mutual informations I(X 0 ; X d ), capturing how the Markovian dependence changes as the fraction of revealed coordinates varies. Proposition 6 (Stationary Markov chains). Let (X t )t∈Z be a stationary time-homogeneous Markov chain on X with transition matrix P and stationary law µ, and let πN be the law of its restriction (X 1 , . . . , X N ), so that N −1 Y 1 N 1 πN (x , . . . , x ) = µ(x ) P (xt , xt+1 ). t=1 0
d
0
0
Set h0 := H(X ), hd := H(X | X ) and Id := I(X ; X d ) = h0 − hd . Then N −1 X d2 1 GN (u) , GN (u) := (N − d)Id u2 (1 − u)d−1 . ρN (u) = 2 du N d=1 P Moreover, if d≥1 Id < ∞, then ρN → g uniformly on [0, 1], where d2 [u EIDu ] , du2 and g is extended continuously to u = 0.
Du ∼ Geom(u),
g(u) :=
(15)
u ∈ (0, 1],
Remark 6. This case has also been studied by Luxembourg et al. (2025). Their analysis shows that the factorization error incurred by parallel sampling can be controlled by updating coordinates that are sufficiently well separated. This motivates their dilated unmasking scheme, which uses a logarithmic number of denoising iterations per block. For irreducible and aperiodic finite-state Markov chains, the lag mutual information is summable, so Proposition 6 yields a continuous limiting profile and Theorem 2 applies. For a non-i.i.d. chain this profile is not identically zero, and every fixed smooth schedule has factorization error of order N/K. If, in addition, the limiting profile is strictly positive on [0, 1], Corollary 2 and the subsequent optimal-to-linear comparison of Remark 5 apply. The next result identifies the limiting profile for a stationary process satisfying the uniform convergence assumption, of which the Markov chain is one instance. Stationarity relaxes the conditional independence structure of Markov chains while retaining translation invariance, a regime of interest for sequence models whose dependence is far from a Markovian one. Proposition 7 (Stationary processes). Let (X t )t∈Z be a stationary process on X , that is, for every finite set of times t1 , . . . , tk and every shift s ∈ Z, d
(X t1 , . . . , X tk ) = (X t1 +s , . . . , X tk +s ), and let πN be the law of (X 1 , . . . , X N ). For a finite A ⊆ N write −A := {−a : a ∈ A} and h(A) := H(X 0 | X −A ),
h0 := H(X 0 ),
I(A) := I(X 0 ; X −A ) = h0 − h(A),
iid
with I(∅) = 0. Let ξi ∼ Bern(u) for i ≥ 1 and Sd,u := {i ∈ [d] : ξi = 1}. Then, for every N , # " N −1 d2 u X ρN (u) = 2 EI(Sd,u ) . du N d=0
If moreover Assumption 1 holds, then, with Su := {i ≥ 1 : ξi = 1}, d2 [u EI(Su )] . N →∞ du2
ρN (u) −→
12
(16)
For a Markov chain, I(Sd,u ) = Imin Sd,u depends on the revealed set only through its nearest element, and the d-sum telescopes to N1 GN (u), recovering (15). The last class of examples we consider is that of exchangeable mixtures. In these mixtures the dependence is mediated by a global latent variable P , in contrast to the local dependence of the Markov examples above. As a result, ρN may collapse away from the origin as N → ∞, while concentrating at it. Intuitively, once a positive proportion of coordinates has been revealed, P is almost fully identified and the remaining coordinates become nearly conditionally independent. Proposition 8 (Exchangeable mixtures). Let πN be an exchangeable law of the form P ∼ Q, and iid
X 1 , . . . , X N | P ∼ P , so that πN (x1 , . . . , xN ) =
Z Y N
p(xi ) Q(dp).
i=1
Let ιN (i) := I(X i+1 ; X i+2 | X 1:i ) for i = 0, . . . , N − 2, then ρN (u) = (N − 1) EB [ιN (B)],
B ∼ Bin(N − 2, u),
(17)
If moreover I(X m+1 ; X m+2 | X 1:m ) = o(1/m) and I(X 1 ; X 2 ) > 0, then ρN (u) −→ 0 N →∞
for u ∈ (0, 1],
ρN (0) → +∞.
(18)
Remark 7. The assumption I(X m+1 ; X m+2 | X 1:m ) = o(1/m) holds, for example, in the Dirichletcategorical or Beta-Bernoulli model, in which the conditional mutual information is of order O(1/m2 ). Equation (17) is the exchangeable analogue of the Markov formula (15). In the Markov case the profile averages dependence across random gap lengths, whereas in the exchangeable case it averages the residual dependence between two future coordinates after conditioning on a random number of previously revealed samples. We close by comparing two schedules for a finite-dimensional exchangeable target using the same number of iterations. In this degenerate regime, the choice of schedule affects not only the constant in the factorization error but also its asymptotic order. Proposition 9 (Comparison of different schedules). Suppose Q is a nondegenerate p-dimensional parametric prior satisfying the regularity conditions of Clarke and Barron (1994), and let KN := 1 + ⌈log2 N ⌉. Then the uniform schedule βkunif :=
k , KN
k = 0, . . . , KN ,
satisfies εunif fact (πN , KN ) ≥
CN log N
for a constant C > 0, whereas the geometric schedule β1 =
1 , N
βk = min{1, 2βk−1 },
satisfies ε⋆fact (πN , KN ) ≤ εgeom fact (πN , KN ) = O(log N ). where ε⋆fact (πN , KN ) denotes the minimum of εfact over all schedules with KN steps. Consequently εunif N C′ fact (πN , KN ) ≥ , ⋆ εfact (πN , KN ) (log N )2 for a constant C ′ > 0.
13
5.1
Numerical illustration
We illustrate the theory for a stationary Markov chain (local dependence, Proposition 6) and for a Beta-Bernoulli exchangeable model (global dependence, Proposition 8). The quantities are evaluated numerically from exact expressions for ρN , without Monte Carlo sampling. Setup. For the Markov chain we take a reversible and aperiodic 10-state lazy random walk on a sparse connected randomly weighted graph, so Id = I(X 0 ; X d ) decays geometrically in d. For the exchangeable model we take the Beta(1, 1)-Bernoulli mixture, for which ι(m) := I(X m+1 ; X m+2 | X 1:m ) =
m X
P(Sm = s) I(X m+1 ; X m+2 | Sm = s),
Sm := X 1 +· · ·+X m ,
s=0
is computed by averaging the posterior-predictive mutual information over the Beta-Binomial sufficient statistic Sm . Here ι(m) = ιN (m) for every N ≥ m + 2. Dependence density. Figure 1(a) shows ρN for N = 64, . . . , 512 together with its limiting profile g for the Markov chain. The curves are close, with small finite-N corrections over the displayed range u ∈ [0, 0.15]. Panel (b) shows ρN for the exchangeable model over the same range of N : the mass concentrates near u = 0 as N grows, consistent with the degeneration of ρN in Proposition 8. Optimal versus linear schedule. For each target we compute the εfact -optimal schedule from Corollary 1, solving the recursion (11) by one-dimensional root-finding on β1 (Brent’s method, Brent, 2013). Neither profile is log-concave, so the sufficient condition of Corollary 1 does not apply. Nonetheless the numerical shooting search found a single crossing of the terminal condition βK = 1 in each case, providing numerical evidence of uniqueness for these examples. Panel (c) compares the optimal and linear schedules at N = 256, K = ⌈log2 N ⌉ + 1 = 9. For the Markov chain the optimal schedule departs only mildly from the linear one; for the exchangeable model it is strongly front-loaded, concentrating budget where ρN is largest. opt Local-versus-global dependence. Panel (d) plots the improvement ratio εlin fact /εfact against N , with the step budget KN = ⌈log2 N ⌉ + 1 growing logarithmically. For the Markov chain the gain is modest and the ratio appears to stabilize over the tested dimension at 1.17, consistent with the regime of Remark 5. For the exchangeable model the numerical results appear consistent with the growth rate lower bound N/(log N )2 , predicted by Proposition 9: here optimizing the schedule changes the asymptotic rate of εfact , not merely its constant.
6
Estimators for the information profile
Computing the recursive schedule in Corollary 1 requires ρ, which is determined by the sequence ι(0), . . . , ι(N −2). Assuming that each model call returns predictive distributions for all unrevealed coordinates, estimating the full sequence with n Monte Carlo replications per entry uses O(nN ) model calls. This computation is performed offline, and its cost is amortized over subsequent generated samples, whose cost is O(K) model calls. We estimate these coefficients through the auxiliary information profile introduced by Lavenant and Zanella (2025) and Chen et al. (2025). For a uniform random permutation σ of [N ], define f (i) := −Eσ [Hπ (X σi+1 | X σ≤i )] ,
i = 0, . . . , N − 1.
(19)
The conditional mutual information identity and the symmetry of the uniform permutation give, for i = 0, . . . , N − 2 ι(i) = Eσ [Hπ (X σi+2 | X σ≤i ) − Hπ (X σi+2 | X σ≤i+1 )] = f (i + 1) − f (i) = ∆f (i).
14
(20)
Figure 1: (a) Dependence density ρN for the Markov chain at N = 64, 128, 256, 512, with the limiting profile g. (b) ρN for the Beta(1, 1)-Bernoulli exchangeable model at N = 64, 128, 256, 512. Panels (a)-(b) display u ∈ [0, 0.15]. (c) εfact -optimal versus linear schedule at N = 256, K = 9. opt (d) Improvement ratio εlin fact /εfact versus N with K = ⌈log2 N ⌉ + 1. Thus estimates of f yield estimates of ι and, in turn, of the dependence density ρ. We first compare two estimators of f , then quantify how errors in the estimated coefficients affect the factorization error and the selected schedule. Fix i ∈ {0, . . . , N − 1}. Let X ∼ π and, independently, let Z ⊆ [N ] be uniform among the subsets of [N ] of cardinality i and, conditionally on Z, let J be uniform on [N ] \ Z. By (19), f (i) can be equivalently written as f (i) = EX, Z, J log π(X J | X Z ) . We compare two estimators of f (i) built from the model predictive distribution pjθ . The logprobability estimator is 1 X f˜iθ := log pjθ (X j | X Z ), (21) N −i j ∈Z /
and the entropy-estimator is fˆiθ := −
1 X H pjθ (· | X Z ) . N −i j ∈Z /
We write f˜i , fˆi for the corresponding quantities under perfect learning, pjθ = πj . Proposition 10. Under perfect learning, the estimators (21) and (22) are unbiased E[f˜i ] = E[fˆi ] = f (i), and Var(fˆi ) ≤ Var(f˜i ). 15
(22)
We now drop the assumption of perfect learning and relate the bias of the two estimators to the model training objective. Up to an additive constant independent of θ, the denoising crossentropy loss decomposes as a weighted sum of conditional KL errors (Ou et al., 2024; Kim et al., 2025) defined for i = 0, . . . , N − 1 as i h 1 X DKL (πj (· | X Z )∥pjθ (· | X Z )) = EZ,J,X Z DKL (πJ (· | X Z )∥pJθ (· | X Z )) , Li := EZ,X Z N −i j ∈Z /
where Z and J are distributed as above. Proposition 11. For general learned predictors, the bias of the log-probability estimator is E f˜θ − f (i) = − Li ≤ 0. i
2 If moreover Li ≤ 2 1 − 1/|X | , then the bias of the entropy estimator satisfies p ϕ(T ) := T log(|X | − 1) + H2 (T ), E[fˆiθ ] − f (i) ≤ ϕ Li /2 , where H2 (T ) := −T log T − (1 − T ) log(1 − T ) is the binary entropy. In particular p E[fˆiθ ] − f (i) = O Li log(1/Li ) as Li → 0. Proposition 11 shows that both biases vanish as Li → 0, but provides different guarantees. The bias of the log-probability estimator is exactly −L√ i and is therefore always non-positive. For Li log(1/Li )) on the absolute bias; since the entropy estimator, we obtain the upper bound O( √ x = o( x log(1/x)) as x → 0+ the latter guarantee is weaker in its dependence on Li . This does not imply that the entropy estimator has a larger absolute bias as the bound is obtained through Pinsker’s inequality and need not be tight. Under perfect learning, both estimators are unbiased, and Proposition 10 shows that the entropy estimator has smaller variance than the log-probability estimator. The following provides a Monte Carlo bound for fˆiθ for arbitrary predictors, including imperfectly learned ones. Proposition 12. Fix i ∈ {0, . . . , N − 1} and let (X (m) , Z (m) ), m = 1, . . . , n be independent θ copies of (X, Z) as defined above, and let fˆi,m be the corresponding estimators (22), and set P 1 θ ¯ ˆ fi := n m fi,m . Then (log |X |)2 . Var(f¯i ) ≤ 4n Moreover, for every δ ∈ (0, 1), with probability at least 1 − δ, r log(2/δ) θ ¯ ˆ . fi − E[fi ] ≤ log |X | 2n Given an estimated vector f¯, we set ῑ(i) := f¯i+1 − f¯i ,
ρ̄(u) := (N − 1)
N −2 X
ῑ(i)BiN −2 (u).
i=0
We conclude by quantifying how errors in the estimated coefficients affect the factorization error and the selected schedule. We do not claim (23) to be tight, and leave the development of sharper inequalities with stronger assumptions to future work. Proposition 13. Let BK := {β = (β0 , . . . , βK ) : 0 = β0 ≤ · · · ≤ βK = 1}, and ∆βmax := max1≤k≤K (βk − βk−1 ). For any vector a = (a(0), . . . , a(N − 2)) ∈ RN −1 , set N −2 K Z βk X X ρa (u) := (N − 1) a(i)BiN −2 (u), εfact (a, β) := N (βk − u)ρa (u) du. i=0
k=1
βk−1
For a fixed schedule β, this functional is linear in a. Writing ∥a∥1 :=
PN −2 i=0
|a(i)|, we have
|εfact (ι, β) − εfact (ῑ, β)| ≤ N ∆βmax ∥ι − ῑ∥1 ,
(23)
If β̂ ∈ arg minβ∈BK εfact (ῑ, β) and β ∗ ∈ arg minβ∈BK εfact (ι, β), then ∗ 0 ≤ εfact (ι, β̂) − εfact (ι, β ∗ ) ≤ 2N max{∆β̂max , ∆βmax } ∥ι − ῑ∥1 .
16
(24)
6.1
Numerical experiments on mixtures of products
We illustrate the theoretical predictions on a class of targets for which all the relevant conditional distributions can be evaluated exactly. This allows us to set pjθ = πj , so that the learning-error contribution vanishes. Target distribution. We consider a mixture of r product distributions over N coordinates, each taking value in a vocabulary X of size L: π(x) =
r X z=1
wz
N Y
µz,j (xj ),
x ∈ XN,
(25)
j=1
where w ∈ ∆r−1 and µz,j ∈ P(X ). Equivalently, a latent variable Z ∼ Cat(w) is drawn first, after which the coordinates are conditionally independent, with X j | Z = z ∼ µz,j . The component marginals are sampled from the symmetric Dirichlet distribution µz,j ∼ Dir(α/L, . . . , α/L) and then held fixed throughout profile estimation and Monte Carlo evaluation. The mixture weights are uniform, wz = 1/r. Throughout the experiment, we set r = 5, L = 10, α = 0.3 and consider N ∈ {8, 16, 32, 64, 128, 256}, K = log2 N . Implementation details are in Appendix E.1. Behavior of ρN . Figure 2 (a) shows that the estimated dependence density ρ̄N (u) becomes increasingly concentrated near u = 0 as N increases. For readability, the upper limit of the vertical axis is set to 70. Conditional on the sampled component marginals, the target is generally not exchangeable, since the distributions µz,j depend on the coordinate j. Thus, Proposition 8 does not apply directly. Nevertheless, the observed behavior suggests that a similar concentration phenomenon can occur beyond the conditionally i.i.d. setting covered by that proposition. Figure 2 (b) compares the linear schedule βklin = k/K with a numerically optimized schedule β̂ for N = 128, K = 7, where the latter is computed from the estimated profile. Evaluating the representation in Theorem 1 with this profile gives estimated factorization errors of 21.142 and 0.681, respectively, corresponding to an estimated improvement factor of approximately 31. Monte Carlo comparison. Figure 2 (c) compares two estimates of the factorization error for the linear and numerically optimized schedules. The solid curves are obtained from Theorem 1, using ῑ(i) = f¯i+1 − f¯i to construct the estimated dependence density. The boxplots summarize 100 independent Monte Carlo estimates from (4), each based on 4000 draws. This comparison provides a check of the representation-based estimates against direct Monte Carlo evaluation. bN,K := εfact (ῑ, β lin )/εfact (ῑ, β̂). Finally, Figure 2 (d) reports the estimated improvement ratio R For the target instances considered, the estimated gain increases with N , indicating substantial benefits from schedule optimization over the tested range of dimensions.
Acknowledgements During the preparation of this work, the authors used OpenAI ChatGPT-5.6 to explore proof strategies and assist with the drafting and revision of the manuscript. The authors take full responsibility for the contents of the paper.
References Koenraad MR Audenaert. A sharp Fannes-type inequality for the von Neumann entropy. arXiv preprint quant-ph/0610146, 2006. Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:17981–17993, 2021. 17
Figure 2: (a) Estimated dependence densities ρ̄N (u) for mixture-of-products targets with N ∈ {8, 16, 32, 64, 128, 256}. The vertical axis is truncated at 70. (b) Linear and numerically optimized schedules for N = 128 and K = 7. (c) Representation-based and direct Monte Carlo estimates of εfact for both schedules with K = log2 N . The curves use the estimated profiles, and the boxplots summarize 100 independent Monte Carlo estimates. Both axes are logarithmic. (d) Estimated bN,K versus N , with K = log2 N , on logarithmic axes. improvement ratio R Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, and Brian Karrer. Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking. arXiv preprint arXiv:2505.24857, 2025. Richard P Brent. Algorithms for minimization without derivatives. Courier Corporation, 2013. Changxiao Cai and Gen Li. Confidence-Based Decoding is Provably Efficient for Diffusion Language Models. arXiv preprint arXiv:2603.22248, 2026. Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279, 2022. Sitan Chen, Kevin Cong, and Jerry Li. Optimal inference schedules for masked diffusion models. arXiv preprint arXiv:2511.04647, 2025. Bertrand S Clarke and Andrew R Barron. Jeffreys’ prior is asymptotically least favorable under entropy risk. Journal of Statistical planning and Inference, 41(1):37–60, 1994. Giovanni Conforti, Alain Durmus, Le-Tuyet-Nhi Pham, and Gael Raoul. Non-asymptotic convergence of discrete diffusion models: Masked and random walk dynamics. arXiv preprint arXiv:2512.00580, 2025. Daniil Dmitriev, Zhihan Huang, and Yuting Wei. Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees. arXiv preprint arXiv:2602.15008, 2026. 18
Mark Fannes. A continuity property of the entropy density for spin lattice systems. Communications in Mathematical Physics, 31(4):291–294, 1973. Alberto Foresti, Mustapha Bounoua, Giulio Franzese, Luca Ambrogioni, and Pietro Michiardi. Improved Sampling Schedules for Discrete Diffusion Models. arXiv preprint arXiv:2602.06849, 2026. Daniel T Gillespie. Exact stochastic simulation of coupled chemical reactions. The journal of physical chemistry, 81(25):2340–2361, 1977. Daniel T Gillespie. Approximate accelerated stochastic simulation of chemically reacting systems. The Journal of chemical physics, 115(4):1716–1733, 2001. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.06768, 2025. Hugo Lavenant and Giacomo Zanella. Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations. arXiv preprint arXiv:2510.25544, 2025. Gen Li and Changxiao Cai. Breaking AR’s Sampling Bottleneck: Provable Acceleration via Diffusion Language Models. arXiv preprint arXiv:2505.21400, 2025. Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2023. Omer Luxembourg, Haim Permuter, and Eliya Nachmani. Plan for speed: Dilated scheduling for masked diffusion language models. arXiv preprint arXiv:2506.19037, 2025. Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.03736, 2024. Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:130136–130184, 2024. Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131–103167, 2024. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. Martin J Wainwright. The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity. arXiv preprint arXiv:2608.13520, 2026. Leo Zhang and Saifuddin Syed. The cosine schedule is Fisher-Rao-optimal for masked discrete diffusion models. arXiv preprint arXiv:2508.04884, 2025. Yunxiao Zhao and Changxiao Cai. Adaptation to intrinsic dependence in diffusion language models. arXiv preprint arXiv:2602.20126, 2026. Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint arXiv:2409.02908, 2024.
19
A
Proof of Section 2
Proof of Proposition 1. Since both methods are Markov chains starting from the fully masked state, equality of the stepwise transition kernels implies equality of the laws of (y0 , . . . , yK ), and in particular of the terminal outputs. k −βk−1 . Let y ∈ (X ∗ )N be the current state with masked Fix k ∈ {1, . . . , K} and recall qk := β1−β k−1 set M := {i ∈ [N ] : y i = m}, r := |M |, c
c
and let x ∈ (X ∗ )N be a candidate next state. Both kernels vanish unless xM = y M ; we assume (1) this in what follows and omit the corresponding indicator from all displays. We write Kk and (2) Kk for the step-k kernels of Algorithms 1 and 2, respectively. In Algorithm 1, the coordinates of M are updated independently, so i Yh (1) (1 − qk ) 1{xi =m} + qk piθ (xi | y) . Kk (x | y) = i∈M
For Algorithm 2, define the set of newly revealed coordinates R := {i ∈ M : xi ̸= m},
u := |R|.
Since piθ (m | y) = 0, the event {yk = x} is equal to {zk = R, ykR = xR }. Therefore (2)
Kk (x | y) = P(sk = u)P(zk = R | sk = u)P(ykR = xR | zk = R) Y r u 1 Y i i pθ (x | y) = qku (1 − qk )r−u piθ (xi | y) = qk (1 − qk )r−u r u u i∈R i∈R i Yh (1) i i = (1 − qk ) 1{xi =m} + qk pθ (x | y) = Kk (x | y), i∈M
where the last line uses piθ (m | y) = 0 once more: the bracket equals 1 − qk when xi = m and qk piθ (xi | y) otherwise. Since k, y and x were arbitrary, the two algorithms have the same transition kernels, which proves the claim.
B
Proofs of Section 3
Proof of Lemma 1. The coefficient identity in Lemma 2 gives TC(π) =
N −2 X
(N − i − 1)ι(i) = DKL (π∥⊗N j=1 πj ).
i=0
All coefficients are positive and ι(i) ≥ 0. If π is a product measure, the right-hand side is zero, so every ι(i) is zero and ρ ≡ 0. Otherwise at least one ι(i) is positive. Since every BiN −2 (u) is strictly positive for u ∈ (0, 1), (8) then gives ρ(u) > 0 throughout this interval. Before proving Theorem 1, we recall the discrete factorization-error representation of Lavenant and Zanella (2025) and Chen et al. (2025). We state it using the conditional mutual information profile ι; the relation to their information-profile formulation is given in (20). Lemma 4 (Conditional mutual information representation). Let s = (s1 , . . . , sK ), where sk := Pk |zk |, and set C0 := 0, Ck := ℓ=1 sℓ . For j = 1, . . . , N − 1, define rs (j) := min{Ck : Ck ≥ j}. Then εfact = Eν β (s) [A(s)],
A(s) =
N −2 X
ι(i) rs (i + 1) − (i + 1) ,
i=0
where ν β (s) is the induced law of the block sizes s = (s1 , . . . , sK ) := (|z1 |, . . . , |zK |). 20
(26)
The factor rs (i + 1) − (i + 1) counts the coordinates after position i + 1 that belong to its parallel update block. We will also use the following equivalent form of the dependence density: ρ(u) =
N −2 X
bi (u) := (N − 1)BiN −2 (u) =
ι(i)bi (u),
i=0
ui (1 − u)N −i−2 , B(i + 1, N − i − 1)
(27)
where B(·, ·) is the Beta function. In particular, each bi integrates to one. Proof of Theorem 1. By the decomposition in Lemma 4, it suffices to derive an explicit expression for E[rs (i) − i] for i = 1, . . . , N − 1, where s = (s1 , . . . , sK ) := (|z1 |, . . . , |zK |) is the vector of block sizes of z ∼ ν β ; recall that rs (i) depends on z only through s. Under ν β , the sizes satisfy X β − β k k−1 sk | s1 , . . . , sk−1 ∼ Bin N − sj , , 1 − βk−1 j<k
and their joint law is the multinomial Mult(N ; ∆β1 , . . . , ∆βK ), where ∆βk := βk − βk−1 . We are iid
going to work with the construction of the multinomial that samples U1 , . . . , UN ∼ Unif(0, 1) and sets sk = #{j : Uj ∈ (βk−1 , βk ]}, Ck = #{j : Uj ≤ βk }; Let U(i) denote the i-th order statistic of U1 , . . . , UN , so that U(i) ∼ Beta(i, N − i + 1) by Lemma 5 in Appendix B, and let k(i) = min{k : βk ≥ U(i) } be the index of the block containing position i. The overshoot rs (i) − i counts the indices j > i such that U(j) falls in the same interval (βk(i)−1 , βk(i) ] as U(i) , hence rs (i) − i = Ck(i) − i. Conditionally on U(i) = u, exactly i of the uniforms are at most u, while the remaining order statistics have the law of the ordered values of N − i independent uniforms on (u, 1). Therefore, writing F (u) := min{βk : βk ≥ u} for the right endpoint of the interval containing u, F (u) − u , Ck(i) − i | U(i) = u ∼ Bin N − i, 1−u hence
F (u) − u . 1−u Using (28), the law of U(i) ∼ Beta(i, N − i + 1), and F (u) = βk on (βk−1 , βk ], we obtain E[rs (i) − i | U(i) = u] = (N − i)
(28)
K Z βk X N −i βk − u i−1 u (1 − u)N −i du B(i, N − i + 1) 1 − u k=1 βk−1 Z K X βk ui−1 (1 − u)N −i−1 =N du (βk − u) B(i, N − i) k=1 βk−1 K Z βk X =N (βk − u)bi−1 (u)du,
E[rs (i) − i] =
k=1
βk−1
where we used (N −i)/B(i, N −i+1) = N/B(i, N −i) and the definition of bi in (27). Substitution into (26) yields εfact (β) = N
N −2 X
ι(i)
i=0
=N
K Z βk X k=1
K Z βk X k=1
(βk − u)bi (u) du
βk−1
(βk − u)ρ(u) du,
βk−1
by exchanging the finite sums and using (27). 21
(29)
iid
Lemma 5. Let U1 , . . . , UN ∼ Unif(0, 1), and U(i) denote the i-th order statistic. Then U(i) ∼ Beta(i, N − i + 1). Proof of Lemma 5. Since {U(i) ≤ u} occurs if and only if at least i of the variables Uj are at most u, and #{j : Uj ≤ u} ∼ Bin(N, u), we have N X N k P(U(i) ≤ u) = u (1 − u)N −k . k k=i
N −1 N N −1 Differentiating with respect to u and using k N , k = N k−1 and (N − k) k = N k fU(i) (u) =
N X N k=i
k
kuk−1 (1 − u)N −k − (N − k)uk (1 − u)N −k−1
N −1 X
N −1 X N −1 j N −1 k N −j−1 =N u (1 − u) −N u (1 − u)N −k−1 j k j=i−1 k=i N − 1 i−1 =N u (1 − u)N −i , i−1 by telescoping. Since for positive integers Γ(N + 1) N! N −1 1 , = = =N i−1 B(i, N − i + 1) Γ(i)Γ(N − i + 1) (i − 1)!(N − i)! the final quantity is the density of the Beta(i, N − i + 1) distribution. Proof of Corollary 1. If π is a product measure, then by Proposition 5 εfact (β) = 0 for every schedule β, so all of them are optimal. Assume now that π is not a product measure. By Lemma 1, ρ(u) > 0 for every u ∈ (0, 1). Since εfact is continuous in β on the compact simplex, a minimizer exists. Define Z b Φ(a, b) = (b − u)ρ(u)du, a
PK so that, by Theorem 1, εfact (β1 , . . . , βK−1 ) = N k=1 Φ(βk−1 , βk ). We first show that every minimizer is an interior point. For any a < b < c, Z b Φ(a, c) − Φ(a, b) − Φ(b, c) = (c − b) ρ(u)du > 0, a
since ρ > 0 on (0, 1), hence Φ(a, c) > Φ(a, b) + Φ(b, c). This shows that splitting an interval strictly decreases the cost. Consequently, if a schedule has a collapsed interval, the corresponding redundant point can be moved to split some non-trivial interval and strictly decrease the objective. Therefore the minimizer must satisfy 0 < β1 < · · · < βK−1 < 1. Next, for a < b, Z b ∂b Φ(a, b) =
∂a Φ(a, b) = −(b − a)ρ(a).
ρ(u)du, a
Since βk appears only in Φ(βk−1 , βk ) and Φ(βk , βk+1 ), the stationary condition at an interior minimizer gives Z βk 1 ∂εfact = ρ(u)du − (βk+1 − βk ) ρ(βk ). 0= N ∂βk βk−1 22
which yields the desired recursion. For the proof of uniqueness, the recursion formula (11) can be written as ∆k+1 = Ψ(βk , ∆k ),
βk = βk−1 + ∆k ,
where 1 Ψ(b, δ) := ρ(b)
∆1 = β1 ,
Z b ρ(u)du. b−δ
First observe that Ψ is strictly increasing in its second argument: for δ < δ ′ with b − δ ′ ≥ 0, Ψ(b, δ ′ ) − Ψ(b, δ) =
1 ρ(b)
Z b−δ ρ(u) du > 0, b−δ ′
since ρ > 0 on (0, 1). Suppose now that two schedules with initial values β1 < β̃1 both satisfy the recursion and the ˜ 1 , and if βk < β̃k and ∆k < ∆ ˜ k , then terminal condition βK = β̃K = 1. Then ∆1 = β1 < β̃1 = ∆ ˜ k) = ∆ ˜ k+1 , ∆k+1 = Ψ(βk , ∆k ) ≤ Ψ(β̃k , ∆k ) < Ψ(β̃k , ∆ using first the monotonicity in the first argument and then the strict monotonicity in the second; consequently βk+1 < β̃k+1 . By induction βK < β̃K , contradicting βK = β̃K = 1. It remains to verify the monotonicity in b when ℓ := log ρ is concave. Substituting u = t+b−δ, Z δ Ψ(b, δ) =
eℓ(t+b−δ)−ℓ(b) dt,
0
so differentiating under the integral sign gives Z δ ∂b Ψ(b, δ) =
′ ℓ (t + b − δ) − ℓ′ (b) eℓ(t+b−δ)−ℓ(b) dt ≥ 0,
0
since concavity makes ℓ′ nonincreasing and t + b − δ ≤ b. Proof of Remark 2. First let ρ(u) = u−2 and consider the geometric schedule βk = (1 + a)βk−1 for k = 2, . . . , K, with β1 = (1 + a)−(K−1) so that βK = 1. For every k ≥ 2 Z βk
1 du = 2 u βk−1
Z βk
1 a du = , 2 u β k βk /(1+a)
so, since ρ(βk ) = βk−2 , the increment prescribed by (11) is 1 βk+1 − βk = ρ(βk )
Z βk ρ(u) du = aβk , βk−1
which is exactly the next geometric increment: the relations for k = 2, . . . , K − 1 hold. The Rβ relation for k = 1 involves 0 1 u−2 du = ∞ and therefore cannot hold; this is the only failing relation, and the sense in which the geometric schedule solves the recursion.
B.1
Proofs of subsection 3.1
Proof of Lemma 2. By (19), f (0) = −N −1 X [N ]\{j} ). Hence
PN
j=1 H(X
j
) and f (N − 1) = −N −1
D(π) = f (N − 1) − f (0) =
N −2 X i=0
23
ι(i).
PN
j=1 H(X
j
|
PN −1 The entropy chain rule, averaged over a uniform permutation, gives H(X) = − j=0 f (j). Therefore j−1 N −2 N −1 −1 X X X NX TC(π) = f (j) − f (0) = ι(i) = (N − i − 1)ι(i). j=0
j=0 i=0
i=0
Using DTC(π) = N D(π) − TC(π) then yields DTC(π) =
N −2 X
(i + 1)ι(i).
i=0
For the integral identities, bi is the density of Beta(i + 1, N − i − 1), so Z 1
Z 1 bi (u) du = 1,
N
Z 1 u bi (u) du = i + 1,
0
(1 − u)bi (u) du = N − i − 1.
N
0
0
Multiplying these identities by ι(i), summing over i, and applying (27) proves the claimed representations. Finally, both i + 1 and N − i − 1 are at least one on the summation range, so nonnegativity of ι gives D(π) ≤ min{TC(π), DTC(π)}. Proof of Proposition 2. The intervals (βk−1 , βk ], k = 1, . . . , K, partition (0, 1], and for any u ∈ (βk−1 , βk ], βk − u ≤ ∆βk ≤ ∆βmax . Therefore, by Theorem 1 and Lemma 2, Z 1 εfact (β) ≤ N ∆βmax
ρ(u)du = N ∆βmax D(π), 0
The constant step size ∆βk = 1/K gives ∆βmax = 1/K, from which we deduce the second display of the statement. Finally, I(X i ; X [N ]−i ) ≤ H(X i ) ≤ log |X | for every i, since conditional entropy is nonnegative and X i takes values in X ; averaging over i yields D(π) ≤ log |X |. Proof of Proposition 3. For the geometric schedule, split the first interval from the others as Z β1 (β1 − u)ρ(u) du + N
εfact = N 0
K Z βk X k=2
(βk − u)ρ(u) du.
βk−1
For the first term, since β1 = a/N , Z β1
Z β1 (β1 − u)ρ(u) du ≤ N β1
N
ρ(u) du ≤ a D(π).
0
0
For k ≥ 2 and u ∈ [βk−1 , βk ], we have βk − u ≤ βk − βk−1 ≤ aβk−1 ≤ au, hence N
K Z βk X k=2
Z 1 (βk − u)ρ(u) du ≤ aN
βk−1
u ρ(u) du ≤ a DTC(π), β1
using Lemma 2 and ρ ≥ 0. Combining the two estimates gives εfact (β) ≤ a(DTC(π) + D(π)), and D(π) ≤ DT C(π) yields the stated bound. For the reverse geometric schedule split the last interval instead. Since 1 − βK−1 = a/N and βK = 1, Z 1 Z 1 Z 1 N (1 − u) ρ(u) du ≤ N (1 − βK−1 ) ρ(u) du = a ρ(u) du ≤ a D(π). βK−1
βK−1
24
βK−1
For k ≤ K − 1 and u ∈ (βk−1 , βk ] we have 1 − βk−1 ≤ (1 + a)(1 − βk ), hence βk − u ≤ βk − βk−1 = (1 − βk−1 ) − (1 − βk ) ≤ a(1 − βk ) ≤ a(1 − u), so that N
K−1 X Z βk
Z βK−1 (βk − u) ρ(u) du ≤ aN
(1 − u) ρ(u) du ≤ a TC(π).
βk−1
k=1
0
Combining and using D(π) ≤ TC(π) concludes the proof.
C
Proofs of Section 4
Proof of Lemma 3. Set aN,j := (N − 1)ιN (j) and GN (u) :=
N −2 X
g
j=0
j N −2
BjN −2 (u).
Since the Bernstein basis is nonnegative and sums to one, ∥ρN − GN ∥∞ ≤ max aN,j − g 0≤j≤N −2
j N −2
−→ 0.
Bernstein’s approximation theorem gives ∥GN − g∥∞ → 0 because g is continuous. The triangle inequality proves the claim. Proof of Theorem 2. Let h = 1/K, tk = kh and ∆βk := βk − βk−1 and set Z βk IkN = (βk − u)ρN (u)du, βk−1
PK N so that εN,K k=1 Ik by Theorem 1. fact (β) = N First notice that since β ∈ C 1 ([0, 1]) we have maxk ∆βk ≤ ∥β ′ ∥∞ h = O(K −1 ), and therefore K X
K X (∆βk )2 ≤ max ∆βk ∆βk = O(K −1 ). k
k=1
k=1
The proof proceeds in three steps. Step 1: replacing ρN by g. Let ηN := ∥ρN − g∥∞ , which tends to zero by Assumption 1. Then K X
IkN −
k=1
K Z βk X k=1
(βk − u)g(u) du ≤ ηN
βk−1
K Z βk X k=1
(βk − u) du
βk−1
K
=
ηN X (∆βk )2 = o(K −1 ). 2 k=1
Step 2: freezing g on each interval. As g is uniformly continuous, let ωg denote its modulus of continuity; then for u ∈ [βk−1 , βk ] |g(u) − g(βk−1 )| ≤ ωg (∆βk ), Hence
Z βk (βk − u)g(u) du = βk−1
ωg (δ) −→ 0. δ↓0
1 g(βk−1 )(∆βk )2 + O (∆βk )2 ωg (∆βk ) . 2
Summing over k and using maxk ∆βk = O(K −1 ), gives K Z βk X k=1
βk−1
K
(βk − u)g(u) du =
1X g(βk−1 )(∆βk )2 + o(K −1 ). 2 k=1
25
Step 3: Riemann sum. Let ωβ ′ be the modulus of continuity of β ′ . By the mean value theorem, ∆βk = hβ ′ (ξk ) for some ξk ∈ (tk−1 , tk ), and |β ′ (ξk ) − β ′ (tk−1 )| ≤ ωβ ′ (h), so (∆βk )2 − h2 β ′ (tk−1 )2 ≤ 2∥β ′ ∥∞ ωβ ′ (h) h2 = o(h2 ) uniformly in k. Since g is bounded, summing the K error terms gives o(h), and therefore " K # K X h 1X 2 ′ 2 h g(βk−1 )(∆βk ) = g(β(tk−1 ))β (tk−1 ) + o(h) 2 2 k=1 k=1 Z 1 1 g(β(t))β ′ (t)2 dt + o(K −1 ), = 2K 0 the last step because t 7→ g(β(t))β ′ (t)2 is continuous, so the bracket is a Riemann sum converging to its integral. Combining these estimates and multiplying by N proves the claim. Note that the error term of Step 1 requires only N → ∞, while those of Steps 2 and 3 depend only on K; the two limits therefore decouple, and no relation between N and K is needed. Proof of Corollary 2. Let Z 1 I(β) =
g(β(t))β ′ (t)2 dt.
0
To characterize the minimizer of I(β), observe that Z 1
p g(u) du
2
0
Z 1 =
p
2 Z 1 g(β(t))β ′ (t)2 dt, g(β(t)) β (t) dt ≤ ′
0
0
where the first equality comes from a change of variable, since β is increasing and β(0) = 0, β(1) = 1, while p the second comes from Cauchy-Schwarz. As a consequence equality holds if and only if t 7→ g(β(t)) β ′ (t) is constant, that is, β ′ (t) ∝ g(β(t))−1/2 . It remains to check that this equality case is attained by Ran p admissible schedule. Since g is y continuous and strictly positive on [0, 1], the function G(y) = 0 g(u) du is a strictly increasing √ C 1 bijection from [0, 1] onto [0, G(1)] with G′ = g > 0 on [0, 1], so β ∗ (t) := G−1 (t G(1)) is well defined, of class C 1 on [0, 1], and satisfies β ∗ (0) = 0,
β ∗ (1) = 1,
(β ∗ )′ (t) = q
G(1) g β ∗ (t)
> 0,
p so that g(β ∗ ) (β ∗ )′ ≡ G(1) is constant and I(β ∗ ) = G(1)2 . Substituting into (14) yields the ∗ stated expression for εN,K fact (β ). Proof of Proposition 4. By (13), Lemma 3 gives Assumption 1, and D(πN ) → D∞ . Dividing the rescaled increments in (13) by D(πN ) therefore gives uniform convergence to g/D∞ , which is the normalized profile required by Lavenant and Zanella (2025, Assumption 1); changing between their grid and j/(N − 2) does not affect the limit, by continuity of g. Both limits defining ∆s̄ then exist, by Theorem 2 and by Lavenant and Zanella (2025, Thm. 15) respectively, after undoing their normalization. Subtracting the two limiting expressions and writing θ(t) = {s̄β ′ (t)}, " # Z 1 Z 1 i h ′ 2 ′ θ(t) 1 − θ(t) ′ ′ 2∆s̄ = g β(t) s̄ β (t) − s̄ hs̄ β (t) + β (t) dt = g β(t) β (t) − dt, s̄ 0 0 which is the stated identity. For the sign, since s̄β ′ (t) = ⌊s̄β ′ (t)⌋ + θ(t) with θ(t) ∈ [0, 1), θ(t) 1 − θ(t) ≤ θ(t) ≤ s̄ β ′ (t),
26
so the integrand is nonnegative and ∆s̄ ≥ 0. For the size of the difference, using 0 ≤ θ(1 − θ) ≤ 41 in the identity above, Z 1 1 g β(t) dt ≤ 2∆s̄ ≤ D∞ , D∞ − 4s̄ 0 R1 1 which gives both 0 ≤ ∆s̄ ≤ 12 D∞ and ∆s̄ − 21 D∞ ≤ 8s̄ g(β(t)) dt ≤ ∥g∥∞ /(8s̄). 0
D
Proofs of Section 5
We first express the conditional mutual information profile in terms of average subset entropies, then derive a representation of ρ useful for stationary targets. Lemma 6. For m = 0, . . . , N , define H̄m :=
1
X
N m
H(X A ).
A⊆[N ], |A|=m
Then f (i) = H̄i − H̄i+1 for i = 0, . . . , N − 1, and ι(i) = 2H̄i+1 − H̄i − H̄i+2 for i = 0, . . . , N − 2. Proof of Lemma 6. The pair (A, j) := (σ≤i , σi+1 ) is uniform over the ordered pairs with |A| = i and j ∈ / A, so by the chain rule H(X A∪{j} ) = H(X A ) + H(X j | X A ), X 1 H(X A∪{j} ) − H(X A ) . Eσ H(X σi+1 | X σ≤i ) = N i (N − i) |A|=i, j ∈A / N Each B with |B| = i + 1 arises from exactly i + 1 such pairs A ∪ {j}, and Ni (N − i) = i+1 (i + 1), so the first sum equals H̄i+1 ; the second equals H̄i , since each A occurs for N − i choices of j. Hence f (i) = H̄i − H̄i+1 , and the expression for ι(i) = ∆f (i) = f (i + 1) − f (i) follows. iid
Lemma 7. For u ∈ [0, 1] let Tu := {t ∈ [N ] : ξt = 1} with ξ1 , . . . , ξN ∼ Bern(u), and define H(u) := E H(X Tu ) with the entropy evaluated for each fixed subset. Then ρ(u) = − N1 H′′ (u). m N N Proof of Lemma 7. Recall that Bm := m u (1 − u)N −m is the m-th Bernstein basis polynomial n of degree N , with the convention Bj ≡ 0 for j ∈ / {0, . . . , n}. Conditionally on |Tu | = m, the set Tu is uniform among the subsets of size m, so X 1 H̄m := E H(X Tu ) |Tu | = m = N H(X A ) m
A⊆[N ], |A|=m
and H(u) = E E H(X Tu ) |Tu | N N X X N m N = u (1 − u)N −m H̄m = Bm (u) H̄m . m m=0 m=0 P P d n n−1 Differentiating twice the Bernstein identity du m Bm (u)cm = n m Bm (u)(cm+1 − cm ) and re-indexing, N −2 X H′′ (u) = N (N − 1) H̄i+2 − 2H̄i+1 + H̄i BiN −2 (u). i=0
By Lemma 6 the bracket equals −ι(i), and by (8) we obtain H′′ (u) = −N (N − 1)
N −2 X
ι(i)BiN −2 (u) = −N ρ(u).
i=0
27
Proof of Proposition 7. Write HN (u) := E[H(X Tu )] with Tu ⊆ [N ] as in Lemma 7. For each realization of Tu , the chain rule in increasing order of the indices gives H(X Tu ) =
N X
ξt H(X t | X Tu ∩[t−1] ).
t=1
Since ξt is independent of Tu ∩ [t − 1], HN (u) = u
N X E H(X t | X Tu ∩[t−1] ) . t=1
The lag set {t − a : a ∈ Tu ∩ [t − 1]} has the law of St−1,u , and by stationarity H(X t | X Tu ∩[t−1] ) = h(St−1,u ) in distribution. Using h(A) = h0 − I(A) and I(∅) = 0, HN (u) = u
N N −1 X X E h(St−1,u ) = N h0 u − u E I(St,u ) . t=1
(30)
t=1
Since the first term is linear in u, Lemma 7 gives the finite-N formula. iid For the limit, realize all the Bernoulli sets on one probability space: let ξd ∼ Bern(u) for d ≥ 1 and put St,u := {d ∈ [t] : ξd = 1}, Su := {d ≥ 1 : ξd = 1}, so that St,u = Su ∩ [t] ↑ Su . For A ⊆ N set I(A) := limm I(A ∩ [m]), the limit existing because m 7→ I(A ∩ [m]) is nondecreasing and bounded by h0 . Then I(St,u ) ↑ I(Su ) surely, so EI(St,u ) → EI(Su ) by monotone convergence, and by Cesàro ΦN (u) :=
N −1 u X E I(St,u ) −→ Φ∞ (u) := u E I(Su ) N t=1
for every u ∈ [0, 1].
By the above, ρN = Φ′′N , and under Assumption 1 ρN → g uniformly. Since ΦN (0) = 0, Taylor’s formula with integral remainder gives Z u ΦN (u) = Φ′N (0) u + (u − s) Φ′′N (s) ds. 0
Taking u = 1 and using the uniform convergence of Φ′′N together with ΦN (1) → Φ∞ (1), the Φ′N (0) converges, say to c. Passing to the limit in the display yields Φ∞ (u) = cu + Rsequence u (u − s)g(s) ds, so Φ∞ is twice continuously differentiable with Φ′′∞ = g, that is, 0 ρN (u) −→
d2 u EI(Su ) . 2 du
Proof of Proposition 6. By the Markov property, conditioning on the whole revealed past is equivalent to conditioning on the nearest revealed coordinate, so I(A) = Imin A for every nonempty finite A ⊆ N, and I(∅) = 0. Since P(min St,u = j) = u(1 − u)j−1 for j = 1, . . . , t, t X E I(St,u ) = u(1 − u)j−1 Ij . j=1
Substituting into (30) and exchanging the order of summation, each j ∈ {1, . . . , N − 1} occurring for the N − j indices t ∈ {j, . . . , N − 1}, u
N −1 X t=1
−1 NX E I(St,u ) = (N − j) Ij u2 (1 − u)j−1 = GN (u), j=1
so that HN (u) = N h0 u − GN (u), and (15) follows from Lemma 7. PN −1 For the limit, set wj (u) := u2 (1 − u)j−1 , so that ρN = j=1 1 − Nj Ij wj′′ . Writing m = j − 1, wj′′ (u) = 2(1 − u)m − 4mu(1 − u)m−1 + m(m − 1)u2 (1 − u)m−2 , 28
2 m−2 and since maxu u(1 − u)m−1 ≤ 1/m ≤ 4/m2 , we get supu |wj′′ (u)| ≤ 10 for u u (1 − u) P and max ′′ every j. Therefore, with ρ∞ := j≥1 Ij wj , N −1
ρN − ρ ∞ ∞ ≤
X 10 X j Ij + 10 Ij . N j=1 j≥N
P
Both terms vanish as N → ∞: the second because j Ij < ∞, the first by Kronecker’s lemma. The convergence is therefore uniform, and since the differentiated series converges uniformly, P term-by-term differentiation of j≥1 Ij wj (u) = u E[IDu ], Du ∼ Geom(u), is legitimate, giving d2 ρ∞ = du 2 u E[IDu ] . Proof of Proposition 8. For every fixed permutation, exchangeability gives I(X σi+1 ; X σi+2 | X σ≤i ) = I(X i+1 ; X i+2 | X 1:i ), so that ιN (m) = ι(m) for m ≤ N − 2. Substitution into (9) proves (17). For the asymptotics, fix u ∈ (0, 1), let BN ∼ Bin(N − 2, u) and AN := {BN ≥ u2 (N − 2)}. By a Chernoff bound there is cu > 0, nondecreasing in u, with P(AcN ) ≤ e−cu N . Since X is finite, 0 ≤ ι(m) ≤ log |X |, so (N − 1)E[ι(BN )1AcN ] ≤ (N − 1) log |X | e−cu N → 0. Given ε > 0, choose mε with ι(m) ≤ ε/m for m ≥ mε ; for N large enough, BN ≥ u2 (N − 2) ≥ mε on AN , so 2(N − 1) . (N − 1)E ι(BN )1AN ≤ ε u(N − 2) Hence lim supN ρN (u) ≤ 2ε/u for every ε > 0, so ρN (u) → 0. For u = 1, BN = N − 2 almost surely and ρN (1) = (N − 1)ι(N − 2) → 0 since ι(m) = o(1/m). For u = 0, BN = 0 almost surely and ρN (0) = (N − 1)ι(0) = (N − 1)I(X 1 ; X 2 ) → +∞. Proof of Proposition 9. Write ι(m) = I(X m+1 ; X m+2 | X 1:m ) and KN = 1 + ⌈log2 N ⌉; nondegeneracy of Q gives ι(0) = I(X 1 ; X 2 ) > 0. Uniform schedule. Since ι(m) ≥ 0, keeping only the m = 0 term in (17) gives ρN (u) ≥ (N − 1)ι(0)(1 − u)N −2 , and keeping only k = 1 in (10), εunif fact ≥ N (N − 1)ι(0)
Z 1/KN 0
1 − u (1 − u)N −2 du. KN
For N large enough, KN ≤ N/2, so 1/N ≤ 1/(2KN ); restricting the integral to [0, 1/N ] and using 1 1 1 1 N −2 ≥ (1 − 1/N )N −2 ≥ 13 there, KN − u ≥ KN − N ≥ 2KN together with (1 − u) εunif fact ≥ N (N − 1)ι(0) ·
1 1 1 (N − 1)ι(0) CN · · = ≥ 2KN 3 N 6KN log N
for a constant C > 0, since KN = O(log N ). Geometric schedule. The schedule β1 = 1/N , βk = min{1, 2βk−1 } is the case a = 1 of Proposition 3 and reaches 1 in at most KN steps, so ε⋆fact (πN , KN ) ≤ εgeom fact (πN , KN ) ≤ 2 DTC(πN ). By exchangeability H(X i | X [N ]\{i} ) = H(X N | X 1:N −1 ) ≥ H(X N | X 1:N −1 , P ) = E[H(P )], and H(X 1:N | P ) = N E[H(P )] because the coordinates are conditionally i.i.d.; hence DTC(πN ) = H(X 1:N ) − N H(X N | X 1:N −1 ) ≤ H(X 1:N ) − H(X 1:N | P ) = I(P ; X 1:N ). For a p-dimensional parametric prior satisfying the regularity conditions of Clarke and Barron (1994), I(P ; X 1:N ) = p2 log N + O(1), so ε⋆fact (πN , KN ) = O(log N ). Combining the two bounds gives the stated ratio.
29
E
Proofs of Section 6
Proof of Proposition 10. Conditionally on Z, the index J is uniform on [N ] \ Z and independent of X, i h P j Z ˜ (31) f (i) = E log π(X J | X Z ) = E N1−i j ∈Z / log πj (X | X ) = E fi . Conditionally on (Z, X Z ), each unrevealed coordinate X j has distribution πj (· | X Z ). Therefore E[f˜i | Z, X Z ] =
1 X E log πj (X j | X Z ) Z, X Z N −i j ∈Z /
1 X =− H πj (· | X Z ) = fˆi . N −i
(32)
j ∈Z /
Taking expectation with respect to (Z, X Z ) and using the tower property together with (31) yields h h ii E fˆi = E E f˜i | Z, X Z = E f˜i = f (i). Finally, the law of total variance applied to f˜i and (32) give Var(f˜i ) = E[Var(f˜i | Z, X Z )] + Var(E[f˜i | Z, X Z ]) ≥ Var(fˆi ). Thus, fˆi is the Rao-Blackwellization of f˜i with respect to (Z, X Z ). Proof of Proposition 11. Conditionally on (Z, X Z ), write pj := πj (· | X Z ),
qj := pjθ (· | X Z ),
j∈ / Z,
so that DKL (pj ∥qj ) = −E log qj (X j ) + E log pj (X j ) . Averaging over j ∈ / Z and taking expectation over (Z, X Z ), the second term is f (i) by (31) and the first is −E[f˜iθ ], giving E[f˜iθ ] − f (i) = −Li . The sign comes from the non-negativity of KL divergence. For the entropy estimator, let Tj := 21 ∥pj − qj ∥1 . The Fannes-Audenaert inequality (Fannes, 1973; Audenaert, 2006) gives |H(pj ) − H(qj )| ≤ ϕ(Tj ). (33) Since E[fˆiθ ] − f (i) = E[H(pJ ) − H(qJ )], and ϕ is concave on [0, 1], being the sum of a linear term and the concave binary entropy, Jensen’s inequality yields E[fˆiθ ] − f (i) ≤ E |H(pJ ) − H(qJ )| ≤ E[ϕ(TJ )] ≤ ϕ(E[TJ ]). By Pinsker’s inequality, the concavity of the square root and the assumption on Li , hp i p p 0 ≤ E[TJ ] ≤ E DKL (pJ ∥qJ )/2 ≤ E[DKL (pJ ∥qJ )]/2 = Li /2 ≤ 1 − 1/|X | ) As ϕ′ (T ) = log (|X |−1)(1−T > 0 for T < 1 − 1/|X |, ϕ is nondecreasing on (0, 1 − 1/|X |), so we T conclude that p E[fˆθ ] − f (i) ≤ ϕ(E[TJ ]) ≤ ϕ Li /2 . i
For p the rate, H2 (T ) = T log(1/T ) + O(T ) as T → 0, so ϕ(T ) = O(T log(1/T )). Substituting T = Li /2 gives the stated rate.
30
θ Proof of Proposition 12. Since 0 ≤ H(p) ≤ log |X | for every p ∈ P(X ), each fˆi,m takes values in 2 (log |X |) θ ˆ and, by independence and [− log |X |, 0]. Popoviciu’s inequality therefore gives Varfi,m ≤ 4 identical distribution, 1 (log |X |)2 θ Var(f¯i ) = Var(fˆi,1 )≤ . n 4n Hoeffding’s inequality for bounded independent summands gives, for every t > 0,
P
f¯i − E[fˆiθ ] ≥ t ≤ 2 exp −
2nt2 , (log |X |)2
q and taking t = log |X | log(2/δ) proves the claimed probability bound. 2n Proof of Proposition 13. For i = 0, . . . , N − 2, set wi (β) := N
K Z βk X k=1
(βk − u)bi (u) du,
βk−1
bi (u) = (N − 1)BiN −2 (u), as in (27). The definition of ρa gives εfact (a, β) =
N −2 X
a(i)wi (β),
i=0
proving linearity. Since bi is nonnegative and integrates to one, 0 ≤ wi (β) ≤ N ∆βmax , consequently, |εfact (ι, β) − εfact (ῑ, β)| ≤ ∥ι − ῑ∥1 ∥w(β)∥∞ ≤ N ∆βmax ∥ι − ῑ∥1 , which proves (23). For the optimization bound, applying (23) at β̂ gives εfact (ι, β̂) ≤ εfact (ῑ, β̂) + N ∆β̂max ∥ι − ῑ∥1 , while optimality of β̂ and another application of (23), ∗ εfact (ῑ, β̂) ≤ εfact (ῑ, β ∗ ) ≤ εfact (ι, β ∗ ) + N ∆βmax ∥ι − ῑ∥1 .
Combining these inequalities yields εfact (ι, β̂) − εfact (ι, β ∗ ) ≤ N ∆max (β̂) + ∆max (β ∗ ) ∥ι − ῑ∥1 , which implies the stated upper bound. The lower bound follows from the optimality of β ∗ for the true profile.
E.1
Implementation details for the mixture target
Let U ⊆ [N ] be the revealed set of coordinates with values X U . By Bayes’ rule, the posterior distribution of the latent variable is Q wz j∈U µz,j (xj ) U Q P (Z = z | x ) = Pr . (34) j ′ ′ z ′ =1 wz j∈U µz ,j (x ) Conditional independence given Z then yields, for every B ⊆ [N ] \ U , π(xB | xU ) = πi (xi | xU ) =
r X z=1 r X
P(Z = z | X U = xU )
Y
µz,i (xi ),
(35)
i∈B
P(Z = z | X U = xU ) µz,i (xi ),
z=1
31
i∈ / U.
(36)
Estimating ι and ρ through the auxiliary profile. We estimate the information profile f using the entropy estimator in (22). For each Monte Carlo draw, we sample X ∼ π and, independently, a uniformly random permutation σ of [N ]. For i = 0, . . . , N − 1, let Ui = {σ1 , . . . , σi } be the set of coordinates revealed after i steps. Conditional on X Ui , the contribution to the estimate of f (i) is 1 X H πj ( · | X Ui ) . fˆi = − N −i j ∈U / i
Starting from the prior weights P(Z = z | X U0 ) = wz , we update the latent log-posterior after revealing coordinate σi+1 as log P Z = z | X Ui+1 = log P Z = z | X Ui + log µz,σi+1 (X σi+1 ) + const. Thus, a single pass through the permutation produces the entire vector (fˆ0 , . . . , fbN −1 ). Direct evaluation costs O(N 2 rL) per draw: for each of the N revealed-set sizes, we compute the conditional distributions and entropies of the remaining coordinates, at a cost of O(rL) per coordinate. We average the profiles over 4000 independent draws. Let f MC denote the resulting raw average. To enforce the monotonicity of the true profile, we compute its isotonic regression, that is its least-squares projection onto the set of non-decreasing sequences N −1 X 2 f¯ := arg min vi − fiMC . v0 ≤···≤vN −1
i=0
We then set ῑ(i) := f¯i+1 − f¯i ≥ 0,
i = 0, . . . , N − 2,
and construct ρ̄ using (8). Monte Carlo estimate of εfact . Independently of the information-profile calculation, for each target and schedule β, we estimate εfact (β) directly from (4) using 4000 draws. For each draw, we sample X ∼ π and, following the multinomial coupling used in the proof of Theorem 1, independently assign each coordinate i to a reveal step distributed as Cat(∆β1 , . . . , ∆βK ), realizing the distribution ν β . We then proceed block by block and accumulate the difference between the log block conditional in (35) and the sum of the corresponding log one-coordinate conditionals in (36).
32