Latent Causal Invariant Model

1 INTRODUCTION . Current data-driven deep learning models , revolutionary in various tasks though , heavily rely on i.i.d data to exploit all types of correlations to fit data well . Among such correlations , there can be spurious ones corresponding to biases ( e.g. , selection or confounding bias due to coincidence of the presence of the third factor ) inherited from the data provided . Such data-dependent spurious correlations can erode the ( i ) interpretability of decision-making , ( ii ) ability of out-of-distribution ( OOD ) generalization , i.e. , extrapolation from observed to new environments , which is crucial especially in safety-critical tasks such as healthcare , and ( iii ) robustness to small perturbation ( Goodfellow et al. , 2014 ) . Recently , there is a Renaissance of causality in machine learning , expected to pursue causal prediction ( Schölkopf , 2019 ) . The so-called “ causality ” is pioneered by Judea Pearl ( Pearl , 2009 ) , as a mathematical formulation of this metaphysical concept grasped in the human mind . The incorporation of a priori about cause and effect endows the model with the ability to identify the causal structure which entails not only the data but also the underlying process of how they are generated . For causal prediction , the old-school methods ( Peters et al. , 2016 ; Bühlmann , 2018 ) causally related the output label Y to the observed input X , which however is NOT conceptually reasonable in scenarios with sensory-level observed data ( e.g . modeling pixels as causal factors of Y does not make much sense ) . For such applications , we rather adopt the manner in Bengio et al . ( 2013 ) ; Biederman ( 1987 ) to relate the causal factors of Y to unobserved abstractions denoted by S , i.e. , Y ← fy ( S , εy ) via mechanism fy . We further assume existence of additional latent components denoted as Z , that together with S generates the input X via mechanism fx as X ← fx ( S , Z , εx ) . Taking image classification as an example , the S and Z respectively refer to object-related abstractions ( e.g. , contour , texture , color ) and contextual information ( e.g. , light , view ) . Such an assumption is similarly adopted in the literature of nonlinear Independent Components Analysis ( ICA ) ( Hyvarinen and Morioka , 2016 ; Hyvärinen et al. , 2019 ; Khemakhem , Kingma and Hyvärinen , 2020 ; Teshima et al. , 2020 ) and latent generative models ( Suter et al. , 2019 ) , which are however without separation of output ( y ) -causative factors ( a.k.a , S ) and other correlating factors ( a.k.a , Z ) that can both be learned in data-fitting process . We encapsulate these assumptions into a novel causal model , namely Latent Causal Invariance Model ( LaCIM ) as illustrated in Fig . 1 , in which we assume the structural equations fx ( associated with S , Z → X ) , fy ( associated with S → Y ) to be the Causal Invariant Mechanisms ( CIMe ) that hold under any circumstances with P ( S , Z ) allowed to be varied across domains . The incorporation of these priories can explain the spurious correlation embedded in the back-door path from Z to Y ( contextual information to the class label in image classification ) . To avoid learning spurious correlations , our goal is to identify the intrinsic CIMe fx , fy . Specifically , we first prove the identifiability ( i.e. , the possibility to be precisely inferred up to an equivalence relation ) of the CIMe . Notably , far beyond the scope in existing literature ( Khemakhem , Kingma and Hyvärinen , 2020 ) , our results can implicitly , and are the first to disentangle the output-causative factors ( a.k.a , S ) from others ( a.k.a , Z ) for prediction , to ensure the isolation of undesired spurious correlation . Guaranteed by such , we propose to estimate the CIMe by extending the Variational Auto-encoder ( VAE ) ( Kingma and Welling , 2014 ) to the supervised scenario . For OOD prediction , we propose to optimize over latent space under the identified CIMe . To verify the correctness of our identifiability claim , we conduct a simulation experiment . We further demonstrate the utility of our LaCIM via high explainable learned semantic features , improved prediction power on various OOD scenarios ( including tasks with confounding and selection bias , healthcare ) , and robustness on security . We summarize our contribution as follows : ( i ) Methodologically , we propose in section 4.1 a latent causal model in which only a subset of latent components are causally related to the output , to avoid spurious correlation and benefit OOD generalization ; ( ii ) Theoretically , we prove the identifiability ( in theorem 4.3 ) of CIMe fx , fy from latent variables to observed data , which disentangles outputcausative factors from others ; ( iii ) Algorithmically , guided by the identifiability , we in section 4.3 reformulate Variational Bayesian method to estimate CIMe during training and optimize over latent space during the test ; ( iv ) Experimentally , LaCIM outperforms others in terms of prediction power on OOD tasks and interpretability in section 5.2 , and robustness to tiny perturbation in section 5.3 . 2 RELATED WORK . The invariance/causal learning proposes to learn the assumed invariance for transferring . For the invariance learning methods in Krueger et al . ( 2020 ) and Schölkopf ( 2019 ) , the “ invariance ” can refer to stable correlation rather than causation , which lacks the interpretability and impedes its generalization to a broader set of domains . For causal learning , Peters et al . ( 2016 ) ; Bühlmann ( 2018 ) ; Kuang et al . ( 2018 ) ; Heinze-Deml and Meinshausen ( 2017 ) assume causal factors as observed input , which is inappropriate for sensory-level observational data . In contrast , our LaCIM introduces latent components as causal factors of the input ; more importantly , we explicitly separate them into the output-causative features and others , to avoid spurious correlation . Further , we provide the identifiability claim of causal invariant mechanisms . In independent and concurrent works , Teshima et al . ( 2020 ) and Ilse et al . ( 2020 ) also explore latent variables in causal relation . As comparisons , Teshima et al . ( 2020 ) did not differentiate S from Z ; and Ilse et al . ( 2020 ) proposed to augment intervened data , which can be intractable in real cases . Other works which are conceptually related to us , as a non-exhaustive review , include ( i ) transfer learning which also leverages invariance in the context of domain adaptation ( Schölkopf et al. , 2011 ; Zhang et al. , 2013 ; Gong et al. , 2016 ) or domain generalization ( Li et al. , 2018 ; Shankar et al. , 2018 ) ; and ( ii ) causal inference ( Pearl , 2009 ; Peters et al. , 2017 ) which proposes a structural causal model to incorporate intervention via “ do-calculus ” for cause-effect reasoning and counterfactual learning ; ( iii ) latent generative model which also assumes generation from latent space to observed data ( Kingma and Welling , 2014 ; Suter et al. , 2019 ) but aims at learning generator in the unsupervised scenario . 3 PRELIMINARIES . Problem Setup & Notation LetX , Y respectively denote the input and output variables . The training data { De } e∈Etrain are collected from the set of multiple environments Etrain , where each domain e is associated with a distribution Pe ( X , Y ) over X × Y and De = { xei , yei , de } i∈ [ ne ] i.i.d∼ Pe with [ k ] : = { 1 , ... , k } for any k ∈ Z+ . The de∈ { 0 , 1 } m denotes the one-hot encoded domain index for e , where 1 ≤ m : = ∣∣Etrain∣∣ ≤ n : = ∑e∈Etrain ne . Our goal is to learn a model f : X 7→ Y that learns output ( y ) -causative factors for prediction and performs well on the set of all environments E ⊃ Etrain , which is aligned with existing OOD generalization works ( Arjovsky et al. , 2019 ; Krueger et al. , 2020 ) . We use respectively upper , lower case letter and Cursive letter to denote the random variable , the instance and the space , e.g. , a is an instance in the space A of random variable A . The [ f ] A denotes the f restricted on dimensions of A . The Sobolev space W k , p ( A ) contains all f such that ∫ A ∣∣∂Afα∣∣A=a∣∣pda < ∞ , ∀α ≤ k. Structural Causal Model . The structural causal model ( SCM ) is defined as the causal graph assigned with structural equations . The causal graph encodes the assumptions in missing arrows in a directed acylic graph ( DAG ) : G = ( V , E ) with V , E respectively denoting the node set and the edge set . The Pa ( k ) denotes the set of parent nodes of Vk for each Vk ∈ V and the X → Y ∈ E indicates the causal effect of X on Y . The structural equations { Vk ← fk ( Pa ( k ) , εk ) } Vk∈V , quantify the causal effects shown in the causal graph G. By assuming independence among exogenous variables { εk } k , the Causal Markov Condition states that P ( { Vk = vk } Vk∈V ) = ΠkP ( Vk = vk|Pa ( k ) = pa ( k ) ) . A back-door path from Va to Vb is defined as a path that ends with an arrow pointing to Va ( Pearl , 2009 ) . 4 METHODOLOGY . We build our causal model associated with Causal Invariant Mechanism ( CIMe , i.e. , fx , fy ) and a priori about the generating process in section 4.1 , followed by our identifiability result for CIMe in section 4.2 . Finially , we introduce our learning method to estimate CIMe in section 4.3 . 4.1 LATENT CAUSAL INVARIANCE MODEL We introduce latent variables to model the abstractions/concepts that play as causal factors that generate the observed variables ( X , Y ) , which is more reasonable than assuming the X as the direct cause of Y in scenarios with sensory-level data . We explicitly separate the latent variables into two parts : the S and Z that respectively denote the y ( output ) -causative and y-non-causative factors , as shown by the arrow S → Y in Fig . 1 . Besides , the X and Y are respectively generated by S , Z and S , via structural equations ( with noise ) fx , fy , which are denoted as Causal Invariant Mechanisms ( CIMe ) that hold across all domains . The output Y denotes the label generated by human knowledge , e.g. , the semantic shape , the contour to discern the object , etc . Hence , we assume the Y as the outcome/effect of these high-level abstractions ( Biederman , 1987 ) rather than the cause ( detailed comparison with Y → S is left in supplementary 7.7.1 ) . We call the model associated with the causal graph in Fig . 1 as Latent Causal Invariance Model ( LaCIM ) , with formal definition given in Def . 4.1 . As an illustration , we consider the image classification in which X , Y denote the image and the class label . Instead of X , i.e. , the pixels , it is more reasonable to assume the causal factors ( of X , Y ) as latent concepts ( S , Z ) that can denote light , angle , the shape of the object to generate X following the physical mechanisms . Among these concepts , only the ones that are causally related to the object , i.e. , S ( e.g. , shape ) are causal factors of the object label , i.e. , Y . Following the physical or natural law , the mechanisms S , Z → X , S → Y invariantly hold across domains . The S : = Rqs , Z : = Rqz denote the space of S , Z , with Pe ( S , Z ) ( that characterizes the correlation between S and Z ) varying across E ( e.g. , the object is more associated with a specific scene than others ) . We assume that the y-non-causative factor ( i.e. , Z ) is associated with ( but not causally related to ) S , Y through the confounder C , which is allowed to take a specific value for each sample unit . Therefore , the back-door path Z←C→S→Y induces the correlation between Z and Y in each single domain . Rather than invariant causation , this correlation is data-dependent and can vary across domains , which is known as “ spurious correlation ” . In real applications , this spurious correlation corresponds to the bias inherited from data , e.g . the contextual information in object classification . This domain-specific S-Z correlation , can be explained by the source variable D , which takes a specific and fixed value for each domain and functions the prior of distribution of the confounder C , as illustrated in Fig . 1 . This source variable D can refer to attributes/parameters that characterize the distribution of S , Z in each domain . When such attributes are unobserved , we use the domain index as a substitute . Consider the cat/dog classification task as an illustration , the animal in each image is either associated with the snow or grass . The S , Z respectively denote the concepts of animals and scenes . The D denotes the sampler , which can be described by the proportions of scenes associated with the cat and those associated with the dog . The D generates the C that denotes the ( time , weather ) to go outside and collect samples . Since each sampler may have a fixed pattern ( e.g . gets used to going out in the sunny morning ( or in the snowy evening ) ) , the data he/she collects , may have sample selection bias ( e.g . with dogs ( cats ) more associated with grass ( snow ) in the sunny morning ( or snowy evening ) ) . In this regard , the scene concepts Z can be correlated with the animal concepts S , and also the label Y . Definition 4.1 ( LaCIM ) . The Latent Causal Invariance Model ( LaCIM ) for e ∈ E is defined as a SCM characterized by ( i ) the causal graph , i.e. , the G = ( V , E ) with V = { C , S , Z , X , Y } and E = { C → S , C → Z , Z → X , S → X , S → Y } ; and ( ii ) structural equations with causal mechanisms { fc , fz , fs , fx , fy } embodying the quantitative causal information : c ← fc ( de , εc ) , z ← fz ( c , εz ) , s← fs ( c , εs ) ; x← fx ( s , z , εx ) ; y ← fy ( s , εy ) , in which { εc , εz , εs , εx , εy } are independent exogenous variables that induce pfc ( c|de ) , pfz ( z|c ) , pfs ( s|c ) , pfx ( x|s , z ) , pfy ( y|s ) . The CIMe fx , fy are assumed to be invariant across E . We call the environment-dependent parts : Pe ( S , Z ) and Pe ( S , Z|X ) as S , Z-prior and S , Z-inference in the following . Remark 1 . We denote LaCIM-ds and LaCIM-d as two versions of LaCIM , with the source variable ds with practical meaning ( e.g . attributes or parameters of P ( S , Z ) ) observed or not . The observation of ds can be possible in some applications ( e.g. , age , gender that characterizes population in medical diagnosis ) . As for the LaCIM-d with ds unobserved , we use domain index D as a substitute . Denote C as the space of C. We assume that the C is finite union of disjoint sets { Cr } Rr=1 , i.e . C : = ∪Rr=1Cr , such that for any cr , i 6= cr , j ∈ Cr , it holds that p ( s , z|cr , i ) = p ( s , z|cr , j ) for any ( s , z ) . Returning to the cat/dog classification example , the C denotes the range of time to collect samples , i.e. , 00 : 00-24 : 00 . The C can be divided into several time periods C1 , ... , CR , such that the proportion of concepts of ( animal , scene ) given any c in the same period is unchanged , e.g. , the dog often comes up on the grass in the morning . Further , since p ( x , y|s , z ) = p ( x|s , z ) p ( y|s ) is invariant , we have for each Cr that p ( x , y|cr , i ) = ∫ p ( x , y|s , z ) p ( s , z|cr , i ) dsdz = ∫ p ( x , y|s , z ) p ( s , z|cr , j ) dsdz = p ( x , y|cr , j ) for any ( x , y ) . That is , the { p ( x , y|cr } cr∈Cr for each ( x , y ) collapse to a single point , namely p ( x , y|cr ) . In this regard , we have pe ( x , y ) : = p ( x , y|de ) = ∑R r=1 p ( x , y|cr ) p ( cr|de ) . Besides , we assume the Additive Noise Model ( ANM ) for X , i.e. , fx ( s , z , εx ) = f̂x ( s , z ) + εx ( we replace f̂x with fx without loss of generality ) , which has been widely adopted to identify the causal factors ( Janzing et al. , 2009 ; Peters et al. , 2014 ; Khemakhem , Kingma and Hyvärinen , 2020 ) . We need to identify the CIMe ( i.e. , fx , fy ) , guaranteed by the identifiability that ensures the learning method to distinguish S from Z to avoid spurious correlation , as presented in section 4.2 . Traditionally speaking , the identifiability means the parameter giving rise to the observational distribution pθ ? ( x , y|de ) can be uniquely determined , i.e. , pθ ( x , y|de ) = pθ̃ ( x , y|de ) =⇒ θ = θ̃ . Instead of strict uniqueness , we rather identify an equivalent class of θ ? ( in Def . 4.2 ) that suffices to disentangle the y-causative features S from Z to avoid learning spurious correlation . To achieve this goal , we first narrow our interest in case when p ( s , z|c ) is exponential family in Eq . ( 1 ) , in which we can respectively identify the S , Z up to linear and point-wise transformations given by theorem 4.3 ; then we generalize to any p ( s , z|c ) as long as it belongs to Sobolev space , as explained in theorem 4.4 . A reformulated VAE is proposed to learn the CIMe practically . For generalization , note that the gap between two environments in terms of prediction given x , i.e. , ∣∣Epe2 [ Y |X = x ] − Epe1 [ Y |X = x ] ∣∣ = ∫S ∣∣pe2 ( s|x ) − pe1 ( s|x ) ∣∣pfy ( y|s ) ds , is mainly due to the inconsistency of S , Z-inference , i.e. , pe ( s , z|x ) 6= pe′ ( s , z|x ) for e′ 6= e ( for details please refer to theorem 7.1 in supplement 7.1 ) . Therefore , one can not directly apply the trained { pe ( s , z|x ) , pe ( y|x ) } e∈Etrain to the inference model of new environment , i.e . pe ′ ( s , z|x ) , pe′ ( y|x ) for e′ /∈ Etrain . To solve this problem and generalize to new environment , we note that since pfx ( x|s , z ) and pfy ( y|s ) are shared among all environments , we propose to inference s , z that give rise to the test sample x via maximizing the identified pfx ( x|s , z ) , as a pseudo-likelihood of x given ( s , z ) , rather than using S , Z-inference model which is inconsistent among environments . Then , we feed estimated s into invariant predictor pfy ( y|s ) for prediction .

Summary: This paper proposes a VAE based model for learning latent causal factors given data from multiple domains. Similar to [Kingma and Hyv¨arinen, 2020], it utilizes additional labels as supervision signals and learns the model using a Bayesian optimization approach given a fixed hypothetical causal structure. The identifiability is obtained by assuming the casual mechanism to be domain invariant, which is partially supported by some empirical experiments.