Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

2026-05-31

AI Summary

Existing masked diffusion language models struggle to correct generation errors due to standard samplers freezing already revealed tokens. This work proposes D3IM, a parameter-free sampler that enables direct self-correction of visible tokens and identifies a model bias toward retaining its own erroneous outputs. To address this, the authors introduce SCOPE, a lightweight post-training strategy that simulates the sampling process to co-optimize the sampler and the model. Evaluated on LLaDA-8B with 64-step denoising, the approach achieves significant improvements: 68.3% on GSM8K, 23.6% on MATH-500, 29.3% on HumanEval, and 30.8% on MBPP, with performance consistently increasing as the number of denoising steps grows.

Abstract

Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.

Problem

• masked diffusion language models
• token revision
• preservation bias
• sampler design
• self-correction

Innovation

• masked diffusion language models
• self-correction
• sampler-matched training
• preservation bias
• D3IM