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

AI summary

Overview

Research area: Large language model reliability, specifically confidence calibration of preference-aligned LLMs; training-time optimization and bilevel optimization.

Technical level: Advanced. The paper assumes familiarity with temperature scaling, cross-entropy and entropy objectives, fine-tuning of LLMs, and bilevel optimization (including hypergradients and the implicit function theorem).

Scope: The paper proposes CALM (Calibration for Large Models via Bilevel Optimization), a first-order bilevel training framework that uses an entropy-maximization upper-level objective to reduce overconfidence in aligned LLMs, and evaluates it on multiple-choice and open-ended question answering in both in-domain and out-of-domain settings.

What This Paper Is About

Preference alignment (RLHF and DPO) makes LLMs overconfident, and the standard fix, post-hoc temperature scaling, is dataset-dependent: a temperature fitted on one domain does not transfer to another. The authors instead change the model during training so that a temperature-like adjustment generalizes across domains, using a bilevel formulation in which the lower level fine-tunes the model under a parametric loss and the upper level selects loss hyperparameters to maximize the entropy of the predictive distribution.

Key Contributions

1. The authors identify the dataset dependence of post-hoc temperature scaling and propose a training-time perspective whose goal is to learn models that admit a more general calibration adjustment across datasets.
2. They formulate calibration-oriented loss design as a bilevel optimization problem, with the lower level training the model and the upper level optimizing temperature-like loss hyperparameters using an entropy-based objective on held-out inputs.
3. They propose an entropy-maximization upper-level objective to mitigate overconfidence in aligned LLMs without requiring additional hard labels, and they make it practical via a BOME-inspired first-order approximation that avoids second-order computation and Hessian inversion.
4. They demonstrate empirically that the method improves calibration across multiple-choice and open-ended generative QA, with the clearest gains in out-of-domain settings, against baselines including calibration-aware fine-tuning.

Main Findings

• Post-hoc temperature scaling does not generalize: The scalar temperature learned on one dataset is specific to that dataset; in the OOD MCQA setting, temperature scaling actually worsened conf-ECE on Llama-3.1 (0.1784 → 0.2513).
• Best OOD conf-ECE on three of four models: On Llama-3.1, CALM reduced conf-ECE from 0.1784 to 0.1050 (reported as a 41.1% reduction) with a 2.35 percentage-point drop in accuracy; on Vicuna-7B it reached the lowest conf-ECE (0.0380) and cw-ECE (0.0576) while marginally improving accuracy (44.95% vs. 44.70%); on OLMo-2-7B it reached the best cw-ECE (0.0912, with a 0.75 pp accuracy change); on Mistral-7B it reached the best conf-ECE (0.0822) with 51.60% accuracy.
• Competing methods show accuracy–calibration trade-offs: The Iterate ablation achieved low ECE on Mistral (cw-ECE 0.0518) but collapsed accuracy from 59.80% to 25.80%; Regularization preserved accuracy on some models but degraded it on Mistral (44.20%); CFT left conf-ECE at or above 0.17 on all four models in the OOD setting.
• The bilevel coupling, not entropy alone, matters: Regularization is exactly CALM's entropy objective used as a grid-tuned single-level penalty, and it did not match CALM at any grid point, which the authors take as evidence that the coupling term is necessary.
• In-domain results: Temperature scaling achieved the best ECE on most models in the matched-distribution setting, but cannot improve accuracy. CALM was the only method that consistently improved accuracy over the baseline across all four models (reported as +1.3, +1.0, +2.5, +1.1 percentage points) while also improving calibration. Regularization achieved the highest accuracy gains (+5 to +10 percentage points) but worsened calibration.
• Generative domain shift: Under the PopQA-to-TriviaQA transfer with Sem-ECE, CALM had the lowest mean Sem-ECE (0.0671), beating DPO/RLHF (0.0846) and every training-time baseline. CFT improved from 0.1330 at fixed α = 0.5 to 0.1109 at its best per-model α from {0.1, 0.5, 0.9}, but still trailed both CALM and the uncalibrated baseline. Label smoothing collapsed to Sem-ECE values up to 1.0.
• Parameterization ablation: On Mistral-7B in the OOD setting, the per-vocabulary parameters (d_y, l_y) gave conf-ECE 0.0822 versus 0.3140 for a collapsed scalar (learned-TS) variant.
• Language ability retention: CALM had the smallest average OOD-to-ID gap (2.3) compared with Iterate (21.5) and Regularization (4.0). Under OOD, CALM averaged 66.99 (−5.8), 64.59 (−0.3), 71.29 (−0.0), and 66.94 (−2.1) across the four models, while Iterate regressed by 16.4, 20.6, 13.7, and 34.4 points, and Label Smoothing degraded severely.
• Seed stability: A three-seed study (reported in Appendix D.1, Table 11) shows mean conf-ECE improving over the uncalibrated baseline on all four models; Vicuna-7B and OLMo-2-7B were stable (conf-ECE standard deviation at or below 0.002), while Llama-3.1 and Mistral-7B showed larger run-to-run variance, so their single-run margins should be read with more caution.
• Reliability diagrams: On Mistral-7B (OOD), the DPO baseline was overconfident in medium- and high-confidence regions; temperature scaling and CFT reduced this only partially; CALM aligned bins closely with the diagonal.

