PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models

Zhejun Jiang, Mijung Park

arXiv:2610.05125v1cs.LGcs.AI

TL;DR

Masked diffusion language models must decide which masked tokens to commit, but common confidence-based ordering can be miscalibrated. BayesER uses a checkpoint-centered Bayesian posterior and predictive entropy to guide decoding. Across several benchmarks it reduces calibration error while preserving or improving task accuracy at selected operating points.

  • Problem

    In MDLMs, token commitment ordering affects generation quality and inference cost, while confidence-based commitment can rely on a single model estimate despite possible miscalibration.

    Source: §1 · p. 1

  • Method

    BayesER uses posterior predictive entropy to order masked positions, posterior predictive means to select tokens, and posterior confidence to determine commitment counts, with a curvature-only posterior at the released checkpoint.

    Source: §1 · p. 2

  • Results

    Across code, mathematical reasoning, planning, and molecule-generation benchmarks, BAYESER reduces expected calibration error while preserving or improving task accuracy relative to common confidence-based decoding schemes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    At selected operating points, posterior prediction can improve sequence-confidence calibration while preserving or improving task performance.

    Source: §6 · p. 9

    Improved benchmark calibration does not guarantee reliability under distribution shift or safe high-stakes deployment, and molecular-generation metrics do not establish biological efficacy, toxicity, or safety.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

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

In MDLMs, confidence affects both whether predictions can be trusted and which masked positions are committed at each denoising step. BayesER uses posterior predictive uncertainty to guide commitment, with experiments reporting lower calibration error while preserving or improving accuracy across benchmarks.

Motivation

  • MDLM decoding must choose which masked positions to commit because committing many in parallel reduces model evaluations, while committing too conservatively reduces parallel-decoding benefits.Committed tokens become context for later denoising steps, so commitment ordering affects both generation quality and inference cost; common rules rely on single-checkpoint softmax confidence.Source: §1 · pp. 1–2

Method

  • BAYESER averages predictions over an approximate parameter posterior, then uses predictive entropy to order masked positions, posterior predictions to select tokens, and posterior confidence to set how many positions to commit.This adjusts both the order and pace of decoding.Source: §1 · p. 2
  • The posterior is centered at the released checkpoint and uses curvature estimation without adapter training, leaving the deterministic model unchanged at the posterior mean.Kronecker products of input and output-gradient second moments represent Hessian blocks to reduce curvature storage and expression costs.Source: §1 · p. 2

Results

  • Across code, math, planning, and molecule-generation benchmarks, BAYESER reduces expected calibration error while preserving or improving accuracy against common confidence-based decoders.The posterior’s prior precision and number of Laplace-LoRA layers govern trade-offs among calibration, task performance, and decoding cost; a posterior fitted on one code dataset also reduces calibration error on another without refitting.Source: §1 · p. 2

2 BACKGROUND

The background introduces masked diffusion decoding, Laplace-LoRA, and calibration metrics used in the paper. It describes iterative masked-token prediction, confidence-based commitment, and binned measures of confidence–accuracy agreement.

2 BACKGROUND

  • The background section introduces decoding, Laplace-LoRA, and expected calibration error as the concepts needed for the paper’s method and calibration evaluation.The supplied description identifies these topics but gives no additional substantive claim about them.Source: §2 · p. 2

2.1 MASKED DIFFUSION DECODING

  • MDLMs repeatedly predict and reveal masked tokens; confidence-based decoding selects the most confident base tokens or those meeting a threshold at each denoising step.A minimum of one token is ensured to be unmasked at a time.Source: §2.1 · pp. 2–3

2.2 LAPLACE-LORA

  • Laplace-LoRA places a Gaussian posterior over downstream-optimized LoRA adapters, while BAYESER uses Fisher and Kronecker-factored curvature approximations without adapter fine-tuning at a fixed checkpoint.LoRA parameterizes transformer weight updates with two low-rank matrices.Source: §2.2 · p. 3

2.3 CALIBRATION AND SEQUENCE-LEVEL CONFIDENCE

  • Expected calibration error groups predictions into confidence bins and measures the gap between accuracy and confidence; the paper also uses the Brier score for binary task-correctness evaluation.For the Brier score, sequence confidence is compared with a binary indicator of task correctness.Source: §2.3 · p. 3

3 METHOD: BAYESER

BAYESER uses a checkpoint-centered posterior to guide MDLM token commitments, and evaluates sequence confidence using commitment-time probabilities. Its method contrasts total predictive entropy with alternative ordering signals.

3 METHOD: BAYESER

  • BAYESER makes position, token, and commitment-budget decisions from a checkpoint-centered posterior predictive distribution.It ranks positions by predictive entropy, selects token identities from the predictive mean, and sets the commitment budget from posterior confidence.Source: §3 · p. 3

3.1 PHASE I: CURVATURE-ONLY LAPLACE-LORA POSTERIOR

  • Unlike standard Laplace-LoRA, BAYESER does not optimize adapter means downstream; training data estimates local curvature around the released checkpoint.The posterior uses KFAC to approximate Fisher information for adapter matrices, with blocks independent across matrices, projections, and layers.Source: §3.1 · p. 4

