PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2207.07411v1cs.LGstat.ML

TL;DR

Pretrained models achieve strong performance yet can fail unpredictably, motivating evaluation of reliability across real-world decision-making settings. The paper introduces Plex extensions for vision and language, finds improved reliability across tasks, and shows that scaling model and pretraining data sizes further improves reliability.

  • Problem

    Reliable AI models must work consistently across real-world settings, beyond achieving strong predictive performance on individual tasks.

  • Method

    The paper develops ViT-Plex and T5-Plex and evaluates model reliability through tasks involving confidence, distribution shifts, and adaptation.

  • Results

    Plex significantly improves reliability across deep-learning tasks and unifies reliability performance under general vision and language models; larger models and pretraining datasets further improve reliability.

  • Takeaways & Limitations

    Plex provides a broadly applicable framework and model extensions for stress-testing and improving reliability across diverse tasks and datasets.

  • Takeaways & Limitations

    The evaluation does not cover semantic class shift directly and restricts class-shift studies to open-set recognition because measuring predictive performance requires an open-vocabulary model.

Abstract

from arXiv · show

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

1. Reliability as a Goal for Artificial Intelligence

The paper frames reliability as consistent performance across uncertainty, robust generalization, and adaptation, rather than accuracy on narrowly matched test data. It evaluates this broader goal across vision and language using diverse tasks and datasets.

  • 1. Reliability as a Goal for Artificial Intelligence: Reliability means working consistently across real-world settings, including uncertainty, robust generalization, and adaptation.The definition is motivated by deployment settings that extend beyond standard in-distribution predictive accuracy.
  • 1. Reliability as a Goal for Artificial Intelligence: A reliable model should perform well across these areas out-of-the-box without customization for individual tasks.The paper treats the predictive probability distribution p(y|x) as the object to stress-test across benchmark suites.
  • 1. Reliability as a Goal for Artificial Intelligence: Prior research often targets narrow benchmarks or single capabilities, fragmenting reliability techniques across tasks.Examples include accuracy and calibration on ImageNet variants and specialized open set recognition methods.
  • 1.3 Contributions: The evaluation framework uses 10 task types and 40 downstream datasets spanning vision and natural language.The benchmark includes 14 finetuning datasets and 26 out-of-distribution evaluation datasets.
  • 1.3 Contributions: The paper introduces ImageNet Real-H for label uncertainty and NaLUE for uncertainty in conversational language understanding.These datasets address areas identified as missing from the existing literature.
  • 1.3 Contributions: ViT-Plex and T5-Plex extend large pretrained vision and language models, achieving strong performance across the reliability suite.The models are trained at scales up to 1 billion parameters and pretraining datasets up to 4 billion examples.

2. Evaluating Reliability

The paper evaluates reliability through a pipeline that pretrains, finetunes or adapts, and then measures downstream performance. Its task suite covers uncertainty, generalization under distribution shifts, and efficient adaptation across vision and language datasets.

  • 2. Evaluating Reliability: The evaluation pipeline pretrains a model, finetunes it on a downstream training split, and evaluates in-distribution and out-of-distribution metrics.Methods to improve reliability are tested during both pretraining and finetuning.
  • 2. Evaluating Reliability: The benchmark covers calibration, selective prediction, open set recognition, label uncertainty, generalization, distribution shifts, active learning, and few-shot settings.These tasks operationalize uncertainty, robust generalization, and adaptation across multiple datasets.
  • 2. Evaluating Reliability: Few-shot evaluation tests 1-, 5-, 10-, and 25-shot prediction, with additional calibration and open set recognition analyses.Zero-shot open set recognition uses Mahalanobis distance scoring from the model representation layer.
  • 2.2.1 Images: The study spans 19 vision distribution-shift datasets and introduces ImageNet ReaL-H for large-scale label-uncertainty evaluation.The vision suite includes natural web images, specialized domains, and datasets of varying sizes.
  • 2.2.2 Text: NaLUE addresses the lack of a large conversational-language corpus for uncertainty quantification and supports new long-tail and out-of-scope shifts.The proposed shifts include NaLUE-tail, Standard-OOS, and Near-OOS.

