Organizations: School of Computer Science and Engineering, Northeastern University, China · Department of Computer Science and Technology, Institute for AI, Tsinghua University, China · Alibaba Group, Hangzhou, China
Masked Diffusion Models (MDMs) have recently emerged as a promising non-autoregressive paradigm for sequence generation. However, their performance is highly sensitive to the choice of decoding strategy. In this work, we reveal that prevalent uncertainty-based decoding strategies induce two decoding biases in MDMs: rigid boundary bias and trivial token bias. These biases limit the model's reasoning ability and ultimately degrade generation quality. To address these challenges, we propose UNmasking Calibration for DecOding DEbiasing (Uncode), a decoding calibration framework that regularizes uncertainty-based decoding by incorporating two complementary priors to shape global decoding trajectories and promote content informativeness. Extensive experiments on three advanced MDMs across seven reasoning- and planning-intensive benchmarks demonstrate that Uncode consistently outperforms existing decoding strategies by more than 7%, while achieving performance comparable to autoregressive models of similar parameter scales. The code is publicly available at https://github.com/NEUIR/Uncode.
Figures & tables
Figure 1: Illustration of decoding biases during reasoning. Uncertainty-based decoding prioritizes (A) rigid boundaries and (B) trivial tokens, causing systematic deviations from optimal problem-solving paths. In contrast, panel (C) presents the ideal unmasking pattern for producing coherent logical chains to achieve reliable inference.
Figure 2: Visualization of the Rigid Boundary Bias and its impact on downstream performance. (a) Unmasking probability for each token position across decoding steps on GSM8K, with both sequence length and decoding steps set to 256, where darker blue intensities denote higher unmasking probabilities. (b) Accuracy comparison of different decoding strategies on GSM8K (reasoning) and Sudoku (planning).
Figure 3: Verification of trivial token bias and trivial-token suppression efficacy. (a) The proportion of trivial tokens unmasked at each decoding step under the uncertainty-based sampler consistently exceeds the AR baseline (Qwen-2.5-7B-Instruct, dashed line). (b) GSM8K accuracy improves monotonically with suppression probability p .
Methods & LLMs
HumanEval
MBPP
GSM8K
MATH500
GPQA
Countdown
Sudoku
Avg.
Autoregressive LLMs
LLaMA-3.1-8B
53.1
56.7
83.9
23.8
31.0
27.0
0.0
39.4
Mistral-7B
43.9
37.0
49.4
7.2
28.1
22.7
0.0
26.9
Qwen-2.5-7B
78.1
62.8
71.9
64.2
32.8
7.7
0.0
45.3
LLaDA-8B-Instruct
Uniform ( 2021a )
15.2
24.6
48.8
15.0
29.0
14.4
2.2
21.3
Table 1: Experimental results on seven different benchmarks. The best score is highlighted in bold , and the second-best is underlined . Following prior practices Nie et al. (2025b) , we use 4-shot settings for GSM8K and MATH500, 5-shot for GPQA and Sudoku, 0-shot for HumanEval and MBPP, and 3-shot for Countdown. Methods marked with † denote samplers using the Semi-AR strategy. Results on Dream-7B are provided in Table 10 (Appendix A.15 ).
Figure 4: Ablation results of individual modules on LLaDA-8B-Instruct and LLaDA-1.5-8B, reporting the average performance across all evaluation benchmarks.
Figure 5: Analysis of uncertainty dynamics during decoding. (a) Answer Entropy vs. Decoding Step. While the baseline (Confidence, blue) often prematurely unmasks answer tokens with high uncertainty, Uncode (red) delays these answer tokens to later steps (bottom-right cluster), ensuring they are generated with high confidence. (b) Uncode achieves faster global uncertainty reduction than the baseline.
Methods
HumanEval
MBPP
GSM8K
MATH500
GPQA
Countdown
Sudoku
Avg.
τ -leaping ( 2023 )
17.7
38.9
55.2
16.2
28.9
32.1
3.2
27.5 (2.14x)
+ Uncode
22.6
40.5
75.4
30.8
29.5
28.4
25.4
36.1 (1.99x)
EB-sampler ( 2025 )
26.8
43.3
61.2
11.6
29.5
34.1
24.2
33.0 (2.32x)
+ Uncode
41.5
46.6
79.3
35.2
28.4
36.2
25.6
41.8 (2.28x)
Fast-dLLM ( 2025a )
35.0
45.9
77.4
25.2
27.8
24.4
0.0
33.9 (4.21x)
+ Uncode
36.0
48.2
77.8
29.2
28.4
36.3
0.6
36.7 (3.99x)
Table 2: Performance and efficiency analysis of Uncode with integrated efficient decoding strategies. The last column reports the average performance, with the speedup factor relative to vanilla decoding shown in parentheses.
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Unmasking dynamics under anchor insertion. A trivial high-frequency token (e.g., space, “the”, or “.”) is manually unmasked before decoding at an intermediate position (25% or 75% of the sequence length). The resulting unmasking patterns reveal a strong local positional bias, with decoding trajectories propagating outward from the inserted anchor.
Figure 7: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on GSM8K. (d) Corresponding performance comparison.
Figure 8: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on MBPP. (d) Corresponding performance comparison.
Figure 9: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on HumanEval. (d) Corresponding performance comparison.
Figure 10: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on GPQA. (d) Corresponding performance comparison.
Figure 11: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on Countdown. (d) Corresponding performance comparison.
Figure 12: Consistent boundary-first unmasking patterns across various uncertainty metrics. (a-c) Visualization of unmasking probability for Confidence, Entropy, and Margin samplers on Sudoku. (d) Corresponding performance comparison.
Definition
MDM
AR
Conservative definition (Ours)
41.7%
20.9%
Expanded w/ NLTK stopwords
49.0%
33.7%
Whitespace-only
8.9%
7.2%
Appendix
Table 3: Robustness of the trivial-token bias observation across different definitions. Values denote the average ratio of trivial tokens generated by MDM and AR.
Dataset
Task Type
#Test
Metric
HumanEval
Code Generation
164
Pass@1
MBPP
Code Generation
427
Pass@1
GSM8K
Math Reasoning
1,319
Acc.
MATH500
Math Reasoning
500
Acc.
GPQA
Science QA
448
Acc.
Countdown
Planning
1,000
Acc.
Appendix
Table 4: Statistics of evaluation datasets.
Figure 13: Analysis of decoding behavior under different λ and α settings. We find that λ substantially affects the induced decoding trajectory, whereas performance remains largely stable across a broad range of α , suggesting that Uncode is not sensitive to the clipping threshold. To examine task-specific effects, we report average results across non-Sudoku datasets (blue) and separately for Sudoku (red).
Method
HumanEval
MBPP
GSM8K
Math500
GPQA
Countdown
Sudoku
Avg.
Confidence
27.4
42.4
59.1
20.8
27.9
34.0
23.8
33.6
L-2-R + Sem. Prior
40.2
45.2
78.2
32.8
27.9
35.7
0.0
37.1
Semi-AR + Sem. Prior
40.2
44.7
78.8
33.2
27.5
32.8
0.0
36.7
Uncode
42.1
47.8
79.2
34.8
29.2
36.3
29.8
42.7
Appendix
Table 5: Comparison of Uncode against fixed-order baselines augmented with the semantic prior.
Method
HumanEval
MBPP
GSM8K
Math500
GPQA
Countdown
Sudoku
Avg.
Confidence
18.9
40.5
58.2
24.0
27.2
22.8
0.0
27.4
Semi-AR
37.8
39.1
75.4
33.2
26.8
27.3
0.0
34.2
Fast-dLLM
38.4
40.3
75.7
34.2
26.8
33.8
0.0
35.6
Uncode
43.3
48.2
79.9
34.4
30.4
35.7
26.0
42.6
Appendix
Table 6: Performance under an extended maximum generation length of 512.
Method
MBPP
Math500
Confidence
1.03
48.88
Uncode
1.04
49.08
Appendix
Table 7: Latency comparison (seconds per iteration) on LLaDA-8B-Instruct with 4 × A100 GPUs.
Figure 14: Correlation analysis between Mean Predictive Entropy and task performance across varying λ values. The blue curve denotes mean predictive entropy (left axis), while the red curve indicates performance (right axis) on (a) GSM8K and (b) HumanEval. The strong negative correlation ( r≈−0.99 ) empirically validates that lower mean predictive entropy serves as a reliable indicator of higher reasoning quality.
Calibration Corpus
HumanEval
MBPP
GSM8K
MATH500
GPQA
Countdown
Sudoku
Avg.
LLaDA-Instruct-8B
Book & Code ( 2015 ; 2025 )
42.1
47.3
79.2
35.0
27.9
36.3
24.4
41.7
SlimPajama ( 2023 )
40.9
47.1
79.1
35.2
28.4
36.3
30.8
42.5
C4 (16 GB) ( 2020 )
42.1
47.8
79.2
34.8
29.2
36.3
29.8
42.7
C4 (50 GB) ( 2020 )
42.2
47.8
79.2
36.0
29.0
36.3
30.0
42.9
C4 (200 GB) ( 2020 )
42.9
47.8
79.2
36.2
29.0
36.3
30.0
43.1
Appendix
Table 8: Evaluation of Uncode across various calibration corpora. ‘Book & Code’ refers to a combination of BookCorpusOpen and OpenR1-Math-220K datasets. C4 ( x GB) denotes a sampled subset of the C4 dataset, where x represents the size of the subset in gigabytes.
λ
HumanEval
MBPP
GSM8K
Math500
GPQA
Countdown
Sudoku
Avg.
0.00
7.93 (2.397)
33.5 (1.538)
59.7 (2.565)
4.0 (2.482)
29.0 (3.852)
34.0 (2.532)
29.8 (1.288)
28.28
0.25
42.1 (1.465)
47.8 (0.495)
79.2 (2.014)
34.8 (1.187)
29.2 (3.821)
34.7 (1.641)
0.2 (2.366)
38.29
0.50
41.5 (1.556)
47.8 (0.497)
78.9 (2.017)
34.0 (1.318)
28.8 (3.871)
36.3 (1.529)
0.0 (2.791)
38.19
1.00
40.2 (1.575)
45.4 (0.543)
77.7 (2.041)
34.4 (1.266)
28.4 (3.895)
35.7 (1.506)
0.0 (2.474)
37.93
Appendix
Table 9: Adaptive λ selection with only 10 unlabeled calibration instances. Values in parentheses are mean predictive entropy measured on the 10-instance calibration subset.
Methods & LLMs
HumanEval
MBPP
GSM8K
MATH500
GPQA
Countdown
Sudoku
Avg.
Autoregressive LLMs
LLaMA-3.1-8B-Instruct
53.1
56.7
83.9
23.8
31.0
27.0
0.0
39.4
Mistral-7B-Instruct
43.9
37.0
49.4
7.2
28.1
22.7
0.0
26.9
Qwen-2.5-7B-Instruct
78.1
62.8
71.9
64.2
32.8
7.7
0.0
45.3
LLaDA-8B-Instruct
Uniform ( 2021a )
15.2
24.6
48.8
15.0
29.0
14.4
2.2
21.3
Appendix
Table 10: Experimental results on seven different benchmarks. The best score is highlighted in bold , and the second-best is underlined . Following prior practices Nie et al. (2025b) , we use 4-shot settings for GSM8K and MATH500, 5-shot for GPQA and Sudoku, 0-shot for HumanEval and MBPP, and 3-shot for Countdown. Methods marked with † denote samplers using the Semi-AR strategy. All settings follow prior works Nie et al. (2025b) ; Zhao et al. (2025) ; Wang et al. (2025) for fair comparison.
Figure 15: A case study on the GSM8K dataset illustrating how the generation order impacts mathematical reasoning. The baseline’s premature commitment to an incorrect answer highlights a critical failure mode, whereas Uncode ’s coherent process leads to the correct solution. The generation timing is color-coded as follows: Early Stage , Middle Stage , and Late Stage .
Figure 16: A case study on the GSM8K dataset illustrating how the trivial token bias impacts mathematical reasoning. The baseline decodes trivial tokens at critical steps (e.g., the comma is decoded at step 121, at position 76, instead of the 7 at position 75), locking the reasoning space (e.g., the comma, used as a thousand separator, causes failure as 0,700 is invalid), and leading to incorrect predictions. In contrast, Uncode suppresses the selection of trivial tokens, resulting in the correct solution. The generation timing is color-coded as follows: Early Stage and Middle Stage .
Figure 17: Case study on the HumanEval dataset for code generation. The baseline method exhibits catastrophic failure due to the premature generation of trivial tokens, whereas Uncode , through explicit trajectory control, produces robust and complete code. The generation timing is color-coded as follows: Early Stage , Middle Stage , and Late Stage .
Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the Info-Gain Sampler, a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.
Kaisen Yang, Jayden Teoh, Kaicheng Yang +2
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2
Masked Diffusion Language Models (MDLMs) generate sequences by iteratively replacing masked tokens with model predictions. At each denoising step, the decoder chooses which positions are sufficiently confident to commit. Existing decoding methods typically rely on softmax confidence, which can be miscalibrated. We introduce BayesER (BAYESian Entropy-based Reordering), a post-hoc Bayesian decoding framework that uses predictive uncertainty to guide token commitment. In BayesER, we construct a lightweight approximate posterior centered at the pretrained checkpoint, similar to Laplace-LoRA but without training LoRA adapters. We average predictions over posterior samples and use predictive entropy to prioritize reliable positions. We examine how posterior predictions affect position ordering and token selection across benchmarks spanning code generation, mathematical reasoning, planning, and molecular generation. We show that BayesER reduces sequence-level calibration error while preserving or improving accuracy relative to common decoding schemes, including confidence-threshold decoding. Additionally, a posterior fitted on one code-generation dataset reduces calibration error on another without refitting, suggesting that Bayesian uncertainty may provide a transferable signal for more reliable MDLM decoding.
Chain-of-thought reasoning helps autoregressive models solve complex problems by generating intermediate steps that support later predictions. Masked diffusion models (MDMs) offer a similar opportunity through arbitrary-order generation: they can ideally reveal intermediate results along logical dependencies. In practice, however, standard decoding simply prioritizes high-confidence tokens, which need not align with this dependency order. We identify this discrepancy as the confidence shortcut: models commit with high certainty to plausible tokens while neglecting long-range dependencies. In multi-digit addition, models predict higher-order digits without properly tracking carries through long chains. Controlled pretraining across diverse reasoning tasks confirms that confidence-guided ordering often selects suboptimal sequences, and confidence-aligned training schemes can exacerbate these failures---for example, increasing addition error rates by an order of magnitude. Our findings caution against relying solely on confidence to choose generation orders and against training objectives that reinforce this preference. The experimental code is available at https://github.com/jinha2536/mdm-arithmetic.
Dueun Kim, Albert No
Department of Artificial Intelligence, Yonsei University