Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models
cs.LG, cs.AI
Submitted: 2026-10-04
Updated: 2026-10-04
License: http://creativecommons.org/licenses/by/4.0/
The gist: Masked Diffusion Language Models (MDLMs) generate sequences by iteratively replacing masked tokens with model predictions.
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.
Sources
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
- DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Discrete Diffusion Models
- Bayesian Active Learning for Classification and Preference Learning
- LoRA: Low-Rank Adaptation of Large Language Models
- Generative Uncertainty in Diffusion Models
- BayesDiff: Estimating Pixel-wise Uncertainty in Diffusion via Bayesian Inference
- Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models
- Large Language Diffusion Models
- Training-Free Bayesianization for Low-Rank Adapters of Large Language Models
- Deferred Commitment Decoding for Diffusion Language Models
- Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
- Improving Sampling for Masked Diffusion Models via Information Gain
- Dream 7B: Diffusion Large Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks