# Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

arXiv: https://arxiv.org/abs/2602.11590 · HTML: https://mlsanity.com/paper/2602.11590/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 50 of its 71 references that are in the map.
- First posted on arXiv 12 Feb 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 30 of 50)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers](https://mlsanity.com/paper/2102.01454/) (2021)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Models](https://mlsanity.com/paper/2505.19223/) (2025)
- [Molecular Transformer - A Model for Uncertainty-Calibrated Chemical Reaction Prediction](https://mlsanity.com/paper/1811.02633/) (2018)
- [Continuous diffusion for categorical data](https://mlsanity.com/paper/2211.15089/) (2022)

## Built on by (most cited first, 13 of 13)

- [Discrete-WAM: Unified Discrete Vision-Action Token Editing for World-Policy Learning](https://mlsanity.com/paper/2606.05645/) (2026)
- [Self-Generated Error Training for Token Editing in Diffusion Language Models](https://mlsanity.com/paper/2606.17175/) (2026)
- [Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decoding](https://mlsanity.com/paper/2607.01775/) (2026)
- [Don't Commit Alone: Joint Token Commitment in Diffusion Language Models](https://mlsanity.com/paper/2607.04469/) (2026)
- [Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models](https://mlsanity.com/paper/2606.01026/) (2026)
- [Re-evaluating Confidence Remasking in Masked Diffusion Language Models](https://mlsanity.com/paper/2606.12232/) (2026)
- [From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models](https://mlsanity.com/paper/2607.26504/) (2026)
- [Hesitation-Aware On-Policy Distillation for Diffusion Language Models](https://mlsanity.com/paper/2609.33301/) (2026)
- [Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models](https://mlsanity.com/paper/2605.26436/) (2026)
- [A Ticket from Marginals to Joints: Coupled-Noise Distillation for One-Step Block Generation in Diffusion Language Models](https://mlsanity.com/paper/2609.06324/) (2026)
- [Less Uniform Discrete Diffusion is More Powerful and Scalable](https://mlsanity.com/paper/2609.35817/) (2026)
- [Blackboard Intelligence Can Surpass Autoregressive on Globally Constrained Problems](https://mlsanity.com/paper/2609.38806/) (2026)
- [Correction Training for Diffusion Models](https://mlsanity.com/paper/2610.14136/) (2026)

## Abstract (arXiv)

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 method~yields better quality-efficiency trade-offs (up to ~4x faster sampling) and enables inference-time compute scaling to further increase sample quality beyond standard MDMs (up to ~1.2x improvement on benchmarks).