# Task-Aware Calibration: Provably Optimal Decoding in LLMs

arXiv: https://arxiv.org/abs/2605.10202 · HTML: https://mlsanity.com/paper/2605.10202/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 41 of its 51 references that are in the map.
- First posted on arXiv 11 May 2026.

## Builds on (most cited first, 30 of 41)

- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [A Survey on LLM-as-a-Judge](https://mlsanity.com/paper/2411.15594/) (2024)
- [Hierarchical Neural Story Generation](https://mlsanity.com/paper/1805.04833/) (2018)
- [SemEval-2017 Task 1: Semantic Textual Similarity - Multilingual and Cross-lingual Focused Evaluation](https://mlsanity.com/paper/1708.00055/) (2017)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [Bias and Fairness in Large Language Models: A Survey](https://mlsanity.com/paper/2309.00770/) (2023)
- [Contrastive Decoding: Open-ended Text Generation as Optimization](https://mlsanity.com/paper/2210.15097/) (2022)
- [Measuring Calibration in Deep Learning](https://mlsanity.com/paper/1904.01685/) (2019)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration](https://mlsanity.com/paper/1910.12656/) (2019)
- [Judging the Judges: A Systematic Study of Position Bias in LLM-as-a-Judge](https://mlsanity.com/paper/2406.07791/) (2024)
- [Beam Search Strategies for Neural Machine Translation](https://mlsanity.com/paper/1702.01806/) (2017)
- [Evaluating model calibration in classification](https://mlsanity.com/paper/1902.06977/) (2019)
- [A Thorough Examination of Decoding Methods in the Era of LLMs](https://mlsanity.com/paper/2402.06925/) (2024)
- [SteerLM: Attribute Conditioned SFT as an (User-Steerable) Alternative to RLHF](https://mlsanity.com/paper/2310.05344/) (2023)
- [Calibrating Predictions to Decisions: A Novel Approach to Multi-Class Calibration](https://mlsanity.com/paper/2107.05719/) (2021)
- [Outcome Indistinguishability](https://mlsanity.com/paper/2011.13426/) (2020)
- [Sampling-Based Approximations to Minimum Bayes Risk Decoding for Neural Machine Translation](https://mlsanity.com/paper/2108.04718/) (2021)
- [Linguistic Calibration of Long-Form Generations](https://mlsanity.com/paper/2404.00474/) (2024)
- [Calibrating Language Models with Adaptive Temperature Scaling](https://mlsanity.com/paper/2409.19817/) (2024)
- [Understanding the Properties of Minimum Bayes Risk Decoding in Neural Machine Translation](https://mlsanity.com/paper/2105.08504/) (2021)
- [SimpleQA Verified: A Reliable Factuality Benchmark to Measure Parametric Knowledge](https://mlsanity.com/paper/2509.07968/) (2025)
- [A Consistent and Differentiable Lp Canonical Calibration Error Estimator](https://mlsanity.com/paper/2210.07810/) (2022)
- [When2Call: When (not) to Call Tools](https://mlsanity.com/paper/2504.18851/) (2025)
- [Mind the Confidence Gap: Overconfidence, Calibration, and Distractor Effects in Large Language Models](https://mlsanity.com/paper/2502.11028/) (2025)
- [Calibrating Long-form Generations from Large Language Models](https://mlsanity.com/paper/2402.06544/) (2024)
- [How to Correctly Report LLM-as-a-Judge Evaluations](https://mlsanity.com/paper/2511.21140/) (2025)

## Abstract (arXiv)

LLM decoding often relies on the model's predictive distribution to generate an output. Consequently, misalignment with respect to the true generating distribution leads to suboptimal decisions in practice. While a natural solution is to calibrate the model's output distribution, for LLMs, this is ill-posed at the combinatorially vast level of free-form language. We address this by building on the insight that in many tasks, these free-form outputs can be interpreted in a semantically meaningful latent structure, for example, discrete class labels, integers, or sets. We introduce task calibration as a paradigm to calibrate the model's predictive distribution in the task-induced latent space. We apply a decision-theoretic result to show that Minimum Bayes Risk (MBR) decoding on the task-calibrated latent distribution is the optimal decoding strategy on latent model beliefs. Empirically, it consistently improves generation quality across different tasks and baselines. We also introduce Task Calibration Error (TCE), an application-aware calibration metric that quantifies the excess loss due to miscalibration. Our work demonstrates that task calibration enables more reliable model decisions across various tasks and applications.