The Risks of Invariant Risk Minimization

1 INTRODUCTION . Prediction algorithms are evaluated by their performance on unseen test data . In classical machine learning , it is common to assume that such data are drawn i.i.d . from the same distribution as the data set on which the learning algorithm was trained—in the real world , however , this is often not the case . When this discrepancy occurs , algorithms with strong in-distribution generalization guarantees , such as Empirical Risk Minimization ( ERM ) , can fail catastrophically . In particular , while deep neural networks achieve superhuman performance on many tasks , there is evidence that they rely on statistically informative but non-causal features in the data ( Beery et al. , 2018 ; Geirhos et al. , 2018 ; Ilyas et al. , 2019 ) . As a result , such models are prone to errors under surprisingly minor distribution shift ( Su et al. , 2019 ; Recht et al. , 2019 ) . To address this , researchers have investigated alternative objectives for training predictors which are robust to possibly egregious shifts in the test distribution . The task of generalizing under such shifts , known as Out-of-Distribution ( OOD ) Generalization , has led to many separate threads of research . One approach is Bayesian deep learning , accounting for a classifier ’ s uncertainty at test time ( Neal , 2012 ) . Another technique that has shown promise is data augmentation—this includes both automated data modifications which help prevent overfitting ( Shorten & Khoshgoftaar , 2019 ) and specific counterfactual augmentations to ensure invariance in the resulting features ( Volpi et al. , 2018 ; Kaushik et al. , 2020 ) . A strategy which has recently gained particular traction is Invariant Causal Prediction ( ICP ; Peters et al . 2016 ) , which views the task of OOD generalization through the lens of causality . This framework assumes that the data are generated according to a Structural Equation Model ( SEM ; Bollen 2005 ) , which consists of a set of so-called mechanisms or structural equations that specify variables given their parents . ICP assumes moreover that the data can be partitioned into environments , where each environment corresponds to interventions on the SEM ( Pearl , 2009 ) , but where the mechanism by which the target variable is generated via its direct parents is unaffected . Thus the causal mechanism of the target variable is unchanging but other aspects of the distribution can vary broadly . As a result , learning mechanisms that are the same across environments ensures recovery of the invariant features which generalize under arbitrary interventions . In this work , we consider objectives that attempt to learn what we refer to as the “ optimal invariant predictor ” —this is the classifier which uses and is optimal with respect to only the invariant features in the SEM . By definition , such a classifier does not overfit to environment-specific properties of the data distribution , so it will generalize even under major distribution shift at test time . In particular , we focus our analysis on one of the more popular objectives , Invariant Risk Minimization ( IRM ; Arjovsky et al . ( 2019 ) ) , but our results can easily be extended to similar recently proposed alternatives . Various works on invariant prediction ( Muandet et al. , 2013 ; Ghassami et al. , 2017 ; Heinze-Deml et al. , 2018 ; Rojas-Carulla et al. , 2018 ; Subbaswamy et al. , 2019 ; Christiansen et al. , 2020 ) consider regression in both the linear and non-linear setting , but they exclusively focus on learning with fully or partially observed covariates or some other source of information . Under such a condition , results from causal inference ( Maathuis et al. , 2009 ; Peters et al. , 2017 ) allow for formal guarantees of the identification of the invariant features , or at least a strict subset of them . With the rise of deep learning , more recent literature has developed objectives for learning invariant representations when the data are a non-linear function of unobserved latent factors , a common assumption when working with complex , high-dimensional data such as images . Causal discovery and inference with unobserved confounders or latents is a much harder problem ( Peters et al. , 2017 ) , so while empirical results seem encouraging , these objectives are presented with few formal guarantees . IRM is one such objective for invariant representation learning . The goal of IRM is to learn a feature embedder such that the optimal linear predictor on top of these features is the same for every environment—the idea being that only the invariant features will have an optimal predictor that is invariant . Recent works have pointed to shortcomings of IRM and have suggested modifications which they claim prevent these failures . However , these alternatives are compared in broad strokes , with little in the way of theory . In this work , we present the first formal analysis of classification under the IRM objective under a fairly natural and general model which carefully formalizes the intuition behind the original work . Our results show that despite being inspired by invariant prediction , this objective can frequently be expected to perform no better than ERM . In the linear setting , we present simple , exact conditions under which solving to optimality succeeds or , more often , breaks down in recovering the optimal invariant predictor . We also demonstrate another major failure case—under mild conditions , there exists a feasible point that uses only non-invariant features and achieves lower empirical risk than the optimal invariant predictor ; thus it will appear as a more attractive solution , yet its reliance on non-invariant features mean it will fail to generalize . As corollaries , we present similar settings where all recently suggested alternatives to IRM likewise fail . Futhermore , we present the first results in the non-linear regime : we demonstrate the existence of a classifier with exponentially small suboptimality which nevertheless heavily relies on non-invariant features on most test inputs , resulting in worse-than-chance performance on distributions that are sufficiently dissimilar from the training environments . These findings strongly suggest that existing approaches to ICP for high-dimensional latent variable models do not cleanly achieve their stated objective and that future work would benefit from a more formal treatment . 2 RELATED WORK . Works on learning deep invariant representations vary considerably : some search for a domaininvariant representation ( Muandet et al. , 2013 ; Ganin et al. , 2016 ) , i.e . invariance of the distribution p ( Φ ( x ) ) , typically used for domain adaptation ( Ben-David et al. , 2010 ; Ganin & Lempitsky , 2015 ; Zhang et al. , 2015 ; Long et al. , 2018 ) , with assumed access to labeled or unlabeled data from the target distribution . Other works instead hope to find representations that are conditionally domain-invariant , with invariance of p ( Φ ( x ) | y ) ( Gong et al. , 2016 ; Li et al. , 2018 ) . However , there is evidence that invariance may not be sufficient for domain adaptation ( Zhao et al. , 2019 ; Johansson et al. , 2019 ) . In contrast , this paper focuses instead on domain generalization ( Blanchard et al. , 2011 ; Rosenfeld et al. , 2021 ) , where access to the test distribution is not assumed . Recent works on domain generalization , including the objectives discussed in this paper , suggest invariance of the feature-conditioned label distribution . In particular , Arjovsky et al . ( 2019 ) only assume invariance of E [ y | Φ ( x ) ] ; follow-up works rely on a stronger assumption of invariance of higher conditional moments ( Krueger et al. , 2020 ; Xie et al. , 2020 ; Jin et al. , 2020 ; Mahajan et al. , 2020 ; Bellot & van der Schaar , 2020 ) . Though this approach has become popular in the last year , it is somewhat similar to the existing concept of covariate shift ( Shimodaira , 2000 ; Bickel et al. , 2009 ) , which considers the same setting . The main difference is that these more recent works assume that the shifts in p ( Φ ( x ) ) occur between discrete , labeled environments , as opposed to more generally from train to test distributions . Some concurrent lines of work study different settings yet give results which are remarkably similar to ours . Xu et al . ( 2021 ) show that an infinitely wide two-layer network extrapolates linear functions when the training data is sufficiently diverse . In the context of domain generalization specifically , Rosenfeld et al . ( 2021 ) prove that ERM remains optimal for both interpolation and extrapolation in the linear setting and that the latter is exponentially harder than the former . These results mirror our findings that none of the studied objectives outperform ERM . 3 MODEL AND INFORMAL RESULTS . We consider an SEM with explicit separation of invariant features zc , whose joint distribution with the label is fixed for all environments , and environmental features ze ( “ non-invariant ” ) , whose distribution can vary . This choice is to ensure that our model properly formalizes the intuition behind invariant prediction techniques such as IRM , whose objective is to ensure generalizing predictors by recovering only the invariant features—we put off a detailed description of these objectives until after we have introduced the necessary terminology . We assume that data are drawn from a set of E training environments E = { e1 , e2 , . . . , eE } and that we know from which environment each sample is drawn . For a given environment e , the data are defined by the following process : first , a label y ∈ { ±1 } is drawn according to a fixed probability : y = { 1 , w.p . η −1 , otherwise . ( 1 ) Next , both invariant features and environmental features are drawn according to a Gaussian:1 zc ∼ N ( y · µc , σ2cI ) , ze ∼ N ( y · µe , σ2eI ) , ( 2 ) with µc ∈ Rdc , µe ∈ Rde—typically , for complex , high-dimensional data we would expect E < dc de . Finally , the observation x is generated as a function of the latent features : x = f ( zc , ze ) . ( 3 ) The complete data generating process is displayed in Figure 3.1 . We assume f is injective , so that it is in principle possible to recover the latent features from the observations , i.e . there exists a function Φ such that Φ ( f ( zc , ze ) ) = [ zc , ze ] T . We remark that this our only assumption on f , even when it is non-linear . Further , note that we model class-conditional means as direct opposites merely for clarity , as it greatly simplifies the calculations . None of our proofs require this condition : it is straightforward to extend our results to arbitrary means , and the non-linear setting also allows for arbitrary covariances . In fact , our proof technique for non-linear f could be applied to any distribution that sufficiently concentrates about its mean ( e.g. , sub-Gaussian ) . We write the joint and marginal distributions as pe ( x , y , zc , ze ) . When clear from context , we omit the specific arguments . Remarks on the model . This model is natural and flexible ; it generalizes several existing models used to analyze learning under the existence of adversarial distribution shift or non-invariant correlations ( Schmidt et al. , 2018 ; Sagawa et al. , 2020 ) . The fundamental facet of this model is the constancy of the invariant parameters η , µc , σ2c , f across environments—the dependence of µe , σe on the environment allows for varying distributions , while the true causal process remains unchanged . Here we make a few clarifying remarks : • We do not impose any constraints on the model parameters . In particular , we do not assume a prior over the environmental parameters . Observe that µc , σ2c are the same for all environments , 1Note the deliberate choice to have ze depend on y . Much work on this problem models spurious features which correlate with the label but are not causal . However , the term “ spurious ” is often applied incongruously ; in recent work , the term has been co-opted to refer to any feature that correlates with the label but does not cause it . Thus there is a subtle distinction : if we allow for anti-causality , i.e . the label causing the features , the resulting correlation is not spurious . We therefore avoid using the term “ spurious ” in this work . hence the subscript indicates the invariant relationship . In contrast , with some abuse of notation , the environmental subscript is used to indicate both dependence on the environment and the index of the environment itself ( e.g. , µi represents the mean specific to environment i ) . • While we have framed the model as y causing zc , the causation can just as easily be viewed in the other direction . The log-odds of y are a linear function of zc—this matches logistic regression with an invariant regression vector βc = 2µc/σ2c and bias β0 = log η 1−η . We present the model as above to emphasize that the causal relationships between y and the zc , ze are a priori indistinguishable , and because we believe this direction is more intuitive . We consider the setting where we are given infinite samples from each environment ; this allows us to isolate the behavior of the objectives themselves , rather than finite-sample effects . Upon observing samples from this model , our objective is thus to learn a feature embedder Φ and classifier2 β̂ to minimize the risk on an unseen environment e : Re ( Φ , β̂ ) : = E ( x , y ) ∼pe [ ` ( σ ( β̂TΦ ( x ) ) , y ) ] . The function ` can be any loss appropriate to classification : in this work we consider the logistic and the 0-1 loss . Note that we are not hoping to minimize risk in expectation over the environments ; this is already accomplished via ERM or distributionally robust optimization ( DRO ; Bagnell 2005 ; Ben-Tal et al . 2009 ) . Rather , we hope to extract and regress on invariant features while ignoring environmental features , such that our predictor generalizes to all unseen environments regardless of their parameters . In other words , the focus is on minimizing risk in the worst-case . We refer to the predictor which will minimize worst-case risk under arbitrary distribution shift as the optimal invariant predictor . To discuss this formally , we define precisely what we mean by this term . Definition 1 . Under the model described by Equations 1-3 , the optimal invariant predictor is the predictor defined by the composition of a ) the featurizer which recovers the invariant features and b ) the classifier which is optimal with respect to those features : Φ∗ ( x ) : = [ I 0 0 0 ] ◦ f−1 ( x ) = [ zc ] , β̂∗ : = [ βc β0 ] : = [ 2µc/σ 2 c log η1−η ] . Observe that this definition closely resembles Definition 3 of Arjovsky et al . ( 2019 ) ; the only difference is that here the optimal invariant predictor must recover all invariant features . As Arjovsky et al . ( 2019 ) do not posit a data model , the concept of recovering “ all invariant features ” is not welldefined for their setting ; technically , a featurizer which outputs the empty set would elicit an invariant predictor , but this would not satisfy the above definition . The classifier β̂∗ is optimal with respect to the invariant features and so it achieves the minimum possible risk without using environmental features . Observe that the optimal invariant predictor is distinct from the Bayes classifier ; the Bayes classifier uses environmental features which are informative of the label but non-invariant ; the optimal invariant predictor explicitly ignores these features . With the model defined , we can informally present our results ; we defer the formal statements to first give a background on the IRM objective in the next section . With a slight abuse of notation , we identify a predictor by the tuple Φ , β̂ which parametrizes it . First , we show that the usefulness of IRM exhibits a “ thresholding ” behavior depending on E and de : Theorem 3.1 ( Informal , Linear ) . For linear f , consider solving the IRM objective to learn a linear Φ with invariant optimal classifier β̂ . If E > de , then Φ , β̂ is precisely the optimal invariant predictor ; it uses only invariant features and generalizes to all environments with minimax-optimal risk . If E ≤ de , then Φ , β̂ relies upon non-invariant features . In fact , when E ≤ de it is even possible to learn a classifier solely relying on environmental features that achieves lower risk on the training environments than the optimal invariant predictor : 2Following the terminology of Arjovsky et al . ( 2019 ) , we refer to the regression vector β̂ as a “ classifier ” and the composition of Φ , β̂ as a “ predictor ” . Theorem 3.2 ( Informal , Linear ) . For linear f and E ≤ de there exists a linear predictor Φ , β̂ which uses only environmental features , yet achieves lower risk than the optimal invariant predictor . Finally , in the non-linear case , we show that IRM fails unless the training environments approximately “ cover ” the space of possible environments , and therefore it behaves similarly to ERM : Theorem 3.3 ( Informal , Non-linear ) . For arbitrary f , there exists a non-linear predictor Φ , β̂ which is nearly optimal under the penalized objective and furthermore is nearly identical to the optimal invariant predictor on the training distribution . However , for any test environment with a mean sufficiently different from the training means , this predictor will be equivalent to the ERM solution on nearly all test points . For test distributions where the environmental feature correlations with the label are reversed , this predictor has almost 0 accuracy . Extensions to other objectives . Many follow-up works have suggested alternatives to IRM—some are described in the next section . Though these objectives perform better on various baselines , there are few formal guarantees and no results beyond the linear case . Due to their collective similarities , we can easily derive corollaries which extend every theorem in this paper to these objectives , demonstrating that they all suffer from the same shortcomings . Appendix E contains example corollaries for each of the results presented in this work .

Summary: * The work gives extended theoretical analysis on the effect of invariant risk minimization scheme, which is an increasingly popular framework for robust prediction. The work considerably extends the results in the original IRM paper. The results seem reasonable, and clarify implausible beliefs on the framework. The intuitions and proving techniques could also inspire new methods to improve the current framework.