Beyond Post-Hoc Temperature Scaling: Bilevel Optimization for LLM Calibration

2026-08-07

AI Summary

This work addresses the poor calibration and overconfidence commonly observed in large language models after preference alignment, a problem exacerbated by the limited cross-domain generalizability of conventional post-hoc calibration methods such as temperature scaling. The authors propose a scalable bilevel optimization framework that integrates calibration directly into training: the lower-level optimizes model parameters using a parameterized loss function, while the upper-level automatically selects the loss hyperparameters by maximizing the entropy of the predicted distribution to discourage overly peaked predictions. By incorporating the maximum predictive entropy objective into the training process—eschewing reliance on post-processing—and leveraging first-order gradient approximations for computational efficiency, the method achieves significant improvements in calibration across both multiple-choice and open-ended question-answering tasks, with notably enhanced out-of-domain generalization.

Abstract

Preference alignment often makes large language models (LLMs) overconfident and poorly calibrated. Traditional post-hoc temperature scaling is inherently domain-dependent: a temperature fitted on one domain does not generalize across domains. This motivates us to modify model parameters during training to improve calibration. We propose maximizing the entropy of predictive distributions as the calibration objective, which directly targets overconfidence by discouraging overly concentrated predictions. Inspired by temperature scaling, we realize this through a bilevel optimization formulation, where the lower level trains the model under a parametric loss and the upper level selects loss hyperparameters to maximize entropy. To make the framework practical at LLM scale, we adopt an efficient first-order approximation that avoids explicit second-order computation. Across both multiple-choice and open-ended generative question answering, experiments demonstrate that our method yields well-calibrated LLMs with particular advantages in out-of-domain generalization.

Problem

• LLM calibration
• overconfidence
• domain generalization
• temperature scaling
• preference alignment

Innovation

• bilevel optimization
• LLM calibration
• entropy maximization
• temperature scaling
• out-of-domain generalization