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

Oct 8, 2026

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

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.

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

Ask PDF

Feb 20, 2026

Improving Sampling for Masked Diffusion Models via Information Gain

Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the Info-Gain Sampler, a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.

Kaisen Yang, Jayden Teoh, Kaicheng Yang +2

Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2

Diffusion Sampling Diffusion Language Model Decoding Masked Diffusion Models

Ask PDF

Mar 2, 2026

DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking

Masked diffusion models (MDMs) generate text by iteratively selecting positions to unmask and then predicting tokens at those positions. Yet MDMs lack proper likelihood evaluation: the evidence lower bound (ELBO) is not only a loose bound on log-likelihood, but, as we show, is also computed under the training distribution rather than the test-time distribution. We resolve this within our DUEL framework, which unifies leading MDM sampling strategies that employ

\textit{deterministic}

position selection. We prove that DUEL samplers admit

\textbf{exact likelihood computation under the test-time distribution}

-- giving MDMs

\textit{proper}

likelihood, and hence proper perplexity, for the first time. This proper perplexity is the natural analogue of autoregressive perplexity and lets us revisit key questions about MDMs.

\textbf{MDMs are substantially better than previously thought}

: the MDM-autoregressive perplexity gap shrinks by up to

32\%

on in-domain data and

82\%

on zero-shot benchmarks. DUEL enables the first principled comparison of fast,parallel samplers across compute budgets -- an analysis impossible with the ELBO and unreliable with generative perplexity -- identifying a strong default method. Finally, oracle search over position orderings reveals MDMs can far surpass autoregressive models -- achieving

36.47

vs.

52.11

perplexity on AG News -- demonstrating the ceiling of MDM performance has not yet been reached.

Gilad Turok, Chris De Sa, Volodymyr Kuleshov

Department of Computer Science, Cornell University, New York, New York, USA.

Masked Diffusion Models Diffusion Language Models Diffusion Sampling Diffusion Language Model Decoding

Ask PDF

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