Source-linked AI summary
CoDR: Training-Free Confidence-Drift Remasking for Diffusion Language Models
Yue Wu, Qinghe Zhang, Yu Zhang, Jian Huang
TL;DR
Masked diffusion samplers rarely revisit committed tokens, although their support can fall as context grows. CoDR uses k-partition probing to identify and selectively refill drifted commitments. Across two backbones, three samplers, and math and code tasks, it improves the accuracy–compute trade-off.
Problem
Existing samplers rarely revisit committed tokens, and confidence drift captures when a token loses support as context grows.
Source: §1 · pp. 1–2Method
CoDR compares commit-time and later-context confidence, estimates drift with k-partition probing, and selectively remasks and refills affected tokens.
Source: §1 · p. 2, §3.2 · pp. 3–4Results
Across two backbones, three samplers, and math and code tasks, CoDR improves the accuracy–compute trade-off.
Source: §5 · p. 9Takeaways & Limitations
Reversible decoding can correct early commitments without retraining or replacing the base sampler.
Source: §5 · p. 9Refill quality remains bounded by the underlying MDLM and sampler, so occasional regressions may occur, especially in arithmetic-heavy cases.
Source: Limitations · p. 10
Abstract
from arXiv · showhide
Masked diffusion language models (MDLMs) decode by repeatedly committing tokens to masked positions, but these commitments are usually irreversible. A token chosen under sparse, partial context is kept fixed, even when later context no longer supports it. Existing samplers mainly decide when to commit a token, but rarely check whether an already committed token should still be kept, allowing early mistakes to propagate. We trace this issue to confidence drift, where the model's confidence in a committed token drops from its sparse commit-time context to the denser context available later. Based on this signal, we propose CoDR (Confidence Drift Remasking), a training-free and sampler-agnostic refinement pass. CoDR estimates drift for all committed positions in only k forward passes via k-partition probing, then remasks and regenerates only the tokens the model no longer endorses. Across two backbones, four reasoning and coding tasks, and three base samplers, CoDR improves average accuracy across all evaluated model-sampler configurations and improves most individual task settings with modest overhead. Controlled experiments show that the gains come from targeted confidence-drift remasking rather than extra compute alone, and that CoDR uses far fewer forward passes than prior remasking methods. Code is available at https://github.com/YueWu0301/CoDR.
1 Introduction
Masked diffusion samplers improve token commitment but usually leave committed tokens fixed, even as denser context changes their support. CoDR detects these confidence drops and selectively revises affected tokens.
1 Introduction
- Tokens committed under sparse context can become incompatible with later context, and co-committed tokens may introduce inconsistencies that propagate through subsequent predictions.The problem is especially pronounced when several positions are committed in parallel.Source: §1 · p. 1
- Existing samplers use confidence, entropy, or attention to choose masked positions for commitment, while remasking methods do not directly measure changes in token support as decoding progresses.ReMDM revisits commitments and CORE probes robustness to perturbed contexts.Source: §1 · pp. 1–2
- CoDR is a training-free refinement pass that estimates confidence drift for committed tokens with k-partition probing, then selectively remasks and refills tokens the model no longer supports.It is designed as a plug-in that does not modify the base model or sampler.Source: §1 · p. 2
- CoDR improves the accuracy–compute trade-off across models, tasks, and samplers, with controlled comparisons attributing gains to targeted revision rather than unguided extra computation.The paper identifies confidence drift as a temporal symptom of decoding instability.Source: §1 · p. 2
2 Related Work
Related work develops masked diffusion models and decoding rules, while reversible-generation methods revisit committed tokens using signals other than confidence drift.
2 Related Work
- Discrete diffusion work established the mask state for language, and later objectives and parameterizations supported increasingly effective masked diffusion language models.D3PM introduced discrete-state diffusion; SEDD, MDLM, and RADD developed subsequent objectives and parameterizations.Source: §2 · p. 2
- MDLM samplers use signals such as entropy, confidence, denoising history, KL criteria, or attention to choose commitments, but typically keep committed tokens fixed.ReMDM iteratively remasks tokens, while CORE tests their robustness under perturbed contexts.Source: §2 · p. 2
- CoDR belongs to reversible generation but differs by using confidence drift: a drop in support for the same committed token as surrounding context becomes denser.This signal targets tokens whose support decreases as decoding progresses.Source: §2 · p. 2
3 Method
CoDR compares a committed token’s confidence at commitment with its confidence under later, denser context, then selectively remasks and refills commitments that lose support. It estimates drift for all committed positions with k-partition probing in k forward passes.
3 Method
- CoDR compares commit-time and later-context confidence, estimates drift with k-partition probing, and selectively remasks and refills drifted tokens using the original sampler.The method leaves other tokens fixed.Source: §3 · pp. 2–3, Figure 2 · p. 3
3.1 Preliminaries: MDLM Decoding
- At each decoding step, the sampler commits selected masked positions using predictions conditioned on the current partial sequence.A co-committed group is treated as conditionally independent given that sequence.Source: §3.1 · p. 3
3.2 Confidence-Drift Detection via k-Partition Probing
- Confidence drift is the discrepancy detected by rescoring a committed token under denser context after it was chosen under sparser context.A large positive drift score signals that the token has lost support; k-partition probing estimates drift for all committed positions in k passes, with k = 2 by default.Source: §3.2 · pp. 3–4
- Drift is a temporal instability signal, not proof that a token is incorrect or a universally optimal revision rule.The method’s practical informativeness is supported by the paper’s trigger comparisons and revision diagnostics.Source: §3.2 · p. 4
3.3 Confidence-Drift Remasking and Refill
- CoDR remasks commitments whose support has decreased beyond τdrift, then invokes the base sampler to refill them under denser bidirectional context.If no drifted position is detected, the sequence remains unchanged; refinement can occur after every block or on less frequent schedules.Source: §3.3 · pp. 4–5
4 Experiments
The experiments test CoDR across models, tasks, samplers, mechanisms, compute use, and robustness. Results report accuracy gains across model-sampler configurations and examine drift-guided revision and compute trade-offs.
4.2 Main Results (RQ1)
- Across both models and all three base samplers, CoDR improves accuracy; LLaDA fd+CoDR averages 48.93% at 123 NFE versus ReMDM’s 48.44% at 512 NFE.On Dream-7B, top1+CoDR reaches 60.62% average accuracy at 324 NFE, while fd+CoDR reaches 60.20% at 135 NFE.Source: §4.2 · pp. 5–6
4.3 Effectiveness of Core Mechanisms (RQ2)
- At matched remasking budgets, conf_drift achieves the best accuracy on both datasets; random remasking adds compute without gains, while bottom-m is less accurate and more expensive.The bottom-m baselines select by current confidence alone, without comparing it with commit-time confidence.Source: §4.3 · p. 6, Table 2 · p. 6
- At k=2, random partitions raised remasking (51.6→83.4 on HumanEval; 57.2→99.4 on GSM8K) and lowered accuracy (43.90%→43.53%; 77.40%→76.40%).The authors use deterministic stride probing by default.Source: §4.3 · pp. 6–7, Table 3 · p. 6
4.4 What Confidence Drift Captures (RQ3)
- Edited tokens had substantially larger median drift than unchanged or non-remasked tokens across B1–B8; Q5 edit rates reached 16.18% on HumanEval and 14.08% on GSM8K.Wrong→Correct examples had the largest maximum and edited-token drift on GSM8K (0.734/0.432); HumanEval showed the same direction, but only two examples made that transition.Source: §4.4 · p. 7, Figure 3 · p. 7
4.5 Practical Compute Efficiency (RQ4)
- Sweeping sampler compute knobs did not improve best baseline accuracy, while CoDR raised Fast-dLLM from 42.68% to 43.90% on HumanEval and 76.00% to 77.80% on GSM8K.The comparison suggests gains come from targeted revision rather than additional NFE alone.Source: §4.5 · p. 7, Figure 4 · p. 8
- Every-two-block refinement preserves HumanEval accuracy, slightly improves GSM8K, cuts NFE by 14–15%, and adds latency without increasing peak memory.Latency rises from 4.40 to 5.55 seconds without caching and from 2.77 to 3.90 seconds with KV drafting; this schedule is faster than every-block refinement in both settings.Source: §4.5 · p. 8, Table 5 · p. 8
4.6 Ablation Study (RQ5)
- Ablations favor k=2, τdrift=0, one refinement pass, and b=32; larger blocks degrade the base sampler, while CoDR improves over it at most block lengths.The authors retain b=32 for the best GSM8K result and a strong HumanEval result at lower NFE than b=16.Source: §4.6 · pp. 8–9
- Across generation lengths 128, 256, and 512, CoDR improves average accuracy for both samplers, though length 128 includes a MATH-500 regression.At length 512, gains are 1.90 points for Fast-dLLM and 0.90 for EB-Sampler; NFE overhead increases with sequence length.Source: §4.6 · p. 9
5 Conclusion
CoDR selectively revisits MDLM commitments whose confidence decreases as context grows, using k-partition probing and the original sampler to refill unstable tokens. The paper reports improved accuracy–compute trade-offs and identifies adaptive probing and cache-aware refinement as future directions.
5 Conclusion
- CoDR is a training-free method that uses k-partition probing to check commitments in a few forward passes, then selectively refills unstable tokens under richer context.The method revisits tokens whose confidence decreases as context grows.Source: §5 · p. 9
- Across two backbones, three samplers, and math and code tasks, CoDR improves the accuracy–compute trade-off; diagnostics associate larger drift with token edits and successful corrections.Controlled comparisons distinguish targeted revision from extra decoding alone.Source: §5 · p. 9
- Figure 6 reports four-task average accuracy and task-level gains across generation lengths for Fast-dLLM and EB sampling; future work may explore adaptive probing and cache-aware refinement.Source: Figure 6 · p. 9, §5 · p. 9
Limitations
CoDR is a targeted inference-time revision method, not a complete error detector. Its effectiveness depends on the base sampler, and refill can introduce fresh errors.
Limitations
- CoDR revises only tokens whose support decreases under denser context, leaving other tokens unchanged.The authors frame confidence drift as a targeted revision signal rather than a complete error detector.Source: Limitations · p. 10
- Revised-token quality remains bounded by the MDLM and sampler used for generation.The method’s effectiveness depends on the underlying base sampler used for refill.Source: Limitations · p. 10
- Refill can introduce fresh local errors and occasional regressions, especially in arithmetic-heavy cases.The authors suggest adaptive refinement schedules, stronger refill constraints, and broader evaluation as future work.Source: Limitations · p. 10
A Detailed Experiment Setting
Experiments follow the settings specified in Section 4.1 unless noted otherwise. Paired bootstrap resampling finds significant average accuracy gains for each LLaDA sampler.
A Detailed Experiment Setting
- Unless otherwise noted, experiments use the model, decoding, precision, prompting, and hardware settings described in Section 4.1.Sampler-specific implementation details are provided in later appendix sections.Source: A Detailed Experiment Setting · p. 12
B.1 Paired Bootstrap Analysis
- All confidence intervals exclude zero, indicating a significant average gain for each LLaDA sampler.The analysis resamples paired baseline/CoDR predictions 10,000 times per task and averages the four task gains equally.Source: §B.1 · p. 12, Table 6 · p. 12
B.2 Additional Benchmarks
CoDR improves accuracy over both Fast-dLLM and EB-Sampler on BBH and GPQA. Additional analyses relate larger drift to token edits and successful corrections, while EB refill reduces NFE without changing the reported HumanEval accuracy.
B.2 Additional Benchmarks
- CoDR improves accuracy over both Fast-dLLM and EB-Sampler on BBH and GPQA.These benchmarks extend the main results beyond the original four tasks.Source: §B.2 · p. 12, Table 7 · p. 12
C.1 Token-Level Drift Diagnostics
- Tables 8 and 9 provide the numerical results underlying Figure 3: per-block median drift by refinement outcome and edit rates across positive-drift quintiles.The supplied passages describe these diagnostics but do not state their specific numerical findings.Source: §C.1 · p. 12
C.2 Output-Level Correctness Transitions
- Wrong→Correct examples show higher maximum and edited-token drift on GSM8K, with the same direction on HumanEval, where there are only two such examples.The results associate larger drift with actual token revision and successful output correction; the accompanying tables report per-block drift and edit rates by positive-drift quintile.Source: §C.2 · p. 12, §C.1 · p. 12
C.3 Top1 Setting and Analysis
- 0.4268 pass@1 on LLaDA-8B HumanEval and 0.5915 on Dream-7B HumanEval are identical for top-1 and EB refill; EB reduces NFE from 377 to 303 and from 393 to 317, respectively.With phase-1 top-1 generation held fixed, EB refill saves about 20% NFE and is used as the default refill sampler.Source: §C.3 · pp. 12–13
C.4 CoRE Setting and Analysis
In the CoRE comparisons, block-wise decoding outperforms single-block decoding on LLaDA’s four-task average, while single-block decoding collapses on Dream. The analysis also motivates scaling candidate-set size to block length.
C.4 CoRE Setting and Analysis
- The CoRE protocol uses length 256, top-1 decoding, revision every E = 8 steps in the mid-trajectory window, k = 1 token per step, and a 0.3 replacement-confidence threshold.For both models, the authors use block-wise decoding with block = 32, candidate-set size m = 4, and protect = 8.Source: §C.4 · p. 13
- Dream’s single-block decoding collapses to near-empty generations, with accuracy only 0.189 in humaneval; its validated fastdllm and top1 samplers also collapse.The authors attribute the failure to high confidence in an early end-of-sequence token before any context is committed; block-wise decoding removes the collapse.Source: §C.4 · p. 13
- With block size 32, m = 32 over-perturbs the committed context; reducing to m = 4 restores normal-length outputs and baseline-level accuracy on both models.On Dream HumanEval, pass@1 falls from 0.5732 to 0.3598 with m = 32; the authors scale m by perturbation ratio rather than absolute count.Source: §C.4 · p. 13
- At matched NFE ≈272, block-wise decoding raises LLaDA’s four-task average from 0.394 to 0.469 and wins on three of four tasks, with the largest gains on reasoning.Single-block is marginally stronger on HumanEval, but block-wise substantially improves MBPP, MATH-500, and GSM8K.Source: §C.4 · p. 13, Table 12 · p. 14
D Comprehensive Ablation Results
The appendices detail CoDR’s hyperparameter evaluations and case studies, showing successful code and math corrections alongside regressions from dropped context and fresh hallucinations.
D Comprehensive Ablation Results
- The ablation section evaluates generation length, probe stride, drift threshold, refinement-pass count, and block length on LLaDA-8B with Fast-dLLM in a 0-shot setting.The listed tables report accuracy and NFE across these configurations.Source: D Comprehensive Ablation Results · p. 14, Table 13 · p. 14, Table 14 · p. 14
- Additional refinement passes beyond n = 1 add substantial compute for little or no gain.This is the reported pattern in the refinement-pass ablation.Source: Table 16 · p. 14
- The block-length analysis compares baseline and refined outputs; extreme sparse boundaries (b = 128) let early contextual errors compound and limit refinement.The table caption identifies block length as controlling intervention frequency and the accuracy–NFE trade-off.Source: Table 17 · p. 15
E Detailed Case Studies
- The detailed case studies contrast the base Fast-dLLM sampler at τ = 0.9 with CoDR on LLaDA-8B.The case-study overview introduces mathematics and code-generation trajectories.Source: E Detailed Case Studies · p. 15
E.1 Success Cases
- CoDR repairs localized code and reasoning errors, including malformed rotation logic, a missing import, an incorrect quotient, and a distorted algebraic substitution.The case studies report recovery of the +4 rotation, insertion of import math, correction of 1764/12 to 147, and recovery of P = 50.Source: §E.1 · p. 15
E.2 Regression Cases
- Regressions arise from dropped context during refill or fresh hallucinations, and diagnostic analysis attributes them to refill variation rather than systematic drift-detection failure.The authors report virtually zero code-generation regressions on HumanEval (0 out of 164), while arithmetic tasks are more vulnerable to fresh hallucinations.Source: §E.2 · p. 15
F AI Ethics Statement
The ethics statement discloses limited LLM assistance during manuscript preparation and says potential ethical concerns were checked. The supplied examples also include both a code correction and a combinatorics regression.
F AI Ethics Statement
- The authors report using LLM assistants only for context refinement and some code refinement, and say they checked for potential ethical concerns.The disclosure concerns manuscript preparation.Source: F AI Ethics Statement · p. 15
- On HumanEval/89, CoDR changes a failed base-sampler rotation solution into a passing +4 rotation solution.The base code has malformed shift logic; the refined code uses the corrected +4 rotation.Source: F AI Ethics Statement · p. 16