3. Plex: Pretrained Large model Extensions

Plex combines large Transformer backbones with scalable ensembling and reliability-oriented last-layer modifications. The vision and language variants build on ViT and T5 while exploring model scale, pretraining data, and adaptation choices.

  • Base Transformer architecture: ViT-Plex uses an encoder-only Vision Transformer, while T5-Plex uses an encoder-decoder Transformer for text.Images are tokenized into patches and text is tokenized with SentencePiece.
  • Model size: ViT-Plex is evaluated from approximately 22 million to 325 million parameters, while T5-Plex ranges from approximately 77 million to 880 million.These scales support experiments on how model size affects reliability.
  • Pretraining dataset size: Vision pretraining scales from ImageNet-21K to JFT with up to 4 billion images, while language uses the web-scraped C4 dataset.T5-Plex focuses exclusively on finetuning-only experiments because of compute constraints.
  • Efficient ensembling: Efficient ensembling uses BatchEnsemble in attention and feedforward layers, including configurations restricted to later layers for faster training.The approach is motivated by the effectiveness of ensembles and Bayesian neural networks for uncertainty and robustness.
  • Last layer changes: Plex modifies the final layer with Gaussian-process and heteroscedastic approaches to improve distance awareness and model input-dependent label noise.The GP layer increases uncertainty far from training representations, while the heteroscedastic method targets many-class settings.
  • Training and adaptation choices: The study tests efficient ensembling during pretraining and finetuning, while restricting some last-layer methods to finetuning.Few-shot experiments also compare frozen-representation logistic regression with gradient descent over all parameters.

4. Summary of Results and Scaling Trends

Plex improves reliability across diverse vision and language tasks, with gains from larger models, larger pretraining datasets, and specialized ensembling or last-layer methods. Pretraining loss strongly predicts downstream reliability, while dataset distribution shift appears more informative than downstream dataset size.

  • Plex sets new state-of-the-art on many reliability tasks while unifying performance across diverse tasks with one general vision or language model.The evaluation spans 139 vision metrics and 54 language metrics, and compares Plex with existing specialized state-of-the-art methods.
  • Scaling model size: Larger models improve overall reliability across uncertainty, robust generalization, and adaptation tasks.Reliability scores are normalized averages over 139 vision and 54 language task metrics, with separate averages for reliability areas.
  • Scaling pretraining dataset size: Larger pretraining datasets improve reliability: JFT 4B continues improving where JFT 300M reaches a performance limit, and ImageNet 10-shot accuracy consistently increases.The comparison uses separate runs with tuned learning-rate schedules and the same training-step settings.
  • Model ingredients: BatchEnsemble is consistently best during pretraining, while the strongest finetuning variants combine BatchEnsemble with Heteroscedastic or Gaussian-process last-layer methods.The best-ranked vision model uses Heteroscedasticity over a pretrained BatchEnsemble, and the best language variants use Gaussian processes or BatchEnsemble.
  • Downstream dataset size: Reliability shows no clear relationship with downstream training-set size, but lower reliability tends to occur when downstream data differ more from the JFT distribution.Examples include remote sensing, histology, and dog-breed datasets whose distributions differ from natural images in JFT.
  • Relationship between reliability tasks: Pretraining loss is highly predictive of reliability, with Pearson correlations above 0.97 for prediction and adaptation areas and 0.76 for uncertainty.Calibration error and calibration AUROC are exceptions among the analyzed uncertainty metrics, correlating less strongly with pretraining performance.

5. Reliability Task Results

