Distance-aware Calibration for Pre-trained Language Models (2024)

Abstract: Language Models for text classification often produce overconfident predictions for both in-distribution and out-of-distribution samples, i.e., the model’s output probabilities do not match their accuracy. Prior work showed that simple post-hoc approaches are effective for mitigating this issue, but are not robust in noisy settings, e.g., when the distribution shift is caused by spelling mistakes. In this work, we propose Distance Aware Calibration (DAC), a post-hoc approach that changes the confidence scores of a Language Model leveraging the distance between new samples been evaluated and the in-domain training set. We show that using DAC on top of a Language Model can improve in-domain calibration, robustness to different kind of distribution shift and also the model’s ability to detect out-of-distribution samples. We provide an extensive evaluation on common text classification benchmark for both calibration and out-of-distribution detection tasks.

## introduction fine-tuning language models such as bert @xcite and roberta @xcite has been established as the de-facto standard approach for text classification whenever a large amount of annotated training examples is available. even in the era of large language models (llms), smaller (and more efficient) fine-tuned models have been shown to still be able to outperform in context learning (icl) methods for much larger llms @xcite . one of the main open challenges with language models (lms) of any size is how to obtain reliable uncertainty estimates for their predictions @xcite ; this is particularly important in safety critical applications where lms are deployed to † work done while working at amazon web services, berlin, germany the clean test sample on top is correctly classified with high confidence (green background), but once it's corrupted the model's prediction flips but the confidence is unchanged (red background). dac leverages the samples' distances from the training set to decrease the confidence on the predictions, leaving the correctly classified one almost unchanged while yielding low confidence for the sample being far from the training dataset (white circle). to mitigate this problem, prior work has focused on either designing new training strategies to achieve better calibration @xcite , or on the use of post-hoc methods @xcite , such as temperature scaling @xcite . the latter can be applied on top of any model, independently from its architecture and its training setup, which makes them highly appealing for practitioners. despite their simplicity, post-hoc methods like temperature scaling (ts) have been shown to be extremely effective for both in-domain (ind) and out-of-domain (ood) † calibration in text classification @xcite , outperforming more complex (and time-consuming) bayesian methods that rely on some kind of (multiple) models average in order to estimate the uncertainty around a prediction. nevertheless, ts is not reliable for all kinds of distribution-shift, as it suffers in the presence of grammatical errors, typos and other common mistakes which are easily found in real world applications @xcite . in this work, we leverage recent findings on distance awareness @xcite @xcite and its connections with high quality uncertainty quantification and propose distance-aware calibration (dac), a post-hoc method to improve model's calibration and its robustness to distribution shifts. our method is built on the belief that the confidence over a model's prediction should decrease as the inputs move away from the training domain. thus, we propose to adjust a model's confidence using sample-specific information, namely, the distance between the embeddings of a new sample and the representation of the samples in the training set. we hypothesised that this would mitigate lms's overconfidence, benefiting ood calibration and, to a smaller extent, in-domain calibration. figure 1 shows an illustrative example of dac's adjusted confidence for test points located at various distances from the training set. in this work we first introduce our calibration method dac ( §3.2.1) and detail the procedure used to optimize it for both ind and ood calibration ( §3.2.2). then, we show the efficacy of our method in both traditional settings -where model calibration is measured on datasets that share the same task (e.g., sentiment analysis) but differ for domain or writing style -and noisy settings, where the ood datatset is a corrupted version of the indomain one ( §5.1). moreover, we will show that a lm equipped with dac not only achieves better † in this work we will use the terms in-domain (out-ofdomain) and in-distribution (out-of-distribution) interchangeably calibration but it's confidence scores can also be used to distinguish between in-domain and out-ofdomain samples, boosting performance in the ood detection task ( §5.2). ## related work calibration in language models several methods have been developed to mitigate language model's overconfidence, thus providing more calibrated predictions. the methods in the literature can be broadly classified in two groups: i) posthoc methods and ii) methods that require model re-training. to the first class of methods belongs temperature scaling (ts) @xcite , which has been shown to be beneficial for model calibration in text classification settings, both ind and ood @xcite . more evidence in favour of ts is presented in @xcite where ts is shown to be superior to other uncertainty quantification methods such as mc dropout @xcite , (last layer) variational inference @xcite and deep ensembles @xcite . more recent works focused on a second class of model calibration techniques, those that require retraining. park and caragea (2022) uses mixup to improve the calibration of fine-tuned language models, showing improvement in expected calibration error (ece) w.r.t. base ts and label smoothing. @xcite leverage model explanations to enhance calibration. the proposed method is used in conjunction with ts to improve model performance for ind and ood calibration. recently, @xcite showed that pretrained-language models (plms) are better calibrated on the masked language modelling (mlm) task than their finetuned counterpart. the authors hypothesised that lm's overconfidence is due to catastrophic forgetting and showed that preserving the pre-trained features while fine-tuning lms can improve ind and ood calibration. finally, @xcite discovered that ts performs poorly when the indomain samples are perturbed with noise. thus, they propose an extension of ts based on the measured distribution-shift between the in-domain dataset and the noisy one. this work closely relates to ours as they also used some notion of "distance" to the in-domain dataset. however, unlike their work, dac's confidence is calibrated sample-wise and does not require access to samples coming from the same ood distribution that will be later use for evaluation. also, we do not focus exclusively on noisy-settings like @xcite . ## experimental setup in this work we aim at improving the model's uncertainty estimates, thus we designed different experiments aimed at assessing i) the model's calibration under different distribution shifts and ii) the model's ood detection capabilities by means of its confidence scores. in this section we will introduce the datasets we used for evaluation, as well as the metrics and the baselines. ## conclusions in this work we introduce dac, a new post-hoc calibration approach that leverages per-sample information in order to compute new calibrated confidence scores. unlike previous work, our approach improves the performance of the underlying language models on a variety of settings, including noisy ones, which have been shown to be particularly challenging for transformers based models. additionally, when coupled with dac, the underlying language model shows significant gains in pure ood settings, where the new calibrated con-12441 fidence allows to better distinguish between indomain and out-of-domain samples. finally, we performed an ablation study to further understand whether dac's performance is sensitive to some core design choices, concluding that, the use of a distribution-free distance such as knn is the one yielding the best overall performance.