[Paper Review] D-VAE: A Variational Autoencoder for Directed Acyclic Graphs

Muhan Zhang, Shali Jiang+3 more|arXiv (Cornell University)|Apr 24, 2019

Advanced Graph Neural Networks73 references87 citations

TL;DR

D-VAE introduces a DAG variational autoencoder with asynchronous message passing to encode computations on DAGs, enabling latent-space Bayesian optimization for neural architecture search and Bayesian network structure learning.

ABSTRACT

Graph structured data are abundant in the real world. Among different graph types, directed acyclic graphs (DAGs) are of particular interest to machine learning researchers, as many machine learning models are realized as computations on DAGs, including neural networks and Bayesian networks. In this paper, we study deep generative models for DAGs, and propose a novel DAG variational autoencoder (D-VAE). To encode DAGs into the latent space, we leverage graph neural networks. We propose an asynchronous message passing scheme that allows encoding the computations on DAGs, rather than using existing simultaneous message passing schemes to encode local graph structures. We demonstrate the effectiveness of our proposed DVAE through two tasks: neural architecture search and Bayesian network structure learning. Experiments show that our model not only generates novel and valid DAGs, but also produces a smooth latent space that facilitates searching for DAGs with better performance through Bayesian optimization.

Motivation & Objective

• Motivate and address the challenge of optimizing directed acyclic graphs (DAGs) in applications like neural architecture search (NAS) and Bayesian network structure learning (BNSL).
• Propose a DAG-structured variational autoencoder (D-VAE) that encodes computations on DAGs rather than just graph structure.
• Develop an asynchronous message passing scheme within a graph neural network to respect DAG dependencies during encoding and decoding.
• Enable optimization in a continuous latent space to facilitate efficient DAG exploration via Bayesian optimization.

Proposed method

• Define a computation as a DAG composed of elementary operations.
• Use a graph neural network with asynchronous message passing to encode DAGs by waiting for all predecessors before updating a node (topological order).
• Provide an injective encoder mapping computations on DAGs to latent space (Theorem 2) under injective aggregation and update functions.
• Decode DAGs from latent vectors using a similar asynchronous, computation-aware generation process with node- and edge-prediction steps.
• Train with teacher forcing to maximize the variational lower bound and encourage reconstruction of DAGs.
• Compare against baselines adapted for DAGs (S-VAE, GraphRNN, GCN,DeepGMG) on reconstruction, validity, uniqueness, novelty, and downstream optimization.

Experimental results

Research questions

• RQ1 Can D-VAE injectively encode computations on DAGs so that each computation maps to a unique latent embedding?
• RQ2 Do latent representations learned by D-VAE support accurate reconstruction and generation of valid DAGs?
• RQ3 Is the latent space smooth and informative for predicting DAG performance and guiding optimization?
• RQ4 Does Bayesian optimization in the D-VAE latent space yield better neural architectures and Bayesian networks than baselines?

Key findings

• D-VAE achieves near-perfect reconstruction and high prior validity, with good novelty and reasonable uniqueness across tasks.
• Latent embeddings from D-VAE predict DAG performance well (lower RMSE and higher Pearson r) compared to baselines.
• Bayesian optimization in D-VAE latent space discovers higher-performing neural architectures than baselines (top accuracy ~94.80%).
• Bayesian network structures found by D-VAE achieve better BIC scores than those from baselines and approach the true generator.
• The learned latent space is smoother and more capable of differentiating high-score regions, enabling effective exploration.