A Graph Autoencoder Approach to Causal Structure Learning

Ignavier Ng, Shengyu Zhu, Zhitang Chen, Zhuangyan Fang

Introduction

Causal structure learning has important applications in many areas such as genetics , biology , and economics . An effective way to discover causal relations amongst variables is to conduct controlled experiments, which however is expensive and even ethically prohibited in certain scenarios. Learning causal structure from purely observational data has attracted much research attention in the past decades .

Two major classes of structure learning methods from observational data are constraint- and score-based. Constraint-based methods, such as PC and FCI , first use conditional independence tests to learn the skeleton of the underlying casual graph and then orient the edges based on a series of orientation rules . Score-based methods, like GES , assign a score to each causal graph according to some pre-defined score function and then search over the space of causal directed acyclic graphs (DAGs) to find the one with the optimal score. However, finding the causal graph with optimal score is generally NP-hard, largely due to the combinatorial acyclicity constraint on the causal graph. Although the performance of several constraint- and score-based methods is guaranteed with infinite data and suitable assumptions , the inferred graphs are usually less satisfactory in practice. More recently, Zheng et al. 2018 proposed a smooth characterization on the acyclicity constraint, thus transforming the combinatorial optimization problem into a continuous one, provided with a proper loss function. Subsequent work has also combined this approach with graph neural networks to model nonlinear causal relationships.

In this work, we propose a new gradient-based structure learning method which generalizes the recent gradient-based methods to a graph autoencoder framework that allows nonlinear structural equation models and vector-valued variables. We demonstrate that on synthetic datasets, our proposed method outperforms other gradient-based methods significantly, especially on relatively large causal graphs. We analyze the training time of our method and observe a near linear training time when scaling the graph size up to 100 nodes.

Gradient-Based Causal Structure Learning

This section introduces the recently developed gradient-based methods for causal structure learning.

Zheng et al. 2018 are probably the first to transform the combinatorial optimization problem of score-based methods into a continuous one for linear structural equation model (SEM), which adopts the following data generation model

where $n$ is the sample size and $X^{(j)}$ denotes the $j$ -th observed sample of $X$ .

The above problem can then be solved by standard numeric optimization methods such as the augmented Lagrangian method . This approach is then named Non-combinatorial Optimization via Trace Exponential and Augmented lagRangian for Structure learning (NOTEARS) by the authors. Though numeric optimization methods may not produce exact acyclic graphs, i.e., $\Tr(e^{A\odot A})-d$ can be very small (e.g., $10^{-8}$ ) but not exactly zero, post-processing like thresholding can be used to meet the hard acyclicity constraint.

To generalize the above model to nonlinear cases, Yu et al. 2019 proposed DAG-GNN, a generative model given below

where $g_{1}$ and $g_{2}$ are point-wise (nonlinear) functions. Causal structure learning is then formulated under the framework of variational autoencoders with multilayer perceptrons (MLPs) to model the causal relations, and the objective is to maximize the evidence lower bound under the acyclicity constraint. Notice that here $Z$ is regarded as the latent variable and its dimension can be selected to be lower than $d$ , the number of variables.

Graph Autoencoder for Causal Structure Learning

We present an alternative generalization of NOTEARS to handle nonlinear causal relations. We consider the case with each variable $X_{i}$ being scalar-valued (i.e., $l=1$ ) in Section 3.1 and then the the more general case ( $l\geq 1$ ) in Section 3.2.

where $f(X^{(j)},A)$ denotes the data generating model w.r.t. the $j$ -th observed sample of $X$ and $\Theta$ denotes the parameters associated with $f$ . For NOTEARS, we would have $f(X^{(j)},A)=A^{T}X^{(j)}$ for linear SEM.

One way to extend $f(X^{(j)},A)$ in Eq. (3) to handle nonlinear relationships is to set

To handle a more complex nonlinear model as in , we consider to extend Eq. (4) to

2 Connection with Graph Autoencoder

We now draw a connection of Eq. (5) with graph autoencoder (GAE), as outlined in Figure 1. We can rewrite Eq. (5) as

Eq. (6), together with Eq. (3), form a GAE (similar to ) trained with reconstruction error where $g_{1}$ and $g_{2}$ are respectively the variable-wise encoder and decoder, and the message passing operation is applied at the latent representation $H^{(j)}$ . In our case, the message passing operation is a linear transformation $A^{T}H^{(j)}$ , similar to the graph convolutional layer used in .

where $\hat{X}^{(j)}=g_{2}\left(A^{T}g_{1}(X^{(j)})\right)$ is the reconstructed output and $\Theta_{1}$ and $\Theta_{2}$ are the MLP weights associated with $g_{1}$ and $g_{2}$ , respectively.

