Amortized learning of neural causal representations

Causal models can compactly and efficiently encode the data-generating process under all interventions and hence may generalize better under changes in distribution. These models are often represented as Bayesian networks and learning them scales poorly with the number of variables. Moreover, these approaches cannot leverage previously learned knowledge to help with learning new causal models. In order to tackle these challenges, we represent a novel algorithm called \textit{causal relational networks} (CRN) for learning causal models using neural networks. The CRN represent causal models using continuous representations and hence could scale much better with the number of variables. These models also take in previously learned information to facilitate learning of new causal models. Finally, we propose a decoding-based metric to evaluate causal models with continuous representations. We test our method on synthetic data achieving high accuracy and quick adaptation to previously unseen causal models.


INTRODUCTION
Neural networks have achieved astonishing success in many machine learning areas by relying on the assumption of independent and identically distributed (iid) data in training and testing settings. In particular, these methods rely on learning correlations from the data and cannot per se identify causal links (de Haan et al., 2019;Schölkopf, 2019). In many settings of interest, e.g. domain adaptation, transfer learning and multitask learning, this is a problem as the iid assumption does not hold. On the other hand, causal models efficiently and compactly encode the data-generating process under interventions.
Previous work in causal learning often focuses on learning representations for causal models in terms of Bayesian networks, e.g. (Heckerman et al., 1995;Cooper & Yoo, 1999;Spirtes et al., 2000;Chickering, 2002;Tsamardinos et al., 2006;Shimizu et al., 2006;Pearl, 2009;Hoyer et al., 2009;Peters et al., 2011;Hauser & Bühlmann, 2012;Goudet et al., 2017). While such approaches provide an easily interpretable output, they rely on iterating over all possible graph structures and thus scale poorly with the number of variables. In addition, it is challenging for Bayesian networks to utilize previously learned information to more efficiently learn new causal graphs in settings such as metalearning. Hence, we propose a paradigm shift away from learning Bayesian networks to learning continuous representations of causal graphs using deep neural networks, such that the learned model not only scales better with the number of variables, but can also take into account previously learned information to help learn new causal graphs more efficiently.
While most research in causal induction focuses on learning a single causal graph at a time, we explore a meta-learning setting, where the algorithm observes many distinct causal graphs during training and is asked to quickly recover a new causal graph during test time.
Motivated by this scenario, we examine a meta-learning setup where learning is composed of episodes, shown in Figure 1. At the beginning of each episode, a new causal graph is sampled and the graph governs the data generation during that episode. At every step t within an episode, a random and independent intervention is conducted on the causal graph and our algorithm receives the node values of a sample under intervention. Note that we assume that there is no time-delay between the intervention and its effect, such that the effect is observed immediately. an episode of length k our algorithm observes k samples from the underlying causal graph and each under a different interventions. To facilitate learning, we ask our model to predict outcomes of interventions at every step. Note that this is different from the meta-learning setups in Bengio et al. (2019); . In their setting, there is a single causal graph across different episodes and in our setting, there is a distinct causal graph for each episode and hence our algorithm observes many different causal graphs. In order to enable learning of causal graphs in this setup, we propose a novel neural network architecture that efficiently learns to encode causal graphs in a compact, continuous representation. When learning causal graphs explicitly as discrete structures, it is straightforward to ...

Episode i
Adjacency matrix Episode i+1 Step j Step  the beginning of episode i (for i < k), a new causal graph Gi that governs the data generating process for that episode is sampled. Both the structure and the conditionals of the causal graph are resampled per episode. The structure of the graph can be represented using an adjacency matrix. This is followed by k random and independent interventions (represented by the hammer) on the causal graph Gi. The model observes one sample per intervention and k samples in total for an episode, and then evaluated on unseen causal graphs at test time.
evaluate accuracy by counting how many of the inferred edges match the ground truth. Our implicit, continuous representation of a causal graph cannot be directly be compared against the ground truth. One possibility is to compare the predictions of node values under all possible interventions, rather than comparing the edge structure. Unfortunately, comparing node predictions scales poorly and would be prohibitively expensive even for medium-sized graphs. Therefore, we instead learn a decoder from the continuous representation into a graph, and then measure its edge accuracy. Our key contributions can be summarized as follows: • We introduce a novel framework for learning to learn causal discovery models using continuous representations via unsupervised losses. • This novel framework that can take in previously learned information to facilitate learning of new causal graphs. • We introduce a decoding metric to evaluate causal models with continuous representations.