3.2 PHASE II: POSTERIOR PREDICTIVE ENTROPY FOR POSITION ORDERING

  • BAYESER averages probabilities across posterior draws over each position’s top-Kv non-masked candidates, then uses predictive entropy to order positions.Low entropy marks a position as comparatively safe to commit; high entropy defers it under the current partially masked context.Source: §3.2 · p. 5

3.3 EVALUATING SEQUENCE-LEVEL CONFIDENCE

  • Sequence confidence is the geometric mean of token probabilities recorded at commitment, evaluated against task correctness using generated-content or final-answer tokens.The generated-content and final-answer ECE variants use different token scopes but share task-level correctness; low generated-content ECE does not establish that every reasoning sentence is correct.Source: §3.3 · p. 6

3.4 MUTUAL INFORMATION AND BASE ENTROPY

  • The method compares total predictive entropy with mutual information and base entropy; mutual information can miss ambiguity shared across posterior draws.Base entropy is included to distinguish entropy-based ordering from Bayesian marginalization.Source: §3.4 · p. 6

4 RELATED WORK

Prior MDLM decoding approaches use confidence rules, entropy, successor-state evaluation, learned policies, or reward-based ordering. BAYESER instead applies posterior predictive uncertainty to unresolved token positions while retaining the released checkpoint as the posterior mean.

4 RELATED WORK

5 EXPERIMENTS

Across code, mathematical reasoning, planning, and molecular-generation benchmarks, BAYESER’s calibration and accuracy outcomes vary by task and operating point. The experiments also examine hyperparameter trade-offs and the runtime and memory costs of posterior depth.

5.1 SETUP

  • The evaluation uses Dream-v0-Instruct-7B across code generation, mathematical reasoning, and planning, transfers the MBPP posterior to HumanEval without refitting, and additionally tests LLaDA and GenMol V2.Comparisons include confidence-threshold decoding, Base Entropy, and Info-gain; sequence-level calibration uses ECE and binary Brier scores, with molecule ECE assessed against chemical validity.Source: §5.1 · p. 7

5.2.1 MATHEMATICAL REASONING AND PROGRAMMING

  • 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 without refitting lowers generated-content ECE from 0.3090 to 0.2016 at 55.08% accuracy versus 54.88% for threshold.Source: §5.2.1 · p. 8
  • On GSM8K, Base Entropy has the highest accuracy (78.17%) and lowest generated-content Brier score, while calibration-oriented BAYESER has the lowest ECE in both scopes.On MATH500, BAYESER improves accuracy from 42.40% to 43.53% and reduces ECE and Brier scores.Source: §5.2.1 · p. 8

5.2.2 PLANNING

  • Planning results differ by task: on Sudoku, threshold and Base Entropy reach 97.60% accuracy and neither posterior setting improves on threshold’s Brier score.On Countdown, accuracy-oriented BAYESER reaches 32.13% accuracy versus 18.90% for threshold and 23.90% for Base Entropy, below Info-gain’s 41.47%; its generated-content ECE and Brier score are 0.2898 and 0.2671, versus 0.4563 and 0.4240 for Info-gain.Source: §5.2.2 · p. 8, Table 2 · p. 8

5.2.3 MOLECULAR GENERATION

  • BAYESER’s benchmark outcomes are task-dependent: it improves calibration with varying accuracy effects, and its molecular-generation setting achieves the highest mean quality yield.On GenMol V2 scaffold decoration, mean quality yield is 15.33%, versus 11.78% for native GenMol and 14.89% for Base Entropy.Source: §5.2.3 · p. 8, §5.2.1 · p. 8
  • On GenMol V2 scaffold decoration, BAYESER’s ECE is 35.36 versus native GenMol’s 45.08 and Base Entropy’s 35.40, while quality yield requires both drug-likeness and synthetic accessibility.ECE is evaluated against chemical validity; Table 3 defines quality yield using QED ≥0.6 and SA ≤4.0.Source: §5.2.3 · p. 8, Table 3 · p. 8

5.3 HYPERPARAMETER SENSITIVITY AND ACCURACY–CALIBRATION TRADE-OFFS

  • Intermediate prior precision provides MBPP operating points with lower ECE and comparable or higher accuracy, while further weakening the prior can reduce ECE but substantially degrade accuracy.A single uncertain layer reaches 59.0% accuracy at ECE 0.223 for τB = 0.85, λ = 300, and 58.4% at ECE 0.147 for τB = 0.90, λ = 200; additional uncertain layers do not consistently improve the trade-off.Source: §5.3 · p. 9, Figure 1 · p. 9

5.4 COMPUTATIONAL EFFICIENCY

  • On HumanEval with S = 10 and τB = 0.90, increasing uncertain depth from one to four layers raises evaluation time from 17.4–22.1 to 28.8–52.9 minutes, while peak memory remains within 15.54–15.71 GiB.Threshold decoding takes 13.4 minutes and uses 14.38 GiB; memory varies little with depth.Source: §5.4 · p. 9, Figure 2 · p. 9

6 CONCLUSION AND DISCUSSION

