Plex: Towards Reliability Using
Pretrained Large Model Extensions

Dustin Tran^*, Jeremiah Liu^, Michael W. Dusenberry^, Du Phan^,
Mark Collier^, Jie Ren^, Kehang Han^, Zi Wang^, Zelda Mariet^, Huiyi Hu^, Neil Band^, Tim G. J. Rudner^, Karan Singhal^, Zachary Nado^,
Joost van Amersfoort^, Andreas Kirsch^, Rodolphe Jenatton^, Nithum Thain^, Honglin Yuan^†, Kelly Buchanan^†, Kevin Murphy^, D. Sculley^, Yarin Gal^, Zoubin Ghahramani^, Jasper Snoek^, Balaji Lakshminarayanan^

^Google   ^University of Oxford

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 in- and 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, pretrained large model extensions (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 improves the out-of-the-box performance and does not require designing scores or tuning the model for each 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.

$^*$ Corresponding author. Email<EMAIL_ADDRESS>     $^\dagger$ Work done at Google.

Code for training and evaluation is at <https://goo.gle/plex-code> as part of Uncertainty Baselines [Nado et al., 2021]. We also release model checkpoints for ViT-Plex and T5-Plex trained on public data.
Layer implementations use Edward2 [Tran et al., 2018]. Author contributions are detailed in <Ref>.

reliability, large models, uncertainty, robustness, adaptation

§ RELIABILITY AS A GOAL FOR ARTIFICIAL INTELLIGENCE

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., 2021] 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
on applications ranging from medical diagnoses with a human-in-the loop, to using AI to detect online misinformation and toxicity, and to perhaps a pathway for artificial general intelligence.
While hypothesizing about this potential is important, we highlight that the typical 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.

Desiderata for a Reliable model. We propose to simultaneously stress-test the   “out-of-the-box” model performance (i.e. the predictive probability distribution $p(y|x)$) across a suite of uncertainty, robust generalization, and adaptation benchmarks, without any customization for individual tasks.

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 and Proschan, 1975, O'Connor and 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., 2021],
calibration [Dawid, 1982], credibility [D'Amour et al., 2020], and trustworthiness [Avin et al., 2021], each with their own broad and intersecting scopes.

§.§ Desiderata for Reliability

It is common practice in machine learning research to focus 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, and where the tasks only indirectly involve accuracy measures. The emphasis in our paper is on how reliable an AI system is in broad array of 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 they should be able to efficiently 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 (<Ref>):

* Uncertainty involves imperfect or unknown information where it is impossible to exactly describe an existing state [Ghahramani, 2015].
Predictive uncertainty quantification enables practitioners to know when to trust the model’s predictions, thereby enabling graceful failures when the model is likely to be wrong.
A variety of measurements can be used to quantify the quality of uncertainty such as expected calibration error [Naeini et al., 2015, Nixon et al., 2019], which measures how well the model’s confidence

aligns with its accuracy.

Providing a quantification of uncertainty also enables better decision making [Parmigiani and Inoue, 2009]; one popular setting is selective prediction where a model may defer its prediction to human experts when it is not confident.
Another popular task is open set recognition
Earlier work
sometimes referred to this setting as out-of-distribution detection. However, in recent work, the term “out-of-distribution” is used as a unifying term for “non I.I.D” test data which includes any type of shift, e.g., covariate or subpopulation shift. We use the term open set recognition to denote detection for shifts where test inputs belong to semantic classes not encountered in the training data.
, where the model encounters inputs from new classes at test time that were not seen during training, and the goal is to reliably detect that such inputs do not belong to any of the training classes.

* Robust Generalization involves an estimate or forecast about an unseen event [Abraham and Ledolter, 1983, Tran et al., 2020]. Prediction quality 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 and 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 data and efficiently learn with as few labeled examples as possible.  Examples include few-shot learning [Li et al., 2006], 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.

Top row: Examples of Plex's capabilities in vision:
(left) Label uncertainty in ImageNet ReaL-H, demonstrating the ability to capture the inherent ambiguity of image labels assigned by humans.
(middle) Active learning on ImageNet1K, displaying
that Plex achieves significantly higher accuracy

compared to a non-pretrained baseline and with fewer labels.
(right) Zero-shot open set recognition on ImageNet1K vs Places365, showing that Plex can distinguish visually similar images without finetuning.
Bottom row: Examples of Plex's capabilities in language:
(left) Plex enables human+AI collaboration with selective prediction, where the model is able to defer a fraction of test examples to humans. Plex is able to better identify cases where it is likely to be wrong than the baseline, and thus achieves higher Collaborative AUC.
(middle) Plex is robust while a baseline latches onto spurious features such as “destination” and “around the world”.
(right) Plex enables structured open set recognition. This provides nuanced clarifications, where Plex can distinguish cases where only part of a request is not supported.

While there has been initial progress on specific tasks within reliability, there remain critical limitations. First, prior work has typically focused on narrow settings. For example, much focus has been on accuracy and calibration error on ImageNet and its corrupted versions as a benchmark task
[Hendrycks and Dietterich, 2019, Ovadia et al., 2019].

Similarly, the literature on open set recognition has led to the development of methods that focus solely on improving open set recognition, often at the expense of other downstream tasks.

This has led to a fragmentation of research techniques, and it remains unknown how well advances on these individual tasks connect to a broader set of tasks and datasets.

The goal of this work is to move towards a notion of general reliability by seeking models that work well across many decision-making tasks.
rather than needing to design, train, and tune a model for each individual task.

The approach is philosophically inspired by the trend in natural language processing where evaluating on multiple  tasks (cf. GLUE [Wang et al., 2018]) has encouraged the community to focus on developing general-purpose techniques.

ViT-Plex (left) and T5-Plex (right) evaluated on a highlighted set of reliability tasks. We also display the state-of-the-art for each task. ViT-Plex and T5-Plex significantly improve state-of-the-art across multiple tasks. Importantly, Plex unifies reliability performance under one general model for vision and language respectively as opposed to specific techniques for each downstream task.
References are listed in <Ref>.

§.§ Approaches to Reliability

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 size; model inductive biases (e.g., architecture);

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 can be a surprise that combining certain methods is detrimental, e.g., Wen et al., 2020).

We investigate how to combine these dimensions to form an overall model for reliability.

Modern AI is trending towards training a single large model on a large and diverse data set, known as pretraining, and then applying the model to a wide variety of related downstream tasks [Brown et al., 2020, Kolesnikov et al., 2020, Radford et al., 2021, Thoppilan et al., 2022]. 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 pretrained models have also significantly improved state-of-the-art on narrower tasks such as accuracy and calibration under covariate shift [Minderer et al., 2021]

and open set recognition [Fort et al., 2021, Ren et al., 2021]. Given these initial promising results, we use large pretrained models as a building block for reliability.
However, large models are compute-intensive and often operate at a different scale than has been studied in many uncertainty and robustness papers. This warrants revisiting existing recipes in this new context, where we focus on methods that scale efficiently to large models.

§.§ 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. As part of the tasks, we build two new datasets in order to assess areas that we found missing in the literature: ImageNet Real-H for large-scale label uncertainty evaluation; and NaLUE for uncertainty in conversational language understanding.

To improve reliability, we develop ViT-Plex and T5-Plex, building on large pretrained models for vision (ViT; Dosovitskiy et al., 2021) and language (T5; Raffel et al., 2020), respectively.

We train Plex over model sizes up to 1 billion parameters and pretraining dataset sizes of up to 4 billion examples.
<Ref> illustrates Plex’s performance on selected tasks comparing to existing state-of-the-art, which typically is a model specialized for that task.

Plex achieves new state-of-the-art on many of the 40 datasets. Importantly, Plex achieves strong performance across all tasks using the out-of-the-box model output without requiring any custom designing or tuning for each task.

§ EVALUATING RELIABILITY

The model and task pipeline. We experiment with approaches for pretraining; given the pretrained model’s checkpoint, we experiment with finetuning and adaptation; finally, given the finetuned model's checkpoint, we evaluate downstream metrics.

To experiment with a variety of models and tasks, we set up a pipeline for training and evaluation; see <Ref>. It involves three steps: pretrain on a large diverse dataset, finetune on a given dataset’s training split, and evaluate over both in-distribution and out-of-distribution datasets using the appropriate task metrics.

We experiment with a variety of methods to improve reliability during both pretraining and finetuning.

§.§ Tasks for Benchmarking Reliability

We evaluate a model’s reliability using 10 types of tasks, which we categorize as follows. More detailed descriptions of metrics in each task and datasets are provided in <Ref>.

Types of distribution shift
using an illustration of ImageNet dogs. $p_{\textrm{\textsc{ind}}}(\cdot)$ refers to in-distribution during training and $p_{\textrm{\textsc{ood}}}(\cdot)$ is out-of-distribution during testing.

§.§.§ Uncertainty

* Calibration assesses the quality of a model’s predicted confidence over a population [Dawid, 1982]. It quantifies how well model confidence (the predictive probability of correctness) aligns with model accuracy (the observed probability of correctness). We compute expected calibration error [Naeini et al., 2015] and calibration AUROC [Kivlichan et al., 2021]
on 14 image and 10 text datasets.

* Selective prediction jointly assesses a model's predictive performance and quality of uncertainty estimates, by abstaining from making predictions when the model is uncertain [El-Yaniv and Wiener, 2010].

There is no popular metric that's the de facto standard in the literature, so we experiment with two:
Oracle Collaborative Accuracy measures accuracy where any deferred predictions are sent to an oracle as a proxy for human$+$AI collaboration [Kivlichan et al., 2021]; and
selective prediction by rejection rate traces a curve of predictive accuracies over the percentage of examples that can be rejected [Band et al., 2021].
We examine selective prediction 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 [Geng et al., 2020].
This happens in the case of semantic (class) shift, where the test input and label distributions both change, particularly in a structured manner where
the output classes change from train to test.

For an overview of distribution shifts, see <Ref>, where we train on a distribution $p_{\textsc{ind}}(\mathbf{x},y)$ and we evaluate on another distribution.
For example, the training set may consist only of dog images and the new input is a coyote image.
We compute AUROC and use maximum softmax probability as a simple and general detection score.

* Label uncertainty is a type of uncertainty inherent in the data labels. This is a form of data uncertainty—irreducible output noise—which is distinct from uncertainty arising from the choice of model [Dusenberry et al., 2020]. Label uncertainty can arise when, for example, human raters disagree about the label for an ambiguous input.  If human disagreement is encoded as a label distribution, we can directly compare our model's predictive distribution to it.  We analyze label uncertainty on two image datasets.
[We focus on label uncertainty in an OOD evaluation-only context, as collecting label distributions can be expensive and therefore the datasets are small relative to single-label training sets. It is not exclusive to the distribution shift setting as training sets may also carry a label distribution per input.]

§.§.§ Robust Generalization

* In-distribution generalization assesses how well a model can make predictions after finetuning on a downstream dataset. In particular, we examine accuracy, negative log-likelihood, and Brier score on the in-distribution test splits of 5 image and 3 text datasets. For binary classification tasks, we also look at AUROC and AUPRC.

* Covariate shift refers to scenarios where the distribution of inputs changes while the conditional distribution of outputs is unchanged [Sugiyama and Kawanabe, 2012].
For example, the training set may include natural dog images and the new input is a drawing of a dog.
We use the same metrics as those used for assessing in-distribution generalization.

* Subpopulation shift refers to scenarios where the distribution of interest is only part of the full distribution seen during training. A common assumption is that data is sampled from individual subpopulation distributions, which are themselves sampled from a meta population distribution [Santurkar et al., 2020, Yuan et al., 2022]. For example, the training set may include natural dog images and the test set only includes terriers.

In this setting, we aim to improve predictive performance on unseen or long-tail subpopulations.

§.§.§ Adaptation

* Active learning assesses a model’s ability to not only learn over a fixed set of data points, but also participate in knowing which data points to learn from in the first place.
This procedure assesses a model’s label efficiency, where label annotations may be scarce, and so we would like to maximize performance while minimizing the number of labeled data points used

[Settles, 2009]. We assess accuracy over the total number of acquired examples and apply margin sampling for multi-class uncertainty sampling.

* Few-shot learning assesses how well a model can make predictions downstream with only a few training examples [Li et al., 2006].
We use 9 datasets and evaluate the settings of 1-shot, 5-shot, 10-shot, and 25-shot ($x$-shot means $x$ examples per class).

* With few-shot uncertainty, we examine calibration

and open set recognition in the few-shot regime.

We use all 9 datasets for few-shot learning in order to evaluate calibration,

and we use those with OOD datasets

for open set recognition. We also

perform zero-shot open set recognition by using Mahalanobis distance scoring to detect whether an input is out-of-distribution based on the model's representation layer [Lee et al., 2018].

