🤖 machine learning

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

This paper proposes a practical bilevel optimization framework that maximizes predictive entropy during training to calibrate large language models, effectively overcoming the domain limitations of traditional post-hoc temperature scaling and improving out-of-domain generalization.

Original authors: Ruochen Jin, Zhanliang Wang, Zongyu Dai, Jiancong Xiao, Bojian Hou

Published 2026-08-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruochen Jin, Zhanliang Wang, Zongyu Dai, Jiancong Xiao, Bojian Hou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a brilliant student who has read every book in the library. This student is a Large Language Model (LLM), a type of artificial intelligence that can write stories, solve math problems, and chat like a human. But there's a catch: this student is a terrible judge of their own knowledge. They are often wildly overconfident. If they don't know the answer to a question, they might say, "I am 99% sure I'm right!" when they are actually wrong. In the real world, this is dangerous. If a medical AI is 99% sure about a diagnosis but is wrong, people could get hurt.

To fix this, scientists have tried a trick called "Temperature Scaling." Think of the AI's confidence like a volume knob on a stereo. If the music is too loud and distorted (too confident), you turn the knob down to make it softer and more realistic. This works well if you are listening to just one song (one specific dataset). But here's the problem: the "volume knob" setting that works for a rock song doesn't work for a jazz song. If you try to use the same setting for a different type of question, the AI either stays too loud or gets too quiet. The old method is like trying to tune a radio for a whole city with just one dial; it just doesn't work everywhere.

This is where a new study, published as a conference paper at COLM 2026, steps in with a clever new idea. Instead of just turning a knob after the student has finished studying, the researchers decided to change how the student learns in the first place. They created a special training method called CALM (Calibration for Large Models via Bilevel Optimization).

Here is how CALM works, using a simple analogy: Imagine a coach training an athlete.

The Old Way (Post-Hoc Scaling): The athlete runs a race, gets tired, and then the coach says, "Okay, next time, run a little slower." The coach tries to adjust the speed after the fact. But if the terrain changes (a different dataset), the old advice doesn't fit.
The CALM Way (Bilevel Optimization): The coach sets up a two-level training camp.
- Level 1 (The Athlete): The athlete practices running the race to get faster and smarter (this is the model learning to answer questions).
- Level 2 (The Coach): While the athlete runs, the coach watches and adjusts the training rules in real-time. The coach's goal isn't just to make the athlete fast; it's to make sure the athlete knows exactly how fast they are running. If the athlete starts bragging too much ("I'm the fastest!"), the coach gently nudges them to be a bit more humble and realistic.

The researchers found that by training the model with this "two-level" system, the AI learns to be naturally confident but not over-confident. They tested this on four different popular AI models (like Llama-3.1, Vicuna, OLMo, and Mistral) and two types of tests: multiple-choice questions and open-ended creative writing.

The results were quite promising. When the AI was tested on questions it had never seen before (the "out-of-domain" test, like a jazz song for a rock-trained student), CALM worked much better than the old "volume knob" method. The old method often made the AI more confused or overconfident on new topics. CALM, however, kept the AI's confidence in check. For example, on one model called Mistral-7B, the old method left the AI with a calibration error of 0.335 (very overconfident), while CALM brought it down to 0.0822 (much more accurate).

Crucially, the researchers showed that this didn't make the AI "dumber." The models still answered questions correctly; they just stopped lying about how sure they were. The paper suggests that this approach is a strong step forward, especially for situations where the AI needs to handle unexpected questions without needing a human to manually tweak its settings every time. It's a way to teach AI to be humble, honest, and reliable, no matter what topic it's discussing.

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

Problem Statement

Large Language Models (LLMs), particularly those aligned via Reinforcement Learning from Human Feedback (RLHF) or Direct Preference Optimization (DPO), frequently exhibit poor calibration, manifesting as overconfidence. While confidence is often used as a proxy for reliability, miscalibrated models in high-stakes domains (e.g., healthcare, legal) can lead to misplaced trust and poor decision-making.

Existing solutions primarily rely on post-hoc Temperature Scaling (TS), where a single scalar temperature is fitted on a validation set to rescale logits after training. However, the paper identifies a critical limitation: TS is inherently dataset-dependent. A temperature optimized for one domain often fails to generalize to others, making it difficult to find a "universal" temperature that calibrates well across diverse tasks. Furthermore, post-hoc methods leave the model parameters fixed, missing the opportunity to learn a model that is intrinsically more amenable to calibration.

Methodology: CALM (Calibration for Large Models via Bilevel Optimization)

The authors propose CALM, a training-time framework that reformulates calibration as a bilevel optimization problem. Instead of adjusting a post-hoc parameter, CALM jointly optimizes the model parameters and "temperature-like" loss hyperparameters to maximize predictive entropy while preserving task performance.

1. Bilevel Formulation

The problem is structured as follows:

