P LEX : Towards Reliability using Pretrained Large Model Extensions
Abstract
A recent trend in artificial intelligence is the use of pretrained models for language and vision tasks, which have achieved extraordinary performance but also puzzling failures.
Probing these models’ abilities in diverse ways is therefore critical to the field.
In this paper, we explore the reliability of models, where we define a reliable model as one that not only achieves strong predictive performance but also performs well consistently over many decision-making tasks involving uncertainty (e.g., selective prediction, open set recognition), robust generalization (e.g., accuracy and proper scoring rules such as log-likelihood on inand out-of-distribution datasets), and adaptation (e.g., active learning, few-shot uncertainty).
We devise 10 types of tasks over 40 datasets in order to evaluate different aspects of reliability on both vision and language domains.
To improve reliability, we developed ViT-Plex and T5-Plex, p retrained l arge model ex tensions ( PLEX ) for vision and language modalities, respectively.
Plex greatly improves the state-of-the-art across reliability tasks, and simplifies the traditional protocol as it does not require designing scores or tuning the model for each individual task.
We demonstrate scaling effects over model sizes up to 1B parameters and pretraining dataset sizes up to 4B examples.
We also demonstrate Plex’s capabilities on challenging tasks including zero-shot open set recognition, active learning, and uncertainty in conversational language understanding.1 Google 2 University of Oxford *
Work done at Google.
Correspondence to: Dustin Tran < trandustin@google.com > .
Presented at ICML Workshop on Principles of Distribution Shift 2022 , Copyright 2022 by the author(s).
A full version of this paper can be found at https:/goo.gle/plex-paper .
Code for training & evaluation is open-sourced in Uncertainty Baselines (Nado et al., 2021).
Layer and method implementations use Edward2 (Tran et al., 2018). [Figure]
ODEL
[Figure]
OBUST [Figure] ENERALIZATION
IID accuracy
Accuracy on out-of-distribution data (e.g. covariate or subpopulation shift)
Scoring rules
[Figure]
NCERTAINTY
[Figure]
uncertainty [Figure]
DAPTATION
[Figure]
Reliability as a Goal for AI
Over the past few years, the deep learning approach to artificial intelligence (AI) has made significant progress on benchmark tasks across domains such as computer vision (Dosovitskiy et al., 2020) and natural language processing (Raffel et al., 2020; Brown et al., 2020).
With this progress, there is unfettered excitement about the potential of AI to have a transformative impact.
While hypothesizing about this potential is important, we highlight that the tasks where deep learning has been most successful have been carefully devised to fit within narrow boundaries—for example, a focus on predictive performance with test inputs close to the data on which the model was trained.
To go beyond these limitations, we argue that the ability of models to make reliable decisions is critical to the deeper integration of AI in the real world.
Here, we define reliability as the ability for a model to work consistently across real-world settings.
We borrow the term from reliability engineering (Barlow & Proschan, 1975; O’Connor & Kleyner, 2012), a discipline of engineering involving risk assessment, testability, and fault tolerance.
Related nomenclature include robustness (Russell et al., 2015), safety (Amodei et al., 2016;
Everitt et al., 2018; Hendrycks et al., 2021b), calibration (Dawid, 1982), credibility (D’Amour et al., 2020) and trustworthiness (Avin et al., 2021), each with their own broad and intersecting scopes.
P LEX Selective
Prediction for Human+AI Collaboration
Structured Open Set Recognition
Can you set my savings pin to 1234?
Vertical=Bank, Domain=
Account, Intent=OOD
Hmm, your Bank Account app doesn't support that.
Would you like to call the customer service?
Please change the radio to local sports.
Vertical=Media, Domain=OOD, Intent=OOD
Sorry, we don't support this Media yet.
Robustness to
Spurious Correlations
Will I make it to my current destination on this fuel tank?
[Figure]
What
[Figure]
s happening around the
[Figure]
pu
[Figure]
Hen-of-the-wood
Florist shop Aquarium
Pelican
Zero-Shot Open Set Recognition
[Figure]
Label Uncertainty
Active Learning
[Figure]
Rainforest
Flowerpot
[Figure]
C o ll a b
o r a t i v e
A U C
[Figure]
Desiderata for Reliability The majority of machine learning research focuses on measures of performance based on the accuracy on a test set drawn from the same distribution as the training set, the so-called independent and identically distributed (i.i.d.) assumption.
However, this does not capture the real-world deployment of AI systems, where often the testing environment is very different from the training environment.
The emphasis in our paper is on how reliable an AI system is in such novel scenarios.
We posit three general categories of desiderata for reliable AI systems: they should represent their own uncertainty, they should generalize robustly to new scenarios, and their learning procedures should be able to adapt to new data.
Importantly, the aim for a reliable model is to do well in all of these areas simultaneously out-of-the-box without requiring any customization for individual tasks (Figure 1):
Uncertainty involves imperfect or unknown information where it is impossible to exactly describe an existing state (Ghahramani, 2015).
Predictive uncertainty quantification allows one to compute optimal decisions (Parmigiani & Inoue, 2009), and enables practitioners to know when to trust the model’s predictions, thereby enabling graceful failures when the model is likely to be wrong.
In the latter case, which is often referred to as selective prediction , the model may defer its prediction to human experts when it is not confident.
Robust Generalization involves an estimate or forecast about an unseen event (Abraham & Ledolter, 1983; Dawid, 1982).
The quality of prediction is typically measured using accuracy (e.g. top-1 error
for classification problems and mean squared error for regression problems) and proper scoring rules such as log likelihood and
Brier score (Gneiting & Raftery, 2007).
In the real world, we care not only about metrics on new data obtained from the same distribution the model was trained on (i.i.d.), but also about robustness , as measured by metrics on data under out-of-distribution shifts such as covariate or subpopulation shift.
Adaptation involves probing the model’s abilities over the course of its learning process.
Benchmarks typically evaluate on static datasets with pre-defined train-test splits.
However in many applications, we are interested in models that can quickly adapt
to
new datasets and efficiently learn
with as few labeled examples as possible.
Exam-
ples include few
-shot learning
(Ravi & Larochelle, 2017),
where the model learns from a small set of examples; active learning (Settles, 2009), where the model not only learns but also participates in acquiring the data to learn from; and lifelong learning (Thrun, 1998), where the model learns over a sequence of tasks and must not forget about relevant information for previous tasks.
Contributions First, we define and evaluate reliability in a comprehensive fashion.
We use 10 types of tasks in order to capture the three reliability areas—uncertainty, robust generalization, and adaptation—and so that the tasks measure a diverse set of desirable properties in each area.
Together the tasks comprise 40 downstream datasets across vision and natural language modalities: 14 datasets for finetuning (including few-shot and active learning-based adaptation) and 26 datasets for out-of-distribution evaluation (Appendix A).
To improve reliability, we develop ViT-Plex and T5-Plex, building on large pretrained models on vision (ViT (Dosovitskiy et al., 2020)) and language (T5 (Raffel et al., 2020)) respectively.
We train variants of Plex over multiple model sizes and pretraining dataset sizes on up to 4 billion examples.
Figure 3 illustrates Plex’s performance on a select set of tasks comparing to existing state-of-the-art, which typically use models specialized for that task.
Plex greatly improves the state-of-the-art over the total of 40 datasets.
Importantly, Plex achieves impressive performance across all tasks using out-of-the box model output without requiring any custom designing or tuning for each individual task.
Tasks for Benchmarking Reliability
We evaluate a model’s reliability using 10 types of tasks, which we define below.
We selected a broad suite ofdownstream datasets under the tasks, each ranging from several hundred to a million examples; see Appendix A.
Uncertainty: Calibration assesses how well a model’s predicted confidence is reflected over a population (Dawid, 1982).
We compute expected calibration error (Naeini et al., 2015) on 14 image and 10 text datasets.
Selective prediction jointly assesses the predictive performance and quality of uncertainty estimates of a model, by abstaining from making predictions on examples for which a model’s predictive uncertainty estimates are above a given threshold and recording predictive accuracy on the remaining examples.
We compute two metrics, Calibration AUC and Oracle Collaborative Accuracy (Kivlichan et al., 2021), on 4 image and 10 text datasets.
Open set recognition assesses how well a model can detect examples belonging to none-of-the training classes.
We use AUROC and experiment with maximum softmax probability as the detection score.
(We use Mahalanobis distance for zero-shot open set recognition.)
Label uncertainty is a type of uncertainty inherent in the data labels.
This is a form of irreducible data uncertainty, e.g. noise, which is considered distinct from uncertainty arising from the choice over models (Dusenberry et al., 2020b)
We use two datasets: CIFAR-10H (Peterson et al., 2019) and
ImageNet ReaL
(Beyer et al., 2020).
Robust Generalization: We assess in-distribution generalization , i.e. how well a model can make predictions after finetuning, by examining accuracy, negative log-likelihood, and Brier score on the in-distribution test splits of 5 image and 3 text datasets.
With out-of-distribution data , we assess how robustly a model’s predictions generalize to input distributions it was not trained on.
We use the same metrics measured for in-distribution, and we investigate 4 types of out-of-distribution data: covariate shift, semantic (class) shift, data uncertainty, and subpopulation shift.
Adaptation: Few-shot learning assesses how well a model can make predictions downstream with only a few training examples.
We use 9 datasets and apply multiple few-shot settings: 1-shot, 5-shot, 10-shot, and 25-shot (x-shot means x examples per class).
We also evaluate few-shot uncertainty , where we examine calibration, selective prediction, and open set recognition in the few-shot regime.
We use all 9 datasets for few-shot learning in order to evaluate calibration and selective prediction, and we use those with OOD datasets (ImageNet and CIFAR-100) for open set recognition.
We also perform zero-shot open set recognition by using the Mahalanobis distance scoring to detect whether an input is out-of-distribution based on the model’s representation layer.
Active learning assesses how well a model knows what it does not know by selecting informative samples to label using uncertainty.
We assess accuracy over a total number of acquired examples and apply margin sampling (Settles, 2009) for multi-class uncertainty sampling.
P LEX : P retrained L arge model Ex tensions
Plex is the result of an extensive study of the reliability of large pretrained models and their complementarity with existing reliability methods.
In particular, ViT Plex and T
Plex use several key ingredients:
• Base Transformer architecture.
We adopt the Transformer standard of an alternating sequence of attention and feedforward layers.
We build on T5 1.1 (Raffel et al., 2020) for text as a Transformer in an encoder-decoder setup where the raw text is tokenized with SentencePiece, and on Vision Transformer (Dosovitskiy et al., 2020) for images in an encoder-only setup where the raw images are effectively tokenized under 32 × 32 patches.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
• Model size.
We investigate 3 scales of the model size in ViT Plex (S, B, L) and 3 scales of the model size in T Plex (Small, Base, Large).
• Pretraining dataset size.
For vision, we scale pretraining from ImageNet21K to the JFT web dataset on up to 4B images.
This mirrors recent work on scaling vision models (Zhai et al., 2021; Pham et al., 2021).
For language, we use the C4 dataset which consists of hundreds of gigabytes of English text scraped from the web (Raffel et al., 2020).
• Efficient ensembling.
Ensembles and Bayesian neural nets have shown to be very effective for uncertainty and robustness (Ovadia et al., 2019;
Dusenberry et al., 2020a; Band et al., 2021).
To do so scalably, we use BatchEnsemble (BE) (Wen et al., 2020) and experiment with its use on both the attention and feedforward layers or on only the feedforward layer.
For faster training, we only apply BatchEnsemble at a select number of later layers, similar to mixture of experts models (Riquelme et al., 2021).
• Last layer changes.
We experiment with two approaches that modify the model’s final layer to improve reliability, given a fixed representation (a.k.a.
deterministic uncertainty quantification setting (Van Amersfoort et al., 2020)).
First, we use Gaussian process (GP) last-layer, which improve distance-awareness of the decision surface by increasing uncertainty far away from the training representations.
We use the GP layer implementation proposed by Liu et al.
In addition, pretraining uses increasingly noisier datasets with a large number of output classes, and the ability to model input-dependent label noise becomes more important.
We apply the Heteroscedastic (Het) method of Collier et al.
• What to apply in pretraining versus finetuning.
We
[Figure]
R e li a b ili t y S
c
o r e
[Figure]
experiment with both pretraining and finetuning for vision models.
Due to compute constraints, we exclusively focus on the finetuning-only setting for T5-Plex.
That is, T5Plex models are initialized from the official pretrained T checkpoints, and we apply efficient ensembling and last layer changes during finetuning.
• Few-shot protocol.
As an alternative to logistic regression on the final layer of frozen representations, we experiment with gradient descent over all parameters.
We also experiment with a GP or Heteroscedastic last layer.
Summary of Results and Scaling Trends
Figure 3 displays our model’s overall performance comparing reliability task performance to existing specialized state-of-the-art.
Here, we validate several takeaways as we ablate to understand the ingredients behind Plex.
Scaling model size improves reliability.
Figure 4 displays
ViT-Plex and T5-Plex over varying model sizes. ViT-Plex is pretrained with JFT by default; I21K denotes pretraining on ImageNet21K. We compute a reliability score which is [Table]
[Figure]
[Table] an average over all 52 task metrics (see Appendix B for details).
Classical machine learning theory would suggest that a larger model translates to more overfitting and might therefore be less reliable as it may be overconfident and less robust.
However, we find that scale improves overall performance across tasks.
Scaling pretraining dataset size improves reliability.
ViTPlex L with JFT performs better than ViT-Plex L with ImageNet-21K (Figure 4). In Table 1, we also perform an ablation by comparing pretraining on JFT with 300M examples to JFT with 4B examples. We pretrain on up to 8M steps with batch size 4096, which is up to 8X more steps than we typically use for pretraining; each result is a separate run using a tuned learning rate schedule.
ImageNet 10-shot accuracy is always better on JFT 4B under the same number of training steps. The models also converge faster with the smaller JFT 300M, reaching a performance limit, whereas JFT 4B keeps improving.
BatchEnsemble improves pretraining.
For vision, we run ablations at the fixed setting of L pretrained with JFT, and we use both B and L sizes for text, which are highly competitive settings.
Figure 5 displays the ranking across tasks for each model.
Methods are applied either during both pretraining and finetuning, or only during finetuning given a pretrained model checkpoint (“BE → Het” means pretraining with BE and finetuning with Het on top).
All the methods displayed improve over a baseline without ensembling or last layer changes (None).
BatchEnsemble is consistently the best for pretraining.
For T5 on text, Plex L outperforms None L and also outperforms Plex B; this indicates the benefit of scale not only in Figure 4’s normalized average score but in their average ranking.
Last-layer methods improve finetuning.
The best ranked models for the vision and language tasks use all of Plex’s ingredients: Het on top of a pretrained BE for vision and
GP on top of a BE for language.
In particular, for T5Plex ablations, BE + GP and BE tend to have the strongest performance.
From more detailed per-dataset analysis in
Appendix E, BE + GP and BE perform well on MNLI and
NaLUE with BE + GP performing slightly better; notably, they outperform even an expensive deep ensemble baseline which also performs well on MNLI and NaLUE.
BE + GP outperforms None on Toxic Comments while a Monte Carlo
Dropout baseline performs best on that task.
A Setup and Downstream Datasets
[Figure] 2. Finetune [Figure] 3. Evaluate
downstream tasks
Ex: JFT, C
Ex: ImageNet, NaLUE
Checkpoint Checkpoint
Ex: ImageNet OOD Calibration,
NaLUE-tail Selective Prediction
Results
Figure 6 describes our overall experimental setup.
We selected a broad suite of 40 downstream datasets under the tasks, each ranging from several thousand to a million examples.
We outline the datasets for each modality.
A.1 Images
We’re motivated to capture datasets spanning natural web images, specialized domains that are likely rare or unseen in large pretrained models, and with both small and large sizes. To do so, we use 11 datasets which we describe below.
• CIFAR-10 is a dataset of web images with a training set of 50,000 examples and a test set of 10,000 examples (Krizhevsky et al., 2009).
Following Dosovitskiy et al.
, we use 99% of the training set for training and 1% for validation.
• CIFAR-100 is a dataset of web images with a training set of 50,000 examples and a test set of 10,000 examples.
Following Dosovitskiy et al.
, we use 99% of the training set for training and 1% for validation.
• ImageNet1K is an image dataset organized according to the WordNet hierarchy, with a training set of roughly 1.2 million examples and a test set of 50,000 examples (Deng et al., 2009).
Following Dosovitskiy et al.
(2020), we use 98% of the training set for training and 2% for validation.
• RETINA is a set of benchmarking tasks with training and evaluation datasets containing retina scans exhibiting varying degrees of diabetic retinopathy, a medical condition that can result in a loss of eyesight (Band et al., 2021).
We choose RETINA as an example of a difficult transfer task, since retina images are meaningfully different from natural web images used for pretraining RETINA
includes two types of tasks : (i) A “Country Shift” task with an in-distribution evaluation set and an evaluation set exhibiting covariate shift in the input data and (ii) a “Severity Shift” task with an in-distribution evaluation set and an evaluation set containing labels not included in the training data, representing more severe types of diabetic retinopathy than the training labels.
• We use 7 datasets with a range from 1,880 to 8,144 training examples: Describable Textures Dataset (Cimpoi et al., 2014), UC Merced (Yang & Newsam, 2010), Caltech 101 (Fei-Fei et al., 2004), Oxford-IIIT Pets (Parkhi et al., 2012), Colorectal Histology (Kather et al., 2016), Caltech-UCSD Birds 200 (Welinder et al., 2010), and Cars196 (Krause et al., 2013).
Distribution shift is a common challenge for image problems, and so we cover multiple types for a total of 19 datasets.
Table 2 provides an outline.
Most notable, there is little work for evaluating label uncertainty, so we propose a large-scale dataset which we call ImageNet ReaL-H . ImageNet ReaL recollects human ratings for the original ImageNet test set (Beyer et al., 2020), and we use its raw data of individual ratings to construct a label distribution representing rater uncertainty for each image. Covariate shift .
• CIFAR-10: CIFAR-10-C (Hendrycks & Dietterich, 2019). • CIFAR-100: CIFAR-100-C (Hendrycks & Dietterich, 2019). [Table]
• ImageNet1K: ImageNet-A (Hendrycks et al., 2021c), ImageNet-C (Hendrycks & Dietterich, 2019), ImageNetV (Recht et al., 2019), ImageNet-Vid-Robust, YTTB Robust (Shankar et al., 2021), ObjectNet (Barbu et al., 2019), and ImageNet-R (Hendrycks et al., 2021a).
• RETINA: RETINA’s
Country Shift dataset (Band et al., 2021).
We train models on images of retinas obtained from patients in the United States (EyePACS) and evaluate trained models on images of retinas obtained from patients in India using different collection equipment (APTOS). . Semantic (class) shift [Table]
• CIFAR-10: CIFAR-100, SVHN. • CIFAR-100: CIFAR-10, SVHN.
• ImageNet1K: Places365.
• RETINA: RETINA’s Severity Shift dataset (Band et al., 2021).
We train models on images of retinas exhibiting no worse than mild diabetic retinopathy, and consider a shifted evaluation dataset with images of moderate diabetic retinopathy or worse. The evaluation data contains features not contained in the training images, such as vitreous hemorrhages. The motivation for this shift is that images of retinas with more severe retinopathy are relatively scarce and that it is likely for a model to be trained only on more widely-available images of retinas exhibiting mild cases of diabetic retinopathy.
Label uncertainty . We use CIFAR-10H (Peterson et al., 2019) which captures human uncertainty over labels for CIFAR-10 dataset.
We also construct a larger-scale variant, which we call ImageNet ReaL-H. Individual human ratings were recollected for the original ImageNet test set, available as raw data from ImageNet ReaL
(Beyer et al., 2020), and we use them to newly construct soft label targets. Subpopulation shift .
We use Semantically Partitioned CIFAR-10/100 (Yuan et al., 2022) for vision subpopulation shift.
CIFAR-10/100 test data is partitioned into semantically similar subpopulations, where each subpopulation has its own data-generating distribution sampled from a meta subpopulation distribution.
We aim to improve predictive performance on tail subpopulations.
A.2 Text
For text, we consider real-world decision making tasks that are known to deploy machine learning models: natural language inference, toxic comments detection, and conversational language understanding. Natural language inference and toxic comments are binary classification tasks that map a (pair of) natural language sentences to a binary category: entailment or no entailment, and toxic or non-toxic, respectively. Conversational language understanding is a task common in chatbot design, where the model maps a natural language query to a multi-token prediction of user intents: for example, “I want to order dinner using Uber Eats” → 3-token prediction of (FoodDelivery
, Uber, Order • For natural language inference, we use the Multi-Genre Natural Language Inference (MNLI) corpus which consists of 433k sentence pairs from a diverse collection of genres (fiction, government report, news magazine articles, etc.) (Williams et al., 2017).
• For toxic comments detection, we use the WikipediaTalk corpus (Wulczyn et al., 2017) which is composed of roughly 200k English Wikipedia talk page comments between Wikipedia editors across the world.
• For conversational language understanding, a large-scale corpus for evaluating uncertainty quantification is lacking.
We propose a new dataset Natural Language understanding Uncertainty Evaluation (NaLUE) that is a relabelled and aggregated version of three large NLU corpuses: CLINC150 (Larson et al., 2019), (Zhang et al., 2021) and [Table]
HWU64 (Liu et al., 2021).
NaLUE contains 50k+ utterances spanning 18 verticals, 77 domains, and roughly 260 intents.
For this task, the model needs to map each utterance to a 3-token sequence of (vertical name, domain name, intent name).
In terms of data distribution, MNLI has a balanced distribution both across the genre and across the label class.
NaLUE exhibits a slight skewness toward some popular domains for chatbot development (e.g, banking customer service requests).
On the other hand, the toxic comments datasets often exhibit extreme label imbalance.
For example, ∼ 10% of the examples in Wikipedia Talk Corpus examples have positive labels, since most online content is not toxic (Kivlichan et al., 2021).
Natural language is diverse, fast evolving, and rich in long-tail linguistic phenomena.
Therefore out-of-distribution examples, particularly long-tail subpopulations, are pervasive in the real-world deployment environment.
In Table 3, we outline a total of 7 out-of-distribution challenge sets. Most notably, we construct three new out-of-distribution shifts for NaLUE.
NaLUE-tail contains utterances from 28 low-frequency intents categories in NaLUE. NaLUE Standard-OOS and NaLUE Near-OOS contain utterances that describe out of the scope services, differing in their closeness in distribution to NaLUE.
• the MNLI-mismatched (Williams et al., 2017) data as the OOD set for NLI, which contains sentence pairs fromgenres that are distinct from those in MNLI training data.
• the CivilComments corpus (Borkan et al., 2019) as the OOD set for toxic comment prediction, which consists of one million public comments appearing on approximately 50 English-language news sites across the world.
• HANS
(McCoy et al., 2019) eval datasets for NLI, which contains template-generated examples attacking the surfacelevel heuristics that the neural models are found to rely on when predicting entailment relationships.
• CivilCommentsIdentity (Borkan et al., 2019) for toxic comments, which is a subset of CivilComments that has explicit mention of social identities (e.g., muslim, LGBTQ, etc) that the model are often found to generate mispredictions.
• NaLUE-tail dataset for CLU, which is a subset of NaLUE corresponding to utterances from 28 low-frequency intents categories.
B Details of Overall Reliability Score
In Figure 4, we aggregate all task metrics under a single scalar between 0 and 100.
In order to do this, we must normalize all metrics to be between 0 and 100; we then compute an unweighted average.
Most metrics are already bounded between 0 and 100: for example, accuracy, expected calibration error (we do 100 − ECE so higher is better), calibration AUC, and AUROC.
The one exception are scoring rules such as log-loss and Brier score.
Because the output distributions are discrete, log-loss has a lower bound of 0 and an upper bound given by the highest entropy distribution (uniform).
Therefore we rescale scoring rule values based on their lower and upper bounds so that they’re now between 0 and 100 and so that higher is better.
C Details of Plex ingredients
In this work, we focus on two domains: images and text.
For images, we use a base architecture of Vision Transformer that performs image classification (Dosovitskiy et al., 2020).
For text, we use T5 which uses an encoder-decoder architecture to treat text problems as text input and text output (Raffel et al., 2020).
On top of these architectures, we experiment with the following methods.
BatchEnsemble (BE) .
BatchEnsembles (Wen et al., 2020) approximate deep ensembles (Lakshminarayanan et al., 2017), but reduce their computational and memory costs by sharing weights across the ensemble members.
The weight matrix
W i of any given ensemble member i is written as the Hadamard product of a shared weight matrix W 0 and a local rankmatrix r
i s ⊤ i :
[Equation]
The vectors r and s are commonly referred to as fast weights.
Unless otherwise stated, Plex applies BE to all layers in the last 2 residual blocks of the network.
This idea follows work for mixture of experts (Riquelme et al., 2021).
Spectral-normalized Neural Gaussian Process (SNGP) .
Unlike ensemble approaches, SNGP proposed by Liu et al.
(2020) focuses on improving the uncertainty quality of a neural network given a fixed representation (
a.k.a.
deterministic uncertainty quantification setting (Van Amersfoort et al., 2020).
When applied to a DNN without pretraining, SNGP enhances the DNN uncertainty property by applying spectral normalization to the hidden weights, and replaces the output layer from a dense layer to a random-feature Gaussian process (GP) layer.
That is, given hidden representations h ( x ) , the GP layer enables scalable computation of a GP posterior by applying a random feature approximation ϕ to the predictive function and then a
Laplace approximation to the predictive variance:
[Equation]
where
[Equation]
where ( W , b ) are frozen random weights of the random feature embedding ϕ ( x ) = cos( W h ( x ) + b ) , and Φ ⊤ Φ = (cid:80) i ϕ ( x i ) ϕ
( x i )
⊤ is the covariance of the random feature embedding estimated using the training data.
Liu et al.
(2020; 2022) show that this combined technique improves the model’s awareness of the semantic distance between the test and train examples on the data manifold, leading to improved performance in calibration and out-of-domain detection.
When applied to a large pretrained DNN, we find it sufficient to only use the last-layer Gaussian process (i.e., omit the spectral normalization regularization), as the pre-trained embedding has already provided a semantic-distance-aware representation of the data.
Heteroscedastic last layer (Het) .
Heteroscedastic last layers are designed to model input-dependent label noise/data uncertainty (a.k.a.
aleatoric uncertainty (Kendall & Gal, 2017)) that is present in the data.
We use the Heteroscedastic (Het) last layer introduced by Collier et al.
(2020; 2021) who place a multivariate Gaussian distribution over the logits in a standard
DNN classifier.
A low-rank approximation to the K × K covariance matrix ( K = number of classes/outputs) is made when K is large and (Collier et al., 2021) further develop a parameter efficient version of the method with parameterization inspired by BE to enable scaling to tens of thousands of classes.
Naming of different methods We apply these modifications either during both pretraining and finetuning, or only during finetuning given a pretrained model checkpoint.
None refers to the baseline without ensembling or last layer changes.
“None → GP” means standard pretraining (without any modifications) and just applying GP layer during finetuning.
“BE” means using BE during both pretraining and finetuning.
“BE → Het” means pretraining with BE and finetuning with Het on top.
D Related Work
Prior work has investigated a variety of approaches to improve narrower definitions of reliability.
From the literature, several overarching dimensions arise (Tran et al., 2020)—such as the importance of model and data size (e.g. pretraining); model inductive biases (e.g. architecture and data augmentation); and the combination of multiple models (e.g. ensembles and
Bayesian neural networks).
There is not yet an understanding of how these dimensions interact (and within current literature, it is no surprise that there are contradicting messages) and which of these dimensions provide complementary benefits.
We investigate how each of these dimensions improve reliability and how they can be “composed” to maximize performance.
Modern AI is trending towards training a single large model on a large data set, known as pretraining, and then applying the model to a wide variety of related downstream tasks (Radford et al., 2021;
Brown et al., 2020; Thoppilan et al., 2022;
Kolesnikov et al., 2020).
This often improves over task-specific state-of-the-art in predictive performance, with many considering such large scale models to represent a “paradigm shift” in ML (Bommasani et al., 2021).
Large-scale pre-trained models have also significantly improved state-of-the-art on narrower tasks such as accuracy and calibration under covariate shift, see (Minderer et al., 2021;
Hendrycks et al., 2019a;b) (as well as (Bommasani et al., 2021, Section 4.8) for additional references) and open set recognition (cf.
(Fort et al., 2021; Ren et al., 2021)).
Given these initial promising results, we use large-scale pre-trained models as a building block for investigating reliability.
However, large models can be compute [Table] [Figure] intensive, which warrants revisiting existing recipes; for instance, vanilla deep ensembles, which work well in previous benchmarks (Ovadia et al., 2019; Gustafsson et al., 2020;
Band et al., 2021), might be computationally expensive.
Hence, we focus on scalable modifications to large models such as efficient ensembles and last-layer variants, detailed in Appendix C.
E Summarization of Language Results
We first compare the performance across types of uncertainty methods, fixing the architecture size to T5-base.
We compare performances in prediction, uncertainty calibration, and human-model collaboration, across all datasets (MNLI, NaLUE and Toxic Comments) and all splits (In-domain, OOD, and tail-population).
Table 5 reports the full results, and Figuresummarizes the rankings of uncertainty methods under each type of population shift (in-domain v.s. OOD v.s. tailpopulation).
Among all methods, DE + GP, Plex (i.e., BE + GP), BE, and MC Dropout tend to have the strongest performance.
In particular, DE + GP almost always dominates the other methods on MNLI and NaLUE, and remains competitive in the case of label imbalance (i.e., Toxic Comments).
However, DE + GP is an expensive method that costs x10 more in memory and compute and therefore is not competitive in scale (a more thorough analysis is in ?? ).
On the other hand, among the more efficient, single-model methods, BE and Plex perform well on MNLI and NaLUE (notably, outperform the most expensive DE), while MCD stands out in the Toxic Comments.
The above observations suggest that, when the training example has a simple distribution, quantifying output-layer uncertainty alone is sufficient to attain strong performance.
However, when there are pathologies in the data distribution (e.g., extreme label imbalance), quantifying the uncertainty within the model’s intermediate representations (e.g., via some form of perturbation like BE) becomes important.
We then investigate how a model’s uncertainty performance is impacted by the architecture size.
For model size scaling, we evaluate Plex, None, and MC Dropout, the three best-performing and efficient methods in the previous study.
We evaluate the performance of each method under three progressively larger architectures: T5 S, T5 B, and T5 L, and observe how the behavior changes across the method and with respect to the architecture size.
Table 5 reports the full results, and Figuresummarizes the rankings of uncertainty methods organized by the sizes of the architecture.
As shown, comparing across architecture sizes, we see a larger architecture almost always leads to stronger performance in collaborative performance.
This trend remains largely consistent even when out-of-distribution.
[Table] [Figure]
[Figure]
[Table] [Figure]
[Table] [Figure] [Table]
[Figure] [Table] [Table] [Table]
[Table] [Table]
[Table]
(h) Calibration, ALL. N o n e S M C D S P l e x S N o n e B M C D B P l e x B N o n e
L M C D L P l e x L     (i) Selective Pred., IND. N o n e S M C D S P l e x [Figure]
N o n e
[Figure]
M C D B P
[Figure]
e x [Figure] e L M C D L P l e x
[Figure] (j
[Figure]
Selective Pred., OOD.
[Figure] [Figure]
N
[Figure]
(l) Selective Pred., ALL.