§.§ Datasets

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.

§.§.§ 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., 2021, 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., 2021, 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., 2021, we use 98% of the training set for training and 2% for validation.

* RETINA is a set of benchmarking 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 chose RETINA as an example of difficult transfer, since retina images are quite different from natural web images used for pretraining.
RETINA includes two types of splits: a “Country Shift” split with an in-distribution test set and a test set exhibiting covariate shift;

and a “Severity Shift” split with an in-distribution test set and a test set containing labels not included in the training data, representing more severe types of diabetic retinopathy.

* 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 and 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. <Ref> 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. For ImageNet1K, we use 7 datasets for different covariate shifts: ImageNet-A, ImageNet-C, ImageNet-R, ImageNetV2, ImageNet-Vid-Robust, ObjectNet, and YTBB Robust.

Covariate shift     Semantic shift     Label uncertainty     Subpopl. shift

CIFAR-10     CIFAR-10-C     CIFAR-100, SVHN     CIFAR-10H     SP-CIFAR-10

CIFAR-100     CIFAR-100-C     CIFAR-10, SVHN     —     SP-CIFAR-100

ImageNet1K     7 datasets     Places365     ImageNet ReaL-H     —

RETINA     Country Shift     Severity Shift     —     —

Vision datasets for evaluation on distribution shift.

§.§.§ 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” $\to$ 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], Banks77 [Zhang et al., 2021] and 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, $\sim$10% of the examples in Wikipedia Talk Corpus examples have positive labels, since most online content is not toxic [Kivlichan et al., 2021].

Covariate shift     Subpopulation shift     Semantic (class) shift

MNLI     MNLI-mismatched     HANS     —

WikipediaTalk     CivilComments     CivilCommentsIdentity     —

NaLUE     —     NaLUE-tail     Standard-OOS, Near-OOS

Language datasets for evaluation on distribution shift.

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 <Ref>, 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.

§ PLEX: PRETRAINED LARGE MODEL EXTENSIONS

Plex is the result of an extensive study of the reliability of large pretrained models and their complementarity with existing reliability methods. ViT-Plex and T5-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., 2021] for images in an encoder-only setup where the raw images are effectively tokenized into patches.

* Model size. We investigate 3 scales of the model size in ViT-Plex: Small (ViT-Plex S; $\sim$22 million parameters) has patch size 32, 384 embedding size, 1536 feedforward size, 12 residual blocks, and 6-headed attention; Base (ViT-Plex B; $\sim$87 million parameters) has patch size 32, 768 embedding size, 3072 feedforward size, 12 residual blocks, and 12-headed attention; and Large (ViT-Plex L; $\sim$325 million parameters) has patch size 32, 1024 embedding size, 4096 feedforward size, 24 residual blocks, and 16-headed attention.

For T5-Plex, we consider three model sizes: Small (T5-Plex S; $\sim$77 million parameters) has 512 embedding size, 8 encoder / decoder blocks, and 6-headed attention; Base (T5-Plex B; $\sim$250 million parameters) has 768 embedding size, 12 encoder / decoder blocks, and 12-headed attention; and Large (T5-Plex L; $\sim$880 million parameters) has 1024 embedding size, 24 encoder / decoder blocks, and 16-headed attention.

* Pretraining dataset size. For vision, we scale pretraining from ImageNet-21K 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., 2020, Band et al., 2021]. To apply ensembling 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]. In both ViT-Plex and T5-Plex, we apply no dropout.

* 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 a Gaussian process (GP) last-layer, which improves 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., 2020.

In addition, to model input-dependent label noise in datasets with many output classes we apply the Heteroscedastic (Het) method of Collier et al., 2021.

For detailed background, see <Ref>.

* What to apply in pretraining versus finetuning.
We apply efficient ensembling during both pretraining and finetuning.

For last-layer methods, we find pretraining benefits can be obtained primarily from only applying the method during finetuning, so we restrict them to that setting (detailed in <Ref>).
In addition, due to compute constraints, we exclusively focus on the finetuning-only setting for T5-Plex. That is, T5-Plex models are initialized from the official pretrained T5 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; on 5-shot and higher, training the full model
with gradient descent
can lead to significant gains (for example, 2-3% accuracy boost on ImageNet).
We also experiment with a GP or Heteroscedastic last layer as an alternative to a linear last layer.

§ SUMMARY OF RESULTS AND SCALING TRENDS

<Ref> displays the largest variants of Plex's performance compared to existing specialized state-of-the-art on a diverse collection of reliability tasks. Plex not only sets new state-of-the-art on many tasks

but Plex also unifies reliability performance under one general model for vision and language respectively. Here, we validate several takeaways as we ablate to understand Plex's ingredients.

Plex's performance aggregated across (top) 139 vision task metrics and (middle) 54 language task metrics.
Subplots (bottom) display average performance for individual reliability areas for vision (first 3) and language (last 2).
Compute is the total # of training days for a single TPUv3 core.
Dots represent size (Small, Base, Large), pretrained on JFT for images and C4 for text; I21K denotes pretraining on ImageNet-21K.

Scaling model size improves reliability.

<Ref> displays ViT-Plex and T5-Plex over varying model sizes.

We compute a reliability score which is a normalized average over all task metrics: 139 for vision and 54 for language (see <Ref> for details).
We also display reliability scores for individual reliability areas, which are the averages separately over the uncertainty, robust generalization, and adaptation tasks.

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.

ViT-Plex L with JFT performs better than ViT-Plex L with ImageNet-21K

In <Ref>, 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.

Pretraining Dataset     1M steps     2M steps     4M steps     8M steps

JFT 300M      72.1%     72.9%     73.0%     73.0%

JFT 4B     72.7%     74.4%     74.6%     —

ImageNet 10-shot accuracy consistently improves with a larger pretraining dataset.

