Invariant Causal Representation Learning for Out-of-Distribution Generalization

1 INTRODUCTION . Modern machine learning algorithms still lack robustness , and may fail to generalize outside of a specific training distribution because they learn easy-to-fit spurious correlations which are prone to change between training and testing environments . We recall the widely used example of classifying images of camels and cows ( Beery et al. , 2018 ) . Here , the training dataset has a selection bias , i.e. , many pictures of cows are taken on green pastures , while most pictures of camels happen to be in deserts . After training , it is found that the model builds on spurious correlations , i.e. , it relates green pastures with cows and deserts with camels , and fails to recognize images of cows on the beach . To address this problem , a natural idea is to identify which features of the training data present domain-varying spurious correlations with labels and which features describe true correlations of interest that are stable across domains . In the example above , the former are the features describing the context ( e.g. , pastures and deserts ) , whilst the latter are the features describing the animals ( e.g. , animal shape ) . By exploiting the varying degrees of spurious correlation naturally present in training data collected from multiple environments , one can try to identify stable features and build invariant predictors . Invariant risk minimization ( IRM ) seeks to find data representations ( Arjovsky et al. , 2019 ) or features ( Rojas-Carulla et al. , 2018 ) for which the optimal predictor is invariant across all environments . The general formulation of IRM is a challenging bi-leveled optimization problem , and theoretical guarantees require constraining both data representations and classifiers to be linear ( Arjovsky et al. , 2019 , Theorem 9 ) , or considering the special case of feature selection ( Rojas-Carulla et al. , 2018 , Theorem 4 ) . Ahuja et al . ( 2020a ) study the problem from the perspective of game theory , with an approach termed invariant risk minimization games ( IRMG ) . They show that the set of Nash equilibria for a proposed game is equivalent to the set of invariant predictors for any finite number of environments , even with nonlinear data representations and nonlinear classifiers . However , these theoretical results in the nonlinear setting only guarantee that one can learn invariant predictors from training environments , but do not guarantee that the learned invariant predictors can generalize well across all environments including unseen testing environments . We propose invariant Causal Representation Learning ( iCaRL ) , a novel approach that enables outof-distribution ( OOD ) generalization in the nonlinear setting ( i.e. , nonlinear representations and nonlinear classifiers1 ) . We achieve this by extending and using methods from representation learning and graphical causal discovery . In more detail , we first introduce our main general assumption : when conditioning on the target ( e.g. , labels ) and the environment ( represented as an index ) , the prior over the data representation ( i.e. , a set of latent variables encoding the data ) belongs to a general exponential family . Unlike the conditionally factorized prior assumed in recent identifiable variational autoencoders ( iVAE ) ( Khemakhem et al. , 2020a ) , this is a more flexible conditionally non-factorized prior , which can actually capture complicated dependences between the latent variables . We then extend iVAE to the case in which the latent variable prior belongs to such a general exponential family . The combination of this result and the previous general assumption allows us to a guarantee that the data representation can be identified up to simple transformations . We then theoretically show that the direct causes of the target can be fully discovered by analyzing all possible graphs in a structural equation model setting . Once they are discovered , the challenging bi-leveled optimization problem in IRM and IRMG can be reduced to two simpler independent optimization problems , that is , learning the data representation and learning the optimal classifier can be performed separately . This leads to a practical algorithm and enables us to obtain generalization guarantees in the nonlinear setting . Overall , we make a number of key contributions : ( 1 ) We propose a general framework for out-ofdistribution generalization in the nonlinear setting with the theoretical guarantees on both identifiability and generalizability ; ( 2 ) We propose a general assumption on the underlying causal diagram for prediction ( Assumption 1 and Fig . 1c ) , which covers many real-world scenarios ( Section 3.2 ) ; ( 3 ) We propose a general assumption on the prior over the latent variables ( Assumption 2 ) , i.e. , a more flexible conditionally non-factorized prior ; ( 4 ) We prove that an extended iVAE with this conditionally non-factorized prior is also identifiable ( Theorems 1 , 2 & 3 ) ; ( 5 ) We prove that our framework has the theoretical guarantees for OOD generalization in the nonlinear setting ( Proposition 1 ) . 2 PRELIMINARIES . 2.1 IDENTIFIABLE VARIATIONAL AUTOENCODERS . Variational autoencoders ( VAEs , see Appendix B ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) lack identifiability guarantees . Consider a VAE model where O ∈ Rd stands for the observed variables ( data ) and X ∈ Rn for the latent variables . Khemakhem et al . ( 2020a ) show that a VAE with an unconditional prior distribution pθ ( X ) over the latent variables is unidentifiable . However , they also show that it is possible to obtain an identifiable model if one posits a conditionally factorized prior distribution over the latent variables , pθ ( X|U ) , where U ∈ Rm is an additional observed variable ( Hyvärinen et al. , 2019 ) . Specifically , let θ = ( f , T , λ ) ∈ Θ be the parameters of the conditional generative model pθ ( O , X|U ) = pf ( O|X ) pT , λ ( X|U ) , ( 1 ) where pf ( O|X ) = p ( O − f ( X ) ) in which is an independent noise variable with probability density function p ( ) . Importantly , the prior pT , λ ( X|U ) is assumed to be conditionally factorial , where each element of Xi ∈ X has a univariate exponential family distribution given U . The conditioning onU is through an arbitrary function λ ( U ) ( e.g. , a neural net ) that outputs the individual exponential family parameters λi ( U ) for each Xi . The prior probability density thus takes the form pT , λ ( X|U ) = ∏ i Qi ( Xi ) /Zi ( U ) exp [ ∑k j=1 Ti , j ( Xi ) λi , j ( U ) ] , ( 2 ) where Qi is the base measure , Xi the i-th dimension of X , Zi ( U ) the normalizing constant , Ti = ( Ti,1 , . . . , Ti , k ) the sufficient statistics , λi ( U ) = ( λi,1 ( U ) , . . . , λi , k ( U ) ) the corresponding natural parameters depending on U , and k the dimension of each sufficient statistic that is fixed in advance . It is worth noting that this prior is restrictive as it is factorial and therefore can not capture dependencies . As in VAEs , the model parameters are estimated by maximizing the corresponding evidence lower bound , LiVAE ( θ , φ ) : =EpD [ Eqφ ( X|O , U ) [ log pf ( O|X ) + log pT , λ ( X|U ) − log qφ ( X|O , U ) ] ] , ( 3 ) where we denote by pD the empirical data distribution given by the dataset D = { ( O ( i ) , U ( i ) ) } N i=1 and qφ ( X|O , U ) denotes an approximate conditional distribution for X given by a recognition 1In fact , we are not restricted to the classification case and allow the target to be either continuous or categorical , which will be formally defined in Section 2.2. network with parameters φ . This approach is called identifiable VAE ( iVAE ) . Most importantly , it can be proved that under the conditions stated in Theorem 2 of ( Khemakhem et al. , 2020a ) , iVAE can identify the latent variablesX up to a permutation and a simple componentwise transformation , see Appendix F . 2.2 INVARIANT RISK MINIMIZATION . Arjovsky et al . ( 2019 ) introduced invariant risk minimization ( IRM ) , whose goal is to construct an invariant predictor f that performs well across all environments Eall by exploiting data collected from multiple environments Etr , where Etr ⊆ Eall . Technically , they consider datasets De : = { ( oei , yei ) } ne i=1 from multiple training environments e ∈ Etr , where oei ∈ O ⊆ Rd is the input observation and its corresponding label is yei ∈ Y ⊆ Rs.2 The datasetDe , collected from environment e , consists of examples identically and independently distributed according to some probability distribution P ( Oe , Y e ) . The goal of IRM is to use these multiple datasets to learn a predictor Y = f ( O ) that performs well for all the environments . Here we define the risk reached by f in environment e as Re ( f ) = EOe , Y e [ ` ( f ( Oe ) , Y e ) ] , where ` ( · ) is a loss function . Then , the invariant predictor can be formally defined as follows : Definition 1 ( Invariant Predictor ( Arjovsky et al. , 2019 ) ) . We say that a data representation Φ ∈ HΦ : O → C elicits an invariant predictor w ◦ Φ across environments E if there is a classifier w ∈ Hw : C → Y simultaneously optimal for all environments , that is , w ∈ arg minw̄∈Hw Re ( w̄◦Φ ) for all e ∈ E , where ◦ means function composition . Mathematically , IRM can be phrased as the following constrained optimization problem : min Φ∈HΦ , w∈Hw ∑ e∈Etr Re ( w ◦ Φ ) s.t . w ∈ arg min w̄∈Hw Re ( w̄ ◦ Φ ) , ∀e ∈ Etr . ( 4 ) Since this is a generally infeasible bi-leveled optimization problem , Arjovsky et al . ( 2019 ) rephrased it as a tractable penalized optimization problem by transfering the inner optimization routine to a penalty term . The main generalization result ( Theorem 9 in Arjovsky et al . ( 2019 ) ) states that if both Φ and w come from the class of linear models ( i.e. , HΦ = Rn×n and Hw = Rn×1 ) , under certain conditions on the diversity of training environments ( Assumption 8 in Arjovsky et al . ( 2019 ) ) and the data generation , the invariant predictor w ◦ Φ obtained by solving Eq . ( 4 ) remains invariant in Eall . 3 PROBLEM SETUP . 3.1 A MOTIVATING EXAMPLE . In this section , we extend the example which was introduced by Wright ( 1921 ) and discussed by Arjovsky et al . ( 2019 ) , and provide a further in-depth analysis . Model 1 . Consider a structural equation model ( SEM ) with a discrete environment variable E that modulates the noises in the structural assignments connecting the other variables ( cf . Fig . 1a below ) : X1 ← Gaussian ( 0 , σ1 ( E ) ) , Y ← X1 + Gaussian ( 0 , σ2 ( E ) ) , X2 ← Y + Gaussian ( 0 , σ3 ( E ) ) , where Gaussian ( 0 , σ ) denotes a Gaussian random variable with zero mean and standard deviation σ , and σ1 , . . . , σ3 are functions of the value e ∈ Eall taken by the environment variable E. To ease exposition , here we consider the simple scenario in which Eall only contains all modifications varying the noises of X1 , X2 and Y within a finite range , i.e. , σi ( e ) ∈ [ 0 , σ2max ] . Then , to predict Y from ( X1 , X2 ) using a least-square predictor Ŷ e = α̂1Xe1 + α̂2X e 2 for environment e , we can • Case 1 : regress from Xe1 , to obtain α̂1 = 1 and α̂2 = 0 , • Case 2 : regress from Xe2 , to obtain α̂1 = 0 and α̂2 = σ1 ( e ) +σ2 ( e ) σ1 ( e ) +σ2 ( e ) +σ3 ( e ) , • Case 3 : regress from ( Xe1 , X e 2 ) , to obtain α̂1 = σ3 ( e ) σ2 ( e ) +σ3 ( e ) and α̂2 = σ2 ( e ) σ2 ( e ) +σ3 ( e ) . In the generic scenario ( i.e. , σ1 ( e ) 6= 0 , σ2 ( e ) 6= 0 , and σ3 ( e ) 6= 0 ) , the regression using X1 in Case 1 is an invariant correlation : it is the only regression whose coefficients do not vary with e. By contrast , the regressions in both Case 2 and Case 3 have coefficients that depend on e. Therefore , only the invariant correlation in Case 1 will generalize well to new test environments . 2The setup applies to both continuous and categorical data . If any observation or label is categorical , we one-hot encode it . Another way to understand Model 1 is through its graphical representation3 , as shown in Fig . 1a . We treat the environment as a random variable E , where E could be any information specific to the environment ( Storkey , 2009 ; Peters et al. , 2015 ; Zhang et al. , 2017 ; Huang et al. , 2020 ) . For simplicity , we let E be the environment index , i.e. , E ∈ { 1 , . . . , N } , where N is the number of training environments . A more realistic version appearing in many settings is shown in Fig . 1b , where the true variables { X1 , X2 } are unobserved and we can only observe their transformationO . In this case , Invariant Causal Prediction ( ICP ) ( Peters et al. , 2015 ) will fail when applied toO , even when Y is not affected by E ( i.e. , the edge E → Y is removed ) . The reason is that each variable ( i.e. , each dimension ) ofO is jointly influenced by both X1 and X2 so that ICP is unable to find the variables containing the information only about X1 by searching for a subset of variablesO . By contrast , both IRM and IRMG work , as long as the transformation is linear . These findings are also empirically illustrated in Section 5.1 . We now go even further and consider a more general causal graph in which Y can have more than one parent or child .

Summary: This paper proposes a general framework for so-called "out-of-distribution generalization" that can handle non-linear associations and causal links between variables. It provides proofs for identifiability and generalization. It proposes a general causal model that is sensible for many prediction problems. It proposes a conditionally non-factorized prior. It provides a tractable method that can be used in practice.