Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Yair Schiff; Omer Belhasin; Roy Uziel; Guanghan Wang; Marianne Arriola; Gilad Turok; Ran Zilberstein; Michael Elad

cs.LGFeb 12, 2026

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Authors: Yair Schiff, Omer Belhasin, Roy Uziel, Guanghan Wang, Marianne Arriola, Gilad Turok, Ran Zilberstein, Michael Elad,

Organizations: Cornell · NVIDIA

Abstract

Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models, enabling parallel token generation while achieving competitive performance. Despite these advantages, MDMs face a fundamental limitation: once tokens are unmasked, they remain fixed, leading to error accumulation and ultimately degrading sample quality. We address this by proposing a framework that trains a model to perform both unmasking and correction. By reusing outputs from the MDM denoising network as inputs for corrector training, we train a model to recover from potential mistakes. During generation we apply additional corrective refinement steps between unmasking ones in order to change decoded tokens and improve outputs. We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence, including already generated tokens. We conduct extensive experimental validation across multiple conditional and unconditional tasks, demonstrating that ProSeCo yields better quality-efficiency trade-offs (up to ~4x faster sampling while still outperforming the best accuracy baseline method) and enables further gains through inference-time compute scaling (up to ~1.2x improvement on benchmarks when allowing for slower, high-NFE decoding).

Figures & tables

Figure 1: ( Left ) Training

\mathpzc{ProSeCo}

: The original process trains for generation via unmasking. For every timestep in the masking process, we also train for self-correction: undoing corruptions that arise from sampling from unmasking predictions. (Right) Using our method for SFT and Max correction sampling with LLaDA [ 34 ] significantly outperforms vanilla masked diffusion SFT.

Panel of Figure 2 — Figure 2: Illustrative example demonstrating the benefits of self-correction. (Left) For standard MDM, during parallel unmasking, errors occur. Mistakes accumulate and lead to sample collapse. (Middle) For $\mathpzc{ProSeCo}$ , a short correction loop steers generation ‘back on track.’ (Right) $\mathpzc{ProSeCo}$ can thus produce a final high quality output even with parallel generation.

			Code		Math
		Corrector Sampling	HumanEval	MBPP	GSM8K	MATH
Off-the-Shelf 8B Models
Prior Work	Llama3.1-Instruct [ 18 ]	✗	63.41	70.90	81.05	47.38
	LLaDA-Base [ 34 ]	✗	34.15	37.20	46.78	17.04
	LLaDA-Instruct [ 34 ]	✗	45.73	47.09	83.40	43.14
	+ ReMDM [ 56 ]	✓	43.90	45.50	83.93	43.76

Table 1: Pass@1 on code and math benchmarks. ‘Corrector Sampling’ ✗ / ✓ symbol indicates whether a correction algorithm was applied. † indicates values taken from Huang et al. [22] . Other "Prior work" values obtained via evaluation with open-sourced weights. Best value is bolded .

Figure 3: Quality-efficiency trade-offs for

\mathpzc{ProSeCo}

. Standard MDMs (Baseline; gray dot ) attain best performance when decoding a single token every step.

\mathpzc{ProSeCo}

can vary number of corrector steps and attain comparable performance more efficiently ( Ours: Fast; green star ), it achieve even better quality for modest increase in compute budget ( Ours: Balanced; orange star ), or it maximize quality by further scaling inference-time compute ( Ours: Max; blue star ).

Figure 4: Pareto frontier of parallel decoding and quality. Applying a modest number of corrector steps, allows

\mathpzc{ProSeCo}

models to recover from paralell decoding errors and extend this frontier.

Figure 5:

\mathpzc{ProSeCo}

better navigates the novelty-property maximization Pareto frontier. Novel samples (

x

-axis) vs. mean property value of novel samples (

y

-axis) for classifier-free guidance [ 48 ] , with varying unmasking steps

T

(line style) and guidance strength

\gamma

(marker size). (Left) Maximizing the ring count property. (Right) Maximizing the drug likeness (QED) property.

	MAUVE ( $\uparrow$ )				Gen. PPL ( $\downarrow$ )				Entropy ( $\uparrow$ )
