Research map / Image generation / Fast diffusion sampling

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Masked diffusion language model decoding · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes ProSeCo, a masked diffusion training and sampling method that jointly trains unmasking and self-correction by reusing the denoiser's argmax outputs as corrector inputs, improving quality-speed trade-offs on math, code, molecule and text generation.

Problem

Masked diffusion models cannot revise tokens once unmasked, so errors from parallel decoding accumulate and degrade sample quality.

Approach

Tie corrector and denoiser weights and add a cross-entropy self-correction loss on denoiser argmax outputs as corrupted inputs, then interleave corrector refinement loops (which can change already-decoded tokens) with standard unmasking steps at inference.

Key results

• ProSeCo Max sampling reaches 72.56 HumanEval, 69.31 MBPP, 92.19 GSM8K and 55.06 MATH, vs 58.54/56.88/88.86/46.60 for vanilla SFT on the same LLaDA-Base 8B recipe.
• Balanced/Fast configurations yield roughly 2-4x speed-ups over LLaDA decoding without accuracy loss, per the paper.
• On OpenWebText with T=1024, ProSeCo reaches MAUVE 0.604 and gen. PPL 10.9, compared with ReMDM 0.403/28.6 and MDLM 0.042/51.3.

Surprising findings

• Applying corrector sampling to a vanilla MDM SFT model does not improve results (ReMDM-style correction on vanilla SFT slightly decreases HumanEval and MBPP), indicating standard MDMs never learn to revise unmasked tokens.
• The corrector loss weight reuses the MDM weight alpha_dot/(1-alpha), and biasing training toward heavier masking (min_t=0.1) helped SFT.

Limitations

• Training requires a second forward pass, adding training compute cost.
• Corrector and denoiser weights are tied, so the two roles are not disentangled.
• Some baseline numbers (ReMeDi) are taken from the original paper because the authors could not reproduce them.

Experimental setup

Datasets: Llama-Nemotron-Post-Training (modified) OpenWebText QM9

Benchmarks: HumanEval MBPP GSM8K MATH QM9 guidance (ring count, QED) OpenWebText unconditional generation

Metrics: pass@1 validity/uniqueness/novelty mean property value MAUVE generative perplexity entropy NFEs

Baselines: Llama3.1-Instruct 8B LLaDA-Base LLaDA-Instruct ReMDM LLaDA1.5 ReMeDi-Instruct Vanilla SFT AR UDLM MDLM PRISM GIDD

Builds on: masked diffusion models LLaDA MDLM ReMDM predictor-corrector sampling for discrete diffusion

Closest papers by meaning

• Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.93)
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.92)
• Remasking Discrete Diffusion Models with Inference-Time Scaling (NeurIPS 2025 · similarity 0.92)
• Re-evaluating Confidence Remasking in Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.91)
• Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions (ICML 2025 · similarity 0.91)
• CORE: Context-Robust Remasking for Diffusion Language Models (ICML 2026 · similarity 0.91)