Fine-Tuning Masked Diffusion for Provable Self-Correction
read the original abstract
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).
This paper has not been read by Pith yet.
Forward citations
Cited by 11 Pith papers
-
Flow Reasoning Models: Scaling Reasoning Through Iterative Self-Refinement
Flow models reach 99.2% Sudoku accuracy in 7 passes and 96.1% on out-of-distribution Sudoku-Extreme by selecting dynamically stable candidates and training with self-conditioning plus DPO to avoid failed outputs.
-
VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing
MDM-VGB augments masked diffusion with backtracking-style reward-guided remasking to achieve quadratic-complexity high-reward generation and sample editing, with proofs of noise robustness.
-
Inference-Time Scaling in Diffusion Models through Iterative Partial Refinement
IPR improves valid solution rates on MNIST Sudoku from 55.8% to 75.0% by iteratively refining partial regions in sequential diffusion models without external verifiers or reward models.
-
Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Large Language Models
Token-to-Mask remasking improves self-correction in diffusion LLMs by resetting erroneous commitments to masks rather than overwriting them, yielding +13.33 points on AIME 2025 and +8.56 on CMATH.
-
Locally Coherent Parallel Decoding in Diffusion Language Models
CoDiLA adds a compact auxiliary AR model on diffusion latents to enforce local sequential validity during parallel token sampling in discrete diffusion language models.
-
Discrete Stochastic Localization for Non-autoregressive Generation
Discrete Stochastic Localization lets a single trained network support an entire family of per-token SNR paths for discrete sequence generation, with masked diffusion as a special case, and improves MAUVE scores when ...
-
A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding
A2D2 derives the Radon-Nikodym derivative for joint insertion-unmasking paths in discrete diffusion to enable reward-tilted fine-tuning and introduces the Adaptive Joint Decoding loss.
-
NAVIRA: Decoupled Stochastic Remasking for Masked Diffusion Language Models
NAVIRA decouples quality scoring from regeneration via stochastic remasking in masked diffusion LMs, improving fluency and LLM-judge scores on a 170M model.
-
Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
Presents D3IM sampler and SCOPE post-training that enable visible-token revision in masked diffusion LMs, reporting double-digit gains on GSM8K and HumanEval for LLaDA-8B.
-
DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs
Adapting LLaDA-8B-Instruct via Discrete Stochastic Localization with continuous per-token Gaussian noise yields continuous denoising that achieves top ROUGE-1 on zero-shot summarization at low step budgets and adds se...
-
Re-evaluating Confidence Remasking in Masked Diffusion Language Models
Re-evaluation finds post-hoc remasking (WINO) yields little-to-no gain over confidence unmasking in standard dLLM settings and can worsen diversity collapse under stochastic decoding.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.