Research map / Large language models / Language, safety and interpretability
Uncertainty-Aware Decoding with Minimum Bayes Risk
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Generalizes MBR decoding by taking the expectation over a posterior of model weights (via IVON unimodal posteriors and deep ensembles), yielding sequence- and token-level uncertainty-aware decoding that improves translation, summarization, and data-to-text quality and reduces hallucinations.
Problem
Standard MBR decoding assumes a single model distribution is a good substitute for the target distribution and ignores weight (parameter) uncertainty, which is linked to hallucinations and poor predictions in language generation.
Approach
Replace the model distribution in the MBR expected-risk with a predictive posterior averaged over model parameters, estimated either by sequence-level combination (averaging utilities across sampled models' hypotheses, usable with black-box LLMs) or token-level averaging of next-token probabilities, and use the resulting expected utility also for selective prediction.
Key results
- On finetuned Gemma-2B, sequence-level Deep Ensemble MBR raised IWSLT17 En-De BLEU from 19.93 (MBR AdamW) to 21.24 and COMET from 76.62 to 77.94 at matched MBR comparisons.
- On from-scratch WMT14 En-De with sampling, token-level Deep Ensemble reached 23.95 BLEU vs 23.37 for MBR@Mean at matched compute; beam search with token-level Deep Ensemble reached 28.98 BLEU and 76.08 COMET vs 27.56/75.23.
- Unimodal IVON posteriors gave improvements at no training overhead for finetuning, performing on par with Deep Ensembles, while unimodal posteriors did not consistently beat the single model baseline when training from scratch.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add sequence-level uncertainty-aware decoding with a Deep Ensemble | Gemma-2B-it · IWSLT17 En-De | BLEU: 19.93 -> 21.24 | Helped |
| add sequence-level uncertainty-aware decoding with a Deep Ensemble | Gemma-2B-it · WMT18 Tr-En | BLEU: 14.75 -> 15.69 | Helped |
| add sequence-level uncertainty-aware decoding with a Deep Ensemble | Gemma-2B-it · XSUM | FactCC: 27.50 -> 32.86 | Helped |
| add sequence-level uncertainty-aware decoding with a Deep Ensemble | Gemma-2B-it · STS-B | RMSE: 0.330 -> 0.269 | Helped |
| add sequence-level Deep Ensemble combination | Transformer big · WMT14 En-De, sampling | BLEU: 23.37 -> 24.70 | Helped |
| add sequence-level Deep Ensemble combination | Transformer big · WMT14 En-De, beam search | BLEU: 27.56 -> 28.99 | Helped |
| average token probabilities across a Deep Ensemble | Transformer base · IWSLT14 De-En, beam search | BLEU: 35.90 -> 38.56 | Helped |
| replace the Deep Ensemble with a unimodal posterior (ablation) | Transformer base · IWSLT14 De-En, sampling | BLEU: 36.03 -> 34.59 | Hurt |
| replace the Deep Ensemble with a unimodal posterior (ablation) | Transformer base · IWSLT14 De-En, beam search | BLEU: 38.30 -> 35.78 | Hurt |
| use a unimodal token-level posterior instead of a Deep Ensemble (ablation) | Transformer base · IWSLT14 De-En, beam search | BLEU: 38.56 -> 35.94 | Hurt |
| use a unimodal sequence-level posterior instead of a Deep Ensemble (ablation) | Gemma-2B-it · IWSLT17 En-De | BLEU: 21.24 -> 20.89 | Hurt |
| replace token-level combination with sequence-level output union (ablation) | Transformer base · IWSLT14 De-En, beam search | BLEU: 38.56 vs 38.30 | Mixed |
Surprising findings
- Unimodal posteriors did not consistently outperform the single-model baseline when training from scratch, whereas Deep Ensembles gave significant gains.
- For ancestral sampling, sequence-level combination of a unimodal posterior only improved up to 4 models, and token-level combination was unreliable for beam-search selective prediction.
Experimental setup
Datasets: WMT14 IWSLT14 afroMT IWSLT17 WMT18 WMT19 XSUM SAMSum E2E-NLG STS-B
Benchmarks: IWSLT17 WMT18 XSUM SAMSum E2E NLG STS-B WMT14 IWSLT14 afroMT WMT19
Metrics: BLEU COMET-22 LaBSE chrF ROUGE FactCC RMSE
Baselines: MBR (AdamW) MBR@Mean single model MBR ancestral sampling beam search
Builds on: minimum bayes risk decoding IVON LoRA Deep Ensembles
Closest papers by meaning
- Model-Based Minimum Bayes Risk Decoding for Text Generation (ICML 2024 · similarity 0.93)
- Graph-based Uncertainty Metrics for Long-form Language Model Generations (NeurIPS 2024 · similarity 0.88)
- Dynamically Structured Diffusion Language Model Decoding via Bayesian Inference (NeurIPS 2026 · similarity 0.88)
- MBR and QE Finetuning: Training-time Distillation of the Best and Most Expensive Decoding Methods (ICLR 2024 · similarity 0.87)
- Efficient Minimum Bayes Risk Decoding using Low-Rank Matrix Completion Algorithms (NeurIPS 2024 · similarity 0.87)
- Decomposing Uncertainty for Large Language Models through Input Clarification Ensembling (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.