A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms

Abstract

This paper proposes exploiting a meta-learning objective that maximizes the speed of transfer to a changed distribution in order to learn how to acquire knowledge in a modular fashion. In particular, we care about how to factor the joint distribution into appropriate conditionals consistent with the causal relationships. If this is achievable, changes in the distribution are assumed to be localized (distributions are localized) (e.g., one of the marginals changes due to an intervention on one of the variables). We show that, under such localized changes of the assumed causal mechanisms, the correct causal graph will tend to have only a few parameters with non-zero gradients, i.e., the parameters that need to be adjusted (the parameters of the modified variables). This is observed experimentally to lead to faster adaptation, and we exploit this property to define a meta-learning surrogate score which, in addition to a continuous parametrization of the graph, will also favor the correct causal graph. Finally, considering AI agents (e.g., a robot autonomously discovering its environment), we consider how the same objective can discover the causal variables themselves, since the observed low-level variables carry no causal meaning. Experiments on the bivariate case validate the proposed ideas and theoretical results.

Introduction

Data is assumed to be independent and identically distributed (IID). Likewise, a model’s performance is typically evaluated with test samples drawn from the same distribution, on the assumption that they represent how the learned system will be used. Although these assumptions are well analyzed from a statistical point of view, they are hard to apply in many practical situations. For example: a medical diagnosis system trained on historical data from one hospital may perform poorly on patients from another hospital because the distribution has shifted. Ideally, we would like our models to generalize well and to adapt quickly to out-of-distribution data.

However, successfully transferring to a new distribution calls for more data. In this paper, rather than focusing on assumptions about the data distribution itself, we focus on how the data distribution changes (e.g., the change from the training distribution to the transfer distribution may be caused by the actions of some agents). The assumption we care about is that when knowledge is represented in an appropriately modular way, only one or a few modules change — such changes are sparse. This is especially important when the distribution change is caused by the action of one or more agents, because an agent intervenes at a particular place and time, which takes the form of the interventions discussed in the causal literature, where one of the causal variables is clamped to a particular value or to a random variable. In general, it is hard for an agent to affect more than one of the underlying causal variables at once. Although this paper is not about agent learning per se, this is a property we explore to help discover these variables and the causal relationships between them. In this respect, the causal graph is a powerful tool, because it tells us how a perturbation of the distribution of an intervened variable will propagate to all the other variables and affect their distributions.

Usually, the causal structure is not known in advance. The problem of causal discovery is generally needed to obtain the causal graph, but this is usually achievable only under strong assumptions. One such assumption is that a learner that has captured the correct structure of the true underlying data-generating process should generalize to perturbations of that structure in a specific way. This can be illustrated by considering temperature and altitude: roughly speaking, a learner trained on data from Switzerland should still be valid when tested on data drawn from the distribution of a less mountainous country such as the Netherlands. It is thus proposed to use the out-of-distribution robustness of the predictive model to guide inference about the true causal structure.

So how can we exploit the assumption of local changes? As we explain theoretically and validate experimentally here, if we have the correct representation of knowledge, then starting from a model pre-trained on the training distribution we should be able to adapt very quickly to the transferred data distribution. This happens because we assume the data-generating process is acquired as a set of independent mechanisms, and moving from the training distribution to the transfer distribution requires changing very few of the mechanisms and parameters. A model corresponding to the correct decomposition of knowledge therefore needs only a few updates and examples to adapt to the transfer distribution. As shown below, the expected gradient on the invariant parameters will be close to 0 (if the model has already been well trained on the training distribution), so the effective search space is greatly reduced during adaptation to the transfer distribution, which speeds up adaptation, as the experiments show. Therefore, based on the assumption of small changes in the space of correct knowledge representations, we can define a meta-learning objective that measures the speed of adaptation — a form of regret — to optimize how knowledge is represented, decomposed, and structured

Back to the earlier temperature and altitude example: if we receive out-of-distribution data from the Netherlands, we expect the model to adapt faster because a small number of transfer samples from the Netherlands have been collected. Similarly to robustness, adaptation speed can be used to guide inference about the true causal structure of the problem at hand, possibly together with other sources of signal about the causal structure.

Main contributions: We first validate on synthetic data that, after certain interventions are performed on the true bivariate causal graph (unknown to the learner), a model that correctly captures the underlying causal structure adapts faster when samples are provided. This shows that adaptation speed can serve as a score that adequately assesses how well a learner fits the underlying causal graph. We then use a smooth parametrization of the causal graph to directly optimize this score in an end-to-end manner. Finally, we show that, when the mixing of variables is unknown, this score can be used to distinguish the correct causal variables.

Which is Cause and Which is Effect?

To be continued