Beyond Post-Hoc Temperature Scaling: Bilevel Optimization for LLM Calibration
2026-08-07 • Machine Learning
Machine Learning
AI summaryⓘ
The authors found that large language models often become too confident in their answers and don't know when they're unsure. Instead of adjusting the model after training, they change the training process to make the model less overconfident by encouraging it to be more uncertain when needed. They do this by maximizing the randomness (or entropy) in the model’s predictions using a special two-step training method. Their approach works well, especially when the model faces questions from new topics it wasn’t trained on.
Large Language ModelsCalibrationOverconfidenceTemperature ScalingEntropyBilevel OptimizationOut-of-Domain GeneralizationHyperparametersFirst-order Approximation
Authors
Ruochen Jin, Zhanliang Wang, Zongyu Dai, Jiancong Xiao, Bojian Hou
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.