Making Pre-trained Language Models both Task-solvers and Self-calibrators

TL;DR

This paper introduces LM-TOAST, a training algorithm that enables pre-trained language models to serve as both effective task-solvers and self-calibrators, especially under limited data and distribution shifts, improving confidence estimation without sacrificing task accuracy.

Contribution

The paper proposes LM-TOAST, a novel training method that enhances PLMs' ability to self-calibrate with limited data, addressing challenges like data imbalance and distribution shifts.

Findings

LM-TOAST improves confidence calibration in PLMs.

It maintains original task performance while enhancing calibration.

Effective in downstream applications like selective classification and adversarial defense.

Abstract

Pre-trained language models (PLMs) serve as backbones for various real-world systems. For high-stake applications, it's equally essential to have reasonable confidence estimations in predictions. While the vanilla confidence scores of PLMs can already be effectively utilized, PLMs consistently become overconfident in their wrong predictions, which is not desirable in practice. Previous work shows that introducing an extra calibration task can mitigate this issue. The basic idea involves acquiring additional data to train models in predicting the confidence of their initial predictions. However, it only demonstrates the feasibility of this kind of method, assuming that there are abundant extra available samples for the introduced calibration task. In this work, we consider the practical scenario that we need to effectively utilize training samples to make PLMs both task-solvers and self-calibrators. Three challenges are presented, including limited training samples, data imbalance, and distribution shifts. We first conduct pilot experiments to quantify various decisive factors in the calibration task. Based on the empirical analysis results, we propose a training algorithm LM-TOAST to tackle the challenges. Experimental results show that LM-TOAST can effectively utilize the training data to make PLMs have reasonable confidence estimations while maintaining the original task performance. Further, we consider three downstream applications, namely selective classification, adversarial defense, and model cascading, to show the practical usefulness of LM-TOAST. The code will be made public at \url{https://github.com/Yangyi-Chen/LM-TOAST}.