Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
Title: Prioritize Revision Over Freezing: Parameter-Free Sampling for Self-Correcting Masked Diffusion Language Models
Abstract
While masked diffusion language models (MDLMs) are designed to re-evaluate every token position during each denoising step, conventional sampling strategies typically lock in tokens as soon as they are revealed, thereby neglecting this inherent capacity for revision. Current methods attempt to mitigate this issue by either incorporating heuristic or learned revision mechanisms or by remasking tokens to [MASK] prior to re-prediction. However, a principled sampling approach that enables direct revision of visible tokens without relying on auxiliary modules remains largely uninvestigated. To address this gap, we propose D3IM, a parameter-free sampler based on a corrector-style reverse update. D3IM facilitates direct visible-to-visible token revision without the need for additional modules or extra passes.
Furthermore, our investigation uncovers a model-side challenge we term "preservation bias," where models exhibit a tendency to replicate their previously incorrect committed tokens rather than correcting them. We counteract this bias through SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training technique that mimics the D3IM sampling dynamics. When applied to LLaDA-8B with 64 denoising steps, the combination of SCOPE and D3IM yields significant performance improvements over the original LLaDA-8B using standard unmasking. Specifically, we observe gains of +13.0 on GSM8K (reaching 68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%). Notably, these performance benefits become more pronounced as the number of denoising steps increases, particularly in math and HumanEval benchmarks.
Source: arXiv Generated at: 2026-06-02 00:00:00 UTC





