Calibrating Language Models with Adaptive Temperature Scaling

TL;DR

This paper introduces Adaptive Temperature Scaling (ATS), a post-hoc calibration method that predicts token-level temperature parameters to improve the confidence calibration of large language models after fine-tuning.

Contribution

The paper presents ATS, a novel adaptive calibration technique that adjusts confidence scores at the token level, addressing calibration degradation post-RLHF fine-tuning.

Findings

ATS improves calibration by 10-50% across benchmarks.

ATS maintains the performance benefits of RLHF.

It predicts token-specific temperature parameters based on features.

Abstract

The effectiveness of large language models (LLMs) is not only measured by their ability to generate accurate outputs but also by their calibration-how well their confidence scores reflect the probability of their outputs being correct. While unsupervised pre-training has been shown to yield LLMs with well-calibrated conditional probabilities, recent studies have shown that after fine-tuning with reinforcement learning from human feedback (RLHF), the calibration of these models degrades significantly. In this work, we introduce Adaptive Temperature Scaling (ATS), a post-hoc calibration method that predicts a temperature scaling parameter for each token prediction. The predicted temperature values adapt based on token-level features and are fit over a standard supervised fine-tuning (SFT) dataset. The adaptive nature of ATS addresses the varying degrees of calibration shift that can occur after RLHF fine-tuning. ATS improves calibration by over 10-50% across three downstream natural language evaluation benchmarks compared to prior calibration methods and does not impede performance improvements from RLHF.