Out-of-distribution Prediction with Invariant Risk Minimization: The Limitation and An Effective Fix

1 INTRODUCTION E Xc YXs Figure 1 : The causal graph in OOD prediction : P ( Y |Xc ) is invariant across domains . The spurious correlation P ( Y |Xs ) may vary . A directed ( bidirected ) edge is a causal relationship ( correlation ) . Strong empirical results have demonstrated the efficacy of deep neural networks ( DNNs ) in a variety of areas including computer vision , natural language processing and speech recognition . However , such positive results overwhelmingly rely on the assumption that the training , validation and test data consist of independent identical samples of the same underlying distribution . In contrast , in the setting of outof-distribution ( OOD ) prediction where ( 1 ) the training data is from multiple domains and ( 2 ) the test set has different distribution from the training , the performance of DNNs can be dramatically degraded . This is because DNNs are prone to pick up spurious correlations which do not hold beyond the training data distribution ( Beery et al. , 2018 ; Arjovsky et al. , 2019 ) . For example , when most camel pictures in a training set have a desert in the background , DNNs will pick up the spurious correlation between a desert and the class label , leading to failures when camel pictures come with different backgrounds in a test set . Therefore , OOD prediction remains an extremely challenging problem for DNNs . The invariant causal relationships across different domains turns out to be the key to address the challenge of OOD prediction . The causal graph in Fig . 1 describes the relationships of the variables in the OOD prediction problem . Although spurious correlations learned in one do- main are unreliable in another , invariant causal relationships enable DNNs that capture causal relationships to generalize to unseen domains . In practice , it is extremely difficult to know whether an input feature is causal or spurious . Thus , a recipe for training DNNs that capture causal relationships is learning causal feature representations that hold invariant causal relationships with the class label . Thus , the main challenge of OOD prediction becomes learning causal feature representations given training data from multiple domains . Invariant Risk Minimization ( IRM ) ( Arjovsky et al. , 2019 ) is proposed to learn causal feature representations for OOD prediction . IRM formulates the invariance of causal relationships as a constraint . It requires that causal feature representations must result in the same optimal classifier across all domains . IRM is written as the conditional independence , Y ⊥⊥ E|F ( X ) in ( Chang et al. , 2020 ; Zeng et al. , 2019 ) which can be derived from the causal graph in Fig . 1 when F ( X ) is a mapping of causal features ( Xc in Fig . 1 ) with no information loss . A detailed discussion on IRM and Y ⊥⊥ E|F ( X ) is in Appendix A . Despite IRM ’ s success in the colored MNIST dataset ( Arjovsky et al. , 2019 ) , in this work , we find an important issue of IRM that has not been discussed . Specifically , we consider the situation where strong spurious correlations among the spurious features , the class label and the domain label only hold for training data . We name this strong Λ spuriousness using the shape of the structure among Xs , E and Y in Fig . 1 . Under strong Λ spuriousness , IRM regularized empirical risk can have low values with spurious feature representations that accurately predict the domain label . This is because , in this setting , picking up spurious features can achieve high accuracy in predicting both domain and class in the training set , but not in the test . However , the colored MNIST dataset can not expose this issue because the strong similarity between the two training domains makes it difficult to pick up the weak spurious correlation between the domain label and the spurious features . To illustrate this problem , we design a new dataset – the colored MNIST plus . As shown in Fig . 2 , in this dataset , under strong Λ spuriousness , the performance of IRM models is significantly degraded . Moreover , to resolve this issue of IRM , we propose an effective solution , which combines IRM with conditional distribution matching ( CDM ) . The CDM constraint means that the representation distribution of instances from the same class should be invariant across domains . Theoretically , we show that ( 1 ) causal feature representations can satisfy CDM and IRM at the same time and ( 2 ) CDM can prevent DNNs from learning spurious feature representations that accurately predict the domain label . Empirically , on our newly introduced dataset , the proposed method achieves significant performance improvement over IRM under strong Λ spuriousness . 2 PRELIMINARIES AND IRM . Notations . We use lowercase ( e.g. , x ) , uppercase ( e.g. , X ) and calligraphic uppercase ( e.g. , X ) letters for values , random variables and spaces . We let X ∈ X , Y ∈ Y and E ∈ E denote raw input , the class label and the domain label where X , Y and E are the spaces of input , class labels and domains . A DNN model consists of a feature learning function F and a classifier G. A feature learning function F : X → Rd maps raw input X to its d-dimensional representations F ( X ) . A classifier G : Rd → Y maps a feature representation to a class label . We denote their parameters by θF and θG , respectively . Let θ = Concat ( θF , θG ) denote the concatenation of them . A domain e of ne instances is denoted by De = { xei , yei } ne i=1 . Let Etr and Ets denote the set of training and test domains , in OOD prediction , we have ( 1 ) |Etr| > 1 and ( 2 ) Etr ∩ Ets = ∅ . Problem Statement . Given data from multiple training domains { De } e∈Etr . We aim to predict the label ye ′ i of each instance with features x e′ i from a test domain { xe ′ i , y e′ i } ne′ i=1 , e ′ ∈ Ets . Invariant Risk Minimization . IRM ( Arjovsky et al. , 2019 ) is a recently proposed method to impose the causal inductive bias : the causal relationships between causal features and the label should be invariant across different domains . It not only aims to address the challenging OOD prediction problem , but also is a pioneer work that guides causal machine learning research towards the development of inductive bias imposing causal constraints . The effectiveness of IRM and its variants have been demonstrated across various areas including computer vision ( Ahuja et al. , 2020 ) , natural language processing ( Chang et al. , 2020 ) , CTR prediction ( Zeng et al. , 2019 ) , reinforcement learning ( Zhang et al. , 2020a ) and financial forecasting ( Krueger et al. , 2020 ) . Arjovsky et al . ( 2019 ) propose the original formulation of IRM as a two-stage optimization problem : arg min θF , θG ∑ e∈Etr E ( x , y ) ∼De [ R e ( G ( F ( x ) ) , y ) ] s.t . θG ∈ arg min θ′G Re ( G ( F ( x ) ; θ′G ) , y ) , ( 1 ) where Re denotes the loss function of domain e. Then , they show that the IRM constraint can be imposed by adding a regularizer into the loss function . It turns out the IRM constraint can be written as the conditional independence Y ⊥⊥ E|F ( X ) ( Chang et al. , 2020 ; Zeng et al. , 2019 ) . This can be derived from the causal graph in Fig . 1 . For brevity , we refer to Y ⊥⊥ E|F ( X ) as the IRM constraint . 3 THE LIMITATION OF IRM . In this section , we find that IRM fails under strong Λ spuriousness – when the spurious correlations among spurious features , the domain label and the class label are strong . To show it , we design a new dataset – the colored MNIST plus ( CMNIST+ ) to expose this limitation of IRM . Then , we try to answer two crucial research questions : ( 1 ) why does IRM fail in CMNIST+ ? ( 2 ) Why does IRM work for the original colored MNIST ( CMNIST ) dataset ? 3.1 WHY DOES IRM FAIL ? . Despite its previous success , results in Fig . 2 show that IRM fails in our newly designed colored MNIST plus dataset , under strong Λ spuriousness ( ρ ≥ 0.8 ) . Here , we show that , under strong Λ spuriousness , at least two types of feature representations , F ( X ) = Y and F ( X ) = E , satisfy the IRM constraint . When F ( X ) = E , the feature representations pick up spurious relationships . Satisfying IRM by Perfectly Predicting the Label . First , we consider a case where there exists a feature representation that perfectly predicts the class label , i.e. , F ( X ) = Y . In this case , we know that the feature representation satisfies the IRM constraint because Y ⊥⊥ E|F ( X ) = Y always holds . This is reasonable because the causal relationship between Y and itself must be invariant across different domains . Satisfying IRM by Perfectly Predicting the Domain Does Not Lead to OOD Generalizable Feature Representations . However , there is another way to satisfy the IRM constraint . If there exists a feature representation F ( X ) = E , then the IRM constraint can be satisfied by F ( X ) . This is due to the fact Y ⊥⊥ E|F ( X ) = E. However , this is not desired behavior of IRM as F ( X ) = E is a type of spurious feature representation . This is because P ( Y |E ) can vary across domains . Using the camel picture example , if domains are the individuals who take the picture , we may find more camel pictures taken by person e than from person e′ ( P ( Y = camel|E = e ) > P ( Y = camel|E = e′ ) ) . Why is F ( X ) = E a good solution for the original IRM optimization problem ( Eq . 1 ) ? This is because F ( X ) = E also leads to low values of the loss function Re since the direct causal influence of the domain label E on the class label Y is also strong when Λ spuriousness is strong . So , with a proper classifier G , we can achieve low values of Re with F ( X ) = E. Then , we describe the CMNIST+ dataset and show experimental and theoretical results to support our claim . The Colored MNIST Plus ( CMNIST+ ) Dataset . We follow CMNIST ( Arjovsky et al. , 2019 ) to create CMNIST+ by resampling and adding colors to instances of MNIST . The digits 0− 4 ( 5− 9 ) are class Y = 1 ( Y = 0 ) . We randomly flip 25 % of class labels so that spuriousness can be stronger than the causal relationship between the shape S ( causal features ) and the class label Y . Table 1 describes the dataset . The variable C ∈ { G , B , R } denotes the color , which represents the spurious feature Xs . We explain why we use three colors for CMNIST+ in Appendix C.2 . The parameter ρ ∈ ( 0.5 , 1 ) controls the strength of the spurious correlations between the color ( spurious features ) C and the class label Y through their common cause , the domain label E. The larger the value of ρ , the stronger the spurious correlations . Intuitively , in the two training domains of CMNIST+ , in addition to the strong spurious correlation between the class label Y and color C , we set the spurious correlation between the domain label E and color C to be strong , too . Thus , the CMNIST+ dataset can expose the problem of IRM : it would not penalize the models that fit the domain variable as the feature representation ( F ( X ) = E ) . To verify this claim , we show analysis results to support the experimental results in Fig . 2 . The data generating process of CMNIST+ can be found in Append C.2 . Analysis Results . Here , we show the analysis results on the CMNIST+ to answer two questions : what features will be learned by ERM and IRM under under different Λ spuriousness ? What is the expected test accuracy of these models ? For simplicity , we assume that the classifier is deterministic , which always predicts the majority class given the feature . If the number of instances from each class is the same , then it would predict a random label . We analyze three types of spurious feature representations : ( 1 ) those that fit color , ( 2 ) those that fit the domain label and ( 3 ) those that use a combination of them for prediction . The first case mimics the behavior of ERM that picks up the spurious correlation between color and the class label . The second one represents a model satisfying IRM by F ( X ) = E. The third one stands for a model satisfying IRM by F ( X ) = Concat ( E , C ) . One may argue that IRM can be fixed if we simply balance the two classes in each domain . Theoretical analysis shows that this is an invalid solution . We summarize results in Table 2 where Ê is the domain label predicted by a deterministic classifier using color as the feature . This is because in an unseen test domain , we can not directly use the domain label for prediction . The highlighted numbers show which type of features would be learned by ERM and IRM . As ρ increases , ERM and IRM are more likely to fit spurious features . This explains results in Fig . 2 . As a DNN may still pick up some causal features ( shape in CMNIST+ ) even when ρ = 0.9 , the test accuracy of IRM in practice would be greater than 0.35 without label balancing and 0.2 with label balancing . Similarly , the test accuracy of ERM would be greater than 0.2 . From Fig . 2 , we can see the test accuracy of IRM , IRM with label balancing and ERM are slightly better than the aforementioned lower bounds . The derivations can be found in Appendix B . Experimental Setup . We use a LeNet-5 ( LeCun et al. , 1998 ) instead of a three-layer MLP . To take input with three colors , we modify the first CNN layer to have three channels . LeNet-5 has more predictive power such that it can pick up the three types of spurious correlations ( color – domain , color – class , and domain – class ) or the causal relationship ( shape – class ) . We randomly split the instances from the training domains into 80 % training and 20 % validation . The model selection is done by picking the one with the lowest validation loss in each run . We report the average test accuracy of the selected models in 10 runs . It is crucial to ensure only data from training domains are used in model selection , since the test domain should be unseen during training and validation . During training , we begin applying the IRM penalty at iteration KIRM , set to 200 , 400 , or 600 . By varying KIRM , we aim to examine the following hypothesis : IRM works by pushing the spurious features out of the representations learned by the standard ERM training before the IRM penalty is applied ( Krueger et al. , 2020 ) . More details on the setup can be found in Appendix C.1 . Experimental Results . Fig . 2 shows the performance of IRM and IRM with balanced classes in each domain . We make the following observations : first , the test accuracy of IRM drops dramatically under strong Λ spuriousness . Second , we show that a naı̈ve fix for IRM , balancing the two classes in each domain by oversampling the minority class , does not lead to improvement in performance .

Summary: In this work, the authors focus on the out-of-distribution generalization problem. The input is dataset from multiple environments and the goal is to learn a model that generalizes well to an unseen test environment. The work is based on recent line of works on invariant risk minimization (IRM) (Arjovsky et al.). The authors show that under a certain type of structure for the generative model, where the domain/environment label itself has a strong correlation with the spurious factors and the target label, IRM fails. The authors propose an extension of the colored MNIST dataset to highlight this problem. Finally, the authors build a method that works better than IRM on the extension of colored MNIST dataset.