Source-linked AI summary
Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg
TL;DR
Standard masked diffusion decoding freezes tokens after committing them, leaving the models’ ability to re-predict visible positions unused. The paper pairs D3IM’s direct revision sampler with SCOPE post-training to address preservation bias; on LLaDA-8B, the pairing improves results across four benchmarks at 64 steps.
Problem
Method
D3IM re-predicts and reselects visible tokens by confidence, while SCOPE post-trains the model on self-generated error states created by that sampling process.
Source: §1 · p. 2, §3.1 · pp. 3–4Results
SCOPE+D3IM improves over original LLaDA-8B with standard unmasking across GSM8K, MATH-500, HumanEval, and MBPP at 64 steps, including +13.0 on GSM8K.
Source: §1 · p. 2Takeaways & Limitations
The results indicate that committed-token revision depends on reliable confidence ranking under self-generated context, which SCOPE aligns with D3IM’s clean-slate top-K rule.
Source: §7 · p. 8All results are on LLaDA-8B, so generalization to other masked diffusion language models and larger scales remains unverified.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Standard MDLM decoding freezes committed tokens despite re-predicting every position, so early mistakes can persist. The paper pairs D3IM, which enables direct revision, with SCOPE training to address preservation bias and improve benchmark performance.
1 Introduction
- Existing approaches either remask committed tokens before revising them or add heuristics, verification passes, or learned modules.A direct visible-to-visible sampler without those additions remains underexplored.Source: §1 · p. 1
- +13.0 on GSM8K, +4.8 on MATH-500, +15.3 on HumanEval, and +10.4 on MBPP: SCOPE+D3IM improves over original LLaDA-8B with standard unmasking at 64 steps.The gains increase with more denoising steps on math and HumanEval.Source: §1 · p. 2
- D3IM is a parameter-free corrector-style sampler that directly revises visible tokens without remasking them or adding auxiliary modules or passes.It re-predicts all positions and retains the top-K predictions by confidence.Source: §1 · pp. 1–2, Figure 1 · p. 2
- The paper identifies preservation bias: models trained only on standard masked-token objectives tend to reproduce their own wrong committed predictions instead of correcting them.SCOPE simulates D3IM sampling and trains recovery from those errors without architectural changes or complex multi-objective training.Source: §1 · p. 2
2 Preliminaries
Absorbing masked diffusion models treat visible tokens as fixed, making committed-token correction impossible under their standard reverse process. LLaDA-style training removes the formal constraint but can still produce behavioral absorption under self-generated context.
2 Preliminaries
- In absorbing masked diffusion models, once a coordinate becomes visible, its clean-token posterior collapses to the current token, preventing later replacement of an early mistake.This makes monotonic denoising stable but leaves no mechanism to correct committed tokens.Source: §2 · pp. 2–3
- LLaDA-style objectives supervise masked positions only, imposing no loss on already-visible tokens and leaving visible-position predictions unconstrained by the absorbing identity.Standard inference nevertheless freezes committed tokens.Source: §2 · p. 3
- LLaDA-8B-Instruct preserves its own wrong visible token 70.1% of the time and recovers the ground-truth token only 9.2% in the wrong-commit stress test.The paper calls the learned tendency to preserve wrong predictions under self-generated context preservation bias.Source: §2 · p. 3
- ReMDM introduces remasking for absorbing-state models, while other revision methods use confidence thresholds, verification, or targeted token-to-mask heuristics.The paper asks whether LLaDA-style models can instead re-predict every position and globally reselect by current confidence.Source: §2 · p. 3
- Self-correcting inference requires both a sampler that re-evaluates committed tokens and a model that revokes wrong visible predictions under self-generated context.D3IM addresses sampler-level absorption; SCOPE addresses behavioral absorption.Source: §2 · p. 3
3 Method
The method pairs D3IM’s clean-slate confidence-based sampler with SCOPE training on self-generated error states. This aligns the model’s confidence ranking with D3IM’s revision process.
3.1 D3IM: Clean-Slate Remasking
- D3IM re-predicts all active positions at each denoising step, retains the top-K current predictions by confidence, and resets the rest to [MASK].A previously committed token can survive, be demoted to [MASK], or be overwritten by a new prediction.Source: §3.1 · p. 3
- D3IM sets the corrector update’s preservation weight ct to zero, enabling visible-to-visible revision without routing through [MASK].Its validity requires the noise-conservation constraint, and effectiveness depends on confidence reflecting correctness under self-generated context.Source: §3.1 · pp. 3–4
3.2 SCOPE: Training on Self-Generated Error States
- SCOPE constructs training inputs by committing the top-ρ fraction of sampled predictions by confidence, leaving other positions masked and exposing the model to high-confidence errors.The constructed context also includes ground-truth visible tokens.Source: §3.2 · p. 4
- SCOPE supervises recovery to ground truth on committed self-predictions and remaining masks, while excluding already-correct ground-truth context positions from the loss.The procedure includes temperature sampling and, with probability 0.5, substitutes standard MDLM training for the self-conditioning branch.Source: §3.2 · p. 4
3.3 Why the Pairing Matters
- SCOPE changes model behavior under self-generated contexts so D3IM’s top-K confidence ranking becomes reliable; its benefit is sampler-specific and does not automatically transfer to other revokable samplers.D3IM reselects every visible token at each step, exposing the original model’s preservation bias.Source: §3.3 · p. 4
4 Experiments
Experiments show that pairing SCOPE training with D3IM improves LLaDA-8B benchmark performance and makes its committed-token revisions more reliable. Analyses show that these gains depend on sampler alignment, vary with task and step budget, and reflect contributions from both SCOPE design choices and D3IM revision channels.
4.1 Setup
- Experiments fine-tuned LLaDA-8B-Instruct with LoRA for 500 FineWeb-Edu steps and evaluated fixed benchmark prompts and protocols at 64, 128, or 256 D3IM steps.Generation lengths, decoding budgets, and evaluation protocols were fixed across benchmark rows.Source: §4.1 · p. 4
4.2 Mechanism Diagnostics
- SCOPE reduces wrong-token preservation from 70.1% to 17.0% and raises ground-truth recovery from 9.2% to 50.2% among wrong self-commits.It also reduces masked-token Expected Calibration Error relative to Original across tested mask rates.Source: §4.2 · p. 5, Figure 2 · p. 6
4.3 Main Results: D3IM Requires Training-Inference Alignment
- Without SCOPE, D3IM harms all four benchmarks at every step count; SCOPE under standard unmasking is comparable on math but weaker on code.The failure is attributed to unreliable confidence ranking under self-generated context, not an insufficient step budget.Source: §4.3 · p. 5
- At 64 steps, SCOPE+D3IM raises GSM8K from 55.3% to 68.3%, MATH-500 from 18.8% to 23.6%, HumanEval from 14.0% to 29.3%, and MBPP from 20.4% to 30.8%.On math and HumanEval, gains grow with more steps; MBPP reaches 36.6% at 256 steps but remains below Original standard unmasking (39.2%).Source: §4.3 · p. 5, Table 1 · p. 6
4.4 Trajectory Analysis
- 274 corrections on one Original GSM8K example oscillate between content and whitespace, while SCOPE corrections target semantic errors and typically revise positions once.The reported SCOPE examples include unit, digit, and operator errors.Source: §4.4 · pp. 5–6, Table 2 · p. 7
4.5 Ablations and Additional Analysis
- Removing focused loss costs 10.8 and confidence-based commit costs 3.2; MDM-only continued training reaches 46.2%, below Original+D3IM’s 48.0%.These are leave-one-out drops relative to full SCOPE; MDM-only is a matched-budget baseline, not a strict leave-one-out variant.Source: §4.5 · p. 6, Table 3 · p. 7
- Disabling either D3IM correction channel reduces SCOPE performance, while disabling token-to-token revision raises Original GSM8K from 47.0% to 63.0%.With SCOPE, full D3IM reaches 70.5% on GSM8K-200, making that channel productive.Source: §4.5 · p. 6, Table 4 · p. 7
- SCOPE gains are sampler-specific: it adds 6–24 pp to D3IM, hurts ReMDM by 1–4 pp, and leaves Tolerator essentially neutral.On GSM8K, SCOPE+D3IM scores 68.3%, exceeding Tolerator’s 65.4%.Source: Table 5 · p. 7, §4.5 · p. 7
5 Related Work
Prior work trains on self-generated predictions or revises committed tokens at inference, but uses different interfaces and revision rules. SCOPE aligns confidence ranking specifically with D3IM’s clean-slate top-K rule.
Training on self-generated states
- SCOPE aligns confidence ranking with D3IM’s clean-slate top-K rule, unlike DMax’s soft embedding-space decoding, ProSeCo’s correction passes, or LLaDA2.1’s thresholded editing.These methods share training on self-generated predictions but differ in interface and goal.Source: §5 · p. 7
Inference-time revokable decoding
- D3IM reconstructs intermediate states through global top-K ranking, unlike methods using remasking schedules, local detection rules, or verification passes.Related methods include ReMDM [Wang et al., 2025], WINO [Hong et al., 2026], Tolerator [Tian et al., 2025], CoRe [Zhai et al., 2026], and T2M [Yao, 2026].Source: §5 · p. 8
Confidence, calibration, and token-quality signals
- SCOPE’s confidence improvements are sampler-specific and do not automatically transfer to arbitrary revokable samplers.SCOPE aligns confidence with D3IM’s clean-slate top-K competition.Source: §5 · p. 8
6 Discussion
D3IM revises visible tokens through direct replacement or demotion to [MASK], while bidirectional context can reinforce wrong committed tokens. The discussion connects revision to a possible way to interrupt that error propagation.
Two revision pathways
- A visible token can be overwritten by a higher-confidence prediction or demoted to [MASK] when it loses the top-K competition.Direct replacement preserves a signal that a similar token is needed; demotion signals that the model lacks a reliable candidate and should re-predict from updated context.Source: §6 · p. 8
Exposure bias in bidirectional models
- In bidirectional MDLMs, wrong committed tokens remain visible to all positions and may gain confidence as surrounding tokens accommodate the error.The paper contrasts this with autoregressive models, where a committed token’s confidence is fixed once generated.Source: §6 · p. 8
- SCOPE+D3IM scores 68.0% versus 66.0% in the reasoning-first answer-first stress test, consistent with revision mitigating error propagation under sparse context.The authors describe this result as consistent with their hypothesis.Source: §6 · p. 8
7 Conclusion
The conclusion identifies confidence reliability under self-generated context as the bottleneck for committed-token revision and says SCOPE aligns model behavior with D3IM. It also outlines possible directions for step-level supervision and reinforcement learning.
7 Conclusion
- Committed-token revision is bottlenecked by confidence-ranking reliability under self-generated context, rather than by the ability to revisit tokens.Without SCOPE, D3IM’s revision is harmful; SCOPE post-training yields gains specific to D3IM’s token-to-token pathway that scale with the step budget.Source: §7 · p. 8
- D3IM’s revision trajectory could provide step-level supervision analogous to process reward models.The authors identify this as one open direction.Source: §7 · p. 8
- D3IM’s token-to-token and token-to-mask actions at every step may offer a higher reinforcement-learning ceiling than rollouts under fixed sampling strategies.The authors frame this as a possibility for future RL approaches.Source: §7 · p. 8
Limitations
The paper identifies boundaries in the scope, cost, and evaluation of its results, as well as an unresolved theoretical question.
Limitations
- Token-level remasking handles local errors but cannot restructure reasoning when a problem requires a fundamentally different solution path.The passage says reasoning-level strategies such as best-of-N or diffusion-based reasoning scaling are needed in that case.Source: Limitations · p. 9
- All results are on LLaDA-8B, so generalization to other MDLMs and larger scales remains unverified.Source: Limitations · p. 9
- SCOPE training doubles per-step cost with two forward passes, although its total budget is 500 LoRA steps and about 20 minutes on one A100.Source: Limitations · p. 9
- Hyperparameters τ=1.5 and ρ=0.3 were fixed without systematic search, and reported numbers are single-run results without error bars.Source: Limitations · p. 9
- The paper lacks a theoretical account of why single-step self-conditioning generalizes to 64-step inference.Source: Limitations · p. 9
A Formal Details: Absorption, LLaDA, and D3IM
This appendix explains why standard absorbing samplers freeze revealed tokens, why LLaDA can in principle revise them, and how D3IM provides a direct revision path. It also describes inference details, including EOS handling.
A.1 Standard masked diffusion models
- In absorbing masked diffusion, once a token is unmasked, its posterior collapses to the current token, so subsequent reverse steps cannot change it.The masked endpoint and schedule define the absorbing setup.Source: §A.1 · p. 11
A.2 LLaDA’s training objective
- LLaDA supervises only masked positions, so its objective does not enforce the absorbing identity at visible positions.This leaves open the possibility of different predictions at visible positions, though preservation bias limits correction in practice.Source: §A.2 · p. 11
A.3 Preservation bias despite non-absorbing predictions
- Preservation bias arises because the model has not seen its own erroneous predictions as context during training, and it tends to reproduce them.The bias is not caused by the objective enforcing token preservation.Source: §A.3 · p. 11
A.4 D3IM as a corrector sampler
- D3IM sets the corrector’s preservation weight to zero, re-evaluates each position, and retains the top-K predictions while resetting the rest to [MASK].This creates a direct visible-token-to-visible-token revision path, unlike samplers that route revision through [MASK].Source: §A.4 · pp. 11–12
B.1 D3IM Inference (semi-AR with last-step accept)
- Semi-autoregressive inference processes blocks left-to-right; earlier blocks are frozen, while D3IM omits the override that preserves committed tokens in the active block.With one block, the schedule reduces to the global form.Source: §B.1 · p. 12
- During non-final steps, EOS/EOT candidates are excluded from top-K survival by setting their confidence to −∞; the final step accepts the whole block.Without suppression, 20% of generations place EOS before position 128; more aggressive logit suppression can cause truncation.Source: §B.1 · pp. 12–13, §B.2 · p. 13
B.2 EOS / EOT Policy Ablation
- Confidence suppression gives the best accuracy among the tested EOS policies (GSM8K 70.5, HumanEval 46.0) while delaying EOS without removing it from the output distribution.The table covers SCOPE+D3IM at 64 steps on bounded task subsets; EOS timing and truncation are pooled over 350 generations.Source: §B.2 · p. 13, Table 6 · p. 13
B.3 SCOPE Training: Implementation Details
SCOPE training simulates prediction errors and commits selected predictions, while retaining standard masked-token training. The appendix reports temperature and commit-rate choices and their benchmark trade-offs.
B.3 SCOPE Training: Implementation Details
- Prompt positions remain unmasked, so both self-prediction and the training pass condition on the full prompt; only response positions enter the mask–commit–retrain cycle.Source: §B.3 · p. 13
- With probability 1 − pself-cond, SCOPE falls back to standard masked-token training to preserve ordinary mask-filling ability and avoid catastrophic forgetting.The implementation also tracks separate losses for correcting wrong commits and preserving correct ones.Source: §B.3 · p. 13
- SCOPE samples from a masked context, commits the top-ρ predictions by confidence, and focuses cross-entropy on committed and masked positions.The recipe uses Beta(2,2) mask rates and τ=1.5 to increase self-prediction errors toward those encountered by D3IM.Source: §B.3 · p. 13
C.1 Temperature
- Table 7 shows higher temperature raises standard-unmasking GSM8K accuracy from 42.2% to 55.2%, while the effect on D3IM is only +2.7 points.The authors use τ=1.5; τ=2.0 is marginally stronger.Source: Table 7 · p. 14, §C.1 · p. 14
B.3 SCOPE Training: Implementation Details
- On 50% mask-rate training, Lwrong decreases from approximately 6.3 to 1.5 and Lright from approximately 0.8 to 0.05 over 500 steps.Source: §B.3 · p. 14
C.2 Commit-Rate Schedule
- None of the three tested commit-rate schedules consistently improves across benchmarks, so the authors retain the fixed rate ρ=0.3.Each schedule trades a small gain on one task for a regression on another.Source: §C.2 · p. 14
C.3 D3IM Channel Ablation: Extended Analysis
The extended analysis compares D3IM with inference-only revision baselines and contrasts SCOPE’s input-side training changes with post-hoc calibration.
C.3 D3IM Channel Ablation: Extended Analysis
- On GSM8K-200 and MBPP-50, WINO is unchanged by SCOPE training (57.0% to 57.0%), ReMDM-conf gains 1 point, and ReMDM-cap is hurt.Tolerator reaches 70.0% on SCOPE, nearly matching full D3IM at 70.5%, but uses a more complex two-phase inference loop.Source: Table 8 · p. 14, §C.3 · p. 14
C.4 Why Not Standard Calibration?
- Post-hoc logit scaling and label smoothing do not change the self-generated contexts that challenge D3IM; SCOPE changes training inputs while retaining ground-truth targets.Source: §C.4 · p. 14
- SCOPE simulates one D3IM step despite 64-step inference, and the reported trajectory analysis suggests this is sufficient in the tested setting.The proposed explanation is that local token-context ranking generalizes across progressive inference stages.Source: §C.4 · p. 14
C.5 Comparison with XDLM Training
The paper finds that D3IM also benefits a model trained with XDLM, though less than SCOPE+D3IM, and tests output-order sensitivity as an exposure-bias stress case.
C.5 Comparison with XDLM Training
- XDLM’s native decoder replaces a token when the new prediction exceeds its confidence, whereas D3IM uses global clean-slate survival; the experiment pairs the XDLM checkpoint with both samplers.Source: §C.5 · pp. 14–15
- On GSM8K, XDLM+D3IM improves over XDLM’s native sampler by +9.5 points (50.8% to 60.3%) but remains below SCOPE+D3IM at 68.3%.The comparison uses the full 1,319-example GSM8K set at 64 steps.Source: Table 9 · p. 15, §C.5 · p. 15
C.6 Exposure Bias Stress Test
- In the GSM8K-200 stress test, SCOPE+D3IM has only a 2-point gap between answer-first and reasoning-first formats, while Original+D3IM drops 9 points on reasoning-first.The Original+D3IM reasoning-first score is 46.0%, compared with 55.0% answer-first.Source: Table 10 · p. 15, §C.6 · p. 15
D Additional Related Work: Remasking Dynamics
- The paper frames its trajectory analysis as addressing a gap in prior work: examining semantic correction versus superficial reshuffling and how revision frequency relates to calibration.Source: D Additional Related Work: Remasking… · p. 15