A Fine-Grained Analysis on Distribution Shift

1 INTRODUCTION . If machine learning models are to be ubiquitous in critical applications such as driverless cars ( Janai et al. , 2020 ) , medical imaging ( Erickson et al. , 2017 ) , and science ( Jumper et al. , 2021 ) , it is pivotal to build models that are robust to distribution shifts . Otherwise , models may fail surprisingly in ways that derail trust in the system . For example , Koh et al . ( 2020 ) ; Perone et al . ( 2019 ) ; AlBadawy et al . ( 2018 ) ; Heaven ( 2020 ) ; Castro et al . ( 2020 ) find that a model trained on images from one set of hospitals may not generalise to the imaging conditions of another ; Alcorn et al . ( 2019 ) ; Dai & Van Gool ( 2018 ) find that a model for driverless cars may not generalise to new lighting conditions or object poses ; and Buolamwini & Gebru ( 2018 ) find that a model may perform worse on subsets of the distribution , such as different ethnicities , if the training set has an imbalanced distribution . Thus , it is important to understand when we expect a model to generalise and when we do not . This would allow a practitioner to have confidence in the system ( e.g . if a model is shown to be robust to the imaging conditions of different hospitals , then it can be deployed in new hospitals with confidence ) . While domain generalization is a well studied area , Gulrajani & Lopez-Paz ( 2021 ) ; Schott et al . ( 2021 ) have cast doubt on the efficacy of existing methods , raising the question : has any progress been made in domain generalization over a standard expectation risk minimization ( ERM ) algorithm ? Despite these discouraging results , there are many examples that machine learning models do generalise across datasets with different distributions . For example , CLIP ( Radford et al. , 2021 ) , with well engineered prompts , generalizes to many standard image datasets . Taori et al . ( 2020 ) found that models trained on one image dataset generalise to another , albeit with some drop in performance ; in particular , higher performing models generalise better . However , there is little understanding and experimentation on when and why models generalise , especially in realistic settings inspired by real-world applications . This begs the following question : Can we define the important distribution shifts to be robust to and then systematically evaluate the robustness of different methods ? To answer the above question , we present a grounded understanding of robustness to distribution shifts . We draw inspiration from disentanglement literature ( see section 6 ) , which aims to separate images into an independent set of factors of variation ( or attributes ) . In brief , we assume the data is composed of some ( possibly extremely large ) set of attributes . We expect models , having seen some distribution of values for an attribute , to be able to learn invariance to that attribute and so to generalise to unseen examples of the attribute and different distributions over that attribute . Using a simple example to clarify the setup , assume our data has two attributes ( shape and color ) among others . Given data with some distribution over the set of possible colors ( e.g . red and blue ) and the task of predicting shape ( e.g . circle or square ) , we want our model to generalise to unseen colors ( e.g . green ) or a different distribution of colors ( e.g . there are very few red circles in the training set , but the samples at evaluation are uniformly sampled from the set of possible colors and shapes ) . Using this framework , we evaluate models across three distribution shifts : spurious correlation , lowdata drift , and unseen data shift ( illustrated in figure 1 ) and two additional conditions ( label noise and dataset size ) . We choose these settings as they arise in the real world and harm generalization performance . Moreover , in our framework , these distribution shifts are the fundamental blocks of building more complex distribution shifts . We additionally evaluate models when there is varying amounts of label noise ( as inspired by noise arising from human raters ) and when the total size of the train set varies ( to understand how models perform as the number of training examples changes ) . The unique ability of our framework to evaluate fine-grained performance of models across different distribution shifts and under different conditions is of critical importance to analyze methods under a variety of real-world settings . This work makes the following contributions : • We propose a framework to unify different distribution shifts , defining how they arise and how different common approaches promote robustness to these shifts . We use this framework to define three , real world inspired distribution shifts . We then use this framework to create a systematic evaluation setup across real and synthetic datasets for different distribution shifts . Our evaluation framework is easily extendable to new distribution shifts , datasets , or methods to be evaluated . We will open source the code upon acceptance . • We evaluate and compare 19 different methods ( training more than 85K models ) in these settings . These methods span 5 common techniques to improve robustness : architecture choice , data augmentation , domain generalization , adaptive algorithms , and representation learning . This allows for a direct comparison across different areas in machine learning . • We find that simple techniques , such as data augmentation and pretraining are often effective and that domain generalization algorithms do work for certain datasets and distribution shifts . However , there is no easy way to select the best approach a-priori and results are inconsistent over different datasets and attributes , demonstrating there is still much work to be done to improve robustness in real-world settings . 2 FRAMEWORK TO EVALUATE GENERALIZATION . In this section we introduce our robustness framework for characterizing distribution shifts in a principled manner . We then relate three common , real world inspired distribution shifts . 2.1 LATENT FACTORISATION . We assume a joint distribution p of inputs x and corresponding attributes y1 , y2 , . . . , yK ( denoted as y1 : K ) with yk 2 Ak where Ak is a finite set . One of these K attributes is a label of interest , denoted as yl ( in a mammogram , the label could be cancer/benign and a nuisance attribute yi with i 6= l could be the identity of the hospital where the mammogram was taken ) . Our aim is to build a classifier f that minimizes the risk R. However , in real-world applications , we only have access to a finite set of inputs and attributes of size n. Hence , we minimize the empirical risk R̂ instead : R ( f ) = E ( x , yl ) ⇠p ⇥ L ( yl , f ( x ) ) ⇤ R̂ ( f ; p ) = 1 n X { ( yli , xi ) ⇠p } ni=1 L ( yli , f ( xi ) ) . where L is a suitable loss function . Here , all nuisance attributes yk with k 6= l are ignored and we work with samples obtained from the marginal p ( yl , x ) . In practice , however , due to selection bias or other confounding factors in data collection , we are only able to train and test our models on data collected from two related but distinct distributions : ptrain , ptest . For example , ptrain and ptest may be concentrated on different subsets of hospitals and this discrepancy may result in a distribution shift ; for example , hospitals may use different equipment , leading to different staining on their cell images . While we train f on data from ptrain by minimizing R̂ ( f ; ptrain ) , we aim to learn a model that generalises well to data from ptest ; that is , it should achieve a small R̂ ( f ; ptest ) . While generalization in the above sense is desirable for machine learning models , it is not clear why a model f trained on data from ptrain should generalise to ptest . It is worth noting that while ptrain and ptest can be different , they are both related to the true distribution p. We take inspiration from disentanglement literature to express this relationship . In particular , that we can view data as being decomposed into an underlying set of factors of variations . We formalise various distribution shifts using a latent variable model for the true data generation process : z ⇠ p ( z ) yi ⇠ p ( yi|z ) i = 1 . . .K x ⇠ p ( x|z ) ( 1 ) where z denotes latent factors . By a simple refactorization , we can write p ( y1 : K , x ) = p ( y1 : K ) Z p ( x|z ) p ( z|y1 : K ) dz = p ( y1 : K ) p ( x|y1 : K ) . Thus , the true distribution can be expressed as the product of the marginal distribution of the attributes with a conditional generative model . We assume that distribution shifts arise when a new marginal distribution for the attributes is chosen , such as p ( y1 : K ) 6= ptrain ( y1 : K ) 6= ptest ( y1 : K ) , but otherwise the conditional generative model is shared across all distributions , i.e. , we have ptest ( y1 : K , x ) = ptest ( y1 : K ) R p ( x|z ) p ( z|y1 : K ) dz , and similarly for ptrain . This is a special case of covariate shift ( Schölkopf et al. , 2012 ) where the input x is the result of an unknown ( but fixed ) generating process conditioned on the attributes . This is realistic in practice , as the physical process ( e.g . imaging in a hospital ) is constrained but attributes such as age , sex , imaging equipment , hospital , etc . may vary . To provide more context , as a running example , we use the color DSPRITES dataset ( Matthey et al. , 2017 ) ; where in our notation y1 defines the color with A1 = { red , green , blue } , and y2 defines the shape with A2 = { ellipse , heart , square } . We can imagine that a data collector ( intentionally or implicitly ) selects some marginal distribution over attributes ptrain ( y1 : K ) when training ; for example they select mostly blue ellipses and red hearts . This induces a new joint distribution over latent factors and attributes : ptrain ( z , y1 : K ) = p ( z|y1 : K ) ptrain ( y1 : K ) . Consequently , during training , we get images with a different joint distribution ptrain ( x , y1 : K ) = R p ( x|z ) ptrain ( z , y1 : K ) . This similarly applies when collecting data for the test distribution . We focus on common cases of distribution shifts visualized in figure 1 ; we discuss these in more detail in section 2.2 . The goal of enforcing robustness to distribution shifts is to maintain performance when the data generating distribution ptrain changes . In other words , we would like to minimize risk on p , ptest given only access to ptrain . We can achieve robustness in the following ways : 1 . Weighted resampling . We can resample the training set using importance weights W ( y1 : K ) = p ( y1 : K ) /ptrain ( y1 : K ) . Given the attributes , the i-th data point ( y1 : Ki , xi ) in the training set is used with probability W ( y1 : Ki ) / Pn i0=1 W ( y 1 : K i0 ) rather than 1/n . We refer to this empirical distribution as preweight . This requires access to the true distribution p ( y1 : K ) , so to avoid bias and improve fairness , it is often assumed that all combinations of attributes happen uniformly at random . 2 . Heuristic Data Augmentation . Weighted resampling can lead to samples being reused many times , leading to overfitting . To mitigate overfitting , we can augment the training data with heuristic choices , such as color jitter . While not the typical interpretation , we can view this setup as encouraging invariance to nuisance attributes in the generative model , which leads us to the next approach . 3 . Learned Data Augmentation : We can learn a generative model p̂ ( x|y1 : K ) from the training data that aims to approximate R p ( x|z ) p ( z|y1 : K ) dz , as the true conditional generator is by our assumption the same over all ( e.g . train and test ) distributions . If such a conditional generative model can be learned , we can sample new synthetic data at training time ( e.g . according to the true distribution p ( y1 : K ) ) to correct for the distribution shift . More precisely , we can generate data from the augmented distribution paug = ( 1 ↵ ) ptrain+↵p̂ ( x|y1 : K ) p ( y1 : K ) and train a supervised classifier on this augmented dataset . Here , ↵ 2 [ 0 , 1 ] is the percentage of synthetic data used for training . 4 . Representation Learning : An alternative factorization of a data generating distribution ( e.g . train ) is ptrain ( y1 : K , x ) = R p ( z|x ) ptrain ( y1 : K |z ) dz . We can learn an unsupervised representation that approximates p ( z|x ) using the training data only , and attach a classifier to learn a task specific head that approximates ptrain ( yl|z ) . Again , by our assumption p ( z|x ) / p ( x|z ) p ( z ) . Given a good guess of the true prior , the learned representation would not be impacted by the specific attribute distribution and so generalise to ptest , p .

Summary: The authors perform an extensive study of different types of distribution shifts to judge which methods perform best on which dataset and/or method. For this, they carefully define which distribution shifts they are interested in and perform a large scale study for each of the shifts. They propose a standardized framework (which can easily be extended to new methods) to train and evaluate new models and methods on the distribution shifts.