A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms

1 INTRODUCTION . The data used to train our models is often assumed to be independent and identically distributed ( iid . ) , according to some unknown distribution . Likewise , the performance of a machine learning model is typically evaluated using test samples from the same distribution , assumed to be representative of the learned system ’ s usage . While these assumptions are well analyzed from a statistical point of view , they are rarely satisfied in many real-world applications . For example , an accident on a major highway could completely perturb the trajectories of cars , and a driving policy trained in a static way might not be robust to such changes . Ideally , we would like our models to generalize well and adapt quickly to out-of-distribution data . However , this comes at a price – in order to successfully transfer to a novel distribution , one might need additional information about these distributions . In this paper , we are not considering assumptions on the data distribution itself , but rather on how it changes ( e.g. , when going from a training distribution to a transfer distribution , possibly resulting from some agent ’ s actions ) . We focus on the assumption that the changes are sparse when the knowledge is represented in an appropriately modularized way , with only one or a few of the modules having changed . This is especially relevant when the distributional change is due to actions by one or more agents , because agents intervene at a particular place and time , and this is reflected in the form of the interventions discussed in the causality literature ( Pearl , 2009 ; Peters et al. , 2016 ) , where a single causal variable is clamped to a particular value or a random variable . In general , it is difficult for agents to influence many underlying causal variables at a time , and although this paper is not about agent learning as such , this is a property of the world that we propose to exploit here , to help discovering these variables 1 Université de Montréal , 2 CIFAR Senior Fellow , 3 École Polytechnique Montréal , 4 Max-Planck Institute for Intelligent Systems , Tübingen , 5 Canada CIFAR AI Chair and how they are causally related to each other . In this context , the causal graph is a powerful tool because it tells us how perturbations in the distribution of intervened variables will propagate to all other variables and affect their distributions . As expected , it is often the case that the causal structure is not known in advance . The problem of causal discovery then entails obtaining the causal graph , a feat which is in general achievable only with strong assumptions . One such assumption is that a learner that has learned to capture the correct structure of the true underlying data-generating process should still generalize to the case where the structure has been perturbed in a certain , restrictive way . This can be illustrated by considering the example of temperature and altitude from Peters et al . ( 2017 ) : a learner that has learned to capture the mechanisms of atmospheric physics by learning that it makes more sense to predict temperature from the altitude ( rather than vice versa ) given training data from ( say ) Switzerland , will still remain valid when tested on out-of-distribution data from a less mountainous country like ( say ) the Netherlands . It has therefore been suggested that the out-of-distribution robustness of predictive models can be used to guide the inference of the true causal structure ( Peters et al. , 2016 ; 2017 ) . How can we exploit the assumption of localized change ? As we explain theoretically and verify experimentally here , if we have the right knowledge representation , then we should get fast adaptation to the transfer distribution when starting from a model that is well trained on the training distribution . This arises because of our assumption that the ground truth data generative process is obtained as the composition of independent mechanisms , and that very few ground truth mechanisms and parameters need to change when going from the training distribution to the transfer distribution . A model capturing a corresponding factorization of knowledge would thus require just a few updates , a few examples , for this adaptation to the transfer distribution . As shown below , the expected gradient on the unchanged parameters would be near 0 ( if the model was already well trained on the training distribution ) , so the effective search space during adaptation to the transfer distribution would be greatly reduced , which tends to produce faster adaptation , as found experimentally . Thus , based on the assumption of small change in the right knowledge representation space , we can define a meta-learning objective that measures the speed of online adaptation in order to optimize the way in which knowledge should be represented , factorized and structured . This is the core idea presented in this paper . Returning to the example of temperature and altitude : when presented with out-of-distribution data from the Netherlands , we expect the correct model to adapt faster given a few transfer samples of actual weather data collected in the Netherlands . Analogous to the case of robustness , the adaptation speed can then be used to guide the inference of the true causal structure of the problem at hand , possibly along with other sources of signal about causal structure . Contributions . We first verify on synthetic data that the model that correctly captures the underlying causal structure adapts faster when presented with data sampled after a performing certain interventions on the true two-variable causal graph ( which is unknown to the learner ) . This suggests that the adaptation speed can indeed function as a score to assess how well the learner fits the underlying causal graph . We then use a smooth parameterization of the considered causal graph to directly optimize this score in an end-to-end gradient-based manner . Finally , we show in a simple setting that the score can be exploited to disentangle the correct causal variables given an unknown mixture of the said variables . 2 WHICH IS CAUSE AND WHICH IS EFFECT ? . As an illustrative example of the proposed ideas , let us consider two discrete random variables A and B , each taking N possible values . We assume that A and B are correlated , without any hidden confounder . Our goal is to determine whether the underlying causal graph is A→ B ( A causes B ) , or B → A . Note that this underlying causal graph can not be identified from observational data from a single ( training ) distribution p only , since both graphs are Markov equivalent for p ( Verma & Pearl , 1991 ) ; see Appendix A . In order to disambiguate between these two hypotheses , we will use samples from some transfer distribution p̃ in addition to our original samples from the training distribution p . 2.1 THE ADVANTAGE OF THE CORRECT CAUSAL MODEL . Without loss of generality , we can fix the true causal graph to be A → B , which is unknown to the learner . Moreover , to make the case stronger , we will consider a setting called covariate shift ( Rojas-Carulla et al. , 2018 ; Quionero-Candela et al. , 2009 ) , where we assume that the change ( again , whose nature is unknown to the learner ) between the training and transfer distributions occurs after an intervention on the cause A . In other words , the marginal of A changes , while the conditional p ( B | A ) does not , i.e . p ( B | A ) = p̃ ( B | A ) . Changes on the cause will be most informative , since they will have direct effects on B . This is sufficient to fully identify the causal graph ( Hauser & Bühlmann , 2012 ) . In order to demonstrate the advantage of choosing the causal model A → B over the anti-causal B → A , we can compare how fast the two models can adapt to samples from the transfer distribution p̃ . We quantify the speed of adaptation as the log-likelihood after multiple steps of fine-tuning via ( stochastic ) gradient ascent , starting with both models trained on a large amount of data from the training distribution . In Figure 1 ( see Section 3.3 for the experimental setup ) , we can see that the model corresponding to the underlying causal model adapts faster . Moreover , the difference is more significant when adapting on a small amount of data , of the order of 10 to 30 samples from the transfer distribution . We will make use of this property as a noisy signal to infer the direction of causality , which here is equivalent to choosing how to modularize the joint distribution . 2.2 PARAMETER COUNTING ARGUMENT . A simple parameter counting argument can help us understand what we are observing in Figure 1 . Since we are using gradient ascent for the adaptation , let ’ s first inspect how the gradients of the log-likelihood wrt . each module behave under the transfer distribution . Proposition 1 . Let G be a causal graph , and p a ( training ) distribution that factorizes according to G , with parameters θ . Let p̃ be a second ( transfer ) distribution that also factorizes according to G. If the training and transfer distributions have the same conditional probability distributions for all Vi but a subset C ( e.g . the transfer distribution is the result of an intervention on the nodes in C ) : p ( Vi | PaG ( Vi ) ) d= p̃ ( Vi | PaG ( Vi ) ) ∀Vi /∈ C ( 1 ) then the expected gradient w.r.t . the parameters θi such that Vi /∈ C of the log-likelihood under the transfer distribution will be zero ∀Vi /∈ C , EV∼p̃ [ ∂ log p ( V ) ∂θi ] = 0 . ( 2 ) Proposition 1 ( see proof in Appendix B.1 ) suggests that if both distributions factorize according to the correct causal graph , then only the parameters of the mechanisms that changed between the training and transfer distributions need to be updated . This effectively reduces the number of parameters that need to be adapted compared to any other factorization over a different graph . It also affects the number of examples necessary for the adaptation , since the sample complexity of a model grows approximately linearly with the VC-dimension ( Ehrenfeucht et al. , 1989 ; Vapnik & Chervonenkis , 1971 ) , which itself also grows approximately linearly with the number of parameters ( for linear models and neural networks ; Shalev-Shwartz & Ben-David , 2014 ) . Therefore we argue that the performance on the transfer distribution ( in terms of log-likelihood ) will tend to improve faster if it factorizes according to the correct causal graph , an assertion which may not be true for every graph but that we can test by simulations . Recall that in our example on two discrete random variables ( each taking say N values ) , we assumed that the underlying causal model isA→ B , and the transfer distribution is the result of an intervention on the cause A . If the model we learn on the training distribution factorizes according to the correct graph , then only N − 1 free parameters should be updated to adapt to the shifted distribution , accounting for the change in the marginal distribution p̃ ( A ) , since the conditional p̃ ( B | A ) = p ( B | A ) stays invariant . On the other hand , if the model factorizes according to the anti-causal graph B → A , then the parameters for both the marginal p̃ ( B ) and the conditional p̃ ( A | B ) must be adapted . Assuming there is a linear relationship between sample complexity and the number of free parameters , the sample complexity would be O ( N2 ) for the anti-causal graph , compared to only O ( N ) for the true underlying causal graph A→ B .

Summary: The paper first shows that, in a very simple two-variable task, the model with the correct underlying structure will adapt faster to a causal intervention than the model with the incorrect structure. This idea is used to develop a “meta-transfer” objective function for which gradient ascent on a continuous representation of the model structure allows learning of that structure. The paper shows that optimizing with respect to this objective with a simple model is guaranteed to converge to the correct structure, and also presents experimental results on toy problems to demonstrate.