Diffusion language models learn to take back a token: remasking confidence drops lifts four-task average by 0.85 to 1.44 points

Decoding in masked diffusion language models no longer treats every committed token as final; the model can go back and revise an earlier choice once fuller context exists. Existing samplers mainly decide when to commit a token and usually keep committed tokens fixed, so early mistakes made under sparse context keep conditioning later predictions. Across two backbones, LLaDA-8B and Dream-7B, three base samplers, top1, Fast-dLLM and EB-Sampler, and four zero-shot tasks, HumanEval, MBPP, MATH-500 and GSM8K, adding CoDR improves the four-task average accuracy of every evaluated model–sampler configuration.

Whether a committed token should be reopened is judged by the gap between its commit-time confidence and its confidence under denser context, not by its current confidence level. Prior remasking methods revisit committed positions but do not directly measure how support for the same token changes as decoding progresses. With Fast-dLLM on LLaDA-8B, triggering on confidence drift reaches 43.90% on HumanEval and 77.40% on GSM8K, above random remasking, bottom-m selection by lowest current confidence, and the top1_change trigger.

Drift need not cost one forward pass per committed position: splitting the committed set into k residue classes by stride and masking each class as a whole yields drift for every position in that class from a single pass. Computing each position's later confidence exactly requires as many forward passes as there are committed positions, an overhead too expensive for generation. The default is k=2, separating even and odd indices; at the same k, deterministic stride partitioning remasks fewer positions and scores higher than random partitioning, 43.90% versus 43.53% on HumanEval and 77.40% versus 76.40% on GSM8K.

Extra compute alone is not the source of the gain; spending it on positions flagged by drift is. Sweeping the Fast-dLLM threshold and the EB-Sampler bound shows that additional base-sampler evaluations do not improve the best HumanEval or GSM8K accuracy. On LLaDA-8B, fd+CoDR raises HumanEval from 42.68% to 43.90% and GSM8K from 76.00% to 77.80%; fd+CoDR reaches a 48.93% four-task average with 123 forward evaluations, above ReMDM's 48.44% average at its 512-forward-evaluation setting.

The demonstrated setting is two open-weight masked diffusion models, LLaDA-8B-Instruct and Dream-v0-Instruct-7B, evaluated with Pass@1 on four zero-shot benchmarks, HumanEval, MBPP, MATH-500 and GSM8K, at generation length 256, block length 32, k=2 and drift threshold 0; the authors also report LLaDA results on BBH and GPQA. The method requires a base sampler that exposes commit-time probabilities and supports masked refill.

The authors note that CoDR's effectiveness depends on the base sampler used for refill, that revised tokens remain bounded by the same model and sampler, and that refill occasionally introduces fresh local errors, especially in arithmetic-heavy cases; confidence drift is framed as a targeted revision signal rather than a complete error detector. At generation length 128 gains are smaller and MATH-500 regresses, so average improvement is not uniform per-task gain. KV caching currently accelerates only initial sampling, because remasking invalidates cached states for probing and refill.