RELATED WORK
Most approaches to causal induction are focuses on learning a single causal graph and can not easily adapt to the meta-learning setting. For a comprehensive review on most of these methods, please refer to Heinze-Deml et al. (2018). We will focus our discussions on causal induction methods that explores a meta-learning setting.
The works of Bengio et al. (2019) and  explore a meta-learning setting for causal induction from interventional data. However, their setting is different in that there is a single underlying causal model across the different episodes for both training and test. In our setting, there is a distinct causal model per episode and hence the algorithm observes many different causal models during training and the model is asked to recover unseen causal models during test. The work of Nair et al. (2019) uses a meta-learning setup similar to ours, however, their model is provided a supervised signal from the groundtruth graph during training, whereas ours can infer the causal structure from an unsupervised learning signal. Most similar to our work, (Dasgupta et al., 2019) uses meta-learning to learn to make predictions under interventions that suggest (but do not prove) that some causal reasoning is occurring. However, their approach does not induce a causal graph, neither explicitly nor via a decoding, thus it cannot be used for general causal discovery, but rather is more focused on the reinforcement learning setting.

CAUSAL RELATIONAL NETWORKS
We propose a general framework for learning causal models from interventional data by leveraging the power of neural networks. Unlike most previous methods that learn Bayesian network structure, we employ continuous representations to represent causal graphs. Furthermore, we train our model on many distinct causal graphs which allows it to learn new causal graphs very quickly. In order to evaluate the effectiveness of our method, we measure the ability of our model to predict the outcome of interventions as well as decode the structure of the underlying causal Bayesian network.

PROBLEM FORMULATION
We propose a meta-learning setup to learn causal graphs. As shown in Figure 1, training occurs in episodes, with a distinct ground-truth causal graph that governs the data generation for the duration of each episode. Episodes have a fixed length k. At the beginning of each episode, we first draw a new causal graph G, followed by generating one sample for a random intervention for a total of k samples per episode. The resulting intervention samples x along with the index of the intervened variable are fed into the model sequentially. Note that there is no time-delay between the intervention and its effects, such that all effects are observed immediately. The model is asked to predict the intervention outcome and hence learning the causal relationships between variables. For a detailed description, please refer to Section 3.2.
To evaluate the performance of our model, we cannot directly compare the inferred graph structure to the ground truth, as is usually done when learning Bayesian networks. We have continuous representations that implicitly encode the underlying causal graph. Given that evaluating the predictions under all possible interventions does not scale to large graphs and state spaces, we propose to use a decoding-based evaluation metric for continuous representations of causal graphs, similar to (Gregor et al., 2019). Specifically, we propose to evaluate if the discrete graph structure can be decoded from the continuous representation by training a neural network decoder. Thus, at every time step t, we feed the hidden state of our model into a decoder trained to output the right causal graph. Note that we do not pass gradients from the ground-truth supervision signal back into the decoder.  nents: attentive relational model that models the data and serves to preserve the relationship between variables; the sequence model accumulates information through time and the graph decoder decodes the graph from the sequence model. Our tasks contains sequence data and hence a natural choice of model would be sequence models such as recurrent neural networks (RNNs). However, after some experimentation, we found that it is difficult for vanilla RNN models such as LSTM (Hochreiter & Schmidhuber, 1997) to model causal relationships, shown in Section 4.2.1. Hence, we propose a new architecture named causal relational networks that can better model causal relationships. The model comprises three modules shown in Figure 2. 1. The attentive relational model models the observed data through a latent representation and serves to preserve the relationship between variables.

