Paper deep dive
Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models
Zhejun Jiang, Mijung Park
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.05125v1
- Canonical: https://arxiv.org/abs/2610.05125v1
Trouble viewing inline? Open PDF directly →
Full Text
68,405 characters extracted from source content.
Expand or collapse full text
Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models Zhejun Jiang Mijung Park Affiliation: The University of British Columbia Affiliation: jzj03@cs.ubc.ca mijungp@cs.ubc.ca 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. 1 Introduction Language models can make wrong predictions with high confidence. This failure is relatively prevalent, as users often mistake stylistic confidence for evidence without justified certainty. A confidence score should mean what it claims: predictions with confidence q should be correct about a fraction q of the time; this is known as calibration. Formally, if the model M is perfectly calibrated, the following holds: ℙ(′s prediction is correct∣′s confidence is q)=qP(M s prediction is correct s confidence is q)=q, according to Guo et al. (2017). This definition of calibration does not eliminate hallucinations or reasoning errors common in modern language models. However, it makes uncertainty actionable, enabling a user or downstream system to decide whether to trust an answer, verify it, defer to a human, or abstain. Calibration in masked diffusion language models (MDLMs) plays an unusually direct role. Unlike autoregressive models that generate tokens from left to right, MDLMs iteratively denoise a partially masked sequence. At each step, a bidirectional forward pass predicts a distribution over the vocabulary at every masked position, and the decoder selects a subset of these predictions to commit before continuing. This process enables parallel generation, arbitrary-order decoding, and infilling, and has recently scaled to models such as LLaDA and Dream (Nie et al., 2025; Ye et al., 2025). Consequently, confidence serves two distinct purposes in MDLMs: it should indicate whether the final sequence is likely correct (the usual confidence in the model’s prediction), and more interestingly, it can guide which token predictions are reliable enough to commit at a given time-step during generation. This creates a central decoding problem that does not arise in standard autoregressive generation: Which masked positions should be committed at each denoising step? Committing many positions in parallel reduces model evaluations, but once committed, those tokens provide context for all subsequent tokens being unmasked at later denoising steps. Committing too conservatively avoids such problems but reduces the computational advantage of parallel decoding. The token commitment ordering therefore controls both generation quality and inference cost. Existing decoders commonly commit to the tokens at certain positions with the largest top-token probabilities or all positions whose confidence exceeds a fixed threshold (Chang et al., 2022; Wu et al., 2025). These rules are efficient and often effective, but they rely on the softmax distribution of a single parameter estimate. Neural networks are often overconfident (Guo et al., 2017), and point-estimate confidence does not reveal whether a prediction remains stable under model uncertainty. In this paper, we investigate whether Bayesian predictive uncertainty provides a better signal for commitment ordering. Given an approximate posterior q(|)q( θ|D) over model parameters θ, the posterior predictive distribution over the output y of a model averages token probabilities across posterior samples, given a noisy state t x_t of the denoising process: p¯(|t)=∼q(|)[p(|t,)]. p( y| x_t)=E_ θ q( θ|D) [p( y| x_t, θ) ]. Its entropy, H(p¯(|t))H( p( y| x_t)), measures the total uncertainty of the Bayesian prediction. Our method, BayesER (BAYESian Entropy-based Reordering), is a post-hoc decoding framework for MDLMs that uses Bayesian predictive entropy. At each denoising step, BayesER forms a posterior predictive distribution for every currently masked position, prioritizes positions with low posterior predictive entropy, selects tokens from the posterior predictive mean, and uses posterior confidence to determine how many positions to commit. In this way, BayesER adjusts both the order and pace of decoding based on the Bayesian predictive distribution. For typical MDLMs such as LLaDa (Nie et al., 2025) and Dream (Ye et al., 2025), with 7 billion parameters or more, identifying posterior distributions over all parameters is computationally challenging. Laplace-LoRA first fine-tunes LoRA adapters to obtain a task-specific maximum a posteriori solution and then approximates a posterior around the adapted model (Yang et al., 2024). While appropriate for downstream adaptation, this procedure changes the mean predictor and makes it difficult to distinguish the benefit of Bayesian uncertainty from that of fine-tuning. We instead construct a curvature-only Laplace-LoRA posterior directly at the released checkpoint, without adapter training. The LoRA update is zero at the posterior mean, so the deterministic model is unchanged by construction. We fit the posterior covariance with training data, representing each block of a Hessian matrix using a Kronecker product of input and output-gradient second moments to reduce the cost of storing and expressing curvature. Across several benchmark datasets spanning code generation, mathematical reasoning, planning, and molecule generation, BayesER substantially reduces expected calibration error (ECE) while preserving—and in several settings improving—task accuracy relative to commonly used confidence-based decoding schemes. Our results show that prior precision and the number of Laplace-LoRA layers govern a three-way trade-off among calibration, task performance, and decoding cost. We further find that a posterior constructed on one code-generation dataset can reduce calibration error on another without refitting, hinting that the learned uncertainty signal may transfer across related tasks. 2 Background We first describe background information on a common decoding scheme in MDLMs, Laplace-LoRA, and expected calibration error (ECE) for evaluating calibration performance. 2.1 Masked diffusion decoding MDLMs generate text by repeatedly predicting and revealing masked tokens (Sahoo et al., 2024; Nie et al., 2025; Ye et al., 2025). We denote a partially decoded, noisy sequence of length L by t(1:L) x_t^(1:L), where t(l) x_t^(l) is a one-hot vector of length |||V| (vocabulary size), representing the l-th position of the sequence with one entry equaling 1 and the rest equaling zero. We also denote which position has masked tokens at time t by ℳt=l:t(l)=[MASK]M_t=\l: x_t^(l)=[MASK]\. We denote a pre-trained model by f0f_ θ_0 with fixed trained parameter values 0 θ_0. In MDLMs, at each time step t, the trained model f0f_ θ_0 outputs the token probability over the vocabulary for each position in the sequence: p0(l)(|t(1:L))p_ θ_0^(l)( v| x_t^(1:L)). The base token is selected by choosing the argmax of this probability: y^t(l)=argmaxp0(l)(|t(1:L)) y_t^(l)= _ vp_ θ_0^(l)( v| x_t^(1:L)) and confidence is: ct(l)=p0(l)(y^t(l)|t(1:L))c_t^(l)=p_ θ_0^(l)( y_t^(l)| x_t^(1:L)). Confidence-based decoding. MDLMs commonly decode tokens via confidence-based decoding. Given a threshold τ, at each denoising step t, KtK_t most confident base tokens are selected by: Kt=max(1,∑l∈ℳtct(l)≥τ),K_t= \! (1, _l _t 1_\c_t^(l)≥τ\ ), where 1 is an indicator function which returns 11, if ct(l)≥τ\c_t^(l)≥τ\, and returns 00, if not. The minimum of one ensures that at least one token is unmasked at a time. 2.2 Laplace-LoRA LoRA (Low-Rank Adaptation) parameterizes the update of each of the Q,K,VQ,K,V matrices of a transformer architecture at each layer through two low-rank matrices (Hu et al., 2021) (omitting the specification of Q,K,VQ,K,V or layer index for simplicity): =0+αr W= W^0+ αr B A, with ∈ℝr×din A ^r× d_ in and ∈ℝdout×r, B ^d_ out× r, where r is a rank of the low-rank matrices and α is a scaling hyper-parameter. Laplace-LoRA To gain posterior uncertainty about the model, Laplace-LoRA first optimizes the LoRA adapter parameters A and B on downstream tasks, then places a Gaussian distribution over A and B, while keeping the underlying weights 0 W^0 fixed (Yang et al., 2024). Laplace-LoRA uses the Laplace approximation to estimate the covariance of the adapter parameters by learning the curvature of the negative log posterior at a posterior mode (MacKay, 1992). For decent-sized large language models, computing the exact curvature is intractable. Laplace-LoRA proposes two levels of approximations: (a) rather than computing the exact curvature, it uses the Fisher information, :=∑n=1Ny∼P(⋅∣f(n))[∇logP(y∣f(n))(∇logP(y∣f(n)))⊤] F:= _n=1^NE_y P (· f_ θ(x_n) ) [ _ θ P (y f_ θ(x_n) ) ( _ θ P (y f_ θ(x_n) ) ) ] ; (b) furthermore, the Fisher information is computed on each linear layer separately using Kronecker-factored approximate curvature (KFAC), forming a block structure of the Fisher information matrix. KFAC represents each layer’s Fisher information using a Kronecker product of input covariance and output covariance, reducing the cost of storing and manipulating curvature (Martens & Grosse, 2015; Ritter et al., 2018). Our method, BayesER, also specifies the curvature using these approximations, but without fine-tuning adapters, but at a fixed checkpoint as shown in Section 3.1. 2.3 Calibration and sequence-level confidence Expected calibration error (ECE) (Naeini et al., 2015) measures how well the model’s prediction accuracy matches its confidence. Large ECE indicates a mismatch between confidence and accuracy, reflecting overconfidence or underconfidence. With N number of samples, predicted labels y^n y_n for each sample n, and true label yny_n, ECE groups confidences into M bins (each bin is ℬmB_m) and measures ECE ECE =∑m=1M|ℬm|N∣acc(ℬm)−conf(ℬm)∣, = _m=1^M |B_m|N acc(B_m)-conf(B_m) , (1) where acc(ℬm)=1|ℬm|∑n∈ℬm(y^n=yn),conf(ℬm)=1|ℬm|∑n∈ℬmp^nacc(B_m)= 1|B_m| _n _m 1( y_n=y_n),conf(B_m)= 1|B_m| _n _m p_n, and p^n p_n is the model’s confidence for sample n: p^n=maxcp(yn=c|xn) p_n= _cp(y_n=c|x_n). Directly applying ECE to text sequences is not straightforward. In Section 3.3, we show two ways to evaluate ECE on the generated sequence. Another commonly used probabilistic prediction metric which we will also use is the Brier score (Brier, 1950). For our binary task-correctness evaluation, BS=1N∑n=1N(cn−zn)2BS= 1N _n=1^N(c_n-z_n)^2, where cnc_n is the sequence confidence score and zn∈0,1z_n∈\0,1\ indicates task correctness. 3 Method: BayesER A decoder in MDLMs makes three decisions at every denoising step: which positions to commit, which tokens to write, and how many positions to commit. BayesER makes all three decisions from a checkpoint-centered posterior predictive distribution. Specifically, it ranks positions by posterior predictive entropy, selects token identities from the posterior predictive mean, and determines the commitment budget from posterior confidence. BayesER consists of two stages: (1) First, we fit the posterior distribution over LoRA adapter parameters using training data, from which we obtain the posterior samples; and (2) Then, using the posterior samples, we compute the predictive posterior entropy at every denoising step t to determine how many tokens to unmask and which positions to fill with which token. We describe each of these two stages below. 3.1 Phase I: Curvature-Only Laplace-LoRA Posterior Checkpoint-centered LoRA parameterization. We place LoRA adapters on selected query and value matrices in the selected J transformer layers. Below, we omit the indices for Q and V and show only the layer index j, since the LoRA operation is the same for Q and V. For a selected pretrained matrix j,0∈ℝdout×din W_j,0 ^d_out× d_in, the effective weight is j(ϕj)=j,0+αrjj,j∈ℝr×din,j∈ℝdout×r, W_j( _j)= W_j,0+ αr B_j A_j, A_j ^r× d_in, B_j ^d_out× r, (2) where we let ϕ0=j,0,j,0j, φ_0=\ A_j,0, B_j,0\_j, and j,0=0, B_j,0=0, with j,0 A_j,0 given by the standard LoRA initialization. Because j,0j,0=0 B_j,0 A_j,0=0, the posterior-mean model is exactly the released checkpoint: f0,ϕ0=f0.f_ θ_0, φ_0=f_ θ_0. Unlike standard Laplace-LoRA, we do not optimize the adapter means on a downstream objective. We use the training data only to estimate local curvature. Approximating Fisher information with KFAC For each LoRA adapter matrix index m, where m can vary between layers and Q and V, we denote its input activation for each observation ~n x_n by hm,nh_m,n, and let gm,n=∇am,nlogp0,ϕ0(Yn∣~n)g_m,n= _a_m,n p_ θ_0, φ_0 (Y_n x_n ) be the gradient with respect to the corresponding pre-activation am,na_m,n, where YnY_n is distributed according to the model predictive distribution. We approximate the Fisher information matrix separately for each adapter matrix, using Kronecker-factored approximate curvature (KFAC): Fm≈Nm(Gm⊗Cm), where Cm=[hmhm⊤], and Gm=[gmgm⊤], F_m≈ N_m (G_m C_m ), where C_m=E\! [h_mh_m ], and \;G_m=E\! [g_mg_m ], (3) where NmN_m counts the contributing source-position observations represented by the sampled rows for block m, rather than calibration sequences or masked targets. The factor CmC_m captures correlations among matrix inputs, whereas GmG_m captures correlations among output-gradient coordinates. Using an isotropic Gaussian prior precision λ, we construct the checkpoint-centered, Laplace-style Gaussian approximation qλ(ϕ)=(ϕ0,[blockdiag(F1,…,FR)+λI]−1),q_λ( φ)=N ( φ_0,\, [blockdiag(F_1,…,F_R)+λ I ]^-1 ), (4) where R is the total number of individual LoRA adapter matrices across the J we consider. For instance, if we have Q,VQ,V matrices per layer and apply LoRA to 10 layers, the posterior covariance consists of small, manageable 4040 blocks (since each Q and V has two adapter matrices, A and B). Together, these blocks define a single joint Gaussian over all selected adapter parameters. Blocks are independent across matrices, projections, and layers; however, within-matrix KFAC correlations remain, and the leading eigendirections of large factors are retained. Directions omitted by the truncated eigenspace retain prior variance. Uncertainty in the LoRA factors At the initialization ϕ0 φ_0, we have j,0 B_j,0 initialized as 0. Changing j,0 A_j,0 alone therefore does not affect j(ϕj) W_j( φ_j), so the gradient with respect to j A_j and its Fisher block is zero. Under our Gaussian approximation, this gives: vec(j)∼(vec(j,0),λ−1).vec( A_j) \! (vec( A_j,0),\,λ^-1 I ). Therefore, uncertainty in j A_j is entirely prior-controlled, while uncertainty in j B_j depends on both the calibration curvature and the prior. At prediction time, each posterior draw ϕ(s)∼qλ(ϕ) φ^(s) q_λ( φ) gives a corresponding pair of sampled LoRA matrices j(s) A_j^(s) and j(s) B_j^(s). We use these matrices to construct the sampled model weights: j(ϕj(s))=j,0+αrj(s)j(s) W_j( φ_j^(s))= W_j,0+ αr B_j^(s) A_j^(s). Although j A_j has zero local Fisher curvature at initialization, its sampled variation still affects predictions because j(s) B_j^(s) is generally non-zero. Writing j(s)=j,0+δj(s) A_j^(s)= A_j,0+δ A_j^(s) and j(s)=δj(s) B_j^(s)=δ B_j^(s), the sampled update includes the bilinear term αrδj(s)δj(s) αrδ B_j^(s)δ A_j^(s) in addition to αrδj(s)j,0 αrδ B_j^(s) A_j,0. Thus, predictive variation reflects both data-informed curvature and prior-controlled variation in locally unconstrained directions. Algorithm 1 BayesER 0: Pre-trained model:f0f_ θ_0, confidence threshold: τB∈[0,1] _B∈[0,1], how many posterior samples to draw:S≥1S≥ 1, initial sequence T(1:L)x_T^(1:L), and the number of selected tokens per position:Kv≥1K_v≥ 1 1: Phase I: Fit the parameter posterior qλq_λ using Algorithm 2 2: Then, sample parameters (s)∼qλ θ^(s) q_λ for s=1,…,Ss=1,…,S 3: Phase I: 4: t←Tt← T; ←∅C← ; ℳT←l:T(l)=[MASK]M_T←\l: x_T^(l)=[MASK]\ 5: while ℳt≠∅M_t≠ do 6: Run base model f0(t(1:L))f_ θ_0(x_t^(1:L)) to obtain token logits and confidences ct(l)c_t^(l) for each position l 7: For each masked position l, define a set t(l)V_t^(l) containing the KvK_v highest-scoring tokens 8: for s=1,…,Ss=1,…,S do 9: Given each posterior sample, run the sampled model f(s)(t(1:L))f_ θ^(s)( x_t^(1:L)) 10: pt,s(l)←p_t,s^(l)← softmax of sampled-model logits over tokens in t(l)V_t^(l) 11: end for 12: p¯t(l)←S−1∑spt,s(l) p_t^(l)← S^-1 _sp_t,s^(l) 13: Ht(l)←−∑v∈t(l)p¯t(l)(v)logp¯t(l)(v)H_t^(l)←- _v _t^(l) p_t^(l)(v) p_t^(l)(v) 14: y^t(l)←argmaxv≠[MASK]p¯t(l)(v) y_t^(l)← _v≠[MASK] p_t^(l)(v); and ct(l)←p¯t(l)(y^t(l))c_t^(l)← p_t^(l)( y_t^(l)) 15: Kt←max(1,∑l∈ℳtct(l)≥τB)K_t← (1, _l _t 1_c^(l)_t≥ _B) 16: t←P_t← the KtK_t masked positions with lowest entropy Ht(l)H_t^(l) 17: Record (t,l,y^t(l),ct(l))(t,l, y^(l)_t,c^(l)_t) in C for each l∈tl _t 18: t−1(l)←y^t(l) x_t-1^(l)← y^(l)_t for each l∈tl _t 19: t←t−1t← t-1 20: ℳt←l:t(l)=[MASK]M_t←\l: x_t^(l)=[MASK]\ 21: end while 22: return Updated sequence t(l) x_t^(l) and saved tokens and their probabilities C 3.2 Phase I: Posterior Predictive Entropy for Position Ordering At step t, the pre-trained model produces base logits ℓ0,t(l)∈ℝ|| ^(l)_0,t ^|V| for each location l. For each position, we define a candidate set containing the top-KvK_v highest-scoring non-masked tokens, i.e., t(l)=TopKv∈∖[MASK](ℓ0,t(l)(v);Kv).V_t^(l)=TopK_v \ [MASK]\ ( ^(l)_0,t(v);K_v ). For S independent posterior draws ϕ(s)∼qλ(ϕ) φ^(s) q_λ( φ), let ℓt,s(l) ^(l)_t,s be the sampled-model logit. We first normalize each draw over the candidate set t(l)V_t^(l): pt,s(l)(v)=exp(ℓt,s(l)(v))∑u∈t(l)exp(ℓt,s(l)(u)),p^(l)_t,s(v)= \! ( ^(l)_t,s(v) ) _u _t^(l) \! ( ^(l)_t,s(u) ), for v∈t(l)v _t^(l). Importantly, we average probabilities rather than logits: p¯t(l)(v)=1S∑s=1Spt,s(l)(v). p_t^(l)(v)= 1S _s=1^Sp^(l)_t,s(v). Our position-ordering criterion is the posterior predictive entropy given by Htpred,(l)=H(p¯t(l))=−∑v∈t(l)p¯t(l)(v)logp¯t(l)(v).H_t^pred,(l)=H\! ( p_t^(l) )=- _v _t^(l) p_t^(l)(v) p_t^(l)(v). (5) Low entropy indicates that the posterior predictive distribution concentrates on a small number of candidates, making the position comparatively safe to commit. High entropy indicates that no clear token is supported under the current partially masked context, so the position should be deferred. Commitment Budgets Given a posterior-confidence threshold τB _B, BayesER sets its per-step commitment budget to Kt=max(1,∑l∈ℳtc(l)t≥τB).K_t= (1,\, _l _t 1_ \c^(l)_t≥ _B \ ). It then chooses the KtK_t positions with the lowest posterior predictive entropy: t=argmin⊆ℳt||=Kt∑l∈Htpred(l).P_t= subarraycP _t\\ |P|=K_t subarray _l H^pred(l)_t. (6) For each l∈tl _t, we commit t−1(l)=y^t(l), where y^t(l)←argmaxv≠[MASK],p¯t(l)(v) x_t-1^(l)= y^(l)_t, where y_t^(l)← _v≠[MASK], p_t^(l)(v), while all unselected positions remain unchanged. Algorithm 1 summarizes our algorithm. 3.3 Evaluating Sequence-Level Confidence To evaluate ECE on a generated sequence, we first compute the geometric mean as a length-normalized summary of commitment confidence, for a token scope ⊆S : πseq()=exp(1||∑(t,l)∈logct(l)). _seq(S)= ( 1|S| _(t,l) c^(l)_t ). (7) using the token probability available at the moment of commitment, before the committed token becomes part of the context for later denoising steps, denoted by =(t,l,yt(l),ct(l))C= \ (t,l,y^(l)_t,c^(l)_t ) \, with ct(l)=p¯t(l)(yt(l)).c^(l)_t= p_t^(l) (y^(l)_t ). This aggregation is a candidate confidence estimator, which we use for πseq(n) _seq(S_n) as the confidence p^n p_n for each generated output and compute ECE using equation 1, with correctness determined by the benchmark’s task criterion. In Section 5, we report the score over two scopes: all generated-content tokens (ECEgenECE_gen) and task-specific final-answer tokens (ECEansECE_ans). Note that low ECEgenECE_gen does not establish that every sentence in a reasoning trace is correct. It means that confidence aggregated over that trace agrees, at the resolution measured, with the frequency of final task success. The generated-content and final-answer variants use different token scopes but share the same task-level correctness event. This geometric mean is different from Li et al. (2026), where their ECE is an arithmetic mean of token probabilities extracted from a single final forward pass under teacher forcing on the completed sequence. Our geometric mean summarizes the model’s confidence at the commitment decisions made along the decoding trajectory. 3.4 Mutual Information and Base Entropy Mutual Information For comparison, one might consider the token–parameter mutual information (MI) given by It(l)=Htpred(l)−1S∑s=1SH(pt,s(l)).I_t^(l)=H^pred(l)_t- 1S _s=1^SH\! (p^(l)_t,s ). MI isolates disagreement among posterior draws (Houlsby et al., 2011), but it can be small even when every draw assigns a diffuse token distribution. This makes MI insufficient for deferred commitment: posterior samples may agree that a position remains ambiguous because the current state does not yet provide enough context. We therefore use total posterior predictive entropy as the default ordering score and retain MI-based ordering as an ablation. We empirically compare MI and predictive entropy in Appendix E. Base Entropy To distinguish the effect of entropy-based ordering from that of Bayesian marginalization, we can also compare our predictive entropy-based selection from base entropy given by: H0,t(l)=H(p0,t(l))H^(l)_0,t=H\! (p^(l)_0,t ), where p0,t(l)(v)=exp(ℓ0,t(l)(v))∑u∈t(l)exp(ℓ0,t(l)(u)).p^(l)_0,t(v)= \! ( ^(l)_0,t(v) ) _u _t^(l) \! ( ^(l)_0,t(u) ). This baseline, Base Entropy, ranks positions by H0,t(l)H^(l)_0,t while retaining the model’s tokens and confidence-based commitment budget. 4 Related Work Decoding and unmasking policies. MaskGIT and Fast-dLLM (Chang et al., 2022; Wu et al., 2025) present confidence-based decoders ; Deferred Commitment Decoding adapts the active window (Shu et al., 2026), and Entropy based sampler uses a point estimate of entropy to control parallel commitments (Ben-Hamu et al., 2025). Lookahead Unmasking and the Information-Gain (Info-Gain) Sampler evaluate candidate successor states (Lee et al., 2025; Yang et al., 2026). Other approaches learn unmasking policies through reinforcement learning (Jazbec et al., 2026) or derive ordering corrections from terminal rewards, as in DPRM (Bu et al., 2026). None of the above methods uses posterior predictive uncertainty to choose positions, tokens, and commitment counts. Bayesian inference for pretrained models. Laplace-LoRA fits a posterior around fine-tuned adapters (Yang et al., 2024), BLoB (Wang et al., 2024) learns a variational posterior during adapter training , and Training-Free Bayesianization (Shi et al., 2025) assigns uncertainty to already-trained adapters. Building on structured Laplace approximations (Ritter et al., 2018; Daxberger et al., 2021), our curvature-only construction fits uncertainty around zero-update adapters at the released checkpoint. This preserves the deterministic mean model and allows us to study posterior prediction without conflating it with downstream fine-tuning. Uncertainty in generation. For continuous diffusion, BayesDiff propagates uncertainty using last-layer Laplace (Kou et al., 2023), while Jazbec et al. (2025) develop a Bayesian measure for assessing synthetic samples. Our focus is utilizing the uncertainty at unresolved discrete token positions to guide commitment during generation. 5 Experiments 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 (%) ↑ ECEgen (×100× 100) ↓ ECEans (×100× 100) ↓ Briergen (×100× 100) ↓ MBPP Threshold (Fast-dllm) 57.6057.60 29.18 31.24 31.8531.85 Base Entropy 57.00(−1.04)57.00\,(-1.04) 28.99(+0.65)28.99\,(+0.65) 31.53(−0.93)31.53\,(-0.93) 31.23(+1.97)31.23\,(+1.97) Info-gain 54.001.71(−6.25)54.00_ 1.71\,(-6.25) 31.601.71(−8.29)31.60_ 1.71\,(-8.29) 33.731.71(−7.97)33.73_ 1.71\,(-7.97) 33.151.23(−4.07)33.15_ 1.23\,(-4.07) BayesER (best accuracy) 57.53¯0.12(−0.12) 57.53_ 0.12\,(-0.12) 23.54¯0.11(+19.33) 23.54_ 0.11\,(+19.33) 25.70¯0.09(+17.73) 25.70_ 0.09\,(+17.73) 28.40¯0.02(+10.84) 28.40_ 0.02\,(+10.84) BayesER (best calibration) 56.731.01(−1.51)56.73_ 1.01\,(-1.51) 16.321.21(+44.07)16.32_ 1.21\,(+44.07) 18.171.02(+41.84)18.17_ 1.02\,(+41.84) 25.570.36(+19.73)25.57_ 0.36\,(+19.73) HumanEval Threshold (Fast-dllm) 54.88 30.90 30.59 33.2733.27 Base Entropy 57.32(+4.45)57.32\,(+4.45) 27.25(+11.81)27.25\,(+11.81) 26.71(+12.68)26.71\,(+12.68) 31.80(+4.42)31.80\,(+4.42) Info-gain 51.220.61(−6.67)51.22_ 0.61\,(-6.67) 33.680.57(−9.00)33.68_ 0.57\,(-9.00) 33.280.53(−8.79)33.28_ 0.53\,(-8.79) 35.820.82(−7.66)35.82_ 0.82\,(-7.66) BayesER (best accuracy) 56.30¯3.01(+2.59) 56.30_ 3.01\,(+2.59) 24.77¯3.04(+19.84) 24.77_ 3.04\,(+19.84) 24.61¯3.13(+19.55) 24.61_ 3.13\,(+19.55) 30.26¯1.80(+9.05) 30.26_ 1.80\,(+9.05) BayesER (best calibration) 55.081.96(+0.36)55.08_ 1.96\,(+0.36) 20.162.04(+34.76)20.16_ 2.04\,(+34.76) 20.331.98(+33.54)20.33_ 1.98\,(+33.54) 27.761.05(+16.56)27.76_ 1.05\,(+16.56) GSM8K Threshold (Fast-dllm) 76.19 10.01 17.91 17.4617.46 Base Entropy 78.17(+2.60)78.17\,(+2.60) 7.14¯(+28.67) 7.14\,(+28.67) 15.98¯(+10.78) 15.98\,(+10.78) 16.39(+6.13)16.39\,(+6.13) Info-gain 69.700.44(−8.52)69.70_ 0.44\,(-8.52) 10.170.18(−1.60)10.17_ 0.18\,(-1.60) 18.800.06(−4.97)18.80_ 0.06\,(-4.97) 16.620.13(+4.81)16.62_ 0.13\,(+4.81) BayesER (best accuracy) 77.98¯0.38(+2.35) 77.98_ 0.38\,(+2.35) 7.170.41(+28.37)7.17_ 0.41\,(+28.37) 16.100.20(+10.11)16.10_ 0.20\,(+10.11) 16.53¯0.18(+5.31) 16.53_ 0.18\,(+5.31) BayesER (best calibration) 76.890.52(+0.92)76.89_ 0.52\,(+0.92) 6.500.55(+35.06)6.50_ 0.55\,(+35.06) 15.750.52(+12.06)15.75_ 0.52\,(+12.06) 16.980.26(+2.73)16.98_ 0.26\,(+2.73) MATH500 Threshold (Fast-dllm) 42.40¯ 42.40 45.50¯ 45.50 50.08¯ 50.08 45.33¯ 45.33 Base Entropy 40.00(−5.66)40.00\,(-5.66) 46.27(−1.69)46.27\,(-1.69) 53.31(−6.45)53.31\,(-6.45) 45.77(−0.98)45.77\,(-0.98) Info-gain 40.730.58(−3.94)40.73_ 0.58\,(-3.94) 45.960.70(−1.01)45.96_ 0.70\,(-1.01) 52.041.13(−3.91)52.04_ 1.13\,(-3.91) 45.720.46(−0.87)45.72_ 0.46\,(-0.87) BayesER (best accuracy) 43.531.17(+2.67)43.53_ 1.17\,(+2.67) 42.711.12(+6.13)42.71_ 1.12\,(+6.13) 49.551.09(+1.06)49.55_ 1.09\,(+1.06) 43.190.84(+4.72)43.19_ 0.84\,(+4.72) 5.1 Setup We evaluate Dream-v0-Instruct-7B (Ye et al., 2025) on code generation (MBPP, HumanEval), mathematical reasoning (GSM8K, MATH500), and planning (4×4 Sudoku, Countdown). HumanEval transfers the MBPP posterior without refitting. We also include LLaDA results in Appendix D. We additionally evaluate molecular scaffold decoration with GenMol V2. Detailed configurations are available in Appendix A.2. We compare against confidence-threshold decoding, Base Entropy, and Info-gain (Yang et al., 2026). Info-gain runs use N=8N=8 candidate paths with fixed K=2K=2 tokens per step. We report task accuracy and sequence-level calibration. We compare the sequence scores in Eq. 7 with binary task correctness using 15 equal-width bins for ECE (Section 2.3) and report binary Brier scores. We compute scores over generated-content tokens and final-answer tokens, yielding ECEgen and ECEans. For molecule generation, ECE is evaluated against chemical validity, while quality yield measures the fraction of attempts meeting the QED and SA criteria in Table 3. 5.2 Main results 5.2.1 Mathematical reasoning and programming Table 1 shows that BayesER improves sequence-confidence calibration on several tasks, with task-dependent changes in accuracy. On MBPP, the accuracy-oriented setting approximately preserves threshold accuracy (57.53% versus 57.60%) while reducing generated-content ECE from 0.2918 to 0.2354. The calibration-oriented setting lowers ECE further to 0.1632 at 56.73% accuracy. Transferring the MBPP posterior to HumanEval also reduces ECE and Brier scores without refitting: generated-content ECE falls from 0.3090 to 0.2016 in the calibration-oriented setting, with 55.08% accuracy versus 54.88% for threshold. Base Entropy retains the highest HumanEval accuracy at 57.32%. On GSM8K, Base Entropy achieves the highest accuracy (78.17%) and lowest generated-content Brier score, while the calibration-oriented BayesER setting achieves the lowest ECE in both scopes. On MATH500, BayesER improves accuracy from 42.40% to 43.53% and reduces ECE and Brier scores. The accuracy–calibration tradeoff will be further elaborated in Section 5.3. 5.2.2 Planning Table 2: Evaluation on planning tasks. Dataset Method Accuracy (%) ↑ ECEgen (×100× 100) ↓ ECEans (×100× 100) ↓ Briergen (×100× 100) ↓ Sudoku Threshold (Fast-dllm) 97.6097.60 3.13 3.13 2.242.24 Base Entropy 97.60(+0.00)97.60\,(+0.00) 2.98¯(+4.79) 2.98\,(+4.79) 2.98¯(+4.79) 2.98\,(+4.79) 2.28¯(−1.54) 2.28\,(-1.54) Info-gain 96.070.50(−1.57)96.07_ 0.50\,(-1.57) 2.540.08(+18.85)2.54_ 0.08\,(+18.85) 2.540.08(+18.85)2.54_ 0.08\,(+18.85) 3.390.37(−51.35)3.39_ 0.37\,(-51.35) BayesER (best accuracy) 97.53¯0.31(−0.07) 97.53_ 0.31\,(-0.07) 3.960.03(−26.52)3.96_ 0.03\,(-26.52) 3.960.03(−26.52)3.96_ 0.03\,(-26.52) 2.400.17(−7.00)2.40_ 0.17\,(-7.00) BayesER (best calibration) 97.400.20(−0.20)97.40_ 0.20\,(-0.20) 2.990.27(+4.47)2.99_ 0.27\,(+4.47) 2.990.27(+4.47)2.99_ 0.27\,(+4.47) 2.390.16(−6.56)2.39_ 0.16\,(-6.56) Countdown Threshold (Fast-dllm) 18.90 66.46 61.54 58.5658.56 Base Entropy 23.90(+26.46)23.90\,(+26.46) 60.07(+9.61)60.07\,(+9.61) 54.67(+11.16)54.67\,(+11.16) 53.21(+9.14)53.21\,(+9.14) Info-gain 41.472.03(+119.42)41.47_ 2.03\,(+119.42) 45.631.77(+31.34)45.63_ 1.77\,(+31.34) 43.561.08(+29.22)43.56_ 1.08\,(+29.22) 42.401.16(+27.60)42.40_ 1.16\,(+27.60) BayesER (best accuracy) 32.13¯0.67(+70.00) 32.13_ 0.67\,(+70.00) 28.98¯0.21(+56.39) 28.98_ 0.21\,(+56.39) 27.920.68(+54.63)27.92_ 0.68\,(+54.63) 26.71¯0.18(+54.39) 26.71_ 0.18\,(+54.39) BayesER (best calibration) 27.800.78(+47.09)27.80_ 0.78\,(+47.09) 25.700.50(+61.33)25.70_ 0.50\,(+61.33) 30.24¯0.81(+50.86) 30.24_ 0.81\,(+50.86) 22.830.08(+61.01)22.83_ 0.08\,(+61.01) Table 2 shows contrasting outcomes on the two planning tasks. Sudoku offers little advantage for Bayesian decoding: threshold and Base Entropy reach 97.60% accuracy, and neither posterior setting improves on threshold’s Brier score. Countdown exhibits the strongest accuracy–calibration trade-off. The accuracy-oriented BayesER setting reaches 32.13% accuracy, exceeding threshold (18.90%) and Base Entropy (23.90%), but falling below Info-gain (41.47%). Its generated-content ECE and Brier score are lower than Info-gain’s: 0.2898 versus 0.4563 and 0.2671 versus 0.4240, respectively. The calibration-oriented setting further reduces generated-content ECE to 0.2570 at 27.80% accuracy. 5.2.3 Molecular generation Table 3: Evaluation on GenMol V2 Scaffold Decoration. Quality Yield measures the percentage of generated attempts satisfying both drug-likeness (QED≥0.6QED≥ 0.6) and synthetic accessibility (SA≤4.0SA≤ 4.0). Method Quality Yield (%) ↑ Frag. Valid (%) ↑ QED ↑ SA ↓ ECE ↓ Native GenMol 11.78±0.6911.78± 0.69 99.33±0.3399.33± 0.33 0.6369±0.01040.6369± 0.0104 4.3932±0.03624.3932± 0.0362 45.08±0.2945.08± 0.29 Base Entropy 14.89¯±0.69 14.89± 0.69 99.89¯±0.19 99.89± 0.19 0.6599±0.00110.6599± 0.0011 4.2809±0.01064.2809± 0.0106 35.40±0.5935.40± 0.59 Info-gain 3.00±0.673.00± 0.67 100.00±0.00100.00± 0.00 0.6704±0.02060.6704± 0.0206 4.1528±0.05954.1528± 0.0595 24.56±1.7124.56± 1.71 DPRM (Bu et al., 2026) 12.44±0.6912.44± 0.69 100.00±0.00100.00± 0.00 0.6619¯±0.0117 0.6619± 0.0117 4.2682¯±0.0742 4.2682± 0.0742 30.62¯±1.33 30.62± 1.33 BayesER (Ours) 15.33±0.3315.33± 0.33 99.67±0.0099.67± 0.00 0.6602±0.00060.6602± 0.0006 4.2900±0.00624.2900± 0.0062 35.36±0.7335.36± 0.73 On GenMol V2 scaffold decoration (Table 3), BayesER achieves the highest mean quality yield at 15.33%, compared with 11.78% for native GenMol and 14.89% for Base Entropy. Its ECE is lower than native GenMol’s (35.36 versus 45.08 on the table’s scale), but nearly identical to Base Entropy’s (35.40). These metrics capture different aspects of generation: ECE concerns chemical validity, whereas quality yield requires both drug-likeness and synthetic accessibility. 5.3 Hyperparameter sensitivity and accuracy–calibration trade-offs Figure 1: Accuracy–calibration trade-off for MBPP with different hyperparameter choices. We additionally evaluate how prior precision λ, the number of uncertain trailing layers LsL_s, and the commitment budget jointly affect the accuracy–calibration tradeoff. In Figure 1, each panel shows five prior precisions (λ∈100,200,300,500,1000λ∈\100,200,300,500,1000\) at four depths. Intermediate precision provides operating points with lower ECE and comparable or higher accuracy. Further weakening the prior can reduce ECE while substantially degrading accuracy. Depth and commitment budget also interact. A single uncertain layer reaches 59.0% at ECE 0.223 for τB=0.85 _B=0.85, λ=300λ=300, and 58.4% at ECE 0.147 for τB=0.90 _B=0.90, λ=200λ=200. These examples show that a shallow posterior can offer a favorable accuracy–calibration trade-off and that additional uncertain layers do not consistently improve it. 5.4 Computational efficiency Figure 2: Dream HumanEval cost versus uncertain depth. Figure 2 shows HumanEval cost on A100 GPUs with S=10S=10 samples and τB=0.90 _B=0.90. Across tested priors, increasing uncertain depth from one to four layers raises evaluation time from 17.4–22.1 to 28.8–52.9 minutes, versus 13.4 minutes for threshold. Peak allocated memory remains within 15.54–15.71 GiB, versus 14.38 GiB for threshold. Deeper posteriors require longer sampled forward passes, while weaker priors can prolong decoding. Shallow posteriors therefore limit runtime overhead. Memory varies little with depth. 6 Conclusion and Discussion We introduced BayesER, a post-hoc Bayesian decoding scheme that uses posterior predictive uncertainty to guide position ordering, token selection, and commitment budgets in masked diffusion language models. Its curvature-only Laplace-LoRA construction preserves the released checkpoint at the posterior mean and requires no adapter fine-tuning. The experiments show that posterior prediction can improve sequence-confidence calibration while preserving or improving task performance at selected operating points. In future work, we plan to extend the single Gaussian posterior to multiple Gaussian posteriors across different adapter-parameter layers to better capture parameter uncertainties and enhance our entropy-based decoding scheme. Reproducibility Statement The pretrained checkpoints used in this work (Dream-v0-Instruct-7B, LLaDA-8B-Instruct, and GenMol V2) and the MBPP, HumanEval, GSM8K, and MATH500 benchmarks are publicly available. Planning evaluations use fixed public Sudoku and Countdown test sets; separate synthetic training sets are generated deterministically with test inputs excluded. Molecular-generation experiments use public GenMol fragments and calibration data drawn from the training split of datamol-io/safe-gpt. Section 3 and Appendix A describe the method, posterior construction, and predictive computation. Section 5 describes the benchmarks and metrics; Tables 4 and 5 specify the Dream configurations, and Appendix D reports the additional LLaDA evaluation settings. Code, posterior artifacts, data-preparation scripts, and evaluation configurations will be open-sourced upon publication. Use of AI Statement Generative AI tools and large language models were used to assist with code implementation, finding related work, manuscript editing, grammar polishing, and typesetting. AI-assisted code was checked by the authors. The authors take responsibility for the final manuscript, including its code, experimental results, and claims. Ethics Statement This work studies uncertainty-aware decoding using pretrained models and public benchmark datasets. Improved benchmark calibration does not guarantee reliable predictions under distribution shift or safe deployment in high-stakes settings. The molecular-generation experiments evaluate computational properties; the reported metrics do not establish biological efficacy, toxicity, or safety. Generated molecules would require independent expert assessment and experimental validation before practical use. Acknowledgments We thank the Digital Research Alliance of Canada (Compute Canada) for its computational resources and services. Z. Jiang and M. Park were supported in part by the Natural Sciences and Engineering Research Council of Canada (NSERC) and the Canada CIFAR AI Chairs program. References Ben-Hamu et al. (2025) Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, and Brian Karrer. Accelerated sampling from masked diffusion models via entropy bounded unmasking. In Advances in Neural Information Processing Systems, 2025. URL https://arxiv.org/abs/2505.24857. Brier (1950) Glenn W. Brier. Verification of forecasts expressed in terms of probability. Monthly Weather Review, 78(1):1–3, 1950. Bu et al. (2026) Dake Bu, Wei Huang, Andi Han, Si Wu, Hau-San Wong, Qingfu Zhang, Taiji Suzuki, and Atsushi Nitanda. DPRM: A plug-in Doob h-transform-induced token-ordering module for discrete diffusion models, 2026. URL https://arxiv.org/abs/2604.24357. Chang et al. (2022) Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Daxberger et al. (2021) Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux — effortless Bayesian deep learning. In Advances in Neural Information Processing Systems, 2021. Guo et al. (2017) Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, p. 1321–1330. PMLR, 2017. URL https://proceedings.mlr.press/v70/guo17a.html. Houlsby et al. (2011) Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning, 2011. URL https://arxiv.org/abs/1112.5745. Hu et al. (2021) Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models, 2021. URL https://arxiv.org/abs/2106.09685. Jazbec et al. (2025) Metod Jazbec, Eliot Wong-Toi, Guoxuan Xia, Dan Zhang, Eric Nalisnick, and Stephan Mandt. Generative uncertainty in diffusion models. arXiv preprint arXiv:2502.20946, 2025. URL https://arxiv.org/abs/2502.20946. Jazbec et al. (2026) Metod Jazbec, Theo X. Olausson, Louis Béthune, Pierre Ablin, Michael Kirchhof, João Monteiro, Victor Turrisi, Jason Ramapuram, and Marco Cuturi. Learning unmasking policies for diffusion language models. In International Conference on Machine Learning, 2026. URL https://arxiv.org/abs/2512.09106. Kou et al. (2023) Siqi Kou, Lei Gan, Dequan Wang, Chongxuan Li, and Zhijie Deng. BayesDiff: Estimating pixel-wise uncertainty in diffusion via Bayesian inference. arXiv preprint arXiv:2310.11142, 2023. URL https://arxiv.org/abs/2310.11142. Lee et al. (2025) Sanghyun Lee, Seungryong Kim, Jongho Park, and Dongmin Park. Lookahead unmasking elicits accurate decoding in diffusion language models, 2025. URL https://arxiv.org/abs/2511.05563. Li et al. (2026) Yansi Li, Gongshen Liu, and Zhuosheng Zhang. The confidence paradox: Unveiling the latent discriminative power of diffusion large language models in mathematical reasoning. In Maria Liakata, Viviane P. Moreira, Jiajun Zhang, and David Jurgens (eds.), Findings of the Association for Computational Linguistics: ACL 2026, p. 43179–43196, San Diego, California, United States, July 2026. Association for Computational Linguistics. ISBN 979-8-89176-395-1. doi: 10.18653/v1/2026.findings-acl.2142. URL https://aclanthology.org/2026.findings-acl.2142/. MacKay (1992) David J. C. MacKay. A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3):448–472, 1992. Martens & Grosse (2015) James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning, p. 2408–2417, 2015. Naeini et al. (2015) Mahdi Pakdaman Naeini, Gregory F. Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the Twenty-Ninth AAAI Conference on Artificial Intelligence, p. 2901–2907, 2015. doi: 10.1609/aaai.v29i1.9602. Nie et al. (2025) Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025. URL https://arxiv.org/abs/2502.09992. Ritter et al. (2018) Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation for neural networks. In International Conference on Learning Representations, 2018. Sahoo et al. (2024) Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems, 2024. Shi et al. (2025) Haizhou Shi, Yibin Wang, Ligong Han, Huan Zhang, and Hao Wang. Training-free Bayesianization for low-rank adapters of large language models. In Advances in Neural Information Processing Systems, 2025. URL https://arxiv.org/abs/2412.05723. Shu et al. (2026) Yingte Shu, Yuchuan Tian, Chao Xu, Yunhe Wang, and Hanting Chen. Deferred commitment decoding for diffusion language models with confidence-aware sliding windows. arXiv preprint arXiv:2601.02076, 2026. URL https://arxiv.org/abs/2601.02076. Wang et al. (2024) Yibin Wang, Haizhou Shi, Ligong Han, Dimitris Metaxas, and Hao Wang. BLoB: Bayesian low-rank adaptation by backpropagation for large language models. In Advances in Neural Information Processing Systems, volume 37, p. 67758–67794, 2024. Wu et al. (2025) Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dLLM: Training-free acceleration of diffusion LLM by enabling KV cache and parallel decoding. arXiv preprint arXiv:2505.22618, 2025. URL https://arxiv.org/abs/2505.22618. Yang et al. (2024) Adam X. Yang, Maxime Robeyns, Xi Wang, and Laurence Aitchison. Bayesian low-rank adaptation for large language models. In International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=FJiUyzOF1m. Yang et al. (2026) Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, and Alex Lamb. Improving sampling for masked diffusion models via information gain, 2026. URL https://arxiv.org/abs/2602.18176. Ye et al. (2025) Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508.15487, 2025. URL https://arxiv.org/abs/2508.15487. Appendix A Posterior and implementation details A.1 Posterior fitting, sampling and computation We use the notation of Section 3.1: 0 θ_0 denotes the fixed pretrained parameters, ϕ φ collects the LoRA adapter parameters, j indexes selected transformer layers, and m indexes individual adapter-matrix blocks. We fit model-Fisher curvature at ϕ0 φ_0, where the LoRA update is zero, using task-matched calibration prompts and suppressing the mask-token logit before softmax. Algorithm 2 summarizes the procedure. For posterior draws we use the exact factorization implied by Eq. 4. Let ϕm φ_m be the vectorized adapter matrix for block m, with mean ϕ0,m φ_0,m; here vecvec stacks rows, consistent with the ordering Gm⊗CmG_m C_m. Draw a matrix zm(s)z_m^(s) of independent standard normal entries in the shape of that adapter matrix and expand it in the retained basis, zm(s)=∑a,bξm,ab(s)um,aG(um,bC)⊤+zm,⟂(s)z_m^(s)= _a,b _m,ab^(s)u^G_m,a(u^C_m,b) +z_m, ^(s). The adapter draw is ϕm(s)=ϕ0,m+vec[∑a,bξm,ab(s)λ+Nmgm,acm,bum,aG(um,bC)⊤+1λzm,⟂(s)], φ_m^(s)= φ_0,m+vec\! [ _a,b _m,ab^(s) λ+N_mg_m,ac_m,b\,u^G_m,a(u^C_m,b) + 1 λz_m, ^(s) ], (8) where cm,bc_m,b and gm,ag_m,a are the retained eigenvalues of CmC_m and GmG_m and NmN_m is the represented source-position observation count defined in the main text. The coefficients ξm,ab(s) _m,ab^(s) are independent standard normal draws. The retained-basis perturbation is the data- and prior-informed component; the last term is the prior-only complement corresponding to truncated directions. Each block draws its own independent zm(s)z_m^(s), realizing the block-independence assumption above. Substituting Eq. 8 into Eq. 2 gives the nonlinear weight-space predictive p0,ϕ(s)(l)(⋅∣t(1:L))p_ θ_0, φ^(s)^(l)(· x_t^(1:L)) at position l. Normalizing the sampled logits on t(l)V_t^(l) gives pt,s(l)p_t,s^(l) from Section 3.2. Algorithm 2 Curvature-only KFAC Laplace-LoRA posterior 0: Released model f0f_ θ_0; calibration set calD_cal; selected projections; LoRA rank r and scale α; KFAC rank k; prior precision λ; per-block row-sampling cap NsN_s 1: Wrap each selected projection as j(ϕj)=j,0+αrjj W_j( φ_j)= W_j,0+ αr B_j A_j with j,0=0 B_j,0=0 (Eq. 2); let ϕ0 φ_0 collect j,0,j,0j\ A_j,0, B_j,0\_j, with projection indices omitted as in the main text 2: Initialize weighted sums SmC=0S^C_m=0, SmG=0S^G_m=0 and counts Nm=0N_m=0 for each block m 3: for each calibration batch of masked sequences ~n x_n and target masks do 4: Forward the masked sequence; suppress the mask-token logit before softmax and sample one categorical target from the model distribution at each masked target position 5: Backpropagate the summed cross-entropy against these sampled targets 6: for each adapter-matrix block m (j A_j and j B_j separately) do 7: Collect paired rows (hm,n,gm,n)(h_m,n,g_m,n) at positions with nonzero gradient; skip blocks with no eligible rows 8: Uniformly subsample up to NsN_s rows; set each retained row’s weight wnw_n to the number of eligible rows divided by the number retained 9: SmC←SmC+∑nwnhm,nhm,n⊤S^C_m← S^C_m+ _nw_nh_m,nh_m,n , SmG←SmG+∑nwngm,ngm,n⊤S^G_m← S^G_m+ _nw_ng_m,ng_m,n , and Nm←Nm+∑nwnN_m← N_m+ _nw_n 10: end for 11: end for 12: for each block m do 13: if Nm=0N_m=0 then 14: Set Fm=0F_m=0 (prior-only block) 15: else 16: Normalize Cm=SmC/NmC_m=S^C_m/N_m and Gm=SmG/NmG_m=S^G_m/N_m 17: Eigendecompose CmC_m and GmG_m; retain the top-k directions of large factors 18: Store (um,bC,cm,b,um,aG,gm,a,Nm)(\u^C_m,b\,\c_m,b\,\u^G_m,a\,\g_m,a\,N_m) with curvature Fm=Nm(Gm⊗Cm)F_m=N_m(G_m C_m) 19: end if 20: end for 21: Form FKFAC=blockdiag(F1,…,FR)F_KFAC=blockdiag(F_1,…,F_R) and qλ(ϕ)=(ϕ0,(FKFAC+λI)−1)q_λ( φ)=N( φ_0,(F_KFAC+λ I)^-1) 22: return qλq_λ and the per-block eigenbasis; λ is applied at inference All reported evaluations are cache-free. Within each decoding step, activations before the first uncertain layer are computed once and reused across posterior draws; each draw evaluates the remaining layers with its sampled weights. This reuse is exact because the preceding layers are unchanged across draws. Activations are recomputed at every decoding step, and draws are batched where memory permits. A.2 Execution protocols All rows in Tables 1 and 2 use exact cache-free forward passes. Planning evaluations use the task geometry described in Section 5. Task-specific decoding geometry. For the Dream experiments, MBPP, HumanEval, GSM8K, and MATH500 use sliding-window decoding with an initial window width of 16 tokens and a maximum width of 128 tokens. Sudoku uses in-place infilling, whereas Countdown uses a 48-token block. These planning-specific settings replace the sliding-window configuration used for code generation and mathematical reasoning. 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. 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=10k=10) Target modules trailing-layer query/value projections Posterior draws S=10S=10 Vocabulary support Kv=256K_v=256 Predictive backend nonlinear posterior weight sampling Commit threshold τ=0.9τ=0.9; posterior budget gate τB=0.85 _B=0.85 or 0.90.9 Ranking posterior predictive entropy Commit budget posterior confidence gate for all selected BayesER rows Window (code/math) Sliding, initial 16 tokens, maximum 128 tokens Geometry (Sudoku) In-place infilling Geometry (Countdown) 48-token block Token proposal posterior argmax for all selected BayesER rows Sequence confidence geometric mean of commit-time probabilities, generated and answer scopes ECE 15 equal-width bins Table 5: Hyperparameters for the evaluation in the main results. HumanEval reuses the MBPP-trained posterior; the remaining rows are task-fitted. Task Target LsL_s Rank Alpha Prior τB _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 MATH500 Accuracy / calibration 2 8 16 2,000 0.90 Sudoku Accuracy 2 8 16 500 0.90 Sudoku Calibration 1 8 16 1,000 0.90 Countdown Accuracy 4 8 16 200 0.85 Countdown Calibration 1 8 16 100 0.90 Appendix B Additional Dream accuracy–calibration sweeps Figures 3–6 extend Figure 1 to four additional benchmarks. Figure 3: HumanEval accuracy versus calibration. Figure 4: GSM8K accuracy versus calibration. Figure 5: MATH500 accuracy versus calibration. Figure 6: Countdown accuracy versus calibration. Appendix C Additional measured cost Figure 7 compares BayesER costs with Confidence-thresholding baseline on H100 GPU. Figure 7: Measured Dream GSM8K cost versus uncertain depth. Full test set (n=1,319n=1,319), GSM8K-fitted rank-8 q/v Laplace-LoRA with S=10S=10, on H100 GPUs. Color denotes prior precision λ. Dashed lines show the confidence-threshold run. Wall time includes loading and scoring; the loose prior increases decoding steps, especially at four layers. Appendix D LLaDA results Table 6 reports full-test results for GSAI-ML/LLaDA-8B-Instruct. The recorded MBPP, GSM8K and MATH500 curvature-only sweeps use rank-8 q/v adapters (α=16α=16), KFAC rank 10, and vocabulary support Kv=256K_v=256. HumanEval retains the MBPP transfer configuration described above. Info-gain uses N=8N=8, fixed K=2K=2. 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. Method Accuracy (%) ↑ ECEgen (×100× 100) ↓ ECEans (×100× 100) ↓ Briergen (×100× 100) ↓ MBPP (500) Threshold (Fast-dllm) 39.4039.40 47.6347.63 50.0150.01 46.1446.14 Base Entropy 41.20(+4.57)41.20\,(+4.57) 44.51¯(+6.53) 44.51\,(+6.53) 46.84¯(+6.33) 46.84\,(+6.33) 43.48¯(+5.76) 43.48\,(+5.76) Info-gain 41.27¯0.61(+4.74) 41.27_ 0.61\,(+4.74) 44.960.50(+5.59)44.96_ 0.50\,(+5.59) 48.010.52(+4.00)48.01_ 0.52\,(+4.00) 43.520.26(+5.67)43.52_ 0.26\,(+5.67) BayesER 41.330.58(+4.91)41.33_ 0.58\,(+4.91) 43.570.67(+8.52)43.57_ 0.67\,(+8.52) 45.890.66(+8.23)45.89_ 0.66\,(+8.23) 42.700.45(+7.46)42.70_ 0.45\,(+7.46) HumanEval (164) Threshold (Fast-dllm) 18.2918.29 69.0569.05 66.5066.50 62.2662.26 Base Entropy 22.56¯(+23.33) 22.56\,(+23.33) 63.34¯(+8.26) 63.34\,(+8.26) 60.48¯(+9.04) 60.48\,(+9.04) 57.46¯(+7.70) 57.46\,(+7.70) Info-gain 15.652.31(−14.44)15.65_ 2.31\,(-14.44) 70.512.11(−2.11)70.51_ 2.11\,(-2.11) 68.381.89(−2.84)68.38_ 1.89\,(-2.84) 62.951.39(−1.11)62.95_ 1.39\,(-1.11) BayesER 23.170.00(+26.67)23.17_ 0.00\,(+26.67) 62.760.02(+9.11)62.76_ 0.02\,(+9.11) 59.810.03(+10.06)59.81_ 0.03\,(+10.06) 56.990.07(+8.46)56.99_ 0.07\,(+8.46) GSM8K (1,319) Threshold (Fast-dllm) 76.12¯ 76.12 17.4917.49 20.59¯ 20.59 20.6320.63 Base Entropy 75.51(−0.80)75.51\,(-0.80) 17.11(+2.21)17.11\,(+2.21) 20.98(−1.89)20.98\,(-1.89) 20.64(−0.05)20.64\,(-0.05) Info-gain 78.090.53(+2.59)78.09_ 0.53\,(+2.59) 15.870.50(+9.27)15.87_ 0.50\,(+9.27) 20.210.58(+1.87)20.21_ 0.58\,(+1.87) 18.880.45(+8.53)18.88_ 0.45\,(+8.53) BayesER 75.920.37(−0.27)75.92_ 0.37\,(-0.27) 16.67¯0.37(+4.68) 16.67_ 0.37\,(+4.68) 20.640.40(−0.21)20.64_ 0.40\,(-0.21) 20.24¯0.30(+1.89) 20.24_ 0.30\,(+1.89) MATH500 (500) Threshold (Fast-dllm) 38.00¯ 38.00 53.68¯ 53.68 55.23¯ 55.23 51.4451.44 Base Entropy 38.20(+0.53)38.20\,(+0.53) 52.13(+2.90)52.13\,(+2.90) 55.01(+0.41)55.01\,(+0.41) 49.67(+3.43)49.67\,(+3.43) Info-gain 37.730.42(−0.70)37.73_ 0.42\,(-0.70) 53.810.43(−0.23)53.81_ 0.43\,(-0.23) 57.350.60(−3.84)57.35_ 0.60\,(-3.84) 51.08¯0.41(+0.69) 51.08_ 0.41\,(+0.69) BayesER 35.400.72(−6.84)35.40_ 0.72\,(-6.84) 54.600.75(−1.71)54.60_ 0.75\,(-1.71) 58.350.57(−5.64)58.35_ 0.57\,(-5.64) 51.480.64(−0.08)51.48_ 0.64\,(-0.08) Sudoku (500) Threshold (Fast-dllm) 85.6085.60 9.579.57 9.579.57 10.42¯ 10.42 Base Entropy 81.80(−4.44)81.80\,(-4.44) 12.85(−34.35)12.85\,(-34.35) 12.85(−34.35)12.85\,(-34.35) 11.99(−15.05)11.99\,(-15.05) Info-gain 78.271.10(−8.57)78.27_ 1.10\,(-8.57) 10.101.14(−5.54)10.10_ 1.14\,(-5.54) 10.101.14(−5.54)10.10_ 1.14\,(-5.54) 14.850.66(−42.48)14.85_ 0.66\,(-42.48) BayesER 85.47¯1.47(−0.16) 85.47_ 1.47\,(-0.16) 9.81¯0.24(−2.49) 9.81_ 0.24\,(-2.49) 9.81¯0.24(−2.49) 9.81_ 0.24\,(-2.49) 10.260.82(+1.57)10.26_ 0.82\,(+1.57) Countdown (1,000) Threshold (Fast-dllm) 35.10¯ 35.10 53.65¯ 53.65 52.1452.14 48.3948.39 Base Entropy 29.70(−15.38)29.70\,(-15.38) 54.08(−0.80)54.08\,(-0.80) 50.65¯(+2.85) 50.65\,(+2.85) 47.92(+0.97)47.92\,(+0.97) Info-gain 39.571.55(+12.73)39.57_ 1.55\,(+12.73) 49.281.23(+8.14)49.28_ 1.23\,(+8.14) 46.991.05(+9.88)46.99_ 1.05\,(+9.88) 45.381.04(+6.22)45.38_ 1.04\,(+6.22) BayesER 28.830.38(−17.85)28.83_ 0.38\,(-17.85) 54.070.33(−0.78)54.07_ 0.33\,(-0.78) 51.020.36(+2.14)51.02_ 0.36\,(+2.14) 47.68¯0.20(+1.46) 47.68_ 0.20\,(+1.46) Appendix E Mutual-information ranking and Bayesian entropy decoding We compare MIRank with BayesER on the first 128 examples of MBPP and GSM8K. The grid uses Ls∈1,2,4L_s∈\1,2,4\, λ∈100,200,500,1000,2000λ∈\100,200,500,1000,2000\, and τB∈0.85,0.90 _B∈\0.85,0.90\ for BayesER. 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, τB=0.85 _B=0.85 BayesER, τB=0.90 _B=0.90 LsL_s λ Acc. ECEg ECEa Steps Acc. ECEg ECEa Steps Acc. ECEg ECEa 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 2 100 28.13 16.87 17.83 20.91 9.38 14.20 16.24 123.17 5.47 18.04 22.09 123.62 2 200 45.31 15.05 16.11 33.48 55.47 13.03 14.74 72.17 53.13 15.10 17.64 76.32 2 500 52.34 29.39 32.08 42.05 60.94 22.64 24.62 42.09 56.25 27.86 29.85 46.04 2 1000 54.69 29.13 32.01 39.08 57.81 27.10 29.00 37.65 57.03 28.62 30.97 43.27 2 2000 53.13 30.67 32.96 40.84 59.38 25.93 27.90 39.11 57.81 28.49 30.61 42.89 4 100 28.91 26.44 27.72 21.88 0.00 38.72 38.58 224.40 0.00 37.81 37.89 222.68 4 200 42.19 12.85 11.29 27.82 46.09 5.04 10.20 112.18 46.09 6.27 10.23 119.48 4 500 53.13 26.61 28.95 40.52 58.59 23.36 25.51 46.13 56.25 26.04 28.29 48.66 4 1000 58.59 24.67 27.11 41.99 59.38 25.22 27.17 39.10 57.81 27.50 29.84 44.98 4 2000 57.81 26.09 28.48 40.41 58.59 26.87 28.74 37.05 58.59 27.65 29.80 42.36 Table 8: Comparison between MIRank and BayesER on 128 examples of GSM8K. Accuracy is in percent; both ECE columns are multiplied by 100. MIRank BayesER, τB=0.85 _B=0.85 BayesER, τB=0.90 _B=0.90 LsL_s λ Acc. ECEg ECEa Steps Acc. ECEg ECEa Steps Acc. ECEg ECEa 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 2 100 44.53 25.79 12.27 27.08 67.19 37.97 12.24 110.55 64.84 38.51 14.77 122.64 2 200 64.06 5.68 17.23 37.99 73.44 8.06 12.95 76.51 75.00 5.96 18.17 88.73 2 500 75.78 6.84 16.15 49.79 80.47 9.69 14.06 57.14 82.03 4.96 14.46 65.91 2 1000 78.13 6.43 16.87 52.21 79.69 11.50 15.85 53.67 81.25 4.04 15.12 61.98 2 2000 72.66 10.91 22.94 54.52 79.69 10.08 14.79 54.41 80.47 5.94 14.97 60.62 4 100 48.44 44.56 33.95 26.02 17.19 25.23 31.09 206.01 19.53 28.30 23.09 206.45 4 200 63.28 25.31 15.01 34.52 67.97 19.21 17.75 103.86 72.66 22.09 12.85 109.77 4 500 71.88 6.91 20.19 52.28 76.56 10.89 17.49 60.30 79.69 9.86 15.90 69.25 4 1000 79.69 10.22 14.45 52.49 79.69 7.53 16.26 55.61 79.69 6.29 17.09 62.16 4 2000 78.13 5.62 17.96 54.41 75.78 10.38 17.55 54.13 78.91 7.69 16.13 62.65