Bayesian uncertainty improves token selection in masked diffusion models
Bayesian Entropy-based Reordering for Calibrated Diffusion Language Models
Machine LearningArtificial Intelligence
Summary
Choosing the right words step-by-step in some AI language models can be tricky because the confidence levels they use to pick words aren’t always trustworthy. The authors introduce a new method that uses a kind of uncertainty calculation—based on Bayesian ideas—to better decide which words to finalize first. This method helps the model make more reliable choices without losing accuracy and can even work well on different tasks without retraining. Overall, it helps make generated sequences more trustworthy and accurate across various areas like coding and planning.
What this means in practice
- •For software developers: Improve code generation tools by using Bayesian uncertainty to select more reliable tokens, reducing errors in generated source code.
- •For drug discovery teams: Enhance molecular design models by improving sequence confidence calibration, leading to more trustworthy molecule generation.
Authors
Zhejun Jiang, Mijung Park
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.