Restoring Calibration for Aligned Large Language Models: A Calibration-Aware Fine-Tuning Approach

TL;DR

This work addresses calibration degradation in LLMs caused by preference alignment methods like RLHF and DPO. It introduces a theoretical framework separating models into calibratable and non-calibratable regimes and proposes calibration-aware fine-tuning (CFT) plus an EM-based ECE regularization (RCFT) to restore calibration without sacrificing alignment performance. The approach is validated across multiple open-source LLMs and benchmarks, showing substantial reductions in calibration error (ECE) from roughly the 14–20% range to about 2–7% in various settings, while preserving or even improving language capabilities. The results demonstrate that calibration can be improved in aligned LLMs, enabling more reliable probability estimates for downstream decision making in high-stakes or distributed-use scenarios.

Abstract

One of the key technologies for the success of Large Language Models (LLMs) is preference alignment. However, a notable side effect of preference alignment is poor calibration: while the pre-trained models are typically well-calibrated, LLMs tend to become poorly calibrated after alignment with human preferences. In this paper, we investigate why preference alignment affects calibration and how to address this issue. For the first question, we observe that the preference collapse issue in alignment undesirably generalizes to the calibration scenario, causing LLMs to exhibit overconfidence and poor calibration. To address this, we demonstrate the importance of fine-tuning with domain-specific knowledge to alleviate the overconfidence issue. To further analyze whether this affects the model's performance, we categorize models into two regimes: calibratable and non-calibratable, defined by bounds of Expected Calibration Error (ECE). In the calibratable regime, we propose a calibration-aware fine-tuning approach to achieve proper calibration without compromising LLMs' performance. However, as models are further fine-tuned for better performance, they enter the non-calibratable regime. For this case, we develop an EM-algorithm-based ECE regularization for the fine-tuning loss to maintain low calibration error. Extensive experiments validate the effectiveness of the proposed methods.

Figures (8)

• Figure 1: Calibration performance comparison between DPO and our approach on Llama3.1-8B-Tulu (a DPO-aligned version of Llama-3.1 touvron2023llama). Left: Model calibration plots after DPO alignment, showing significant overconfidence. Middle: Calibration plots after applying our fine-tuning approach, demonstrating improved calibration. Right: The evolution of confidence ECE and classwise ECE across different stages (pre-trained, SFT, DPO, and our method) shows how our approach effectively restores calibration errors.
• Figure 2: Illustration of TCE, calibratable and non-calibratable Regimes. Assume that the target probabilistic generative model $\boldsymbol{p}^*$ has a 60% accuracy. Left: The TCE range (yellow) is bounded between $C|60\%-\textnormal{Accuracy}|$ and $2\cdot|60\%-\textnormal{Accuracy}|$. Right: The calibratable regime (red) spans from 0 to TCE when accuracy $\leq 60\%$. The non-calibratable regime (blue) spans from a non-zero lower bound to TCE when accuracy $> 60\%$.
• Figure 3: Calibration Plots of (a) DPO, (b) Temperature Scaling (TS), (c) our CFT, (d) our RCFT on Llama-3.1-8B-Tulu. Each panel plots the model’s predicted probabilities (i.e., confidence) on the $x$-axis against the observed accuracy (fraction correct) on the $y$-axis, binned into ten groups. The diagonal line in each panel represents perfect calibration. The depth of the color indicates the sample density in that column. DPO has the worst calibration performance. Other three methods improve the calibration performance where our CFT has the lowest con-ECE (shown in the parenthesis). This figure omits the first two bins because the model selects an answer with the largest predicted probability which is always larger than 0.25 in the four options prediction task (so no samples exist below that threshold).
• Figure 4: (a) Using only $\mathcal{L}_{\text{SFT}_2}$ achieves high accuracy ($\sim$90%) but shows poor calibration. (b) Using only $\mathcal{L}_{\text{ECE}}$ drives the model toward random guessing, achieving near-zero ECE but with only 25% accuracy.
• Figure 5: Calibration Plots of (a, e) DPO, (b, f) Temperature Scaling (TS), (c, g) our CFT, (d, h) our RCFT on Llama-3.1-8B-Tulu. (a-d) are the classwise calibration curve and (e-h) are the confidence calibration curve. Each panel plots the model’s predicted probabilities (i.e., confidence) on the $x$-axis against the observed accuracy (fraction correct) on the $y$-axis, binned into ten groups. The diagonal line in each panel represents perfect calibration. The depth of the color indicates the sample density in that column. DPO has the worst calibration performance. Other three methods improve the calibration performance where our CFT has the lowest con-ECE (shown in the parenthesis). The figures of conf-ECE (e-h) omit the first two bins because the model selects an answer with the largest predicted probability which is always larger than 0.25 in the four options prediction task (so no samples exist below that threshold).

Theorems & Definitions (13)

• Definition 3.1: Classwise Calibration
• Definition 3.2: Confidence Calibration
• Proposition 4.1
• Definition 4.2: Target probabilistic generative model
• Definition 4.3: Target Calibration Error (TCE).
• Theorem 4.4: Upper Bound of TCE
• Theorem 4.5: Lower Bound of TCE
• Theorem 4.6: Upper bound for ECE
• Definition 1.1: Multiclass Calibration
• proof