Source-linked AI summary
Sharp Convergence Rates for Masked Diffusion Models
Yuchen Liang, Zhiheng Tan, Ness Shroff, Yingbin Liang
TL;DR
Masked diffusion samplers have strong empirical performance but limited theory, especially for tight parameter dependence and FHS error. The paper develops direct TV analyses for Euler and FHS, establishing improved Euler guarantees with matching lower bounds and showing that FHS error is exactly governed by score estimation.
Problem
Existing KL-based analyses for masked diffusion samplers can have loose parameter dependencies, strong score assumptions, and do not characterize FHS theoretically.
Method
The paper develops a direct total-variation analysis for Euler and a structure-exploiting error analysis for FHS.
Results
Euler obtains improved parameter dependencies and matching lower bounds, while FHS incurs no sampling error beyond score-estimation error with a matching lower bound.
Takeaways & Limitations
The analyses provide tight TV convergence guarantees for Euler and dimension- and vocabulary-independent error characterization for FHS.
Takeaways & Limitations
The Euler analysis assumes an L1 score-estimation error condition.
Abstract
from arXiv · showhide
Discrete diffusion models have achieved strong empirical performance in text and other symbolic domains, with masked (absorbing-rate) variants emerging as competitive alternatives to autoregressive models. Among existing samplers, the Euler method remains the standard choice in many applications, and more recently, the First-Hitting Sampler (FHS) has shown considerable promise for masked diffusion models. Despite their practical success, the theoretical understanding of these samplers remains limited. Existing analyses are conducted in Kullback-Leibler (KL) divergence, which often yields loose parameter dependencies and requires strong assumptions on score estimation. Moreover, these guarantees do not cover recently developed high-performance sampler of FHS. In this work, we first develop a direct total-variation (TV) based analysis for the Euler method that overcomes these limitations. Our results relax assumptions on score estimation, improve parameter dependencies, and establish convergence guarantees without requiring any surrogate initialization. Also for this setting, we provide the first convergence lower bound for the Euler sampler, establishing tightness with respect to both the data dimension $d$ and the target accuracy $\varepsilon$. Finally, we analyze the FHS sampler and show that it incurs no sampling error beyond that induced by score estimation, which we show to be tight with a matching lower error bound. Overall, our analysis introduces a direct TV-based error decomposition along the CTMC trajectory and a decoupling-based path-wise analysis for FHS, which may be of independent interest.
1 Introduction
The paper addresses limited and often loose theory for masked diffusion samplers by developing TV-based analyses for Euler and FHS, including matching lower bounds and finite-step guarantees.
- Motivation: Existing discrete-diffusion analyses largely use KL divergence, while tight TV guarantees and matching lower bounds remain missing.The paper motivates TV analysis as a route to sharper parameter dependence, particularly for dimension d.
- Motivation: FHS generates samples in exactly d steps, but its theoretical sampling error and matching lower bound had not been established.Its analysis requires techniques exploiting the sampler’s distinct structure.
- Contributions: The paper gives affirmative answers to both convergence questions through direct TV analysis and new FHS analysis.Its contributions include an Euler-method lower bound and a tight FHS error characterization.
- Euler method: The Euler method improves ε dependence while retaining best-known dependencies on d and S, without bounded-score assumptions or surrogate initialization.The required number of steps is e O(dS/√ε) for a slightly perturbed target distribution, improving prior results by O(1/√ε).
- Euler method: The Euler lower bound is tight in both data dimension d and target accuracy ε under constant step-sizes and without early-stopping.This identifies intrinsic convergence limits for Euler sampling.
- FHS: FHS incurs no sampling error beyond score-estimation error, with a tight lower bound and no additional dependence on d or vocabulary size S.The analysis uses token-wise error decomposition and an intrinsic decoupling structure.
2 Preliminaries of Discrete Diffusion Models
The preliminaries formulate discrete diffusion through CTMCs and related masked-diffusion models, then introduce reverse-time sampling, Euler discretization, and FHS.
- CTMC framework: Discrete diffusion models use CTMCs over [S]^d, where d is the number of tokens and S is vocabulary size.The forward process is specified by a rate matrix governing instantaneous state transitions.
- CTMC framework: Independent and homogeneous token evolution simplifies computation by factorizing the forward conditional distribution.The paper uses a constant noise schedule β_t ≡ 1 in its absorbing-rate setting.
- Reverse process: The reverse process is the exact time reversal of the forward CTMC and is also a CTMC with a reverse rate matrix.Practical sampling may use early-stopping near t = 0, while the concrete score is estimated by a neural network.
- Euler method: Euler sampling discretizes the reverse-time process on a time grid and updates tokens using estimated reverse rates.The step-sizes are η_k = t_{k+1} − t_k, and token-wise transitions replace one coordinate at a time.
- Masked diffusion models: Masked diffusion models are equivalent to absorbing-rate CTMCs under the schedule α_t = e^-t, sharing the same reverse process.This relationship lets the paper use MDMs and absorbing-rate diffusion models interchangeably.
- First-Hitting Sampler: FHS directly simulates individual unmasking events instead of using a predefined time grid, requiring exactly d steps for a complete sequence.The section identifies unresolved questions about how score-estimation error affects FHS convergence.
3 Convergence Guarantees for the Euler Method
The Euler analysis directly decomposes total-variation error, yielding improved convergence guarantees under weaker score-estimation assumptions and a matching lower bound. In the absorbing-rate setting, the resulting step complexity is tight in data dimension and target accuracy up to logarithmic factors.
- General TV decomposition: The direct TV analysis bounds final error by initial TV mismatch plus rate-matrix mismatch accumulated along the sampling path.Unlike KL-based analyses, the accumulation term depends on absolute rate differences rather than a Bregman divergence.
- General TV decomposition: The analysis requires only an averaged L1 score-estimation error over states and time, rather than a uniformly small error at every state and time.This avoids the loose behavior near t ≈ 0 associated with pointwise control when the score diverges.
- Convergence upper bound: With δ ≍√ε/d, T ≍log(d/√ε), and κ ≍√ε/(dS log δ−1), the method reaches TV error ≲√ε in N = e O(dS/√ε) steps.The stated bound follows from the theorem's error decomposition under the absorbing-rate matrix and prescribed Euler step sizes.
- Convergence upper bound: The absorbing-rate Euler method achieves O(1/√ε) sampling complexity while retaining O(dS) dependencies and requiring no bounded score estimates.The method can initialize directly at the all-mask singleton δ[MASK]d instead of using a surrogate distribution.
- Convergence lower bound: The first Euler convergence lower bound requires N = ˜Ω(d/√ε) steps for √ε-TV error, making the upper bound tight in d and ε up to log-factors.The lower bound assumes constant step sizes and no estimation error.
4 Convergence Guarantees for the FHS Algorithm
The FHS analysis bounds sampling error under score-estimation assumptions and shows that its finite-step error is entirely attributable to score estimation. A matching lower bound establishes tightness.
- The integrated score-entropy loss is equivalent up to a data-dependent constant to the standard NELBO training objective for masked diffusion models.
- FHS is analyzed under an assumption on the score-estimation error, with no-[MASK] data required by Theorem 4.
- FHS achieves ε-KL error in exactly d sampling steps, with no additional error beyond score estimation and no further dependence on d or vocabulary size S.
- A decoupling property makes transition time and index independent of estimation error, enabling path-wise error decomposition and conversion to sampling error through data processing.
- Theorem 5 provides the first FHS convergence lower bound and shows that the Theorem 4 upper bound is exactly tight.
5 Conclusion
The paper concludes that direct TV analysis improves convergence guarantees for masked diffusion models and that FHS has tightly characterized error. It also identifies high-order FHS variants as future work.
- The work introduces direct TV analysis with tighter parameter dependencies than prior KL-based results.
- The analysis removes bounded-score and surrogate-initialization requirements and establishes tight lower bounds for the Euler method in TV.
- FHS incurs no sampling error beyond score-estimation error, and a worst-case construction shows this bound is tight.
- The paper identifies high-order variants of FHS as an interesting direction for further acceleration.
C Proof of Proposition 1
The proposition proof exploits the absorbing score’s time structure to relate score-entropy loss to the training objective and bound the time-weighted loss.
- The absorbing score factors into an explicit time-dependent coefficient and a time-independent conditional-distribution estimate.
- The proof compares the time-weighted loss with score-entropy loss under an assumption on the estimated conditional probabilities.
- With step sizes η_k = κ min {1, T − t_k}, the analysis controls the time-weighted loss and the resulting number of steps.
D Proof of Theorem 1
Theorem 1 is proved by decomposing TV error into initialization, estimation, and discretization effects, then bounding these terms for a truncated τ-leaping representation of Euler sampling.
- The discrete proof replaces the unavailable continuous-space Reynolds Transport Theorem with a rate-of-change argument based on probability derivatives.
- The Euler method is analyzed through a truncated τ-leaping sampler that is asymptotically equivalent to Euler and has piecewise-constant sampling rates.
- The TV error is decomposed into estimation and discretization components after bounding the initialization contribution from the all-mask state.
- Under the prescribed step-size schedule, the second discretization term dominates the first in the resulting upper-bound analysis.
G Proof of Theorem 3
This proof constructs an augmented forward process and analyzes its Euler discretization to derive the required sampling-step bound.
- Construction: The proof defines qγ as a γ-perturbation of a delta distribution and couples processes initialized from qγ and δa.The augmentation is key to the coupling argument.
- Euler discretization: The Euler method is applied on equally spaced discretization points over the augmented interval [γ, T + γ].The discretization follows the Euler update defined in the proof.
- Process properties: The augmented absorbing-rate process preserves tokens after they jump from [MASK] to ai, with transition probabilities independent of the token index i.This symmetry implies the resulting distribution pT is independent of i.
- Error and step bound: The proof bounds the perturbation and discretization contributions using γ = ε1/4 and derives the resulting requirement on the number of steps.The final step count is obtained after controlling terms involving γ and κ.
H Proof of Proposition 2
This proof connects score estimation and clean-data prediction through the model parameterization, then combines the resulting bounds to establish the proposition.
- Score–prediction connection: The proof uses the connection between the score estimator st(y, xt) and the predicted clean-data distribution µθ.This relation links the score-based quantities to token prediction probabilities.
- Loss decomposition: The integrated LSE loss is expanded into terms involving conditional token distributions and logarithms of µθ.The derivation uses the definition of the integrated LSE loss and the clean-data distribution connection.
- Termwise analysis: The proof separates the resulting expression into two terms and analyzes each using the model identities and conditional distributions.The subsets Mk contain uniformly sampled token indices, while xMc k 0 provides the corresponding conditional distribution.
- Conclusion: Combining the two bounds yields the proposition's claimed result.The final step combines equations (25) and (26).
I FHS Sampling in CTMC Framework: Proposition 3 and its Proof
This section characterizes FHS and the reverse absorbing-rate CTMC through their unmasking events, establishing the statistical structure needed for comparison.
- Path representation: Each process is represented by token unmasking times, the indices unmasked at those times, and the assigned token values.Because unmasked tokens remain fixed, these variables fully characterize the sampling path.
- Proposition 3: FHS and the reverse absorbing-rate CTMC share the same distributions for unmasking times and token indices, differing only in predicted token values.This is the central decoupling property established by Proposition 3.
- Reverse CTMC characterization: The proof derives reverse-CTMC transition rates, total leaving rates, and conditional distributions for the next unmasked index and token value.These distributions are then matched to the corresponding FHS update rules.
- Unmasking-time dynamics: Given the current state and unmasking time, the next unmasking time is independent of the current sequence state.The proof derives this independence for the reverse CTMC and notes the corresponding FHS property.
- FHS updates: FHS samples the next masked index uniformly and predicts its token value using µθ.The estimator µθ is compared with the ground-truth token-prediction distribution.
J Proof of Theorem 4
This proof analyzes FHS by comparing its full path distribution with the reverse CTMC, using Proposition 3 to isolate token-prediction error.
- Decoupling: Proposition 3 decouples FHS and reverse-CTMC paths: their unmasking times and indices match statistically, while token-value sampling may differ.This structure allows the path comparison to focus on predicted token values.
- Pathwise comparison: The path distributions are factorized and compared through KL(PCTMC(path)||PFHS(path)).The output-distribution bound is then obtained using the data processing inequality.
- Initialization: The common initialization consists of the fully masked state with ατd = 0.This initialization is shared by the two path measures.
- Error decomposition: The resulting pathwise KL divergence decomposes into a sum of expected KL divergences between token-prediction distributions.Thus, estimation error affects token-value sampling rather than unmasking times or indices.
K Proof of Theorem 5
The proof constructs a target and predictive model showing the FHS error bound is tight, then establishes supporting results for the time discretization and reverse CTMC analysis.
- FHS lower bound: The constructed instance yields a lower error bound, demonstrating tightness of the upper bound in Theorem 4.The proof explicitly identifies this special case as tight for the stated upper bound.
- FHS lower bound: A point-mass target q0 and a predictive model that assigns probabilities 1−ρ and ρ to the correct and alternative tokens construct the lower-bound instance.The construction fixes distinct unmasked sequences a and b and sets q0 = δa.
- FHS lower bound: FHS generates the target sequence with probability PFHS(a) = (1−ρ)^d because each token is predicted exactly once.Since q0 = δa, the resulting KL divergence is computed directly from this probability.
- Time discretization: The discretization proof splits the time-step analysis into cases before, after, and at the threshold k∗ where T−t_k crosses 1.The cases establish the step-size behavior and are summed to determine the total number of steps.
- Auxiliary proof: The indicator-function argument shows that its instantaneous change vanishes away from the boundary where the compared distributions are equal.Discreteness preserves the sign of the difference for sufficiently small time increments.
- Reverse CTMC analysis: For the reverse CTMC, the proof characterizes leaving rates, first unmasking times, and next-state distributions through local transition probabilities and an ODE.The instantaneous leaving rate is identified as Λ(t,x), while the survival function is used to analyze the first unmasking time.
L.6 Proof of Lemma 6
The lemma uses order statistics of independent uniform variables and the symmetry of the Beta distribution to derive the distribution of the relevant transformed unmasking times.
- Order-statistic representation: The variables u^(d) > ··· > u^(1) are order statistics of d independent Uniform(0,1) variables, with the k-th statistic following a Beta distribution.Specifically, the k-th order statistic has distribution Beta(k, d + 1 − k).
- Beta transformation: Using 1−X ∼ Beta(b,a) whenever X ∼ Beta(a,b), the proof derives the distribution of ατ_k−1.This Beta-distribution symmetry supplies the final transformation used in the lemma.