(left) Ranking of method ablations over 139 metrics on vision tasks and (right) 54 metrics on language tasks. Each model has a box plot of rankings (lower is better).
Plex's combination of efficient ensembling and last-layer changes ranks best on average.

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. <Ref> 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 (“None$\to$Het” denotes pretraining with no changes 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.

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 T5-Plex ablations,
BE$+$GP and BE tend to have the strongest performance.
From more detailed per-dataset analysis in <Ref>, 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. T5-Plex L also outperforms T5-Plex B, which indicates the benefit of scale not only in <Ref>'s normalized average score but in their average ranking.

Downstream dataset size has no obvious pattern with reliability. <Ref> decomposes the performance of ViT-Plex L by analyzing it as a function of the downstream dataset’s size of training set. That is, we aggregate reliability performance for each training dataset separately, ranging from a size of 1,880 examples with Describable Textures Dataset (dtd) to 1.2 million examples with ImageNet. There is no clear pattern with respect to size. On the other hand, the datasets with lower reliability tend to be more different from the distribution of natural images in JFT: UC Merced is a remote sensing dataset of map areas, and Colectoral Histology is a histology dataset of human colorectal cancer. Pet images are not uncommon in JFT, but most breeds in Oxford IIT Pets aren't in JFT.
This suggests reliability performance is possibly more connected to the distribution shift between pretraining and downstream dataset than the downstream dataset’s size.

Reliability performance of ViT-Plex L, aggregated by each downstream dataset.

§.§ Relationship Between Reliability Tasks

Given that there are many tasks that make up reliability, here we analyze relationships between the individual tasks. We’re motivated by the following question: Is there an inherent underlying evaluation metric that is indicative of reliability? Specifically, can reliability be predicted from pretraining performance, without any downstream finetuning or adaptation?

Correlation of pretraining validation loss with each task (evaluation metric), averaging over all datasets for that task. Most of the tasks correlate highly, meaning upstream performance is a key signal for downstream performance.

In  <Ref>, we analyze the correlation of validation loss on the pretraining dataset with each downstream metric. Most of the tasks correlate highly, meaning that the pretraining performance is an important predictor of downstream performance. The one exception are two uncertainty metrics—calibration error and calibration AUROC—which may make sense intuitively given that calibration as a property about models is not tied to predictive performance.

Interestingly, AUROC, which refers to the metric for open set detection (an uncertainty task), is strongly coupled with predictive performance (Pearson correlation of 0.99!), more than, say, selective prediction performance (OC-AUC); this suggests that for open set recognition, prediction may be more important than uncertainty.
Few-shot accuracy also correlates strongly with pretraining performance, and more examples leads to higher correlation ($25 > 10 \approx 5 > 1$-shot).
Accuracy and NLL correlate less strongly than few-shot: this is likely because they are measured on both in- and out-of-distribution datasets whereas few-shot is only evaluated on in-distribution test splits.

Surprisingly, we can find an even stronger answer to the question: training loss on pretraining data is predictive of reliability. <Ref> shows Pearson correlation of upwards of 0.97 for the reliability score, prediction, and adaptation areas. Uncertainty is the least correlated but still quite strong at 0.76. This suggests that to perform well for reliability, simply fitting the data—that is, having high model capacity and the ability to efficiently train to utilize that capacity—is one of the most essential ingredients.

Score     Score robustness     Score uncertainty     Score adaptation

JFT     -0.977     -0.973     -0.762     -0.982

CIFAR-10     -0.494     -0.909     0.173     -0.376

CIFAR-100     -0.517     -0.965     -0.271     -0.360

ImageNet     -0.566     -0.957     0.038     0.440

Correlation matrix between the training loss on each dataset and the reliability over all vision tasks as well as separately in the three categories. A value of -1.0 means a lower training loss has a perfect linear relationship with a higher score. Training loss on the JFT pretraining dataset has a surprisingly high correlation with reliability. Performance on downstream datasets is predictive of overall robustness but otherwise has less correlation.

§.§ The Effect of Pretraining vs Finetuning

How do the effects we find above differ as model ingredients are applied during the pretraining phase vs finetuning phase?

Here, we analyze the performance of models depending on choices made separately during the two phases.

For example, can the benefits of BatchEnsemble be obtained when only applied during finetuning given a model pretrained without any changes?

To broadcast the weights of a pretrained model $U$ into those of a downstream model $D$, we follow the initialization scheme for $D$ and replace any weights common to $U$ and $D$ by those in $U$.
For example, a BatchEnsemble obtained by broadcasting a pretrained vanilla model will inherit its slow weights from the model, but fast weights will be initialized randomly.

Consistent with all finetuning experiments, the head layer is also reinitialized rather than inherited from the pretrained model.

Upstream and downstream versus downstream-only use of specific architectures on Imagenet. We compare the extent to which a single pretrained deterministic model can be cast into alternative architectures during finetuning.

<Ref> displays results on ImageNet. BatchEnsemble when applied during both pretraining and finetuning strictly outperforms the finetuning-only variant across the 3 metrics. We also ran a naive deep ensemble baseline and find that this is similarly the case. On the other hand, for both GP and Heteroscedastic, applying the method only during finetuning roughly matches the performance of the method when applied during both pretraining and finetuning. Therefore we apply BatchEnsemble for both pretraining and finetuning, and we restrict using last-layer methods for finetuning. See <Ref> for more analysis of the heteroscedastic last layer for vision tasks.

§.§ Ensemble Scaling

Ensembles of neural networks [Hansen and Salamon, 1990, Lakshminarayanan et al., 2017], which aggregate the predictions of several instances of the same model class, provide a remarkably simple, yet effective, way to improve the performance of the base model. This holds true not only for predictive performance

but also, crucially, for robustness and uncertainty quantification [Ovadia et al., 2019, Gustafsson et al., 2020, Wen et al., 2020]. However, ensembles require an increasing amount of compute as one increases the size of the ensemble. We're motivated to understand this axis of model scaling, ensemble size, as it compares to an alternative, popular axis of model scaling: increasing width and depth.

Aggregated reliability score

ImageNet test accuracy

Trading off model size with deep ensemble size. Overall, scaling the single model from Small to Large has a greater influence on the overall performance.

<Ref> shows that naive ensembling consistently provides better performance as the number of ensemble members increases. However, this comes at significant computational cost. Scaling up the model size from S to B or from B to L leads to an $\sim$4x increase in compute, and larger single models tend to outperform the ensembles of 4 smaller models. This motivates the importance of efficient ensembles in Plex: BatchEnsemble adds minimal extra compute and outperforms the baseline without ensembles consistently in the ranking across tasks (<Ref>).

§.§ Distance Awareness

We investigate distance awareness using Gaussian processes.

Key message: ViT-GP does slightly better across tasks (consistently) but the performance improvement is most dominant in low-data (few-shot regime).
Key message: For the choice of applying GP layer to pretraining vs finetuning, there is not a substantial difference (for full-data regime).

We assess ViT vs ViT-GP in three downstream tasks:
Task 1: IND Accuracy, IND ECE
Task 2: Selective prediction measured by rejection vs accuracy diagram
Task 3: OOD detection measured by MSP AUC
All the task evaluations are carried out from low-data regime to full-data regime
X-axis: number of images per class used for finetuning; Y-axis: accuracy against the test split of finetune dataset. Example plot below uses CIFAR-100 as the finetune dataset.

Figure: IID accuracy (left) and ECE (right) vs data size for ViT and ViT-GP
X-axis: number of images per class used for finetuning; Y-axis: accuracy against the validation split of finetune dataset. Example plot below uses CIFAR-100 as the finetune dataset.

Figure 3: Rejection vs Accuracy vs data size for ViT and ViT-GP
X-axis: percent of data rejected by uncertainty; Y-axis: accuracy against a combined test set which consists of the test split of the finetune dataset and the test split of the OOD dataset. Example plot below uses CIFAR-100 as the finetune dataset and CIFAR-10 as the OOD dataset.

(the legend label format: <pretrained>:<finetuned> [number of samples per class])
Figure 4: MSP AUC for OOD detection vs data size for ViT and ViT-GP
X-axis: number of images per class used for finetuning; Y-axis: MSP AUC for OOD detection (IND: the test split of the finetune dataset, OOD: the test split of the OOD dataset). Example plot below uses CIFAR-100 as the finetune dataset and CIFAR-10 as the OOD dataset.

§ RELIABILITY TASK RESULTS

In this section, we examine performance on individual tasks across the areas of uncertainty, robust generalization, and adaptation.

§.§ Uncertainty

To make an accurate assessment of the quality of a model's predictive uncertainty, we consider a set of tasks, each highlighting a different property of models' predictive uncertainty estimates.

§.§.§ Calibration

Expected calibration error averaged over three in-distribution vision evaluation datasets (top) and six out-of-distribution vision evaluation datasets (bottom).

Model types are shown on the $x$-axis.

Expected calibration error averaged over two in-distribution language evaluation datasets (top) and three out-of-distribution language evaluation datasets (bottom).

Model types are shown on the $x$-axis.

Plex improves calibration over a baseline without any changes, especially on out-of-distribution. Scaling model size and pretraining dataset size can improve calibration.

Intuitively, calibration is about the “accuracy” of a model's uncertainty estimates.
It reflects how well the model’s confidence, which quantifies the predictive probability of correctness, aligns with the model's accuracy, which is the observed probability of correctness [Dawid, 1982]. We investigate two different measures of calibration.

Expected Calibration Error
Expected calibration error (ECE) [Naeini et al., 2015] is a binning metric that computes the average difference in confidence and accuracy within different confidence bins.
We evaluate ECE of vision models on 3 in-distribution datasets—CIFAR-10, CIFAR-100, ImageNet—and 6 out-of-distribution—CIFAR-10H, ImageNet ReaL-H, ImageNet-A, ImageNet-C, ImageNet-R, and ImageNet-V2. We also evaluate ECE on T5-Plex over 2 in-distribution datasets—MNLI-matched and NaLUE—and 3 out-of-distribution datasets—MNLI-mismatched, HANS, and NaLUE-tail[We excluded toxicity detection datasets which exhibits severe label imbalance. For these datasets, ECE as a population-average metric is not a suitable measure. For example, a naive model that always predict the majority label will trivially achieve high accuracy and low ECE.].

<Ref> displays the average calibration error of vision models on both in- and out-of-distribution evaluation datasets.
Calibration errors on in-distribution are relatively low in general (less than 2.5% ECE).

Plex improves calibration error over using no changes (None) by roughly a reduction in half on average; on out-of-distribution, it is roughly a 3% improvement whether the models

both use JFT or both use ImageNet21K.

Regarding model size, we find calibration error improves over S, B, and L. Another significant difference is in pretraining dataset size from ImageNet-21K to JFT, where we find JFT models consistently outperform ImageNet-21K pretrained models.

A similar trend can be observed in the language domain.
As shown in <Ref>, compared to the baseline without changes (None), Plex improves the model's calibration performance by roughly 3% on in- and out-of-distribution.

On the other hand, a model's uncertainty performance is impacted more by the type of uncertainty methods than the model size. For example, the Plex B model is on average stronger than the None L model.

Models      C10     C100     ImageNet

Plex S     93.6%     90.9%     86.7%

Plex B     93.6%     92.4%     87.4%

Plex L     89.2%     93.7%     87.7%

Models      MNLI     NaLUE     Toxic Comments

Plex S     76.4%     92.9%     90.5%

Plex B     82.7%     93.4%     90.4%

Plex L     86.0%     88.9%     91.3%

Calibration AUROC, measured on the in-distribution test split of each vision dataset (left) and text dataset (right).

Calibration AUROC.

Calibration AUROC considers the binary classification problem of predicting from a model's uncertainty on a given input whether its prediction for the same input (i.e., the class associated with the highest class probability) is correct.
It is the area under the ROC curve for this binary classification task, and it

aggregates the model's predictive performance over all possible confidence thresholds.
Comparing to ECE, Calibration AUROC only evaluates the uncertainty score’s ranking performance, that is, whether a model consistently assigns higher uncertainty to incorrect predictions rather than correct predictions, and such an uncertainty score can be used as a signal for prediction correctness with good precision-recall and ROC performance [Krishnan and Tickoo, 2020, Kivlichan et al., 2021]. <Ref> presents Calibration AUROC on vision and text datasets. We see a similar pattern as ECE where larger models generally perform better, but the pattern is slightly weaker. For example, Plex L performs best on only 4/6 datasets.

§.§.§ Selective Prediction

The model extensions in Plex provide a significant improvement for selective prediction, enabling the ability to predict at much lower error rates when deferring just a small fraction of predictions.

Deep learning models are typically evaluated under the lens of average predictive accuracy. However, this doesn't account for the real-world cost of mistakes in deployed models today.  Often the risk associated with a mistake can outweigh the benefit of being correct and thus, in expectation, it can be better to not predict at all under some confidence level or defer to a more expensive procedure (e.g., a human expert).

This motivates selective prediction, which includes predictive performance as part of a larger decision-making scenario in which the model may abstain from making certain predictions [El-Yaniv and Wiener, 2010].

Selective prediction performance can be evaluated using several approaches, and there has not been much standardization in the literature. We investigate two that use model uncertainty, through joint human-model collaboration and through rejection rates.

Models      C10     C100     ImageNet

Plex S     0.986     0.960     0.940

Plex B     0.991     0.960     0.960

Plex L     0.999     0.980     0.980

Plex L I21K     0.997     0.970     0.960

Oracle Collaborative AUROC with a review fraction of 0.5% of all predictions.

Default pretraining dataset is JFT and I21K signifies ImageNet21K.

Selective Prediction by Model Uncertainty using Human–Model Collaboration.

Many real world use cases of AI allow for a model to defer a subset of predictions to a human expert.
Oracle Collaborative Accuracy and AUROC measure the performance of an oracle–model collaboration system, where the oracle acts a proxy for human experts in order to automate evaluation.
The model sends predictions with high predictive uncertainty to an oracle subject to a fixed referral budget (e.g., only 1% of all queries can be referred to the oracle).
We compute Oracle Collaborative Accuracy and AUROC over a range of budgets and evaluate on both vision and language modalities.

For vision, <Ref> displays Oracle Collaborative AUROC with a review budget of 0.5% on three vision datasets over different variants of Plex.

Surprisingly, the AUCs are quite high: on ImageNet for example, Plex L attains 0.98, and the models all achieve greater than 0.9 across the datasets.

This implies that the ability to defer can enable lower error, which can be useful for higher-risk applications.  Plex consistently performs best, even when taking into account different model and dataset sizes.  We also examined the metric over review budgets of 1%, 2%, and 5%, and found the results to be unchanged. This suggests that the results on 0.5% are the model's limiting performance with an oracle in the loop.

In-domain, by method

OOD, by method

Subpopulation, by method

In-domain, by size

OOD, by size

Subpopulation, by size

Ranking performance in selective prediction (lower is better) across different methods (<ref>-<ref>) and model sizes (<ref>-<ref>). Each model's box plot is its ranking over multiple datasets: ranking is done by AUROC on MNLI and Toxic Comments, and by Accuracy on NaLUE (AUROC is not well-defined for multi-class classification).

For language, <Ref> reports results across three different model sizes (S, B, L) and for eight methods: a baseline without changes (None), MC Dropout (MCD), Deep Ensemble (DE), Gaussian process (GP), Batch Ensemble (BE), Heterostochastic (Het) and two method combinations Gaussian process ensemble (DE+GP) and BE-GP (Plex) (see <Ref> in the Appendix for detailed results). We show the rankings of these methods under different types of test data distribution (i.e. in-domain, OOD and tail-population). We first see that across different methods, DE$+$GP, BE$+$GP (Plex), BE, and MCD 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 10x more in memory and compute and therefore is not competitive in scale (a more thorough analysis is in <Ref>). On the other hand, among the more efficient, single-model methods (i.e., Plex, BE and MCD), 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 examples are drawn from a relatively 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 and high label noise), quantifying the uncertainty within the model's intermediate representations (e.g., via some form of perturbation like BE) becomes important.

Next we investigate how a model's performance is impacted by the model size.
For model size scaling, we evaluate BE$+$GP (Plex), MCD, and None.

We evaluate the performance of each method under progressively larger architectures, S, B, and L, and observe how the behavior changes across the method and with respect to the architecture size.
summarizes 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 collaboration. This trend remains largely consistent even under distributional shifts and in tail groups. On the other hand, a model's uncertainty performance is impacted more by the type of uncertainty methods. That is, the MCD and BE$+$GP models are on average stronger than None models, regardless of architecture size.
Finally, within larger architectures (i.e., T5 B and T5 L), BE$+$GP generally outperforms MCD, and MCD outperforms None. This trend is broken in two situations, (1) in Toxic Comments, MCD strongly outperforms all other architectures, (2) in small architecture, BE-GP is often the poorer performing model, and None achieves the best performance in multi-token prediction problems (NaLUE).
Finally, between the larger architectures (i.e., T5 B v.s. T5 L), Plex (BE$+$GP) generally outperforms MCD, and MCD outperforms None. This trend is broken in two situations, (1) in Toxic Comments, MCD strongly outperforms all other architectures, (2) among the small architectures, the None baseline achieves the best performance in multi-token prediction problems (i.e., NaLUE).

Selective Prediction Pipeline. Given an input, a model returns a prediction and uncertainty estimate.
If the uncertainty estimate exceeds a threshold $\gamma$ (indicating a level of uncertainty corresponding to a rejection rate $\tau$ that reflects the capacity for expert reviews), the diagnosis is referred to a human expert; otherwise, it is processed with no further review.
The schematic is based on Figure 2 in Band et al., 2021.

Selective Prediction by Model Uncertainty using Rejection Rates.

Rejection AUC measures a model's performance in the scenario where it is permitted to not predict on some fraction of the data for which the model is most uncertain, e.g., up to 10% of all queries.  Unlike the collaboration metric, it does not assume the query is passed to an oracle as rejected inputs may be subject to further review for broader decision making such as to improve data collection.
Rejection AUC is computed from predictive uncertainty (e.g., predictive entropy, predictive variance, confidence) and a chosen predictive performance evaluation metric (e.g., accuracy, AUROC, AUPRC) at different rejection rates.
For a rejection rate of $\tau \in [0, 1]$, the $\tau \times 100$% of data points in the evaluation set for which the model is most uncertain are identified.
Those inputs are then rejected, and we assess the predictions on the remaining $(1 - \tau) \times 100$% of data points.
Accuracy–Rejection AUC is given by computing the area under the Accuracy–Rejection Rate curve.
Additional results for AUROC and AUPRC—Rejection AUCs can be found in <Ref>.

Country Shift: Out-of-Distribution

Severity Shift: Out-of-Distribution

Selective Prediction on RETINA (Out-of-Distribution).
The Country Shift task is concerned with diagnosing diabetic retinopathy from retina scans

obtained using different medical equipment and a different patient population.
The Severity Shift task is concerned with detecting levels of severity of diabetic retinopathy more severe than those included in the training data.
Accuracy is used as the evaluation metric and model predictive entropy is used as the per-example uncertainty estimate.

Plex L
None L
Het L
GP L

2*Country Shift

2*Severity Shift

Accuracy–Rejection AUC on RETINA Country and Severity Shift datasets.

<Ref> displays overall performance for a selection of finetuned models that were pretrained on ImageNet-21K, and <Ref> displays the full accuracy–rejection curves for the Country and Severity Shift out-of-distribution tasks.

We find that that Plex generally improves performance over using no changes on both datasets on three out of four evaluation tasks and that Het L performs as well or better than Plex L on all four tasks.

