[Submitted on 17 Dec 2025 (v1), last revised 5 Nov 2026 (this version, v3)]
Abstract:While Diffusion Language Models (DLMs) are theoretically well-suited for iterative refinement due to their non-causal structure, off-the-shelf masked diffusion language models (MDLMs) often fail to reliably revise incorrect tokens in practice. The key challenge is that the model cannot tell correct from erroneous visible tokens: standard MDLM training supervises only masked positions, undermining confidence-guided refinement. We study corrective behavior in DLMs, the ability to assign lower confidence to incorrect tokens and iteratively refine them while preserving correct content. Under matched compute and data, we show that conventional masked-diffusion fine-tuning induces much less of this capability than a correction-oriented post-training principle that explicitly supervises visible corrupted tokens, which enables discriminative confidence and targeted refinement. We further find that also supervising clean visible tokens, as in BERT- or GIDD-style training, stops remask-based refinement from improving with additional steps. To evaluate it, we introduce the Code Revision Benchmark, a controllable, executable benchmark for error localization and in-place correction. The effect transfers from a 0.5B model to LLaDA-8B and persists, with a smaller margin, when corruptions are drawn from the base model’s own most plausible predictions. The same error-aware confidence also makes parallel decoding more effective. In HumanEval(+) code completion at 0.5B, mathematical reasoning at 16B, and high-entropy parallel decoding with remasking (ParallelBench), models trained with our objective outperform standard masked-diffusion fine-tuning under matched compute, and outperform the base model on HumanEval(+), MATH, and the high-entropy ParallelBench average; at 8B, neither objective improves from-scratch generation over the base model. Code and data: this https URL
Submission history
From: Shuibai Zhang [view email]
[v1] Wed, 17 Dec 2025 17:04:38 UTC (243 KB)
[v2] Thu, 29 Jan 2026 17:52:44 UTC (357 KB)
[v3] Thu, 5 Nov 2026 17:13:47 UTC (661 KB)
