Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models

Zhejun Jiang; Mijung Park

cs.LGOct 4, 2026

Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models

Authors: Zhejun Jiang, Mijung Park

Organizations: The University of British Columbia

Abstract

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.

Figures & tables

Dataset	Method	Accuracy (%) $\uparrow$	ECE gen ( $\times 100$ ) $\downarrow$	ECE ans ( $\times 100$ ) $\downarrow$	Brier gen ( $\times 100$ ) $\downarrow$
MBPP	Threshold (Fast-dllm)	$\mathbf{57.60}$	29.18	31.24	$31.85$
	Base Entropy	$57.00\,(-1.04)$	$28.99\,(+0.65)$	$31.53\,(-0.93)$	$31.23\,(+1.97)$
	Info-gain	$54.00_{\scriptscriptstyle 1.71}\,(-6.25)$	$31.60_{\scriptscriptstyle 1.71}\,(-8.29)$	$33.73_{\scriptscriptstyle 1.71}\,(-7.97)$	$33.15_{\scriptscriptstyle 1.23}\,(-4.07)$
	BayesER (best accuracy)	$\underline{57.53}_{\scriptscriptstyle 0.12}\,(-0.12)$	$\underline{23.54}_{\scriptscriptstyle 0.11}\,(+19.33)$	$\underline{25.70}_{\scriptscriptstyle 0.09}\,(+17.73)$	$\underline{28.40}_{\scriptscriptstyle 0.02}\,(+10.84)$
	BayesER (best calibration)	$56.73_{\scriptscriptstyle 1.01}\,(-1.51)$	$\mathbf{16.32}_{\scriptscriptstyle 1.21}\,(+44.07)$	$\mathbf{18.17}_{\scriptscriptstyle 1.02}\,(+41.84)$	$\mathbf{25.57}_{\scriptscriptstyle 0.36}\,(+19.73)$
HumanEval	Threshold (Fast-dllm)	54.88	30.90	30.59	$33.27$

Table 1: Evaluation on mathematical reasoning and programming tasks. Accuracy is reported in percent; ECE and Brier values and their standard deviations are multiplied by 100. Subscripts are the standard deviations of three seeds on the corresponding column scale. Numbers in parentheses are relative percentage improvements over confidence-thresholding decoding. Bold values indicate the best result and underlined values indicate the second-best result per metric within each benchmark.

Dataset	Method	Accuracy (%) $\uparrow$	ECE gen ( $\times 100$ ) $\downarrow$	ECE ans ( $\times 100$ ) $\downarrow$	Brier gen ( $\times 100$ ) $\downarrow$
Sudoku	Threshold (Fast-dllm)	$\mathbf{97.60}$	3.13	3.13	$\mathbf{2.24}$
	Base Entropy	$\mathbf{97.60}\,(+0.00)$	$\underline{2.98}\,(+4.79)$	$\underline{2.98}\,(+4.79)$	$\underline{2.28}\,(-1.54)$
	Info-gain	$96.07_{\scriptscriptstyle 0.50}\,(-1.57)$	$\mathbf{2.54}_{\scriptscriptstyle 0.08}\,(+18.85)$	$\mathbf{2.54}_{\scriptscriptstyle 0.08}\,(+18.85)$	$3.39_{\scriptscriptstyle 0.37}\,(-51.35)$
	BayesER (best accuracy)	$\underline{97.53}_{\scriptscriptstyle 0.31}\,(-0.07)$	$3.96_{\scriptscriptstyle 0.03}\,(-26.52)$	$3.96_{\scriptscriptstyle 0.03}\,(-26.52)$	$2.40_{\scriptscriptstyle 0.17}\,(-7.00)$
	BayesER (best calibration)	$97.40_{\scriptscriptstyle 0.20}\,(-0.20)$	$2.99_{\scriptscriptstyle 0.27}\,(+4.47)$	$2.99_{\scriptscriptstyle 0.27}\,(+4.47)$	$2.39_{\scriptscriptstyle 0.16}\,(-6.56)$
Countdown	Threshold (Fast-dllm)	18.90	66.46	61.54	$58.56$

Table 2: Evaluation on planning tasks.