Examining the quality of different models' predictive uncertainty estimates for out-of-distribution evaluation more carefully, we find that on the Country Shift out-of-distribution (covariate shift) task, all models exhibit non-monotonically increasing accuracy–rejection rate curves.
Since the rejected examples are selected based on the models' predictive uncertainty estimates, this behavior indicates that beyond certain rejection-rate thresholds, the models tend to be underconfident about correct predictions and overconfident about mistakes.
We do not observe this pattern on the Severity Shift out-of-distribution (semantic shift) task or on the Country and Severity Shift in-distribution tasks (see <Ref>).

§.§.§ Open-set recognition (OSR)

A simple scoring technique (maximum softmax probability) works well across open set recognition problems. Over image and text datasets, Plex provides a consistent improvement for new state-of-the-art, and without custom changes per dataset.

Open-set recognition, also called out-of-distribution detection or anomaly detection, aims to detect samples from new classes that are not included in training. In contrast to OOD generalization where the test example belongs to the same in-distribution training classes $(x, y)$, $y \in Y_\textsc{ind}$, OSR aims to detect test example $(x, y)$, where $y \in Y_\textsc{ood}$.

To evaluate OSR performance, we design an uncertainty score $U_{\theta}(x)$ that indicates the likelihood of an input $x$ being OOD, and $\theta$ denotes the classification model's parameters.
As a default for all OSR problems, we apply the most commonly used uncertainty score: $1-\text{MSP}$, where MSP is the maximum over predicted softmax probabilities [Hendrycks and Gimpel, 2017].

A mixture of in-distribution examples and OOD examples are used as the test set. For each test example, the model generates an uncertainty score.
Comparing the uncertainty score with its ground truth OOD label (0 indicates the example is from in-distribution data, and 1 indicates the example is from OOD data), we use AUROC to measure how well the uncertainty score separates the two groups. Note that the model is not finetuned using any OOD examples; test OOD examples are only used for evaluation.

We discuss results on alternative uncertainty scores in <Ref>.

In-Distribution                                     Out-of-Distribution
3*Far-OOD      2*Vision     CIFAR-10                                                     SVHN

CIFAR-100                                                    SVHN

Language                    2lNaLUE Standard Out-of-Scope
4*Near-OOD     3*Vision     CIFAR-10                                                     CIFAR-100

CIFAR-100                                                    CIFAR-10

ImageNet2012                                                 Places365

Language                    2lNaLUE Near Out-of-Scope

The open set recognition tasks studied in this section.

Comparing models’ OSR performance across tasks.
We evaluate the performance of Plex and multiple baseline models.
Based on the similarity between the in- and out-of-distribution data, we can classify OSR tasks under two categories (<Ref>): (a) far-OOD detection and (b) near-OOD detection.
In <Ref>, we observe that first, larger models have better performance.
Second, among all the models, Plex performs the best for most takes.

Third, for ImageNet2012 vs Places365, models pretrained with ImageNet-21K (I21K) outperform the same models but pretrained on JFT.

2*MSP     2c|Far OOD                                  3cNear-OOD

C-10 vs SVHN     C-100 vs SVHN     C-10 vs C-100     C-100 vs C-10     I2012 vs Places
None                      0.996                          0.932                           0.986                              0.929                              0.828

None (I21K)                 0.996                          0.913                           0.981                              0.924                              0.838

None$\to$GP               0.997                          0.931                           0.986                              0.927                              0.824

None$\to$Het              0.996                          0.925                           0.984                              0.948                              0.827

None$\to$DE          0.996                          0.923                           0.986                              0.940                              0.831

None$\to$BE               0.996                          0.934                           0.982                              0.924                              0.822

GP                       0.995                          0.938                           0.981                              0.927                              0.828

Het                      0.996                          0.943                           0.979                              0.948                              0.828

DE L                  0.996                          0.937                           0.987                              0.946                              0.834

BE L                  0.996                          0.927                           0.987                              0.940                              0.827

Plex L (I21K)           0.996                          0.932                           0.983                              0.934                              0.842

Plex L           0.997                          0.938                           0.988                              0.954                              0.831

Plex B                  0.995                          0.925                           0.975                              0.906                              0.801

Plex S                  0.990                          0.882                           0.957                              0.864                              0.773

The MSP based OOD AUROC of all types of models across the five tasks.

Open-set intent detection based on large language pretrained models.
To study OSR performance of different models in the language domain, we design an intent detection task for detecting natural utterances that are out of the scope (OOS) services. For the in-distribution dataset NaLUE, the model needs to map each utterance input $x$ to a 3-token sequence of $y=(y_1,y_2,y_3) = (\text{vertical name}, \text{domain name}, \text{intent name})$. To study both far-OOD performance and near-OOD performance, we construct two out-of-the-scope datasets for NaLUE:

* NaLUE Standard-OOS: completely out-of-domain queries, based on CLINC150-OOS. The standard-OOS queries do not share any of the (vertical name, domain name, intent name) with the in-domain query.

* NaLUE Near-OOS: in-domain, out-of-scope queries. It is created by sampling 20% intent from each known domain as OOD data. The near-OOS examples can share the same vertical name / domain name (but definitely different intent name) with the in-domain query.

The model outputs a sequence, and we compute the uncertainty score based on the conditional softmax probability $p( y_{l} | y_{<l} ,x, \theta)$, using the conditional entropy. The conditional entropy is a standard uncertainty score for sequences [Malinin and Gales, 2021],

\begin{align}
\operatorname{uncertainty}(y|x,\theta)= - \frac{1}{L} \sum_{l=1}^L \mathbb{H}( y_{l} | y_{<l} ,x, \theta),
\end{align}

$\mathbb{H}( y_{l} | y_{<l} ,x, \theta)= - \sum_{k=1}^K p(y_l=k|y_{<l}, x, \theta) \log\left(p(y_l=l | y_{<l}, x, \theta) \right)$ is the conditional entropy based on the conditional distribution $p(y_l=k|y_{<l},x, \theta)$, and  $k=1\ldots K$ is the total number of intent names.
The uncertainty score is compared with the ground truth OOD label. An input $(x, y_1,y_2,y_3)$ is labeled as OOD if any of $y_l$ belongs to the OOD labels.

In <Ref>, models including None, GP, Het, MCD, BE, BE-GP, Deep Ensemble (DE), DE-GP are evaluated for their detection performance on Standard-OOS and Near-OOS. BE and DE-GP outperform the other models, followed by MCD.

2cStandard-OOS     2cNear-OOS

Model       AUROC               AUPRC              AUROC             AUPRC
None       0.987               0.941              0.877             0.535

Het       0.981               0.932              0.854             0.491

GP        0.997               0.987              0.884             0.554

MCD       0.998               0.989              0.917             0.616

BE        0.999               0.997              0.909             0.627

Plex     0.997               0.986              0.890             0.543

DE        0.991               0.955              0.884             0.558

DE-GP     0.999               0.997              0.916             0.625

OOS intent detection performance with AUROC and AUPRC for different models.

§.§.§ Label Uncertainty

We propose label uncertainty as an important challenge, with a new large-scale dataset and propose an evaluation metric. The heteroscedastic last-layer particularly helps for label uncertainty.

It is typical to encounter label noise in non-academic datasets. However, benchmark datasets are often carefully curated to reduce label noise due to unreliable or disagreeing annotators.

We use the existing CIFAR-10H dataset, which is a relabelled version of the CIFAR-10 test set with on average $>$50 crowd-sourced noisy labels per input [Peterson et al., 2019].

There is a notable lack of label uncertainty datasets in the literature, and we found limiting performance on CIFAR-10 as our models achieve $>$99% accuracy on the original CIFAR-10 test set.

Therefore we build a larger dataset called ImageNet ReaL-H, which leverages the raw annotations of ImageNet ReaL [Beyer et al., 2020]. When raw labels were available for an image we followed the same averaging procedure as for CIFAR-10H to produce a soft label. There are some cases where no raw annotations were available. These cases correspond to an image where a set of ML models all agreed with the original ImageNet label for the image and thus the image was not sent to human annotators for re-labelling. In these cases we took the one-hot ImageNet label as the soft label (equivalent to all human annotators agreeing the original ImageNet label was correct).

To evaluate a model's ability to capture label uncertainty, consider a divergence measure $\mathcal{D}$ between probability distributions,

\begin{align}
\mathbb{E}_{x\sim p_{\textsc{data}}(\mathbf{x})} \Big[ \mathcal{D}\Big( p_{\textsc{data}}(y\mid\mathbf{x})~\|~ p_{\textsc{model}}(y\mid\mathbf{x}) \Big) \Big]
\approx
\frac{1}{N} \sum_{n=1}^N \mathcal{D}\Big( p_{\textsc{data}}(y\mid\mathbf{x}_n)~\|~ p_{\textsc{model}}(y\mid\mathbf{x}_n) \Big),
\end{align}

where $p_{\textsc{data}}(\cdot)$ is the true label distribution for the data point and $p_{\textsc{model}}(\cdot)$ is the model's distribution. This evaluation metric only reaches 0 when the model captures the true label distribution for every example. We use KL divergence and approximate the label distribution using an empirical distribution over the multiple label samples per input.

KL divergence on ImageNet ReaL-H (left) and CIFAR-10H (right) for None, None upstream$\to$heteroscedastic downstream, BatchEnsemble upstream$\to$BatchEnsemble L + heteroscedastic downstream, and Deep Ensemble models.

We observe in <Ref> that the heteroscedastic models (Het in finetuning-only and Plex) outperform the None model.

Plex does best on ImageNet ReaL-H excluding DE which uses more compute, but Plex performs worse on CIFAR-10H.

§.§ Robust Generalization

An important aspect of model reliability is it's ability to make accurate predictions when the test data distribution changes.  In this section, we examine model robustness to different forms of data distribution shift.  We look at performance both in-distribution and under two out-of-distribution shifts: covariate shift and subpopulation shift.[
There are 2 distribution shift types which we do not cover here: semantic (class) shift and label uncertainty. Semantic (class) shift is not possible to measure predictive performance under without an open vocabulary model such as recent image-text models [Radford et al., 2021]. Therefore we restrict studies of class shift to the open set recognition task (<Ref>). Label uncertainty is covered in <Ref>.

§.§.§ In-Distribution Generalization

Ablation experiments show that pretraining does not help for diabetic retinopathy diagnosis on images, but the Plex model extensions do give a benefit.  Across language tasks, pretraining, model scale, and the Plex extensions all contribute to better in-domain generalization.

Predictive performance on downstream datasets, with (top) accuracy and (bottom) negative log-likelihood over CIFAR-10, CIFAR-100, and ImageNet respectively as a function of the model scale (number of model parameters).  The blue dotted line is the Pareto frontier.  Note, while number of parameters is a useful comparison of model scale (e.g. memory footprint) it doesn't capture all aspects of scale such as computational overhead, training time, etc.

We look at accuracy and NLL across the in-distribution test splits of each dataset after finetuning on their respective training splits.

Vision Tasks
For RETINA in <Ref>, we compare Plex to the best-performing ResNet-50 baseline results based on a wide array of uncertainty quantification methods [Gal and Ghahramani, 2016, Blundell et al., 2015, Rudner et al., 2021, Rudner et al., 2022, Dusenberry et al., 2020, Farquhar et al., 2020] reported in Band et al., 2021 and find that pretraining on neither I21K or JFT results in an improvement in predictive performance across all evaluation metrics compared to the state-of-the-art ResNet-50 trained from scratch.
This failure may be due to the fact that the retina scans used for the diagnosis tasks differ significantly in appearance from the images in the pretraining datasets and as such may be too dissimilar to convey meaningful inductive biases into the finetuned neural network.  Nevertheless, the Plex extensions do still provide a benefit over the pretrained model.
See <Ref> for CIFAR and ImageNet.

Plex L
Plex L
None L
None L

(Single Model)

Results on the RETINA Country Shift task (for in-distribution evaluation).
The SoTA (Ensemble) outperforms Plex L pretrained on both ImageNet-21K or JFT on all evaluation metrics but ECE, where Plex L pretrained on ImageNet-21K performs best.

Language Tasks In the Appendix, <Ref> summarizes the performance of different uncertainty methods, and <Ref> evaluates the models' performance across three different scales (i.e., T5$_{\texttt{small}}$, T5$_{\texttt{base}}$, T5$_{\texttt{large}}$). Finally, <Ref> in the Appendix visually summarize the relative performance ranking across methods and scales. Comparing across methods, uncertainty methods (specifically, BE, DE, GP, and their combinations) provides improved generalization performance when compared to the None baseline.
Specifically, DE+GP provides the strongest performance across almost all tasks. Among the more efficient methods, Plex (BE+GP) consistently ranked among the top performing methods, while the performance of GP, MCD and BE varies across the datasets (<Ref>). Comparing across model scale, the methods' generalization performance generally improves as the model scale increases, with Plex (BE+GP) the method benefiting the most strongly from scaling, delivering on average the strongest performance at T5$_{\texttt{large}}$.

