Recovering Latent Causal Factor for Generalization to Distributional Shifts

1 Introduction . Current data-driven deep learning models , revolutionary in various tasks though , often exploit all types of correlations to fit data well . Among such correlations , there can be spurious ones corresponding to biases ( e.g. , confounding bias due to the presence of a third unseen factor ) inherited from the data provided . Such data-dependent spurious correlations can erode the prediction power on unseen domains with distributional shifts , which can cause serious consequences especially in safety-critical tasks such as healthcare . Recently , there is a Renaissance of causality in machine learning , expected to pursue causal relationships [ 59 ] to achieve stable generalization across domains . The so-called area of “ causality ” is pioneered by Structural Causal Models [ 51 ] , as a mathematical formulation of this metaphysical concept grasped in the human mind . The incorporation of these human priors about cause and effect endows the model with the ability to identify the causal structure [ 51 ] which entails not only the data but also the underlying process of how they are generated . To achieve causal modeling , the old-school methods [ 52 , 10 ] directly causally related the output label Y to a subset of covariates X , which is however not conceptually reasonable in applications with sensory-level data ( e.g . model pixels as causal factors of the output does not make sense in image classification [ 11 ] ) . ∗Corresponding author †Work done during an internship at Microsoft Research Asia . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . For such applications , we rather adopt the manner of human visual perception [ 8 , 9 , 80 ] to causally relate the label Y to unobserved abstractions denoted by S , i.e. , Y ← S. We further assume the existence of another non-causal latent factor ( of Y ) denoted as Z , that together with S generate the input X : X ← ( S , Z ) . Such an assumption is similarly adopted in the literature [ 25 , 27 , 35 , 75 , 71 ] . To model shifts across domains , we allow Z to be spuriously correlated with S ( hence also the output ) , as marked by the bidirected arrow in Fig . 1 ( a ) . Taking image classification as an example , the S and Z respectively refer to object-related abstractions ( e.g. , contour , texture ) and contextual information ( e.g. , background , view ) . Due to this correlation , the model can learn contextual information into prediction , which may fail to generalize to the domain such that this correlation is broken . We encapsulate above assumptions into the skeleton illustrated in Fig . 1 ( a ) , in which the spurious correlation between S and Z varies across domains , as marked by the red bi-directed arrow in Fig . 1 ( b ) . Taking a closer inspection , such a domain-dependent spurious correlation is governed by an auxiliary domain variable D in Fig . 1 ( c ) , which causes the domain shifts . We call the set of causal models augmented with D as Latent Causal Invariance Models ( LaCIM ) . Here , the “ Causal Invariance ” refers to P ( Y |S ) , which together with P ( X|S , Z ) , can be proved to be stable to the shifts across domains , under the assumptions embedded in the causal structure of LaCIM . Equipped with such an invariance , we prove that the S and the ground-truth predictor : P ( Y |s ? ) for x generated from ( s ? , z ? ) , are identifiable up to transformations that do not mix the non-causal information . Under such an identifiability guarantee , we propose to learn the P ( Y |S ) and P ( X|S , Z ) by reformulating the Variational Auto-encoder ( VAE ) [ 37 ] to fit the joint distribution of the input and output variables from training domains . During the test stage , we first infer the value of S by optimizing the estimated P ( X|S , Z ) over latent space , followed by the learned P ( Y |S ) for prediction . We first use simulated data to verify the correctness of the identifiability claim . Then , to demonstrate the utility , we test our approach on real-world data , consistently achieving better generalization to the new distribution ; besides , we find that our inferred causal factor can be concentrated in highly explainable semantic regions for the task of image classification . We summarize our contribution as follows : Methodologically ( in sec . 4.1 ) , we propose LaCIM in which the causal assumptions of two latent factors and the distributional shifts are incorporated ; Theoretically ( in theorem 4.4 ) , we prove the identifiability of the causal factor and the ground-truth predicting mechanism ; Algorithmically ( in sec . 4.3 ) , guided by the identifiability , we reformulate Variational Bayesian method to learn P ( X|S , Z ) , P ( Y |S ) for prediction ; Experimentally ( in sec . 5.2 ) , our approach generalizes better to distributional shifts , compared with others . 2 Related Work . Causality for Domain Generalization . Due to its stable transferability , the concept of causality has been introduced in many recent works for domain generalization [ 39 , 59 , 52 , 10 , 40 , 21 , 68 ] . Most of these works learned the assumed ( causal ) invariance for generalizing to unseen domains . However , they suffer from either i ) lacking explicit causal modeling ; or ii ) inappropriate causal relations made for the output . Specifically , for i ) , the [ 39 , 59 ] are still data-driven methods to learn stable correlation ( i.e. , invariance ) without incorporating causal assumptions [ 51 ] beyond data , which may impede its generalization to a broader set of domains ; for ii ) , the [ 52 , 10 , 40 , 21 , 68 ] causally relate the output with covariates , which is inappropriate for sensory-level data . Our Specification . We explicitly incorporate the causal assumptions . Specifically , we introduce i ) latent factors and separate them into the causal and the non-causal factor ; ii ) the domain variable D , as a selecting mechanism to generate the varied S-Z correlation across domains . Such a causal modeling makes it possible to recover the causal factor S for generalization . In independent and concurrent works , [ 75 ] and [ 28 ] also explore latent variables in causal relation . As comparisons , [ 75 ] did not differentiate S from Z . The spurious correlation in [ 28 ] is limited in the correlation between domains and the output ; while it is allowed in our setting to exist in a single domain , which is more aligned with real scenarios , e.g. , the dog is more associated with grass than snow in a domain when most samples are collected in sunny morning . Other Conceptually Related Works : i ) transfer learning that leverages invariance in the context of domain adaptation [ 60 , 81 , 17 ] or domain generalization [ 43 , 63 ] ; ( ii ) causal inference [ 51 , 53 ] which builds structural causal models and define intervention ( a.k.a , “ do-calculus ” ) for cause-effect reasoning and counterfactual learning ; and ( iii ) latent generative model that assumes generation from latent space to observed data [ 37 , 71 ] but aims at learning generator in the unsupervised scenario . 3 Preliminaries . Problem Setting . Let X , Y respectively denote the input and output variables . The training data { De } e∈Etrain are collected from multiple environments e ∈ Etrain , where each e is associated with a distribution Pe ( X , Y ) over X × Y and De : = { xei , yei } i∈ [ ne ] i.i.d∼ Pe with [ k ] : = { 1 , ... , k } for any k ∈ Z+ . Our goal is to learn a robust predictor f : X → Y that only exploit the causal factor for prediction and generalize well to all domains E ⊃ Etrain . 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 spaceA of random variable A. ForA : = f ( X ) ∩B with B : = Rp [ i1 ] ×Rp [ i2 ] × ... ×Rp [ ik ] , the [ f ( x ) ] A denotes the f ( x ) restricted on dimensions of A , i.e. , [ f ( x ) ] A : = [ fi1 ( x ) , ... , fik ( x ) ] . The Sobolev space W k , p ( A ) contains all f such that ∫ A ∣∣∂Afα|A=a∣∣pdµ ( a ) < ∞ , ∀α ≤ k. Structural Causal Model . The structural causal model ( SCM ) is defined as a triplet M : = 〈G , F , P ( ε ) 〉 , in which i ) the causal structure G : = ( V , E ) ( V , E respectively denote the node and edge set ) is described by a directed acyclic graph ( DAG ) ; ii ) the structural equations F : = { fk } Vk∈V are autonomous , i.e. , intervening on Vk does not affect others , based on which we can define the dooperator and calculate the causal effect ; iii ) the P ( ε ) are probability measure for exogenous variables { εk } k. By assuming independence among { εk } k , we obtain according to Causal Markov Condition that each P that is compatible with G has P ( { Vk = vk } Vk∈V ) = ΠkP ( Vk = vk|Pa ( k ) = pa ( k ) ) . An acyclic directed mixed graph ( ADMG ) can further allow the existence of bidirectional arrows↔ , meaning the spurious correlation between two variables connected . 4 Methodology . We first incorporate the causal assumptions into LaCIM in sec . 4.1 . Under such assumptions , we identify the invariant distributions P ( X|S , Z ) and P ( Y |S ) , which are repectively dubbed as generative invariance and causal invariance that are robust to domain shifts . Equipped with these invariances , we in sec . 4.2 show that the causal factor can be identified without mixing information from non-causal one during prediction . Finally , we introduce our learning method in sec . 4.3 to estimate the P ( X|S , Z ) and P ( Y |S ) , which are respectively resorted in the inference and prediction that constitute a robust predictor during test stage . 4.1 Latent Causal Invariance Models . In this section , we introduce a set of structural causal models dubbed as Latent Causal Invariance Model ( LaCIM ) , which incorporates the causal assumptions mentioned above and also the source of distributional shifts . The corresponding causal structure of LaCIM is illustrated in Fig . 1 ( c ) , which we will introduce step-by-step from the skeleton in Fig . 1 ( a ) . Fig . 1 ( a ) . Specifically , the ADMG in Fig . 1 ( a ) introduces latent factors V : = { S , Z } to model the abstractions/concepts that generate the observed variables ( X , Y ) , as similarly assumed in unsupervised latent generative models [ 37 ] for image generation . Further , we explicitly separate the V into S and Z , with only S causally related to the label Y . In image classification , such a causal factor refers to the ( shape , contour ) of the object need to be classified ; while the image X is additionally affected by contextual factor such as light , view . Fig . 1 ( a ) → Fig . 1 ( b ) . In addition , we assume that S is spuriously correlated with Z , as marked by the red “ ↔ ” in Fig . 1 ( a ) . Such a spurious correlation corresponds to the bias inherited from data , e.g . the contextual information in image classification . Therefore , the magnitude of this correlation is distribution-dependent and thus can vary across domains . Statistically , the “ spurious correlation '' implicates the presence of a third unobserved ( we use dot circle to represent unobserved variables ) confounder , which is denoted as C in Fig . 1 ( b ) . The unblocked path from Z to Y induced by C can lead to learning the non-causal factor during data-fitting , which can degrade the performance on unseen domains if the correlation between this non-causal factor and the output is broken . Fig . 1 ( b ) → Fig . 1 ( c ) . Taking a further inspection in Fig . 1 ( b ) , the varying degree of correlation can be either due to the distributional shift of S , Z|C or of the C itself across domains ( we use red color to mean varied distributions ) . As both shifts are domain-dependent , we in Fig . 1 ( c ) ascribe them to a domain variable D , which causes the mutation of its children nodes ’ distribution , i.e. , S , Z and C. Such a domain variable has been similarly introduced in [ 69 , 68 ] to generate mutable variables . In our scenario , we do not require D to be observed ; rather , we only need the domain index d̃e ( one-hot encoded vector with length m : = |Etrain| ) . The set of SCMs augmented with D , with the SCM Markovian compatible to the DAG of C , S , Z , X , Y in Fig . 1 ( c ) , is dubbed as Latent Causal Invariance Models ( LaCIM ) that is formally defined as follows : Definition 4.1 ( LaCIM ) . The LaCIM denotes a set of SCMs augmented with the domain variable D , i.e. , { 〈Me , de〉 } e∈E , in which de denotes the value of D and Me : = 〈G , Fe , P ( ε ) 〉 for e. The G denotes the DAG restricted on C , S , Z , X , Y . For each environment/domain e , the Fe : = { fx , fy , fes , fez , fec } correspond to generating mechanism ofX , Y , S , Z , C , with fec ( εc ) : = gc ( εc , de ) , fes ( c , εs ) : = gs ( c , εs , d e ) and fez ( c , εz ) : = gz ( c , εz , d e ) from some gc , gs , gz . Remark 1 . Different from scenarios in which X generates [ 28 ] nor generated from Y [ 1 ] , we consider the scenario when the X and Y are generated concurrently , which can widely exist but ignored in the literature . For example , the clinicians are recording the disease status while implementing the ultrasound test at the same time , during medical diagnosis . As an illustration , we consider the following example , in which the distributional shifts caused by domain variable D can refer to sampling bias in data . Example 4.1 ( Sampling Bias ) . Consider the cat/dog classification , in which the animal in each image is either associated with the snow or grass . The D refers to the sampler , which generates the C that denotes time and weather to collect each sample . The S , Z respectively refer to the features of animals and context . Since each sampler may have a fixed sampling pattern ( e.g . gets used to going out in the sunny morning ( or in the snowy evening ) ) , the data one collects may have sampling bias : dogs ( cats ) more associated with grass ( snow ) in the sunny morning ( or snowy evening ) . The Def . 4.1 specifies the generating mechanisms across environments and how they differ . Equipped with such a specification , we can identify the invariant mechanisms that are stable to domain shifts : Proposition 4.2 ( Causal Invariance & Generative Invariance ) . For LaCIM in Def . 4.1 , the P ( Y |S ) and P ( X|S , Z ) are invariant to shifts across E , and are respectively denoted as Causal Invariance ( CI ) and Generative Invariance ( GI ) . Remark 2 . The generating process from latent variables to observed variables follows from physical law , e.g. , the shape , contour , color , view , light should satisfy physical constraints to generate a reasonable image . Therefore , it is naturally hold that such generating processes are invariant . The P ( X|S , Z ) and P ( Y |S ) can induce an invariant predicting mechanism . Specifically , for a new sample x← fx ( s ? , z ? , εx ) , y ← fy ( s ? , εy ) , we can first infer the causal factor s ? from pfx ( x|s , z ) by maximizing log-likelihood of pfx ( x|s , z ) over S ×Z and then feed the estimated s into pfy ( y|s ? ) for prediction . To ensure the robustness of such a two-step invariant prediction , we need to answer two following identifiability questions : 1 . Can the inferred causal factor S not mix the information of ( disentangled from ) others ? 2 . Can such an invariant predictor recover the ground-truth predictor P ( Y |s ? ) ? We will answer these questions in the subsequent section , followed by our learning methods to identify the causal factor and the causal/generative invariance for prediction .

Summary: To deal with the issue of distributional shifts between training and testing domains, the authors propose Latent Causal Invariance Models (LaCIM) consisting of both causal latent factors and non-causal latent factors and the extent of the correlations between them is governed by a domain variable. They theoretically show the identifiability of the causal latent factors and the ground-truth predicting mechanism in the proposed LaCIM. Based on the identifiability, they learn the model by reformulating VAE and then verify it on various real-world data.