Methodology in Plain English

The starting point is temperature scaling, which divides logits by a positive scalar to soften predictions, with the temperature fitted by minimizing cross-entropy on a dataset. The authors generalize this in three steps.

First, they replace the single scalar temperature with higher-dimensional parameters: a vector of class-dependent temperatures and a vector of additive logit shifts, which reduces to standard temperature scaling as a special case.

Second, they move from post-hoc adjustment to training time, optimizing the model parameters and these temperature-like parameters jointly, so the model itself becomes one for which a general adjustment works.

Third, they replace the cross-entropy objective with entropy maximization over the predictive distribution. Entropy maximization discourages concentrated predictions, pushing the model away from the overconfidence regime and toward the underconfidence extreme; the authors argue that increasing entropy moves the model continuously between these regimes and may pass through a point where confidence matches accuracy. Because entropy alone could overshoot and hurt discrimination, they keep cross-entropy on the training set as a lower-level constraint requiring the model parameters to be optimal under the current temperature parameters.

This yields a bilevel problem: the upper level minimizes the entropy objective on a held-out calibration set, and the lower level requires the parameters to be optimal for the training cross-entropy. Computing the exact hypergradient would require second-order derivatives and an inverse Hessian, which is infeasible at LLM scale. Following BOME (Liu et al., 2022), the authors reformulate the problem as constrained optimization over the joint variable (α, θ), enforce lower-level optimality through a residual constraint that compares the current training loss to a surrogate value function, and update using a direction that combines the upper-level entropy gradient with the residual gradient, requiring only first-order derivatives.

Experiments fine-tune four aligned models with QLoRA (rank 64, α = 32, float16, NVIDIA A100 40GB GPUs, 5 epochs, batch size 16, weight decay 0.02). They compare against the uncalibrated DPO/RLHF model, temperature scaling, label smoothing, and calibration-aware fine-tuning (CFT), plus two ablations (Regularization and Iterate). Evaluation covers MCQA (trained on Alpaca and tested on MMLU, MedMCQA, OpenBookQA, and ARC-Challenge, 3,000 test samples for OOD; a 2,000 training / 1,000 calibration / 2,000 test split for ID) and open-ended QA using PopQA and TriviaQA, with conf-ECE and cw-ECE for MCQA and Sem-ECE for free-form generation.

Why This Matters

Impact on research: The paper reframes calibration from a post-hoc correction into a training-time objective, and it connects calibration loss design to bilevel optimization, showing that a bilevel constraint can be made practical for multi-billion-parameter models via a first-order approximation. The Regularization ablation provides evidence that the coupling, not entropy maximization alone, drives the gains, which is a specific claim about why bilevel structure helps.

Real-world applications:

• Healthcare decision support, where overconfident predictions can lead to poor decisions and misplaced trust.
• Legal analysis, where confidence is used as a proxy for reliability.
• Any deployment where models are used across domains they were not fine-tuned on, since the reported advantage is largest under distribution shift.
• Selective prediction or abstention systems that rely on well-calibrated confidence to decide when to defer to a human.

Industry relevance: Since the method operates on supervised fine-tuning with QLoRA on 40GB GPUs and does not require an auxiliary calibration head, hidden-state probe, or reward model modification, it is complementary to existing alignment pipelines and can be folded into fine-tuning rather than added as a separate post-processing stage.

Future Directions

• Reducing run-to-run variance: the three-seed study showed larger conf-ECE variance on Llama-3.1 and Mistral-7B than on Vicuna-7B and OLMo-2-7B, so the stability of the calibration gains across seeds and the single-run margins remain open questions.
• Closing the residual in-domain gap: temperature scaling still achieved the best ECE on most models in the ID setting, so matching post-hoc methods under matched distributions while retaining accuracy gains is unresolved.
• Extending the evaluation beyond the reported benchmarks: the training-cost comparisons and detailed hyperparameter grids are placed in the appendices, and the experiments cover four models and the datasets listed above; broader model families, scales, and tasks are not reported.
• Better integration with complementary calibration approaches: the authors position CALM as complementary to auxiliary-head calibrators, hidden-state probes, and RLHF-time reward calibration, leaving combinations of these with bilevel training-time calibration unexplored here.

Target Audience

Researchers and practitioners working on LLM reliability, uncertainty quantification, and calibration; machine learning engineers who fine-tune aligned models with QLoRA and need confidence estimates that transfer across domains; and readers interested in bilevel optimization applied at LLM scale, particularly those familiar with post-hoc temperature scaling who want to understand its training-time generalization.

Authors’ 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.