Efficiently Disentangle Causal Representations

1 INTRODUCTION . Causal reasoning is a fundamental tool that has shown great impact in different disciplines ( Rubin & Waterman , 2006 ; Ramsey et al. , 2010 ; Rotmensch et al. , 2017 ) , and it has roots in work by David Hume in the eighteenth century ( Hume , 2003 ) and in classical AI ( Pearl , 2003 ) . Causality has been mainly studied from a statistical perspective ( Pearl , 2009 ; Peters et al. , 2016 ; Greenland et al. , 1999 ; Pearl , 2018 ) with Judea Pearl ’ s work on the causal calculus leading its statistical development . More recently , there has been a growing interest to integrate statistical techniques into machine learning to leverage their benefits . Welling raises a particular question about how to disentangle correlation from causation in machine learning settings to take advantage of the sample efficiency and generalization abilities of causal reasoning ( Welling , 2015 ) . Although machine learning has achieved important results on a variety of tasks like computer vision and games over the past decade ( e.g. , Mnih et al . ( 2015 ) ; Silver et al . ( 2017 ) ; Szegedy et al . ( 2017 ) ; Hudson & Manning ( 2018 ) ) , current approaches can struggle to generalize when the test data distribution is much different from the training distribution ( common in real applications ) . Further , these successful methods are typically “ data hungry ” , requiring an abundance of labeled examples to perform well across data distributions . In statistical settings , encoding the causal structure in models has been shown to have significant efficiency advantages . In support of the advantages of encoding causal mechanisms , Bengio et al . ( 2020 ) recently introduced an approach to disentangling causal relationships in end-to-end machine learning by comparing the adaptation speeds of separate models that encode different causal structures . With this as the baseline , in this paper , we propose a more efficient approach to learning disentangled representations with causal mechanisms , based on the difference of conditional probabilities in original and new distributions . The key idea is to approximate the difference with model ’ s generalization abilities so that it fits in standard machine learning framework and can be efficiently computed . In contrast to the state-of-the-art baseline approach , which relies on learner ’ s adaptation speed to new distribution , the proposed approach only requires evaluating the generalization ability of the model . Our method is based on the same assumption as the baseline that the conditional distribution P ( B|A ) does not change between the train and transfer distribution . This assumption can be explained with an atmospheric physics example of learning P ( A , B ) , where A ( Altitude ) causes B ( Temperature ) ( Peters et al. , 2016 ) . The marginal probability of A can be changed depending on , for instance , the place ( Switzerland to less mountainous country like Netherlands ) , but P ( B|A ) as the underlining causality mechanism does not change . Therefore , the causality structure can be inferred from the robustness of predictive models on out-of-distribution ( Peters et al. , 2016 ; 2017 ) . The proposed method is more efficient by omitting the adaptation process , and is more robust when the marginal distribution is complicated . We provide theoretical explanation and experimental verification for the advantage of the proposed method . Our experiments show the proposed technique is 1.9– 11.0× more sample efficient and 9.4–32.4× quicker than measuring adaptation speed on various tasks . We also argue that the proposed approach has less hyper parameters and it is straight-forward to implement the approach within the standard machine learning workflows . Our contributions can be summarized as follows . • We propose an efficient approach to disentangling representations for causal mechanisms by measuring generalization . • We theoretically prove that the proposed estimators can identify causal direction and disentangles causal mechanisms . • We empirically show that the proposed approach is significantly quicker and more sample efficient for various tasks . Sample efficiency is important when data size in transfer distribution is small . 2 APPROACH . To begin , we reflect on the tasks and the disentangling approach ( as baseline ) described by previous work ( Bengio et al. , 2020 ) . The invariance of conditional distribution for the correct causal direction P ( B|A ) is the key assumption in their work , and we also follow it in this work . We notice that their baseline approach compares the adaptation speed of models on a transfer data distribution , and hence requires significant time for adaptation . We propose an approach to learn causality mechanisms by directly measuring the changes in conditional probabilities before and after intervention for both P ( B|A ) and P ( A|B ) . Further , we optimize the proposed approach to use generalization loss rather than a divergence metric—because loss can be directly measured in standard machine learning workflows—and we show that it is likely to correctly predict causal direction and disentangle causal mechanisms . 2.1 CAUSALITY DIRECTION PREDICTION . As the first step towards learning disentangled representations for causal mechanisms , we start with the binary classification task . Given two discrete variables A and B , we want to determine whether A causes B , or vice-versa . We assume noiseless dynamics , and A and B do not have hidden confounders . The training ( transfer ) data contains samples ( a , b ) from training ( transfer ) distribution , P1 ( P2 ) . The baseline approach defines models that factor the joint distribution P ( A , B ) into two causality directions PA→B ( A , B ) = PA→B ( B|A ) PA→B ( A ) and PB→A = PB→A ( A|B ) PB→A ( B ) . It then compares their speed of adaptation to transfer distribution . Intuitively , the factorization with correct causality direction should adapt more quickly to the transfer distribution . Suppose A→ B is the ground-truth causal direction . For the correct factorization , they assume that conditional distribution PA→B ( B|A ) does not change between the train and transfer distributions , so that only the marginal PA→B ( A ) needs adaptation . In contrast , for the factorization with incorrect causality direction , both the conditional PB→A ( A|B ) and marginal distributions PB→A ( B ) need adaptation . They analyze that updating a marginal distribution P ( A ) A→B is likely to have lower sample complexity than the conditional distribution PB→A ( A|B ) ( only a part of PB→A ( A , B ) ) because the later has more parameters . Therefore , the model with correct factorization will adapt more quickly , and causality direction can be predicted from adaptation speed . To leverage this observation , the baseline method defines a meta-learning objective . Let LA→B and LB→A be the log-likelihood of PA→B ( A , B ) and PB→A ( A , B ) , respectively . Their baseline approach optimizes a regret , R , to acquire an indicator variable γ ∈ R. If γ > 0 , the prediction is A→ B , otherwise B → A. R = − log [ σ ( γ ) LA→B + ( 1− σ ( γ ) ) LB→A ] ( 1 ) 2.2 PROPOSED MECHANISM : OBSERVE THE MODEL ’ S CONDITIONAL DISTRIBUTION DIVERGENCE . Rather than relying on gradients of end-to-end predictions to identify causal direction , we propose to directly observe the divergence of each model ’ s conditional distribution predictions under the intervention of the transfer dataset . We consider the distribution of a single model ’ s predictions on the training and transfer distributions . These distributions depend on both the marginal and conditional distributions learned by the models . It is expected that a model encoding the correct causal direction will have learned a correct conditional distribution , but the marginal distribution might shift from the training to transfer datasets . Hence , we would like to ignore changes in the model ’ s predictions that are due to marginal distribution differences and focus directly on the conditional distribution . Given access to the model ’ s conditional distribution predictions , we could use the KL-divergence to directly measure the conditional distribution differences between the training and transfer datasets . A model that encodes the correct causal structure should show no change in its conditional distribution predictions , so the divergence of these predictions should be small . On the other hand , a model that encodes the incorrect causal structure is likely to show significantly large divergence . We tie these observations together in the following proposition : Proposition 1 . Given two data distributions with the same directed causality between two random variables , A and B , the difference of the KL-divergences of their conditional distributions is an unbiased estimator of the correct causality direction between A and B . This approach shares the same intuition as the baseline approach . The model with correct causal direction should only witness minimal changes in the conditional distribution of its predictions between train and transfer distributions . Thus , the KL-divergence of the predictions should not change for models with correct causal structure . The proof for Proposition 1 is included in Appendix A . 2.3 PROPOSED SIMPLIFICATION : GENERALIZATION LOSS APPROXIMATES THE DIVERGENCE . Although we can use the KL-divergence of the conditional probabilities of two data distributions as an unbiased estimator of causality direction , it requires additional computation to model conditional distribution in transfer distribution . However , in many practical end-to-end learning settings , it is more efficient and straight-forward to acquire generalization loss , such as the cross-entropy , that can be used to approximate the conditional KL-divergence . More specifically , we compare the generalization gaps between two causal models to approximate the causality direction . Let the generalization gap , G , be the difference of a model ’ s losses on the training and transfer datasets : GA→B = LtransferA→B − LtrainA→B . Here , L ·A→B is the loss on the specified set . Further , we define a directionality score as the difference of generalization gaps : SG = GB→A − GA→B We show that , for an appropriately chosen loss , such as cross-entropy , SG is a biased but reasonable estimator of the correct causal direction . When SG > 0 , it is likely that A→ B is the correct causal direction–the generalization gap for the incorrect causal model dominates the score . We formalize this notion in Proposition 2 in Appendix B : Proposition 2 . Given two data distributions with the same directed causality between two random variables , A and B , the difference of their generalization gaps is a biased estimator of the correct causality direction between A and B . In practice for the tasks we test in next section , we find that although SG is biased , it always indicates the correct causal direction . An intuitive understanding is that this approach measures how well models trained on train distribution can predict in transfer distribution–their generalization ability . Algorithm 1 summarizes the process for identifying a correct causal model . In Appendix C , we describe that the generalization-based approach should converge more quickly than a gradient-based approach , especially in practical settings . Algorithm 1 The proposed approach for causality direction prediction . 1 : Train fA→B on training data , and get train loss LtrainA→B . 2 : Train fB→A on training data , and get train loss LtrainB→A . 3 : Get transfer loss LtransferA→B with fA→B on transfer data . 4 : Get transfer loss LtransferB→A with fB→A on transfer data . 5 : Get generalization loss GA→B = LtransferA→B − LtrainA→B . 6 : Get generalization loss GB→A = LtransferB→A − LtrainB→A . 7 : Compute SG = GB→A − GA→B . 8 : If SG > 0 return A→ B , else return return B → A . The baseline approach and our proposed approach are based on the same intuition of unchanging conditional distributions for correct causality direction . However , the baseline approach must observe gradients during transfer distribution training , while the proposed approach looks directly at the changes in model outputs on the transfer data distribution . Informally , we note that these approaches are roughly equivalent by observing the following : For the correct causal direction , G·→· = Ltransfer·→· − Ltrain·→· ∇G·→· = ∇Ltransfer·→· −∇Ltrain·→· Since ∇Ltrain = 0 , ∇G·→· = ∇Ltransfer·→· . The previous work shows that ∇θiLtransfer·→· = 0 for model parameters , θi , representing correct causal structures ( i.e. , joint distribution ) , so∇Ltransfer·→· should also be small . Similar to that , we expect G·→· to be small relative to the incorrect causal model .

Summary: This paper describes an approach for learning a representation U(X,Y), V(X,Y) of data (X,Y) such that U and V are causally meaningful, and U causes V. The approach relies on observing data from two domains, P and Q, where P(V| U) = Q(V |U) (reflecting the causal structure). The approach is a modification of Bengio et al. The main new idea is that the objective function can be tweaked by replacing a KL divergence term with a term involving domain-shift induced generalization errors.