§.§.§ Covariate shift

In general, Plex significantly improves performance across metrics under covariate shift, even when it doesn't outperform on in-distribution test data. Model size (larger models perform better under covariate shift) and pretraining data are consistently major factors.

In order to be robust under covariate shift, a model should be able to reliably make correct predictions on noisy, corrupted, and otherwise distribution-shifted inputs. In this section, we evaluate robustness using a variety of datasets that exhibit different types of covariate shift.

Plex L
Plex L
None L
None L

(Single Model)

Predictive performance on the RETINA Country Shift task (for out-of-distribution evaluation).
Plex L pretrained on either ImageNet-21K or JFT outperforms the SoTA Ensemble on all evaluation metrics.

RETINA under Country Shift
RETINA's Country Shift task is concerned with diagnosing diabetic retinopathy from retina scans obtained using different medical equipment and a different patient population.
In <Ref>, we compare Plex to the best-performing ResNet-50 baseline results based on a wide array of methods to improve uncertainty quantification [Gal and Ghahramani, 2016, Blundell et al., 2015, Rudner et al., 2021, Rudner et al., 2022, Dusenberry et al., 2020, Farquhar et al., 2020] reported in Band et al., 2021 and find that pretraining on either I21K or JFT results in an improvement in predictive performance across all evaluation metrics compared to the state-of-the-art ResNet-50 trained from scratch.
This result is in contrast with corresponding results on the Country Shift in-distribution task, for which the best-performing models trained from scratch outperform all pretrained models.  In this case, pretraining does improve generalization under covariate shift even when training from scratch yields better predictive performance on in-distribution evaluation tasks.  The Plex entensions also yield a significant improvement both in accuracy and log-likelihood.

Predictive performance with, starting from top row, (1) accuracy,
(2) NLL, and (3) Brier score
across 4 out-of-distribution datasets (columns): ImageNet-C, ImageNet-A, ImageNet-R, and ImageNet-V2 respectively as a function of model scale (number of parameters).

Plex L trained on JFT consistently outperforms other methods and pretraining datasets, and increasing model scale consistently leads to improved metric performance under corrupted shifts.

ImageNet Covariate Shift <Ref> displays performance over multiple types of covariate shift as a function of model scale (number of parameters):
image corruptions (ImageNet-C [Hendrycks and Dietterich, 2019]),
natural adversarial examples (ImageNet-A [Hendrycks et al., 2021]),
artistic renditions (ImageNet-R [Hendrycks et al., 2021]),
and “natural” shift obtained by following ImageNet's data collection procedure to obtain a new test set (ImageNet-V2 [Recht et al., 2019, Taori et al., 2020]).

We find that BE ViT L outperforms None methods, that pretraining on JFT outperforms models trained on other pretraining datasets, and that increasing model scale consistently leads to improved performance under shifts across metrics. Unlike in some sections, for performance under shifts, we consistently find that pretraining on JFT outperforms that on ImageNet-21K.

Stability measured with Anchor Accuracy and Accuracy PMK on ImageNet-Vid-Robust and YTBB-Robust. We find that on the main pm-k accuracy, None, Plex L and BE L models perform similarly, although BE L outperforms slightly.  On anchor accuracy, None and BE perform similarly.  We also see that performance improves with model size.

Stability ImageNet-Vid-Robust and YTBB-Robust [Shankar et al., 2021] are benchmark datasets for measuring robustness to natural perturbations in images by using subsequent frames from video sources with human curation. On the robust accuracy (pm-k accuracy) metric, in <Ref>, None and BE L models perform equally well, and performance improves with model size. Models pretrained on JFT outperform those pretrained on I21K, with BE B JFT even outperforming None L I21K and BE L I21K.

Generalization, by method.

Sel. Predict., by method.

Calibration, by method.

Generalization, by size.

Sel. Predict., by size.

Calibration, by size.

T5-Plex model's ranking performance in generalization under domain shift compared to different methods (<ref>-<ref>) and across architecture sizes (<ref>-<ref>).

Language Style and Topical Drift In natural language applications, it is common for a trained NLP model to be deployed to an environment that exhibits significant style and topical drift when compared to training data. For example, a NLI model trained on written literature is used to analyze in-person dialogues, or a toxic comment detection model trained on U.S. web forums is deployed to non-U.S. news websites [Williams et al., 2017, Kivlichan et al., 2021]. To understand the robustness of large uncertainty models with respect to these common language drifts, we evaluate the models' prediction performance on out-of-distribution splits of the NLI and Toxic Comment tasks (MNLI-mismatched and Civil Comments, respectively). Specifically, MNLI-mismatched represents a low-degree shift scenario where models trained on written and spoken genres (e.g., government report, fiction, phone conversations) are tested on different subtypes of similar genre (e.g., fundraising letters, in-person conversation, 9/11 calls). On the other hand, Civil Comments represents a high-degree shift scenario where models trained on web forum conversations before 2015 were tested on news websites comments from a different period (2015-2017). <Ref> compares the performance of nine models based on T5$_{\texttt{base}}$, and also that of three representative efficient methods (None, MCD and Plex) across model scales (i.e., T5$_{\texttt{small}}$, T5$_{\texttt{base}}$, T5$_{\texttt{large}}$) (see Appendix Table <ref>-<ref> for detailed results). As shown, the models' performance in out-of-distribution generalization correlates well with their in-domain performance, with DE-GP being the most competitive ensemble model, and BE / Plex the most competitive efficient method for MNLI, and MCD the most efficient method for Toxic Comments. Comparing between model scales, larger models in general lead to stronger out-of-distribution performance, with Plex being the most competitive method for moderate-to-large size models (T5$_{\texttt{base}}$, T5$_{\texttt{large}}$).

§.§.§ Subpopulation shift

Plex improves generalization under subpopulation shift relative to models omitting pretraining, efficient ensembling, or last layer changes, for both language and vision datasets.

We next study performance on shifts where data is composed of subpopulations which may be rare or unseen during training. A common assumption is that subpopulation shift data is drawn from a distribution of distributions: subpopulations are drawn from a population distribution, and each subpopulation has its own data-generating distribution [Santurkar et al., 2020, Yuan et al., 2022]. Since data in practice is often generated by individuals or groups who may have different data distributions, and models ideally generalize to unseen individuals and groups, this setting can be used to measure a natural notion of reliability. We study Plex under multiple ways of partitioning data into subpopulations.

Accuracy over subpopulations for CIFAR-10 and CIFAR-100. For each setting, we display the 5th, 25th, 50th, 75th, and 95th percentile accuracy for ViT-Plex L among the subpopulations. See <Ref> for comparison with no pretraining, which performs much worse.

Semantic Partitioning. One method for producing data with subpopulation shift is partitioning a standard dataset into subpopulations such that each subpopulation contains semantically similar examples. We leverage the Semantically Partitioned CIFAR10 and CIFAR100 datasets from Yuan et al., 2022 to study how Plex performs on new subpopulations. In particular, we measure classification accuracy with 30 and 100 subpopulations unseen during training, studying how accuracy varies among subpopulations.