The paper presents BAYESER as post-hoc Bayesian decoding for masked diffusion language models, using posterior predictive uncertainty to guide decoding decisions. It reports calibration improvements at selected operating points and describes the study’s scope and reproducibility materials.

6 CONCLUSION AND DISCUSSION

  • BAYESER 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; future work proposes multiple Gaussian posteriors across adapter-parameter layers.Source: §6 · p. 9

REPRODUCIBILITY STATEMENT

  • The work identifies its pretrained checkpoints and benchmarks as public and says code, posterior artifacts, data-preparation scripts, and evaluation configurations will be open-sourced upon publication.The reproducibility statement also describes fixed public planning test sets and the molecular calibration-data source.Source: REPRODUCIBILITY STATEMENT · p. 10

USE OF AI STATEMENT

  • AI tools assisted with implementation, related-work discovery, manuscript editing, grammar polishing, and typesetting; the authors say they checked AI-assisted code and take responsibility for the final work..Source: USE OF AI STATEMENT · p. 10

ETHICS STATEMENT

  • The study uses pretrained models and public benchmarks; improved benchmark calibration does not establish reliability under distribution shift or safety in high-stakes deployment.Molecular metrics do not establish biological efficacy, toxicity, or safety, and generated molecules require independent expert assessment and experimental validation before practical use.Source: ETHICS STATEMENT · p. 10

A.1 POSTERIOR FITTING, SAMPLING AND COMPUTATION

The appendix details how BayesER fits a curvature-only posterior around zero-update LoRA adapters, samples adapter weights, and computes predictions efficiently across denoising steps.

A.1 POSTERIOR FITTING, SAMPLING AND COMPUTATION

  • Curvature is fit at zero-update adapters using task-matched calibration prompts, with the mask-token logit suppressed before softmax.The procedure wraps selected projections with zero-initialized LoRA updates and samples categorical targets at masked positions before accumulating gradients.Source: §A.1 · p. 13
  • Posterior draws combine data- and prior-informed retained-basis perturbations with a prior-only complement, with independent draws across blocks.The draws use the exact factorization implied by Eq. 4 and feed sampled weights into the nonlinear predictive computation.Source: §A.1 · p. 13
  • Per-block curvature uses weighted sampled rows, retained KFAC eigen-directions, and a Gaussian posterior centered at the zero-update adapter with prior precision λ.Blocks without eligible rows are prior-only; the block-diagonal KFAC curvature and λ define qλ(ϕ) = N(ϕ0, (FKFAC + λI)^−1).Source: §A.1 · p. 13
  • Cache-free decoding reuses activations before the first uncertain layer across posterior draws, but recomputes activations at every decoding step.Draws are batched where memory permits; reuse is exact because the preceding layers are unchanged across draws.Source: §A.1 · p. 14

A.2 EXECUTION PROTOCOLS

The appendix specifies the Dream evaluation configurations, adds accuracy–calibration sweeps on four benchmarks, and reports measured decoding costs on H100 GPUs.

A.2 EXECUTION PROTOCOLS

  • Dream evaluations use cache-free forward passes; code and math use sliding windows, while Sudoku uses in-place infilling and Countdown uses 48-token blocks.HumanEval reuses the MBPP-trained posterior without refitting; the other Dream posteriors are fitted on task-specific or synthetic training data.Source: §A.2 · p. 14, Table 4 · p. 14

B ADDITIONAL DREAM ACCURACY–CALIBRATION SWEEPS

C ADDITIONAL MEASURED COST

  • On H100 GPUs, the measured GSM8K cost comparison varies prior precision and uncertain depth; the loose prior increases decoding steps, especially at four layers.Figure 7 uses the full test set (n = 1,319), a GSM8K-fitted rank-8 q/v Laplace-LoRA posterior with S = 10, and compares against confidence-threshold decoding.Source: Figure 7 · p. 17, C ADDITIONAL MEASURED COST · p. 17

D LLADA RESULTS

The appendix reports full-test LLaDA results and specifies the evaluation setup and metric conventions for those comparisons.

D LLADA RESULTS

  • The LLaDA full-test comparisons cover MBPP, GSM8K, and MATH500 with rank-8 q/v adapters, α = 16, KFAC rank 10, and vocabulary support Kv = 256.HumanEval retains the MBPP transfer configuration; Info-gain uses N = 8 and fixed K = 2.Source: D LLADA RESULTS · p. 18
  • Table 6 reports accuracy, ECE, and generated-content Brier values, with standard deviations across three seeds and relative improvements measured against confidence-threshold decoding.Accuracy is in percent; ECE and Brier values and their standard deviations are multiplied by 100, and bold or underlined values mark best or second-best results per metric within each benchmark, including ties.Source: Table 6 · p. 18

E MUTUAL-INFORMATION RANKING AND BAYESIAN ENTROPY DECODING

The appendix compares MIRANK and BAYESER on small MBPP and GSM8K subsets and specifies the tested BAYESER hyperparameter grid.

E MUTUAL-INFORMATION RANKING AND BAYESIAN ENTROPY DECODING

Loading 2610.05125v1…