Data	1.00				14.8				5.44
AR ( $T=1024$ ) $	0.76				12.1				5.22
$T=$	$128$	$256$	$512$	$1024$	$128$	$256$	$512$	$1024$	$128$	$256$	$512$	$1024$
MDLM $ [ 45 ]	0.015	0.023	0.031	0.042	61.5	55.8	53.0	51.3	5.52	5.49	5.48	5.46
ReMDM $ [ 56 ]	0.057	0.216	0.350	0.403	42.5	30.5	21.1	28.6	5.43	5.34	5.21	5.38
PRISM † [ 24 ]	0.118	0.294	0.423	0.527	21.5	18.0	16.4	15.3	5.18	5.15	5.12	5.10

Table 2: Unconditional generation sample quality for models trained on OpenWebText. † Values reported in Kim et al. [24] . $ Values reported in Wang et al. [56] . ∗ Values obtained with open-sourced weights; sequence length reduced to 512 to match the model trained in von Rütte et al. [55] .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

		HumanEval	MBPP	GSM8K	MATH
Baseline (Vanilla SFT)
	Accuracy (%)	58.54	56.88	88.86	46.60
	Average NFEs	247.8	290.3	214.2	259.4
	Generation length	1024	1024	1024	1024
	Block length	32	32	32	32
	Demasking Tokens/Step	1	1	1	1

Appendix

Table 3: Sampling hyperparameters for results attained with

\mathpzc{ProSeCo}

and reported in Figure 3 .

$T=$	Unmasking steps	Maximum Self Correction Steps	Avg. NFEs
128	64	64	123.9
256	128	128	212.6
512	512	512	601.3
1024	1024	1024	1114.2

Appendix

Table 4: Parameters used for GIDD [ 55 ] unmasking and self-correction budgets corresponding to the results in Table 2 .

Model / configuration	HumanEval	MBPP	GSM8K	MATH
Vanilla SFT	58.54	56.88	88.86	46.60
+ ReMDM (paper-recommended)	56.10	50.00	88.48	46.22
+ ReMDM (additional sweep)	62.20	54.50	88.70	46.50
$\mathpzc{ProSeCo}$ SFT + Max Sampling	72.56	69.31	92.19	55.06

Appendix

Table 5: Additional ReMDM configurations applied to our vanilla LLaDA-Base 8B SFT checkpoint. The paper-recommended configuration corresponds to the ReMDM row under our SFT baselines in Table 1 . The additional sweep uses random remasking and multiple ReMDM loops per block. Values are Pass@1; best values are bolded .

Model	HumanEval	MBPP	GSM8K	MATH
Vanilla SFT	58.54	56.88	88.86	46.60
Vanilla SFT ( $2\times$ steps)	62.80	56.35	89.61	48.08
$\mathpzc{ProSeCo}$ SFT + Max Sampling	72.56	69.31	92.19	55.06

Appendix

Table 6: Pass@1 comparison with extended vanilla SFT. The

2\times

baseline receives twice the optimization steps of the original vanilla SFT baseline. Best values are bolded .

		HumanEval	MBPP	GSM8K	MATH
Baseline (Vanilla SFT)
	Accuracy	58.54	56.88	88.86	46.60
	Avg. NFEs	247.8	290.3	214.2	259.4
	Tokens/Second	123.14	131.04	148.92	153.15
$\mathpzc{ProSeCo}$ Max
	Accuracy (+ gain)	72.56 ( +14.0 )	69.31 ( +12.4 )	92.19 ( +3.3 )	55.06 ( +8.5 )

Appendix

Table 7: Quality vs. Efficiency Metrics. We report Accuracy (%), Average Number of Function Evaluations (NFEs), and Throughput (Tokens/Second). Best accuracy gain and efficiency results are bolded . Improvements over the baseline are highlighted in green .

Figure 6: Pareto frontier analysis with throughput (tokens per second, TPS) reported.

Figure 7: Ablation: Performance across various configurations of corrector steps. Frequency

\omega\in\{1,2,4,8\}

(denoted by color) and

S\in\{1,2,4,8\}

(denoted by marker size).

