On the Opportunities and Risks of Foundation Models

Real-world ML systems need to be robust to distribution shifts — they should work well on test distributions which differ from the train distribution. High-stakes applications such as poverty mapping in under-resourced countries (Xie et al., 2016; Jean et al., 2016), self-driving cars (Yu et al., 2020a; Sun et al., 2020a), and medical diagnosis (AlBadawy et al., 2018; Dai and Gool, 2018) all require models that generalize well to circumstances not seen in the training data, e.g., test examples from different countries, under different driving conditions, or from different hospitals. Prior work has shown that these types of distribution shifts can cause large drops in performance even in state-of-the-art models (Blitzer et al., 2006; Daumé III, 2007; Sugiyama et al., 2007; Ganin and Lempitsky, 2015; Peng et al., 2019; Kumar et al., 2020a; Arjovsky et al., 2019; Szegedy et al., 2014; Hendrycks and Dietterich, 2019; Sagawa et al., 2020a; Recht et al., 2019; Abney, 2007; Ruder and Plank, 2018; Geirhos et al., 2018; Kumar et al., 2020b; Yu et al., 2020b; Geirhos et al., 2020; Xie et al., 2021a; Koh et al., 2021).

In this section, we consider the role of foundation models on robustness to distribution shifts. A foundation model is trained on a large and diverse unlabeled dataset sampled from a distribution ppre and can be adapted to many downstream tasks. For each downstream task 𝒯, the foundation model is adapted to labeled training data sampled from an in-distribution (ID) training distribution pID, and then evaluated on an out-of-distribution (OOD) test distribution pOOD. For example, a poverty prediction model (Xie et al., 2016; Jean et al., 2016) may be pretrained on unlabeled satellite data from across the world to learn useful features for all countries, then fine-tuned on labeled examples from Nigeria, and finally evaluated in Malawi where labeled examples are scarce.

We argue that 1) foundation models are a particularly promising approach to robustness. Existing work shows that pretraining on unlabeled data is an effective, general-purpose way to improve accuracy on OOD test distributions, in contrast to many robustness interventions which are constrained to narrow types of distribution shifts. However, we also discuss why 2) foundation models may not always mitigate distribution shifts, such as shifts due to spurious correlations or changes over time. Finally, 3) we outline several research directions to leverage and improve foundation models for robustness.

We note that one of the ways in which foundation models lead to improved extrapolation is by providing inductive biases (via model initialization) for the adapted model, which are learned on a diverse dataset that extends beyond the downstream training data. However, this same inductive bias can also encode harmful associations from the pretrained data and lead to representational and allocational harms in the presence of distribution shift. See §4.6: data and §5.1: fairness for further discussion of such harms and methods for mitigation.

Advantages

By learning representations on a large and diverse foundation model training distribution ppre, foundation models can improve accuracy of the adapted derivative on the downstream test distribution pOOD. OpenAI’s CLIP model, which is a foundation model trained on a diverse set of images and natural language documents, has been shown to be robust to some benchmark distribution shifts on ImageNet (Radford et al., 2021): for example, both CLIP and a standard ResNet50 obtain 76% accuracy on ImageNet, but CLIP achieves 6% higher accuracy on ImageNetV2 (Recht et al., 2019) and 35% higher accuracy on ImageNet Sketch (Radford et al., 2021), which are both related but different from the original ImageNet training distribution. In contrast, many other robustness interventions, such as adversarial training (Madry et al., 2018), invariant risk minimization (Arjovsky et al., 2019), or using larger models have had little impact on effective robustness (defined as the gap between in-distribution and out-of-distribution performance) on these ImageNet tasks, especially without explicit knowledge of the distribution shift (Taori et al., 2020; Santurkar et al., 2020; Radford et al., 2021; Miller et al., 2021).

Many other works demonstrate that pretraining on large datasets can improve robustness to common image corruptions, label shift, and label corruptions (Hendrycks et al., 2019a; Hendrycks et al., 2019b); to real-world spatial shifts in satellite imagery tasks (Xie et al., 2021a; Kumar et al., 2022); and to shifts across topics in natural language understanding tasks (Hendrycks et al., 2020; Fisch et al., 2019; Yogatama et al., 2019). As another example, diversifying the foundation model training data to include multiple languages (as in multilingual BERT (Liu et al., 2020b)) significantly improves performance in unseen language pairs.

