Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Researchers have developed Progressive Self-Correction (ProSeCo), a framework that enables masked diffusion models to correct previously generated tokens rather than leaving them fixed. The method reuses the model's denoising network outputs to train a corrector, then applies additional refinement steps during generation to fix errors. ProSeCo achieves up to 3x faster sampling and 1.3x improvement in sample quality over standard masked diffusion models across multiple tasks.

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 ~2-3x faster sampling and enables inference-time compute scaling to further increase sample quality beyond standard MDMs up to ~1.3x improvement on benchmarks .