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: ## Very useful study on calibrating OOD models and benchmarks  This paper provides a fine-grained analysis on nature of distribution shifts in image datasets through extensive experiments on various benchmark methods for robustness. They work with an assumption that the underlying factors of variation, including the label, are encoded through discrete attributes. The authors provide a concrete framework to put forth the latent factorization model, including appropriate grouping for various types of benchmark methods, distribution shifts as well as training conditions. Several experimental demonstrations are provided to compare and contrast robustness of existing methods leading to numerous useful observations.