End-to-End Guarantees for Masked Diffusion Models under Conditional Mixing

Fareed Sheriff

cs.LGOct 29, 2026

End-to-End Guarantees for Masked Diffusion Models under Conditional Mixing

Authors: Fareed Sheriff

Organizations: Massachusetts Institute of Technology

Abstract

Masked diffusion models train a denoiser on masked cross-entropy and generate by unmasking many positions in parallel. Existing MDM sampling bounds assume an exact denoiser or leave its error uncontrolled while training bounds cover only the masking distribution on which the model was trained. We present an end-to-end guarantee for a learned denoiser under conditional mixing (CM): stationary sources (e.g. floored order- $r$ Markov and two-sided-Doeblin hidden-Markov sources) for which revealing distant visible tokens moves a masked token's posterior on average exponentially little in the distance with correlation length $τ_c$ . We show the empirical risk minimizer trained under any masking distribution attains per-token excess $\varepsilon'$ on that distribution from $\varepsilon'^{-\mathcal{O}(τ_c\log|\mathcal{V}|)}$ training sequences, independent of sequence length $N$ -- one sample budget serves every decoding schedule when each schedule's denoiser is trained under its own weighting. We further show that the exponent's order $Θ(τ_c\log|\mathcal{V}|)$ is minimax-optimal up to constants; more efficient rates require added structure. With this, we bound the end-to-end KL as $\mathrm{KL}(π\Vert p_{\mathrm{alg}})\leq \varepsilon+N\varepsilon'$ under a dilated schedule of $T=\mathcal{O}(τ_c\log(N|\mathcal{V}|/\varepsilon))$ denoising steps versus $Ω(N)$ steps for a fixed random order on an explicit CM source. Finally, we show that block-autoregressive samplers chain across blocks with no added error and that a $Λ$ -looped order- $r$ message-passing class on order- $r$ sources achieves polynomial sample complexity, with approximation error improving per loop until saturation at $Λr\approx τ_c\log(1/\varepsilon')$ . Experiments on released and retrained models support our results' qualitative predictions.

Figures & tables

Prediction	Result
Influence on MDLM’s conditionals
mean influence decays exponentially ( Assumption 1 )	exp. fit $R^{2}$ $0.989$ vs. power law $0.966$
worst-case influence stays large ( Assumption 2 )	$95$ th pct. $0.72$ at $d=64$
dilated factorization term $\hat{I}\approx 0$ by $b=16$	yes; random order $\neq 0$ at $64$
Block size on a synthetic source ( Corollary 3 )
held-out excess small at every $B$	$\leq 0.031$ nats

Table 1: Theory-motivated checks and results. Single runs; setups and estimators in Appendix B .

Panel of Figure 1 — Figure 1: Left: NLL under GPT-2-large (averaged over sequences) against decoding steps for dilated, random, and confidence-greedy schedules on MDLM; vertical line: the spacing at which the dilated $\hat{I}$ vanishes; gray band: half to twice that spacing; colored shading: bootstrap intervals. Middle: LLaDA-8B, dilated and random only, with MDLM’s threshold. Right: MAUVE on MDLM ( $512$ reference windows). Single runs; colors as on the left.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

	dilated finetune	uniform control	looped model	block models
Initialization	released MDLM	released MDLM	scratch, seed $0$	scratch, seed $0$ , shared across $B$
Data	OWT train, seed $0$	same windows as finetune	OWT train	synthetic, $20{,}000$ seqs
Context	$1024$	$1024$	$1024$	$256$
Tokens / steps	$0.25$ B / $7630$	$0.25$ B / $7630$	$4$ B / $7630$	plateau rule, $1500$ – $8000$ steps
Batch (sequences)	$32$	$32$	$512$	$256$
Optimizer	AdamW $(0.9,0.999)$	AdamW $(0.9,0.999)$	AdamW $(0.9,0.95)$	AdamW $(0.9,0.999)$

Appendix

Table 2: Training settings of the models trained in this paper. Weight decay for the looped model applies to matrices only. The dilated finetune and the uniform control share the same drawn training windows and differ in the masking law and loss target alone.

$R^{2}$ , mean decay exponential fit	$0.989$
$R^{2}$ , mean decay power-law fit	$0.966$
$(\hat{\tau}_{c},\widehat{C}_{0})$ on the tail window; interval for $\hat{\tau}_{c}$	$(81.1,\,1)$ ; $[43,95]$
local decay scale over the first distances	${\approx}26$
$95$ th percentile of $\hat{\kappa}$ at $d=1$ / $d=64$	$0.999$ / $0.722$
$\hat{I}(16)$ , dilated: point/upper interval edge (per summand)	$0.0036$ / $0.042$
$\hat{I}(64)$ , random order (per summand)	$0.087$

Appendix

Table 3: Influence measurement on MDLM’s conditionals.

$T$	$1$	$2$	$4$	$8$	$16$	$24$	$32$	$48$	$64$	$96$	$128$
released, dilated	$7.556$	$4.654$	$3.703$	$3.381$	$3.260$	$3.241$	$3.230$	$3.235$	$3.219$	$3.229$	$3.223$
released, random	$7.556$	$5.210$	$4.192$	$3.708$	$3.482$	$3.409$	$3.376$	$3.339$	$3.324$	$3.292$	$3.286$
dilated finetune, dilated	$7.557$	$4.633$	$3.685$	$3.371$	$3.256$	$3.226$	$3.212$	$3.200$	$3.194$	$3.190$	$3.188$
dilated finetune, random	$7.557$	$5.176$	$4.146$	$3.666$	$3.439$	$3.368$	$3.333$	$3.296$	$3.280$	$3.263$	$3.257$
uniform control, dilated	$7.556$	$4.618$	$3.664$	$3.347$	$3.231$	$3.201$	$3.186$	$3.174$	$3.167$	$3.164$	$3.161$
uniform control, random	$7.556$	$5.162$	$4.130$	$3.644$	$3.415$	$3.344$	$3.309$	$3.272$	$3.256$	$3.236$	$3.230$

Appendix

Table 4: Forward sampling error at all step counts for the six model-schedule pairs.

Panel of Figure 2 — Figure 2: Forward sampling error against decoding steps: dilated order on the released model and both finetunes, and random order on the released model; values from Table 4 . Left: all step counts; right: zoom on $T\geq 8$ .

	dilated		random		confidence
$T$	NLL	MAUVE	NLL	MAUVE	NLL	MAUVE
$1$	$8.460$ $[8.45,8.47]$	$0.005$	$8.461$ $[8.45,8.47]$	$0.006$	$8.460$ $[8.45,8.47]$	$0.006$
$2$	$7.420$ $[7.41,7.43]$	$0.022$	$8.132$ $[8.12,8.14]$	$0.007$	$8.557$ $[8.55,8.57]$	$0.005$
$3$	$6.636$ $[6.62,6.65]$	$0.068$	$7.841$ $[7.83,7.85]$	$0.007$	$8.554$ $[8.54,8.56]$	$0.005$
$4$	$6.116$ $[6.10,6.13]$	$0.158$	$7.583$ $[7.57,7.60]$	$0.011$	$8.539$ $[8.53,8.55]$	$0.006$
$6$	$5.652$ $[5.63,5.67]$	$0.341$	$7.147$ $[7.13,7.17]$	$0.007$	$8.512$ $[8.50,8.53]$	$0.005$

Appendix

Table 5: Generation quality on the released MDLM checkpoint at every step count of the sweep and the extended grid: mean per-sequence NLL under GPT-2-large (nats per token,

95\%

bootstrap intervals over sequences) and MAUVE against the held-out reference, for the three decoding orders at matched per-step commit sizes.

checkpoint	$\hat{\Lambda}^{*}$ (all counts)	$\hat{\Lambda}^{*}$ ( $\Lambda\geq 2$ )	loss minimum at
$1$ B	$2.419$	$3.16$	$\Lambda=6$
$2$ B	$2.528$	$3.42$	$\Lambda=6$
$4$ B	$2.621$ $[2.608,2.635]$	$3.61$	$\Lambda=7$

Appendix

Table 6: Loop-count saturation at the three checkpoints: the fitted saturation point on the full grid, the refit excluding

\Lambda=1

, and the loop count at which held-out loss is lowest.

Masked Diffusion Models

Masked diffusion models generate sequences by repeatedly predicting and revising masked tokens, using bidirectional context for parallel sampling.

See all 186 papers

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
Oct 8, 2026Jingda Wu, Changxiao CaiMasked Diffusion Models Diffusion Sampling
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA.
Improving Sampling for Masked Diffusion Models via Information Gain
Feb 20, 2026Kaisen Yang, Jayden Teoh, Kaicheng Yang +2Diffusion Sampling Diffusion Language Model Decoding
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2
DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking
Mar 2, 2026Gilad Turok, Chris De Sa, Volodymyr KuleshovMasked Diffusion Models Diffusion Language Models
Department of Computer Science, Cornell University, New York, New York, USA.

Oct 8, 2026J/K move · Enter open · S save

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure

Jingda Wu, Changxiao Cai

Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA.

Masked Diffusion Models Diffusion Sampling Diffusion Transformer Statistical Learning Theory

Open paper AskPDF arXiv

Masked diffusion models (MDMs) generate discrete data by iteratively unmasking multiple tokens in parallel using a learned mask predictor. Their output carries two errors: parallelization error due to parallel sampling, and estimation error from the mask predictor learned with finite samples. Existing theory treats these errors separately, and the available statistical bounds typically scale with the number

V^L

of length-

L

sequences over a vocabulary of size

V

, suffering from the curse of dimensionality. In this work, we establish end-to-end statistical guarantees showing that transformer-based MDMs learn structured distributions with a sample complexity governed by their latent structure. Specifically, we model this structure through mixtures of hidden Markov trees, which capture latent dependencies within components and heterogeneity across components. When the data distribution is well approximated by such a mixture, a transformer-based MDM with parallel unmasking schedules achieves

\varepsilon

-accurate sampling in KL divergence with

\widetilde{O}(LK_{\mathrm{par}}/\varepsilon)

training samples, where

K_{\mathrm{par}}

is polynomial in

L

V

, the number of hidden states, and the number of mixture components. The transformer network and sampling procedure require no knowledge of the tree topologies or component parameters, thereby adapting to the unknown latent structure automatically. Overall, our results provide theoretical justification for the empirical effectiveness of MDMs on real-world data.

Abstract

Masked Diffusion Models

Explore similar work

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure

Improving Sampling for Masked Diffusion Models via Information Gain

DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking