[Paper Review] A Graph Autoencoder Approach to Causal Structure Learning

Ignavier Ng, Shengyu Zhu+2 more|arXiv (Cornell University)|Nov 18, 2019

Bayesian Modeling and Causal Inference17 references54 citations

TL;DR

Introduces a graph autoencoder (GAE) based gradient method for causal structure learning that handles nonlinear SEMs and vector-valued variables, showing improved performance and near-linear training time on larger graphs compared to prior gradient-based methods.

ABSTRACT

Causal structure learning has been a challenging task in the past decades and several mainstream approaches such as constraint- and score-based methods have been studied with theoretical guarantees. Recently, a new approach has transformed the combinatorial structure learning problem into a continuous one and then solved it using gradient-based optimization methods. Following the recent state-of-the-arts, we propose a new gradient-based method to learn causal structures from observational data. The proposed method generalizes the recent gradient-based methods to a graph autoencoder framework that allows nonlinear structural equation models and is easily applicable to vector-valued variables. We demonstrate that on synthetic datasets, our proposed method outperforms other gradient-based methods significantly, especially on large causal graphs. We further investigate the scalability and efficiency of our method, and observe a near linear training time when scaling up the graph size.

Motivation & Objective

• Motivate the need for scalable causal structure learning from observational data.
• Generalize gradient-based structure learning to a graph autoencoder framework for nonlinear relationships.
• Enable handling of vector-valued variables within a continuous optimization setup.
• Demonstrate scalability and efficiency benefits on synthetic datasets.
• Compare performance against state-of-the-art gradient-based methods.

Proposed method

• Formulate causal structure learning as a graph autoencoder with encoder g1 and decoder g2 and linear message passing A^T H^(j).
• Use f(X^(j),A)=g2(A^T g1(X^(j))) to model nonlinear relations; share weights across variables.
• Optimize reconstruction loss with an L1 penalty on A: min_{A,Θ1,Θ2} 1/2n sum_j ||X^(j) - X_hat^(j)||^2_F + λ||A||_1.
• Enforce acyclicity via the smooth constraint tr(exp(A ⊙ A)) - d = 0 and solve with Augmented Lagrangian method.
• Employ two MLPs for g1 and g2 with shared weights; update via gradient-based optimization (Adam).
• Compare against NOTEARS and DAG-GNN on synthetic data with scalar and vector-valued variables; assess SHD and TPR.

Experimental results

Research questions

• RQ1 Can a graph autoencoder framework model nonlinear causal relationships and support vector-valued variables in causal structure learning?
• RQ2 Does the GAE-based approach improve accuracy over state-of-the-art gradient-based methods (NOTEARS, DAG-GNN) on synthetic data, especially for larger graphs?
• RQ3 What are the scalability and training-time characteristics of the proposed method as graph size grows?
• RQ4 How does the method perform under scalar vs vector-valued variable settings?

Key findings

• GAE outperforms NOTEARS and DAG-GNN on synthetic datasets, especially for larger graph sizes.
• The method achieves near-linear training time when scaling up to 100 nodes.
• For scalar-valued cases, GAE yields lower SHD and higher TPR across tested graph sizes and nonlinear data-generating models.
• For vector-valued cases (l=5, latent dim l′=3), GAE shows better SHD and TPR, particularly as graph size increases.
• Experiments indicate DAG-GNN may not scale as well on these datasets, while NOTEARS can be competitive in some settings.
• Experiments conducted on GPU show GAE training times remain fast compared to DAG-GNN.