Empirical or Invariant Risk Minimization? A Sample Complexity Perspective

1 INTRODUCTION . A recent study shows that models trained to detect COVID-19 from chest radiographs rely on spurious factors such as the source of the data rather than the lung pathology ( DeGrave et al. , 2020 ) . This is just one of many alarming examples of spurious correlations failing to hold outside a specific training distribution . In one commonly cited example , Beery et al . ( 2018 ) trained a convolutional neural network ( CNN ) to classify camels from cows . In the training data , most pictures of the cows had green pastures , while most pictures of camels were in the desert . The CNN picked up the spurious correlation and associated green pastures with cows thus failing to classify cows on beaches . Recently , Arjovsky et al . ( 2019 ) proposed a framework called invariant risk minimization ( IRM ) to address the problem of models inheriting spurious correlations . They showed that when data is gathered from multiple environments , one can learn to exploit invariant causal relationships , rather than relying on varying spurious relationships , thus learning robust predictors . More recent work suggests that empirical risk minimization ( ERM ) is still state-of-the-art on many problems requiring OOD generalization ( Gulrajani & Lopez-Paz , 2020 ) . This gives rise to a fundamental question : when is IRM better than ERM ( and vice versa ) ? In this work , we seek to answer this question through a systematic comparison of the sample complexity of the two approaches under different types of train and test distributional mismatches . The distribution shifts ( Ptrain ( X , Y ) 6= Ptest ( X , Y ) ) that we consider informally stated satisfy an invariance condition – there exists a representation Φ∗ of the covariates such that Ptrain ( Y |Φ∗ ( X ) ) = Ptest ( Y |Φ∗ ( X ) ) = P ( Y |Φ∗ ( X ) ) . A special case of this occurs when Φ∗ is identity – Ptrain ( X ) 6= Ptest ( X ) but Ptrain ( Y |X ) = Ptest ( Y |X ) – such a shift is known as a covariate-shift ( Gretton et al. , 2009 ) . In many other settings Φ∗ may not be identity ( denoted as I ) , examples include settings with confounders or anti-causal variables ( Pearl , 2009 ) where covariates appear spuriously correlated with the label and Ptrain ( Y |X ) 6= Ptest ( Y |X ) . We use causal Bayesian networks to illustrate these shifts in Figure 1 . Suppose Xe = [ Xe1 , X e 2 ] represents the image , where X e 1 is the shape of the animal and Xe2 is the background color , Y e is the label of the animal , and e is the index of the environment/domain . In Figure 1a ) Xe2 is independent of ( Y e , Xe1 ) , it represents the covariate shift case ( Φ∗ = I ) . In Figure 1b ) Xe2 is spuriously correlated with Y e through the confounder εe . In Figure 1c ) Xe2 is spuriously correlated with Y e as it is anti-causally related to Y e. In both Figure 1b ) and c ) Φ∗ 6= I ; Φ∗ is a block diagonal matrix that selects Xe1 . Our setup assumes we are given data from multiple training environments satisfying the invariance condition , i.e. , P ( Y |Φ∗ ( X ) ) is the same across all of them . Ideally , we want to learn and predict using E [ Y |Φ∗ ( X ) ] ; this predictor has a desirable OOD behavior as we show later where we prove min-max optimality with respect to ( w.r.t . ) unseen test distributions satisfying the invariance condition . Our goal is to analyze and compare ERM and IRM ’ s ability to learn E [ Y |Φ∗ ( X ) ] from finite training data acquired from a fixed number of training environments . Our analysis has two parts . 1 ) Covariate shift case ( Φ∗ = I ) : ERM and IRM achieve the same asymptotic solution E [ Y |X ] . We prove ( Proposition 4 ) that the sample complexity for both the methods is similar thus there is no clear winner between the two in the finite sample regime . For the setup in Figure 1a ) , both ERM and IRM learn a model that only uses Xe1 . 2 ) Confounder/Anti-causal variable case ( Φ∗ 6= I ) : We consider a family of structural equation models ( linear and polynomial ) that may contain confounders and/or anti-causal variables . For the class of models we consider , the asymptotic solution of ERM is biased and not equal to the desired E [ Y |Φ∗ ( X ) ] . We prove that IRM can learn a solution that is within O ( √ ) distance from E [ Y |Φ∗ ( X ) ] with a sample complexity that increases as O ( 1 2 ) and increases polynomially in the complexity of the model class ( Proposition 5 , 6 ) ; ( defined later ) is the slack in IRM constraints . For the setup in Figure 1b ) and c ) , IRM gets close to only using Xe1 , while ERM even with infinite data ( Proposition 17 in the supplement ) continues to use Xe2 . We summarize the results in Table 1 . Arjovsky et al . ( 2019 ) proposed the colored MNIST ( CMNIST ) dataset ; comparisons on it showed how ERM-based models exploit spurious factors ( background color ) . The CMNIST dataset relied on anti-causal variables . Many supervised learning datasets may not contain anti-causal variables ( e.g . human labeled images ) . Therefore , we propose and analyze three new variants of CMNIST in addition to the original one that map to different real-world settings : i ) covariate shift based CMNIST ( CS-CMNIST ) : relies on selection bias to induce spurious correlations , ii ) confounded CMNIST ( CF-CMNIST ) : relies on confounders to induce spurious correlations , iii ) anti-causal CMNIST ( ACCMNIST ) : this is the original CMNIST proposed by Arjovsky et al . ( 2019 ) , and iv ) anti-causal and confounded ( hybrid ) CMNIST ( HB-CMNIST ) : relies on confounders and anti-causal variables to induce spurious correlations . On the latter three datasets , which belong to the Φ∗ 6= I class described above , IRM has a much better OOD behavior than ERM , which performs poorly regardless of the data size . However , IRM and ERM have a similar performance on CS-CMNIST with no clear winner . These results are consistent with our theory and are also validated in regression experiments . 2 RELATED WORKS . IRM based works . Following the original work IRM from Arjovsky et al . ( 2019 ) , there have been several interesting works — ( Teney et al. , 2020 ; Krueger et al. , 2020 ; Ahuja et al. , 2020 ; Chang et al. , 2020 ; Mahajan et al. , 2020 ) is an incomplete representative list — that build new methods inpired from IRM to address the OOD generalization problem . Arjovsky et al . ( 2019 ) prove OOD guarantees for linear models with access to infinite data from finite environments . We generalize these results in several ways . We provide a first finite sample analysis of IRM . We characterize the impact of hypothesis class complexity , number of environments , weight of IRM penalty on the sample complexity and its distance from the OOD solution for linear and polynomial models . Theory of domain generalization and domain adaption . Following the seminal works ( BenDavid et al. , 2007 ; 2010 ) , there have been many interesting works — ( Muandet et al. , 2013 ; Ajakan et al. , 2014 ; Zhao et al. , 2019 ; Albuquerque et al. , 2019 ; Li et al. , 2017 ; Piratla et al. , 2020 ; Matsuura & Harada , 2020 ; Deng et al. , 2020 ; David et al. , 2010 ; Pagnoni et al. , 2018 ) is an incomplete representative list ( see Redko et al . ( 2019 ) for further references ) — that build the theory of domain adaptation and generalization and construct new methods based on it . While many of these works develop bounds on loss over the target domain using train data and unlabeled target data , some ( Ben-David & Urner , 2012 ; David et al. , 2010 ; Pagnoni et al. , 2018 ) analyze the finite sample ( PAC ) guarantees for domain adaptation under covariate shifts . These works ( Ben-David & Urner , 2012 ; David et al. , 2010 ; Pagnoni et al. , 2018 ) access unlabeled data from a target domain , which we do not . Instead , we have data from multiple training domains ( as in domain generalization ) . In these works , the guarantees are w.r.t . a specific target domain , while we provide ( for linear and polynomial models ) worst-case guarantees w.r.t . all the unseen domains satisfying the invariance condition . Also , we consider a larger family of distribution shifts including covariate shifts . The above two categories are not exhaustive – e.g. , there are some recent works that characterize how some inductive biases favor extrapolation Xu et al . ( 2021 ) and can be better for OOD generalization . 3 SAMPLE COMPLEXITY OF INVARIANT RISK MINIMIZATION . 3.1 INVARIANT RISK MINIMIZATION . We start with some background on IRM ( Arjovsky et al. , 2019 ) . Consider a datasetD = { De } e∈Etr , which is a collection of datasets De = { ( xei , yei , e ) } ne i=1 } obtained from a set of training environments Etr , where e is the index of the environment , i is the index of the data point in the environment , ne is the number of points from environment , xei ∈ X ⊆ Rn is the feature value and yei ∈ Y ⊆ R is the corresponding label . Define a probability distribution { πe } e∈Etr , πe is the probability that a training data point is from environment e. Define a probability distribution of points conditional on environment e as Pe , ( Xe , Y e ) ∼ Pe . Define the joint distribution P̄ , ( Xe , Y e , e ) ∼ P̄ , dP̄ ( Xe , Y e , e ) = πedPe ( Xe , Y e ) . D is a collection of i.i.d . samples from P̄ . Define a predictor f : X → R and the space F of all the possible maps from X → R. Define the risk achieved by f in environment e as Re ( f ) = Ee [ ` ( f ( Xe ) , Y e ) ] , where ` is the loss , f ( Xe ) is the predicted value , Y e is the corresponding label and Ee is the expectation conditional on environment e. The overall expected risk across the training environments is R ( f ) = ∑ e∈Etr π eRe ( f ) . We are interested in two settings : regression ( square loss ) and binary-classification ( cross-entropy loss ) . In the main body , our focus is regression ( square loss ) and we mention wherever the results extend to binary-classification ( cross-entropy ) . We discuss these extensions in the supplement . OOD generalization problem . We want to construct a predictor f that performs well across many unseen environments Eall , where Eall ⊇ Etr . For o ∈ Eall\Etr , the distribution Po can be very different from the train environments . Next we state the OOD problem . min f∈F max e∈Eall Re ( f ) ( 1 ) The above problem is very challenging to solve since we only have access to data from training environments Etr but are required to find the robust solution over all environments Eall . Next , we make assumptions on Eall and characterize the optimal solution to equation 1 . Assumption 1 . Invariance condition . There exists a representation Φ∗ that transforms Xe to Ze = Φ∗ ( Xe ) and ∀e , o ∈ Eall , ∀z ∈ Φ∗ ( X ) satisfies Ee [ Y e|Ze = z ] = Eo [ Y o|Zo = z ] . Also , ∀e ∈ Eall , ∀z ∈ Φ∗ ( X ) , Vare [ Y e|Ze = z ] = ξ2 , where Vare is the conditional variance . The above assumption is inspired from causality ( Pearl , 2009 ) . Φ∗ acts as the causal feature extractor and from the definition of causal features , it follows that Ee [ Y e|Ze = z ] does not vary across environments . When a human labels a cow she uses Φ∗ to extract causal features from the pixels to identify cow while ignoring the background . The first part of the above assumption encompasses a large class of distribution shifts including standard covariate shifts ( Gretton et al. , 2009 ) . Covariate shift assumes ∀e , o ∈ Eall , ∀x ∈ X , P ( Y e|Xe = x ) and P ( Y o|Xo = x ) are equal thus implying Ee [ Y e|Xe = x ] = Eo [ Y o|Xo = x ] . Therefore , for covariate shifts , Φ∗ is identity in Assumption 1 . A simple instance illustrating Assumption 1 with Φ∗ = I is when Y e ← g ( Xe ) +εe , where Ee [ εe ] = 0 , Ee [ ( εe ) 2 ] = σ2 , εe ⊥ Xe . Using Assumption 1 , we define the invariant map m : Φ∗ ( X ) → R as follows ∀z ∈ Φ∗ ( X ) , m ( z ) = Ee [ Y e|Ze = z ] , where Ze = Φ∗ ( Xe ) ( 2 ) Assumption 2 . Existence of an environment where the invariant representation is sufficient . ∃ an environment e ∈ Eall such that Y e ⊥ Xe|Ze Assumption 2 states there exists an environment where the information that Xe has about Y e is also contained in Ze . Define a composition m ◦ Φ∗ , ∀x ∈ X , m ◦ Φ∗ ( x ) = Ee [ Y e|Ze = Φ∗ ( x ) ] . Proposition 1 . If ` is the square loss , and Assumptions 1 and 2 hold , then m ◦ Φ∗ solves the OOD problem ( equation 1 ) . The proofs of all the propositions are in the supplement . A similar result holds for the cross-entropy loss ( discussion in supplement ) . For the rest of the paper , we focus on learning m ◦ Φ∗ as it solves the OOD problem . For covariate shifts Φ∗ = I , m ( x ) = Ee [ Y e|Xe = x ] is the OOD solution . In Arjovsky et al . ( 2019 ) , a proof connecting m ◦ Φ∗ and OOD was not stated . Recently , in Koyama & Yamaguchi ( 2020 ) , a result similar to Proposition 1 was shown but with a few differences . The authors assume conditional probabilities are invariant unlike our assumption that only requires conditional expectations and variances to be invariant . However , their result applies to more losses . m ◦Φ∗ is the target we want to learn . Arjovsky et al . ( 2019 ) proposed IRM since standard min-max optimization over the training environments Etr and ERM fail to learn m ◦ Φ∗ in many cases . The authors in Arjovsky et al . ( 2019 ) identify a crucial property of m ◦ Φ∗ and use it to define an object called invariant predictor that we define next . Invariant predictor and IRM optimization . Define a representation map Φ : X → Z from feature space to representation space Z ⊆ Rq . Define a classifier map , w : Z → R from representation space to real values . DefineHΦ andHw as the spaces of representations and classifiers respectively . A data representation Φ elicits an invariant predictor w ◦ Φ across environments e ∈ Etr if there is a classifier w that achieves the minimum risk simultaneously for all the environments , i.e. , ∀e ∈ Etr , w ∈ arg minw̄∈Hw Re ( w̄ ◦ Φ ) . Observe that if we we transform the data with representation Φ∗ then m will achieve the minimum risk simultaneously in all the environments . If Φ∗ ∈ HΦ and m ∈ Hw , then m ◦ Φ∗ is an invariant predictor . IRM selects the invariant predictor with least sum risk across environments ( results presented later can be adapted if invariant predictor was selected based on the worst-case risk over the environments as well ) as follows : min Φ∈HΦ , w∈Hw R ( w ◦ Φ ) = ∑ e∈Etr πeRe ( w ◦ Φ ) s.t . w ∈ arg min w̄∈Hw Re ( w̄ ◦ Φ ) , ∀e ∈ Etr ( 3 ) From the above discussion we know m ◦ Φ∗ is a feasible solution to equation 3 . It is also the ideal solution we want IRM to find since it solves equation 1 . Later in Propositions 4 , 5 , and 6 , we show that IRM actually solves equation equation 1 . For the setups in Proposition 5 , and 6 , conventional ERM based approaches fail thus justifying the need for above formulation .

Summary: This paper considers a learning scenario where training data $(X,Y)$ comes from a mixture, where membership in each mixture component (“environment”) is clearly labeled. Generalization is required not just under the same mixture, but potentially under changing mixing distributions. This is captured via several alternative formulations. The key underlying assumption that makes the extrapolation possible is the existence of a a representation $\Phi(X)$ that leads to an environment-invariant statistics (by focusing on quadratic loss and linear predictors, the paper asks for invariant $E[Y|\Phi(X)]$ and constant $var[Y|\Phi(X)]$). A range of other assumptions are made to enable analysis, such as sufficiency of the representation in at least one environment and boundedness of the loss and its gradient.