Task-Aware Calibration: Provably Optimal Decoding in LLMs

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 level of free-form language. We build on the recent insight that in many tasks, these free-form outputs can be interpreted in terms of semantically meaningful latent structures, for example, discrete class labels, real numbers, or sets. We introduce task calibration as a paradigm to calibrate the model’s predictive distribution in the task-specific latent space. We show that Minimum Bayes Risk (MBR) decoding on a task-calibrated latent distribution is the optimal decoding strategy on latent model beliefs. We empirically confirm this to consistently improve generation quality across various tasks and latent structures. We also introduce Task Calibration Error (TCE), an application-aware calibration metric that quantifies the excess loss due to miscalibration. Our work demonstrates that using task-aligned latent structures enables semantically meaningful calibration for LLMs, which improves downstream decision making, such as decoding.