# Fine-Tuning Masked Diffusion for Provable Self-Correction

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

## Key facts

- Cited by 39 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 40 of its 51 references that are in the map.
- First posted on arXiv 1 Oct 2025.

## Placed in

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

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [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)
- [CRITIC: Large Language Models Can Self-Correct with Tool-Interactive Critiquing](https://mlsanity.com/paper/2305.11738/) (2023)
- [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)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [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)
- [Training Language Models to Self-Correct via Reinforcement Learning](https://mlsanity.com/paper/2409.12917/) (2024)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Mercury: Ultra-Fast Language Models Based on Diffusion](https://mlsanity.com/paper/2506.17298/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling](https://mlsanity.com/paper/2409.02908/) (2024)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [A Reparameterized Discrete Diffusion Model for Text Generation](https://mlsanity.com/paper/2302.05737/) (2023)
- [Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://mlsanity.com/paper/2505.24857/) (2025)
- [OpenCoder: The Open Cookbook for Top-Tier Code Large Language Models](https://mlsanity.com/paper/2411.04905/) (2024)
- [Energy-Based Diffusion Language Models for Text Generation](https://mlsanity.com/paper/2410.21357/) (2024)

## Built on by (most cited first, 30 of 39)

- [ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs](https://mlsanity.com/paper/2510.04767/) (2025)
- [Language Modeling with Hyperspherical Flows](https://mlsanity.com/paper/2605.11125/) (2026)
- [Learn from Your Mistakes: Self-Correcting Masked Diffusion Models](https://mlsanity.com/paper/2602.11590/) (2026)
- [The Diffusion Duality, Chapter II: Psi-Samplers](https://mlsanity.com/paper/2602.21185/) (2026)
- [Corrective Diffusion Language Models](https://mlsanity.com/paper/2512.15596/) (2025)
- [CORE: Context-Robust Remasking for Diffusion Language Models](https://mlsanity.com/paper/2602.04096/) (2026)
- [Flow Reasoning Models: Turning Flows Into Efficient Recurrent Reasoners](https://mlsanity.com/paper/2606.29150/) (2026)
- [Guided Star-Shaped Masked Diffusion](https://mlsanity.com/paper/2510.08369/) (2025)
- [Provably adaptive sampling with uniform and remasking discrete diffusion models](https://mlsanity.com/paper/2608.23554/) (2026)
- [Locally Coherent Parallel Decoding in Diffusion Language Models](https://mlsanity.com/paper/2603.20216/) (2026)
- [Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models](https://mlsanity.com/paper/2604.01622/) (2026)
- [Training-Free Self-Correction for Multimodal Masked Diffusion Models](https://mlsanity.com/paper/2602.02927/) (2026)
- [Forward-Free Diffusion Language Models with BPTT-Free Looped Refinement](https://mlsanity.com/paper/2606.08357/) (2026)
- [A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding](https://mlsanity.com/paper/2606.13565/) (2026)
- [Don't Commit Alone: Joint Token Commitment in Diffusion Language Models](https://mlsanity.com/paper/2607.04469/) (2026)
- [BackPlay: Head-Only Look-Back Self-Correction for Diffusion Language Models](https://mlsanity.com/paper/2601.06428/) (2026)
- [Generalized Discrete Diffusion with Self-Correction](https://mlsanity.com/paper/2603.02230/) (2026)
- [Latent-DARM: Bridging Discrete Diffusion And Autoregressive Models For Reasoning](https://mlsanity.com/paper/2603.09184/) (2026)
- [Self-Aware Markov Models for Discrete Reasoning](https://mlsanity.com/paper/2603.16661/) (2026)
- [Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Language Models](https://mlsanity.com/paper/2604.18738/) (2026)
- [The Confidence Shortcut: A Reasoning Failure Mode of Masked Diffusion Models](https://mlsanity.com/paper/2605.29123/) (2026)
- [DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs](https://mlsanity.com/paper/2606.01024/) (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)
- [VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing](https://mlsanity.com/paper/2606.28301/) (2026)
- [Concurrent Image Understanding and Generation: Self-Correcting Coupled Markov Jump Processes](https://mlsanity.com/paper/2607.13188/) (2026)
- [CaRE Compute-aware Remasking Evaluation Protocol for Masked Diffusion Language Models](https://mlsanity.com/paper/2607.24763/) (2026)
- [From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models](https://mlsanity.com/paper/2607.26504/) (2026)
- [Reconsidering Positional Supervision in Masked Diffusion Language Model Training](https://mlsanity.com/paper/2601.22947/) (2026)
- [Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models](https://mlsanity.com/paper/2605.26436/) (2026)

## Abstract (arXiv)

A natural desideratum for generative models is self-correction--detecting and revising low-quality tokens at inference. While Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces, their capacity for self-correction remains poorly understood. Prior attempts to incorporate self-correction into MDMs either require overhauling MDM architectures/training or rely on imprecise proxies for token quality, limiting their applicability. Motivated by this, we introduce PRISM--Plug-in Remasking for Inference-time Self-correction of Masked Diffusions--a lightweight, model-agnostic approach that applies to any pretrained MDM. Theoretically, PRISM defines a self-correction loss that provably learns per-token quality scores, without RL or a verifier. These quality scores are computed in the same forward pass with MDM and used to detect low-quality tokens. Empirically, PRISM advances MDM inference across domains and scales: Sudoku; unconditional text (170M); and code with LLaDA (8B).