Lower Level (Task Optimization): Trains the model parameters $\theta$ to minimize a standard parametric loss (cross-entropy) on the training set $S_T$ , given a set of hyperparameters $\alpha$ .
$\theta(\alpha) \in \arg\min_{\theta'} \mathcal{L}_{tr}(\theta'; \alpha)$
Upper Level (Calibration Optimization): Optimizes the hyperparameters $\alpha$ to maximize the entropy of the predictive distribution on a held-out calibration set $S_V$ , subject to the constraint that $\theta$ remains optimal for the lower-level task.
$\min_{\alpha} \mathcal{L}_{ent}(\alpha, \theta(\alpha))$

The hyperparameters $\alpha$ generalize the concept of temperature scaling. Instead of a single scalar $\tau$ , CALM uses a higher-dimensional parameter $\alpha = (l, \tau)$ , where $l \in \mathbb{R}^K$ introduces additive logit shifts and $\tau \in \mathbb{R}^K_{>0}$ provides class-dependent temperatures. This allows for per-vocabulary adjustments rather than a global rescaling.

2. Objective Functions

Upper-Level Objective: Maximizes Shannon entropy $H(f_{\alpha, \theta}(\cdot | x))$ to discourage overly concentrated (overconfident) predictions. This provides a controllable direction to move the model from overconfidence toward underconfidence, potentially finding a point where confidence matches accuracy.
Lower-Level Constraint: Ensures the model retains discriminative ability by minimizing cross-entropy on the training data.

3. Efficient First-Order Approximation

Exact bilevel optimization requires computing hypergradients involving second-order derivatives (Hessian inverses), which is computationally prohibitive for LLMs with billions of parameters. To address this, CALM adopts a first-order approximation inspired by BOME (Bilevel Optimization Made Easy):

It reformulates the problem as a constrained optimization over the joint variable $u = (\alpha, \theta)$ .
It approximates the lower-level value function by running $T$ inner gradient steps from the current $\theta$ to generate a surrogate $\tilde{\theta}$ .
It constructs a "plug-in residual" constraint and computes a joint update direction using only first-order derivatives, avoiding explicit Hessian inversion.

Key Contributions

Training-Time Perspective: The paper shifts the calibration paradigm from post-hoc dataset-specific adjustments to training-time optimization, aiming to learn models that admit general calibration adjustments across datasets.
Bilevel Loss Design: It formulates calibration-oriented loss design as a bilevel problem where the upper level optimizes temperature-like hyperparameters using an entropy-based objective, and the lower level trains the model.
Entropy Maximization: The authors propose an entropy-maximization upper-level objective to mitigate overconfidence in aligned LLMs without requiring additional hard labels or auxiliary heads.
Scalable Framework: They introduce an efficient first-order bilevel optimization scheme (CALM) that avoids expensive second-order computations, making the approach practical for large-scale LLM fine-tuning.

Experimental Results

The method was evaluated on four aligned LLMs (Llama-3.1-Tulu-8B, Vicuna-7B, OLMo-2-7B, Mistral-7B) across multiple-choice QA (MCQA) and open-ended generative QA.

Out-of-Domain (OOD) Generalization

Superior Calibration: CALM achieved the best confidence Expected Calibration Error (conf-ECE) on three of the four models in the OOD setting, significantly reducing error compared to baselines. For example, on Llama-3.1, it reduced conf-ECE from 0.1784 to 0.1050.
Accuracy Preservation: Unlike some ablations (e.g., "Iterate" or "Regularization" without the bilevel constraint) that suffered catastrophic accuracy collapse, CALM maintained high task accuracy while improving calibration.
Robustness: In open-ended generative QA (cross-domain transfer between PopQA and TriviaQA), CALM achieved the lowest mean Semantic-ECE (0.0671), outperforming both the uncalibrated baseline and strong training-time baselines like Calibration-Aware Fine-Tuning (CFT).

In-Domain (ID) Performance

In the ID setting, where training and evaluation distributions match, Temperature Scaling performed best on ECE (as expected for post-hoc methods). However, CALM was the only training-based method that consistently improved both accuracy and calibration across all four models compared to the baseline.

Ablation Studies

Bilevel Structure: Removing the bilevel constraint (Regularization) or the coupling term (Iterate) resulted in either worse calibration or significant accuracy loss, demonstrating that the coupling between task performance and calibration is essential.
Parameterization: Using per-vocabulary vector parameters ( $d_y, l_y$ ) was far superior to a collapsed scalar parameter (learned TS), highlighting the need for fine-grained adjustments to handle asymmetric miscalibration across answer options.

Significance and Claims

The paper claims that CALM provides a robust, training-time solution to the overconfidence problem in aligned LLMs. Its primary significance lies in:

Generalization: Unlike post-hoc TS, CALM learns a calibration mechanism that generalizes across domains, addressing the "dataset dependence" of current methods.
Trade-off Management: It achieves a superior accuracy-calibration trade-off, avoiding the utility collapse often seen when attempting to improve calibration via simple regularization.
Practicality: By utilizing a first-order approximation, it makes bilevel optimization feasible for modern LLMs, offering a path to "calibrated-by-design" models rather than "calibrated-by-post-processing."

The authors conclude that their method is particularly advantageous in out-of-domain settings, where alternative methods often become unstable or suffer substantial utility loss, making it a viable approach for deploying reliable LLMs in diverse, real-world scenarios.