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: This paper introduces the "Latent Causal Invariance Model" (LaCIM), a structural variational autoencoder-like approach for learning to isolate causal factors from "spurious" ones in prediction to improve generalization to new environments. The authors assume a latent generative structure for the data, with separate latent constructs for causal and contextual factors. Under strong parametric assumptions and requirements on the number of source environments the authors prove identifiability of the causal factors. Experiments on imaging datasets (2 benchmarks and a small medical dataset) show LaCIM outperforms some existing domain adaptation baselines as well as Invariant Risk Minimization (IRM). Further, qualitative assessment of images generated by ablating the causal and contextual factors show disentanglement of these factors.