Method	Quality Yield (%) $\uparrow$	Frag. Valid (%) $\uparrow$	QED $\uparrow$	SA $\downarrow$	ECE $\downarrow$
Native GenMol	$11.78\pm 0.69$	$99.33\pm 0.33$	$0.6369\pm 0.0104$	$4.3932\pm 0.0362$	$45.08\pm 0.29$
Base Entropy	$\underline{14.89}\pm 0.69$	$\underline{99.89}\pm 0.19$	$0.6599\pm 0.0011$	$4.2809\pm 0.0106$	$35.40\pm 0.59$
Info-gain	$3.00\pm 0.67$	$\mathbf{100.00\pm 0.00}$	$\mathbf{0.6704\pm 0.0206}$	$\mathbf{4.1528\pm 0.0595}$	$\mathbf{24.56\pm 1.71}$
DPRM ( Bu et al., 2026 )	$12.44\pm 0.69$	$\mathbf{100.00\pm 0.00}$	$\underline{0.6619}\pm 0.0117$	$\underline{4.2682}\pm 0.0742$	$\underline{30.62}\pm 1.33$
BayesER (Ours)	$\mathbf{15.33\pm 0.33}$	$99.67\pm 0.00$	$0.6602\pm 0.0006$	$4.2900\pm 0.0062$	$35.36\pm 0.73$