MODEL DESCRIPTION
2. The sequence model computes the belief state at every step t, such that the belief state at t models the available information up to t.
3. The graph decoder outputs the estimated adjacency matrix of the underlying causal graph.
As in a meta-learning setup, training is divided into episodes. Each episode begins with sampling a new causal graph. Data from each episode is sampled from the same causal graph. The samples are read as a sequence as shown in Figure 2. Note that the data is not time-series data, since there is no time-delay between interventions and effects from it. For a causal graph of size n, at each time step, one sample vector of size n along with the intervention information is passed into the model.
Attentive Relational Model. A natural choice for encoding the input would be to use a monolithic fully connected encoder with a sequence model (such as a LSTM). However, in our experiments (in section 4.2.1), we observed such monolithic model does not manage to learn the causal relationships between variables. To tackle this, we propose a novel relational network architecture inspired by the relational networks used in . This architecture is trained with unsupervised learning. In particular, the model can more appropriately models causal relationships between variables because it is trained to predict, given a particular variable in the sample, the value of all other variables. Figure 3: The attentive relational model for inputs with 3 variables. The 3 distinct input MLP takes in each variable x t i of the inputs separately and outputs all other variables x t k , where k = i. The variables that are direct parents or child of the predicted variable would be able to better predict the variable. Hence, the attention MLP computes a weight based on how well x i can predict x k and summarize them into the immediate output O t , which is then passed to the sequence model. The attentive relational model receives an unsupervised loss based on its ability to predict each variable.
Within an episode, let x t 0 , x t 1 , ..., x t n denote a particular sample from the causal graph under an intervention at time t, x t 0 denote the index of the intervened variable and x t i denote value of the i-th variable in the t-th sample. For every variable, we instantiate a feedforward neural network that takes the values of this variable as input. Note that there are n MLPs in the model that are identical in shape but do not share parameters (see Figure 3). The task of each MLP is to predict the value of the remaining variables (i.e. output size n − 1). Specifically, let f k be the MLP that takes as input x t i and attempts to predict all x t k , such that k = i, i.e. x 0 , x 1 , ...x k,k =i = f k (x i ) and hence, there are n − 1 predictions of each variable x t i . We we augment our architecture with an attention model, based on the intuition that a variable can better predict its direct parents or children (in comparison with unconnected or distantly connected variables). The attention model can learn how well variable i can predict variable j. Thus, for a given variable x t i , the MLP returns n − 1 predictions (for all x k,k =i ), which are fed into the attention MLP (see Figure 3). This then outputs attention weight w t i,k , with k denoting the prediction of variable X t i from f k . Next, we use the attention weights to compute the intermediate prediction O t i , where this prediction is being passed through a MLP, where it outputs the final predictions of the variable X t i . Taking all of this together, the MLPs together with the attention weight encode the relationship between a single variable and the rest of the causal graph. The output for the autoencoder predicts the sample from the intervention distributionX t 0 , ...,X t n−1 . Note that the model ignores the prediction for the variable that has been intervened upon, since that variable doesn't give us information about the underlying causal graph.
Sequence model. We use a simple sequence model to accumulate information produced by the encoder and update the belief state of the model. In fact, in our setup, we use a simple cumulative summation operator with no learned parameter, the se-quence model simply adds current information to already accumulated information. The sequence model at time t takes the immediate output O t from the attentive relational model and add it to the hidden state of the sequence model. Then the hidden state h t of the sequence model at time t is updated as h t = h t−1 + O t . This simple additive update guarantees that the final representation h t is invariant with respect to permutations in the order of presentation of the data points. A more complex recurrent model, such as LSTM ( (Hochreiter & Schmidhuber, 1997)), could also be used but it is not trivial to preserve permutation invariance in this case.
Graph decoder. Having computed the hidden state h t of the sequence model, we use a LSTM (Hochreiter & Schmidhuber, 1997) to decode the structure of the causal graph. The structure of the causal graph is presented as a flattened adjacency matrix (a sequence), and the graph decoder is trained to predict the adjacency matrix conditioned on hidden state h t from the sequence model. The model is trained using teacher forcing during training and free-running for evaluation. We also ensure that no gradient with supervision signal are passed back into the sequence model.

EXPERIMENTAL SETUP AND RESULTS
Our experiments aim to evaluate our proposed method for causal discovery in a meta-learning setting. We study our proposed method for causal induction on synthetic data and we perform experiments to answer the following questions: Q1. How well can our proposed method recover causal structures?
Q2. Can our model leverage previously learned information to more efficiently recover causal structures? Q3. Does the structured relational autoencoder help to extract causal relationships?
Q4. Does the accuracy of the decoder reflect the model's ability to recover causal structures?
Q5. How much would supervision signal improve the results?

