WILDS: A Benchmark of in-the-Wild Distribution Shifts
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, Percy Liang
Introduction
Distribution shifts—where the training distribution differs from the test distribution—can significantly degrade the accuracy of machine learning (ML) systems deployed in the wild. In this work, we consider two types of distribution shifts that are ubiquitous in real-world settings: domain generalization and subpopulation shift (Figure 1). In domain generalization, the training and test distributions comprise data from related but distinct domains. This problem arises naturally in many applications, as it is often infeasible to collect a training set that spans all domains of interest. For example, in medical applications, it is common to seek to train a model on patients from a few hospitals, and then deploy it more broadly to hospitals outside the training set (Zech et al., 2018); and in wildlife monitoring, we might seek to train an animal recognition model on images from one set of camera traps and then deploy it to new camera traps (Beery et al., 2018). In subpopulation shift, we consider test distributions that are subpopulations of the training distribution, with the goal of doing well even on the worst-case subpopulation. For example, it is well-documented that standard models often perform poorly on under-represented demographics (Buolamwini and Gebru, 2018; Koenecke et al., 2020), and so we might seek models that can perform well on all demographic subpopulations.
Despite their ubiquity in real-world deployments, these types of distribution shifts are under-represented in the datasets widely used in the ML community today (Geirhos et al., 2020). Most of these datasets were designed for the standard i.i.d. setting, with training and test sets from the same distribution, and prior work on retrofitting them with distribution shifts has focused on shifts that are cleanly characterized but not always likely to arise in real-world deployments. For instance, many recent papers have studied datasets with shifts induced by synthetic transformations, such as changing the color of MNIST digits (Arjovsky et al., 2019), or by disparate data splits, such as generalizing from cartoons to photos (Li et al., 2017a). Datasets like these are important testbeds for systematic studies, but they do not generally reflect the kinds of shifts that are likely to arise in the wild. To develop and evaluate methods for real-world shifts, we need to complement these datasets with benchmarks that capture shifts in the wild, as model robustness need not transfer across shifts: e.g., models can be robust to image corruptions but not to shifts across datasets (Taori et al., 2020; Djolonga et al., 2020), and a method that improves robustness on a standard vision dataset can even consistently harm robustness on real-world satellite imagery datasets (Xie et al., 2020).
In this paper, we present Wilds, a curated benchmark of 10 datasets with evaluation metrics and train/test splits representing a broad array of distribution shifts that ML models face in the wild (Figure 2). With Wilds, we seek to complement existing benchmarks by focusing on datasets with realistic shifts across a diverse set of data modalities and applications: animal species categorization (Beery et al., 2020a), tumor identification (Bandi et al., 2018), bioassay prediction (Wu et al., 2018; Hu et al., 2020b), genetic perturbation classification (Taylor et al., 2019), wheat head detection (David et al., 2020), text toxicity classification (Borkan et al., 2019b), land use classification (Christie et al., 2018), poverty mapping (Yeh et al., 2020), sentiment analysis (Ni et al., 2019), and code completion (Raychev et al., 2016; Lu et al., 2021). These datasets reflect natural distribution shifts arising from different cameras, hospitals, molecular scaffolds, experiments, demographics, countries, time periods, users, and codebases.
Wilds builds on extensive data-collection efforts by domain experts, who are often forced to grapple with distribution shifts to make progress in their applications. To design Wilds, we worked with them to identify, select, and adapt datasets that fulfilled the following criteria:
Distribution shifts with performance drops. The train/test splits reflect shifts that substantially degrade model performance, i.e., with a large gap between in-distribution and out-of-distribution performance.
Real-world relevance. The training/test splits and evaluation metrics are motivated by real-world scenarios and chosen in conjunction with domain experts. In Appendix A, we further discuss the framework we use to assess the realism of a dataset.
Potential leverage. Distribution shift benchmarks must be non-trivial but also possible to solve, as models cannot be expected to generalize to arbitrary distribution shifts. We constructed each Wilds dataset to have training data from multiple domains, with domain annotations and other metadata available at training time. We hope that these can be used to learn robust models: e.g., for domain generalization, one could use these annotations to learn models that are invariant to domain-specific features (Sun and Saenko, 2016; Ganin et al., 2016), while for subpopulation shift, one could learn models that perform uniformly well across each subpopulation (Hu et al., 2018; Sagawa et al., 2020a).
We chose the Wilds datasets to collectively encompass a diverse set of tasks, data modalities, dataset sizes, and numbers of domains, so as to enable evaluation across a broad range of real-world distribution shifts. In Section 8, we further survey the distribution shifts that occur in other application areas—algorithmic fairness and policing, medicine and healthcare, genomics, natural language and speech processing, education, and robotics—and discuss examples of datasets from these areas that we considered but did not include in Wilds, as their distribution shifts did not cause an appreciable performance drop.
To make the Wilds datasets more accessible, we have substantially modified most of them to clarify the distribution shift, standardize the data splits, and preprocess the data for use in standard ML frameworks. In Section 10, we introduce our accompanying open-source Python package that fully automates data loading and evaluation. The package also includes default models appropriate for each dataset, allowing all of the baseline results reported in this paper to be easily replicated. To track the state-of-the-art in training algorithms and model architectures that are robust to these distribution shifts, we are also hosting a public leaderboard; we discuss guidelines for developers in Section 9. Code, leaderboards, and updates are available at https://wilds.stanford.edu.
Datasets are significant catalysts for ML research. Likewise, benchmarks that curate and standardize datasets—e.g., the GLUE and SuperGLUE benchmarks for language understanding (Wang et al., 2019a, b) and the Open Graph Benchmark for graph ML (Hu et al., 2020b)—can accelerate research by focusing community attention, easing development on multiple datasets, and enabling systematic comparisons between approaches. In this spirit, we hope that Wilds will facilitate the development of ML methods and models that are robust to real-world distribution shifts and can therefore be deployed reliably in the wild.
Existing ML benchmarks for distribution shifts
Distribution shifts have been a longstanding problem in the ML research community (Hand, 2006; Quiñonero-Candela et al., 2009). Earlier work studied shifts in datasets for tasks including part-of-speech tagging (Marcus et al., 1993), sentiment analysis (Blitzer et al., 2007), land cover classification (Bruzzone and Marconcini, 2009), object recognition (Saenko et al., 2010), and flow cytometry (Blanchard et al., 2011). However, these datasets are not as widely used today, in part because they tend to be much smaller than modern datasets.
Instead, many recent papers have focused on object recognition datasets with shifts induced by synthetic transformations, such as ImageNet-C (Hendrycks and Dietterich, 2019), which corrupts images with noise; the Backgrounds Challenge (Xiao et al., 2020) and Waterbirds (Sagawa et al., 2020a), which alter image backgrounds; or Colored MNIST (Arjovsky et al., 2019), which changes the colors of MNIST digits. It is also common to use data splits or combinations of disparate datasets to induce shifts, such as generalizing to photos solely from cartoons and other stylized images in PACS (Li et al., 2017a); generalizing to objects at different scales solely from a single scale in DeepFashion Remixed (Hendrycks et al., 2020b); or using training and test sets with disjoint subclasses in BREEDS (Santurkar et al., 2020) and similar datasets (Hendrycks and Dietterich, 2019). While our treatment here is necessarily brief, we discuss other similar datasets in Appendix B.
These existing benchmarks are useful and important testbeds for method development. As they typically target well-defined and isolated shifts, they facilitate clean analysis and controlled experimentation, e.g., studying the effect of backgrounds on image classification (Xiao et al., 2020), or showing that training with added Gaussian blur improves performance on real-world blurry images (Hendrycks et al., 2020b). Moreover, by studying how off-the-shelf models trained on standard datasets like ImageNet perform on different test datasets, we can better understand the robustness of these widely-used models (Geirhos et al., 2018b; Recht et al., 2019; Hendrycks and Dietterich, 2019; Taori et al., 2020; Djolonga et al., 2020; Hendrycks et al., 2020b).
However, as we discussed in the introduction, robustness to these synthetic shifts need not transfer to the kinds of shifts that arise in real-world deployments (Taori et al., 2020; Djolonga et al., 2020; Xie et al., 2020), and it is thus challenging to develop and evaluate methods for training models that are robust to real-world shifts on these datasets alone. With WILDS, we seek to complement existing benchmarks by curating datasets that reflect natural distribution shifts across a diverse set of data modalities and application.
Problem settings
Each Wilds dataset is associated with a type of domain shift: domain generalization, subpopulation shift, or a hybrid of both (Figure 2). We focus on these types of distribution shifts because they collectively capture the structure of most of the shifts in the applications we studied; see Section 8 for more discussion. In each setting, we can view the overall data distribution as a mixture of domains . Each domain corresponds to a fixed data distribution over , where is the input, is the prediction target, and all points sampled from have domain . We encode the domain shift by assuming that the training distribution has mixture weights for each domain , while the test distribution is a different mixture of domains with weights . For convenience, we define the set of training domains as , and likewise, the set of test domains as .
At training time, the learning algorithm gets to see the domain annotations , i.e., the training set comprises points . At test time, the model gets either or drawn from , depending on the application.
In domain generalization, we aim to generalize to test domains that are disjoint from the training domains , i.e., . To make this problem tractable, the training and test domains are typically similar to each other: e.g., in Camelyon17-wilds, we train on data from some hospitals and test on a different hospital, and in iWildCam2020-wilds, we train on data from some camera traps and test on different camera traps. We typically seek to minimize the average error on the test distribution.
2 Subpopulation shift (Figure 1-Bottom)
In subpopulation shift, we aim to perform well across a wide range of domains seen during training time. Concretely, all test domains are seen at training, with , but the proportions of the domains can change, with . We typically seek to minimize the maximum error over all test domains. For example, in CivilComments-wilds, the domains represent particular demographics, some of which are a minority in the training set, and we seek high accuracy on each of these subpopulations without observing their demographic identity at test time.
3 Hybrid settings
The categories of domain generalization and subpopulation shift provide a general framework for thinking about domain shifts, and the methods that have been developed for each setting have been quite different, as we will discuss in Section 6. However, it is not always possible to cleanly define a problem as one or the other; for example, a test domain might be present in the training set but at a very low frequency. In Wilds, we consider some hybrid settings that combine both domain generalization and subpopulation shift. For example, in FMoW-wilds, the inputs are satellite images and the domains correspond to the year and geographical region in which they were taken. We simultaneously consider domain generalization across time (the training/test sets comprise images taken before/after a certain year) and subpopulation shift across regions (there are images from the same regions in the training and test sets, and we seek high performance across all regions).
Wilds datasets
We now briefly describe each Wilds dataset, as summarized in Figure 2. For each dataset, we consider a problem setting—domain generalization, subpopulation shift, or a hybrid—that we believe best reflects the real-world challenges in the corresponding application area; see Appendix A for more discussion of these considerations. To avoid confusion between our modified datasets and their original sources, we append -wilds to the dataset names. We provide more details and context on related distribution shifts for each dataset in Appendix E.
Animal populations have declined 68% on average since 1970 (Grooten et al., 2020). To better understand and monitor wildlife biodiversity loss, ecologists commonly deploy camera traps—heat or motion-activated static cameras placed in the wild (Wearn and Glover-Kapfer, 2017)—and then use ML models to process the data collected (Weinstein, 2018; Norouzzadeh et al., 2019; Tabak et al., 2019; Beery et al., 2019; Ahumada et al., 2020). Typically, these models would be trained on photos from some existing camera traps and then used across new camera trap deployments. However, across different camera traps, there is drastic variation in illumination, color, camera angle, background, vegetation, and relative animal frequencies, which results in models generalizing poorly to new camera trap deployments (Beery et al., 2018).
We study this shift on a variant of the iWildCam 2020 dataset (Beery et al., 2020a), where the input is a photo from a camera trap, the label is one of 182 animal species, and the domain specifies the identity of the camera trap (Figure 3). The training and test sets comprise photos from disjoint sets of camera traps. As leverage, we include over 200 camera traps in the training set, capturing a wide range of variation. We evaluate models by their macro F1 scores, which emphasizes performance on rare species, as rare and endangered species are the most important to accurately monitor. Appendix E.1 provides additional details and context.
1.2 Camelyon17-wilds: Tumor identification across different hospitals
Models for medical applications are often trained on data from a small number of hospitals, but with the goal of being deployed more generally across other hospitals. However, variations in data collection and processing can degrade model accuracy on data from new hospital deployments (Zech et al., 2018; AlBadawy et al., 2018). In histopathology applications—studying tissue slides under a microscope—this variation can arise from sources like differences in the patient population or in slide staining and image acquisition (Veta et al., 2016; Komura and Ishikawa, 2018; Tellez et al., 2019).
We study this shift on a patch-based variant of the Camelyon17 dataset (Bandi et al., 2018), where the input is a 96x96 patch of a whole-slide image of a lymph node section from a patient with potentially metastatic breast cancer, the label is whether the patch contains tumor, and the domain specifies which of 5 hospitals the patch was from (Figure 4). The training and test sets comprise class-balanced patches from separate hospitals, and we evaluate models by their average accuracy. Prior work suggests that staining differences are the main source of variation between hospitals in similar datasets (Tellez et al., 2019). As we have training data from multiple hospitals, a model could use that as leverage to learn to be robust to stain variation. Appendix E.2 provides additional details and context.
1.3 RxRx1-wilds: Genetic perturbation classification across experimental batches
High-throughput screening techniques that can generate large amounts of data are now common in many fields of biology, including transcriptomics (Harrill et al., 2019), genomics (Echeverri and Perrimon, 2006; Zhou et al., 2014), proteomics and metabolomics (Taylor et al., 2021), and drug discovery (Broach et al., 1996; Macarron et al., 2011; Swinney and Anthony, 2011; Boutros et al., 2015). Such large volumes of data, however, need to be created in experimental batches, or groups of experiments executed at similar times under similar conditions. Despite attempts to carefully control experimental variables such as temperature, humidity, and reagent concentration, measurements from these screens are confounded by technical artifacts that arise from differences in the execution of each batch. These batch effects make it difficult to draw conclusions from data across experimental batches (Leek et al., 2010; Parker and Leek, 2012; Soneson et al., 2014; Nygaard et al., 2016; Caicedo et al., 2017).
We study the shift induced by batch effects on a variant of the RxRx1 dataset (Taylor et al., 2019), where the input is a 3-channel image of cells obtained by fluorescent microscopy (Bray et al., 2016), the label indicates which of the 1,139 genetic treatments (including no treatment) the cells received, and the domain specifies the batch in which the imaging experiment was run. As summarized in Figure 5, the training and test sets consist of disjoint experimental batches. As leverage, the training set has images from 33 different batches, with each batch containing one sample for every class. We assess a model’s ability to normalize batch effects while preserving biological signal by evaluating how well it can classify images of treated cells in the out-of-distribution test set. Appendix E.3 provides additional details and context.
1.4 OGB-MolPCBA: Molecular property prediction across different scaffolds
Accurate prediction of the biochemical properties of small molecules can significantly accelerate drug discovery by reducing the need for expensive lab experiments (Shoichet, 2004; Hughes et al., 2011). However, the experimental data available for training such models is limited compared to the extremely diverse and combinatorially large universe of candidate molecules that we would want to make predictions on (Bohacek et al., 1996; Sterling and Irwin, 2015; Lyu et al., 2019; McCloskey et al., 2020). This means that models need to generalize to out-of-distribution molecules that are structurally different from those seen in the training set.
We study this shift on the OGB-MolPCBA dataset, which is directly adopted from the Open Graph Benchmark (Hu et al., 2020b) and originally from MoleculeNet (Wu et al., 2018). As summarized in Figure 6, it is a multi-label classification dataset, where the input is a molecular graph, the label is a 128-dimensional binary vector where each component corresponds to a biochemical assay result, and the domain specifies the scaffold (i.e., a cluster of molecules with similar structure). The training and test sets comprise molecules with disjoint scaffolds; for leverage, the training set has molecules from over 40,000 scaffolds. We evaluate models by averaging the Average Precision (AP) across each of the 128 assays. Appendix E.4 provides additional details and context.
1.5 GlobalWheat-wilds: Wheat head detection across regions of the world
Models for automated, high-throughput plant phenotyping—measuring the physical characteristics of plants and crops, such as wheat head density and counts—are important tools for crop breeding (Thorp et al., 2018; Reynolds et al., 2020) and agricultural field management (Shi et al., 2016). These models are typically trained on data collected in a limited number of regions, even for crops grown worldwide such as wheat (Madec et al., 2019; Xiong et al., 2019; Ubbens et al., 2020; Ayalew et al., 2020). However, there can be substantial variation between regions, due to differences in crop varieties, growing conditions, and data collection protocols. Prior work on wheat head detection has shown that this variation can significantly degrade model performance on regions unseen during training (David et al., 2020).
We study this shift in an expanded version of the Global Wheat Head Dataset (David et al., 2020, 2021), a large set of wheat images collected from 12 countries around the world (Figure 7). It is a detection dataset, where the input is a cropped overhead image of a wheat field, the label is the set of bounding boxes for each wheat head visible in the image, and the domain specifies an image acquisition session (i.e., a specific location, time, and sensor with which a set of images was collected). The data split captures a shift in location, with training and test sets comprising images from disjoint countries. As leverage, we include images from 18 acquisition sessions over 5 countries in the training set. We evaluate model performance on unseen countries by measuring accuracy at a fixed Intersection over Union (IoU) threshold, and averaging across acquisition sessions to account for imbalances in the numbers of images in them. Additional details are provided in Appendix E.5.
2 Subpopulation shift datasets
Automatic review of user-generated text is an important tool for moderating the sheer volume of text written on the Internet. We focus here on the task of detecting toxic comments. Prior work has shown that toxicity classifiers can pick up on biases in the training data and spuriously associate toxicity with the mention of certain demographics (Park et al., 2018; Dixon et al., 2018). These types of spurious correlations can significantly degrade model performance on particular subpopulations (Sagawa et al., 2020a).
We study this problem on a variant of the CivilComments dataset (Borkan et al., 2019b), a large collection of comments on online articles taken from the Civil Comments platform (Figure 8). The input is a text comment, the label is whether the comment was rated as toxic, and the domain is a 8-dimensional binary vector where each component corresponds to whether the comment mentions one of the 8 demographic identities male, female, LGBTQ, Christian, Muslim, other religions, Black, and White. The training and test sets comprise comments on disjoint articles, and we evaluate models by the lowest true positive/negative rate over each of these 8 demographic groups; these groups overlap with each other, deviating slightly from the standard subpopulation shift framework in Section 3. Models can use the provided domain annotations as leverage to learn to perform well over each demographic group. Appendix E.6 provides additional details and context.
3 Hybrid datasets
ML models for satellite imagery can enable global-scale monitoring of sustainability and economic challenges, aiding policy and humanitarian efforts in applications such as deforestation tracking (Hansen et al., 2013), population density mapping (Tiecke et al., 2017), crop yield prediction (Wang et al., 2020b), and other economic tracking applications (Katona et al., 2018). As satellite data constantly changes due to human activity and environmental processes, these models must be robust to distribution shifts over time. Moreover, as there can be disparities in the data available between regions, these models should ideally have uniformly high accuracies instead of only doing well on data-rich regions and countries.
We study this problem on a variant of the Functional Map of the World dataset (Christie et al., 2018), where the input is an RGB satellite image, the label is one of 62 building or land use categories, and the domain represents the year the image was taken and its geographical region (Africa, the Americas, Oceania, Asia, or Europe) (Figure 9). The different regions have different numbers of examples, e.g., there are far fewer images from Africa than the Americas. The training set comprises data from before 2013, while the test set comprises data from 2016 and after; years 2013 to 2015 are reserved for the validation set. We evaluate models by their test accuracy on the worst geographical region, which combines both a domain generalization problem over time and a subpopulation shift problem over regions. As we provide both time and region annotations, models can leverage the structure across both space and time to improve robustness. Appendix E.7 provides additional details and context.
3.2 PovertyMap-wilds: Poverty mapping across different countries
Global-scale poverty estimation is a specific remote sensing application which is essential for targeted humanitarian efforts in poor regions (Abelson et al., 2014; Espey et al., 2015). However, ground truth measurements of poverty are lacking for much of the developing world, as field surveys for collecting the ground truth are expensive (Blumenstock et al., 2015). This motivates the approach of training ML models on countries with ground truth labels and then deploying them on different countries where we have satellite data but no labels (Xie et al., 2016; Jean et al., 2016; Yeh et al., 2020).
We study this shift through a variant of the poverty mapping dataset collected by Yeh et al. (2020), where the input is a multispectral satellite image, the output is a real-valued asset wealth index from surveys, and the domain represents the country the image was taken in and whether the image is of an urban or rural area (Figure 10). The training and test set comprise data from disjoint sets of countries, and we evaluate models by the correlation of their predictions with the ground truth. Specifically, we take the lower of the correlations over the urban and rural subpopulations, as prior work has shown that accurately predicting poverty within these subpopulations is especially challenging. As poverty measures are highly correlated across space (Jean et al., 2018; Rolf et al., 2020), methods can utilize the provided location coordinates, and the country and urban/rural annotations, to improve robustness. Appendix E.8 provides additional details and context.
3.3 Amazon-wilds: Sentiment classification across different users
In many consumer-facing ML applications, models are trained on data collected on one set of users and then deployed across a wide range of potentially new users. These models can perform well on average but poorly on some users (Tatman, 2017; Caldas et al., 2018; Li et al., 2019b; Koenecke et al., 2020). These large performance disparities across users are practical concerns in consumer-facing applications, and they can also indicate that models are exploiting biases or spurious correlations in the data (Badgeley et al., 2019; Geva et al., 2019).
We study this issue on a variant of the Amazon review dataset (Ni et al., 2019), where the input is the review text, the label is the corresponding 1-to-5 star rating, and the domain identifies the user who wrote the review (Figure 11). The training and test sets comprise reviews from disjoint sets of users; for leverage, the training set has reviews from 5,008 different users. As our goal is to train models with consistently high performance across users, we evaluate models by the 10th percentile of per-user accuracies. Appendix E.9 provides additional details and context. We discuss other distribution shifts on this dataset (e.g., by category) in Appendix F.4.
3.4 Py150-wilds: Code completion across different codebases
Code completion models—autocomplete tools used by programmers to suggest subsequent source code tokens, such as the names of API calls—are commonly used to reduce the effort of software development (Robbes and Lanza, 2008; Bruch et al., 2009; Nguyen and Nguyen, 2015; Proksch et al., 2015; Franks et al., 2015). These models are typically trained on data collected from existing codebases but then deployed more generally across other codebases, which may have different distributions of API usages (Nita and Notkin, 2010; Proksch et al., 2016; Allamanis and Brockschmidt, 2017). This shift across codebases can cause substantial performance drops in code completion models. Moreover, prior studies of real-world usage of code completion models have noted that they can generalize poorly on some important subpopulations of tokens such as method names (Hellendoorn et al., 2019).
We study a variant of the Py150 Dataset (Raychev et al., 2016; Lu et al., 2021), where the goal is to predict the next token (e.g., "environ", "communicate" in Figure 12) given the context of previous tokens. The input is a sequence of source code tokens, the label is the next token, and the domain specifies the repository that the source code belongs to. The training and test sets comprise code from disjoint GitHub repositories. As leverage, we include over 5,300 repositories in the training set, capturing a wide range of source code variation. We evaluate models by their accuracy on the subpopulation of class and method tokens. Additional dataset and model details are provided in Appendix E.10.
Performance drops from distribution shifts
For a dataset to be appropriate for Wilds, the distribution shift reflected in its official train/test split should cause significant performance drops in standard models. How to measure the performance drop due to a distribution shift is a crucial but subtle question. In this section, we discuss our approach and the results on each of the Wilds datasets. To construct Wilds, we selected datasets with large performance drops; in Section 8, we discuss other datasets with real-world shifts that did not show large performance drops and were therefore not included in the benchmark.
Our general approach is to measure the difference between the out-of-distribution (OOD) and in-distribution (ID) performance of standard models trained via empirical risk minimization (ERM). Concretely, we first measure the OOD performance using the official train/test splits described in Section 4. We then construct an appropriate in-distribution (ID) setting to measure ID performance, typically by modifying the official train/test splits. However, practical constraints often prevent us from constructing an ID setting in exactly the way we want, which makes the choice of appropriate ID setting for each dataset a case-by-case issue.
Choosing an appropriate in-distribution (ID) setting is the crux of measuring how much a distribution shift affects performance. But what distribution should “in-distribution” be taken with respect to? Consider a distribution shift from a training distribution to a test distribution . It is common to measure ID performance by taking a model trained on and evaluating it on additional held-out data from .For example, in domain generalization, we might train a model on the training domains and then report its ID performance on held-out examples from the same domains; and in subpopulation shift, we might report average performance on as the ID performance. This is useful for checking if the model can generalize well on both the training and the shifted test distributions. However, it fails to isolate the effect of the distribution shift since it does not control for the data distribution on which the model is evaluated: the ID setting evaluates on data from , whereas the OOD setting evaluates on data from . As a result, the performance gap might also be due to other factors such as differences in the difficulty of fitting a model to versus .
For illustration, consider the task of wheat head detection on GlobalWheat-wilds. The shift from to , which contain images of wheat fields in Europe and North America respectively, involves changes in factors such as wheat genotype, illumination, and growing conditions. These changes mean that the task can be more challenging in some regions than others: for example, wheat is grown in higher densities in certain regions than others, and it is harder to detect wheat heads reliably when they are more densely packed together. If, for example, the task is harder in the regions in , then we might see especially low performance on compared to . However, this performance gap would overestimate the actual gap caused by the distribution shift, in the sense that performance on would still be lower even if we could train a model purely on data from .
To isolate the gap caused by the distribution shift, it is therefore important to keep the evaluation data distribution fixed between the ID and OOD settings by evaluating on in the ID setting. For example, we could measure ID performance by training on and evaluating on and then compare this with the standard OOD setting of training on and evaluating on . However, there is a practical drawback: we generally have much more data from rather than , and training and evaluating on would require us to have a substantial number of labeled examples from each test domain. In contrast, the standard ID setting of training and evaluating on is typically much more feasible, and it is also more convenient as we can reuse the same model trained on for both ID and OOD evaluations.
In Wilds, we take the approach of measuring ID performance on whenever practically feasible, and we lean on standard ID evaluations on otherwise. In either case, we generally provide held-out data from in order to track model performance on .
2 Types of in-distribution settings
To measure the performance drop on each Wilds dataset, we picked the most appropriate ID setting(s) that were feasible. We now describe five specific ways of constructing ID settings and their pros and cons. The first two ID settings (test-to-test and mixed-to-test) control for the evaluation distribution and thus isolate the performance drops due to distribution shifts, as discussed in Section 5.1. However, these procedures require substantial training data from test domains, so in cases where such data is not practically available, we consider the other ID settings (train-to-train, average, and random split). Appendix E describes dataset-specific rationales for the selected ID settings and additional details for each dataset.
Below, we denote the training and OOD test sets of the official Wilds splits as and , sampled from distributions and , respectively.
To control for the evaluation distribution, we can hold the test set fixed and train on a separate but identically-distributed training set drawn from . The ID performance reported in this setting is directly comparable to OOD performance, which is also evaluated on . The main drawback is that for a fair comparison to the OOD setting, where we train a model on , we would require to match the size of . This is not feasible in our datasets, as typically comprises the bulk of the available data. We therefore do not use the test-to-test comparison for any of the Wilds datasets and instead consider the more practical alternative below, which still controls for the evaluation data distribution.
In the mixed-to-test setting, we train a model on a mixture of data from and and then evaluate it only on . This is a more practical version of the test-to-test setting, as it retains the advantage of controlling for the evaluation distribution, while mitigating the need for large amounts of labeled data from to use for training.In practice, we typically split up and use some of it for training by replacing examples in (so that the size of the training set is similar to the OOD setting). This still requires to be large enough to support using a sufficient number of examples for training while also having enough examples left over for accurate evaluation. We use the mixed-to-test comparison for the Wilds datasets wherever feasible, except when we expect the train-to-train comparison to give similar results as described in the below discussion on train-to-train setting (e.g., for iWildCam2020-wilds and Py150-wilds).
One downside is that compared to the test-to-test setting, the mixed-to-test setting might underestimate ID performance, since it trains a model that simultaneously fits both and , instead of just focusing on . However, this is useful as a sanity check that we can learn a model that can simultaneously fit both and ; if such a model were not possible to learn, then it suggests that the distribution shift in the dataset is intractable for the model family.
In the train-to-train setting, we train a model on and evaluate on a separate but identically-distributed test set drawn from . As discussed in Section 5.1, this is practical—it does not require large amounts of data from , and we can reuse the model for OOD evaluation—but has the drawback of not controlling for the evaluation distribution.
This drawback is less of an issue when we expect and to be of equal difficulty in the sense of Section 5.1. This may be the case when the dataset has a relatively large number of training and test domains that are drawn from the same distribution, and they are thus roughly interchangeable. For instance, in iWildCam2020-wilds and Py150-wilds, there are many available domains (camera traps and GitHub repositories, respectively) randomly split across and , so we use the train-to-train comparison for them. For most of the other datasets, we also include train-to-train comparisons to track model performance on (i.e., the official splits typically also include a held-out ; we report results on these in Appendix E), but we complement them whenever feasible with other ID settings that better isolate the effect of the distribution shift.
In subpopulation shift datasets, we measure the OOD performance of a model by reporting the performance on the worst-case subpopulation, and we can measure ID performance by simply reporting the average performance. This average comparison corresponds to a special case of the train-to-train setting,In subpopulation shifts, the training distribution reflects the empirical make-up over the pre-defined subpopulations, whereas the test distribution of interest corresponds to the worst-case subpopulation. so they share the same pros and cons. In particular, the average comparison is much more practical than running a test-to-test comparison on each subpopulation, as it can be especially difficult to obtain sufficient training examples from minority subpopulations. In Table 1, we use this average comparison for the CivilComments-wilds and Amazon-wilds datasets, which both consider a large number of subpopulations that are individually quite small.
Another standard approach to measuring ID performance is to shuffle all of the data in into i.i.d. training, validation, and test splits, while keeping the size of the training set constant. We use this in OGB-MolPCBA to be consistent with prior work from the Open Graph Benchmark (Hu et al., 2020b). As with the train-to-train comparison, the random split comparison is simple to implement and does not require large amounts of data from , but it does not control for the evaluation distribution.
3 Model selection
We used standard model architectures for each dataset: ResNet and DenseNet for images (He et al., 2016; Huang et al., 2017), DistilBERT for text (Sanh et al., 2019), a Graph Isomorphism Network (GIN) for graphs (Xu et al., 2018), and Faster-RCNN (Ren et al., 2015) for detection. As our goal is high OOD performance, we use a separate OOD validation set for early stopping and hyperparameter selection.This means that while the ERM models do not make use of any additional metadata (e.g., domain annotations) during training, this metadata is still implicitly (but very mildly) used for model selection. Relative to the training set, this OOD validation set reflects a distribution shift similar to, but distinct from, the test set. For example, in iWildCam2020-wilds, the training, validation, and test sets each comprise photos from distinct sets of camera traps. We detail experimental protocol in Appendix D and models and hyperparameters for each dataset in Appendix E.
For the ID comparisons, we use the same hyperparameters optimized on the OOD validation set, so our ID results are slightly lower than if we had optimized hyperparameters for ID performance (Appendix D). In other words, the ID-OOD gaps in Table 1 are slightly underestimated.
4 Results
Table 1 shows that for each dataset, OOD performance is consistently and substantially lower than the corresponding ID performance. Moreover, on the datasets that allow for mixed-to-test ID comparisons, we show that models trained on a mix of the ID and OOD distributions can simultaneously achieve high ID and OOD performance, indicating that lower OOD performance is not due to the OOD test sets being intrinsically more difficult than the ID test sets. Overall, these results demonstrate that the real-world distribution shifts reflected in the Wilds datasets meaningfully degrade standard model performance. Additional results for datasets that admit multiple ID comparisons are described for each dataset in Appendix E.
Baseline algorithms for distribution shifts
Many algorithms have been proposed for training models that are more robust to particular distribution shifts than standard models trained by empirical risk minimization (ERM), which trains models to minimize the average training loss. Unlike ERM, these algorithms tend to utilize domain annotations during training, with the goal of learning a model that can generalize across domains. In this section, we evaluate several representative algorithms from prior work and show that the out-of-distribution performance drops shown in Section 5 still remain.
Methods for domain generalization typically involve adding a penalty to the ERM objective that encourages some form of invariance across domains. We include two such methods as representatives:
CORAL (Sun and Saenko, 2016), which penalizes differences in the means and covariances of the feature distributions (i.e., the distribution of last layer activations in a neural network) for each domain. Conceptually, CORAL is similar to other methods that encourage feature representations to have the same distribution across domains (Tzeng et al., 2014; Long et al., 2015; Ganin et al., 2016; Li et al., 2018c, b).
IRM (Arjovsky et al., 2019), which penalizes feature distributions that have different optimal linear classifiers for each domain. This builds on earlier work on invariant predictors (Peters et al., 2016).
Other techniques for domain generalization include conditional variance regularization (Heinze-Deml and Meinshausen, 2017); self-supervision (Carlucci et al., 2019); and meta-learning-based approaches (Li et al., 2018a; Balaji et al., 2018; Dou et al., 2019).
2 Subpopulation shift baselines
In subpopulation shift settings, our aim is to train models that perform well on all relevant subpopulations. We test the following approach:
Group DRO (Hu et al., 2018; Sagawa et al., 2020a), which uses distributionally robust optimization to explicitly minimize the loss on the worst-case domain during training. Group DRO builds on the maximin approach developed in Meinshausen and Bühlmann (2015).
Other methods for subpopulation shifts include reweighting methods based on class/domain frequencies (Shimodaira, 2000; Cui et al., 2019); label-distribution-aware margin losses (Cao et al., 2019); adaptive Lipschitz regularization (Cao et al., 2020); slice-based learning (Chen et al., 2019b; Ré et al., 2019); style transfer across domains (Goel et al., 2020); or other DRO algorithms that do not make use of explicit domain information and rely on, for example, unsupervised clustering (Oren et al., 2019; Sohoni et al., 2020) or upweighting high-loss points (Nam et al., 2020; Liu et al., 2021a).
Subpopulation shifts are also connected to the well-studied notions of tail performance and risk-averse optimization (Chapter 6 in Shapiro et al. (2014)). For example, optimizing for the worst case over all subpopulations of a certain size, regardless of domain, can guarantee a certain level of performance over the smaller set of subpopulations defined by domains (Duchi et al., 2020; Duchi and Namkoong, 2021).
3 Setup
We trained CORAL, IRM, and Group DRO models on each dataset. While Group DRO was originally developed for subpopulation shifts, for completeness, we also experiment with using it for domain generalization. In that setting, Group DRO models aim to achieve similar performance across domains: e.g., in Camelyon17-wilds, where the domains are hospitals, Group DRO optimizes for the training hospital with the highest loss. Similarly, we also test CORAL and IRM on subpopulation shifts, where they encourage models to learn invariant representations across subpopulations. As in Section 5, we used the same OOD validation set for early stopping and to tune the penalty weights for the CORAL and IRM algorithms. More experimental details are in Appendix D, and dataset-specific hyperparameters and domain choices are discussed in Appendix E.
4 Results
Table 2 shows that models trained with CORAL, IRM, and Group DRO generally fail to improve over models trained with ERM. The exception is the CivilComments-wilds subpopulation shift dataset, where the worst-performing subpopulation is a minority domain. By upweighting the minority domain, Group DRO obtains an OOD accuracy of 70.0% (on the worst-performing subpopulation) compared to 56.0% for ERM, though this is still substantially below the ERM model’s ID accuracy of 92.2% (on average over the entire test set). CORAL and IRM also perform well on CivilComments-wilds, though the gains there stem from the fact that our implementation heuristically upsamples the minority domain (see Appendix E.6). All other datasets involve domain generalization; the failure of the baseline algorithms here is consistent with other recent findings on standard domain generalization datasets (Gulrajani and Lopez-Paz, 2020).
These results indicate that training models to be robust to distribution shifts in the wild remains a significant open challenge. However, we are optimistic about future progress for two reasons. First, current methods were mostly designed for other problem settings besides domain generalization, e.g., CORAL for unsupervised domain adaptation and Group DRO for subpopulation shifts. Second, compared to existing distribution shift datasets, the Wilds datasets generally contain diverse training data from many more domains as well as metadata on these domains, which future algorithms might be able to leverage.
Empirical trends
We end our discussion of experimental results by briefly reporting on several trends that we observed across multiple datasets.
Prior work has shown that there is often insufficient information at training time to distinguish models that would generalize well under distribution shift; many models that perform similarly in-distribution (ID) can vary substantially out-of-distribution (OOD) (McCoy et al., 2019a; Zhou et al., 2020; D’Amour et al., 2020a). In Wilds, we attempt to alleviate this issue by providing multiple training domains in each dataset as well as an OOD validation set for model selection. Perhaps as a result, we do not observe significantly higher variance in OOD performance than ID performance in Table 1, with the exception of Amazon-wilds and CivilComments-wilds, where the OOD performance is measured on a smaller subpopulation and is therefore naturally more variable. Excluding those datasets, the average standard deviation from Table 1 is 2.6% for OOD performance and 2.0% for ID performance, which is comparable. These results raise the question of when underspecification, as reported in prior work, could be more of an issue.
2 Model selection with in-distribution versus out-of-distribution validation sets
All of the baseline results reported in this paper use an OOD validation set for model selection, as discussed in Section 5.3. To facilitate research into comparisons of ID versus OOD performance, most Wilds datasets also provide an ID validation and/or test set. For example, in iWildCam2020-wilds, the ID validation set comprises photos from the same set of camera traps used for the training set. These ID sets are not used for model selection nor official evaluation.
Gulrajani and Lopez-Paz (2020) showed that on the DomainBed domain generalization datasets, selecting models with an ID validation set leads to higher OOD performance than using an OOD validation set. This contrasts with our approach of using OOD validation sets, which we find to generally provide a good estimate of OOD test performance. Specifically, in Appendix D.1, we show that for our baseline models, model selection using an OOD validation set results in comparable or higher OOD performance than model selection using an ID validation set. This difference could stem from many factors: for example, Wilds datasets tend to have many more domains, whereas DomainBed datasets tend to have fewer domains that can be quite different from each other (e.g., cartoons vs. photos); and there are some differences in the exact procedures for comparing performance using ID versus OOD validation sets. Further study of the effects of these different model selection procedures and choices of validation sets would be a useful direction for future work.
3 The compounding effects of multiple distribution shifts
Several Wilds datasets consider hybrid settings, where the goal is to simultaneously generalize to unseen domains as well as to certain subpopulations. We observe that combining these types of shifts can exacerbate performance drops. For example, in PovertyMap-wilds and FMoW-wilds, the shift to unseen domains exacerbates the gap in subpopulation performance (and vice versa). Notably, in FMoW-wilds, the difference in subpopulation performance (across regions) is not even manifested until also considering another shift (across time). While we do not always observe the compounding effect of distribution shifts—e.g., in Amazon-wilds, subpopulation performance is similar whether we consider shifts to unseen users or not—these observations underscore the importance of evaluating models on the combination of distribution shifts that would occur in practice, instead of considering each shift in isolation.
Distribution shifts in other application areas
Beyond the datasets currently included in Wilds, there are many other applications where it is critical for models to be robust to distribution shifts. In this section, we discuss some of these applications and the challenges of finding appropriate benchmark datasets in those areas. We also highlight examples of datasets with distribution shifts that we considered but did not include in Wilds, because their distribution shifts did not lead to a significant performance drop. Constructing realistic benchmarks that reflect distribution shifts in these application areas is an important avenue of future work, and we would highly welcome community contributions of benchmark datasets in these areas.
Distribution shifts which degrade model performance on minority subpopulations are frequently discussed in the algorithmic fairness literature. Geographic inequities are one concern (Shankar et al., 2017; Atwood et al., 2020): e.g., publicly available image datasets overrepresent images from the US and Europe, degrading performance in the developing world (Shankar et al., 2017) and prompting the creation of more geographically diverse datasets (Atwood et al., 2020). Racial disparities are another concern: e.g., commercial gender classifiers are more likely to misclassify the gender of darker-skinned women, likely in part because training datasets overrepresent lighter-skinned subjects (Buolamwini and Gebru, 2018), and pedestrian detection systems fare worse on darker-skinned pedestrians (Wilson et al., 2019). As in Section 4.2.1, NLP models can also show racial bias.
Unfortunately, publicly available algorithmic fairness benchmarks (Mehrabi et al., 2019)—e.g., the COMPAS recidivism dataset (Larson et al., 2016)—suffer from several limitations. First, the datasets are often quite small by the standards of modern ML: the COMPAS dataset has only a few thousand rows (Larson et al., 2016). Second, they tend to have relatively few features, and disparities in subgroup performance are not always large (Larrazabal et al., 2020), limiting the benefit of more sophisticated approaches: on COMPAS, logistic regression performs comparably to a black-box commercial algorithm (Jung et al., 2020; Dressel and Farid, 2018). Third, the datasets sometimes represent “toy” problems: e.g., the UCI Adult Income dataset (Asuncion and Newman, 2007) is widely used as a fairness benchmark, but its task—classifying whether a person will have an income above $50,000—does not represent a real-world application. Finally, because many of the domains in which algorithmic fairness is of most concern—e.g., criminal justice and healthcare—are high-stakes and disparities are politically sensitive, it can be difficult to make datasets publicly available.
Creating algorithmic fairness benchmarks which do not suffer from these limitations represents a promising direction for future work. In particular, such datasets would ideally have: 1) information about a sensitive attribute like race or gender; 2) a prediction task which is of immediate real-world interest; 3) enough samples, a rich enough feature set, and large enough disparities in group performance that more sophisticated machine learning approaches would plausibly produce improvement over naive approaches.
Predictive policing is a prominent example of a real-world application where fairness considerations are paramount: algorithms are increasingly being used in contexts such as predicting crime hotspots (Lum and Isaac, 2016) or a defendant’s risk of reoffending (Larson et al., 2016; Corbett-Davies et al., 2016, 2017; Lum and Shah, 2019). There are numerous concerns about these applications (Larson et al., 2016; Corbett-Davies et al., 2016, 2017; Lum and Shah, 2019), one of which is that these ML models might not generalize beyond the distributions that they were trained on (Corbett-Davies and Goel, 2018; Slack et al., 2019). These distribution shifts include shifts over locations—e.g., a criminal risk assessment trained on several hundred defendants in Ohio was eventually used throughout the United States (Latessa et al., 2010)—and shifts over time, as sentencing and other criminal justice policies evolve (Corbett-Davies and Goel, 2018). There are, of course, also subpopulation shift concerns around whether models are biased against particular demographic groups.
We investigated these shifts using a dataset of pedestrian stops made by the New York City Police Department under its “stop-and-frisk” policy, where the task is to predict whether a pedestrian who was stopped on suspicion of weapon possession would in fact possess a weapon (Goel et al., 2016). This policy had a pronounced racial bias: Black people stopped by the police on suspicion of possessing a weapon were 5 less likely to actually possess one than their White counterparts (Goel et al., 2016). We emphasize that we oppose stop-and-frisk (and any “improved” ML-powered stop-and-frisk) since there is overwhelming evidence that the policy was racially discriminatory (Gelman et al., 2007; Goel et al., 2016; Pierson et al., 2018) and such massive inequities require more than algorithmic fixes. Rather, we use the dataset as a realistic example of the phenomena that arise in real policing contexts, including 1) substantial heterogeneity across locations and racial groups and 2) distributions that arise in part because of biased policing practices.
Overall, we found large performance disparities across race groups and locations. Interestingly, however, we also found that these disparities cannot be attributed to the distribution shift, as the disparities were not reduced when we trained models specifically on the race groups or locations that suffer the worst performance. Indeed, the groups that see the worst performance—Black and Hispanic pedestrians—comprise large majorities of the dataset, making up more than 90% of the stops. This contrasts with the typical setting in algorithmic fairness where models perform worse on minority groups in the training data. Our results suggest the disparities are due to the dataset being noisier for some race and location groups, potentially as a result of the biased policing practices underlying the dataset. We provide further details in Appendix F.1.
2 Medicine and healthcare
Substantial evidence indicates the potential for distribution shifts in medical settings (Finlayson et al., 2021). One concern is demographic subpopulation shifts (e.g., across race, gender, or socioeconomic status), since historically-disadvantaged populations are underrepresented in many medical datasets (Chen et al., 2020). Another concern is heterogeneity across hospitals; this might include differences in imaging, as in Section 4.1.2, and other operational protocols such as lab tests (D’Amour et al., 2020a; Subbaswamy et al., 2020). Finally, changes over time can also produce distribution shifts: for example, Nestor et al. (2019) showed that switching between two electronic health record (EHR) systems produced a drop in performance, and the COVID-19 epidemic has affected the distribution of chest radiographs (Wong et al., 2020).
Creating medical distribution shift benchmarks thus represents a promising direction for future work, if several challenges can be overcome. First, while there are large demographic disparities in healthcare outcomes (e.g., by race or socioeconomic status), many of them are not due to distribution shifts, but to disparities in non-algorithmic factors (e.g., access to care or prevalence of comorbidities (Chen et al., 2020)) or to algorithmic problems unrelated to distribution shift (e.g., choice of a biased outcome variable (Obermeyer et al., 2019)). Indeed, several previous investigations have found relatively small disparities in algorithmic performance (as opposed to healthcare outcomes) across demographic groups (Chen et al., 2019a; Larrazabal et al., 2020); Seyyed-Kalantari et al. (2020) finds larger disparities in true positive rates across demographic groups, but this might reflect the different underlying label distributions between groups.
Second, many distribution shifts in medicine arise from concept drifts, in which the relationship between the input and the label changes, for example due to changes in clinical procedures and the definition of the label (Widmer and Kubat, 1996; Beyene et al., 2015; Futoma et al., 2020). It can be difficult to ensure that a potential benchmark has sufficient leverage for models to learn how to handle, e.g., an abrupt change in the way a particular clinical procedure is carried out.
A last challenge is data availability, as stringent medical privacy laws often preclude data sharing (Price and Cohen, 2019). For example, EHR datasets are fundamental to medical decision-making, but there are few widely adopted EHR benchmarks—with the MIMIC database being a prominent exception (Johnson et al., 2016)—and relatively little progress in predictive performance has been made on them (Bellamy et al., 2020).
3 Genomics
Advances in high-throughput genomic and molecular profiling platforms have enabled systematic mapping of biochemical activity of genomes across diverse cellular contexts, populations, and species (Consortium et al., 2012; Ho et al., 2014; Kundaje et al., 2015; Regev et al., 2017; Consortium, 2019; Moore et al., 2020; Consortium et al., 2020). These datasets have powered ML models that have been fairly successful at deciphering functional DNA sequence patterns and predicting the consequences of genetic perturbations in cell types in which the models are trained (Libbrecht and Noble, 2015; Zhou and Troyanskaya, 2015; Kelley et al., 2016; Ching et al., 2018; Eraslan et al., 2019; Jaganathan et al., 2019; Avsec et al., 2021b). However, distribution shifts pose a significant obstacle to generalizing these predictions to new cell types.
A concrete example is the prediction of genome-wide profiles of regulatory protein-DNA binding interactions across cell types and tissues (Srivastava and Mahony, 2020). These regulatory maps are critical for understanding the fundamental mechanisms of dynamic gene regulation across healthy and diseased cell states, and predictive models are an essential complement to experimental approaches for comprehensively profiling these maps.
Regulatory proteins bind regulatory DNA elements in a sequence-specific manner to orchestrate gene expression programs. These proteins often form different complexes with each other in different cell types. These cell-type-specific protein complexes can recognize distinct combinatorial sequence syntax and thereby bind to different genomic locations in different cell types, even if all of these cell types share the same genomic sequence. Hence, ML models that aim to predict protein-DNA binding landscapes across cell types typically integrate DNA sequence and additional context-specific input data modalities, which provide auxiliary information about the regulatory state of DNA in each cell type (Srivastava and Mahony, 2020). The training cell-type specific sequence determinants of binding induce a distribution shift across cell types, which can in turn degrade model performance on new cell types (Balsubramani et al., 2017; Li et al., 2019a; Li and Guan, 2019; Keilwagen et al., 2019; Quang and Xie, 2019).
We studied the above problem in the context of the ENCODE-DREAM in-vivo Transcription Factor Binding Site Prediction Challenge (Balsubramani et al., 2020), which is an open community challenge introduced to systematically benchmark ML models for predicting genome-wide DNA binding maps of many regulatory proteins across cell types.
For each regulatory protein, regions of the genome are associated with binary labels (bound/unbound). The task is to predict these binary binding labels as a function of underlying DNA sequence and chromatin accessibility signal (an experimental measure of cell type-specific regulatory state) in test cell types that are not represented in the training set.
A systematic evaluation of the top-performing models in this challenge highlighted a significant gap in prediction performance across cell types, relative to cross-validation performance within training cell types (Li et al., 2019a; Li and Guan, 2019; Keilwagen et al., 2019; Quang and Xie, 2019). This performance gap was attributed to distribution shifts across cell types, due to regulatory proteins forming cell-type-specific complexes that can recognize different combinatorial sequence syntax. Hence, the same DNA sequence can be associated with different binding labels for a protein across contexts.
We investigated these distribution shifts in more detail for a restricted subset of the challenge’s prediction tasks for two regulatory proteins, using a total of 14 genome-wide binding maps across different cell types. While we generally found a performance gap between in- and out-of-distribution settings, we did not include this dataset in the official benchmark for several reasons. For example, we were unable to learn a model that could generalize across all the cell types simultaneously, even in an in-distribution setting, which suggested that the model family and/or feature set might not be rich enough to fit the variation across different cell types. Another major complication was the significant variation in intrinsic difficulty across different splits, as measured by the performance of models we train in-distribution. Further work will be required to construct a rigorous benchmark for evaluating distribution shifts in the context of predicting regulatory binding maps. We discuss details in Appendix F.2.
4 Natural language and speech processing
Subpopulation shifts are an issue in automated speech recognition (ASR) systems, which have been shown to have higher error rates for Black speakers than for White speakers (Koenecke et al., 2020) and for speakers of some dialects (Tatman, 2017). These disparities were demonstrated using commercial ASR systems, and therefore do not have any accompanying training datasets that are publicly available. There are many public speech datasets with speaker metadata that could potentially be used to construct a benchmark, e.g., LibriSpeech (Panayotov et al., 2015), the Speech Accent Archive (Weinberger, 2015), VoxCeleb2 (Chung et al., 2018), the Spoken Wikipedia Corpus (Baumann et al., 2019), and Common Voice (Ardila et al., 2020). However, these datasets have their own challenges: some do not have a sufficiently diverse sample of speaker backgrounds and accents, and others focus on read speech (e.g., audiobooks) instead of more natural speech.
In natural language processing (NLP), a current focus is on challenge datasets that are crafted to test particular aspects of models, e.g., HANS (McCoy et al., 2019b), PAWS (Zhang et al., 2019), and CheckList (Ribeiro et al., 2020). These challenge datasets are drawn from test distributions that are often (deliberately) quite different from the data distributions that models are typically trained on. Counterfactually-augmented datasets (Kaushik et al., 2019) are a related type of challenge dataset where the training data is modified to make spurious correlates independent of the target, which can result in more robust models. Others have studied train/test sets that are drawn from different sources, e.g., Wikipedia, Reddit, news articles, travel reviews, and so on (Oren et al., 2019; Miller et al., 2020; Kamath et al., 2020).
Several synthetic datasets have also been designed to test compositional generalization, such as CLEVR (Johnson et al., 2017), SCAN (Lake and Baroni, 2018), and COGS (Kim and Linzen, 2020). The test sets in these datasets are chosen such that models need to generalize to novel combinations of parts of training examples, e.g., familiar primitives and grammatical roles (Kim and Linzen, 2020). CLEVR is a visual question-answering (VQA) dataset; other examples of VQA datasets that are formulated as challenge datasets are the VQA-CP v1 and v2 datasets (Agrawal et al., 2018), which create subpopulation shifts by intentionally altering the distribution of answers per question type between the train and test splits.
These NLP examples involve English-language models; other languages typically have fewer and smaller datasets available for training and benchmarking models. Multi-lingual models and benchmarks (Conneau et al., 2018; Conneau and Lample, 2019; Hu et al., 2020a; Clark et al., 2020) are another source of subpopulation shifts with corresponding disparities in performance: training sets might contain fewer examples in low-resource languages (Nekoto et al., 2020), but we would still hope for high model performance on these minority groups.
In addition to user shifts on the Amazon Reviews dataset (Ni et al., 2019), we also looked at category and time shifts on the same dataset, as well as user and time shifts on the Yelp Open Datasethttps://www.yelp.com/dataset. However, for many of those shifts, we only found modest performance drops. We provide additional details on Amazon in Appendix F.4 and on Yelp in Appendix F.5.
5 Education
ML models can help in educational settings in a variety of ways: e.g., assisting in grading (Piech et al., 2013; Shermis, 2014; Kulkarni et al., 2014; Taghipour and Ng, 2016), estimating student knowledge (Desmarais and Baker, 2012; Wu et al., 2020), identifying students who need help (Ahadi et al., 2015), or automatically generating explanations (Williams et al., 2016; Wu et al., 2019a). However, there are substantial distribution shifts in these settings as well. For example, automatic essay scoring has been found to be affected by rater bias (Amorim et al., 2018) and spurious correlations like essay length (Perelman, 2014), leading to problems with subpopulation shift. Ideally, these systems would also generalize across different contexts, e.g., a model for scoring grammar should work well across multiple different essay prompts. Recent attempts at predicting grades algorithmically (BBC, 2020; Broussard, 2020) have also been found to be biased against certain subpopulations.
Unfortunately, there is a general lack of standardized education datasets, in part due to student privacy concerns and the proprietary nature of large-scale standardized tests. Datasets from massive open online courses are a potential source of large-scale data (Kulkarni et al., 2015). In general, dataset construction for ML in education is an active area—e.g., the NeurIPS 2020 workshop on Machine Learning for Educationhttps://www.ml4ed.org/ has a segment devoted to finding “ImageNets for education”—and we hope to be able to include one in the future.
6 Robotics
Robot learning has emerged as a strong paradigm for automatically acquiring complex and skilled behaviors such as locomotion (Yang et al., 2019; Peng et al., 2020), navigation (Mirowski et al., 2017; Kahn et al., 2020), and manipulation (Gu et al., 2017; et al, 2019). However, the advent of learning-based techniques for robotics has not convincingly addressed, and has perhaps even exasperated, problems stemming from distribution shift. These problems have manifested in many ways, including shifts induced by weather and lighting changes (Wulfmeier et al., 2018), location changes (Gupta et al., 2018), and the simulation-to-real-world gap (Sadeghi and Levine, 2017; Tobin et al., 2017). Dealing with these challenging scenarios is critical to deploying robots in the real world, especially in high-stakes decision-making scenarios.
For example, to safely deploy autonomous driving vehicles, it is critical that these systems work reliably and robustly across the huge variety of conditions that exist in the real world, such as locations, lighting and weather conditions, and sensor intrinsics. This is a challenging requirement, as many of these conditions may be underrepresented, or not represented at all, by the available training data. Indeed, prior work has shown that naively trained models can suffer at segmenting nighttime driving scenes (Dai and Van Gool, 2018), detecting relevant objects in new or challenging locations and settings (Yu et al., 2020; Sun et al., 2020a), and, as discussed earlier, detecting pedestrians with darker skin tones (Wilson et al., 2019).
Creating a benchmark for distribution shifts in robotics applications, such as autonomous driving, represents a promising direction for future work. Here, we briefly summarize our initial findings on distribution shifts in the BDD100K driving dataset (Yu et al., 2020), which is publicly available and widely used, including in some of the works listed above.
We investigated the task of multi-label binary classification of the presence of each object category in each image. In general, we found no substantial performance drops across a wide range of different test scenarios, including user shifts, weather and time shifts, and location shifts. We provide additional details in Section F.3.
Our findings contrast with previous findings that other tasks, such as object detection and segmentation, can suffer under the same types of shifts on the same dataset (Yu et al., 2020; Dai and Van Gool, 2018). Currently, Wilds consists of datasets involving classification and regression tasks. However, most tasks of interest in autonomous driving, and robotics in general, are difficult to formulate as classification or regression. For example, autonomous driving applications may require models for object detection or lane and scene segmentation. These tasks are often more challenging than classification tasks, and we speculate that they may suffer more severely from distribution shift.
7 Feedback loops
Finally, we have restricted our attention to settings where the data distribution is independent of the model. When the data distribution does depend on the model, distribution shifts can arise from feedback loops between the data and the model. Examples include recommendation systems and other consumer products (Bottou et al., 2013; Hashimoto et al., 2018); dialogue agents (Li et al., 2017b); molecular compound optimization (Cuccarese et al., 2020; Reker, 2020); decision systems (Liu et al., 2018; D’Amour et al., 2020b); and adversarial settings like fraud or malware detection (Rigaki and Garcia, 2018). While these adaptive settings are outside the scope of our benchmark, dealing with these types of distribution shifts is an important area of ongoing work.
Guidelines for method developers
We now discuss some community guidelines for method development using Wilds. More specific submission guidelines for our leaderboard can be found at https://wilds.stanford.edu.
Wilds is primarily designed as a benchmark for developing and evaluating algorithms for training models that are robust to distribution shifts. To facilitate systematic comparisons of these algorithms, we encourage algorithm developers to use the standardized datasets (i.e., with no external data) and default model architectures provided in Wilds, as doing so will help to isolate the contributions of the algorithm versus the training dataset or model architecture. Our primary leaderboard will focus on submissions that follow these guidelines.
Moreover, we encourage developers to test their algorithms on all applicable Wilds datasets, so as to assess how well they do across different types of data and distribution shifts. We emphasize that it is still an open question if a single general-purpose training algorithm can produce models that do well on all of the datasets without accounting for the particular structure of the distribution shift in each dataset. As such, it would still be a substantial advance if an algorithm significantly improves performance on one type of shift but not others; we aim for Wilds to facilitate research into both general-purpose algorithms as well as ones that are more specifically tailored to a particular application and type of distribution shift.
2 Methods beyond training algorithms
Beyond new training algorithms, there are many other promising directions for improving distributional robustness, including new model architectures and pre-training on additional external data beyond what is used in our default models. We encourage developers to test these approaches on Wilds as well, and we will track all such submissions on a separate leaderboard from the training algorithm leaderboard.
3 Avoiding overfitting to the test distribution
While each Wilds dataset aims to benchmark robustness to a type of distribution shift (e.g., shifts to unseen hospitals), practical limitations mean that for some datasets, we have data from only a limited number of domains (e.g., one OOD test hospital in Camelyon17-wilds). As there can be substantial variability in performance across domains, developers should be careful to avoid overfitting to the specific test sets in Wilds, especially on datasets like Camelyon17-wilds with limited test domains. We strongly encourage all model developers to use the provided OOD validation sets for development and model selection, and to only use the OOD test sets for their final evaluations.
4 Reporting both ID and OOD performance
Prior work has shown that for many tasks, ID and OOD performance can be highly correlated across different model architectures and hyperparameters (Taori et al., 2020; Liu et al., 2021b; Miller et al., 2021). It is reasonable to expect that methods for improving ID performance could also give corresponding improvements in OOD performance in Wilds, and we welcome submissions of such methods. To better understand the extent to which any gains in OOD performance can be attributed to improved ID performance versus a model that is more robust to (i.e., less affected by) the distribution shift, we encourage model developers to report both ID and OOD performance numbers. See Miller et al. (2021) for an in-depth discussion of this point.
5 Extensions to other problem settings
In this paper, we focused on the domain generalization and subpopulation shift settings. In Appendix C, we discuss how Wilds can be used in other realistic problem settings that allow training algorithms to leverage additional information, such as unlabeled test data in unsupervised domain adaptation (Ben-David et al., 2006). These sources of leverage could be fruitful approaches to improving OOD performance, and we welcome community contributions towards this effort.
Using the Wilds package
Finally, we discuss our open-source PyTorch-based package that exposes a simple interface to our datasets and automatically handles data downloads, allowing users to get started on a Wilds dataset in just a few lines of code. In addition, the package provides various data loaders and utilities surrounding domain annotations and other metadata, which supports training algorithms that need access to these metadata. The package also provides standardized evaluations for each dataset. More documentation and installation information can be found at https://wilds.stanford.edu.
The Wilds package provides a simple, standardized interface for all datasets in the benchmark as well as their data loaders, as summarized in Figure 13. This short code snippet covers all of the steps of getting started with a Wilds dataset, including dataset download and initialization, accessing various splits, and initializing the data loader. We also provide multiple data loaders in order to accommodate a wide array of algorithms, which often require specific data loading schemes.
To allow algorithms to leverage domain annotations as well as other groupings over the available metadata, the Wilds package provides Grouper objects. Grouper objects (e.g., grouper in Figure 14) extract group annotations from metadata, allowing users to specify the grouping scheme in a flexible fashion.
Finally, the Wilds package standardizes and automates the evaluation for each dataset. As summarized in Figure 15, invoking the eval method of each dataset yields all metrics reported in the paper and on the leaderboard.
Reproducibility
An executable version of our paper, hosted on CodaLab, can be found at https://wilds.stanford.edu/codalab. This contains the exact commands, code, environment, and data used for the experiments reported in our paper, as well as all trained model weights. The WILDS package is open-source and can be found at https://github.com/p-lambda/wilds.
Acknowledgements
Many people generously volunteered their time and expertise to advise us on Wilds. We are grateful for all of the helpful suggestions and constructive feedback from: Aditya Khosla, Andreas Schlueter, Annie Chen, Aleksander Madry, Alexander D’Amour, Allison Koenecke, Alyssa Lees, Ananya Kumar, Andrew Beck, Behzad Haghgoo, Charles Sutton, Christopher Yeh, Cody Coleman, Dan Hendrycks, Dan Jurafsky, Daniel Levy, Daphne Koller, David Tellez, Erik Jones, Evan Liu, Fisher Yu, Georgi Marinov, Hongseok Namkoong, Irene Chen, Jacky Kang, Jacob Schreiber, Jacob Steinhardt, Jared Dunnmon, Jean Feng, Jeffrey Sorensen, Jianmo Ni, John Hewitt, John Miller, Kate Saenko, Kelly Cochran, Kensen Shi, Kyle Loh, Li Jiang, Lucy Vasserman, Ludwig Schmidt, Luke Oakden-Rayner, Marco Tulio Ribeiro, Matthew Lungren, Megha Srivastava, Nelson Liu, Nimit Sohoni, Pranav Rajpurkar, Robin Jia, Rohan Taori, Sarah Bird, Sharad Goel, Sherrie Wang, Shyamal Buch, Stefano Ermon, Steve Yadlowsky, Tatsunori Hashimoto, Tengyu Ma, Vincent Hellendoorn, Yair Carmon, Zachary Lipton, and Zhenghao Chen.
The design of the WILDS benchmark was inspired by the Open Graph Benchmark (Hu et al., 2020b), and we are grateful to the Open Graph Benchmark team for their advice and help in setting up our benchmark.
This project was funded by an Open Philanthropy Project Award and NSF Award Grant No. 1805310. Shiori Sagawa was supported by the Herbert Kunzel Stanford Graduate Fellowship. Henrik Marklund was supported by the Dr. Tech. Marcus Wallenberg Foundation for Education in International Industrial Entrepreneurship, CIFAR, and Google. Sang Michael Xie and Marvin Zhang were supported by NDSEG Graduate Fellowships. Weihua Hu was supported by the Funai Overseas Scholarship and the Masason Foundation Fellowship. Sara Beery was supported by an NSF Graduate Research Fellowship and is a PIMCO Fellow in Data Science. Jure Leskovec is a Chan Zuckerberg Biohub investigator. Chelsea Finn is a CIFAR Fellow in the Learning in Machines and Brains Program.
We also gratefully acknowledge the support of DARPA under Nos. N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), IIS-2030477 (RAPID); Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Chan Zuckerberg Biohub, Amazon, JPMorgan Chase, Docomo, Hitachi, JD.com, KDDI, NVIDIA, Dell, Toshiba, and UnitedHealth Group.
References
A Dataset realism
In this section, we discuss the framework we use to assess the realism of a benchmark dataset. Realism is subtle to pin down and highly contextual, and assessing realism often requires consulting with domain experts and practitioners. As a general framework, we can view a benchmark dataset as comprising the data, a task and associated evaluation metric, and a train/test split that potentially reflects a distribution shift. Each of these components can independently be more or less realistic:
The data—which includes not just the inputs but also any associated metadata (e.g., the domain that each data point came from)— is realistic if it accurately reflects what would plausibly be collected and available for a model to use in a real application. The realism of data also depends on the application context; for example, using medical images captured with state-of-the-art equipment might be realistic for well-equipped hospitals, but not necessarily for clinics that use older generations of the technology, or vice versa. Extreme examples of unrealistic data include the Gaussian distributions that are often used to cleanly illustrate the theoretical properties of various algorithms.
The task and evaluation metric is realistic if the task is relevant to a real application and if the metric measures how successful a model would be in that application. Here and with the other components, realism lies on a spectrum. For example, in a wildlife conservation application where the inputs are images from camera traps, the real task might be to estimate species populations (Parham et al., 2017), i.e., the number of distinct individual animals of each species seen in the overall collection of images; a task that is less realistic but still relevant and useful for ecologists might be to classify what species of animal is seen in each image (Tabak et al., 2019). The choice of evaluation metric is also important. In the wildlife example, conservationists might care more about rare species than common species, so measuring average classification accuracy would be less realistic than a metric that prioritizes classifying the rare species correctly.
The distribution shift (train/test split) is realistic if it reflects training and test distributions that might arise in deployment for that dataset and task. For example, if a medical algorithm is trained on data from a few hospitals and then expected to be deployed more widely, then it would be realistic to test it on hospitals that are not in the training set. On the other hand, an example of a less realistic shift is to, for instance, train a pedestrian classifier entirely on daytime photos and then test it only on nighttime photos; in practice, any reasonable dataset for pedestrian detection that is used in a real application would include both daytime and nighttime photos.
Through the lens of this framework, existing ML benchmarks tend to focus on object recognition tasks with realistic data (e.g., photos) but not necessarily with realistic distribution shifts. With Wilds, we seek to address this gap by selecting datasets that represent a wide variety of tasks (with realistic evaluation metrics and data) and that reflect realistic distribution shifts, i.e., train/test splits that are likely to arise in real-world deployments.
To elaborate on the realism of the distribution shift, we associate each dataset in Wilds with the distribution shift (i.e., problem setting) that we believe best reflects the real-world challenges in the corresponding application area. For example, domain generalization is a realistic setting for the Camelyon17-wilds dataset as medical models are typically trained on data collected from a handful of hospitals, but with the goal of general deployment across different hospitals. On the other hand, subpopulation shift is appropriate for the CivilComments-wilds dataset, as the real-world challenge is that some demographic subpopulations (domains) are underrepresented, rather than completely unseen, in the training data. The appropriate problem setting depends on many dataset-specific factors, but some common considerations include:
Domain type. Certain types of domains are generally more appropriate for a particular setting. For example, if the domains represent time, as in FMoW-wilds, then domain generalization is suitable as a common challenge is to generalize from past data to future data. On the other hand, if the domains represent demographics and the goal is to improve performance on minority subpopulations, as in CivilComments-wilds, then subpopulation shift is typically more appropriate.
Data collection challenges. When collecting data from a new domain is expensive, domain generalization is often appropriate, as we might want to train on data from a limited number of domains but still generalize to unseen domains. For example, it is difficult to collect patient data from multiple hospitals, as in Camelyon17-wilds, or survey data from new countries, as in PovertyMap-wilds.
Continuous addition of new domains. A special case of the above is when new domains are continuously created. For example, in Amazon-wilds, where domains correspond to users, new users are constantly signing up for the platform; and in iWildCam2020-wilds, where domains correspond to camera traps, new cameras are constantly being deployed. These are natural domain generalization settings.
B Prior work on ML benchmarks for distribution shifts
In this section, we discuss existing ML distribution shift benchmarks in more detail, categorizing them by how they induce their respective distribution shifts. We focus here on work that has appeared in ML conferences and journals; we discuss related work from other research communities in Section 8 and Appendix E. We also restrict our attention to publicly-available datasets. While others have studied some proprietary datasets with realistic distribution shifts, such as the StreetView StoreFronts dataset (Hendrycks et al., 2020b) or diabetic retinopathy datasets (D’Amour et al., 2020a), these datasets are not publicly available due to privacy and other commercial reasons.
Some of the most widely-adopted benchmarks induce distribution shifts by synthetically transforming the data. Examples include rotated and translated versions of MNIST and CIFAR (Worrall et al., 2017; Gulrajani and Lopez-Paz, 2020); surface variations such as texture, color, and corruptions like blur in Colored MNIST (Gulrajani and Lopez-Paz, 2020), Stylized ImageNet (Geirhos et al., 2018a), ImageNet-C (Hendrycks and Dietterich, 2019), and similar ImageNet variants (Geirhos et al., 2018b); and datasets that crop out objects and replace their backgrounds, as in the Backgrounds Challenge (Xiao et al., 2020) and other similar datasets (Sagawa et al., 2020a; Koh et al., 2020). Benchmarks for adversarial robustness also fall in this category of distribution shifts from transformations (Goodfellow et al., 2015; Croce et al., 2020). Though adversarial robustness is not a focus of this work, we note that recent work on temporal perturbations with the ImageNet-Vid-Robust and YTBB-Robust datasets (Shankar et al., 2019) represents a different form of distribution shift that also impacts real-world applications. Outside of visual object recognition, other work has used synthetic datasets and transformations to explore compositional generalization, e.g., SCAN (Lake and Baroni, 2018). We discuss this more in Section 8.
Fully synthetic datasets such as SYNTHIA (Ros et al., 2016) and StreetHazards (Hendrycks et al., 2020a) have been adopted for out-of-distribution detection as well as domain adaptation and generalization, e.g., by testing robustness to transformations in the seasons, weather, time, or architectural style (Hoffman et al., 2018; Volpi et al., 2018). While the data is synthetic, it can still look realistic if a high-fidelity simulator is used. In particular, synthetic benchmarks that study transfers from synthetic to real data (Ganin and Lempitsky, 2015; Richter et al., 2016; Peng et al., 2018) can be important tools for tackling real-world problems: even though the data is synthesized and by definition, not real, the synthetic-to-real distribution shift can still be realistic in contexts where real data is much harder to acquire than synthetic data (Bellemare et al., 2020). In this work, we do not study these types of synthetic distribution shifts; instead, we focus on distribution shifts that occur in the wild between real data distributions.
Other benchmarks do not rely on transformations but instead split the data in a way that induces particular distribution shifts. These benchmarks have realistic data, e.g., the data points are derived from real-world photos, but they do not necessarily reflect distribution shifts that would arise in the wild. For example, BREEDS (Santurkar et al., 2020) and a related dataset (Hendrycks and Dietterich, 2019) test generalization to unseen ImageNet subclasses by holding out subclasses specified by several controllable parameters; similarly, NICO (He et al., 2020) considers subclasses that are defined by their context, such as dogs at home versus dogs on the beach; DeepFashion-Remixed (Hendrycks et al., 2020b) constrains the training set to include only photos from a single camera viewpoint and tests generalization to unseen camera viewpoints; BDD-Anomaly (Hendrycks et al., 2020a) uses a driving dataset but with all motorcycles, trains, and bicycles removed from the training set only; and ObjectNet (Barbu et al., 2019) comprises images taken from a few pre-specified viewpoints, allowing for systematic evaluation for robustness to camera angle changes but deviating from natural camera angles.
A well-studied special case of the above category is the class of distribution shifts obtained by combining several disparate datasets (Torralba and Efros, 2011), training on one or more of them and then testing on the remaining datasets. A recent influential example is the ImageNetV2 dataset (Recht et al., 2019), which was constructed to be similar to the original ImageNet dataset. Unlike ImageNetV2, however, many of these distribution shifts were constructed to be more drastic than might arise in the wild. For example, standard domain adaptation benchmarks include training on MNIST but testing on SVHN street signs (LeCun et al., 1998; Yuval et al., 2011; Tzeng et al., 2017; Hoffman et al., 2018), as well as transfers across datasets containing different renditions (e.g., photos, clipart, sketches) in DomainNet (Peng et al., 2019) and the Office-Home dataset (Venkateswara et al., 2017).
The main difference between domain adaptation and domain generalization is that in the latter, we do not assume access to unlabeled data from the test distribution. This makes it straightforward to use domain adaptation benchmarks for domain generalization, e.g., in DomainBed (Gulrajani and Lopez-Paz, 2020); we focus on domain generalization in this work, but further discuss unsupervised domain adaptation in Section C. Other similar benchmarks that have been proposed for domain generalization include VLCS (Fang et al., 2013), which tests generalization across similar visual object recognition datasets; PACS (Li et al., 2017a), which (like DomainNet) tests generalization across datasets with different renditions; and ImageNet-R (Hendrycks et al., 2020b) and ImageNet-Sketch (Wang et al., 2019c), which also test generalization across different renditions by collecting separate datasets from Flickr and Google Image queries.
C Potential extensions to other problem settings
In this paper, we have focused on two problem settings involving domain shifts: domain generalization and subpopulation shifts. Here, we discuss other problem settings within the framework of domain shifts that could also apply to Wilds datasets. Using Wilds to benchmark and develop algorithms for these settings is an important avenue for future work, and we welcome community contributions towards this effort.
Within the general framework of domain shifts, specific problem settings can differ along the following axes of variation:
Seen versus unseen test domains. Test domains may be seen during training time (), as in subpopulation shift, or unseen (), as in domain generalization. The domain generalization and subpopulation shift settings mainly differ on this factor.
Train-time domain annotations. The domain identity may be observed for none, some, or all of the training examples. Train-time domain annotations are straightforward to obtain in some settings, e.g., we should know which patients in the training sets came from which hospitals, but can be harder to obtain in some settings, e.g., we might only have demographic information on a subset of training users. In our domain generalization and subpopulation shift settings, is always observed at training time.
Test-time domain annotations. The domain identity may be observed for none, some, or all of the test examples. Test-time domain annotations allow models to be domain-specific, e.g., by treating domain identity as a feature if the train and test domains overlap. For example, if the domains correspond to continents and the data to satellite images from a continent, we would presumably know what continent each image was taken from. On the other hand, if the domains correspond to demographic information, this might be hard to obtain at test time (as well as training time, as mentioned above). In domain generalization, may be observed at test time, but it is not helpful by itself as all of the test domains are unseen at training time. However, when combined with test-time unlabeled data, observing the domain at test time could help with adaptation. In subpopulation shift, we typically assume that is unobserved at test time, though this need not always be true.
Test-time unlabeled data. Varying amounts of unlabeled test data—samples of drawn from the test distribution —may be available, from none to a small batch to a large pool. This affects the degree to which models can adapt to test distributions. For example, if the domains correspond to locations and the data points to photos taken at those locations, we might assume access to some unlabeled photos taken at the test locations.
Each combination of the above four factors corresponds to a specific problem setting with a different set of applicable methods. In the current version of the Wilds benchmark, we focus on domain generalization and subpopulation shifts, which represent specific configurations of these factors. We briefly discuss a few other problem settings in the remainder of this section.
C.2 Unsupervised domain adaptation
In the presence of distribution shift, a potential source of leverage is observing unlabeled test points from the test distribution. In the unsupervised domain adaptation setting, we assume that at training time, we have access to a large amount of unlabeled data from each test distribution of interest, as well as the resources to train a separate model for each test distribution. For example, in a satellite imagery setting like FMoW-wilds, it might be appropriate to assume that we have access to a large set of unlabeled recent satellite images from each continent and the wherewithal to train a separate model for each continent.
Many of the methods for domain generalization discussed in Section 6 were originally methods for domain adaptation, since methods for both settings share the common goal of learning models that can transfer between domains. For example, methods that learn features that have similar distributions across domains are equally applicable to both settings (Ben-David et al., 2006; Long et al., 2015; Sun et al., 2016; Ganin et al., 2016; Tzeng et al., 2017; Shen et al., 2018; Wu et al., 2019b). In fact, the CORAL algorithm that we use as a baseline in this work was originally developed for, and successfully applied in, unsupervised domain adaptation (Sun and Saenko, 2016). Other methods rely on knowing the test distribution and are thus specific to domain adaptation, e.g., learning to map data points from source to target domains (Hoffman et al., 2018), or estimating the test label distribution from unlabeled test data (Saerens et al., 2002; Zhang et al., 2013; Lipton et al., 2018; Azizzadenesheli et al., 2019; Alexandari et al., 2020; Garg et al., 2020).
C.3 Test-time adaptation
A closely related setting to unsupervised domain adaptation is test-time adaptation, which also assumes the availability of unlabeled test data. For datasets where there are many potential test domains (e.g., in iWildCam2020-wilds, we want a model that can ideally generalize to any camera trap), it might be infeasible to train a separate model for each test domain, as unsupervised domain adaptation would require. In the test-time adaptation setting, we assume that a model is allowed to adapt to a small amount of unlabeled test data in a way that is computationally much less intensive than typical domain adaptation methods. This is a difference of degree and not of kind, but it can have significant practical implications. For example, domain adaptation approaches typically require access to the training set and a large unlabeled test set at the same time, whereas test-time adaptation methods typically only require the learned model (which can be much smaller than the original training set) as well as a smaller amount of unlabeled test data.
A number of test-time adaptation methods have been recently proposed (Li et al., 2017c; Sun et al., 2020b; Wang et al., 2020a). For example, adaptive risk minimization (ARM) is a meta-learning approach that adapts models to each batch of test examples under the assumption that all data points in a batch come from the same domain (Zhang et al., 2020). Many datasets in Wilds are suitable for the test-time adaptation setting. For example, in iWildCam2020-wilds, images from the same domain are highly similar, sharing the same location, background, and camera angle, and prior work has shown inferring these shared features can improve performance considerably (Beery et al., 2020b).
C.4 Selective prediction
A different problem setting that is orthogonal to the settings described above is selective prediction. In the selective prediction setting, models are allowed to abstain on points where their confidence is below a certain threshold. This is appropriate when, for example, abstentions can be handled by backing off to human experts, such as pathologists for Camelyon17-wilds, content moderators for CivilComments-wilds, wildlife experts for iWildCam2020-wilds, etc. Many methods for selective prediction have been developed, from simply using softmax probabilities as a proxy for confidence (Cordella et al., 1995; Geifman and El-Yaniv, 2017), to methods involving ensembles of models (Gal and Ghahramani, 2016; Lakshminarayanan et al., 2017; Geifman et al., 2018) or jointly learning to abstain and classify (Bartlett and Wegkamp, 2008; Geifman and El-Yaniv, 2019; Feng et al., 2019).
Intuitively, even if a model is not robust to a distribution shift, it might at least be able to maintain high accuracies on some subset of points that are close to the training distribution, while abstaining on the other points. Indeed, prior work has shown that selective prediction can improve model accuracy under distribution shifts (Pimentel et al., 2014; Hendrycks and Gimpel, 2017; Liang et al., 2018; Ovadia et al., 2019; Feng et al., 2019; Kamath et al., 2020). However, distribution shifts still pose a problem for selective prediction methods; for instance, it is difficult to maintain desired abstention rates under distribution shifts (Kompa et al., 2020), and confidence estimates have been found to drift over time (e.g., Davis et al. (2017)).
D Additional experimental details
For each hyperparameter setting, we used early stopping to pick the epoch with the best OOD validation performance (as measured by the specified metrics for each dataset described in Section 4), and then picked the model hyperparameters with the best early-stopped validation performance. We found that this gave similar or slightly better OOD test performance than selecting hyperparameters using the ID validation set (Table 3).
Using the OOD validation set for early stopping means that even if the training procedure does not explicitly use additional metadata, as in ERM, the metadata might still be implicitly (but mildly) used for model selection in one of two related ways. First, the metric might use the metadata directly (e.g., by computing the accuracy over different subpopulations defined in the metadata). Second, the OOD validation set is generally selected according to this metadata (e.g., comprising data from a disjoint set of domains as the training set). We expect that implicitly using the metadata in these ways should increase the OOD performance of each model. Nevertheless, as Sections 5 and 6 show, there are still large gaps between OOD and ID performance.
In general, we selected model hyperparameters with ERM and used the same hyperparameters for the other algorithm baselines (e.g., CORAL, IRM, or Group DRO). For CORAL and IRM, we did a subsequent grid search over the weight of the penalty term, using the defaults from Gulrajani and Lopez-Paz (2020). Specifically, we tried penalty weights of for CORAL and penalty weights of for IRM. We fixed the step size hyperparameter for Group DRO to its default value of 0.01 (Sagawa et al., 2020a).
D.2 Replicates
We typically use a fixed train/validation/test split and report results averaged across 3 replicates (random seeds for model initialization and minibatch order), as well as the unbiased standard deviation over those replicates. There are three exceptions to this. For PovertyMap-wilds, we report results averaged over 5-fold cross validation, as model training is relatively fast on this dataset. For Camelyon17-wilds, results vary substantially between replicates, so we report results averaged over 10 replicates instead. Similarly, for CivilComments-wilds, we report results averaged over 5 replicates.
D.3 Baseline algorithms
For all classification datasets, we train models against the cross-entropy loss. For the PovertyMap-wilds regression dataset, we use the mean-squared-error loss.
We adapted the implementations of CORAL from Gulrajani and Lopez-Paz (2020); IRM from Arjovsky et al. (2019); and Group DRO from Sagawa et al. (2020a). We note that CORAL was originally proposed in the context of domain adaptation (Sun and Saenko, 2016), where it was shown to substantially improve performance on standard domain adaptation benchmarks, and it was subsequently adapted for domain generalization (Gulrajani and Lopez-Paz, 2020).
Following these implementations, we use minibatch stochastic optimizers to train models under each algorithm, and we sample uniformly from each domain regardless of the number of training examples in it. This means that the CORAL and IRM algorithms optimize for their respective penalty terms plus a reweighted ERM objective that weights each domain equally (i.e., effectively upweighting minority domains). The Group DRO objective is unchanged, as it still optimizes for the domain with the worst loss, but the uniform sampling improves optimization stability.
Both CORAL and IRM are designed for models with featurizers, i.e., models that first map each input to a feature representation and then predict based on the representation. To estimate the feature distribution for a domain, these algorithms need to see a sufficient number of examples from that domain in a minibatch. However, some of our datasets have large numbers of domains, making it infeasible for each minibatch to contain examples from all domains. For these algorithms, our data loaders form a minibatch by first sampling a few domains, and then sampling examples from those domains. For consistency in our experiments, we used the same total batch size for these algorithms and for ERM and Group DRO, with a default of 8 examples per domain in each minibatch (e.g., if the batch size was 32, then in each minibatch we would have 8 examples 4 domains).
For Group DRO, as in Sagawa et al. (2020a), each example in the minibatch is sampled independently with uniform probabilities across domains, and therefore each minibatch does not need to only comprise a small number of domains. We note that reweighting methods like Group DRO are effective only when the training loss is non-vanishing, which we achieve through early stopping (Byrd and Lipton, 2019; Sagawa et al., 2020a, b).
E Additional dataset details and results
In this section, we discuss each Wilds dataset in more detail. For completeness, we start by repeating the motivation behind each dataset from Section 4. We then describe the task, the distribution shift, and the evaluation criteria, and present baseline results that elaborate upon those in Sections 5 and 6. We also discuss the broader context behind each dataset and how it connects with other distribution shifts in similar applications. Finally, we describe how each dataset was modified from its original version in terms of the evaluation, splits, and data. Unless otherwise specified, all experiments follow the protocol laid out in Appendix D.
Animal populations have declined 68% on average since 1970 (Grooten et al., 2020). To better understand and monitor wildlife biodiversity loss, ecologists commonly deploy camera traps—heat or motion-activated static cameras placed in the wild (Wearn and Glover-Kapfer, 2017)—and then use ML models to process the data collected (Weinstein, 2018; Norouzzadeh et al., 2019; Tabak et al., 2019; Beery et al., 2019; Ahumada et al., 2020). Typically, these models would be trained on photos from some existing camera traps and then used across new camera trap deployments. However, across different camera traps, there is drastic variation in illumination, color, camera angle, background, vegetation, and relative animal frequencies, which results in models generalizing poorly to new camera trap deployments (Beery et al., 2018).
We study this shift on a variant of the iWildCam 2020 dataset (Beery et al., 2020a).
We consider the domain generalization setting, where the domains are camera traps, and we seek to learn models that generalize to photos taken from new camera deployments (Figure 3). The task is multi-class species classification. Concretely, the input is a photo taken by a camera trap, the label is one of 182 different animal species, and the domain is an integer that identifies the camera trap that took the photo.
The dataset comprises 203,029 images from 323 different camera traps spread across multiple countries in different parts of the world. The original camera trap data comes from the Wildlife Conservation Society (http://lila.science/datasets/wcscameratraps). These images tend to be taken in short bursts following motion-activation of a camera trap, so the images can be additionally grouped into sequences of images from the same burst, though our baseline models do not exploit this information and our evaluation metric treats each image individually. Each image is associated with the following metadata: camera trap ID, sequence ID, and datetime.
As is typical for camera trap data, approximately 35% of the total number of images are empty (i.e., do not contain any animal species); this corresponds to one of the 182 class labels. The ten most common classes across the full dataset are “empty” (34%), ocellated turkey (8%), great curassow (6%), impala (4%), black-fronted duiker (4%), white-lipped peccary (3%), Central American agouti (3%), ocelot (3%), gray fox (2%) and cow (2%).
We note that the labels in this dataset can be somewhat noisy, as is typical of camera trap data. Some ecologists might label all images in a sequence as the same animal (which can result in empty/dark frames being labeled as an animal), whereas other ecologists might try to label it frame-by-frame. This label noise imposes a natural ceiling on model performance, though the label noise is equally present in ID vs. OOD data.
We split the dataset by randomly partitioning the data by camera traps:
Training: 129,809 images taken by 243 camera traps.
Validation (OOD): 14,961 images taken by 32 different camera traps.
Test (OOD): 42,791 images taken by 48 different camera traps.
Validation (ID): 7,314 images taken by the same camera traps as the training set, but on different days from the training and test (ID) images.
Test (ID): 8,154 images taken by the same camera traps as the training set, but on different days from the training and validation (ID) images.
The camera traps are randomly distributed across the training, validation (OOD), and test (OOD) sets. The number of examples per location vary widely from 1 to 8494, with a median of 194 images (Figure 16). All images from the same sequence (i.e., all images taken in the same burst) are placed together in the same split. See Appendix E.1.4 for more details.
We evaluate models by their macro F1 score (i.e., we compute the F1 score for each class separately, then average those scores). We also report the average accuracy of each model across all test images, but primarily use the macro F1 score to better capture model performance on rare species. In the natural world, protected and endangered species are rare by definition, and are often the most important to accurately monitor. However, common species are much more likely to be captured in camera trap images; this imbalance can make metrics like average accuracy an inaccurate picture of model effectiveness.
Though the problem is challenging for existing ML algorithms, adapting to photos from different camera traps is simple and intuitive for humans. Repeated backgrounds and habitual animals, which cause each sensor to have a unique class distribution, provide a strong implicit signal across data from any one location. We anticipate that approaches that utilize the provided camera trap annotations can learn to factor out these common features and avoid learning spurious correlations between particular backgrounds and animal species.
E.1.2 Baseline results
For all experiments, we use ResNet-50 models (He et al., 2016) that were pretrained on ImageNet, using a learning rate of 3e-5 and no -regularization. As input, these models take in images resized to 448 by 448. We trained these models with the Adam optimizer and a batch size of 16 for 12 epochs. To pick hyperparameters, we did a grid search over learning rates and regularization strengths . We report results aggregated over 3 random seeds.
Model performance dropped substantially and consistently going from the train-to-train in-distribution (ID) setting to the official out-of-distribution (OOD) setting (Table 4), with a macro F1 score of 47.0 on the ID test set but only 31.0 on the OOD test set. We note that macro F1 and average accuracy both differ between the OOD validation and test sets: this is in part due to the difference in class balance between them, which in turn is due to differences in the proportion of classes across camera traps. In particular, macro F1 can vary between splits because we take the average F1 score across all classes that are present in the evaluation split, and not all splits have the same classes present (e.g., a rare species might be in the OOD validation set but not OOD test set, or vice versa). In additional, average accuracy can differ between splits due in part to variation in the fraction of empty images per location (e.g., the camera traps that were randomly assigned to the OOD validation set have a smaller proportion of empty images).
We only ran a train-to-train comparison because there are a relatively large number of domains (camera traps) split i.i.d. between the training and test sets, which suggests that the training and test sets should be “equally difficult”. The size of the ID-OOD gap in macro F1 is large enough that we expect it should hold up even in a test-to-test comparison. However, the results in Table 4 and Figure 17 show that there is substantial variability between domains, and it would be useful for future work to establish the magnitude of the ID-OOD gap under the test-to-test or mixed-to-test comparisons.
We trained models with CORAL, IRM, and Group DRO, treating each camera trap as a domain, and using the same model hyperparameters as ERM. These did not improve upon the ERM baseline (Table 4). The IRM models performed especially poorly on this dataset; we suspect that this is because the default estimator of the IRM penalty term can be negatively biased when examples are sampled without replacement from small domains, but further investigation is needed. We also tried reweighting the training data so that each label had equal weight, but this did not improve over ERM either.
Across locations, there is drastic variation in illumination, camera angle, background, vegetation, and color. This variation, coupled with considerable differences in the distribution of animals between camera traps, likely encourages the model to overfit to specific animal species appearing in specific locations, which may account for the performance drop.
The original iWildCam 2020 competition allows users to use MegaDetector (Beery et al., 2019), which is an animal detector trained on a large set of data beyond what is provided in the training set. Using an animal detection model like MegaDetector typically improves classification performance on camera traps (Beery et al., 2018). However, we intentionally do not use MegaDetector in our baselines for iWildCam2020-wilds for two reasons. First, though the trained MegaDetector model is publicly available, the MegaDetector training set is not, which makes it difficult to build on top of it and run controlled experiments. Second, bounding box annotations are costly and harder to obtain, and there is much more data with image-level species label, so it would be useful to be able to train models that do not have to rely on bounding box annotations.
We still welcome leaderboard submissions that use MegaDetector, as it is useful to see how much better models can perform when they use MegaDetector or other similar animal detectors, but we will distinguish these submissions from others that only use what is provided in the training set.
A different source of leverage comes from the temporal signal in the camera trap images, which are organized into sequences that each correspond to a burst of images from a single motion trigger. Using this sequence information (e.g., by taking the median prediction across a sequence) can also improve model performance (Beery et al., 2018), and we welcome submissions that explore this direction.
E.1.3 Broader context
Differences across data distributions at different sensor locations is a common challenge in automated wildlife monitoring applications, including using audio sensors to monitor animals that are easier heard than seen such as primates, birds, and marine mammals (Crunchant et al., 2020; Stowell et al., 2019; Shiu et al., 2020), and using static sonar to count fish underwater to help maintain sustainable fishing industries (Pipal et al., 2012; Vatnehol et al., 2018; Schneider and Zhuang, 2020). As with camera traps, each static audio sensor has a specific species distribution as well as a sensor specific background noise signature, making generalization to new sensors challenging. Similarly, static sonar used to measure fish escapement have sensor-specific background reflectance based on the shape of the river bottom. Moreover, since species are distributed in a non-uniform and long-tailed fashion across the globe, it is incredibly challenging to collect sufficient samples for rare species to escape the low-data regime. Implicitly representing camera-specific distributions and background features in per-camera memory banks and extracting relevant information from these via attention has been shown to help overcome some of these challenges for static cameras (Beery et al., 2020b).
More broadly, shifts in background, image illumination and viewpoint have been studied in computer vision research. First, several works have shown that object classifiers often rely on the background rather than the object to make its classification (Rosenfeld et al., 2018; Shetty et al., 2019; Xiao et al., 2020). Second, common perturbations such as blurriness or shifts in illumination, tend to reduce performance (Dodge and Karam, 2017; Temel et al., 2018; Hendrycks and Dietterich, 2019). Finally, shifts in rotation and viewpoint of the object has been shown to degrade performance (Barbu et al., 2019).
E.1.4 Additional details
We generate the data splits in three steps. First, to generate the OOD splits, we randomly split all locations into three groups: Validation (OOD), Test (OOD), and Others. Then, to generate the train-to-train ID splits, we split the Others group uniformly by date at random into three sets: Training, Validation (ID), and Test (ID).
When doing the ID split, some locations only ended up in some of but not all of Training, Validation (ID), and Test (ID). For instance, if there were very few dates for a specific location (camera trap), it may be that no examples from that location ended up in the train split. This defeats the purpose of the ID split, which is to test performance on locations that were seen during training. We therefore put these locations in the train split. Finally, any images in the test set with classes not present in the train set were removed.
The original iWildCam 2020 Kaggle competition similarly split the dataset by camera trap, though the competition focused on average accuracy. We consider a smaller subset of the data here. Specifically, the Kaggle competition uses a held-out test set that we are not utilizing, as the test set is intended to be reused in a future competition and is not yet public. Instead, we constructed our own test set by splitting the Kaggle competition training data into our own splits: train, validation (ID), validation (OOD), test (ID), test (OOD).
Images are organized into sequences, but we treat each image separately. In the iWildCam 2020 competition, the top participants utilized the sequence data and also used a pretrained MegaDetector animal detection model that outputs bounding boxes over the animals. These images are cropped using the bounding boxes and then fed into a classification network. As we discuss above, we intentionally do not use MegaDetector in our experiments.
In addition, compared to the iWildCam 2020 competition, the iWildCam 2021 competition changed several class definitions (such as removing the “unknown” class) and removed some images that were taken indoors or had humans in the background. We have applied these updates to iWildCam2020-wilds as well.
E.2 Camelyon17-wilds
Models for medical applications are often trained on data from a small number of hospitals, but with the goal of being deployed more generally across other hospitals. However, variations in data collection and processing can degrade model accuracy on data from new hospital deployments (Zech et al., 2018; AlBadawy et al., 2018). In histopathology applications—studying tissue slides under a microscope—this variation can arise from sources like differences in the patient population or in slide staining and image acquisition (Veta et al., 2016; Komura and Ishikawa, 2018; Tellez et al., 2019).
We study this shift on a patch-based variant of the Camelyon17 dataset (Bandi et al., 2018).
We consider the domain generalization setting, where the domains are hospitals, and our goal is to learn models that generalize to data from a hospital that is not in the training set (Figure 4). The task is to predict if a given region of tissue contains any tumor tissue, which we model as binary classification. Concretely, the input is a 96x96 histopathological image, the label is a binary indicator of whether the central 32x32 region contains any tumor tissue, and the domain is an integer that identifies the hospital that the patch was taken from.
The dataset comprises 450,000 patches extracted from 50 whole-slide images (WSIs) of breast cancer metastases in lymph node sections, with 10 WSIs from each of 5 hospitals in the Netherlands. Each WSI was manually annotated with tumor regions by pathologists, and the resulting segmentation masks were used to determine the labels for each patch. We also provide metadata on which slide (WSI) each patch was taken from, though our baseline algorithms do not use this metadata.
We split the dataset by domain (i.e., which hospital the patches were taken from):
Training: 302,436 patches taken from 30 WSIs, with 10 WSIs from each of the 3 hospitals in the training set.
Validation (OOD): 34,904 patches taken from 10 WSIs from the 4th hospital. These WSIs are distinct from those in the other splits.
Test (OOD): 85,054 patches taken from 10 WSIs from the 5th hospital, which was chosen because its patches were the most visually distinctive. These WSIs are also distinct from those in the other splits.
Validation (ID): 33,560 patches taken from the same 30 WSIs from the training hospitals.
We do not provide a Test (ID) set, as there is no practical setting in which we would have labels on a uniformly randomly sampled set of patches from a WSI, but no labels on the other patches from the same WSI.
We evaluate models by their average test accuracy across patches. Histopathology datasets can be unwieldy for ML models, as individual images can be several gigabytes large; extracting patches involves many design choices; the classes are typically very unbalanced; and evaluation often relies on more complex slide-level measures such as the free-response receiver operating characteristic (FROC) (Gurcan et al., 2009). To improve accessibility, we pre-process the slides into patches and balance the dataset so that each split has a 50/50 class balance, making average accuracy is a reasonable measure of performance (Veeling et al., 2018; Tellez et al., 2019).
Prior work has shown that differences in staining between hospitals are the primary source of variation in this dataset, and that specialized stain augmentation methods can close the in- and out-of-distribution accuracy gap on a variant of the dataset based on the same underlying slides (Tellez et al., 2019). However, the general task of learning histopathological models that are robust to variation across hospitals (from staining and other sources) is still an open research question. In this way, the Camelyon17-wilds dataset is a controlled testbed for general-purpose methods that can learn to be robust to stain variation between hospitals, given a training set from multiple hospitals.
E.2.2 Baseline results
For all experiments, we use DenseNet-121 models (Huang et al., 2017) models trained from scratch on the 96 96 patches, following prior work (Veeling et al., 2018). These models used a learning rate of , -regularization strength of , a batch size of 32, and SGD with momentum (set to 0.9), trained for 5 epochs with early stopping. We selected hyperparameters by a grid search over learning rates , , , and -regularization strengths . We report results aggregated over 10 random seeds.
Table 5 shows that the model was consistently accurate on the train-to-train in-distribution (ID) validation set and to a lesser extent on the out-of-distribution (OOD) validation set, which was from a held-out hospital. However, it was wildly inconsistent on the test set, which was from a different held-out hospital, with a standard deviation of 6.4% in accuracies across 10 random seeds. There is a large gap between train-to-train ID validation and OOD validation accuracy, and between OOD validation and OOD test accuracy (in part because we early stop on the highest OOD validation accuracy). Nevertheless, we found that using the OOD validation set gave better results than using the ID validation set; see Appendix D.1 for more discussion.
We ran an additional mixed-to-test comparison, where we moved 1 of the 10 slidesThis slide was randomly chosen and corresponded to about 6% of the test patches; some slides contribute more patches than others because they contain larger tumor regions. from the test hospital to the training set and tested on the patches from the remaining 9 slides. The mixed-to-test setting gives significantly higher accuracy on the reduced test set (Table 6), suggesting that the observed performance drop is due to the distribution shift, as opposed to the intrinsic difficulty of the examples from the test hospital. We note that this mixed-to-test comparison mixes in only a small amount of test data and is therefore likely to be an underestimate of in-distribution performance on the test set; we opted to only mix in 1 slide so as to preserve enough test examples to be able to accurately estimate model performance.
We trained models with CORAL, IRM, and Group DRO, treating each hospital as a domain. However, they performed comparably or worse than the ERM baseline. For the CORAL and IRM models, our grid search selected the lowest values of their penalty weights (0.1 and 1, respectively) based on OOD validation accuracy.
These results demonstrate a subtle failure mode when considering out-of-distribution accuracy: there are models (i.e., choices of hyperparameters and random seeds) that do well both in- and out-of-distribution, but we cannot reliably choose these models from just the training/validation set. Due to the substantial variability in test accuracy on Camelyon17-wilds (see Figure 18), we ask researchers to submit leaderboard submissions with results from 10 random seeds, instead of the 3 random seeds required for other datasets.
Many specialized methods have been developed to handle stain variation in the context of digital histopathology. These typically fall into one of two categories: data augmentation methods that perturb the colors in the training images (e.g., Liu et al. (2017); Bug et al. (2017); Tellez et al. (2018)) or stain normalization methods that seek to standardize colors across training images (e.g., Macenko et al. (2009); BenTaieb and Hamarneh (2017)). These methods are reasonably effective at mitigating stain variation, at least in some contexts (Tellez et al., 2019; Miller et al., 2021), though the general problem of learning digital histopathology models that can be effectively deployed across multiple hospitals/sites is still an open challenge.
To facilitate more controlled experiments, we will have two leaderboard tracks for Camelyon17-wilds. For the first track, which focuses on general-purpose algorithms, submissions should not use color-specific techniques (e.g., color augmentation) and should also train their models from scratch, instead of fine-tuning models that are pre-trained from ImageNet or other datasets. For the second track, submissions can use any of those techniques, including specialized methods for dealing with stain variation. These separate tracks will help to disentangle the contributions of more general-purpose learning algorithms and model architectures from the contributions of specialized augmentation techniques or additional training data.
E.2.3 Broader context
Other than stain variation, there are many other distribution shifts that might occur in histopathology applications. For example, patient demographics might differ from hospital to hospital: some hospitals might tend to see patients who are older or more sick, and patients from different backgrounds and countries vary in terms of cancer susceptibility (Henderson et al., 2012). Some cancer subtypes and tissues of origin are also more common than others, leading to potential subpopulation shift issues, e.g,. a rare cancer subtype in one context might be more common in another; or even if it remains rare, we would seek to leverage the greater quantity of data from other subtypes to improve model accuracy on the rare subtype (Weinstein et al., 2013).
Beyond histopathology, variation between different hospitals and deployment sites has also been shown to degrade model accuracy in other medical applications such as diabetic retinopathy (Beede et al., 2020) and chest radiographs (Zech et al., 2018; Phillips et al., 2020), including recent work on COVID-19 detection (DeGrave et al., 2020). Even within the same hospital, process variables like which scanner/technician took the image can significantly affect models (Badgeley et al., 2019).
In these medical applications, the gold standard is to evaluate models on an independent test set collected from a different hospital (e.g., Beck et al. (2011); Liu et al. (2017); Courtiol et al. (2019); McKinney et al. (2020)) or at least with a different scanner within the same hospital (e.g., Campanella et al. (2019)). However, this practice has not been ubiquitous due to the difficulty of obtaining data spanning multiple hospitals (Esteva et al., 2017; Bejnordi et al., 2017; Codella et al., 2019; Veta et al., 2019). The baseline results reported above show that even evaluating on a single different hospital might be insufficient, as results can vary widely between different hospitals (e.g., between the validation and test OOD datasets). We hope that the Camelyon17-wilds dataset, which has multiple hospitals in the training set and independent hospitals in the validation and test sets, will be useful for developing models that can generalize reliably to new hospitals and contexts (Chen et al., 2020).
E.2.4 Additional details
The Camelyon17-wilds dataset is adapted from whole-slide images (WSIs) of breast cancer metastases in lymph nodes sections, obtained from the CAMELYON17 challenge (Bandi et al., 2018). Each split is balanced to have an equal number of positive and negative examples. The varying number of patches per slide and hospital is due to this class balancing, as some slides have fewer tumor (positive) patches. We selected the test set hospital as the one whose patches were visually most distinct; the difference in test versus OOD validation performance shows that the choice of OOD hospital can significantly affect performance.
From these WSIs, we extracted patches in a standard manner, similar to Veeling et al. (2018). The WSIs were scanned at a resolution of 0.23m–0.25m in the original dataset, and each WSI contains multiple resolution levels, with approximately 10,00020,000 pixels at the highest resolution level (Bandi et al., 2018). We used the third-highest resolution level, corresponding to reducing the size of each dimension by a factor of 4. We then tiled each slide with overlapping 9696 pixel patches with a step size of 32 pixels in each direction (such that none of the central 3232 regions overlap), labeling them as the following:
Tumor patches have at least one pixel of tumor tissue in the central 3232 region. We used the pathologist-annotated tumor annotations provided with the WSIs.
Normal patches have no tumor and have at least 20% normal tissue in the central 3232 region. We used Otsu thresholding to distinguish normal tissue from background.
We discarded all patches that had no tumor and <20% normal tissue in the central 3232 region.
To maintain an equal class balance, we then subsampled the extracted patches in the following way. First, for each WSI, we kept all tumor patches unless the WSI had fewer normal than tumor patches, which was the case for a single WSI; in that case, we randomly discarded tumor patches from that WSI until the numbers of tumor and normal patches were equal. Then, we randomly selected normal patches for inclusion such that for each hospital and split, there was an equal number of tumor and normal patches.
The task in the original CAMELYON17 challenge (Bandi et al., 2018) was the patient-level classification task of determining the pathologic lymph node stage of the tumor present in all slides from a patient. In contrast, our task is a lesion-level classification task. Patient-level, slide-level, and lesion-level tasks are all common in histopathology applications. As mentioned above, the original dataset provided WSIs and tumor annotations, but not a standardized set of patches, which we provide here. Moreover, it did not consider distribution shifts; both of the original training and test splits contained slides from all 5 hospitals.
The Camelyon17-wilds patch-based dataset is similar to one of the datasets used in Tellez et al. (2019), which was also derived from the CAMELYON17 challenge; there, only one hospital is used as the training set, and the other hospitals are all part of the test set. Camelyon17-wilds is also similar to PCam (Veeling et al., 2018), which is a patch-based dataset based on an earlier CAMELYON16 challenge; the data there is derived from only two hospitals.
The full, original CAMELYON17 dataset contains 1000 WSIs from the same 5 hospitals, although only 50 of them (which we use here) have tumor annotations. The other 950 WSIs may be used as unlabeled data. Beyond the CAMELYON17 dataset, the largest source of unlabeled WSI data is the Cancer Genome Atlas (Weinstein et al., 2013), which typically has patient-level annotations (e.g., patient demographics and clinical outcomes).
E.3 RxRx1-wilds
High-throughput screening techniques that can generate large amounts of data are now common in many fields of biology, including transcriptomics (Harrill et al., 2019), genomics (Echeverri and Perrimon, 2006; Zhou et al., 2014), proteomics and metabolomics (Taylor et al., 2021), and drug discovery (Broach et al., 1996; Macarron et al., 2011; Swinney and Anthony, 2011; Boutros et al., 2015). Such large volumes of data, however, need to be created in experimental batches, or groups of experiments executed at similar times under similar conditions. Despite attempts to carefully control experimental variables such as temperature, humidity, and reagent concentration, measurements from these screens are confounded by technical artifacts that arise from differences in the execution of each batch. These batch effects make it difficult to draw conclusions from data across experimental batches (Leek et al., 2010; Parker and Leek, 2012; Soneson et al., 2014; Nygaard et al., 2016; Caicedo et al., 2017).
We study the shift induced by batch effects on a variant of the RxRx1-wilds dataset (Taylor et al., 2019). As illustrated in Figure 5, there are significant visual differences between experimental batches, making recognizing siRNA perturbations in OOD experiments in the RxRx1-wilds dataset a particularly challenging task for existing ML algorithms.
We consider the domain generalization setting, where the domains are experimental batches and we seek to generalize to images from unseen experimental batches. Concretely, the input is a 3-channel image of cells obtained by fluorescent microscopy, the label indicates which of the 1,139 genetic treatments (including no treatment) the cells received, and the domain specifies the experimental batch of the image.
RxRx1-wilds was created by Recursion (recursion.com) in its automated high-throughput screening laboratory in Salt Lake City, Utah. It is comprised of fluorescent microscopy images of human cells in four different cell lines: HUVEC, RPE, HepG2, and U2OS. These were acquired via fluorescent microscopy using a 6-channel variant of the Cell Painting assay (Bray et al., 2016). Figure 19 shows an example of the cellular contents of each of these 6 channels: nuclei, endoplasmic reticuli, actin, nucleoli and cytoplasmic RNA, mitochondria, and Golgi. To make the dataset smaller and more accessible, we only included the first 3 channels in RxRx1-wilds.
The images in RxRx1-wilds are the result of executing the same experimental design 51 different times, each in a different batch of experiments. The design consists of four 384-well plates, where individual wells are used to isolate populations of cells on each plate (see Figure 20).
The wells are laid out in a 1624 grid, but only the wells in the inner 1422 grid are used since the outer wells are most susceptible to environmental factors. Of these 308 usable wells, one is left untreated to provide a negative control phenotype, while the rest are treated with small interfering ribonucleic acid, or siRNA, at a fixed concentration. Each siRNA is designed to knockdown a single target gene via the RNA interference pathway, reducing the expression of the gene and its associated protein (Tuschl, 2001). However, siRNAs are known to have significant but consistent off-target effects via the microRNA pathway, creating partial knockdown of many other genes as well. The overall effect of siRNA transfection is to perturb the morphology, count, and distribution of cells, creating a phenotype associated with each siRNA. The phenotype is sometimes visually recognizable, but often the effects are subtle and hard to detect.
In each plate, 30 wells are set aside for 30 positive control siRNAs. Each has a different gene as its primary target, which together with the single untreated well already mentioned, provides a set of reference phenotypes per plate. Each of the remaining 1,108 wells of the design (277 wells 4 plates) receives one of 1,108 treatment siRNA, respectively, so that there is at most one well of each treatment siRNA in each experiment. We say at most once because, although rare, it happens that either an siRNA is not correctly transferred into the designated destination well, resulting in an additional untreated well, or an operational error is detected by quality control procedures that render the well unsuitable for inclusion in the dataset.
Each experiment was run in a single cell type, and of the 51 experiments in RxRx1-wilds, 24 are in HUVEC, 11 in RPE, 11 in HepG2, and 5 in U2OS. Figure 21 shows the phenotype of the same siRNA in each of these four cell types.
We split the dataset by experimental batches, with the training and test splits having roughly the same composition of cell types:
Training: 33 experiments (16 HUVEC, 7 RPE, 7 HepG2, 3 U2OS), site 1 only = 40,612 images.
Validation (OOD): 4 experiments (1 HUVEC, 1 RPE, 1 HepG2, 3 U2OS), sites 1 and 2 = 9,854 images.
Test (OOD): 14 experiments (7 HUVEC, 3 RPE, 3 HepG2, 1 U2OS), sites 1 and 2 = 34,432 images.
Test (ID): same 33 experiments as in the training set, site 2 only = 40,612 images.
In addition to the class (siRNA), each image is associated with the following metadata: cell type, experiment, plate, well, and site. We emphasize that all the images of an experiment are found in exactly one of the training, validation (OOD) or test (OOD) splits. See Appendix E.3.4 for more data processing details.
We evaluate models by their average accuracy across test images. Note that there are two images per well in the test set, which we evaluate independently.
The cell types are not balanced in the training and test sets. Correspondingly, we observed higher performance on the HUVEC cell type, which is over-represented, and lower performance on the U2OS cell type, which is under-represented. While maintaining high performance on minority (or even unseen) cell types is an important problem, for RxRx1-wilds, we opt to measure the average accuracy across all experiments instead of, for example, the worst accuracy across cell types. This is because the relatively small amount of training data available from the minority cell type (U2OS) makes it challenging to cast RxRx1-wilds as a tractable subpopulation shift problem. We also note that the difference in performance across cell types leads to the validation performance being significantly lower than the test performance, as there is a comparatively smaller fraction of HUVEC and a comparatively higher fraction of U2OS.
By design, there is usually one sample per class per experiment in the training set, with the following exceptions: 1) there are usually four samples per positive control, though 2) samples may be missing, as described above. Moreover, while batch effects can manifest themselves in many complicated ways, it is the case that the training set consists of a large number of experiments selected randomly amongst all experiments in the dataset, hence we expect models to be able to learn what is common amongst all such samples per cell type, and for that ability to generalize to to test batches. We emphasize that, whether in the training or test sets, the same cell types are perturbed with the same siRNA, and thus the phenotypic distributions for each batch share much of the same generative process.
We also note that, while not exploited here, there is quite a bit of structure in the RxRx1-wilds dataset. For example, except in the case of errors, all treatment siRNA appear once in each experiment, and all control conditions appear once per plate, so four times per experiment. Also, due to the operational efficiencies gained, the 1,108 treatment siRNAs always appear in the same four groups of 277 per experiment. So while the particular well an siRNA appears in is randomized, it will always appear with the same group of 276 other siRNAs. This structure can be exploited for improving predictive accuracy via post-prediction methods such as linear sum assignment. However, such methods do not represent improved generalization to OOD samples, and should be avoided.
E.3.2 Baseline results
For all experiments, we train the standard ResNet-50 model (He et al., 2016) pretrained on ImageNet, using a learning rate of and -regularization strength of . We trained these models with the Adam optimizer, using default parameter values and , with a batch size of 75 for 90 epochs, linearly increasing the learning rate for 10 epochs, then decreasing it following a cosine learning rate schedule. We selected hyperparameters by a grid search over learning rates , , , -regularization strengths , and numbers of warmup epochs . We report results aggregated over 3 random seeds.
Model performance dropped significantly going from the train-to-train in-distribution (ID) setting to the official out-of-distribution (OOD) setting (Table 7), with an average accuracy of 35.9% on the ID test set but only 29.9% on the OOD test set for ERM.
We ran an additional mixed-to-test comparison, where we moved half of the OOD test set into the training set, while keeping the overall amount of training data the same. Specifically, we moved one site per experiment from the OOD test set into the training set, and discarded an equivalent number of training sites, while leaving the validation set unchanged. While the test set in the mixed-to-test setting is effectively half as large as in the standard split, we expect it to be distributed similarly, since the two test set versions comprise the same 14 experiments.
Table 8 shows that there is a large gap between the OOD test accuracies in the official setting (29.9%) and the test accuracies in the mixed-to-test setting (39.8%). We note that the latter is higher than the train-to-train ID test accuracy of 35.9% reported in Table 7. This difference mainly stems from the slight difference in cell type composition between the test sets in the train-to-train and mixed-to-test settings; in particular, the train-to-train test set has a slightly higher proportion of the minority cell type (U2OS), on which performance is worse, and a slightly lower proportion of the majority cell type (HUVEC), on which performance is better. In this sense, the mixed-to-test result of 39.8% is a more accurate reflection of in-distribution performance on this dataset, and the results in Table 7 therefore understate the magnitude of the distribution shift.
We also trained models with CORAL, IRM, and group DRO, treating each experiment as a domain, and using the same model hyperparameters as ERM. However, the models trained using these methods all performed poorly compared to the ERM model (Table 7). One complication with these methods is that the experiments in the training set comprise different cell types, as mentioned above; this heterogeneity can pose a challenge to methods that treat each domain equivalently.
An important observation about batch effects in biological experiments: it is often the case that batch effects are mediated via biological mechanisms. For example, an increase in cellular media concentration may lead to cell growth and proliferation, while the upregulation of proliferation genes will do the same. Thus the “nuisance” factors associated with batch effects are often correlated with the biological signal we are attempting to observe, and cannot be disentangled from the biological factors that explain the data. Correction algorithms should take account of such trade-offs and attempt to optimize for both correction and signal preservation.
E.3.3 Broader context
As previously mentioned, high-throughput screening techniques are used broadly across many areas of biology, and therefore batch effects are a common problem in fields such as genomics, transcriptomics, proteomics, metabolomics, etc., so a particular solution in one such area may prove to be applicable in many areas of biology (Goh et al., 2017).
There are other datasets that are used in studying batch effects. The one most comparable to RxRx1-wilds is the BBBC021 dataset (Ljosa et al., 2012), which contains 13,200 3-channel fluorescent microscopy images of MCF7 cells acquired across 10 experimental batches. A subset of 103 treatments from 38 drug compounds belonging to 12 known mechanism of action (MoA) groups was first studied in Ando et al. (2017), and has been the subject of subsequent studies (Caicedo et al., 2018; Godinez et al., 2018; Tabak et al., 2020). Note that this dataset differs dramatically from RxRx1, in that there are fewer images, treatments, batches, and cell types, and each batch contains only a small subset of the total treatments.
E.3.4 Additional details
RxRx1-wilds contains two non-overlapping 256256 fields of view per well. Therefore, there could be as many as 125,664 images in the dataset (= 51 experiments 4 plates/experiment 308 wells/plate 2 images/well). 154 images were removed based on data quality, leaving a total dataset of 125,510 images.
The underlying raw dataset consists of 2048 2048 pixel, 6 channel, 16bpp images. To fit within the constraints of the Wilds benchmark, images for RxRx1-wilds were first downsampled to 1024 1024 and 8bpp, cropped to the center 256 256 pixels, and only the first three channels (nuclei, endoplasmic reticuli, actin) were retained. The original RxRx1 dataset, available at rxrx.ai and described in Taylor et al. (2019), provides 512 512 center crops of the downsampled images with all 6 channels retained.
The original RxRx1 dataset was also used in a NeurIPS 2019 competition hosted on Kaggle. The validation (OOD) and test (OOD) splits in RxRx1-wilds correspond to the public and private test sets from the Kaggle competition. The original RxRx1 dataset did not have an additional test (ID) split, and thus the original training split had both sites 1 and 2, for a total of 81,442 images. The Kaggle competition also aggregated predictions from both sites to form a single prediction per well, whereas in RxRx1-wilds, we treat each site separately.
As described in Section E.3.1, each plate in both the training and test sets contains the same 31 control conditions (one untreated well, and 30 positive control siRNAs). The Kaggle competition provided the labels for these control conditions in the test set, expecting that competitors would use them for various domain alignment techniques such as CORAL. However, these labels were instead used by the top competitors to bootstrap pseudo-labeling techniques. For RxRx1-wilds, for consistency with the other datasets and the typical domain generalization setting, we have opted not to release these control test labels for training.
The poor performance reported here on RxRx1-wilds may seem surprising in light of the fact that the top finishers of the Kaggle competition achieved near perfect accuracy on the test (OOD) set. This difference is due to a number of factors, including:
Adjustments made to the original RxRx1 dataset for RxRx1-wilds, as detailed in this subsection.
Differences in the network architectures used. To make training on RxRx1-wilds more accessible, we used a less compute-intensive architecture than typical in the competition.
Differences in training techniques used like pseudo-labeling (using the test control labels, as described above) and batch-level dataset augmentations or ensembling.
Differences in the way accuracy is measured. In the Kaggle competition, accuracy was measured for each well, meaning site-level predictions were aggregated to well-level predictions, and only for treatment classes, whereas in RxRx1-wilds, for convenience, accuracy is measured at each site and for both treatment and control classes.
The use of post-prediction methods like linear sum assignment that exploited the particular structure of the experiments in the RxRx1 dataset, as described under Potential Leverage in Section E.3.1.
E.4 OGB-MolPCBA
Accurate prediction of the biochemical properties of small molecules can significantly accelerate drug discovery by reducing the need for expensive lab experiments (Shoichet, 2004; Hughes et al., 2011). However, the experimental data available for training such models is limited compared to the extremely diverse and combinatorially large universe of candidate molecules that we would want to make predictions on (Bohacek et al., 1996; Sterling and Irwin, 2015; Lyu et al., 2019; McCloskey et al., 2020). This means that models need to generalize to out-of-distribution molecules that are structurally different from those seen in the training set.
We study this issue through the OGB-MolPCBA dataset, which is directly adopted from the Open Graph Benchmark (Hu et al., 2020b) and originally curated by MoleculeNet (Wu et al., 2018).
We consider the domain generalization setting, where the domains are molecular scaffolds, and our goal is to learn models that generalize to structurally distinct molecules with scaffolds that are not in the training set (Figure 6). This is a multi-task classification problem: for each molecule, we predict the presence or absence of 128 kinds of biological activities, such as binding to a particular enzyme. In addition, we cluster the molecules into different scaffold groups according to their two-dimensional structure, and annotate each molecule with the scaffold group that it belongs to. Concretely, the input is a molecular graph, the label is a 128-dimensional binary vector where each component corresponds to a biochemical assay result, and the domain specifies the scaffold. Not all biological activities are measured for each molecule, so can have missing values.
OGB-MolPCBA contains more than 400K small molecules with 128 kinds of prediction labels. Each small molecule is represented as a graph, where the nodes are atoms and the edges are chemical bonds. The molecules are pre-processed using RDKit (Landrum et al., 2006). Input node features are 9-dimensional, including atomic number, chirality, whether the atom is in the ring. Input edge features are 3-dimensional, including bond type and bond stereochemistry.
We split the dataset by scaffold structure. This scaffold split (Wu et al., 2018) is also used in the Open Graph Benchmark (Hu et al., 2020b). By attempting to separate structurally different molecules into different subsets, it provides a realistic estimate of model performance in prospective experimental settings. We assign the largest scaffolds to the training set to make it easier for algorithms to leverage scaffold information, and the smallest scaffolds to the test set to ensure that it is maximally diverse in scaffold structure:
Training: The largest 44,930 scaffolds, with an average of 7.8 molecules per scaffold.
Validation (OOD): The next largest 31,361 scaffolds, with an average of 1.4 molecules per scaffold.
Test (OOD): The smallest 43,793 scaffolds, which are all singletons.
In Figure 22 (A), we plot the statistics of the scaffolds in terms of the number of molecules belonging to each scaffold. We see that the scaffold sizes are highly skewed, with the test set containing (by design) the scaffolds with the least molecules. However, the differences in scaffold sizes do not significantly change the statistics of the molecules in each split. In Figures 22 (B) and (C), we see that the label statistics remain very similar across train/validation/test splits, suggesting that the main distribution shift comes from the difference in the input molecular graph structure.
We evaluate models by their average Average Precision (AP) across tasks (i.e., we compute the average precision for each task separately, and then average those scores), following Hu et al. (2020b). This accounts for the extremely skewed class balance in OGB-MolPCBA (only 1.4% of data is positive). Not all labels are available for each molecule; when calculating the AP for each task, we only consider the labeled molecules for the task.
We provide the scaffold grouping of molecules for training algorithms to leverage. Finding generalizable representations of molecules across different scaffold groups is useful for models to make accurate extrapolation on unseen scaffold groups. In fact, very recent work (Jin et al., 2020) has leveraged scaffold information of molecules to improve the extrapolation performance of molecular property predictors.
One notable characteristic of the scaffold group is that the size of each group is rather small; on the training split, each scaffold contains only 7.8 molecules on average. This also results in many scaffold groups: 44,930 groups in the training split. In Figure 22, we show that these scaffold groups are well-behaved in the sense that the train/validation/test splits contain contain similar ratios of positive labels as well as missing labels.
E.4.2 Baseline results
For all experiments, we use Graph Isomorphism Networks (GIN) (Xu et al., 2018) combined with virtual nodes (Gilmer et al., 2017), as this is currently the model with the highest performance in the Open Graph Benchmark (Hu et al., 2020b). We follow the same hyperparameters as in the Open Graph Benchmark: 5 GNN layers with a dimensionality of 300; the Adam optimizer (Kingma and Ba, 2015) with a learning rate of 0.001; and training for 100 epochs with early stopping. For each of the baseline algorithms (ERM, CORAL, IRM, and Group DRO), we separately tune the dropout rate from ; in addition, for CORAL and IRM, we tune the penalty weight as in Appendix D.
We first compare the generalization performance of ERM on the official scaffold split against the conventional random split, in which the entire molecules are randomly split into train/validation/test sets with the same split ratio as the scaffold split (i.e., 80/10/10). Results are in Table 10. The test performance of ERM drops by 7.2 points AP when the scaffold split is used, suggesting that the scaffold split is indeed harder than the random split.
To maintain consistency with the Open Graph Benchmark, and because the number of examples (molecules) per domain (scaffold) is relatively small compared to other datasets, we opted not to split off a portion of the training set into Validation (ID) and Test (ID) sets. We therefore do not run a train-to-train comparison for OGB-MolPCBA. Moreover, as the official scaffold split assigns the largest scaffolds to the training set and the smallest scaffolds to the test set, the test scaffolds all only have one molecule per scaffold, which precludes running test-to-test and mixed-to-test comparisons.
A potential issue with the random split ID comparison is that it does not measure performance on the same test distribution as the official split, and therefore might be confounded by differences in intrinsic difficulty. However, we believe that the random split setting provides a reasonable measure of ID performance for OGB-MolPCBA, as Figure 22 shows that the distribution of scaffolds assigned to the training versus test sets are similar. As the random split contains many singleton scaffolds in its test set that do not have corresponding molecules in the training set, we believe that it is likely to be an underestimate of the ID-OOD gap in OGB-MolPCBA.
Table 9 also shows that ERM performs better than CORAL, IRM, and Group DRO, all of which use scaffolds as the domains. For CORAL and IRM, we find that smaller penalties give better generalization performance, as larger penalty terms make the training insufficient. We use the penalty for CORAL and for IRM.
The primary issue with these existing methods is that they make the model significantly underfit the training data even when dropout is turned off. For instance, the training AP of CORAL and IRM is 20.0% and 15.9%, respectively, which are both lower than the 36.1% that ERM obtains even with 0.5 dropout. Also, these methods are primarily designed for the case when each group contains a decent number of examples, which is not the case for the OGB-MolPCBA dataset.
E.4.3 Broader context
Because of the very nature of discovering new molecules, out-of-distribution prediction is prevalent in nearly all applications of machine learning to chemistry domains. Beyond drug discovery, a variety of tasks and their associated datasets have been proposed for molecules of different sizes.
For small organic molecules, the scaffold split has been widely adopted to stress-test models’ capability for out-of-distribution generalization. While OGB-MolPCBA primarily focuses on predicting biophysical activity (e.g., protein binding), other datasets in MoleculeNet (Wu et al., 2018) include prediction of quantum mechanical properties (e.g., HOMO/LUMO), physical chemistry properties (e.g., water solubility), and physiological properties (e.g., toxicity prediction (Attene-Ramos et al., 2013)).
Besides small molecules, it is also of interest to apply machine learning over larger molecules such as catalysts and proteins. In the domain of catalysis, using machine learning to approximate expensive quantum chemistry simulation has gotten attention. The OC20 dataset has been recently introduced, containing 200+ million samples from quantum chemistry simulations relevant to the discovery of new catalysts for renewable energy storage and other energy applications (Becke, 2014; Chanussot et al., 2020; Zitnick et al., 2020). The OC20 dataset explicitly provides test sets with qualitatively different materials. In the domain of proteins, the recent trend is to use machine learning to predict 3D structure of proteins given their amino acid sequence information. This is known as the protein folding problem, and has sometimes been referred to as the Holy Grail of structural biology (Dill and MacCallum, 2012). CASP is a bi-annual competition to benchmark the progress of protein folding (Moult et al., 1995), and it evaluates predictions made on proteins whose 3D structures are identified very recently, presenting a natural temporal distribution shift. Recently, the AlphaFold2 deep learning model obtained breakthrough performance on the CASP challenge (Jumper et al., 2020), demonstrating exciting avenues of machine learning for structural biology.
E.4.4 Additional details
Data processing. The OGB-MolPCBA dataset contains 437,929 molecules annotated with 128 kinds of labels, each representing a bioassay curated in the PubChem database (Kim et al., 2016b). More details are provided in the MoleculeNet (Wu et al., 2018) and the Open Graph Benchmark (Hu et al., 2020b), from which the dataset is adopted.
E.5 GlobalWheat-wilds
Models for automated, high-throughput plant phenotyping—measuring the physical characteristics of plants and crops, such as wheat head density and counts—are important tools for crop breeding (Thorp et al., 2018; Reynolds et al., 2020) and agricultural field management (Shi et al., 2016). These models are typically trained on data collected in a limited number of regions, even for crops grown worldwide such as wheat (Madec et al., 2019; Xiong et al., 2019; Ubbens et al., 2020; Ayalew et al., 2020). However, there can be substantial variation between regions, due to differences in crop varieties, growing conditions, and data collection protocols. Prior work on wheat head detection has shown that this variation can significantly degrade model performance on regions unseen during training (David et al., 2020).
We study this shift in an expanded version of the Global Wheat Head Dataset (David et al., 2020, 2021), a large set of wheat images collected from 12 countries around the world.
We consider the domain generalization setting, where the goal is to learn models that generalize to images taken from new countries and acquisition sessions (Figure 7). The task is wheat head detection, which is a single-class object detection task. Concretely, the input is an overhead outdoor image of wheat plants, and the label is a set of bounding box coordinates that enclose the wheat heads (the spike at the top of the wheat plant containing grain), excluding the hair-like awns that may extend from the head. The domain specifies an acquisition session, which corresponds to a specific location, time, and sensor for which a set of images were collected. Our goal is to generalize to new acquisition sessions that are unseen during training. In particular, the dataset split captures a shift in location, with training and test sets comprising images from disjoint countries as discussed below.
The dataset comprises 6,515 images containing 275,187 wheat heads. These images were collected over 47 acquisition sessions in 16 research institutes across 12 countries. We describe the metadata and statistics of each acquisition session in Table 11.
Many factors contribute to the variation in wheat appearance across acquisition sessions. In particular, across locations, there is substantial variation due to differences in wheat genotypes, growing conditions (e.g., planting density), illumination protocols, and sensors. We study the effect of this location shift by splitting the dataset by country and assigning acquisition sessions from disjoint continents to the training and test splits:
Training: Images from 18 acquisition sessions in Europe (France 13, Norway 2, Switzerland, United Kingdom, Belgium), containing 131,864 wheat heads across 2,943 images.
Validation (OOD): Images from 7 acquisition sessions in Asia (Japan 4, China 3) and 1 acquisition session in Africa (Sudan), containing 44,873 wheat heads across 1,424 images.
Test (OOD): Images from 11 acquisition sessions in Australia and 10 acquisition sessions in North America (USA 6, Mexico 3, Canada), containing 66,905 wheat heads across 1,434 images.
Validation (ID): Images from the same 18 training acquisition sessions in Europe, containing 15,733 wheat heads across 357 images.
Test (ID): Images from the same 18 training acquisition sessions in Europe, containing 16,093 wheat heads across 357 images.
We evaluate models by first computing the average accuracy of bounding box detection within each image; then computing the average accuracy for each acquisition session by averaging its per-image accuracies; and finally averaging the accuracies of each acquisition session. The accuracy of a bounding box detection is measured at a fixed Intersection over Union (IoU) threshold of 0.5. The accuracy of an image is computed as , where is the number of true positives, which are ground-truth bounding boxes that can be matched with some predicted bounding box at IoU above the threshold; is the number of false negatives, which are ground-truth bounding boxes that cannot be matched as above; and is the number of false positives, which are predicted bounding boxes that cannot be matched with any ground-truth bounding box. We use accuracy rather than average precision, which is a common metric for object detection, because it was used in previous Global Wheat Challenges with the dataset (David et al., 2020, 2021). We use a permissive IoU threshold of 0.5 because there is some uncertainty regarding the precise outline of wheat head instances due to the stem and awns extending from the head. We measure the average accuracy across acquisition sessions because the number of images varies significantly across acquisition sessions, from 17 to 200 images in the test set, and we use average accuracy instead of worst-case accuracy because the wheat images are more difficult for some acquisition sessions.
The appearance of wheat heads in the images taken from different acquisition sessions can vary significantly, due to differences in the sensors used; illumination conditions, due to differences in illumination protocols, or the time of day and time of year that the images were taken; wheat genotypes; growth stages; growing conditions; and planting strategies. For example, different locations might feature a mix of different varieties of wheat (with different genotypes) with different appearances. Likewise, wheat planting strategies and growing conditions vary between regions and can contribute to differences between sessions, e.g., higher planting density may result in more closely packed plants and more occlusion between wheat head instances.
To provide leverage for models to learn to generalize across these conditions, we include images from 5 countries and 18 acquisition sessions in the training set. These training sessions cover all growth stages and include significant variation among all of the other factors. While the test domains include unseen conditions (e.g., sensors and genotypes not seen in the training set), our hope is that the variation in the training set will be sufficient to learn models that are robust to changes in these conditions.
E.5.2 Baseline results
For all experiments, we use the Faster-RCNN detection model (Ren et al., 2015), which has been successfully applied to the wheat head localization problem (Madec et al., 2019; David et al., 2020). To train, we fine-tune a model pre-trained with ImageNet, using a batch size of 4, a learning rate of , and weight decay of for 10 epochs with early stopping. The hyperparameters were chosen from a grid search over learning rates and weight decays . We report results aggregated over 3 random seeds.
We ran both train-to-train and mixed-to-test comparisons. For the train-to-train comparison, which uses the data splits described in the previous subsection, the Test (ID) accuracy is substantially higher than the Test (OOD) accuracy (77.1 (0.5) vs. 51.2 (1.8); Table 12). However, the Test (ID) and Test (OOD) sets come from entirely different regions, so this performance gap could also reflect a difference in the difficulty of the wheat head detection task in different regions (e.g., wheat heads that are more densely packed are harder to tell apart).
The mixed-to-test comparison controls for the test distribution by randomly splitting each test domain (acquisition session) into two halves, and then assigning one half to the training set. In other words, we randomly take out half of the test set and use it to replace existing examples in the training set, so that the total training set size is the same, and we retain the other half of the test set for evaluation. We also evaluated the ERM model trained on the official split on this subsampled test set. On this subsampled test set, the mixed-to-test ID accuracy is significantly higher than the OOD accuracy of the ERM model trained on the official split (63.3 (1.7) vs. 49.6 (1.9); Table 13).
We also compared the per-domain accuracies of the models trained in the mixed-to-test and official settings (Table 14) on the subsampled test set. The accuracy drop is not evenly distributed across each domain, though some of the domains have a relatively small number of images, so there is some variance across random replicates. The location/site of the acquisition session—which is correlated with factors like wheat genotype and the sensor used—has a large effect on performance (e.g., the KSU and Terraref sessions displayed a larger drop than the other sessions), but beyond that, it is not clear what factors are most strongly driving the accuracy drop. The Terraref sessions were particularly difficult even in the mixed-to-test setting, because of the strong contrast in its photos and the presence of hidden wheat heads under leaves. On the other hand, the KSU sessions had comparatively high accuracies in the mixed-to-test setting, but still displayed a large accuracy drop in the official OOD setting. As the KSU sessions differed primarily in their development stages and had largely similar ID and OOD accuracies, development stage does not seem to be a main driver of the accuracy drop. Finally, we note that the especially high variance across replicates for UQ_2 and UQ_3 is due to the proportion of empty images in those domains (88% for UQ_2 and 57% for UQ_3). Empty images are scored as either having 0% or 100% accuracy and therefore can have a large impact on the overall domain accuracy.
We also trained models with group DRO, treating each acquisition session as a domain, and using the same model hyperparameters as ERM. However, the group DRO models perform poorly compared to the ERM model as reported in Table 12. We leave the investigation of CORAL and IRM for future work because it is not straightforward to apply these algorithms to detection tasks.
Our baseline models were trained without any data augmentation, in contrast to baselines reported in the original dataset (David et al., 2020). Data augmentation could reduce the performance gap and warrants further investigation in future work, although David et al. (2020) still observed performance gaps on models trained with data augmentation in the original version of the dataset. Moreover, while we evaluated models by their average performance across acquisition sessions, we noticed a large variability in performance across domains. It is possible that some domains are more challenging or suffer from larger performance drops than others, and characterizing and mitigating these variations is interesting future work.
E.5.3 Broader context
Wheat head localization, while being an important operational trait for wheat breeders and farmers, is not the only deep learning application in plant phenotyping that suffers from lack of generalization. Other architectural traits such as plant segmentation (Sadeghi-Tehran et al., 2017; Kuznichov et al., 2019), plant and plant organ detection (Fan et al., 2018; Madec et al., 2019), leaves and organ disease classification (Fuentes et al., 2017; Shakoor et al., 2017; Toda and Okura, 2019), and biomass and yield prediction (Aich et al., 2018; Dreccer et al., 2019) would also benefit from plant phenotyping models that generalize to new deployments. In many of these applications, field images exhibit variations in illumination and sensors, and there has been work on mitigating biases across sensors (Ayalew et al., 2020; Gogoll et al., 2020). Finally, developing models that generalize across plant species would benefit the breeding and growing of specialized crops that are presently under-represented in plant phenotyping research worldwide (Ward and Moghadam, 2020). We hope that GlobalWheat-wilds can foster the development of general solutions to plant phenotyping problems, increase collaboration between plant scientists and computer vision scientists, and encourage the development of new multi-domain plant datasets to ensure that plant phenotyping results are generalizable to all crop growing regions of the world.
E.5.4 Additional details
The data is taken directly from the 2021 Global Wheat Challenge (David et al., 2021), which is an expanded version of the 2020 Global Wheat Challenge dataset (David et al., 2020). Compared to the challenge, the dataset splits are different: we split off part of the training set to form the Validation (ID) and Test (ID) sets, and we rearranged the Validation (OOD) and Test (OOD) sets so that they split along disjoint continents. Finally, we note that the 2021 challenge differs from the 2020 challenge in that images from North America were in the training set in the 2020 challenge, but were used for evaluation in the 2021 challenge, and are consequently assigned to the test set in GlobalWheat-wilds.
E.6 CivilComments-wilds
Automatic review of user-generated text is an important tool for moderating the sheer volume of text written on the Internet. We focus here on the task of detecting toxic comments. Prior work has shown that toxicity classifiers can pick up on biases in the training data and spuriously associate toxicity with the mention of certain demographics (Park et al., 2018; Dixon et al., 2018). These types of spurious correlations can significantly degrade model performance on particular subpopulations (Sagawa et al., 2020a).
We study this issue through a modified variant of the CivilComments dataset (Borkan et al., 2019b).
We cast CivilComments-wilds as a subpopulation shift problem, where the subpopulations correspond to different demographic identities, and our goal is to do well on all subpopulations (and not just on average across these subpopulations). Specifically, we focus on mitigating biases with respect to comments that mention particular demographic identities, and not comments written by members of those demographic identities; we discuss this distinction in the broader context section below.
The task is a binary classification task of determining if a comment is toxic. Concretely, the input is a comment on an online article (comprising one or more sentences of text) and the label is whether it is rated toxic or not. In CivilComments-wilds, unlike in most of the other datasets we consider, the domain annotation is a multi-dimensional binary vector, with the 8 dimensions corresponding to whether the comment mentions each of the 8 demographic identities male, female, LGBTQ, Christian, Muslim, other religions, Black, and White.
CivilComments-wilds comprises 450,000 comments, each annotated for toxicity and demographic mentions by multiple crowdworkers. We model toxicity classification as a binary task. Toxicity labels were obtained in the original dataset via crowdsourcing and majority vote, with each comment being reviewed by at least 10 crowdworkers. Annotations of demographic mentions were similarly obtained through crowdsourcing and majority vote.
Each comment was originally made on some online article. We randomly partitioned these articles into disjoint training, validation, and test splits, and then formed the corresponding datasets by taking all comments on the articles in those splits. This gives the following splits:
We evaluate a model by its worst-group accuracy, i.e., its lowest accuracy over groups of the test data that we define below.
As mentioned above, toxicity classifiers can spuriously latch onto mentions of particular demographic identities, resulting in a biased tendency to flag comments that innocuously mention certain demographic groups as toxic (Park et al., 2018; Dixon et al., 2018). To measure the extent of this bias, we define subpopulations based on whether they mention a particular demographic identity, compute the sensitivity (a.k.a. recall, or true positive rate) and specificity (a.k.a. true negative rate) of the classifier on each subpopulation, and then report the worst of these two metrics over all subpopulations of interest. This is equivalent to further dividing each subpopulation into two groups according to the label, and then computing the accuracy on each of these two groups.
Specifically, for each of the 8 identities we study (e.g., “male”), we form 2 groups based on the toxicity label (e.g., one group of comments that mention the male gender and are toxic, and another group that mentions the male gender and are not toxic), for a total of 16 groups. These groups overlap (a comment might mention multiple identities) and are not a complete partition (a comment might not mention any identity).
We then measure a model’s performance by its worst-group accuracy, i.e., its lowest accuracy over these 16 groups. A high worst-group accuracy (relative to average accuracy) implies that the model is not spuriously associating a demographic identity with toxicity. We can view this subpopulation shift problem as testing on multiple test distributions (corresponding to different subsets of the test set, based on demographic identities and the label) and reporting the worst performance over these different test distributions.
We use 16 groups (8 identities 2 labels) instead of just 8 groups (8 identities) to capture the desire to balance true positive and true negative rates across each of the demographic identities. Without splitting by the label, it would be possible for two different groups to have equal accuracies, but one group might be much more likely to have non-toxic comments flagged as toxic, whereas the other group might be much more likely to have toxic comments flagged as non-toxic. This would be undesirable from an application perspective, as such a model would still be biased against a particular demographic. In Appendix E.6.4, we further discuss the motivation for our choice of evaluation metric as well as its limitations.
As variability in performance over replicates can be high due to the small sizes of some demographic groups (Table 17), we report results averaged over 5 random seeds, instead of the 3 seeds that we use for most other datasets.
Since demographic identity annotations are provided at training time, we have an i.i.d. dataset available at training time for each of the test distributions of interest (corresponding to each group). Moreover, even though demographic identity annotations are unavailable at test time, they are relatively straightforward to predict.
E.6.2 Baseline results
For all experiments, we fine-tuned DistilBERT-base-uncased models (Sanh et al., 2019), using the implementation from Wolf et al. (2019) and with the following hyperparameter settings: batch size 16; learning rate using the AdamW optimizer (Loshchilov and Hutter, 2019) for 5 epochs with early stopping; an -regularization strength of ; and a maximum number of tokens of 300, since 99.95% of the input examples had 300 tokens. The learning rate was chosen through a grid search over , and all other hyperparameters were simply set to standard/default values.
The ERM model does well on average, with 92.2% average accuracy (Table 15). However, it does poorly on some subpopulations, e.g., with 57.4% accuracy on toxic comments that mention other religions. Overall, accuracy on toxic comments (which are a minority of the dataset) was lower than accuracy on non-toxic comments, so we also trained a reweighted model that balanced toxic and non-toxic comments by upsampling the toxic comments. This reweighted model had a slightly worse average accuracy of 89.8% and a better worst-group accuracy of 69.2% (Table 15, Reweighted (label)), but a significant gap remains between average and worst-group accuracies.
We note that the relatively small size of some of the demographic subpopulations makes it infeasible to run a test-to-test comparison, i.e., estimate how well a model could do on each subpopulation (corresponding to demographic identity) if it were trained on just that subpopulation. For example, Black comments comprise only <4% of the training data, and training just on those Black comments is insufficient to achieve high in-distribution accuracy. Without running the test-to-test comparison, it is possible that the gap between average and worst-group accuracies can be explained at least in part by differences in the intrinsic difficulty of some of the subpopulations, e.g., the labels of some subpopulations might be noisier because human annotators might disagree more frequently on comments mentioning a particular demographic identity. Future work will be required to establish estimates of in-distribution accuracies for each subpopulation that can account for these differences.
The CORAL, IRM, and group DRO baselines involve partitioning the training data into disjoint domains. We study the following partitions, corresponding to different rows in Table 15:
Label Black: 4 domains, 1 for each combination of class and Black.
On the Label partition, we used Group DRO to train a model that seeks to balance the losses on the positive and negative examples. This performs similarly to the standard reweighted models described above (Table 15, Group DRO (label)). We found that the worst-performing demographic for non-toxic comments was the Black demographic (Table 16), which motivated the Label Black partition. There, we used CORAL, IRM, and Group DRO to train models. However, these models did not perform significantly better (Table 15, label Black). While there were slight improvements on the Black groups, accuracy degraded on some other groups like non-toxic LBGTQ comments.
We note that our implementations of CORAL and IRM are built on top of the standard reweighting algorithm, i.e., they sample equally from each group. As these two algorithms perform similarly to reweighting, it indicates that the additional penalty term is not significantly affecting performance. Indeed, our grid search for the penalty weights selected the lowest value of the penalties ( for CORAL and for IRM).
Adapting the baseline methods to handle multiple overlapping groups, which were not studied in their original settings, could be a potential approach to improving accuracy on this task. Another potential approach is using baselining to account for different groups having different intrinsic levels of difficulty (Oren et al., 2019). For example, comments mentioning different demographic groups might differ in terms of how subjective classifying them is. Others have also explored specialized data augmentation techniques for mitigating demographic biases in toxicity classifiers (Zhao et al., 2018).
Adragna et al. (2020) recently used a simplified variant of the CivilComments dataset, with artificially-constructed training and test environments, to show a proof-of-concept that IRM can improve performance on minority groups. Methods such as IRM and group DRO rely heavily on the choice of groups/domains/environments; investigating the effect of different choices would be a useful direction for future work. Other recent work has studied methods that try to automatically learn groups, for example, through unsupervised clustering (Oren et al., 2019; Sohoni et al., 2020) or identifying high-loss points (Nam et al., 2020; Liu et al., 2021a).
Toxicity classification is one application where human moderators can work together with an ML model to handle examples that the model is unsure about. However, Jones et al. (2021) found that using selective classifiers—where the model is allowed to abstain if it is unsure—can actually further worsen performance on minority subpopulations. This suggests that in addition to having low accuracy on minority subpopulations, standard models can be poorly calibrated on them.
Another important consideration for toxicity detection in practice is shifts over time, as online discourse changes quickly, and what is seen as toxic today might not have even appeared in the dataset from a few months ago. We do not study this distribution shift in this work. One limitation of the CivilComments-wilds dataset is that it is fixed to a relatively short period in time, with most comments being written in the span of a year; this makes it harder to use as a dataset for studying temporal shifts.
Finally, we note that collecting “ground truth” human annotation of toxicity is itself a subjective and challenging process; recent work has studied ways of making it less biased and more efficient (Sap et al., 2019; Han and Tsvetkov, 2020).
E.6.3 Broader context
The CivilComments-wilds dataset does not assume that user demographics are available; instead, it uses mentions of different demographic identities in the actual comment text. For example, we want models that do not associate comments that mention being Black with being toxic, regardless of whether a Black or non-Black person wrote the comment. This setting is particularly relevant when user demographics are unavailable, e.g., when considering anonymous online comments.
A related and important setting is subpopulation shifts with respect to user demographics (e.g., the demographics of the author of the comment, regardless of the content of the comment). Such demographic disparities have been widely documented in natural language and speech processing tasks (Hovy and Spruit, 2016), among other areas. For example, NLP models have been shown to obtain worse performance on African-American Vernacular English compared to Standard American English on part-of-speech tagging (Jørgensen et al., 2015), dependency parsing (Blodgett et al., 2016), language identification (Blodgett and O’Connor, 2017), and auto-correct systems (Hashimoto et al., 2018). Similar disparities exist in speech, with state-of-the-art commercial systems obtaining higher word error rates on particular races (Koenecke et al., 2020) and dialects (Tatman, 2017).
These disparities are present not just in academic models, but in large-scale commercial systems that are already widely deployed, e.g., in speech-to-text systems from Amazon, Apple, Google, IBM, and Microsoft (Tatman, 2017; Koenecke et al., 2020) or language identification systems from IBM, Microsoft, and Twitter (Blodgett and O’Connor, 2017). Indeed, the original CivilComments dataset was developed by Google’s Conversation AI team, which is also behind a public toxicity classifier (Perspective API) that was developed in partnership with The New York Times (NYTimes, 2016).
E.6.4 Additional details
The evaluation metric used in the original competition was a complex weighted combination of various metrics, including subgroup AUCs for each demographic identity, and a new pinned AUC metric introduced by the original authors (Borkan et al., 2019b); conceptually, these metrics also measure the degree to which model accuracy is uniform across the different identities. After discussion with the original authors, we replace the composite metric with worst-group accuracy (i.e., worst TPR/FPR over identities) for simplicity. Measuring subgroup AUCs can be misleading in this context, because it assumes that the classifier can set separate thresholds for different subgroups (Borkan et al., 2019b, a).
One downside is that measuring worst-group accuracy treats false positives and false negatives equally. In deployment systems, one might want to weight these differently, e.g., using cost-sensitive learning or by simply raising or lowering the classification threshold, especially since real data is highly imbalanced (with a lot more negatives than positives). One could also binarize the labels and identities differently: in this benchmark, we simply use majority voting from the annotators.
Perhaps more fundamentally, even if TPR and FPR were balanced across different identities, this need not imply unambiguously equitable performance, because different subpopulations might have different intrinsic levels of noise and difficulty. See Corbett-Davies and Goel (2018) for more discussion of this problem of infra-marginality.
In practice, models might also do poorly on intersections of groups (Kearns et al., 2018), e.g., on comments that mention multiple identities. Given the size of the dataset and comparative rarity of some identities and of toxic comments in general, accuracies on these intersections are difficult to estimate from this dataset. A potential avenue of future work is to develop methods for evaluating models on such subgroups, e.g., by generating data in particular groups through templates (Park et al., 2018; Ribeiro et al., 2020).
The CivilComments-wilds dataset comprises comments from a large set of articles from the Civil Comments platform, annotated for toxicity and demographic identities (Borkan et al., 2019b). We partitioned the articles into disjoint training, validation, and test splits, and then formed the corresponding datasets by taking all comments on the articles in those splits. In total, the training set comprised 269,038 comments (60% of the data); the validation set comprised 45,180 comments (10%); and the test set comprised 133,782 (30%).
The original datasetwww.kaggle.com/c/jigsaw-unintended-bias-in-toxicity-classification/ also had a training and test split with disjoint articles. These splits are related to ours in the following way. Let the number of articles in the original test split be . To form our validation split, we took articles (sampled uniformly at random) from the original training split, and to form our test split, we took articles (also sampled uniformly at random) from the original training split and added it to the existing test split. We added a fixed validation set to allow other researchers to be able to compare methods more consistently, and we tripled the size of the test set to allow for more accurate worst-group accuracy measurement.
Similarly, we combined some of the demographic identities in the original dataset to obtain larger groups (for which we could more accurately estimate accuracy). Specifically, we created an aggregate LGBTQ identity that combines the original homosexual_gay_or_lesbian, bisexual, other_sexual_orientation, transgender, and other_gender identities (e.g., it is 1 if any of those identities are 1), and an aggregate other_religions identity that combines the original jewish, hindu, buddhist, atheist, and other_religion identities. We also omitted the psychiatric_or_mental_illness identity, which was evaluated in the original Kaggle competition, because of a lack of sufficient data for accurate estimation; but we note that baseline group accuracies for that identity seemed higher than for the other groups, so it is unlikely to factor into worst-group accuracy. In our new split, each identity we evaluate on (male, female, LGBTQ, Christian, Muslim, other_religions, Black, and White) has at least 500 positive and 500 negative examples. In Table 17 we show the sizes of each subpopulation in the test set; the training and validation sets follow similar proportions.
For convenience, we also add an identity_any identity; this combines all of the identities in the original dataset, including psychiatric_or_mental_illness and related identities.
We also trained a group DRO model using domains, 1 for each combination of class and the 8 identities. This model performed similarly to the other group DRO models.
All of the data, including the data with identity annotations that we use and the data with just label annotations, are also annotated for additional toxicity subtype attributes, specifically severe_toxicity, obscene, threat, insult, identity_attack, and sexual_explicit. These annotations can be used to train models that are more aware of the different ways that a comment can be toxic; in particular, using the identity_attack attribute to learn which comments are toxic because of the use of identities might help the model learn how to avoid spurious associations between toxicity and identity. These additional annotations are included in the metadata provided through the Wilds package.
The original CivilComments dataset (Borkan et al., 2019b) also contains 1.5M training examples that have toxicity (label) annotations but not identity (group) annotations. For simplicity, we have omitted these from the current version of CivilComments-wilds. These additional data points can be downloaded from the original data source and could be used, for example, by first inferring which group each additional point belongs to, and then running group DRO or a similar algorithm that uses group annotations at training time.
E.7 FMoW-wilds
ML models for satellite imagery can enable global-scale monitoring of sustainability and economic challenges, aiding policy and humanitarian efforts in applications such as deforestation tracking (Hansen et al., 2013), population density mapping (Tiecke et al., 2017), crop yield prediction (Wang et al., 2020b), and other economic tracking applications (Katona et al., 2018). As satellite data constantly changes due to human activity and environmental processes, these models must be robust to distribution shifts over time. Moreover, as there can be disparities in the data available between regions, these models should ideally have uniformly high accuracies instead of only doing well on data-rich regions and countries.
We study this problem on a variant of the Functional Map of the World dataset (Christie et al., 2018).
We consider a hybrid domain generalization and subpopulation shift problem, where the input is a RGB satellite image (resized to 224 224 pixels), the label is one of 62 building or land use categories, and the domain represents both the year the image was taken as well as its geographical region (Africa, the Americas, Oceania, Asia, or Europe). We aim to solve both a domain generalization problem across time and improve subpopulation performance across regions.
FMoW-wilds is based on the Functional Map of the World dataset (Christie et al., 2018), which collected and categorized high-resolution satellite images from over 200 countries based on the functional purpose of the buildings or land in the image, over the years 2002–2018 (see Figure 9). We use a subset of this data and split it into three time range domains, 2002–2013, 2013–2016, and 2016–2018, as well as five geographical regions as subpopulations (Africa, Americas, Oceania, Asia, and Europe). For each example, we also provide the timestamp and location coordinates, though our baseline models only use the coarse time ranges and geographical regions instead of these additional metadata.
Training: 76,863 images from the years 2002–2013.
Validation (OOD): 19,915 images from the years from 2013–2016.
Test (OOD): 22,108 images from the years from 2016–2018.
Validation (ID): 11,483 images from the years from 2002–2013.
Test (ID): 11,327 images from the years from 2002–2013.
The original dataset did not evaluate models under distribution shifts. Our training split is a subset of the original training dataset, filtered for images in the appropriate time range; similarly, our OOD and ID validation splits are subsets of the original validation dataset, and our OOD and ID test splits are subsets of the original test dataset. See Appendix E.7.4 for more dataset details.
The train/val/test data splits contain images from disjoint location coordinates, and all splits contain data from all 5 geographic regions. The ID and OOD splits within the test and validation sets may have overlapping locations, but have non-overlapping time ranges. There is a disparity in the number of examples in each region, with Africa and Oceania having the least examples (Figure 23); this could be due to bias in sampling and/or a lack of infrastructure and land data in certain regions.
We evaluate models by their average and worst-region OOD accuracies. The former measures the ability of the model to generalize across time, while the latter additionally measures how well models do across different regions/subpopulations under a time shift.
FMoW-wilds considers both domain generalization across time and subpopulation shift across regions. As we provide both time and region annotations, models can leverage the structure across both space and time to improve robustness. For example, one hypothesis is that infrastructure development occurs smoothly over time. Utilizing this gradual shift structure with the timestamp metadata may enable adaptation across longer time periods (Kumar et al., 2020). The data distribution may also shift smoothly over spatial locations, and so enforcing some consistency with respect to spatial structure may improve predictions (Rolf et al., 2020; Jean et al., 2018). Furthermore, to mitigate the fact that some regions (e.g., Africa) have less labeled data, one could potentially transfer knowledge of other regions with similar economies and infrastructure. The location coordinate metadata allows for transfer learning across similar locations at any spatial scale.
E.7.2 Baseline results
For all experiments, we follow Christie et al. (2018) and use a DenseNet-121 model (Huang et al., 2017) pretrained on ImageNet and with no regularization. We use the Adam optimizer (Kingma and Ba, 2015) with an initial learning rate of that decays by 0.96 per epoch, and train for 50 epochs for with early stopping and with a batch size of 64. All reported results are averaged over 3 random seeds.
In the train-to-train comparison, Table 20 shows that average accuracy drops by 6.7% when evaluated on the OOD test set () compared to the ID test set setting. The drop in average accuracy is especially large (11.6%) on images from the last year of the dataset (2017), furthest in the future from the training set. In addition, there is a substantial 26.0% drop in worst-region accuracy, with the model performing much worse in Africa than other regions (Table 19).
We also ran a mixed-to-test comparison where we mixed in some data from the OOD period (2013–2018) into the training set, while keeping the overall training set size constant. A model trained on this mixed split had a much smaller drop in performance under the time and region shifts (Table 20). While the magnitude of the ID-OOD gap in worst-region accuracy shrinks from 26.0% in the train-to-train setting to 16.3% in the mixed-to-test setting, the gap remains significant, implying that the drop in performance is largely due to the distribution shift across time and region instead of a change in the intrinsic difficulty of the OOD data.
We compare ERM against CORAL, IRM, and Group DRO, using examples from different years as distinct domains. Table 18 shows that many of these methods are comparable or worse than ERM in terms of both ID and OOD test performance. As with most other datasets, our grid search selected the lowest values of the penalty weights for CORAL () and IRM ().
Intriguingly, a large subpopulation shift across regions only occurs with a combination of time and region shift. This is corroborated by the mixed-split region shift results (Table 20), which do not have a time shift between training and test sets, and correspondingly do not display a large disparity in performance across regions. This drop in performance may be partially due to label shift: from Figure 24, we see that the label distributions between Africa and other regions are very different, e.g., with a large drop in recreational facilities and a sharp increase in single residential units. We do not find a similarly large label shift between and splits of the dataset.
Despite having the smallest number of training examples (Figure 23), the baseline models do not suffer a drop in performance in Oceania on validation or test sets (Table 19). We hypothesize that infrastructure in Oceania is more similar to regions with a large amount of data than Africa. In contrast, Africa may be more distinct and may have changed more drastically over 2002-2018, the time extent of the dataset. This suggests that the subpopulation shift is not merely a function of the number of training examples.
We note that our dataset splits can separate on particular factors such as the introduction of new sensors, which is natural with progression over time. For example, the WorldView-3 sensor came online in 2014. Future work should look into the role of auxiliary factors such as new sensors that are associated with time but may be controllable. We did not find a sharp difference in performance due to the introduction of WorldView-3; we found that the performance decays gradually over time, suggesting that the performance drop comes from other factors.
As with PovertyMap-wilds, there are important ethical considerations associated with remote sensing applications, e.g., around surveillance and privacy issues, as well as the potential for systematic biases that negatively affect particular populations. As an example of the latter, the poor model performance on satellite images from Africa that we observe in FMoW-wilds raises issues of bias and fairness. With regard to privacy, we note that the image resolution in FMoW-wilds is lower than that of other public and easily-accessible satellite data such as that from Google Maps. We refer interested readers to the UNICEF discussion paper by Berman et al. (2018) for a more in-depth discussion of the ethics of remote sensing especially as it pertains to development and humanitarian endeavors.
E.7.3 Broader context
Recognizing infrastructure and land features is crucial to many remote sensing applications. For example, in crop land prediction Wang et al. (2020b), recognizing gridded plot lines, plot circles, farm houses, and other visible features are important in recognizing crop fields. However, farming practices and equipment evolve over time and vary widely across the world, requiring both robust object recognition and synthesis of their different usage patterns.
Although the data is typically limited, we desire generalization on a global scale without requiring frequent large-scale efforts to gather more ground-truth data. It is natural to have labeled data with limited temporal or spatial extent since ground truth generally must be verified on the ground or requires manual annotations from domain experts (i.e., they are often hard to be crowdsourced). A number of existing remote sensing datasets have limited spatial or temporal scope, including the UC Merced Land Use Dataset (Yang and Newsam, 2010), TorontoCity (Wang et al., 2017), and SpaceNet (DigitalGlobe and Works, 2016). However, works based on these datasets generally do not systematically study shifts in time or location.
E.7.4 Additional details
The FMoW-wilds dataset is derived from Christie et al. (2018), which collected over 1 million satellite images from over 200 countries over 2002-2018. We use the RGB version of the original dataset, which contains 523,846 total examples, excluding the multispectral version of the images. Methods that can utilize a sequence of images can group the images from the same location across multiple years together as input, but we consider the simple formulation here for our baseline evaluation.
The original dataset from Christie et al. (2018) is provided as a set of hierarchical directories with JPEG images of varying sizes. To reduce download times and I/O usage, we resize these images to 224 224 pixels, and then store them as PNG images. We also collect all the metadata into CSV format for easy processing.
The original dataset is posed as a image time-series classification problem, where the model has access to a sequence of images at each location. For simplicity, we treat each image as a separate example, while making sure that the data splits all contain disjoint locations. We use the train/val/test splits from the original dataset, but separate out two OOD time segments: we treat the original validation data from 2013-2016 as OOD val and the original test data from 2016-2018 as OOD test. We remove data from after 2013 from the training set, which reduces the size of the training set in comparison to the original dataset.
Compared to PovertyMap-wilds, FMoW-wilds contains much higher resolution images (sub-meter resolution vs. 30m resolution) and contains a larger variety of viewpoints/tilts, both of which could present computational or algorithmic challenges. For computational purposes, we resized all images to (following Christie et al. (2018)), but raw images can be thousands of pixels wide. Some recent works have tried to balance this tradeoff between viewing overall context and the fine-grained detail (Uzkent and Ermon, 2020; Kim et al., 2016a), but how best to do this is an open question. FMoW-wilds also contains additional information on azimuth and cloud cover which could be used to correct for the variety in viewpoints and image quality.
E.8 PovertyMap-wilds
A different application of satellite imagery is poverty estimation across different spatial regions, which is essential for targeted humanitarian efforts in poor regions (Abelson et al., 2014; Espey et al., 2015). However, ground-truth measurements of poverty are lacking for much of the developing world, as field surveys are expensive (Blumenstock et al., 2015; Xie et al., 2016; Jean et al., 2016). For example, at least 4 years pass between nationally representative consumption or asset wealth surveys in the majority of African countries, with seven countries that had either never conducted a survey or had gaps of over a decade between surveys (Yeh et al., 2020). One approach to this problem is to train ML models on countries with ground truth labels and then deploy them to different countries where we have satellite data but no labels.
We study this problem through a variant of the poverty mapping dataset collected by Yeh et al. (2020).
We consider a hybrid domain generalization and subpopulation shift problem, where the input is a multispectral LandSat satellite image with 8 channels (resized to 224 224 pixels), the output is a real-valued asset wealth index computed from Demographic and Health Surveys (DHS) data, and the domain represents the country the image was taken in and whether the image is of an urban or rural area. We aim to solve both a domain generalization problem across country borders and improve subpopulation performance across urban and rural areas.
PovertyMap-wilds is based on a dataset collected by Yeh et al. (2020), which assembles satellite imagery and survey data at 19,669 villages from 23 African countries between 2009 and 2016 (Figure 10). Each input image has 8 channels: 7 from the LandSat satellite and an 8th channel for nighttime light intensity from a separate satellite, as prior work has established that these night lights correlate with poverty measures (Noor et al., 2008; Elvidge et al., 2009).
There are domains corresponding to the 23 countries and whether the location is urban or rural. Each example comes with metadata on its location coordinates, survey year, and its urban/rural classification.
In contrast to other datasets, which have a single fixed ID/OOD split, the relatively small size of PovertyMap-wilds allows us to use 5 different folds, where each fold defines a different set of OOD countries. In each fold, we use the following splits of the data (the number of countries and images in each split varies slightly from fold to fold):
Training: 10000 images from 13–14 countries.
Validation (OOD): 4000 images from 4–5 different countries (distinct from training and test (OOD) countries).
Test (OOD): 4000 images from 4–5 different countries (distinct from training and validation (OOD) countries).
Validation (ID): 1000 images from the same 13–14 countries in the training set.
Test (ID): 1000 images from the same 13–14 countries in the training set.
All splits contain images of both urban and rural locations, with the countries assigned randomly to each split in each fold.
The distribution of wealth may shift across countries due to differing levels economic development, agricultural practices, and other factors. For example, Abelson et al. (2014) use thatched vs. metal roofs to distinguish between poor and wealthy households, respectively in Kenya and Uganda. However, other countries may have a different mapping of roof type to wealth where metal roofs signify more poor households. Similar issues can arise when looking at the health of crops (related to vegetation indices such as NDVI that are simple functions of the multispectral channels in the satellite image) as a sign for wealth in rural areas, since crop health is related to climate and the choice of crops, which vary upon region.
Asset wealth may also shift dramatically between countries. Figure 25 shows the mean asset wealth per country, as well as urban vs. rural asset wealth per country. Mean asset wealth ranges from -0.4 to +0.8 depending on the country. There is a stark difference between mean asset wealth in urban and rural areas, with urban asset wealth being positive in all countries while rural mean asset wealth being mostly negative.
As is standard in the literature (Jean et al., 2016; Yeh et al., 2020), the models are evaluated on the Pearson correlation () between their predicted and actual asset wealth indices. We measure the average correlation, to test generalization under country shifts, and also the lower of the correlations on the urban and rural subpopulations, to test generalization between urban and rural subpopulations. We report the latter as previous works on poverty prediction from satellite imagery have noted that a significant part of model performance relies on distinguishing urban vs. rural areas, and improving performance within these subpopulations is an ongoing challenge, with rural areas generally faring worse under existing models (Jean et al., 2016; Yeh et al., 2020).
We average all correlations across the 5 different folds, using 1 random seed per fold. The resulting standard deviations reflect the fact that different folds have different levels of difficulty (e.g., depending on how similar the ID and OOD countries are). For the purposes of comparing different algorithms and models, we note that these standard deviations might make the comparisons appear noisier than they are, since a model might perform similarly across random seeds but still have a high standard deviation if it has different performances on different folds on the data. In contrast, other Wilds datasets report results on the same data split but averaged across different random seeds.
Large socioeconomic differences between countries makes generalization across borders challenging. However, some indicators of wealth are known to be robust and are able to be seen from space. For example, roof type (e.g. thatched or metal roofing) has been shown to be a reliable proxy for wealth (Abelson et al., 2014), and contextual factors such as the health of nearby croplands, the presence of paved roads, and connections to urban areas are plausibly reliable signals for measuring poverty. Poverty measures are also known to be highly correlated across space, meaning nearby villages will likely have similar poverty measures, and methods can utilize this spatial structure (using the provided location coordinate metadata) to improve predictions (Jean et al., 2018; Rolf et al., 2020). We show the correlation with distance in Figure 26, which plots the distance between pairs of data points against the absolute differences in asset wealth between pairs.
E.8.2 Baseline results
For all experiments, we follow Yeh et al. (2020) and train a ResNet-18 model (He et al., 2016) to minimize squared error. We use the Adam optimizer (Kingma and Ba, 2015) with an initial learning rate of that decays by 0.96 per epoch, and train for 200 epochs for with early stopping (on OOD ) and with a batch size of 64.
When shifting across country borders, Table 21 shows that ERM suffers a 0.04 drop in average in the official OOD setting compared to the train-to-train ID setting. Moreover, the drop in performance is exacerbated when looking at urban and rural subpopulations, even though all splits contain urban and rural examples; the difference in worst over the urban and rural subpopulations triples from 0.04 to 0.12 compared to the difference in average . Correlation is consistently lower on the rural subpopulation than the urban subpopulation.
We ran an additional mixed-to-test comparison where we considered an alternative training set with data that was uniformly sampled from all countries, while keeping the overall training set size constant (i.e., compared to the standard training set, it has fewer examples from each country, but data from more countries). A model trained on this mixed split had a much smaller drop in performance between the ID and OOD test sets (Table 22), which implies that the performance drop between the ID and OOD test sets is largely due to the distribution shift from seen to unseen countries.
We trained models with CORAL, IRM, and Group DRO, taking examples from different countries as coming from distinct domains. Table 21 shows that these baselines are generally comparable to ERM and that they continue to be susceptible to shifts across countries and urban/rural areas. As with most other datasets, our grid search selected the lowest values of the penalty weights for CORAL () and IRM ().
These results corroborate performance drops seen in previous out-of-country generalization tests for poverty prediction from satellite imagery (Jean et al., 2016). In general, differences in infrastructure, economic development, agricultural practices, and even cultural differences can cause large shifts across country borders. Differences between urban and rural subpopulations have also been well-documented (Jean et al., 2016; Yeh et al., 2020). Models based on nighttime light information could suffer more in rural areas where nighttime light intensity is uniformly low or even zero.
Since survey years are also available, we could also investigate the robustness of the model over time. This would enable the models to be used for a longer time before needing more updated survey data, and we leave this to future work. Yeh et al. (2020) investigated predicting the change in asset wealth for individual villages in the World Bank Living Standards Measurement Surveys (LSMS), which is a longitudinal study containing multiple samples from the same village. PovertyMap-wilds only contains cross-sectional samples which do not provide direct supervision for changes over time at any one location, but it is still possible to consider aggregate shifts across years.
As with FMoW-wilds, there are important ethical considerations associated with remote sensing applications, e.g., around surveillance and privacy issues, as well as the potential for systematic biases that negatively affect particular populations. As we describe in Section E.8.4, noise has been added to the location metadata in PovertyMap-wilds to protect privacy. The distribution shifts across country and urban/rural boundaries that we study in PovertyMap-wilds are an example of a bias that affects model performance and therefore could have adverse policy consequences. We refer interested readers to the UNICEF discussion paper by Berman et al. (2018) for a more in-depth discussion of the ethics of remote sensing especially as it pertains to development and humanitarian endeavors.
E.8.3 Broader context
Computational sustainability applications in the developing world also include tracking child mortality (Burke et al., 2016; Osgood-Zimmerman et al., 2018; Reiner et al., 2018), educational attainment (Graetz et al., 2018), and food security and crop yield prediction (You et al., 2017; Wang et al., 2020b; Xie et al., 2020). Remote sensing data and satellite imagery has the potential to enable high-resolution maps of many of these sustainability challenges, but as with poverty measures, ground truth labels in these applications come from expensive surveys or observations from human workers in the field. Some prior works consider using spatial structure (Jean et al., 2018; Rolf et al., 2020), unlabeled data (Xie et al., 2016; Jean et al., 2018; Xie et al., 2020), or weak sources of supervision (Wang et al., 2020b) to improve global models despite the lack of ground-truth data. We hope that PovertyMap-wilds can be used to improve the robustness of machine learning techniques on satellite data, providing an avenue for cheaper and faster measurements that can be used to make progress on a general set of computational sustainability challenges.
E.8.4 Additional details
The PovertyMap-wilds dataset is derived from Yeh et al. (2020), which gathers LandSat imagery and Demographic and Health Surveys (DHS) data from 19669 villages across 23 countries in Africa . The images are pixels large over 7 multispectral channels and an eighth nighttime light intensity channel. The LandSat satellite has a 30m resolution, meaning that each pixel of the image covers a spatial area. The location metadata is perturbed by the DHS as a privacy protection scheme; urban locations are randomly displaced by up to 2km and rural locations are perturbed by up to 10km. While this adds noise to the data, having a large enough image can guarantee that the location is in the image most of the time. The target is a real-valued composite asset wealth index computed as the first principal component of survey responses about household assets, which is thought to be a less noisy measure of households’ longer-run economic well-being than other welfare measurements like consumption expenditure (Sahn and Stifel, 2003; Filmer and Scott, 2011). Asset wealth also has the advantage of not requiring adjustments for inflation or for purchasing power parity (PPP), as it is not based on a currency.
We normalize each channel by the pixel-wise mean and standard deviation for each channel, following (Yeh et al., 2020). We also do a similar data augmentation scheme, adding random horizontal and vertical flips as well as color jitter (brightness factor 0.8, contrast factor 0.8, saturation factor 0.8, hue factor 0.1).
The data download process provided by Yeh et al. (2020) involves downloading and processing imagery from Google Earth Engine. We process each image into a compressed NumPy array with 8 channels. We also provide all the metadata in a CSV format.
We also ran an ablation where we removed the nighttime light intensity channel. This resulted in a drop in OOD of 0.04 on average and 0.06 on the rural subpopulation, demonstrating the usefulness of the nightlight data in asset wealth estimation.
We report a much larger drop in correlation due to spatial shift than in Yeh et al. (2020). To explain this, we note that our data splitting method is slightly different from theirs. They have two separate experiments (with different data splits) to test in-distribution vs. out-of-distribution generalization. In contrast, our data splits on both held-out in-distribution and out-of-distribution points at the same time with respect to the same training set, thus allowing us to compare both metrics simultaneously on one model as a more direct comparison. We use the same OOD country folds as the original dataset. However, Yeh et al. (2020) split the ID train/val/test while making sure that the spatial extent of the images between each split never overlap, while we simply take uniformly random splits of the ID data. This means that between our ID train/val/test splits, we may have images that have share some overlapping spatial extent, for example for two very nearby locations. Thus, a model can utilize some memorization here to improve ID performance. We believe this is reasonable since, with more ID data, more of the spatial area will be labeled and memorization should become an increasingly viable strategy for generalization in-domain.
E.9 Amazon-wilds
In many consumer-facing ML applications, models are trained on data collected on one set of users and then deployed across a wide range of potentially new users. These models can perform well on average but poorly on some individuals (Tatman, 2017; Caldas et al., 2018; Li et al., 2019b; Koenecke et al., 2020). These large performance disparities across users are practical concerns in consumer-facing applications, and they can also indicate that models are exploiting biases or spurious correlations in the data (Badgeley et al., 2019; Geva et al., 2019). We study this issue of inter-individual performance disparities on a variant of the Amazon-wilds Reviews dataset (Ni et al., 2019).
We consider a hybrid domain generalization and subpopulation problem where the domains correspond to different reviewers. The task is multi-class sentiment classification, where the input is the text of a review, the label is a corresponding star rating from 1 to 5, and the domain is the identifier of the reviewer who wrote the review. Our goal is to perform consistently well across a wide range of reviewers, i.e., to achieve high tail performance on different subpopulations of reviewers in addition to high average performance. In addition, we consider disjoint set of reviewers between training and test time.
The dataset comprises 539,502 customer reviews on Amazon taken from the Amazon Reviews dataset (Ni et al., 2019). Each input example has a maximum token length of 512. For each example, the following additional metadata is also available at both training and evaluation time: reviewer ID, product ID, product category, review time, and summary.
To reliably measure model performance on each reviewer, we include at least 75 reviews per reviewer in each split. Concretely, we consider the following splits, where reviewers are randomly assigned to either in-distribution or out-of-distribution sets:
Training: 245,502 reviews from 1,252 reviewers.
Validation (OOD): 100,050 reviews from another set of 1,334 reviewers, distinct from training and test (OOD).
Test (OOD): 100,050 reviews from another set of 1,334 reviewers, distinct from training and validation (OOD).
Validation (ID): 46,950 reviews from 626 of the 1,252 reviewers in the training set.
Test (ID): 46,950 reviews from 626 of the 1,252 reviewers in the training set.
The reviewers in the train and in-distribution splits; the validation (OOD) split; and the test (OOD) split are all disjoint, which allows us to test generalization to unseen reviewers. See Appendix E.9.4 for more details.
To assess whether models perform consistently well across reviewers, we evaluate models by their accuracy on the reviewer at the 10th percentile. This follows the federated learning literature, where it is standard to measure model performance on devices and users at various percentiles in an effort to encourage good performance across many devices (Caldas et al., 2018; Li et al., 2019b).
We include more than a thousand reviewers in the training set, capturing variation across a wide range of reviewers. In addition, we provide reviewer ID annotations for all reviews in the dataset. These annotations could be used to directly mitigate performance disparities across reviewers seen during training time.
E.9.2 Baseline results
For all experiments, we finetuned DistilBERT-base-uncased models (Sanh et al., 2019), using the implementation from Wolf et al. (2019), and with the following hyperparameter settings: batch size 8; learning rate with the AdamW optimizer (Loshchilov and Hutter, 2019); -regularization strength ; 3 epochs with early stopping; and a maximum number of tokens of 512. We selected the above hyperparameters based on a grid search over learning rates , and all other hyperparameters were simply set to standard/default values.
A DistilBERT-base-uncased model trained with the standard ERM objective performs well on average, but performance varies widely across reviewers (Figure 27, Table 23). Despite the high average accuracy of 71.9%, per-reviewer accuracies vary widely between 100.0% and 12.0%, with accuracy at the 10th percentile of 53.8%. The above variation is larger than expected from randomness: a random binomial baseline with equal average accuracy would have a 10th percentile accuracy of 65.4%. We observe low tail performance on both previously seen and unseen reviewers, with low 10th percentile accuracy on in-distribution and out-of-distribution sets (Table 23). In addition, we observe drops on both average and 10th percentile accuracies upon evaluating on unseen reviewers, as evident in the performance gaps between the in-distribution and the out-of-distribution sets.
As with CivilComments-wilds, the relatively small number of reviews per reviewer makes it difficult to run a test-to-test comparison (e.g., training a model on just the reviewers in the bottom 10th percentile). Without running the test-to-test comparison, it is possible that the gap between average and 10th percentile accuracies can be explained at least in part by differences in the intrinsic difficulty of reviews from different reviewers, e.g., some reviewers might not write text reviews that are informative of their star rating. Future work will be required to establish in-distribution accuracies that account for these differences.
We now consider models trained by existing robust training algorithms and show that these models also perform poorly on tail reviewers, failing to mitigate the performance drop (Table 23). We observe that reweighting to achieve uniform class balance fails to improve the 10th percentile accuracy, showing that variation across users cannot be solved simply by accounting for label imbalance. In addition, CORAL, IRM, and Group DRO fail to improve both average and 10th percentile accuracies on both ID and OOD sets. Our grid search selected for the CORAL penalty and for the IRM penalty.
The distribution shift and the evaluation criteria for Amazon-wilds focus on the tail performance, unlike the other datasets in Wilds. Because of this, Amazon-wilds might have distinct empirical trends or be conducive to different algorithms compared to other datasets. Potential approaches include extensions to algorithms for worst-group performance, for example to handle a large number of groups, as well as adaptive approaches that yield user-specific predictions.
E.9.3 Broader context
Performance disparities across individuals have been observed in a wide range of tasks and applications, including in natural language processing (Geva et al., 2019), automatic speech recognition (Koenecke et al., 2020; Tatman, 2017), federated learning (Li et al., 2019b; Caldas et al., 2018), and medical imaging (Badgeley et al., 2019). These performance gaps are practical limitations in applications that call for good performance across a wide range of users, including many user-facing applications such as speech recognition (Koenecke et al., 2020; Tatman, 2017) and personalized recommender systems (Patro et al., 2020), tools used for analysis of individuals such as sentiment classification in computational social science (West et al., 2014) and user analytics (Lau et al., 2014), and applications in federated learning. These performance disparities have also been studied in the context of algorithmic fairness, including in the federated learning literature, in which uniform performance across individuals is cast as a goal toward fairness (Li et al., 2019b; Dwork et al., 2012). Lastly, these performance disparities can also highlight models’ failures to learn the actual task in a generalizable manner; instead, some models have been shown learn the biases specific to individuals. Prior work has shown that individuals—technicians for medical imaging in this case—can not only be identified from data, but also are predictive of the diagnosis, highlighting the risk of learning to classify technicians rather than the medical condition (Badgeley et al., 2019). More directly, across a few natural language processing tasks where examples are annotated by crowdworkers, models have been observed to perform well on annotators that are commonly seen at training time, but fail to generalize to unseen annotators, suggesting that models are merely learning annotator-specific patterns and not the task (Geva et al., 2019).
E.9.4 Additional details
We consider a modified version of the Amazon reviews dataset (Ni et al., 2019). We consider disjoint reviewers between the training, OOD validation, and OOD test sets, and we also provide separate ID validation and test sets that include reviewers seen during training for additional reporting. These reviewers are selected uniformly at random from the reviewer pool, with the constraint that they have at least 150 reviews in the pre-processed dataset. Statistics for each split are described in Table 24. Notably, each reviewer has at least 75 reviews in the training set and exactly 75 reviews in the validation and test sets.
To process the data, we first eliminate reviews that are longer than 512 tokens, reviews without any text, and any duplicate reviews with identical star rating, reviewer ID, product ID, and time. We then obtain the 30-core subset of the reviews, which contains the maximal set of reviewers and products such that each reviewer and product has at least 30 reviews; this is a standard preprocessing procedure used in the original dataset (Ni et al., 2019). To construct the dataset for reviewer shifts in particular, we further eliminate the following reviews: (i) reviews that contain HTML, (ii) reviews with identical text within a user in order to ensure sufficiently high effective sample size per reviewer, and (iii) reviews with identical text across users to eliminate generic reviews. Once we have the filtered set of reviews, we consider reviewers with at least 150 reviews and sample uniformly at random until the training set contains approximately 250,000 reviews and each evaluation set contains at least 100,000 reviews. As we construct the training set, we reserve a random sample of 75 reviews for each user for evaluation and put all other reviews in the training set. For the evaluation set, we put a random sample of 75 reviews for each user.
The original dataset does not prescribe a specific task or split. We consider a standard task of sentiment classification, but instead of using a standard i.i.d. split, we instead consider disjoint users between training and evaluation time as described above. In addition, we preprocess the data as detailed above.
E.10 Py150-wilds
Code completion models—autocomplete tools used by programmers to suggest subsequent source code tokens, such as the names of API calls—are commonly used to reduce the effort of software development (Robbes and Lanza, 2008; Bruch et al., 2009; Nguyen and Nguyen, 2015; Proksch et al., 2015; Franks et al., 2015). These models are typically trained on data collected from existing codebases but then deployed more generally across other codebases, which may have different distributions of API usages (Nita and Notkin, 2010; Proksch et al., 2016; Allamanis and Brockschmidt, 2017). This shift across codebases can cause substantial performance drops in code completion models. Moreover, prior studies of real-world usage of code completion models have noted that these models can generalize poorly on some important subpopulations of tokens such as method names (Hellendoorn et al., 2019).
We study this problem using a variant of the Py150 Dataset, originally developed by Raychev et al. (2016) and adapted to a code completion task by Lu et al. (2021).
We consider a hybrid domain generalization and subpopulation shift problem, where the domains are codebases (GitHub repositories), and our goal is to learn code completion models that generalize to source code written in new codebases. Concretely, the input is a sequence of source code tokens taken from a single file, the label is the next token (e.g., "environ", "communicate" in Figure 12), and the domain is an integer that identifies the repository that the source code belongs to. We aim to solve both a domain generalization problem across codebases and improve subpopulation performance on class and methods tokens.
The dataset comprises 150,000 Python source code files from 8,421 different repositories on GitHub (github.com). Each source code file is associated with the repository ID so that code from the same repository can be linked.
We split the dataset by randomly partitioning the data by repositories:
Training: 79,866 code files from 5,477 repositories.
Validation (OOD): 5,160 code files from different 261 repositories.
Test (OOD): 39,974 code files from different 2,471 repositories.
Validation (ID): 5,000 code files from the same repositories as the training set (but different files).
Test (ID): 20,000 code files from the same repositories as the training set (but different files).
The repositories are randomly distributed across the training, validation (OOD), and test (OOD) sets. As we use models pre-trained on the CodeSearchNet dataset (Husain et al., 2019), which partially overlaps with the Py150 dataset, we ensured that all GitHub repositories used in CodeSearchNet only appear in the training set in Py150-wilds and not in the validation/test sets.
Table 25 shows the token statistics of the source code files, as well as the token type breakdown (e.g., class, method, punctuator, keyword, literal). The tokens are defined by the built-in Python tokenizer and the CodeGPT tokenizer, following Lu et al. (2021). Training and evaluation are conducted at the token-level (more details are provided below).
We evaluate models by their accuracy on predicting class and method tokens in the test set code files. This subpopulation metric is inspired by Hellendoorn et al. (2019), which finds that in real-world settings, developers primarily use code completion tools for completing class names and method names; in contrast, measuring average token accuracy would prioritize common tokens such as punctuators, which are often not a problem in real-world settings.
We provide the GitHub repository that each source code files was derived from, which training algorithms can leverage. As programming tools like code completion are expected to be used across codebases in real applications (Nita and Notkin, 2010; Allamanis and Brockschmidt, 2017), it is important for models to learn generalizable representations of code and extrapolate well on unseen codebases. We hope that approaches using the provided repository annotations can learn to factor out common features and codebase-specific features, resulting in more robust models.
Additionally, besides the (integer) IDs of repositories, we also provide the repository names and file names in natural language as extra metadata. While we only use the repository IDs in our baseline experiments described below, the extra natural language annotations can potentially be leveraged as well to adapt models to target repositories/files.
E.10.2 Baseline results
For all experiments, we use the CodeGPT model (Lu et al., 2021) pre-trained on CodeSearchNet (Husain et al., 2019) as our model and finetune it on Py150-wilds, using all the tokens in the training set. We tokenize input source code by the CodeGPT tokenizer and take blocks of length 256 tokens. We then train the CodeGPT model with a batch size of 6 (with tokens), a learning rate of , no regularization, and the AdamW optimizer (Loshchilov and Hutter, 2019) for 3 epochs with early stopping. Using the hyperparameters from Lu et al. (2021) as a starting point, we selected the above hyperparameters by a grid search over learning rates and regularization strength . All other hyperparameters were simply set to standard/default values.
Table 26 shows that model performance on class and method tokens dropped substantially from 75.4% on the train-to-train in-distribution repositories in the Test (ID) set to 67.9% on the out-of-distribution repositories in the Test (OOD) set. This gap shrinks if we evaluate the model on all tokens (instead of class and method tokens): accuracy drops from 74.5% on Test (ID) to 69.6% on Test (OOD). This is because the evaluation across all tokens includes many tokens that are used universally across repositories, such as punctuators and keywords.
We only ran a train-to-train comparison because there are a relatively large number of domains (repositories) split i.i.d. between the training and test sets, which suggests that the training and test sets should be “equally difficult”. We therefore do not expect test-to-test and mixed-to-test comparisons to yield significantly different results.
We trained CORAL, IRM, and Group DRO baselines, treating each repository as a domain. For CORAL and IRM, we find that the smaller penalties give slightly better generalization performance ( for CORAL and for IRM). Compared to the ERM baseline, while CORAL and IRM reduced the performance gap between ID and OOD, neither of them improved upon ERM on the final OOD performance.
E.10.3 Broader context
Machine learning can aid programming and software engineering in various ways: automatic code completion (Raychev et al., 2014; Svyatkovskiy et al., 2019), program synthesis (Bunel et al., 2018; Kulal et al., 2019), program repair (Vasic et al., 2019; Yasunaga and Liang, 2020), code search (Husain et al., 2019), and code summarization (Allamanis et al., 2015). However, these systems face several forms of distribution shifts when deployed in practice. One major challenge is the shifts across codebases (which our Py150-wilds dataset focuses on), where systems need to adapt to factors such as project content, coding conventions, or library or API usage in each codebase (Nita and Notkin, 2010; Allamanis and Brockschmidt, 2017). A second source of shifts is programming languages, which includes adaptation across different domain-specific languages (DSLs), e.g., in robotic environments (Shin et al., 2019); and across different versions of languages, e.g., Python 2 and 3 (Malloy and Power, 2017). Another challenge is the shift from synthetic training sets to real usage: for instance, Hellendoorn et al. (2019) show that existing code completion systems, which are typically trained as language models on source code, perform poorly on the real completion instances that are most commonly used by developers in IDEs, such as API calls (class and method calls).
E.10.4 Additional details
We generate the splits in the following steps. First, to avoid test set contamination, we took all of the repositories in CodeSearchNet (which, as a reminder, is used to pretrain our baseline model) and assigned them to the training set. Second, we randomly split all of the remaining repositories into three groups: Validation (OOD), Test (OOD), and Others. Finally, to generate the ID splits, we randomly split the files in the Others repositories into three sets: Training, Validation (ID), and Test (ID).
The original Py150 dataset (Raychev et al., 2016) splits the total 150k files into 100k training files and 50k test files, regardless of the repository that each file was from. In Py150-wilds, we re-split the dataset based on repositories to construct the aforementioned train, validation (ID), validation (OOD), test (ID), and test (OOD) sets.
Additionally, in the Py150 code completion task introduced in Lu et al. (2021), models are evaluated by the accuracy of predicting every token in source code. However, according to developer studies, this evaluation may include various tokens that are rarely used in real code completion, such as punctuators, strings, numerals, etc. (Robbes and Lanza, 2008; Proksch et al., 2016; Hellendoorn et al., 2019). To define a task closer to real applications, in Py150-wilds we focus on class name and method name prediction (which are used most commonly by developers).
F Datasets with distribution shifts that do not cause performance drops
In this section, we provide more details on the stop-and-frisk dataset discussed in Section 8.1. The original data was provided by the New York City Police Department, and has been widely used in previous ML and data analysis work (Goel et al., 2016; Zafar et al., 2017; Pierson et al., 2018; Kallus and Zhou, 2018; Srivastava et al., 2020). For our analysis, we use the version of the dataset that was processed by Goel et al. (2016). Our problem setting and dataset structure closely follow theirs.
We study a subpopulation shift in a weapons prediction task, where each data point corresponds to a pedestrian who was stopped by the police on suspicion of criminal possession of a weapon. The input is a vector that represents 29 observable features from the UF-250 stop-and-frisk form filled out by the officer after each stop: e.g., whether the stop was initiated based on a radio run or at an officer’s discretion, whether the officer was uniformed, and any reasons the officer gave for the stop (encoded as a categorical variable). Importantly, these features can all be observed by the officer prior to making the stop.When we consider subpopulation shifts over race groups, the input additionally includes 75 one-hot indicators corresponding to the precinct that the stop was made in. We do not include those features when we consider shifts over locations, as they prevent the model from generalizing to new locations. The binary label is whether the pedestrian in fact possessed a weapon (i.e., whether the stop fulfilled its stated purpose). We consider, separately, two types of domains : 1) race groups and 2) locations (boroughs in New York City). We consider location and race as our domains because previous work has shown that they can produce substantial disparities in policing practices and in algorithmic performance (Goel et al., 2016).
Each row of the dataset represents one stop of one pedestrian. Following Goel et al. (2016), we filter for the 621,696 stops where the reason for the stop is suspicion of criminal possession of a weapon. We then filter for rows with complete data for observable features; with stopped pedestrians who are Black, white, or Hispanic; and who are stopped during the years 2009-2012 (the time range used in Goel et al. (2016)). These filters yield a total of 506,283 stops, 3.5% of which are positive examples (in which the officer finds that the pedestrian is illegally possessing a weapon).
The training versus validation split is a random 80%-20% partition of all stops in 2009 and 2010. We test on stops from 2011-2012; this follows the experimental setup in Goel et al. (2016). Overall, our data splits are as follows:
Training: 241,964 stops from 2009 and 2010.
Validation: 60,492 stops from 2009 and 2010, disjoint from the training set.
In the experiments below, we do not use the entire training set, as we observed in our initial experiments that the model performed less well on certain subgroups (Black pedestrians and pedestrians from the Bronx). To determine whether this inferior performance might be ameliorated by training specifically on those groups, we controlled for training set size by downsampling the training set to the size of the disadvantaged population of interest for a given split. Specifically, we consider the following (overlapping) training subsets, each of which is subsampled from the overall training set described above:
Black pedestrians only: 155,929 stops of Black pedestrians from 2009 and 2010.
All pedestrians, subsampled to # Black pedestrians: 155,929 stops of all pedestrians from 2009 and 2010.
Bronx pedestrians only: 69,129 stops of pedestrians in the Bronx from 2009 and 2010.
All pedestrians, subsampled to # Bronx pedestrians: 69,129 stops of all pedestrians from 2009 and 2010.
These amount to running a test-to-test comparison for the subpopulations of Black pedestrians and Bronx pedestrians.
Our metric for classifier performance is the precision for each race group and each borough at a global recall of 60%—i.e., when using a threshold which recovers 60% of all weapons in the test data, similar to the recall evaluated in Goel et al. (2016). The results are similar when using different recall thresholds. Examining the precision for each race/borough captures the fact, discussed in Goel et al. (2016), that very low-precision stops may violate the Fourth Amendment, which requires reasonable suspicion for conducting a police stop; thus, the metric encapsulates the intuition that the police are attempting to avoid Fourth Amendment violations for any race group or borough while still recovering a substantial fraction of the illegal weapons.
F.1.2 Baseline results
For all experiments, we use a logistic regression model trained with the Adam optimizer (Kingma and Ba, 2015) and early stopping. We trained one model on each of the 4 training sets, separately picking hyperparameters through a grid search across 7 learning rates logarithmically-spaced in and batch sizes in . Table 29 provides the hyperparameters used for each training set. All models were trained with a reweighted cross-entropy objective that upsampled the positive examples to achieve class balance.
Performance differed substantially across race and location groups: precision was lowest on Black pedestrians (Table 27, top row) and pedestrians in the Bronx (Table 28, top row). To assess whether in-distribution training would improve performance on these groups, we trained the model only on Black pedestrians (Table 27, bottom row) and pedestrians in the Bronx (Table 28, bottom row). However, this did not substantially improve performance on Black pedestrians or pedestrians from the Bronx; the difference in precision was less than 0.005 for both groups relative to the original model trained on all races and locations. This is consistent with the fact that groups with the lowest performance are not necessarily small minorities of the dataset: for example, more than 90% of the stops are of Black or Hispanic pedestrians, but performance on these groups is worse than that for white pedestrians. The lack of improvement from in-distribution training suggests that approaches like group DRO would be unlikely to further improve performance, and we thus did not assess these approaches.
We observed large disparities in performance across race and location groups. However, the fact that test-to-test in-distribution training did not ameliorate these disparities suggests that they do not occur because some groups comprise small minorities of the original dataset, and thus suffer worse performance. Instead, our results suggest that classification performance on some race and location groups are intrinsically noisier; it is possible, for example, that collection of additional features would be necessary to improve performance on these groups (Chen et al., 2018).
F.1.3 Additional details
The features we use are very similar to those used in Goel et al. (2016). The two primary differences are that 1) we remove features which convey information about a stopped pedestrian’s race, since those might be illegal to use in real-world policing contexts and 2) we do not include a “local hit rate” feature which captures the fraction of historical stops in the vicinity of a stop which resulted in discovery of a weapon; we omit this latter feature because it was unnecessary to match performance in Goel et al. (2016). test-to-test
F.2 ENCODE: Transcription factor binding across different cell types
Here we provide details on the transcription factor binding dataset discussed in Section 8.3. Transcription factors (TFs) are regulatory proteins that bind specific DNA elements in the genome to activate or repress transcription of target genes. There are estimated to be approximately 1,600 human TFs, and the binding landscape of each TF can be highly variable across different cell types (Deplancke et al., 2016). Understanding how these binding patterns change across different cell types and affect cellular function is critical for understanding the mechanics of dynamic gene regulation across cell types and across healthy and diseased cell states.
Several experimental strategies have been developed to profile genome-wide binding landscapes of individual TFs in specific cell types of interest. However, genome-wide profiling of TF binding is challenging in practice, as it requires large numbers of cells and reagents (e.g., high-affinity antibodies) that are difficult and expensive to acquire. Moreover, profiling each individual TF requires a separate experiment, so it can be prohibitively costly to map out even a few different TFs out of the >1000 in the human genome. Therefore, there has been wide interest in computational approaches that can predict the genome-wide binding maps of multiple TFs in new cell types from a single and more practical genome-wide assay.
DNA sequence is one of the principal determinants of where a TF binds along the genome,Most TFs, including the ones we provide in this benchmark, have DNA-binding domains which bind to sequence motifs: short recognition sequences (4-20 bases in length) in the genome with specific binding affinity distributions (Stormo and Zhao, 2010). and many ML models have been developed to predict TF binding as a function of DNA sequence in a particular cell type (Alipanahi et al., 2015; Quang and Xie, 2019; Avsec et al., 2021b). However, even when the DNA sequence is invariant across different cell types (e.g., among cell types from the same organism), the TF binding landscape can still be highly variable (Deplancke et al., 2016). Therefore, TF binding models that only use sequence inputs cannot make different predictions for the same sequence across different cell types; we also need complementary, cell-type-specific inputs to model changes in binding over different cell types.
In this section, we explore the use of genome-wide chromatin accessibility assays such as DNase-seq and ATAC-seq (Boyle et al., 2008; Thurman et al., 2012; Buenrostro et al., 2013), in conjunction with DNA sequence, to predict TF binding. DNA is typically accessible in a highly local and cell-type-specific manner, and in particular, genomic sequences with high accessibility are typically bound by one or more TFs, although the identity of the TF is not directly measured by the experiment (Lee et al., 2004). By measuring chromatin accessibility at each base in the genome in a specific cell type of interest, we can obtain a cell-type-specific profile of binding locations; moreover, these experiments are often cheaper than profiling even a single TF (Minnoye et al., 2021). Our goal is to use this accessibility signal, combined with DNA sequence, to accurately predict the binding patterns of multiple TFs in new cell types.
We study the problem of predicting genome-wide TF binding across different cell types using data from the ENCODE-DREAM Transcription Factor Binding Site Prediction Challenge (Balsubramani et al., 2020).
We consider the domain generalization setting, where the domains are cell types, and we seek to learn models that can generalize to cell types that are not in the training set. The task is to predict if a particular transcription factor (TF) would bind to a particular genomic location in a cell type of interest (Figure 28). The input is DNA sequence (which we assume to be shared across all cell types) and a cell-type-specific biochemical measurement of chromatin accessibility obtained through the DNase-seq assay.
Concretely, we segment the genome into uniformly-sized, overlapping bins that are 200 base pairs (bp) in length, and tiled 50bp apart. Given a TF , each genomic bin in cell type has a binding status . Our goal is to predict each bin’s binding status as a function of the local DNA sequence and the local cell-type specific accessibility profile (Figure 28). We treat each TF separately, i.e., for each , we have separate training and test sets and separate models.
The dataset comprises (a) genome-wide sequence; (b) TF binding maps for two TFs, JUND and MAX, across a total of 6 and 8 cell types respectively; and (c) an accessibility profile for each cell type. As described above and illustrated in Figure 28, we break up these genome-wide data into overlapping 12800 bp windows, which each correspond to a single training example. The central 6400 bp of each 12800 bp window is tiled with overlapping 200 bp bins that each correspond to one coordinate of the corresponding . These 12800 bp windows are tiled 6400 bp apart, such that each genomic location falls within the central 6400 bp region of exactly one window.
We split the examples by domain (cell type) as well as by chromosome (a large contiguous subsequence of the genome) within a cell type. In each split, we use one cell type for the test data, one for the validation data, and the remaining cell types for the training data. These domain-wise splits are listed in Table 30, and are divided into two types:
The ENCODE-DREAM splits follow the original challenge setup (Balsubramani et al., 2020) closely in evaluating only on the cell type liver, which is a primary tissue, in contrast to all of the other cell types, which are immortalized cell lines that have been grown outside the body for many generations. This is a more realistic setting in the sense that it is easier to collect data from immortalized cell lines, which we can then use to train a model that predicts TF binding in harder-to-profile primary tissues. However, the fact that none of the training cell types are primary tissues might limit generalization to primary tissues. Moreover, because cell types are highly variable, conclusions drawn from a single liver cell type might not generalize to other cell types.
We thus also use a round-robin set of splits, where we assign each cell type to test and validation sets in a rotating manner. This round-robin evaluation comprises several splits for each TF.
For each split in Table 30, the data are divided into training, validation, and test sets by chromosome:
Training: 323,894 windows per training cell type, before filtering. To improve class balance, we filter out windows with all 128 bins labeled negative, which typically removes over 3/4 of these training windows; the exact number filtered out depends on the split. The training windows are taken from all chromosomes except .
Validation (OOD): 27,051 windows from 1 validation cell type and from chromosomes .
Test (OOD): 23,109 windows from 1 test cell type and from chromosomes .
Validation (ID): 27,051 windows in total across all training cell types, from chromosomes .
Test (ID): 23,109 windows in total across all training cell types, from chromosomes .
For computational speed, the Validation (OOD) and Test (OOD) sets above were subsampled by a factor of 3 from the available raw data, while the Validation (ID) and Test (ID) sets were subsampled by a factor of 3 the number of training cell types.
We evaluate models by their average precision (AP) in predicting binary binding status (excluding all ambiguous bins). Specifically, we treat each bin as a separate binary classification problem; in other words, we split up each prediction of the 128-dimensional vector into at most 128 separate binary predictions after excluding ambiguous bins. We then compute the average precision of the model on this binary classification problem.
The choice of average precision as an evaluation metric is motivated by the class imbalance (low proportion of bound/positive labels) of this binary classification problem over bins. All splits have more than one hundred times as many unbound bins than bound bins (Table 31).
F.2.2 Baseline results
Our model is a version of the fully convolutional U-Net model for image segmentation (Ronneberger et al., 2015), modified from the architecture in Li and Guan (2019). It is illustrated in Figure 29. We train each model using the average cross-entropy loss over the 128 output bins in each example (after excluding ambiguous bins).
For hyperparameters, we searched over the learning rates , and -regularization strengths . We use 10 replicates (with different random seeds) for all reported results.
Table 32 and Table 34 show the results of ERM models trained on each split. On many individual splits (Table 34), the OOD validation and OOD test performance are very different, reflecting the variability across the cell types. On average across the round robin splits, the OOD validation performance is slightly higher than the OOD test performance, as we selected hyperparameters and did early stopping to maximize the former (Table 32). We also observed high variance in training and test performance across random seeds in a few splits (e.g., the K562 / liver split for the transcription factor JUND), which suggests some optimization instability in our training protocol. We also computed the in-distribution baselines in a train-to-train setting, i.e., on the Validation (ID) and Test (ID) splits described above.
We also ran corresponding in-distribution baselines in a test-to-test setting, i.e., we trained ERM models on data from the training chromosomes in the test cell type, and tested it on the same test set comprising data from the test chromosomes in the test cell type. Prior work has shown that there is minimal variation in performance between chromosomes on this problem (Won et al., 2010; Alipanahi et al., 2015; Keilwagen et al., 2019), so we can approximate these training and test distributions as identical. Table 33 shows these in-distribution results. For the round-robin splits, the difference between the train-to-train and test-to-test settings is that the former trains and tests on mixtures of multiple cell types, whereas the latter trains and tests on individual cell types.
We considered two TFs, MAX and JUND, separately. For the ENCODE-DREAM splits, both TFs showed large ID-OOD performance gaps. However, we opted not to use the ENCODE-DREAM split as a Wilds dataset because the variability between cell types made us cautious about over-interpreting the results on a single cell type. For example, for MAX, we found that the Validation (OOD) and Test (OOD) cell types were so different that their results were anti-correlated across different random seeds, which would have made benchmarking challenging. Moreover, the fact that the test cell type (liver) in the ENCODE-DREAM splits was the only primary tissue might have meant that the training data might have insufficient leverage for a model to learn to close the ID-OOD gap.
We therefore focused on analyzing the round-robin splits. For MAX, using the train-to-train comparison, the average Test (ID) and Test (OOD) AP across the round-robin splits were not significantly different (64.9 (2.1) vs. 59.6 (2.0), respectively; Table 32). For JUND, using the train-to-train comparison, the average Test (ID) and Test (OOD) AP across the round-robin splits showed a larger gap (54.1 (4.2) vs. 42.9 (3.2), respectively; Table 32), but the variability in training performance made these results less reliable. Moreover, the test-to-test ID results were significantly higher than the train-to-train ID results (62.4 (2.6) vs. 54.1 (4.2), respectively), which suggests that either the model capacity or feature set is not rich enough to fit the variation across different cell types. We therefore opted not to include the round-robin splits in Wilds as well.
Even in the in-distribution (test-to-test) setting, the results in Table 35 show how different model performance can be for different domains. For example, the liver domain of primary tissue (from a human donor) is derived from lower-quality data than many of the long-standard cell lines (grown outside the human body) constituting other domains, and is consequently noisier than many of them (Balsubramani et al., 2017). The extent of this variation underscores the importance of accounting for the variability between domains when measuring the effect of distribution shift; for example, a train-to-train comparison could lead to significantly different conclusions than a test-to-test comparison.
The effect of the distribution shift also seems to depend on the particular TF (MAX vs. JUND) used. Biologically, different TFs show different levels of cell-type-specificity, and better understanding which TFs have binding patterns that can be accurately predicted from the cell-type-specific accessibility assays is important future work.
One of the main obstacles preventing us from using this ENCODE dataset as a Wilds benchmark is the instability in optimization that we reported above. We speculate that this instability could, in part, be due to the class imbalance in the data, but more work will be needed to ascertain this and to develop methods for training models more reliably on this type of genomic data.
As we mentioned above, Table 32 reports significantly higher test-to-test ID results than train-to-train ID results for the JUND round-robin split scheme. The main difference between the test-to-test and train-to-train settings in the round-robin splits is that the former trains and tests on a single cell type, whereas the latter trains and tests on a mixture of cell types. The fact that ID performance is significantly higher in the former than the latter suggests that the learned models are not able to fit JUND binding patterns across multiple cell types. This could be due to a model family that is not large or expressive enough, or it could be because the feature set does not have all of the necessary information to accurately predict binding across cell types. In either case, it is unlikely that a training algorithm developed to be robust to distribution shifts will be able to significantly improve OOD performance in this setting, as the issue seems to lie in the model family or the data distribution instead.
Overall, it is commonly understood that distribution shifts between cell types are a significant problem for TF binding prediction, and many methods have been developed to tackle these shifts (Balsubramani et al., 2017; Li et al., 2019a; Li and Guan, 2019; Keilwagen et al., 2019; Quang and Xie, 2019). Nonetheless, we found it challenging to establish a rigorous distribution shift benchmark around this task, as our results were confounded by factors such as optimization issues, large variability between cell types, and the difficulty of learning a model that could fit multiple cell types even in an i.i.d. setting. We hope that future work on evaluating and mitigating distribution shifts in TF binding prediction can build upon our results and address these challenges.
F.2.3 Additional details
The ground-truth labels were derived from high-quality chromatin immunoprecipitation sequencing (ChIP-seq) experiments, which provide a genome-wide track of binding enrichment scores for each TF. Statistical methods based on standardized pipelines (Landt et al., 2012) were used to identify high-confidence binding events across the genome, resulting in a genome-wide track indicating whether each of the windows of sequence in the genome is bound or unbound by the TF, or whether binding is ambiguous but likely (these were ignored in our benchmarking).
Our data include two TFs chosen for their basic importance in cell-type-specific gene regulation: MAX and JUND. MAX canonically recognizes a short, common sequence (the domain CACGTG), but its structure leads it to bind to DNA as a dimer, and facilitates cooperative activity with a range of partners (Grandori et al., 2000) with many weaker and longer-range sequence determinants of binding (Allevato et al., 2017). JUND belongs to a large family of TFs (bZIP) known for binding in cooperation with partners in the family in a variety of modes, all involving a short 7bp sequence (TGA[C/G]TCA) and its two halves.
The network consists of encoder and decoder portions:
Encoder. The encoder is composed of five downscaling convolutional blocks, each consisting of two stride-1 convolutional layers with kernel size 7 (and padding such that the output size is left unchanged), followed by a max-pooling layer with kernel size 2. Each successive block halves the input window size and scales up the number of convolutional filters (by 1.5).
Decoder. Mirroring the encoder, the decoder is composed of five upscaling convolutional blocks, each consisting of two convolutional layers with kernel size 7 and an upsampling layer (a ConvTranspose layer with kernel size 2 and stride 2). Each successive block doubles the input window size. The respective sizes of the decoder layer representations are the same as the encoder in reverse, culminating in a representation that is then run through a convolutional layer (kernel size 200, stride 50) to reduce it to a single channel (with length 253). A final fully-connected layer results in a 128-dimensional output.
Batch normalization is applied after every layer except the last, and each intermediate convolutional layer is padded such that the output and input sizes are equal.
The ENCODE-DREAM prediction challenge contains binding data for many TFs from a large range of cell types, discretized into the same 200-bp windows used in this benchmark. The ENCODE portal (encodeproject.org) contains more ChIP-seq datasets from the 13 challenge cell lines for which we provide DNase accessibility data. DNA shape and gene expression data types were also provided in the original challenge.
DNA shape. Twisting, bending, and shearing of DNA influence local binding in a TF-specific fashion (Rohs et al., 2009).
Gene expression. Expression levels of all human genes were provided using RNA-seq data from ENCODE. This can be used to model the presence of cofactor proteins that can recruit TFs for binding (Ptashne and Gann, 1997).
However, none of the top challenge participants found these data modalities useful (Balsubramani et al., 2017), so they are not provided in this benchmark.
We normalize the distribution of each DNase-seq signal readout to the average of the DNase-seq signals over training cell types. We use a version of quantile normalization (Bolstad et al., 2003) with piecewise polynomial interpolation. Li and Guan (2019) also use this, but instead normalize to the test domain’s DNase distribution. As this technique uses test-domain data, it is out of the scope of our benchmark. However, we note that in genomics settings it is realistic to have relatively cheaply available chromatin accessibility data in the target cell type of interest.
The prediction task of the challenge was a binary classification problem over the 200 bp bins, which did not involve the fixed 12800 bp windows. To predict on a bin, participating teams were free to use as much of the regions surrounding (flanking) the bin as they wished. The winning teams all used at least 1000 bp total for each bin, and further work has shown the efficacy of using much larger flanking regions of tens of thousands of bp (Quang and Xie, 2019; Avsec et al., 2021a). We instead predict on 128 bins at once (following Li and Guan (2019)), which allows for more efficient training and prediction.
Our ERM baselines’ OOD test performance is competitive with the original challenge results, but lower than the state-of-the-art performance of Li and Guan (2019) because of the aforementioned differences in data processing, splits, and architecture, as well as the cross-domain training method employed by that paper and predecessor work (Li et al., 2019a). These and other state-of-the-art models noted that their domain adaptation strategies played a major role in improving performance.
F.3 BDD100K: Object recognition in autonomous driving across locations
As discussed in Section 8.6, autonomous driving, and robotics in general, is an important application that requires effective and robust tools for handling distribution shift. Here, we discuss our findings on a modified version of the BDD100K dataset that evaluates on shifts based on time of day and location. Our results below suggest that more challenging tasks, such as object detection and segmentation, may be more suited to evaluations of distribution shifts in an autonomous driving context.
In line with the other datasets in Wilds, we evaluate using a classification task. Specifically, the task is to predict whether or not 9 different categories appear in the image : bicycles, buses, cars, motorcycles, pedestrians, riders, traffic lights, traffic signs, and trucks. This is a multi-task binary classification problem, and the label is thus a 9-dimensional binary vector.
The BDD100K dataset is a large and diverse driving dataset crowd-sourced from tens of thousands of drivers, covering four different geographic regions and many different times of day, weather conditions, and scenes (Yu et al., 2020). The original dataset contains 80,000 images in the combined training and validation sets and is richly annotated for a number of different tasks such as detection, segmentation, and imitation learning. We use bounding box labels to construct our task labels, and as discussed later, we use location and image tags to construct the shifts we evaluate.
In evaluating the trained models, we consider average accuracy across the binary classification tasks, averaged over each of the validation and test sets separately. We next discuss how we create and evaluate two different types of shift based on time of day and location differences.
F.3.2 Time of day shift
We evaluate two different types of shift, depicted in Figure 30. For time of day shift (Figure 30 top row), we use the original BDD100K training set, which has roughly equal proportions of daytime and non daytime images (Yu et al., 2020). However, we construct a test set using the original BDD100K validation set that only includes non-daytime images. We then split roughly the same number of images randomly from the training set to form an in-distribution validation set, which allows us to do a train-to-train comparison. There are 64,993, 4,860, and 4,742 images in the training, validation, and test splits, respectively.
Table 36 summarizes our findings. For time of day shift, we actually observe slightly higher test performance, on only non daytime images, than validation performance on mixed daytime and non daytime images. We contrast this with findings from Dai and Van Gool (2018); Yu et al. (2020), who showed worse test performance for segmentation and detection tasks, respectively, on non daytime images. We believe this disparity can be attributed to the difference in tasks—for example, it is likely more difficult to draw an accurate bounding box for a car at night than to simply recognize tail lights and detect the presence of a car.
F.3.3 Location shift
For location shift (Figure 30 bottom row), we combine all of the data from the original BDD100K training and validation sets. We construct training and validation sets from all of the images captured in New York, and we use all images from California for the test set. The validation set again is in-distribution with respect to the training set and has roughly the same number of images as the test set. There are 53,277, 9,834, and 9,477 images in the training, validation, and test splits, respectively.
In the case of location shift, we see from Table 36 that there is a small drop in performance, possibly because this shift is more drastic as the locations are disjoint between training and test time. However, the performance drop is relatively small and the test time accuracy is still comparable to validation accuracy. In general, we believe that these results lend support to the conclusion that, for autonomous driving and robotics applications, other more challenging tasks are better suited for evaluating performance. Generally speaking, incorporating a wide array of different applications will likely require a simultaneous effort to incorporate different tasks as well.
F.4 Amazon: Sentiment classification across different categories and time
Our benchmark dataset Amazon-wilds studies user shifts. In Section 7, we discussed empirical trends on other types of distribution shifts on the same underlying 2018 Amazon Reviews dataset (Ni et al., 2019). We now present the detailed setup and empirical results for the time and category shifts.
For all experiments in this section, we finetune BERT-base-uncased models, using the implementation from Wolf et al. (2019), and with the following hyperparameter settings: batch size 8; learning rate ; -regularization strength ; 3 epochs; and a maximum number of tokens of 512. These hyperparameters are taken from the Amazon-wilds experiments.
F.4.2 Time shifts
We consider the domain generalization setting, where the domain is the year in which the reviews are written. As in Amazon-wilds, the task is multi-class sentiment classification, where the input is the text of a review, the label is a corresponding star rating from 1 to 5.
The dataset is a modified version of the Amazon Reviews dataset (Ni et al., 2019) and comprises customer reviews on Amazon. Specifically, we consider the following split:
Training: 1,000,000 reviews written in years 2000 to 2013.
Validation (OOD): 20,000 reviews written in years 2014 to 2018.
Test (OOD): 20,000 reviews written in years 2014 to 2018.
To construct the above split, we first randomly sample 4,000 reviews per year for the evaluation splits. For years in which there are not sufficient reviews, we split the reviews equally between validation and test. After constructing the evaluation set, we then randomly sample from the remaining reviews to form the training set.
To assess whether models generalize to future years, we evaluate models by their average accuracy on the OOD test set.
We only observed modest performance drops due to time shift. Our baseline model performs well on the OOD test set, achieving 76.0% accuracy on average and 75.4% on the worst year (Table 37). To measure performance drops due to distribution shifts, we ran a test-to-test comparison by training a model on reviews written in years 2014 to 2018 (Table 38). The performance gaps between the model trained on the official split and the model trained on the test-to-test split are consistent but modest across the years, with the biggest drop of 1.1% for 2018.
F.4.3 Category shifts
Shifts across categories—where a model is trained on reviews in one category and then tested on another—have been studied extensively (Blitzer et al., 2007; Mansour et al., 2009; Hendrycks et al., 2020c). In line with prior work, we observe that model performance drops upon evaluating on a few unseen categories. However, the observed difference between out-of-distribution and in-distribution baselines varies from category to category and is not consistently large (Hendrycks et al., 2020c). In addition, we find that training on more diverse data with more product categories tends to improve generalization to unseen categories and reduce the effect of the distribution shift; similar phenomena have also been reported in prior work (Mansour et al., 2009; Guo et al., 2018).
We consider the domain generalization setting, where the domain is the product category. As in Amazon-wilds, the task is multi-class sentiment classification, where the input is the text of a review, the label is a corresponding star rating from 1 to 5.
The dataset is a modified version of the Amazon Reviews dataset (Ni et al., 2019) and comprises customer reviews on Amazon. Specifically, we consider the following split for a given set of training categories:
Training: up to 1,000,000 reviews in training categories.
Validation (OOD): reviews in categories unseen during training.
Test (OOD): reviews in categories unseen during training.
Validation (ID): reviews in training categories.
Test (ID): reviews in training categories.
To construct the above split, we first randomly sample 1,000 reviews per category for the evaluation splits (for categories with insufficient number of reviews, we split the reviews equally between validation and test) and then randomly sample from the remaining reviews to form the training set.
To assess whether models generalize to unseen categories, we evaluate models by their average accuracy on each of the categories in the OOD test set.
We first considered training on four categories (Books, Movies and TV, Home and Kitchen, and Electronics) and evaluating on unseen categories. We observed that a BERT-base-uncased model trained via ERM yields a test accuracy of 75.4% on the four in-distribution categories and a wide range of accuracies on unseen categories (Table 39, columns Multiple). While the accuracies on some unseen categories are lower than the train-to-train in-distribution accuracy, it is unclear whether the performance gaps stem from the distribution shift or differences in intrinsic difficulty across categories; in fact, the accuracy is higher on many unseen categories (e.g., All Beauty) than on the in-distribution categories, illustrating the importance of accounting for intrinsic difficulty.
To control for intrinsic difficulty, we ran a test-to-test comparison on each target category. We controlled for the number of training reviews to the extent possible; the standard model is trained on 1 million reviews in the official split, and each test-to-test model is trained on 1 million reviews or less, as limited by the number of reviews per category. We observed performance drops on some categories, for example on Clothing, Shoes, and Jewelry (83.0% in the test-to-test setting versus 75.2% in the official setting trained on the four different categories) and on Pet Supplies (78.8% to 76.8%). However, on the remaining categories, we observed more modest performance gaps, if at all. While we thus found no evidence for significance performance drops for many categories, these results do not rule out such drops either: one confounding factor is that some of the oracle models are trained on significantly smaller training sets and therefore underestimate the in-distribution performance.
In addition, we compared training on four categories (Books, Movies and TV, Home and Kitchen, and Electronics), as above, to training on just one category (Books), while keeping the training set size constant. We found that decreasing the number of training categories in this way lowered out-of-distribution performance: across many OOD categories, accuracies were modestly but consistently higher for the model trained on four categories than for the model trained on a single category (Table 39).
F.5 Yelp: Sentiment classification across different users and time
We present empirical results on time and user shifts in the Yelp Open Datasethttps://www.yelp.com/dataset.
For all experiments in this section, we finetune BERT-base-uncased models, using the implementation from Wolf et al. (2019), and with the following hyperparameter settings: batch size 8; learning rate ; -regularization strength ; 3 epochs with early stopping; and a maximum number of tokens of 512. We select the above hyperparameters based on a grid search over learning rates , using the time shift setup; for the user shifts, we adopted the same hyperparameters.
F.5.2 Time shifts
We consider the domain generalization setting, where the domain is the year in which the reviews are written. As in Amazon-wilds, the task is multi-class sentiment classification, where the input is the text of a review, the label is a corresponding star rating from 1 to 5.
The dataset is a modified version of the Yelp Open Dataset and comprises 1 million customer reviews on Yelp. Specifically, we consider the following split:
Training: 1,000,000 reviews written in years 2006 to 2013.
Validation (OOD): 20,000 reviews written in years 2014 to 2019.
Test (OOD): 20,000 reviews written in years 2014 to 2019.
To construct the above split, we first randomly sample 1,000 reviews per year for the evaluation splits. For years in which there are not sufficient reviews, we split the reviews equally between validation and test. After constructing the evaluation set, we then randomly sample from the remaining reviews to form the training set.
To assess whether models generalize to future years, we evaluate models by their average accuracy on the OOD test set.
We observe modest performance drops due to time shift. A BERT-base-uncased model trained with the standard ERM objective performs well on the OOD test set, achieving 76.0% accuracy on average and 73.9% on the worst year (Table 40). To measure performance drops due to distribution shifts, we run a test-to-test in-distribution comparison by training on reviews written in years 2014 to 2019 (Table 41). While there are consistent performance gaps between the out-of-distribution and the in-distribution baselines in later years, they are modest in magnitude with the largest drop of 3.1% for 2018.
F.5.3 User shift
As in Amazon-wilds, we consider the domain generalization setting, where the domains are reviewers and the task is multi-class sentiment classification. Concretely, the input is the text of a review, the label is a corresponding star rating from 1 to 5, and the domain is the identifier of the user that wrote the review.
The dataset is a modified version of the Yelp Open Dataset and comprises 1.2 million customer reviews on Yelp. To measure generalization to unseen reviewers, we train on reviews written by a set of reviewers and consider reviews written by unseen reviewers at test time. Specifically, we consider the following random split across reviewers:
Training: 1,000,104 reviews from 11,856 reviewers.
Validation (OOD): 40,000 reviews from another set of 1,600 reviewers, distinct from training and test (OOD).
Test (OOD): 40,000 reviews from another set 1,600 reviewers, distinct from training and validation (OOD).
Validation (ID): 40,000 reviews from 1,600 of the 11,856 reviewers in the training set.
Test (ID): 40,000 reviews from 1,600 of the 11,856 reviewers in the training set.
The training set includes at least 25 reviews per reviewer, whereas the evaluation sets include exactly 25 reviews per reviewer. While we primarily evaluate model performance on the above OOD test set, we also provide in-distribution validation and test sets for potential use in hyperparameter tuning and additional reporting. These in-distribution splits comprise reviews written by reviewers in the training set.
To assess whether models perform consistently well across reviewers, we evaluate models by their accuracy on the reviewer at the 10th percentile.
We observe only modest variations in performance across reviewers. A BERT-base-uncased model trained with the standard ERM objective achieves 71.5% accuracy on average and 56.0% accuracy at the 10th percentile reviewer (Table 42). The above variation is modestly larger than expected from randomness; a random binomial baseline with equal average accuracy would have a tenth percentile accuracy of 60.1%.