Calibrated Language Model Fine-Tuning for In- and Out-of-Distribution Data

TL;DR

This work tackles miscalibration in fine-tuned language models for both ID and OOD data by introducing a data-augmentation-inspired calibration framework with on-manifold interpolation and off-manifold entropy regularization. The proposed objective regularizes predictions through pseudo-sample generation inside and outside the data manifold, improving ID calibration and OOD uncertainty. Across six text-classification datasets, the method achieves superior ECE and NBAUCC scores for misclassification and OOD detection, with ablations confirming complementary effects of both regularizers. The approach offers a general, practical calibration technique for large pre-trained language models that preserves accuracy while enhancing reliability.

Abstract

Fine-tuned pre-trained language models can suffer from severe miscalibration for both in-distribution and out-of-distribution (OOD) data due to over-parameterization. To mitigate this issue, we propose a regularized fine-tuning method. Our method introduces two types of regularization for better calibration: (1) On-manifold regularization, which generates pseudo on-manifold samples through interpolation within the data manifold. Augmented training with these pseudo samples imposes a smoothness regularization to improve in-distribution calibration. (2) Off-manifold regularization, which encourages the model to output uniform distributions for pseudo off-manifold samples to address the over-confidence issue for OOD data. Our experiments demonstrate that the proposed method outperforms existing calibration methods for text classification in terms of expectation calibration error, misclassification detection, and OOD detection on six datasets. Our code can be found at https://github.com/Lingkai-Kong/Calibrated-BERT-Fine-Tuning.

Figures (6)

• Figure 1: The reliability diagrams on in-distribution data (the first row) and the histograms of the model confidence on out-of-distribution (OOD) data (the second row) of CNN kim2014convolutional and fine-tuned BERT-MLP classifier devlin2019bert. Though BERT improves classification accuracy, it makes over-confident predictions for both in-distribution and OOD data.
• Figure 2: The on-manifold and off-manifold samples generated by our calibration procedure. Mixup adopts a coarse linear interpolation and the generated data point may deviate from the data manifold.
• Figure 3: Calibration curves of OOD detection and misclassification detection on WOS. Our method can achieve high $F_1$ scores starting from a small threshold which indicates that it indeed provides low confidences for misclassified and OOD samples; the $F_1$ scores of the baselines peak at high thresholds which indicates that they are poorly calibrated.
• Figure 4: Parameter study of $\delta_{\rm on}$, $\delta_{\rm off}$ and $\delta_y$.
• Figure 5: Parameter study of $\delta_{\rm on}$, $\delta_{\rm off}$ and $\delta_y$. We use ${\rm NBAUCC}_{1}$ for OOD and misclassification detection.