Across reliability tasks, Plex generally improves calibration, selective prediction, open-set recognition, label uncertainty, and generalization, often benefiting from larger models or more data. Results also show that uncertainty-method choice and task distribution affect performance, while Plex extensions help beyond pretraining alone.

  • Selective prediction: At a 0.5% review budget, Plex L reaches 0.98 Oracle Collaborative AUROC on ImageNet, while all evaluated variants exceed 0.9 across the vision datasets.These results support using uncertainty-based deferral to reduce error in higher-risk applications.
  • Overall reliability results: Plex improves reliability across calibration, open-set recognition, selective prediction, label uncertainty, and generalization tasks.The reported gains span vision and language evaluations, with Plex often achieving strong or state-of-the-art performance without dataset-specific changes.
  • Label uncertainty: Heteroscedastic models outperform the None baseline for label uncertainty, with Plex best on ImageNet ReaL-H excluding the more-compute-intensive DE but worse on CIFAR-10H.The paper introduces ImageNet ReaL-H because existing datasets and highly accurate standard benchmarks provide limited label-uncertainty evaluation.
  • Generalization and ablations: Plex extensions improve beyond pretraining alone in diabetic retinopathy, while language in-distribution generalization benefits from pretraining, model scale, and Plex extensions.The paper attributes the limited pretraining benefit for retinal scans to their substantial appearance differences from pretraining images.

6. Conclusion

The paper frames reliable deep learning through diverse stress tests and introduces Plex extensions for vision and language. Plex targets consistent confidence, distribution-shift robustness, and adaptation across a broad evaluation suite.

  • 6. Conclusion: The authors develop Plex, pretrained large-model extensions for vision and language that significantly improve reliability in deep learning.The extensions are ViT-Plex and T5-Plex.
  • 6. Conclusion: The framework stress-tests reliability through confidence quantification, distribution-shift robustness, and rapid adaptation to new distributions.The evaluation includes multiple tasks and datasets spanning these reliability dimensions.
  • 6. Conclusion: The techniques are presented as broadly applicable to other large models and wider task suites.The conclusion cites LaMDA, PaLM, and BIG-bench as examples of broader applicability.

Appendix C. Additional Details on Tasks & Datasets

The appendix details the metrics and datasets used to evaluate reliability across vision and language. It covers distribution shifts, label uncertainty, subpopulation shifts, and task-specific evaluation resources.

  • Appendix C. Additional Details on Tasks & Datasets: Table 14 lists references for the task-specific state-of-the-art comparisons used in Figure 3.
  • Appendix C. Additional Details on Tasks & Datasets: Vision evaluation includes corruption, robustness, country-shift, severity-shift, semantic-shift, and subpopulation-shift datasets.Examples include CIFAR-10-C, CIFAR-100-C, ImageNet robustness datasets, RETINA Country Shift, and RETINA Severity Shift.
  • Appendix C. Additional Details on Tasks & Datasets: The evaluation uses task-specific metrics summarized in Table 15, with arrows indicating whether higher or lower values are better.The table organizes metrics by task.
  • Appendix C. Additional Details on Tasks & Datasets: Label uncertainty is evaluated with CIFAR-10H and the newly constructed ImageNet ReaL-H soft-label dataset.ImageNet ReaL-H uses newly collected individual human ratings for the original ImageNet test set.
  • Appendix C. Additional Details on Tasks & Datasets: Language evaluation uses MNLI-mismatched, CivilComments, HANS, CivilCommentsIdentity, and NaLUE-tail to assess covariate and subpopulation shifts.NaLUE-tail contains utterances from 28 low-frequency intent categories.

Appendix D. Details of Plex ingredients

The appendix describes model ingredients used to improve uncertainty quality and computational efficiency. These include BatchEnsemble, heteroscedastic last layers, and Gaussian-process output layers.

  • Appendix D. Details of Plex ingredients: BatchEnsemble shares weights across ensemble members to reduce the computational and memory costs of deep ensembles.Each member combines a shared weight matrix with a local rank-1 matrix through a Hadamard product.
  • Appendix D. Details of Plex ingredients: Plex applies BatchEnsemble to all layers in the last two residual blocks unless otherwise stated.This design follows prior mixture-of-experts work.
  • Appendix D. Details of Plex ingredients: SNGP replaces the dense output layer with a random-feature Gaussian process and uses spectral normalization to improve uncertainty quality for fixed representations.For pretrained networks, the authors find that only the last-layer Gaussian process is sufficient because pretrained embeddings already encode semantic distance.
  • Appendix D. Details of Plex ingredients: Heteroscedastic last layers model input-dependent label noise or label uncertainty using a parameter-efficient covariance approximation.The parameterization is inspired by BatchEnsemble and supports scaling to tens of thousands of classes.
  • Appendix D. Details of Plex ingredients: None, Het, GP, and BE require one forward pass from one model, whereas deep ensembles and Monte Carlo Dropout are more expensive.The comparison concerns both memory and compute requirements.