It is interesting to compare the proposed generalization Eq. (5) or Eq. (6) with Eq. (2) used by DAG-GNN. The linear SEM can be written as $X=A^{T}X+Z$ with $X$ and $Z$ being the vectors concatenating all the observational and noise variables, respectively. We consider $Z$ as additive noises and Eq. (5) or Eq. (6) can be viewed as generalized causal relations taking $X_{pa(i)}$ as inputs for the $i$ -th observational variable. In contrast, DAG-GNN further writes linear SEM as $X=(I-A^{T})^{-1}Z$ and considers it as a generative model which takes random noises $Z$ as input. The experiments conducted show that our alternative generalization performs better, in both efficiency and efficacy, than DAG-GNN on similar datasets used by .

3 Augmented Lagrangian Method

The optimization problem given in Eq. (7) can be solved using the augmented Lagrangian method. The augmented Lagrangian is given by

where $h(A):=\Tr(e^{A\odot A})-d$ , $\alpha$ is the Lagrange multiplier, and $\rho>0$ is the penalty parameter. We then have the following update rules

where $\beta>1$ and $\gamma<1$ are tuning hyperparameters. Problem (8) is first-order differentiable and we apply gradient descent method implemented in Tensorflow with Autograd and Adam optimizer .

Experiments

In this section, we conduct experiments on synthetic datasets to demonstrate the effectiveness of our proposed method. We compare our method against two recent gradient-based methods, NOTEARS and DAG-GNN .

We use the same experimental setup as in . In particular, we generate a random DAG using the Erdős–Rényi model with expected node degree $3$ , then assign uniformly random edge weights to construct the weighted adjacency matrix $A$ . We generate $X$ by sampling from ANM $X=f(A,X)+Z$ with some function $f$ elaborated soon. The noise $Z$ follows standard matrix normal. We report the structural Hamming distance (SHD) and true positive rate (TPR) for each of the method, averaged over four seeds. With sample size $n=3,000$ , we conduct experiments on four different graph sizes $d\in\{10,20,50,100\}$ . We consider scalar-valued variables ( $l=1$ ) and vector-valued case ( $l>1$ ) in Sections 4.1 and 4.2, respectively.

For the proposed GAE, we use a 3-layer MLP with 16 ReLU units for both encoder and decoder. For NOTEARS and DAG-GNN , we use the default hyperparameters found in the authors’ code.

Following the setup in , our first dataset uses the data generating procedure below:

which is a generalized linear model. The results of SHD and TPR are reported in Figure 2a. One can see that GAE outperforms NOTEARS and DAG-GNN, with SHD close to zero for graphs of 100 nodes. We also observe that NOTEARS has better performance than DAG-GNN when $d=100$ , which indicates that DAG-GNN may not scale well on this dataset.

We next consider a more complicated data generating model, where the nonlinearity occurs after the linear combination of the variables:

As shown in Figure 2b, our method has better performance than both NOTEARS and DAG-GNN in terms of SHD and TPR across all graph sizes. We conjecture that with higher nonlinearity, the use of proposed GAE framework results in a much better performance than NOTEARS and DAG-GNN. For both nonlinear data generating procedure Eq. (9) and (10), it is surprising that a linear model such as NOTEARS is on par with DAG-GNN or even has better performance in some cases. We hypothesize that the formulation of DAG-GNN results in the lack of causal interpretability on the adjacency matrix learned.

2 Vector-Valued Case

In particular, we choose $l=5$ and the latent dimension of GAE $l^{\prime}=3$ . The SHD and TPR are reported in Figure 3, which shows that GAE has better SHD and TPR than NOTEARS and DAG-GNN especially when the graph size is large. It is also observed that DAG-GNN slightly outperforms NOTEARS for $d=20$ and $50$ .

3 Training Time

Scalability is one of the important aspect in causal structure learning . To compare the efficiency and scalability of different gradient-based methods, we compute the average training time of GAE and DAG-GNN over all experiments conducted in Section 4.1 and 4.2. The experiments were computed on NVIDIA V100 Tensor Core GPU with 16GB of memory hosted on AWS GPU cloud instances. We do not include the average training time of NOTEARS as it uses only CPU instances.

As shown in Figure 4, the average training time of GAE is less than 2 minutes even for graphs of 100 nodes. On the other hand, DAG-GNN training can be much more time-consuming: the training procedure takes $32.9\pm 4.0$ and $77.0\pm 18.0$ minutes for $d=10$ and $100$ , respectively. We also observe that the training time of GAE seems to scale linearly when increasing the graph size to 100. The training time is fast as deep learning is known to be highly parallelizable on GPU , which yields a promising direction in gradient-based methods for causal structure learning.

Conclusion

The formulation of continuous constrained approach by enables the application of gradient-based methods on causal structure learning. In this work, we propose a new gradient-based method that generalize the recent gradient-based methods to a GAE framework that allows nonlinear SEM and vector-valued datasets. On synthetic datasets, we demonstrate that our proposed method outperforms other state-of-the-art methods significantly especially on large causal graphs. We also investigate the scalability and efficiency of our method, and observe a near linear training time when scaling the graph size up to 100 nodes. Future works include benchmarking our proposed method on different graph types and real world datasets, as well as testing it on synthetic dataset with larger graph size (up to 500 nodes) to verify if the training time still scales linearly.