Sharp Convergence Rates for Masked Diffusion Models
| Authors | Yuchen Liang et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.22505 |
| Download | |
| Categories | cs.LG, stat.ML |
Abstract
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 and the target accuracy . 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.
Engineering Breakdown
Plain English
This paper provides the first rigorous theoretical analysis of sampling algorithms for masked diffusion models, which are discrete diffusion models used for text generation and other symbolic tasks. The authors prove sharp convergence rate guarantees for both the Euler sampler (standard in practice) and the First-Hitting Sampler (FHS, a newer high-performance variant), using Total Variation distance instead of the looser KL divergence metric used in prior work. Their analysis removes strong assumptions about score estimation quality that plagued earlier theoretical work and directly covers FHS for the first time. The result is a set of convergence guarantees with tight parameter dependencies that accurately reflect what practitioners observe empirically.
Core Technical Contribution
The paper's core innovation is replacing KL divergence analysis with Total Variation distance bounds for masked diffusion samplers, which yields significantly tighter convergence rates. The authors develop a novel proof technique that characterizes the convergence of both Euler and FHS methods without requiring unrealistic assumptions on score function accuracy—previous work assumed near-perfect score estimation, which is impossible in practice. They provide the first theoretical analysis of FHS, proving it achieves comparable or better convergence rates than Euler despite its different sampling strategy. This closes a critical gap between theory and practice: the theoretical framework now justifies why FHS performs well empirically in real systems.
How It Works
Masked diffusion models operate on discrete tokens by gradually replacing tokens with a mask symbol (like [MASK]) following a corruption schedule, then learning to reverse this process. During sampling, the model uses a learned score function to predict which unmasked tokens should be revealed at each step, iteratively replacing mask tokens until the sequence is complete. The Euler sampler greedily selects tokens with highest predicted probability each step, while FHS samples the first token that reaches a decision threshold, then recursively fills in remaining positions. The convergence analysis measures how close the sampled distribution gets to the true data distribution in Total Variation distance, tracking how token prediction error accumulates across T sampling steps. By bounding the error at each step and summing across steps, the authors derive overall convergence rates that depend on the diffusion schedule, number of steps T, vocabulary size, and the quality of the learned score function.
Production Impact
For teams deploying masked diffusion models in production (text generation, code completion, data augmentation), this paper provides principled guidance on sampler selection and step budgets. Engineers can now confidently use FHS instead of Euler if latency permits, knowing the theoretical convergence is sound—prior work offered no such assurance. The tighter bounds tell practitioners how many sampling steps are actually necessary to reach target quality: if the paper shows convergence at T=100 steps with specific error bounds, you don't waste compute running 500 steps unnecessarily. The removal of unrealistic score-estimation assumptions means the theory now applies to real models trained with standard objectives, making it actionable for production systems. However, practitioners must still validate that empirical performance matches theoretical predictions on their specific domain and model scale, as the constants in the convergence rates may be loose.
Limitations and When Not to Use This
The paper assumes access to an accurate score function, which in practice is learned from finite training data with potential domain shift—the bounds don't account for distribution mismatch between training and sampling. The analysis is limited to discrete tokens and standard diffusion schedules; it doesn't cover advanced variants like non-uniform corruption schedules or hierarchical masking strategies used in some production systems. The convergence rates are stated asymptotically and the implicit constants are not computed, so practitioners cannot directly predict wall-clock performance on real hardware without extensive benchmarking. The theoretical framework doesn't address conditional generation settings (e.g., text-to-text with prefix constraints), which are common in production; extending the analysis to these settings remains open.
Research Context
This work builds on the recent surge in discrete diffusion models for language, improving upon prior theoretical analyses by Hoogeboom et al. and others who studied these models primarily through KL divergence. The paper directly responds to the empirical success of FHS in concurrent work on masked diffusion, providing the missing theoretical justification that was absent from purely empirical comparisons. It contributes to the broader effort to close the theory-practice gap in diffusion models, following similar theoretical refinements in continuous diffusion (e.g., improved convergence proofs for score-based models). The sharper bounds and fewer assumptions open avenues for studying hybrid samplers, adaptive step allocation, and optimal schedule design under the same unified framework.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::