Persistent challenges

Despite promising signs that foundation models will result in substantial improvements to robustness, we anticipate that foundation models are not a panacea for distribution shifts. We discuss this in the context of two broad categories of distribution shifts below.

Spurious correlations

Spurious correlations are statistical correlations between features and labels with predictive power on the training distribution but not on the test distribution (Heinze-Deml and Meinshausen, 2017; Arjovsky et al., 2019; Sagawa et al., 2020a). Well-known examples include reliance on background color for object recognition (Xiao et al., 2020), surgical markers for medical diagnostics (Winkler et al., 2019), annotator biases in crowdsourced data (Tsuchiya, 2018; Gururangan et al., 2018; Poliak et al., 2018; Geva et al., 2019), and demographic biases (Abid et al., 2021; Nadeem et al., 2021; Gehman et al., 2020). Models learn these spurious correlations largely because the foundation model training and adaptation data exhibit these biases (Nagarajan et al., 2020; Gehman et al., 2020), and this issue cannot simply be addressed with larger models (Sagawa et al., 2020b).

Foundation models may exacerbate or mitigate the effects of spurious correlations, but this depends on the nature of the particular downstream task and its relation to the foundation model training data and algorithm. By training with a diverse dataset, foundation models may improve robustness to spurious correlations that are found only in a subset of the training data: e.g., existing studies find that pretrained language models can avoid spurious correlations by quickly learning from counterexamples to the spurious correlations (Tu et al., 2020). However, foundation models can also exacerbate the issue by introducing biases present in the foundation model training data, as observed for demographic biases in GPT-3 and other NLP models (Abid et al., 2021; Nadeem et al., 2021; Gehman et al., 2020). Moreover, training at scale alone need not fully address the root issue of identifying and not relying on the features that are predictive on the downstream training set but not on the downstream test set (Heinze-Deml and Meinshausen, 2017). Addressing these challenges will require us to understand and manage the inductive bias from foundation model training and develop adaptation algorithms that are resistant to learning spurious correlations.

Extrapolation and temporal drift

Finally, the few- and zero-shot capabilities of foundation models will mean that these models will increasingly be used far beyond the training distribution. While large-scale foundation model training can help with certain forms of extrapolation to new distributions (Papadimitriou and Jurafsky, 2020), there may be limits to their extrapolation capabilities. For example, existing language models cannot handle changes to world knowledge or language change without re-training (Lazaridou et al., 2021; Dhingra et al., 2021), zero-shot transfer in CLIP suffers greatly in satellite image domains (Radford et al., 2021), and ImageNet pretraining does not substantially improve the performance of large models on medical images (Raghu et al., 2019; Ke et al., 2021). We believe that foundation models cannot be assumed to automatically extrapolate within a given modality (e.g., all images), and it will become increasingly important to define and separate the forms of extrapolation that are newly enabled by foundation models from those that remain out of reach. Though existing taxonomies for distribution shifts have been proposed in generality (Quiñonero-Candela et al., 2009; Ye et al., 2021), fully understanding and defining the types of distribution shifts for which foundation models are effective is a major open problem for robustness research.

Opportunities

Foundation models hold substantial promise as a general-purpose robustness intervention for distribution shifts and open new avenues for robustness research. We outline some opportunities and open questions below.

Understanding foundation model representations

Existing studies of the robustness of foundation models have been largely empirical, and there is little understanding of the mechanism behind gains in robustness. Sun et al. (2019b) hypothesize that pretrained representations bring disparate domains (such as ID and OOD distributions) closer together, which can in turn improve generalization from labeled ID data to OOD data (Ben-David et al., 2010). Controlled experimentation on measuring the distance between domain representations with and without pretraining can elucidate this effect. There are initial promising directions in characterizing foundation model training (e.g., contrastive learning as a spectral graph decomposition (HaoChen et al., 2021a)) and their inductive biases (Saunshi et al., 2020a; Lee et al., 2020a; Zhang and Hashimoto, 2020; Xie et al., 2020). However these theories are limited and fail to address other empirically effective foundation models such as fully generative language models (e.g., GPT-3 (Brown et al., 2020) and image-GPT (Chen et al., 2020d)). Further understanding how these inductive biases are useful under distribution shift may lead to a more complete theory (§4.10: theory) of how foundation models improve robustness.

