Calibrating Language Models with Adaptive Temperature Scaling (2024)

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.

## introduction large language models (llms) have become a cornerstone of modern artificial intelligence, offering impressive capabilities in natural language processing tasks. however, the reliability of llms is intertwined with their ability to generate confidence scores that accurately reflect the likelihood of their outputs being correct. this calibration, aligning a model's confidence with its accuracy, is essential, especially when llms are deployed in real-world scenarios where decisions based on incorrect outputs can have significant consequences. while unsupervised pre-training methods have shown success in producing well-calibrated llms, * equal contribution. a challenge arises when these models undergo finetuning through reinforcement learning from human feedback (rlhf). while rlhf fine-tuning is effective in enhancing model performance on specific tasks and aligning outputs with human preferences, recent studies indicate a notable degradation in the calibration of llms post-rlhf finetuning @xcite @xcite . this degradation compromises the model's ability to provide reliable confidence scores, an issue that becomes critical when these models are applied to tasks requiring high levels of trust and accuracy. an important question arises: how can we maintain the performance gains achieved through rlhf fine-tuning while ensuring that the model's confidence scores remain reliable? to address this challenge, our work introduces adaptive temperature scaling (ats), a post-hoc calibration technique that predicts a temperature scaling parameter for each token prediction based on a language model's hidden features. basic temperature scaling is a widely-used calibration method that applies a single temperature parameter across all outputs of a model. this technique, while effective in some contexts, assumes uniform calibration needs across all inputs, which is often not the case for complex models like llms. ats, in contrast, predicts a unique temperature scaling parameter for each set of token predictions. this input-specific approach allows ats to refine the calibration process, addressing the varying degrees of calibration shift that can occur after rlhf fine-tuning. for instance, certain inputs or topics might be more susceptible to miscalibration post-rlhf, and ats can adaptively adjust the scaling for these instances more aggressively than for others where the model's confidence remains relatively well-aligned with its accuracy. importantly, our approach reduces the need for task-specific calibration, which may be difficult to achieve in many cases, given the wide variety of downstream tasks that llms may be used for. we conduct experiments on mmlu, triviaqa, and truthfulqa to evaluate the effectiveness of ats in improving the calibration of llms following rlhf fine-tuning. our findings demonstrate that ats improves the calibration of post-rlhf llms by 10-50% on average, while having no effect on model performance. ## related work recent literature has extensively discussed the challenges of maintaining calibration in llms, particularly highlighting the degradation in calibration post-rlhf @xcite @xcite @xcite . the concept of verbalized confidence has been explored as a way to counteract this degradation @xcite , and dialogue models have been shown to express uncertainty in a well-calibrated manner @xcite . compared to works on improving sentence level calibration given tokenlevel probabilities @xcite , our work aims to directly improve the calibration of token-level probabilities. the calibration of neural networks has been a topic of significant interest, with foundational concepts such as proper scoring rules @xcite laying the groundwork. model mismatch and distribution shift often degrade calibration, commonly quantified with common metrics including expected calibration error (ece) @xcite and brier score @xcite . modern neural networks have been found to exhibit overconfidence @xcite @xcite , especially in the context of image classification @xcite @xcite . various methods have been proposed for calibrating neural networks, including temperature scaling @xcite , platt scaling @xcite , label smoothing @xcite , scaling binning @xcite , and more sophisticated approaches @xcite @xcite . while these methods offer strategies for improving model calibration, our approach uniquely adapts the temperature scaling parameter for each token prediction based on its hidden features, tailoring the method to the problem of language modeling. ## background and problem setting we consider access to a conversation sft dataset of @xmath0, y)} with vocabulary v where @xmath1, denotes the instruction, each with sequence length l x , and @xmath2. we wish to calibrate language model π(y|x). while we do not make any assumptions about the training process of π, we find our calibration method is most useful for language models following an rlhf process where token-level calibration is often significantly degraded compared to base language models which are generally well calibrated @xcite . for a given sample (x, y), we generate a set of unnormalized logits ẑ = π(x) ∈ r lx+ly×|v | where each ẑi defines the unnormalized logits for the i + 1-th token and |v | is the vocabulary size. prior methods @xcite propose various scaling methods for calibrating models by transforming logits. in matrix scaling, a calibration head is used to produce calibrated logits @xmath3, b are learnable parameters. in the case of language modeling where |v | is large, learning a full transform matrix becomes computationally infeasible, so we compare to vector scaling, where w is constrained to a diagonal matrix. temperature scaling is the case when w is constrained further to a scalar matrix and b to the zero-vector. to learn these parameters, these methods minimize the cross-entropy over the sft dataset calculated over response tokens. ## adaptive temperature scaling architecture. temperature scaling, while effective in classification settings, struggles to adapt logits well in language modeling as the confidence scores that are most important (such as those that contain actual answers or facts) account for only a small portion of natural language sequences. therefore, optimizing a single temperature parameter often results in post-rlhf language models still being overconfident post scaling. additionally, language model miscalibration largely varies based on the type of token being predicted following rlhf. matrix and vector scaling can in theory perform adaptive confidence prediction by using logits as features; however, they are prone to overfitting, as we find in section 5. to balance regularization with modeling capac-ity in our calibration head, we instead propose to use a head architecture that predicts a singular temperature for every token prediction. for an input pair (x, y), we first produce input-dependent features ĥ ∈ r lx+ly,h using the language model π. we then learn a calibration head to produce a temperature vector c θ ( ĥ) = τ ∈ r lx+ly . we exponentiate τ to ensure positive values then transform logits to yield calibrated logits @xmath4. in practice, we find that directly using the logits ẑ as features can be inefficient (with a large vocabulary size) and also less effective compared to hidden states. therefore, we use the last hidden state of the language model π as the features for predicting τ . with this architecture formulation, we retain the ability to predict confidences adaptively depending on the context, while also never changing the ranking for the possible next token given specific context, as each set of token logits are scaled by only a single value. ## experiments in this section, we aim to evaluate our proposed method on multiple benchmarks to demonstrate its effectiveness in improving calibration of llms fine-tuned with rlhf. we compare our method to no calibration as well as existing temperature scaling methods. additionally, we ablate the main components of our method including the loss function, loss weighting, and head architecture. evaluation setting. we evaluate using two 7b parameter post-rlhf models llama-2-chat-7b @xcite and qwen-chat-7b. as the calibration dataset, we use the alpaca gpt-4 @xcite instruction tuning dataset, which contains a diverse set of instructions with high quality answers. we then evaluate model calibration on three downstream tasks. we perform multiple choice evaluation on the mmlu @xcite by aggregating statistics across the entire dataset. specifically we concatenate the confidences and correctness labels from all subjects, then calculate the calibration metrics. we also evaluate on two free response datasets, triviaqa @xcite and truthfulqa @xcite . metrics. in multiple choice inference, we have a set of tokens ids o which represent the valid options for a multiple choice answer, so the confidence scores are @xmath5,@xmath6. to calculate confidences over a long sequence of response tokens for an input x, we sample a generation ŷ of length l ŷ from the original language model then concatenate to the instruction to form ẑ and q following calibration. then, we calculate an average over transition probabilities on the response tokens. we use the expected calibration error (ece) @xcite and brier score @xcite to evaluate calibration. we also report accuracy but each method does not significantly affect accuracy. baselines. we compare our method to the post-rlhf model without calibration, temperature scaling, vector scaling, and scaling binning @xcite . we do not evaluate matrix scaling as the full matrix becomes computationally infeasible for large vocabulary sizes, as the projection matrix requires the square of the vocabulary size parameters. ## conclusion in this paper, we introduce adaptive temperature scaling, a calibration technique for post-rlhf llms, offering a significant improvement in model calibration without compromising performance. by predicting a temperature based on token-level features, ats addresses the diverse calibration needs of llms. our results across multiple benchmarks confirm our approach can calibrate llms on a wide array of tasks using existing sft datasets.