Source-linked AI summary
Uncertainty Quantification with Pre-trained Language Models: A Large-Scale Empirical Analysis
Yuxin Xiao, Paul Pu Liang, Umang Bhatt, Willie Neiswanger, Ruslan Salakhutdinov, Louis-Philippe Morency
TL;DR
PLM confidence must track correctness, particularly in high-stakes use. This paper compares pipeline choices across three classification tasks and two domain settings. It finds calibration is generally consistent across tasks and domains, with XLNet less robust to domain shift.
Problem
PLM confidence should reliably estimate correctness, especially when predictions inform high-stakes decisions.
Source: §1 · p. 1Method
The paper systematically compares prevalent PLM-pipeline choices across three NLP classification tasks and two domain settings.
Source: §7 · p. 9Results
Takeaways & Limitations
For existing PLM pipelines, the authors recommend ELECTRA, larger models when possible, Temp Scaling, and Focal Loss.
Source: §7 · p. 9The analysis applies to existing widely used PLMs and does not claim generalization to new pre-training corpus and task combinations.
Source: §8 · p. 9
Abstract
from arXiv · showhide
Pre-trained language models (PLMs) have gained increasing popularity due to their compelling prediction performance in diverse natural language processing (NLP) tasks. When formulating a PLM-based prediction pipeline for NLP tasks, it is also crucial for the pipeline to minimize the calibration error, especially in safety-critical applications. That is, the pipeline should reliably indicate when we can trust its predictions. In particular, there are various considerations behind the pipeline: (1) the choice and (2) the size of PLM, (3) the choice of uncertainty quantifier, (4) the choice of fine-tuning loss, and many more. Although prior work has looked into some of these considerations, they usually draw conclusions based on a limited scope of empirical studies. There still lacks a holistic analysis on how to compose a well-calibrated PLM-based prediction pipeline. To fill this void, we compare a wide range of popular options for each consideration based on three prevalent NLP classification tasks and the setting of domain shift. In response, we recommend the following: (1) use ELECTRA for PLM encoding, (2) use larger PLMs if possible, (3) use Temp Scaling as the uncertainty quantifier, and (4) use Focal Loss for fine-tuning.
1 Introduction
The paper addresses the lack of broad empirical evidence on how to compose well-calibrated PLM pipelines, evaluating pipeline choices across NLP tasks and domain settings. It recommends specific PLMs, model sizes, uncertainty quantifiers, and fine-tuning losses, while reporting several PLM-specific calibration patterns.
1 Introduction
- Reliable confidence estimates matter because PLM predictions may inform high-stakes decisions, including medical diagnosis.A 90% confidence estimate should correspond to a correct prediction 90% of the time in practice.Source: §1 · p. 1
- Prior studies examined limited pipeline considerations and few PLM types, leaving the ideal choices for calibration unclear.Source: §1 · p. 1
- The analysis compares four to five options for each of four pipeline considerations across three classification tasks and in- and out-of-domain settings.It assesses explicit prediction and calibration error plus selective prediction and out-of-domain detection.Source: §1 · p. 2
- The paper recommends ELECTRA, larger PLMs when possible, Temp Scaling for post hoc recalibration, and Focal Loss during fine-tuning.Source: §1 · p. 2
- PLM calibration is generally consistent across tasks and domains, but XLNet is most vulnerable to domain shift; quantifiers help more out-of-domain, and ensembles are less effective.Larger PLMs are better calibrated in-domain for Commonsense Reasoning, unlike the other tasks.Source: §1 · p. 2
- The released evaluation includes 120 PLM pipelines, tested on all three tasks in both domain settings using four metrics.The metrics cover prediction and calibration error, selective prediction, and out-of-domain detection.Source: §1 · p. 2
2 Background
The background defines a PLM-based classification pipeline and its calibration objective, then situates the study among broader uncertainty research. It identifies a gap in prior PLM calibration studies, which considered only a limited set of models and pipeline choices.
2.1 Problem Formulation
- A PLM encodes each input for a fine-tuned classifier that outputs label probabilities, which an optional post hoc quantifier can modify to reduce calibration error.Calibration means predicted confidence matches empirical likelihood; the study uses ECE and examines PLM choice, size, quantifier, and loss.Source: §2.1 · p. 2
2.2 Related Work
- Prior uncertainty work spans varied domains, perspectives, techniques, and related properties, while PLM studies often examine few models and omit quantifiers or losses beyond Temp Scaling and Cross Entropy.PLMs typically use Transformer architectures and self-supervised pretraining.Source: §2.2 · pp. 2–3
3 Which Pre-trained Language Model?
The section compares five PLM encoders across three classification tasks and in- and out-of-domain settings. ELECTRA performs best overall, while calibration generally worsens under domain shift.
3.1 Experiment Setup
- Evaluation spans sentiment analysis, natural language inference, and commonsense reasoning, with in-domain and out-of-domain datasets for each task.The in-/out-of-domain pairs are IMDb/Yelp, MNLI/SNLI, and SWAG/HellaSWAG, respectively.Source: §3.1 · p. 3
- The encoder comparison uses base BERT, XLNet, ELECTRA, RoBERTa, and DeBERTa models with a two-layer classifier.Source: §3.1 · pp. 3–4
3.2 Empirical Findings
- ELECTRA has the lowest ECE across all three tasks and both domains, plus the lowest prediction error and RPP.The authors hypothesize that replaced-token detection contributes to its success by avoiding artificial [MASK] tokens and learning from all input tokens.Source: Figure 1 · p. 4, §3.2 · p. 6
3.1 Experiment Setup
- Calibration is measured by ECE, while selective prediction is assessed with RPP, where lower RPP means confidence favors correct predictions.ECE aggregates the absolute accuracy–confidence gap across confidence bins.Source: §3.1 · p. 6
3.2 Empirical Findings
- Calibration deteriorates out-of-domain for all five PLMs, with generally consistent ECE rankings across tasks and domains; XLNet has the largest ECE increase.The authors suggest that XLNet's permuted language-modeling pretraining may be vulnerable to domain shift.Source: §3.2 · p. 6
4 What Model Size?
Larger PLMs generally calibrate better out-of-domain, despite slightly higher in-domain ECE; Commonsense Reasoning is an exception, where larger models also calibrate better in-domain. Larger models also achieve lower prediction errors in both domains.
4.1 Experiment Setup
- The experiment compares large and base versions of five PLMs while keeping the remaining setup unchanged.The authors recommend larger PLMs if computational budget permits.Source: §4.1 · p. 6, §4.2 · p. 6
4.2 Empirical Findings
- In Commonsense Reasoning, larger PLMs are significantly better calibrated in-domain than their respective base versions.The authors suggest this may indicate greater awareness of uncertainty during reasoning.Source: §4.2 · p. 6
- Larger PLMs tend to have slightly higher in-domain ECE but lower out-of-domain ECE and prediction errors in both domains.Out-of-domain calibration and prediction errors are positively correlated; pipelines with better calibration also tend to be more accurate under domain shift.Source: §4.2 · p. 6
5 Which Uncertainty Quantifier?
The study compares four uncertainty quantifiers and finds that Temp Scaling reduces ECE the most, with greater calibration improvement out-of-domain. It preserves prediction results while adding little computational burden.
5.1 Experiment Setup
- The study compares four post hoc quantifiers: Temp Scaling, MC Dropout, Ensemble, and last-layer stochastic variational inference.MC Dropout averages dropout-enabled forward passes; Ensemble averages independently trained models, while LL SVI averages stochastic passes through Bayesian classification layers.Source: §5.1 · pp. 6–7
5.2 Empirical Findings
- Temp Scaling, MC Dropout, Ensemble, and LL SVI reduce ECE in descending order of effectiveness, with larger improvements out-of-domain.LL SVI can increase ECE, while Temp Scaling preserves prediction results and adds little computational burden.Source: Figure 3 · p. 7, §5.2 · pp. 7–8
- MC Dropout and LL SVI can reduce out-of-domain prediction error in Commonsense Reasoning but usually become overconfident, raising ECE.Ensemble lowers prediction error, but is less effective for calibration in these PLM-based pipelines.Source: §5.2 · p. 7
- Ensemble and MC Dropout increase computational cost during fine-tuning or inference, respectively.Temp Scaling recalibrates predictions post hoc with little added computational burden.Source: §5.2 · p. 8
6 Which Fine-tuning Loss?
The study compares five fine-tuning losses and finds Focal Loss compelling across calibration and out-of-domain detection. Its FAR95 is lowest, while ECE patterns for other losses vary between in-domain and out-of-domain evaluation.
6.1 Experiment Setup
- The comparison fine-tunes BERT base with Cross Entropy, Brier Loss, Focal Loss, Label Smoothing, and MMCE.Focal Loss modulates cross-entropy to focus learning on hard misclassified samples; MMCE is a differentiable calibration-error proxy generally used alongside cross-entropy.Source: §6.1 · p. 8
- FAR95 measures the share of in-domain samples whose confidence falls below the 95th percentile of out-of-domain sample confidence.FAR95 scores are relatively high for Sentiment Analysis and Natural Language Inference, which the authors attribute possibly to those pipelines also predicting well out-of-domain.Source: §6.1 · p. 8, §6.2 · p. 9
6.2 Empirical Findings
- Label Smoothing, Focal Loss, and MMCE produce better-calibrated BERT base models than Cross Entropy and Brier Loss.Cross Entropy, Focal Loss, and MMCE calibrate better in-domain, while Brier Loss and Label Smoothing show decreased ECE out-of-domain.Source: §6.2 · p. 8
- Focal Loss delivers competitively low ECE in- and out-of-domain across all three tasks and the lowest FAR95 among five losses.The authors interpret its FAR95 as indicating that Focal Loss models are most alert to domain shift.Source: §6.2 · p. 9
7 Conclusion
The paper compares prevalent choices for four considerations in PLM-based pipelines across three NLP classification tasks and two domain settings. It recommends a set of pipeline choices and reports several calibration patterns specific to PLM-based pipelines.
7 Conclusion
- The analysis compares prevalent pipeline options across three NLP classification tasks and two domain settings to identify choices for reducing calibration error.The four considerations span PLM encoding, model size, uncertainty quantification, and fine-tuning loss.Source: §7 · p. 9
- The authors recommend ELECTRA encoding, larger PLMs when possible, Temp Scaling for post hoc recalibration, and Focal Loss during fine-tuning.Source: §7 · p. 9
- Calibration quality is generally consistent across tasks and domains, but XLNet is least robust to domain shift; larger models fare better in-domain only for Commonsense Reasoning.Uncertainty quantifiers are generally more effective out-of-domain, while ensembles are less effective at reducing calibration error in PLM pipelines.Source: §7 · p. 9
8 Limitation
Computational constraints limit the pre-training combinations examined, and the paper does not claim generalization to new combinations. It also leaves the effects of domain-shift degree and other NLP task types for future work.
8 Limitation
- Computational constraints prevented pre-training PLMs from scratch with other combinations of pre-training corpora and tasks.Source: §8 · p. 9
- The analysis applies to existing widely used PLMs, but the authors do not claim it generalizes to new pre-training corpus and task combinations.Source: §8 · p. 9
- The authors note that techniques tailored to PLM-based pipelines may further mitigate calibration error.Source: §8 · p. 9
- The study does not inspect how the degree of domain shift affects calibration or evaluate cross-lingual classification and generation.Those task types are left for future work.Source: §8 · p. 9
A Responsible NLP Research
The experiments require substantial computational power, and the reported scope focuses on English PLM-based pipelines assessed on six NLP datasets. The authors release evaluation benchmarks and describe the licenses and content boundaries of the packages and datasets used.
A Responsible NLP Research
- Training multiple large-scale PLMs consumes a fair amount of computational power.Source: A Responsible NLP Research · p. 12
- The work focuses on English PLM-based pipelines and assesses them on six NLP datasets.Source: A Responsible NLP Research · p. 12
- The authors release evaluation benchmarks; the Hugging Face package uses Apache License 2.0 and Bayesian-Torch uses BSD 3-Clause, while the benchmarks contain no uniquely identifying personal information or offensive content.Source: A Responsible NLP Research · p. 12