Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CATTO: Balancing Preferences and Confidence in Language Models

Nisarg Parikh, Kunjal Panchal, Ananya Sai, Pannaga Shivaswamy, Andrew Lan

TL;DR

CATTO introduces a calibration-aware token-level training objective that integrates confidence calibration directly into Direct Preference Optimization, addressing miscalibration induced by alignment. By defining a differentiable per-token calibration loss anchored to a differentiable correctness surrogate based on token margins, CATTO reduces Expected Calibration Error while preserving or improving task accuracy. The framework also adds Confidence@k for test-time, Bayes-optimal candidate selection using calibrated token probabilities. Empirically, CATTO yields robust calibration improvements in- and out-of-distribution across multiple benchmarks and enhances downstream decision-making without sacrificing performance, validating its practical impact for safer and more reliable LLM deployment.

Abstract

Large language models (LLMs) often make accurate next token predictions but their confidence in these predictions can be poorly calibrated: high-confidence predictions are frequently wrong, and low-confidence predictions may be correct. This miscalibration is exacerbated by preference-based alignment methods breaking the link between predictive probability and correctness. We introduce a Calibration Aware Token-level Training Objective (CATTO), a calibration-aware objective that aligns predicted confidence with empirical prediction correctness, which can be combined with the original preference optimization objectives. Empirically, CATTO reduces Expected Calibration Error (ECE) by 2.22%-7.61% in-distribution and 1.46%-10.44% out-of-distribution compared to direct preference optimization (DPO), and by 0.22%-1.24% in-distribution and 1.23%-5.07% out-of-distribution compared to the strongest DPO baseline. This improvement in confidence does not come at a cost of losing task accuracy, where CATTO maintains or slightly improves multiple-choice question-answering accuracy on five datasets. We also introduce Confidence@k, a test-time scaling mechanism leveraging calibrated token probabilities for Bayes-optimal selection of output tokens.

CATTO: Balancing Preferences and Confidence in Language Models

TL;DR

CATTO introduces a calibration-aware token-level training objective that integrates confidence calibration directly into Direct Preference Optimization, addressing miscalibration induced by alignment. By defining a differentiable per-token calibration loss anchored to a differentiable correctness surrogate based on token margins, CATTO reduces Expected Calibration Error while preserving or improving task accuracy. The framework also adds Confidence@k for test-time, Bayes-optimal candidate selection using calibrated token probabilities. Empirically, CATTO yields robust calibration improvements in- and out-of-distribution across multiple benchmarks and enhances downstream decision-making without sacrificing performance, validating its practical impact for safer and more reliable LLM deployment.

Abstract

Large language models (LLMs) often make accurate next token predictions but their confidence in these predictions can be poorly calibrated: high-confidence predictions are frequently wrong, and low-confidence predictions may be correct. This miscalibration is exacerbated by preference-based alignment methods breaking the link between predictive probability and correctness. We introduce a Calibration Aware Token-level Training Objective (CATTO), a calibration-aware objective that aligns predicted confidence with empirical prediction correctness, which can be combined with the original preference optimization objectives. Empirically, CATTO reduces Expected Calibration Error (ECE) by 2.22%-7.61% in-distribution and 1.46%-10.44% out-of-distribution compared to direct preference optimization (DPO), and by 0.22%-1.24% in-distribution and 1.23%-5.07% out-of-distribution compared to the strongest DPO baseline. This improvement in confidence does not come at a cost of losing task accuracy, where CATTO maintains or slightly improves multiple-choice question-answering accuracy on five datasets. We also introduce Confidence@k, a test-time scaling mechanism leveraging calibrated token probabilities for Bayes-optimal selection of output tokens.
Paper Structure (63 sections, 17 theorems, 122 equations, 1 figure, 12 tables)

This paper contains 63 sections, 17 theorems, 122 equations, 1 figure, 12 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Calibration during Supervised Fine-Tuning
  4. Calibration during Preference Optimization
  5. Methodology
  6. Notation.
  7. A Differentiable Per-Token Calibration Loss
  8. A population measure of miscalibration.
  9. Decomposing the per-token calibration risk.
  10. A surrogate for correctness.
  11. From token-level to sequence-level calibration.
  12. Calibration-Aware Preference Optimization
  13. Joint objective for confidence-calibrated DPO.
  14. Confidence-Aware Inference via Confidence@k
  15. Confidence@k decision rule.
  16. ...and 48 more sections

Key Result

Proposition 1.1

The subgradients of the per-token $L_1$ calibration loss with respect to $c_\theta(x)$ satisfy In contrast, $L_2$ gradients scale as $2(c_\theta(x) - \widetilde{z}(x))$. This boundedness ensures that large surrogate errors $\epsilon(x)$ do not destabilize confidence updates.

Figures (1)

  • Figure 1: Reliability diagrams (top) compare predicted confidence with actual accuracy, while confidence histograms (bottom) show how often different confidence levels occur, using the same random seed. The dashed line represents perfect calibration, deviations from it (marked with red bars annotated as "Gap") indicate miscalibration. DPO produces strong overconfidence that remains even after RCFT post-hoc calibration, whereas CATTO brings confidence closer to true accuracy, reducing calibration error (ECE).

Theorems & Definitions (30)

  • Proposition 1.1: Gradient Stability
  • Proposition 1.2: Directional Consistency of the Margin Surrogate
  • Proposition 1.3: Preference Ordering Stability
  • Proposition 1.4: Bayes-Optimality of Confidence@k
  • Proposition 5.1: Ordering-Preserving Surrogate
  • Proposition 5.2: Bayes optimality of the Brier ($L_2$) calibration loss
  • proof
  • Proposition 5.3: Bayes optimality of $L_1$ for deterministic surrogate targets
  • proof
  • Definition 5.4: Surrogate Contamination Model
  • ...and 20 more