SYNTHETIC DATASET
We explore the setting of binary Bayesian networks with all variables observed. We focus on this setting as a proof of concept for our methodology; extensions to categorical Bayesian networks are straightforward. Intuitively, the causal graphs can be thought of as a soft random Boolean network that is a combination of AND and OR gates, which is motivated by discrete variables combined by noisy AND and OR gates. More formally, the ground-truth model defined a joint p(X = x) = i p (X i = x i | x <i , w i ) is formulated as follows, where w i ∈ {0, 1} is an intervention mask sampled randomly (w i = 0 corresponds to an intervention replacing the node distribution by an uniform distribution), 1} is the connectivity mask representing the Bayesian net causal graph, and β = 5 is a concentration parameter. This model simulates a Bayesian network over binary variables with randomized interventions. It can be thought of as a soft random Boolean network. For the ease of sampling, we restrict the ground-truth model to a directed acylic graph (DAG). We have fixed the number of variables in the graph to 5 for the experiments.
Identifiability In our setting, all variables are observed, such that there are no latent confounders. All interventions are random selected and are independent and all interventions are hard interventions. It has been shown in (Eberhardt et al., 2012) that the true causal graph is identifiable in principle under our setting.
We perform the same hyper-parameter search for our baselines as well as our model. We use a batch size of 8, learning rate of 0.0002, all models are trained with the Adam optimizer (Kingma & Ba,  (Nair & Hinton, 2010) in between. The relational encoder final output size is 64. The graph decoder LSTM has a hidden state size of 128.
Evaluation metrics We evaluate our model on the accuracy of the decoded graph with respect to the ground-truth graph. The accuracy is defined to be the percentage of decoded edges that align with the ground-truth edge. The edge is said to be aligned with the ground-truth edge if both of the argmax of the Bernoulli distribution over the learned edge and the ground-truth edge matches.

EXPERIMENTAL RESULTS
Q1. Recovering causal structure. As a distinct causal graph is sampled at the beginning of each episode, the edge accuracy for the models should naturally improve across the course of an episode, for different stages of training. The model has learned to extract information about the causal graph structure from intervention samples, reflected in the improvements in edge accuracy as the model observes more intervention data from the causal graph. In the beginning of the episode, the model only knows the underlying causal graph is a DAG. As the model observes more intervention samples from the same graph, it can quickly infer underlying causal graph. As shown in Figure 4a, towards the end of training (yellow curve), our model achieves 90% in edge accuracy just after seeing 100 samples from the new graph.
Q2. Leveraging previously learned information. Shown in Figure 4a, our model quickly learned to leverage previously learned information to more efficiently learn new causal models. In fact, the model quickly learns that the underlying causal graph is a DAG and hence achieving over 60% accuracy even at the beginning of each episode. The model then learns to decode the true causal graph from intervention samples throughout training, this is reflected in the improvement in edge accuracy during the episode throughout the training process.

BASELINE MODELS
There are 3 main components to our model as shown in Figure 2. First, we evaluate if the attentive relational model is necessary to capture causal relationships. Hence, we evaluate against a monolitic baseline model such as a fully-connected MLP, this then forms our LSTM baseline model. We then evaluate the amount of performance gain the model could achieve by receiving supervision signal from our graph decoder. Lastly, in order to verify that the information for the conditional graph decoder for decoding the graph is indeed contained in our model, we compare the performance of our model to one with zero information in its input to the graph decoder. The sequence model that accumulates information across time is fairly straightforward and hence we have not performed comparisons against this component.
LSTM baseline. Our model uses a relatively sophisticated relational network to extract information in the input. In order to verify the necessity of this component, we compare our encoder to a LSTM with naive fully connected feed-forward neural network (MLP). For simplicity, we call this the LSTM baseline. For this baseline, we no longer need an autoencoder to process the input, as the autoencoder could just be an identity model. As there are no autoencoder (with its own decoder), the baseline is missing the unsupervised learning signal. To compensate for this, we pass the gradients from the supervised graph decoder to the MLP encoder. The baseline model has an easier job, since it could learn directly from supervision signal.
Q3. Is the relational autoencoder necessary? Interestingly, as shown in Figure 5, though the baseline model has an easier task, the edge prediction accuracy is only around 70% throughout training as compared to the 90% reached by our model as shown in Figure 5. The baseline model neither learns within an episode (as it observes more samples), nor does it learn with more training.
Q4. Does the accuracy of the decoder reflect the model's ability to recover causal structures? Next, we verify that the belief state in our model contains useful information for decoding the structured causal graph. This baseline uses the same relational encoder-decoder and graph decoder from our original model; the only difference is that the belief state is set to zero. If this baseline could also learn to decode the structured causal graph, it would indicate that it was not the information in the belief state that helped to decode the structured causal graph, it was rather the graph decoder itself. Figure 5 shows that the decoder was not in fact able to decode the causal graph structure. The accuracy remains around and below 60% throughout training, compared to the 90% that our model achieved at the end of training (Figure 4a), indicating that the belief state was crucial to decoding the causal graph.
Q5. Supervised graph learning. We next evaluate the amount of performance gain the model gets by receiving supervision signal from the graph decoder. This is achieved by removing the stop gradient from the belief state that is passed in as inputs to the graph decoder. The results are shown in Figure 4b. The model with supervision signal achieved 95% accuracy as compared to the 90% achieved by the unsupervised model.

CONCLUSION
In this paper, we presented a new framework for learning and evaluating learned causal models using neural networks. In particular, we have introduced a novel neural architecture causal relational networks, and proposed to use a decoding-based metric for evaluation. We have tested our methodology on synthetic data and have shown that it can accurately learn causal structures. Furthermore, we have shown that causal relational networks can very quickly learn new causal models.