Appendix E. Additional Results for Open Set Recognition

The appendix compares five uncertainty scores for open set recognition across challenging tasks and model types. Mahalanobis distance performs best overall, while lighter alternatives offer task- and model-dependent trade-offs.

  • Appendix E. Additional Results for Open Set Recognition: The five compared OOD measures are MSP, Mahalanobis distance, Relative Mahalanobis distance, entropy, and MaxLogit.
  • Appendix E. Additional Results for Open Set Recognition: Mahalanobis distance outperforms the other OOD scores across the evaluated tasks, with Relative Mahalanobis distance performing second best.The comparison covers MSP, Mahalanobis distance, Relative Mahalanobis distance, entropy, and MaxLogit.
  • Appendix E. Additional Results for Open Set Recognition: Mahalanobis distance becomes computationally expensive as the number of classes grows because it requires fitting and evaluating class-specific Gaussians.For ImageNet2012 versus Places365, this involves 1,000 class Gaussians.
  • Appendix E. Additional Results for Open Set Recognition: For lighter scores, MaxLogit is best among single models except Het and None→Het, while entropy is best for ensemble models.
  • Appendix E. Additional Results for Open Set Recognition: The authors recommend Mahalanobis methods for best performance, with MaxLogit or entropy when computation is limited.MaxLogit is recommended for single models and entropy for ensemble models.

Appendix F. Additional Results for Selective Prediction

The appendix evaluates selective prediction with accuracy and uncertainty-based referral across RETINA shifts, while situating these results among broader reliability methods and sparse mixture-of-experts extensions.

  • Appendix F. Additional Results for Selective Prediction: Model selection uses area under the in-distribution selective-prediction accuracy curve, alongside accuracy, negative log-likelihood, calibration error, and selective-prediction metrics.Evaluation covers in-distribution and distributionally shifted data, with referral rates from 0% to 99%.
  • Appendix F. Additional Results for Selective Prediction: Plex L trails the SoTA ensemble in in-distribution RETINA Country Shift but outperforms it under out-of-distribution evaluation.Table 17 reports AUROC-based selective prediction performance for Plex L pretrained on ImageNet-21K or JFT.
  • Appendix F. Additional Results for Selective Prediction: RETINA Country Shift and Severity Shift both assess diabetic-retinopathy diagnosis, using accuracy and predictive entropy for selective prediction.Severity Shift restricts images to cases containing at worst moderate diabetic retinopathy.
  • Appendix F. Additional Results for Selective Prediction: Sparse mixture-of-experts methods conditionally activate subsets of experts, expanding parameter counts while maintaining constant training and inference costs.V-MoE routes image patches to selected experts, and the appendix uses MoE as shorthand for V-MoE.
  • Appendix F. Additional Results for Selective Prediction: Deep ensembles of V-MoEs combine static ensembling with sparse-MoE adaptivity, with reported benefits across few-shot classification and out-of-distribution detection.The appendix defines notation for four-member MoE ensembles and related upstream/downstream variants.

G.2 Results

The sparse-MoE evaluation extends prior results with the paper’s prediction, uncertainty, and adaptation scores. The combined ensemble [MoE]4 outperforms the standard deep ensemble [None]4 at equal computational cost.

  • G.2 Results: The evaluation extends Allingham et al. (2021) with additional computer-vision metrics and summarizes prediction, uncertainty, and adaptation scores.These results are presented in Table 18.
  • G.2 Results: [MoE]4 outperforms [None]4 while having the same computational cost, supporting the complementarity of static ensembling and sparse-MoE adaptivity.The authors characterize combining ensembling and sparse MoEs as a promising direction for improving model reliability.
  • G.2 Results: MoE and [MoE]4 show strong few-shot-learning performance, while MoE has lower uncertainty scores than None because of weaker calibration.E3 performs comparably to deep ensembles formed from a single pretrained model.