Table 3: Evaluation on GenMol V2 Scaffold Decoration. Quality Yield measures the percentage of generated attempts satisfying both drug-likeness (

\text{QED}\geq 0.6

) and synthetic accessibility (

\text{SA}\leq 4.0

Figure 1: Accuracy–calibration trade-off for MBPP with different hyperparameter choices.

Figure 2: Dream HumanEval cost versus uncertain depth.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Setting
Mean model	Dream-v0-Instruct-7B
Benchmarks	MBPP, HumanEval, GSM8K, MATH500, Sudoku, Countdown
Prompt protocol	task-matched Dream instruction prompts
Posterior mean	released checkpoint; zero LoRA update
Curvature	categorical model Fisher, low-rank KFAC ( $k=10$ )
Target modules	trailing-layer query/value projections

Appendix

Table 4: Main experimental configuration for the six-benchmark Dream study. HumanEval uses the MBPP-trained posterior without refitting; all other posteriors are fitted on task-specific or synthetic training data.

Task	Target	$L_{s}$	Rank	Alpha	Prior	$\tau_{\mathrm{B}}$
MBPP	Accuracy	1	8	16	300	0.85
MBPP	Calibration	1	8	16	200	0.90
HumanEval	Accuracy	1	8	16	300	0.85
HumanEval	Calibration	1	8	16	200	0.90
GSM8K	Accuracy	4	8	16	2,000	0.90
GSM8K	Calibration	1	8	16	500	0.85

Appendix

Table 5: Hyperparameters for the evaluation in the main results. HumanEval reuses the MBPP-trained posterior; the remaining rows are task-fitted.

Figure 3: HumanEval accuracy versus calibration.

Figure 4: GSM8K accuracy versus calibration.

Figure 5: MATH500 accuracy versus calibration.

Figure 6: Countdown accuracy versus calibration.

Figure 7: Measured Dream GSM8K cost versus uncertain depth. Full test set (

n=1{,}319

), GSM8K-fitted rank-8 q/v Laplace-LoRA with

S=10

, on H100 GPUs. Color denotes prior precision

\lambda

. Dashed lines show the confidence-threshold run. Wall time includes loading and scoring; the loose prior increases decoding steps, especially at four layers.

MBPP (500)
Method	Accuracy (%) $\uparrow$	ECE gen ( $\times 100$ ) $\downarrow$	ECE ans ( $\times 100$ ) $\downarrow$	Brier gen ( $\times 100$ ) $\downarrow$
Threshold (Fast-dllm)	$39.40$	$47.63$	$50.01$	$46.14$
Base Entropy	$41.20\,(+4.57)$	$\underline{44.51}\,(+6.53)$	$\underline{46.84}\,(+6.33)$	$\underline{43.48}\,(+5.76)$
Info-gain	$\underline{41.27}_{\scriptscriptstyle 0.61}\,(+4.74)$	$44.96_{\scriptscriptstyle 0.50}\,(+5.59)$	$48.01_{\scriptscriptstyle 0.52}\,(+4.00)$	$43.52_{\scriptscriptstyle 0.26}\,(+5.67)$
BayesER	$\mathbf{41.33}_{\scriptscriptstyle 0.58}\,(+4.91)$	$\mathbf{43.57}_{\scriptscriptstyle 0.67}\,(+8.52)$	$\mathbf{45.89}_{\scriptscriptstyle 0.66}\,(+8.23)$	$\mathbf{42.70}_{\scriptscriptstyle 0.45}\,(+7.46)$
HumanEval (164)

Appendix

Table 6: LLaDA full-test results. Accuracy is reported in percent; ECE and generated-content Brier values and their standard deviations are multiplied by 100. Subscripts are the standard deviations of three seeds on the corresponding column scale. Numbers in parentheses are relative percentage improvements over confidence-thresholding decoding. Bold values indicate the best result and underlined values indicate the second-best result per metric within each benchmark, including ties.

		MIRank				BayesER , $\tau_{\mathrm{B}}=0.85$				BayesER , $\tau_{\mathrm{B}}=0.90$
$L_{s}$	$\lambda$	Acc.	ECE g	ECE a	Steps	Acc.	ECE g	ECE a	Steps	Acc.	ECE g	ECE a	Steps
1	100	28.91	10.89	10.77	25.74	34.38	8.40	6.61	100.70	40.63	12.62	10.44	93.91
1	200	48.44	19.63	21.55	38.02	57.03	16.42	18.14	59.38	55.47	19.27	20.04	67.27
1	500	54.69	27.78	30.58	39.11	57.81	25.77	27.89	43.26	57.81	27.02	29.19	45.13
1	1000	53.13	30.62	33.09	43.25	57.81	27.18	29.16	37.36	58.59	27.30	29.60	42.54
1	2000	57.81	26.34	29.27	40.44	57.81	27.56	29.50	37.62	57.03	29.37	31.60	42.23

Appendix

Table 7: Comparison between MIRank and BayesER on 128 examples of MBPP. Accuracy is in percent; both ECE columns are multiplied by 100.

		MIRank				BayesER , $\tau_{\mathrm{B}}=0.85$				BayesER , $\tau_{\mathrm{B}}=0.90$
$L_{s}$	$\lambda$	Acc.	ECE g	ECE a	Steps	Acc.	ECE g	ECE a	Steps	Acc.	ECE g	ECE a	Steps
1	100	59.38	31.56	12.84	30.39	58.59	31.50	16.28	102.63	63.28	38.08	15.94	118.50
1	200	64.84	10.14	19.34	41.83	77.34	5.56	15.76	70.57	75.78	6.65	16.95	77.88
1	500	78.91	4.61	18.37	49.34	79.69	8.37	13.51	56.29	75.78	9.10	17.81	63.16
1	1000	77.34	8.21	19.78	52.70	76.56	10.08	18.80	53.72	78.13	7.45	16.38	60.42
1	2000	77.34	9.39	17.38	51.38	76.56	9.09	18.56	53.38	78.13	7.63	17.32	61.18

Appendix

Table 8: Comparison between MIRank and BayesER on 128 examples of GSM8K. Accuracy is in percent; both ECE columns are multiplied by 100.

Explore similar work

Masked Diffusion Decoding as $x$ -Prediction Flow
Jun 27, 2026Weitian Wang, Lianlei Shan, Shubham Rai +2Masked Diffusion Models Diffusion Language Model Decoding
Robert Bosch GmbH, Germany · University of the Chinese Academy of Sciences, China · Ruhr University Bochum, Germany
Reliable Parallel Decoding in Masked Diffusion Language Models
Sep 29, 2026Zhenghao He, Bohan Liu, Guangzhi Xiong +1Masked Diffusion Models Diffusion Language Model Decoding
Department of Computer Science, University of Virginia
Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models
Jun 15, 2026Heecheol Yun, Joonhyung Park, Joowon Kim +1Masked Diffusion Models Diffusion Language Model Decoding
1KAIST · 2AITRICS

Jun 27, 2026, cs.CLJ/K move · Enter open · S save

Masked Diffusion Decoding as $x$ -Prediction Flow

Weitian Wang, Lianlei Shan, Shubham Rai, Cecilia De La Parra+1

Robert Bosch GmbH, Germany · University of the Chinese Academy of Sciences, China · Ruhr University Bochum, Germany

Masked Diffusion Models Diffusion Language Model Decoding Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens, but their standard decoder reduces each step to a binary action: a position is either committed to a single token or left fully masked, with no representation of partial belief in between. This all-or-nothing regime discards rich predictive information and forces premature, irrevocable commitments, leading to poor performance under a limited decoding budget. In this paper, we reinterpret mask prediction as clean-state prediction (

x

-prediction) and show that it can be used to induce a continuous flow in input embedding space. Building on this view, we propose a continuous decoding framework for MDLMs where tokens can accumulate partial progress at each diffusion step and remain revisable. To match the uneven contextual constraints across positions in language, we replace the globally synchronous schedule in image diffusion with a confidence-based asynchronous update in which the diffusion progress is token-wise accumulated. Additionally, we introduce a lightweight policy network and formulate its training as a reinforcement learning problem. Applied to pretrained LLaDA, our continuous decoder reaches 97% of its performance on the HumanEval dataset with 25% of decoding budget.

Abstract

Explore similar work

Masked Diffusion Decoding as xxx-Prediction Flow

Reliable Parallel Decoding in Masked Diffusion Language Models

Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

Masked Diffusion Decoding as $x$ -Prediction Flow