RepairMDM: Verifier-Localized Repair Training for Masked Diffusion Language Models
Aryan Keluskar ⋅ Kyung S Lee
Abstract
Masked diffusion language models (MDMs) decode in parallel and can, in principle, revise earlier tokens. Yet a token committed with high confidence that violates a global constraint is invisible to the model's own confidence and cannot be resampled away. This is the multimodality problem of non-autoregressive decoding meeting overconfidence, a regime we call a confident wrong-commit. We introduce RepairMDM, which post-trains the MDM unmasking posterior into a verifier-localized repair conditional: a deterministic verifier builds the training data offline (masking a violated constraint's variables and supervising a minimal-edit fix), so that at inference a cheap constraint check localizes the error and the trained fill repairs it. On 3-SAT, RepairMDM repairs over 80\% of held-out errors, outperforming best-of-$N$ by 3 times while a PRISM-style quality head and verifier search sit near zero. On Sudoku, it repairs $100\%$ at $16$ blanks, whereas best-of-$N$ and the PRISM head repair none. We find that the advantage is regime-specific, since it needs confident errors and a precise verifier, vanishing under a coarse pass/fail code verifier. More broadly, this exposes a limit of intrinsic self-correction: localizing a confident error is not the same as repairing it. Finding it needs an external oracle, and resolving it needs a retrained generator.
Chat is not available.
Successful Page Load