From <Ref>, Plex greatly outperforms the baseline presented in Yuan et al., 2022 for tail subpopulation shift accuracy (we report 25th percentile among subpopulations). The baseline does not leverage large-scale pretraining or other changes Plex introduces. To better understand how Plex's ingredients lead to improvements in accuracy under subpopulation shift, we perform a series of ablations shown in <Ref>.
Comparing Plex to training with BatchEnsemble for both pretraining and finetuning (`BE→BE'), the latter produces a slight degradation in performance for both median and tail subpopulations, especially for CIFAR100, where performance is less saturated. This indicates that last layer changes offer some improvement but do not fully explain Plex's improvement over baseline in this setting. Looking at `None→BE', standard pretraining slightly degrades performance relative to BatchEnsemble pretraining, but using BatchEnsemble for finetuning still offers a significant boost over completely standard training (`None'); ensembling appears to be more important in this setting than last layer changes. Comparing Plex and `None', we see a significant difference in performance, but the difference between Plex's performance and that of the baseline presented in <Ref> remains much larger. This indicates that the difference in pretraining is an important driver of Plex's improvements on unseen and long-tail subpopulations, since the baseline does not pretrain on large-scale data.

Generalization, by method.

Sel. Predict., by method.

Calibration, by method.

Generalization, by size.

Sel. Predict., by size.

Calibration, by size.

T5-Plex model's ranking performance in subpopulations with spurious patterns compared to different methods (<ref>-<ref>) and across architecture sizes (<ref>-<ref>).

Spurious Correlations.

A well-observed issue in machine learning models is their tendency to learn decision rules that rely on spurious, non-causal patterns that exhibit strong statistical correlation with the outcome in the training data, i.e., shortcut learning [Geirhos et al., 2020]. In particular, previous theoretical studies show that a model's tendency to learn spurious patterns is rooted in the overparameterized model's tendency to flexibly adapt to the (biased) empirical distribution of the data, and cannot be fully addressed by increasing model size [Bommasani et al., 2021]. However, recent empirical studies also show that, when the data distribution exhibits suitable diversity such that it contains a small amount of counterexamples where the spurious pattern doesn't hold, large pretrained models can still lead to improved robustness [Tu et al., 2020]. Therefore, we investigate if the large pretrained models' robustness indeed improves with model scale, and if the uncertainty methods bring additional gain on top of the None baselines.

For the language modality, we evaluate the performance of T5 models on two spurious-correlation subpopulations: HANS for natural language inference, and CivilComments-Identities for toxic comment detection [McCoy et al., 2019, Borkan et al., 2019]. Specifically, HANS evaluates the NLI model's robustness against non-causal heuristic patterns that are empirically associated with sentence entailment (e.g., lexical overlap) [McCoy et al., 2019]. On the other hand, CivilComments-Identities contains comments that mention certain gender, sexual orientation, ethnic or religious identities, which empirically exhibits differential distribution of toxicity label with respect to their surface-level textual features. A toxic detection model that relies on these surface-level identity mentions can lead to unintended consequences, unjustly  reinforcing existing social stereotypes to disadvantaged identity groups [Borkan et al., 2019].

<Ref> summarizes the performance of different uncertainty methods, and across three different scales (i.e., T5$_{\texttt{small}}$, T5$_{\texttt{base}}$, T5$_{\texttt{large}}$) (See Appendix <Ref> for detailed results). Compared to the None baseline, the uncertainty methods (BE, DE, GP and their combinations) provide significant improvements across all aspects of model performance (i.e., generalization, selective prediction, and uncertainty calibration). For generalization and selective prediction performance, the ensemble of SNGP models (DE-GP) and Plex (BE-GP) perform the best among ensemble and non-ensemble methods, respectively. Furthermore, within each method class, the subpopulation generalization improves as the model scale increases, with Plex (BE+GP) being the most competitive method for moderate- and large-size models. Finally, the trend in uncertainty calibration is less consistent, we see that among all methods, GP and BE are on average the most calibrated ensemble and non-ensemble method, while the model scale does not seem to have a consistent impact to calibration performance.

§.§ Adaptation

In this section, we examine the models’ ability to adapt to new data. We focus on data-efficient learning (small data generalization), where the goal is to attain high reliability performance with only a small set of examples.
Our study focuses on ViT vision models, since adaptation has not yet been studied for T5 text models [Raffel et al., 2020].

§.§.§ Few-shot Learning

Plex's BatchEnsemble representation improves few-shot and over increasing model and pretraining dataset scales.
Full-model training can work better than linear evaluation if the goal is to maximize performance from limited examples.

In few-shot learning, we examine how well representations learned during pretraining enable fast downstream adaptation. We use a linear evaluation protocol where we extract features from the models’ pre-logits layer and train a multinomial logistic regression model using L-BFGS. Unlike linear regression which is also popular

(e.g., Dosovitskiy et al., 2021), logistic regression produces a categorical distribution, enabling one to measure the quality of the few-shot model's uncertainty estimates; we take advantage of this in <Ref>.

Few-shot performance averaged over 9 image datasets. (top) Comparison of pretrained model choices. (bottom) Comparison of pretrained model and dataset scale, where the default dataset is JFT and I21K signifies ImageNet21K.

BatchEnsemble's representation slightly outperforms others.

Model and pretraining dataset size have the most significant effect, with larger doing better and ImageNet21K doing better than JFT.

Model choice and size
<Ref> displays test set accuracy and negative log likelihood (NLL) averaged over 9 datasets: birds, caltech, cars, cifar100, col hist, dtd, imagenet, pets, and uc merced.
The choice of model translates to a consistent 1-2% improvement in accuracy depending on the few-shot setting. The BatchEnsemble representation is among the top performing that gives high predictive performance.

The size of the model and pretraining dataset also has a noticeable effect on few-shot performance.
While ImageNet-21K is 10 times smaller than JFT, models pretrained on ImageNet-21K consistently outperform those pretrained on JFT across prediction and uncertainty tasks, with the exception of calibration. We hypothesize this is because ImageNet-21K is closer in distribution to the datasets we evaluate on average. It is also surprising that when finetuning over the full ImageNet dataset (<Ref>), the conclusion is reversed with JFT models outperforming ImageNet-21K. This suggests a transition point when the pretrained models are trained on a sufficiently high number of examples.

Linear evaluation versus full-model training on ImageNet. Full-model training consistently outperforms linear evaluation at 5-shot and higher.

Full-model training
Intuitively, linear evaluation is about “representation learning”, assessing how well a pretrained model’s fixed representation layer generalizes to many different scenarios. This setup may not be optimal for real-world transfer if the goal is simply to maximize performance from limited examples, without constraints on the approach to doing so.

In <Ref>, we adopt the usual finetuning protocol and experiment with training the model’s full set of parameters given the pretrained checkpoint.

We set up a few-shot version of ImageNet1k for this study.[

On 1-shot, linear evaluation outperforms full-model training which may not be surprising given the high possibility of overfitting.
Starting from 5-shot (and which is likely to hold for even lower shot settings), full-model training surpasses linear evaluation

with a consistent 2-3% improvement in accuracy.

§.§.§ Few-shot Uncertainty

Plex performs well on both few-shot and zero-shot open set recognition, particularly using larger model sizes and on ImageNet21K.
Few-shot calibration remains a challenging problem.

Few-shot uncertainty performance averaged over 9 image datasets for calibration and selective prediction, and 2 datasets for open set recognition.

(top) Comparison of pretrained model choices. (bottom) Comparison of pretrained model and dataset scale, where the default dataset is JFT and I21K signifies ImageNet21K.
Plex's BatchEnsemble representation works best on few-shot open set recognition, and with ImageNet21K as the pretraining dataset.
Calibration and selective prediction performance are similar across model and scaling choices; they remain challenging problems.

Few-shot learning is an important and practical setting, for which the literature almost exclusively considers accuracy to measure performance. Few-shot uncertainty quantification is also a key problem:

compared to finetuning, few-shot models are more sensitive to the examples they've seen and by extension, much less accurate. Thus few-shot models should be more likely to be uncertain, and the quality of their uncertainty estimates must be good in order to know when to refrain from making a prediction.

Few-shot uncertainty tasks
In <Ref>, we measure uncertainty performance

on three metrics: ECE and Calibration AUROC for calibration, averaged over all 9 datasets used for few-shot learning, and OOD AUROC for open set recognition.
Open set recognition requires an OOD dataset in addition to a test split, so we only evaluate it for two datasets: CIFAR-100 with an OOD dataset of CIFAR-10; and ImageNet1K with an OOD dataset of Places365.

Calibration performance is roughly comparable across methods. BatchEnsemble has a noticeable improvement on open set recognition.
The scale of the pretraining dataset has no significant effect on calibration, where the trend is weaker than what's seen on calibration of finetuned models (<Ref>). The choice of pretraining dataset does have a noticeable effect on open set recognition, where like in <Ref>, Plex on ImageNet21K surprisingly does better than JFT models.

2cFar-OOD              2cNear-OOD

C-10 vs SVHN     C-100 vs SVHN     C-10 vs C-100     C-100 vs C-10
Plex S            0.596                0.449                 0.670                     0.592

Plex B            0.797                0.643                 0.733                     0.679

Plex L            0.717                0.586                 0.716                     0.690

Plex L I21K     0.895                0.956                 0.864                     0.902

5cRelative Mahalanobis
Plex S            0.843                0.814                 0.805                     0.678

Plex B            0.938                0.848                 0.863                     0.769

Plex L            0.967                0.875                 0.890                     0.821

Plex L I21K     0.957                0.932                 0.919                     0.915

Zero-shot open set recognition with Plex, varying scale and detection method.

Zero-shot open set recognition

Given that the pretrained model itself provides rich and robust representations, we can use that as a feature extractor and conduct zero-shot open set recognition without any training examples.

To do so, we can take advantage of Mahalanobis distance (Maha) [Lee et al., 2018] as a detection score: Maha measures the distance between the test input and the fitted training distribution in the embedding space. It operates on a fixed representation layer and does not require operating on softmax outputs with a newly trained last layer.
The training distribution is fitted using a class conditional Gaussian $\gaussk, k=1, 2, \dots, K$ to each of the $K$ in-distribution classes based on the embedding $\vz$. We estimate the mean vectors and covariance matrix as:
$ \mathbf{\vmu}_k = \frac{1}{N_k} \sum_{i:y_i=k} \vz_i$, for $k=1, \dots, K,$
and $\vSigma = \frac{1}{N}\sum_{k=1}^K \sum_{i:y_i=k}\left( \vz_i-\mathbf{\vmu}_k \right)(\vz_i-\mathbf{\vmu}_k)^T$.
Note that class-conditional means $\mathbf{\vmu}_k$ are independent for each classes, while the covariance matrix $\Sigma$ is shared by all classes to avoid under-fitting errors.
For a test input, Mahalanobis distances from a test input to each of the fitted $K$ in-distribution Gaussian distributions $\gaussk$ is computed, and the minimum of the distances over all classes is used as the uncertainty score,
$\text{MD}(\vz) = \min_{k} \{ \left(\vz-\mathbf{\vmu}_k\right)^T \vSigma^{-1} \left(\vz-\mathbf{\vmu}_k \right) \}.$

We also experiment with a Relative Mahalanobis distance variant (RMaha) [Ren et al., 2021], a modified version of Mahalanobis distance which corrects for the background confounding effect using another Gaussian distribution fitted using entire training data ignoring class labels.
The uncertainty score is defined as $\text{RMD}_k(\vz) \nonumber = \text{MD}(\vz)  - \text{MD}_0(\vz)$, where $\text{MD}_0(\vz)$ indicates the Mahalanobis distance to a distribution fitted to the entire training data not considering the class labels: $\mathcal{N}(\mathbf{\vmu}_0, \vSigma_0)$, where $\vmu_0 = \frac{1}{N} \sum_{i=1}^N \vz_i$ and $\vSigma_0 = \frac{1}{N}\sum_{i=1}^N \left( \vz_i-\vmu_0 \right)(\vz_i-\vmu_0)^T$.
This is a good proxy for the background distribution.

Zero-shot results are displayed in <Ref>.

Interestingly, the zero-shot setting achieves performance close to that of the finetuned setting. For example, in the challenging near-OOD task CIFAR-100 vs CIFAR-10, the best zero-shot AUROC achieved by Plex L (I21K) is 0.915, while its performance with finetuning is 0.934, and the best model which uses JFT and finetuning achieves 0.954 (<Ref>).

For the far-OOD task CIFAR-100 vs SVHN, the best AUROC achieved by zero-shot is 0.933, and the best model which uses JFT and finetuning achievs 0.938.

The Relative Mahalanobis distance also significantly outperforms the raw Mahalanobis distance, suggesting that there are background features that confound the raw Mahalanobis distance. In fact, the raw Mahalanobis distance does not work well for most of the pretrained models, mostly having AUROC lower than 0.9 and varies depending on the model types.

§ CONCLUSION

We presented a framework for thinking about reliable deep learning, and provided a number of tasks and datasets for stress-testing the reliability of models through multiple tasks such as the ability to quantify its confidence in predictions, be robust to distribution shifts, and adapt quickly to new distributions. To improve reliability, we also developed Plex, pretrained large model extensions for vision (ViT-Plex) and language (T5-Plex) that significantly improve the reliability in deep learning. Our techniques are broadly applicable for other large models (cf. LaMDA [Thoppilan et al., 2022], PaLM [Chowdhery et al., 2022]) and wider suite of tasks, cf. BIG-bench [Srivastava et al., 2022].

§ ACKNOWLEDGEMENTS

We thank
Ben Adlam,
Dilip Krishnan,
Ed Chi,
Neil Houlsby,
Rif A. Saurous,
Sharat Chikkerur
for helpful discussions and feedback on earlier drafts of the paper.
We also thank Tom Small and Ajay Nainani for assisting with visualizations.

[Abraham and Ledolter, 1983]
Bovas Abraham and Johannes Ledolter.
Statistical methods for forecasting, volume 179.
Wiley Online Library, 1983.

[Allingham et al., 2021]
James Urquhart Allingham, Florian Wenzel, Zelda E Mariet, Basil Mustafa, Joan
Puigcerver, Neil Houlsby, Ghassen Jerfel, Vincent Fortuin, Balaji
Lakshminarayanan, Jasper Snoek, Dustin Tran, Carlos Riquelme Ruiz, and
Rodolphe Jenatton.
Sparse MoEs meet efficient ensembles.
arXiv preprint arXiv:2110.03360, 2021.

[Amodei et al., 2016]
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and
Dan Mané.
Concrete problems in AI safety.
arXiv preprint arXiv:1606.06565, 2016.

[Avin et al., 2021]
Shahar Avin, Haydn Belfield, Miles Brundage, Gretchen Krueger, Jasmine Wang,
Adrian Weller, Markus Anderljung, Igor Krawczuk, David Krueger, Jonathan
Lebensold, et al.
Filling gaps in trustworthy development of ai.
Science, 3740 (6573):0 1327–1329, 2021.

[Band et al., 2021]
Neil Band, Tim G. J. Rudner, Qixuan Feng, Angelos Filos, Zachary Nado,
Michael W Dusenberry, Ghassen Jerfel, Dustin Tran, and Yarin Gal.
Benchmarking Bayesian deep learning on diabetic retinopathy
detection tasks.
In Advances in Neural Information Processing Systems, Datasets
and Benchmarks Track, 2021.

[Barbu et al., 2019]
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan
Gutfreund, Josh Tenenbaum, and Boris Katz.
Objectnet: A large-scale bias-controlled dataset for pushing the
limits of object recognition models.
In Advances in Neural Information Processing Systems,
volume 32, 2019.

[Barlow and Proschan, 1975]
Richard E Barlow and Frank Proschan.
Statistical theory of reliability and life testing: probability
Technical report, Florida State Univ Tallahassee, 1975.

[Beluch et al., 2018]
William H Beluch, Tim Genewein, Andreas Nürnberger, and Jan M Köhler.
The power of ensembles for active learning in image classification.
In Computer Vision and Pattern Recognition, 2018.

[Bengio et al., 2015]
Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup.
Conditional computation in neural networks for faster models.
arxiv preprint arxiv:1511.06297, 2015.

[Bengio et al., 2013]
Yoshua Bengio, Nicholas Léonard, and Aaron Courville.
Estimating or propagating gradients through stochastic neurons for
conditional computation.
arxiv preprint arxiv:1308.3432, 2013.

[Beyer et al., 2020]
Lucas Beyer, Olivier J Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and
Aäron van den Oord.
Are we done with ImageNet?
arXiv preprint arXiv:2006.07159, 2020.

[Blundell et al., 2015]
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.
Weight Uncertainty in Neural Networks.
In International Conference on Machine Learning, 2015.

[Bommasani et al., 2021]
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney
von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma
Brunskill, et al.
On the opportunities and risks of foundation models.
arXiv preprint arXiv:2108.07258, 2021.

[Borkan et al., 2019]
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman.
Nuanced metrics for measuring unintended bias with real data for text
In Companion proceedings of the 2019 world wide web
conference, pages 491–500, 2019.

[Brown et al., 2020]
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla
Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell,
et al.
Language models are few-shot learners.
In Advances in Neural Information Processing Systems, 2020.

[Chowdhery et al., 2022]
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra,
Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian
Gehrmann, et al.
Palm: Scaling language modeling with pathways.
arXiv preprint arXiv:2204.02311, 2022.

[Cimpoi et al., 2014]
M. Cimpoi, S. Maji, I. Kokkinos, S. Mohamed, and A. Vedaldi.
Describing textures in the wild.
In Computer Vision and Pattern Recognition, 2014.

[Citovsky et al., 2021]
Gui Citovsky, Giulia DeSalvo, Claudio Gentile, Lazaros Karydas, Anand
Rajagopalan, Afshin Rostamizadeh, and Sanjiv Kumar.
Batch active learning at scale, 2021.

[Cohn et al., 1996]
David A Cohn, Zoubin Ghahramani, and Michael I Jordan.
Active learning with statistical models.
Journal of Artificial Intelligence Research, 4:0
129–145, 1996.

[Collier et al., 2020]
Mark Collier, Basil Mustafa, Efi Kokiopoulou, Rodolphe Jenatton, and Jesse
A simple probabilistic method for deep classification under
input-dependent label noise.
arXiv preprint arXiv:2003.06778, 2020.

[Collier et al., 2021]
Mark Collier, Basil Mustafa, Efi Kokiopoulou, Rodolphe Jenatton, and Jesse
Correlated input-dependent label noise in large-scale image
In Computer Vision and Pattern Recognition, 2021.

[D'Amour et al., 2020]
Alexander D'Amour, Katherine Heller, Dan Moldovan, Ben Adlam, Babak Alipanahi,
Alex Beutel, Christina Chen, Jonathan Deaton, Jacob Eisenstein, Matthew D
Hoffman, et al.
Underspecification presents challenges for credibility in modern
machine learning.
arXiv preprint arXiv:2011.03395, 2020.

[Dawid, 1982]
A Philip Dawid.
The well-calibrated Bayesian.
Journal of the American Statistical Association, 1982.

[Deng et al., 2009]
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.
Imagenet: A large-scale hierarchical image database.
In Computer Vision and Pattern Recognition, pages 248–255.
Ieee, 2009.

[Dosovitskiy et al., 2021]
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn,
Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg
Heigold, Sylvain Gelly, et al.
An image is worth 16x16 words: Transformers for image recognition at
In International Conference on Learning Representations, 2021.

[Dusenberry et al., 2020]
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snoek,
Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran.
Efficient and scalable Bayesian neural nets with rank-1 factors.
In International Conference on Machine Learning,

[Dusenberry et al., 2020]
Michael W Dusenberry, Dustin Tran, Edward Choi, Jonas Kemp, Jeremy Nixon,
Ghassen Jerfel, Katherine Heller, and Andrew M Dai.
Analyzing the role of model uncertainty for electronic health
In ACM Conference on Health, Inference, and Learning,

[El-Yaniv and Wiener, 2010]
Ran El-Yaniv and Yair Wiener.
On the foundations of noise-free selective classification.
Journal of Machine Learning Research, 110 (5), 2010.

[Emam et al., 2021]
Zeyad Ali Sami Emam, Hong-Min Chu, Ping-Yeh Chiang, Wojciech Czaja, Richard
Leapman, Micah Goldblum, and Tom Goldstein.
Active learning at the imagenet scale.
arXiv preprint arXiv:2111.12880, 2021.

[Evci et al., 2022]
Utku Evci, Vincent Dumoulin, Hugo Larochelle, and Michael C Mozer.
Head2toe: Utilizing intermediate representations for better transfer
arXiv preprint arXiv:2201.03529, 2022.

[Everitt et al., 2018]
Tom Everitt, Gary Lea, and Marcus Hutter.
AGI safety literature review.
arXiv preprint arXiv:1805.01109, 2018.

[Farquhar et al., 2020]
Sebastian Farquhar, Michael A. Osborne, and Yarin Gal.
Radial Bayesian Neural Networks: Beyond Discrete Support In
Large-Scale Bayesian Deep Learning.
In International Conference on Artificial Intelligence and
Statistics, 2020.

[Fedus et al., 2022]
William Fedus, Barret Zoph, and Noam Shazeer.
Switch transformers: Scaling to trillion parameter models with simple
and efficient sparsity.
Journal of Machine Learning Research, 230 (120), 2022.

[Fei-Fei et al., 2004]
Li Fei-Fei, Rob Fergus, and Pietro Perona.
Learning generative visual models from few training examples: An
incremental Bayesian approach tested on 101 object categories.
Computer Vision and Pattern Recognition Workshop, 2004.

[Fort et al., 2021]
Stanislav Fort, Jie Ren, and Balaji Lakshminarayanan.
Exploring the limits of out-of-distribution detection.
In Advances in Neural Information Processing Systems, 2021.

[Gal and Ghahramani, 2016]
Yarin Gal and Zoubin Ghahramani.
Dropout As a Bayesian Approximation: Representing Model Uncertainty
in Deep Learning.
In International Conference on Machine Learning, 2016.

[Gal et al., 2017]
Yarin Gal, Riashat Islam, and Zoubin Ghahramani.
Deep Bayesian active learning with image data.
In International Conference on Machine Learning, 2017.

[Geirhos et al., 2020]
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel,
Wieland Brendel, Matthias Bethge, and Felix A Wichmann.
Shortcut learning in deep neural networks.
Nature Machine Intelligence, 20 (11):0
665–673, 2020.

[Geng et al., 2020]
Chuanxing Geng, Sheng-jun Huang, and Songcan Chen.
Recent advances in open set recognition: A survey.
IEEE transactions on pattern analysis and machine
intelligence, 430 (10):0 3614–3631, 2020.

[Ghahramani, 2015]
Zoubin Ghahramani.
Probabilistic machine learning and artificial intelligence.
Nature, 5210 (7553):0 452–459, 2015.

[Gneiting and Raftery, 2007]
Tilmann Gneiting and Adrian E Raftery.
Strictly Proper Scoring Rules, Prediction, and Estimation.
Journal of the American Statistical Association, 1020
(477):0 359–378, March 2007.

[Gustafsson et al., 2020]
Fredrik K Gustafsson, Martin Danelljan, and Thomas B Schön.
Evaluating scalable Bayesian deep learning methods for robust
computer vision.
In Computer Vision and Pattern Recognition Workshops, 2020.

[Hakkani-Tür et al., 2002]
Dilek Hakkani-Tür, Giuseppe Riccardi, and Allen Gorin.
Active learning for automatic speech recognition.
In IEEE international Conference on Acoustics, Speech, and
Signal Processing, 2002.

[Hansen and Salamon, 1990]
L.K. Hansen and P. Salamon.
Neural network ensembles.
IEEE Transactions on Pattern Analysis and Machine
Intelligence, 1990.

[Hendrycks and Dietterich, 2019]
Dan Hendrycks and Thomas Dietterich.
Benchmarking Neural Network Robustness to Common
Corruptions and Perturbations.
In International Conference on Learning Representations,

[Hendrycks and Gimpel, 2017]
Dan Hendrycks and Kevin Gimpel.
A Baseline for Detecting Misclassified and
Out-of-Distribution Examples in Neural Networks.
In International Conference on Learning Representations,

[Hendrycks et al., 2019]
Dan Hendrycks, Steven Basart, Mantas Mazeika, Mohammadreza Mostajabi, Jacob
Steinhardt, and Dawn Song.
Scaling out-of-distribution detection for real-world settings.
arXiv preprint arXiv:1911.11132, 2019.

[Hendrycks et al., 2021]
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan
Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al.
The many faces of robustness: A critical analysis of
out-of-distribution generalization.
In ICCV, 2021a.

[Hendrycks et al., 2021]
Dan Hendrycks, Nicholas Carlini, John Schulman, and Jacob Steinhardt.
Unsolved problems in ML safety.
arXiv preprint arXiv:2109.13916, 2021b.

[Hendrycks et al., 2021]
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song.
Natural adversarial examples.
In Computer Vision and Pattern Recognition,

[Hu et al., 2019]
Peiyun Hu, Zachary C Lipton, Anima Anandkumar, and Deva Ramanan.
Active learning with partial feedback.
In International Conference on Learning Representations, 2019.

[Kather et al., 2016]
Jakob Nikolas Kather, Cleo-Aron Weis, Francesco Bianconi, Susanne M Melchers,
Lothar R Schad, Timo Gaiser, Alexander Marx, and Frank Gerrit Z"ollner.
Multi-class texture analysis in colorectal cancer histology.
Scientific reports, 6:0 27988, 2016.

[Kendall and Gal, 2017]
Alex Kendall and Yarin Gal.
What uncertainties do we need in Bayesian deep learning for
computer vision?
Advances in neural information processing systems, 30, 2017.

[Kivlichan et al., 2021]
Ian D Kivlichan, Zi Lin, Jeremiah Liu, and Lucy Vasserman.
Measuring and improving model-moderator collaboration using
uncertainty estimation.
arXiv preprint arXiv:2107.04212, 2021.

[Kolesnikov et al., 2020]
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung,
Sylvain Gelly, and Neil Houlsby.
Big transfer (bit): General visual representation learning.
In ECCV, 2020.

[Krause et al., 2013]
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei.
3d object representations for fine-grained categorization.
In IEEE Workshop on 3D Representation and Recognition, 2013.

[Krishnan and Tickoo, 2020]
Ranganath Krishnan and Omesh Tickoo.
Improving model calibration with accuracy versus uncertainty
In Advances in Neural Information Processing Systems,
volume 33, 2020.

[Krizhevsky et al., 2009]
Alex Krizhevsky, Geoffrey Hinton, et al.
Learning multiple layers of features from tiny images.
arXiv preprint, 2009.

[Lakshminarayanan et al., 2017]
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.
Simple and scalable predictive uncertainty estimation using deep
In Advances in Neural Information Processing Systems, 2017.

[Larson et al., 2019]
Stefan Larson, Anish Mahendran, Joseph J Peper, Christopher Clarke, Andrew Lee,
Parker Hill, Jonathan K Kummerfeld, Kevin Leach, Michael A Laurenzano,
Lingjia Tang, et al.
An evaluation dataset for intent classification and out-of-scope
arXiv preprint arXiv:1909.02027, 2019.

[Lee et al., 2018]
Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin.
A simple unified framework for detecting out-of-distribution samples
and adversarial attacks.
In Advances in Neural Information Processing Systems, 2018.

[Lepikhin et al., 2021]
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping
Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.
GShard: Scaling giant models with conditional computation and
automatic sharding.
In International Conference on Learning Representations, 2021.

[Li et al., 2006]
Fei-Fei Li, Rob Fergus, and Pietro Perona.
One-shot learning of object categories.
IEEE transactions on pattern analysis and machine
intelligence, 280 (4):0 594–611, 2006.

[Liu et al., 2020]
Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and
Balaji Lakshminarayanan.
Simple and principled uncertainty estimation with deterministic deep
learning via distance awareness.
In Advances in Neural Information Processing Systems, 2020.

[Liu et al., 2022]
Jeremiah Zhe Liu, Shreyas Padhy, Jie Ren, Zi Lin, Yeming Wen, Ghassen Jerfel,
Zack Nado, Jasper Snoek, Dustin Tran, and Balaji Lakshminarayanan.
A simple approach to improve single-model deep uncertainty via
arXiv preprint arXiv:2205.00403, 2022.

[Liu et al., 2021]
Xingkun Liu, Arash Eshghi, Pawel Swietojanski, and Verena Rieser.
Benchmarking natural language understanding services for building
conversational agents.
In Increasing Naturalness and Flexibility in Spoken Dialogue
Interaction, pages 165–183. Springer, 2021.

[Lou et al., 2021]
Yuxuan Lou, Fuzhao Xue, Zangwei Zheng, and Yang You.
Cross-token modeling with conditional computation, 2021.
URL <https://arxiv.org/abs/2109.02008>.

[Malinin and Gales, 2021]
Andrey Malinin and Mark Gales.
Uncertainty estimation in autoregressive structured prediction.
In International Conference on Learning Representations, 2021.

[Martinez-Cantin et al., 2007]
Ruben Martinez-Cantin, Nando de Freitas, Arnaud Doucet, and José A
Active policy learning for robot planning and exploration under
In Robotics: Science and Systems, volume 3, pages 321–328,

[McCoy et al., 2019]
Tom McCoy, Ellie Pavlick, and Tal Linzen.
Right for the wrong reasons: Diagnosing syntactic heuristics in
natural language inference.
In Annual Meeting of the Association for Computational
Linguistics, 2019.

[Minderer et al., 2021]
Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai,
Neil Houlsby, Dustin Tran, and Mario Lucic.
Revisiting the calibration of modern neural networks.
In Advances in Neural Information Processing Systems, 2021.

[Nado et al., 2021]
Zachary Nado, Neil Band, Mark Collier, Josip Djolonga, Michael W Dusenberry,
Sebastian Farquhar, Qixuan Feng, Angelos Filos, Marton Havasi, Rodolphe
Jenatton, et al.
Uncertainty baselines: Benchmarks for uncertainty & robustness in
deep learning.
arXiv preprint arXiv:2106.04015, 2021.

[Naeini et al., 2015]
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht.
Obtaining well calibrated probabilities using Bayesian binning.
In AAAI, 2015.

[Nixon et al., 2019]
Jeremy Nixon, Michael W Dusenberry, Linchuan Zhang, Ghassen Jerfel, and Dustin
Measuring calibration in deep learning.
In CVPR Workshops, 2019.

[O'Connor and Kleyner, 2012]
Patrick O'Connor and Andre Kleyner.
Practical reliability engineering.
John Wiley & Sons, 2012.

[Ovadia et al., 2019]
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian
Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek.
Can you trust your model's uncertainty? Evaluating predictive
uncertainty under dataset shift.
In Advances in Neural Information Processing Systems, 2019.

[Parkhi et al., 2012]
O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. V. Jawahar.
Cats and dogs.
In Computer Vision and Pattern Recognition, 2012.

[Parmigiani and Inoue, 2009]
G. Parmigiani and Lurdes Inoue.
Decision theory: principles and approaches.
Wiley series in probability and statistics. John Wiley & Sons, 2009.

[Peterson et al., 2019]
Joshua C Peterson, Ruairidh M Battleday, Thomas L Griffiths, and Olga
Human uncertainty makes classification more robust.
In International Conference on Computer Vision, 2019.

[Pham et al., 2021]
Hieu Pham, Zihang Dai, Golnaz Ghiasi, Hanxiao Liu, Adams Wei Yu, Minh-Thang
Luong, Mingxing Tan, and Quoc V Le.
Combined scaling for zero-shot transfer learning.
arXiv preprint arXiv:2111.10050, 2021.

[Radford et al., 2021]
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh,
Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark,
Gretchen Krueger, and Ilya Sutskever.
Learning transferable visual models from natural language
In International Conference on Machine Learning, 2021.

[Raffel et al., 2020]
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael
Matena, Yanqi Zhou, Wei Li, and Peter J Liu.
Exploring the limits of transfer learning with a unified text-to-text
Journal of Machine Learning Research, 2020.

[Recht et al., 2019]
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar.
Do ImageNet classifiers generalize to ImageNet?
In International Conference on Machine Learning, 2019.

[Ren et al., 2021]
Jie Ren, Stanislav Fort, Jeremiah Liu, Abhijit Guha Roy, Shreyas Padhy, and
Balaji Lakshminarayanan.
A simple fix to mahalanobis distance for improving near-ood
arXiv preprint arXiv:2106.09022, 2021.

[Riccardi and Hakkani-Tur, 2005]
Giuseppe Riccardi and Dilek Hakkani-Tur.
Active learning: Theory and applications to automatic speech
IEEE Transactions on Speech and Audio Processing, 130
(4):0 504–511, 2005.

[Riquelme et al., 2021]
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe
Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby.
Scaling vision with sparse mixture of experts.
In Advances in Neural Information Processing Systems, 2021.

[Roth and Small, 2006]
Dan Roth and Kevin Small.
Margin-based active learning for structured output spaces.
In European Conference on Machine Learning, pages 413–424.
Springer, 2006.

[Rudner et al., 2021]
Tim G. J. Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal.
Tractable Function-Space Variational Inference in
Bayesian Neural Networks.
In ICML Workshop on Uncertainty & Robustness in Deep
Learning, 2021.

[Rudner et al., 2022]
Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin
Continual Learning via Sequential Function-Space
Variational Inference.
In International Conference on Machine Learning, 2022.

[Russell et al., 2015]
Stuart Russell, Daniel Dewey, and Max Tegmark.
Research priorities for robust and beneficial artificial
Ai Magazine, 360 (4):0 105–114, 2015.

[Santurkar et al., 2020]
Shibani Santurkar, Dimitris Tsipras, and Aleksander Madry.
Breeds: Benchmarks for subpopulation shift.
arXiv preprint arXiv:2008.04859, 2020.

[Scheffer et al., 2001]
Tobias Scheffer, Christian Decomain, and Stefan Wrobel.
Active hidden markov models for information extraction.
In International Symposium on Intelligent Data Analysis, pages
309–318. Springer, 2001.

[Settles, 2009]
Burr Settles.
Active learning literature survey.
preprint, 2009.

[Shankar et al., 2021]
Vaishaal Shankar, Achal Dave, Rebecca Roelofs, Deva Ramanan, Benjamin Recht,
and Ludwig Schmidt.
Do image classifiers generalize across time?
In International Conference on Computer Vision, 2021.

[Shazeer et al., 2017]
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le,
Geoffrey Hinton, and Jeff Dean.
Outrageously large neural networks: The sparsely-gated
mixture-of-experts layer.
In International Conference on Learning Representations, 2017.

[Siddhant and Lipton, 2018]
Aditya Siddhant and Zachary C Lipton.
Deep Bayesian active learning for natural language processing:
Results of a large-scale empirical study.
In EMNLP, 2018.

[Song et al., 2019]
Shuang Song, David Berthelot, and Afshin Rostamizadeh.
Combining mixmatch and active learning for better accuracy with fewer
arXiv preprint arXiv:1912.00594, 2019.

[Srivastava et al., 2022]
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar
Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià
Garriga-Alonso, et al.
Beyond the imitation game: Quantifying and extrapolating the
capabilities of language models.
arXiv preprint arXiv:2206.04615, 2022.

[Sugiyama and Kawanabe, 2012]
Masashi Sugiyama and Motoaki Kawanabe.
Machine learning in non-stationary environments: Introduction
to covariate shift adaptation.
MIT press, 2012.

[Tamkin et al., 2022]
Alex Tamkin, Dat Nguyen, Salil Deshpande, Jesse Mu, and Noah Goodman.
Active learning helps pretrained models learn the intended task.
arXiv preprint arXiv:2204.08491, 2022.

[Taori et al., 2020]
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht,
and Ludwig Schmidt.
Measuring robustness to natural distribution shifts in image
Advances in Neural Information Processing Systems, 2020.

[Thompson et al., 1999]
Cynthia A Thompson, Mary Elaine Califf, and Raymond J Mooney.
Active learning for natural language parsing and information
In International Conference on Machine Learning, 1999.

[Thoppilan et al., 2022]
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv
Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du,
et al.
Lamda: Language models for dialog applications.
arXiv preprint arXiv:2201.08239, 2022.

[Thrun, 1998]
Sebastian Thrun.
Lifelong learning algorithms.
In Learning to learn, pages 181–209. Springer, 1998.

[Tran et al., 2018]
Dustin Tran, Matthew W Hoffman, Dave Moore, Christopher Suter, Srinivas
Vasudevan, and Alexey Radul.
Simple, distributed, and accelerated probabilistic programming.
Advances in Neural Information Processing Systems, 2018.

[Tran et al., 2020]
Dustin Tran, Jasper Snoek, and Balaji Lakshminarayanan.
Practical uncertainty estimation and out-of-distribution robustness
in deep learning.
Advances in Neural Information Processing Systems tutorial,

[Tran et al., 2019]
Toan Tran, Thanh-Toan Do, Ian Reid, and Gustavo Carneiro.
Bayesian generative active deep learning.
In International Conference on Machine Learning, pages
6295–6304. PMLR, 2019.

[Tu et al., 2020]
Lifu Tu, Garima Lalwani, Spandana Gella, and He He.
An empirical study on robustness to spurious correlations using
pre-trained language models.
Transactions of the Association for Computational Linguistics,
8:0 621–633, 2020.

[Van Amersfoort et al., 2020]
Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal.
Uncertainty estimation using a single deep deterministic neural
In International Conference on Machine Learning, 2020.

[Wang et al., 2018]
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R
Glue: A multi-task benchmark and analysis platform for natural
language understanding.
arXiv preprint arXiv:1804.07461, 2018a.

[Wang et al., 2018]
Zi Wang, Caelan Reed Garrett, Leslie Pack Kaelbling, and Tomás
Active model learning and diverse action sampling for task and motion
In 2018 IEEE/RSJ International Conference on Intelligent Robots
and Systems (IROS), pages 4107–4114. IEEE, 2018b.

[Welinder et al., 2010]
Peter Welinder, Steve Branson, Takeshi Mita, Catherine Wah, Florian Schroff,
Serge Belongie, and Pietro Perona.
Caltech-ucsd birds 200.
preprint, 2010.

[Wen et al., 2020]
Yeming Wen, Ghassen Jerfel, Rafael Muller, Michael W Dusenberry, Jasper Snoek,
Balaji Lakshminarayanan, and Dustin Tran.
Combining ensembles and data augmentation can harm your calibration.
arXiv preprint arXiv:2010.09875, 2020a.

[Wen et al., 2020]
Yeming Wen, Dustin Tran, and Jimmy Ba.
BatchEnsemble: an Alternative Approach to Efficient
Ensemble and Lifelong Learning.
In International Conference on Learning Representations,

[Williams et al., 2017]
Adina Williams, Nikita Nangia, and Samuel R Bowman.
A broad-coverage challenge corpus for sentence understanding through
arXiv preprint arXiv:1704.05426, 2017.

[Wulczyn et al., 2017]
Ellery Wulczyn, Nithum Thain, and Lucas Dixon.
Ex machina: Personal attacks seen at scale.
In WWW, 2017.

[Yang and Newsam, 2010]
Yi Yang and Shawn Newsam.
Bag-of-visual-words and spatial extensions for land-use
In International Conference on Advances in Geographic
Information Systems, 2010.

[Yi et al., 2022]
John Seon Keun Yi, Minseok Seo, Jongchan Park, and Dong-Geol Choi.
Using self-supervised pretext tasks for active learning, 2022.
URL <https://arxiv.org/abs/2201.07459>.

[Yuan et al., 2022]
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal.
What do we mean by generalization in federated learning?
International Conference on Learning Representations, 2022.

[Zhai et al., 2021]
X Zhai, A Kolesnikov, N Houlsby, and L Beyer.
Scaling vision transformers.
arXiv preprint arXiv:2106.04560, 2021.

[Zhang et al., 2021]
Jian-Guo Zhang, Kazuma Hashimoto, Yao Wan, Ye Liu, Caiming Xiong, and Philip S
Are pretrained transformers robust in intent classification? a
missing ingredient in evaluation of out-of-scope intent detection.
arXiv preprint arXiv:2106.04564, 2021.

§ AUTHOR CONTRIBUTIONS

* Andreas Kirsch: Active Learning only – joint with Joost van Amersfoort. Implemented Active Learning loop, various acquisition functions and datasets, and ran experiments on CIFAR-10/CIFAR-100. Advised on experimental setup for larger datasets, and wrote active learning section.

* Balaji Lakshminarayanan: Helped with direction, narrative, getting resources, advising folks on experiment setup (open set recognition, few-shot adaptation, last layer), helped with writing (abstract, intro, contributions) and revising other sections.

* Clara Huiyi Hu: Implemented BE + GP in language domain, executed associated experiments and helped building post-analysis colabs. On the vision side, contributed to sweeping and debugging BE/GP variants.

* Dustin Tran: Came up with work's vision, lead overall writing and experiment design, coordinated individual leads and contributors. Designed and tuned Plex, various figures and plotting code, and core infra setup.

* Du Phan: Lead and designed the infrastructure to evaluate uncertainty methods on the language domain. Implemented methods None, DE, MCD, GP, DE-GP, Het for language and performed most corresponding experiments across MNLI, ToxicComments, NaLUE tasks.

* D. Sculley: Advisor on project, including its direction and the paper narrative.

* Honglin Yuan: Initial experimental setup for CIFAR subpopulation shift experiments, including pipeline for converting data into format to be consumed by uncertainty baselines codebase.

* Jasper Snoek: Helped with direction, narrative, getting resources, helped with writing, editing, aggregating results, various plots.

* Jeremiah Liu:
Lead for language domain efforts. Developed all tasks, datasets, evaluation metrics for language.

Implemented initial t5x infra and baseline (None) experiments. Work with eng co-lead Du and teammates Clara and Jie in finishing executing experiments and generating visualizations. Wrote the sections on language. On the vision side, implemented the ViT-GP model, conducted initial pretraining experiments. Assisted editing the section on SNGP method.

* Jie Ren: Lead work on open set recognition (vision and text), GPs, and assisted with few-shot. Helped with engineering work including building vision finetuning models, running vision and text experiments. Made the demo figures for vision and text. Helped with Mahalanobis distance-based active learning method.

* Joost van Amersfoort: See Andreas Kirsch (joint contribution). Co-wrote compute workflow tutorial for collaborators.

* Karan Singhal: Vision subpopulation shift setup and experiments: integrated code for CIFAR-10/100 subpopulation evaluation into codebase, set up, ran, and tuned experiments with different ingredients (Det, BE, BE+HET, BE+GP). Ran ablations: omitting last-layer changes, omitting ensembling during pretraining, ImageNet-only pretraining, no pretraining. Created ablations figure and added discussion. Made detailed edits to paper.

* Kehang Han: Responsible for finetuning experiments involving GP and few-shot experiments. Created Imagenet2012Fewshot TFDS (<https://www.tensorflow.org/datasets/catalog/imagenet2012_fewshot>). Primary author of <Ref>. Helped with GP model ablations.

* Kelly Buchanan: Contributed to early formation of project and codebase, particularly around tasks involving structured output evaluation and uncertainty.