Training + Sampling	HumanEval	MBPP	GSM8K	MATH
Vanilla SFT + Standard MDM	58.54	56.88	88.86	46.60
Vanilla SFT + $\mathpzc{ProSeCo}$ Balanced	1.83	1.85	60.42	27.06
$\mathpzc{ProSeCo}$ SFT + $\mathpzc{ProSeCo}$ Balanced	72.56	69.31	92.19	55.06

Appendix

Table 8: Ablation: Applying self-correction sampling to models trained with standard MDM loss leads to worse performance, underscoring the importance of the

\mathpzc{ProSeCo}

framework. Pass@1 accuracy (%) performance is reported.

Figure 8: Ablation: Selecting self-correction loss weight

\lambda_{t}

. We retrain the discrete classifier model for maximizing ring count with various

\lambda_{t}

values.

$T=$	MAUVE ( $\uparrow$ )	Gen. PPL ( $\downarrow$ )	Entropy ( $\uparrow$ )
128	0.311 $\pm$ 0.006	23.045 $\pm$ 0.031	5.452 $\pm$ 0.002
256	0.572 $\pm$ 0.006	16.52 $\pm$ 0.014	5.376 $\pm$ 0.005
512	0.622 $\pm$ 0.005	13.196 $\pm$ 0.031	5.293 $\pm$ 0.001
1024	0.592 $\pm$ 0.006	10.92 $\pm$ 0.014	5.221 $\pm$ 0.005

Appendix

Table 9: Mean

\pm

standard error for OWT unconditional generation metrics for

\mathpzc{ProSeCo}

Dataset	Licence
GSM8K [ 12 ]	MIT
HumanEval [ 10 ]	MIT
Llama-Nemotron: Efficient Reasoning Models [ 5 ]	CC BY 4.0
MBPP [ 3 ]	MIT
MinveraMath [ 20 ]	MIT
OpenWebText [ 17 ]	Creative Commons CC0 license (“no rights reserved”)

Appendix

Table 10: Datasets and corresponding licenses.

Library	License
HuggingFace [ 61 ]	Apache 2.0
Hydra [ 63 ]	MIT
Fast-DLLM [ 62 ]	Apache 2.0
Language Model Evaluation Harness [ 15 ]	MIT
Matplotlib [ 23 ]	Matplotib license
Mauve [ 41 ]	GNU General Public License, Version 3

Appendix

Table 11: Software and corresponding licenses.

Panel of Figure 9 — Figure 9: Example of $\mathpzc{ProSeCo}$ ’s generated response on the HumanEval dataset using the max accuracy configuration (see Table 3 ).

Panel of Figure 10 — Figure 10: Example of $\mathpzc{ProSeCo}$ ’s generated response on the GSM8K dataset using the max accuracy configuration (see Table 3 ).

Panel of Figure 11 — Figure 11: Generated sample from $\mathpzc{ProSeCo}$ trained on OWT, with a total budget of $T=256$ : $64$ unmasking steps, corrector frequency $\omega=1$ and $S=3$ steps per loop.

Masked Diffusion Models

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

See all 183 papers

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
May 31, 2026Longxuan Yu, Shaorong Zhang, Yu Fu +3Masked Diffusion Models Diffusion Model Fine-Tuning
University of California, Riverside · Microsoft
Simple Self-Conditioning Adaptation for Masked Diffusion Models
Apr 28, 2026Michael Cardei, Huu Binh Ta, Ferdinando FiorettoMasked Diffusion Models Diffusion Model Fine-Tuning
University of Virginia
Random Remasking Scales Masked Diffusion Inference
Oct 22, 2026Zhaoyi Zhou, Frank Wu, Andrea ZanetteMasked Diffusion Models Diffusion Language Model Inference
Carnegie Mellon University

May 31, 2026J/K move · Enter open · S save

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu+2

University of California, Riverside · Microsoft

Masked Diffusion Models Diffusion Model Fine-Tuning Diffusion Language Model Decoding Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.

Abstract

Masked Diffusion Models

Explore similar work

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Simple Self-Conditioning Adaptation for Masked Diffusion Models

Random Remasking Scales Masked Diffusion Inference