Appendix H. Analysis of Heteroscedastic Last Layer

This analysis adds heteroscedastic output layers to deterministic and BatchEnsemble base models to test whether modeling input-dependent label noise improves performance on noisy-label datasets.

  • Appendix H. Analysis of Heteroscedastic Last Layer: Heteroscedastic output layers are added to both the base deterministic model and the BatchEnsemble base model.The modification targets uncertainty associated with input-dependent label noise.
  • Appendix H. Analysis of Heteroscedastic Last Layer: The experiment tests whether modeling input-dependent label noise improves performance on datasets known to contain noisy labels.The comparison concerns the effect of the output-layer design on downstream performance.

H.1 Heteroscedastic improves pretraining performance on JFT.

Heteroscedastic modeling improves performance on JFT when applied over the base deterministic model, while its effect on ImageNet-21K is neutral.

  • H.1 Heteroscedastic improves pretraining performance on JFT.: Performance gains on JFT and neutral performance on ImageNet-21K are observed for the heteroscedastic model over the base deterministic model.The comparison is shown in Figures 29 and 30.
  • H.1 Heteroscedastic improves pretraining performance on JFT.: A closer inspection finds that [MoE]4 has a slightly worse prediction score than [None]4 because of isolated weaker performance on CIFAR-10.This caveat qualifies the broader sparse-MoE ensemble comparison.
  • H.1 Heteroscedastic improves pretraining performance on JFT.: Figure 31 compares deterministic, heteroscedastic, and BatchEnsemble variants across upstream and downstream configurations on ImageNet1K.It also includes deterministic-upstream/heteroscedastic-downstream and BatchEnsemble-plus-heteroscedastic configurations.
  • H.1 Heteroscedastic improves pretraining performance on JFT.: Figure 32 compares the same model-family variants on CIFAR-10, including heteroscedastic downstream and BatchEnsemble configurations.The listed variants distinguish upstream and downstream model changes.
  • H.1 Heteroscedastic improves pretraining performance on JFT.: Figure 33 compares these deterministic, heteroscedastic, and BatchEnsemble configurations on CIFAR-100.The figure includes both standalone and combined upstream/downstream variants.

H.2 Impact of heteroscedastic last layer downstream

Downstream heteroscedastic heads improve reliability without requiring heteroscedastic pretraining, while avoiding some of its extra cost and tuning requirements.

  • None→Het performs similarly or better than Het→Het across all downstream datasets.This indicates that downstream heteroscedastic gains can be obtained without heteroscedastic upstream pretraining.
  • Heteroscedastic models are more expensive than deterministic models and require tuning temperature and low-rank approximation rank hyperparameters.
  • None→Het outperforms None on all metrics for ImageNet and CIFAR-100, while CIFAR-10 performance is saturated and similar.Applying the heteroscedastic head downstream improves performance beyond in-distribution accuracy, NLL, and ECE.

Appendix I. Summarization of Language Results

Language experiments compare uncertainty methods and model sizes across prediction, calibration, and human-model collaboration under multiple population shifts. Efficient methods such as Plex and Batch Ensemble perform strongly, while larger architectures generally improve collaborative performance.

  • The uncertainty-method comparison fixes the architecture at T5-base and evaluates prediction, calibration, and human-model collaboration across MNLI, NaLUE, and Toxic Comments.Evaluations cover in-domain, out-of-distribution, and tail-population splits; Table 20 contains full results and Figure 34 summarizes rankings.
  • DE+GP, Plex, BE, and MC Dropout tend to have the strongest performance across language uncertainty evaluations.DE+GP nearly dominates MNLI and NaLUE but costs 10x more memory and compute, whereas efficient BE and Plex perform well on those datasets.
  • The architecture-scaling study evaluates Plex, None, and MC Dropout on progressively larger T5 S, T5 B, and T5 L architectures.Table 20 reports full results and Figure 35 organizes rankings by architecture size.
  • Larger architectures almost always lead to stronger collaborative performance, including under out-of-distribution evaluation.
Loading 2207.07411v1…