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: The paper presents a novel approach for determining the causal direction between two random variables $A$ and $B$. The approach is based on the assumption that the conditional distribution $P_{A \rightarrow B}(B \mid B)$ does not change between the train and transfer distribution. As a result, a model that predicts the correct causal direction $A \rightarrow B$ should generalize better from the train to the transfer distribution compared to a model predicting the wrong causal direction $B \rightarrow A$. While previous work has proposed to use this insight to determine the causal direction by comparing the adaptation speed, this paper proposes to directly measure and compare the generalization performance. The results indicate that the proposed approach leads to the same performance in terms of causal relation prediction, but that it is more sample efficient and faster.