Data augmentation in foundation model training

While foundation models trained without knowledge of the downstream tasks can avoid some task-specific biases and often improve robustness, certain statistical biases stemming from how the foundation model was trained may persist. As a concrete example, many contemporary self-supervision algorithms are heavily dependent on choosing an appropriate set of data augmentations (Chen et al., 2020c), which in turn confers different types of robustness in the adaptation phase. For instance, Xiao et al. (2021) show that a foundation model for vision trained with contrastive learning on rotation augmentations may improve OOD performance on adaptation tasks with rotation invariance, but may not improve robustness for tasks where OOD generalization requires other invariances. Further research into what types of data augmentations improve robustness for a wide range of downstream tasks — including data augmentations that are learned from data (Wong and Kolter, 2020; Tamkin et al., 2021c) or designed to be generally applicable across data modalities (Verma et al., 2021) — will inform better foundation model training algorithms (§4.2: training).

Encoding structure in foundation model training

In general, exploring new ways of encoding known structure and invariances in the data is an important path forward for foundation model training. Many real-world tasks have additional metadata (e.g., spatial location coordinates, climate information from auxiliary satellites in our poverty prediction example), which may provide additional structure for OOD generalization (e.g., across geographic areas) (Xie et al., 2021a; Koh et al., 2021). For example, Xie et al. (2021a) show that metadata can be used as targets for pretraining to improve downstream OOD accuracy. In language, modeling the tags in HTML data provides additional downstream-task-adjacent supervision, allows for new forms of prompting (e.g., filling in <title> tags for title suggestion), and improves data efficiency (Aghajanyan et al., 2021). While current data augmentation methods encode hand-crafted knowledge, other avenues such as exploiting metadata could provide more automated ways of determining which structures and invariances to incorporate for foundation model training.

Specialization vs. diversity in foundation model training data

The choice of foundation model training data has downstream effects — training on a more diverse dataset is not always better for downstream performance than a more specialized foundation model (Cole et al., 2021; Chalkidis et al., 2020) (see §4.3: adaptation for a more detailed discussion). In some domains such as satellite images and specialized text topics, continued pretraining on the specialized domain improves the downstream performance significantly (Reed et al., 2021; Gururangan et al., 2020). This is a potential source of tension: on one hand, we might want to train the foundation model on a large, diverse dataset in order to have more robust performance under distribution shifts, while on the other hand, we might need to specialize the foundation model to improve its in-distribution and out-of-distribution performance on downstream tasks. A better understanding of how specialization affects the in-distribution and out-of-distribution performance of foundation models will allow us to design and collect more effective foundation model training sets.

Adaptation methods

Although foundation models provide a strong starting point, how the adaptation method uses the pretrained information can affect robustness. For instance, lightweight tuning methods for language models (e.g., adapter/prefix/prompt tuning (Houlsby et al., 2019; Li and Liang, 2021; Lester et al., 2021)), which adapt the model for a new task by optimizing a small set of parameters (such as a continuous prompt) while keeping the other foundation model parameters frozen, seem to give OOD performance benefits (§4.3: adaptation). Xie et al. (2021b) explain this in a special case, where composing a learned model with a frozen foundation model can reduce the complexity of the learned model, improving generalization both ID and OOD. On vision datasets, Wortsman et al. (2021); Kumar et al. (2022) find that freezing a foundation model and training only the head can lead to better OOD performance than fine-tuning the entire model. Kumar et al. (2022) explain this theoretically by showing that full fine-tuning can distort pretrained features even in a simple setting (two layer linear networks). However, it is still poorly understood in general why freezing parameters seems to improve OOD performance. Finally, while current adaptation methods may suffice for good ID generalization, the methods do not explicitly account for distribution shift in their design. As a first step, we can investigate how methods for distribution shifts such as domain adaptation, domain generalization, and semi-supervised learning methods interact with foundation models when used for adaptation. Progress in these directions can lead to adaptation methods that can better leverage foundation models for robustness.