PaperlayerView exact revision on arXiv

Source-linked AI summary

DAG-GNN: DAG Structure Learning with Graph Neural Networks

Yue Yu, Jie Chen, Tian Gao, Mo Yu

arXiv:1904.10098v1cs.LGcs.AIstat.ML

TL;DR

Learning a faithful DAG from samples is difficult because the graph search space is superexponential and existing continuous methods are limited by linear SEM assumptions. The paper proposes DAG-GNN, a VAE-based deep generative model with graph neural network parameterization and a polynomial acyclicity constraint. DAG-GNN outperforms DAG-NOTEARS on synthetic data and compares favorably with combinatorial Bayesian-information-criterion optimization on benchmark data.

  • Problem

    DAG learning is computationally difficult, and existing continuous optimization methods are limited by linear SEM and least-squares assumptions.

  • Method

    DAG-GNN uses variational inference with novel graph neural network encoder and decoder architectures plus a polynomial acyclicity constraint.

  • Results

    DAG-GNN outperforms DAG-NOTEARS on synthetic data and compares favorably with graphs obtained by combinatorial Bayesian-information-criterion optimization on benchmark data.

  • Takeaways & Limitations

    The model supports complex distributions, discrete variables, and vector-valued variables within DAG structure learning.

  • Takeaways & Limitations

    The graph neural network activation in the decoder must match the type of the variable X.

Abstract

from arXiv · show

Learning a faithful directed acyclic graph (DAG) from samples of a joint distribution is a challenging combinatorial problem, owing to the intractable search space superexponential in the number of graph nodes. A recent breakthrough formulates the problem as a continuous optimization with a structural constraint that ensures acyclicity (Zheng et al., 2018). The authors apply the approach to the linear structural equation model (SEM) and the least-squares loss function that are statistically well justified but nevertheless limited. Motivated by the widespread success of deep learning that is capable of capturing complex nonlinear mappings, in this work we propose a deep generative model and apply a variant of the structural constraint to learn the DAG. At the heart of the generative model is a variational autoencoder parameterized by a novel graph neural network architecture, which we coin DAG-GNN. In addition to the richer capacity, an advantage of the proposed model is that it naturally handles discrete variables as well as vector-valued ones. We demonstrate that on synthetic data sets, the proposed method learns more accurate graphs for nonlinearly generated samples; and on benchmark data sets with discrete variables, the learned graphs are reasonably close to the global optima. The code is available at \url{https://github.com/fishmoon1234/DAG-GNN}.

1. Introduction

DAG structure learning is difficult because exact search is superexponential, while common continuous approaches rely on restricted linear SEM assumptions. DAG-GNN addresses these limits with a VAE-based graph model that supports richer variables and a practical acyclicity constraint.

  • DAG structure learning is NP-hard, with a search space superexponential in the number of graph nodes.
  • Linear SEM and least-squares objectives are convenient but may not reflect real-life data distributions.
  • DAG-GNN uses variational inference with graph neural network encoder and decoder parameterizations to capture complex distributions faithful to the DAG.
  • The VAE framework handles both continuous and discrete variables by matching the decoder likelihood to variable types.
  • Graph neural network parameterization supports vector-valued variables as node features input to and output of the networks.
  • The proposed polynomial acyclicity constraint is more convenient for deep-learning platforms and is reported as numerically stable as the matrix exponential.
  • Across synthetic and benchmark settings, DAG-GNN outperforms DAG-NOTEARS on synthetic data and compares favorably with Bayesian-information-criterion graphs.

2. Background and Related Work

DAG learning recovers a graph from samples when the graph and distribution satisfy faithfulness, but traditional exact search is computationally difficult. Related methods therefore use approximate search, independence tests, or continuous optimization, with recent approaches extending toward neural models.

  • Faithfulness means that a DAG and distribution entail exactly the same conditional independencies, enabling graph recovery from the distribution.
  • Score-based DAG methods evaluate graph fit and search over graphs, while constraint-based methods use conditional-independence tests to assess edges.
  • Approximate methods impose assumptions such as bounded tree-width, tree-like structure, or other DAG constraints to make search tractable.
  • Continuous optimization transforms discrete DAG search into an equality constraint usable by gradient-based methods, but the approach was applied only to linear SEM.
  • Early neural-network approaches included GAN-style models, but the cited approach did not enforce acyclicity.

3. Neural DAG Structure Learning

DAG-GNN generalizes linear SEM-based DAG learning with a deep generative model whose encoder and decoder are parameterized by graph neural networks. Variational inference makes the model tractable while preserving an explicit learnable adjacency matrix.

  • DAG-GNN learns a DAG’s weighted adjacency matrix using a deep generative model that generalizes linear SEM.
  • The model rewrites sampling as X = f_A(Z), viewing graph neural networks as functions that transform node features Z into representations X.
  • Its architecture combines graph-dependent transforms with parameterized functions f1 and f2, yielding a generalized version of linear SEM when f2 is invertible.
  • Variational Bayes replaces intractable log-evidence maximization with an ELBO using a variational posterior q(Z|X).
  • The resulting VAE uses an encoder to infer latent Z and a decoder to reconstruct X, with both densities parameterized by neural networks.

3.4. Architecture and Loss Function

The continuous-variable architecture uses a matrix-normal prior, Gaussian encoder and decoder distributions, and MLP-based transformations around the DAG-dependent triangular solve. Its ELBO is estimated from a closed-form KL term and Monte Carlo reconstruction term.

  • The VAE starts with matrix-valued X and latent Z, both represented as m×d matrices.
  • The encoder computes Gaussian means and standard deviations through (I − A^T) MLP(X, W1, W2).
  • The decoder models p(X|Z) as a factored Gaussian using MLP((I − A^T)^-1Z, W3, W4).
  • The architecture emphasizes nonlinear transformation of the linear-SEM sample (I − A^T)^-1Z, which the authors associate with better nonlinearity capture than an alternative design.
  • The KL-divergence term has a closed form, while reconstruction accuracy is computed with Monte Carlo samples from q(Z|X).

3.5. Discrete Variables

DAG-GNN accommodates discrete variables by encoding each variable as a one-hot row and replacing the Gaussian decoder likelihood with a categorical distribution produced by row-wise softmax.

  • The method naturally handles discrete variables with finite support of cardinality d.
  • Each discrete variable is represented by a one-hot row whose active location identifies its value.
  • For discrete data, p(X|Z) becomes a factored categorical distribution with probability matrix PX.
  • The decoder changes from Gaussian outputs MX, SX to probability outputs PX using a row-wise softmax.
  • The KL term remains unchanged, while the ELBO reconstruction term is modified for the categorical likelihood.

3.6. Connection to Linear SEM

DAG-GNN extends the linear SEM by adding nonlinear mappings and a variational latent representation, while retaining a direct connection to the least-squares formulation. A separate matrix-based equality constraint is used to enforce acyclicity during ELBO optimization.

  • The proposed model applies nonlinear transformations to SEM sampling, treats the result as a decoder, and pairs it with a variational encoder.
  • Stripping away variational components shows that reconstruction error and latent regularization correspond to the ELBO’s reconstruction and KL terms.
  • Removing the nonlinear mappings yields perfect reconstruction and reduces the sample loss to the least-squares loss used by Zheng et al. (2018).
  • Neither ELBO maximization nor least-squares minimization alone guarantees acyclicity, so the objective is paired with an equality constraint.
  • The proposed alternative constraint uses (I + αB)^m and is computationally less numerically difficult than exp(B) when α is chosen judiciously.
  • The hyperparameter α depends on an estimate of B’s spectral radius, which is bounded by its maximum row sum under nonnegativity.

3.8. Training

Training solves the nonlinear equality-constrained learning problem with an augmented Lagrangian. The procedure progressively strengthens the acyclicity penalty while updating the Lagrange multiplier and minimizing sample-based ELBO subproblems.

  • The unknowns are the adjacency matrix A and all VAE parameters θ, optimized under a nonlinear equality constraint.
  • The augmented Lagrangian approach converts the constrained problem into a sequence of unconstrained minimizations.
  • Increasing c drives the minimizer toward h(A) = 0, where the augmented Lagrangian equals the objective function.
  • The Lagrange multiplier λ is updated alongside the increasing penalty so it converges toward the optimality-condition multiplier.
  • Each ELBO subproblem can be solved with blackbox stochastic optimization because the objective is defined on samples.

4. Experiments

Experiments compare DAG-GNN with DAG-NOTEARS on synthetic linear, nonlinear, and vector-valued data, then evaluate discrete benchmarks and applications. DAG-GNN generally improves structure recovery, handles vector-valued and discrete variables, and achieves the lowest SHD on the protein-signaling application.

  • Experimental design: Experiments evaluate DAG-GNN against DAG-NOTEARS on synthetic data, GOPNILP on discrete benchmarks, and DAG-NOTEARS and FSG on applications.Synthetic evaluations report SHD and FDR across graph sizes and repeated trials.
  • Linear case: In the linear case, DAG-GNN learns substantially more accurate graphs than DAG-NOTEARS when graphs are large.
  • Nonlinear case: For h(x) = cos(x + 1), DAG-GNN slightly improves SHD and reduces FDR by approximately a factor of three versus DAG-NOTEARS.The parameter heat map also indicates fewer false alarms and a relatively sparser estimated matrix.
  • Nonlinear case: With higher nonlinearity, DAG-GNN achieves significantly better SHD and FDR than DAG-NOTEARS.
  • Vector-valued case: For vector-valued data with d = 5 and dZ = 1, DAG-GNN significantly outperforms DAG-NOTEARS, recovering all ground-truth edges with similar estimated weights.DAG-NOTEARS barely learns the graph in this setting.
  • Benchmark data sets: On discrete benchmarks, DAG-GNN results are reasonably close to the ground truth, while GOPNILP results are nearly optimal.The BIC-score gap may reflect the simple autoencoder's weaker approximation of multinomial distributions.
  • Applications: On the protein-signaling network, DAG-GNN achieves the lowest SHD and recovers 8 of 20 ground-truth edges in an acyclic estimated graph.It also predicts 5 indirectly connected edges and 3 reverse edges.

5. Conclusion

The paper addresses NP-hard DAG structure learning by combining continuous acyclicity constraints with neural functional approximation. DAG-GNN uses VAE-based graph neural networks to model broader data types and demonstrates practical competitiveness across synthetic, benchmark, and application data.

  • DAG-GNN combines a continuous acyclicity constraint with a deep generative model to recover DAGs from complex data distributions.
  • Variational autoencoders and new graph neural network architectures support scalar, vector-valued, continuous, and discrete variables.
  • Extensive synthetic, benchmark, and application experiments demonstrate the proposal's practical competitiveness.

A. Proofs

The proofs establish the trace-based acyclicity condition and bound the associated exponential expression using monotonicity.

  • A. Proofs: The proof begins by defining B = A ◦ A and observing that B is nonnegative.
  • A. Proofs: A nonnegative adjacency-derived matrix B has no cycle exactly when tr[(I + αB)^m] = m.The argument uses tr(B^k) > 0 as the cycle criterion and the fact that any cycle has length at most m.
  • A. Proofs: Theorem 2’s proof introduces an additional expression, but the supplied passage does not state its subsequent derivation.
  • A. Proofs: For fixed c and |λ|, the function (1 + α|λ|)^m increases with positive m and approaches e^(c|λ|).Therefore, for finite m > 0, (1 + α|λ|)^m ≤ e^(c|λ|).

B. Structure Learning over KB Relations

The KB-relations task represents entities as samples and relations as variables, then reports high-confidence extracted edges while accounting for compound FreeBase value paths.

  • B. Structure Learning over KB Relations: Each sample represents an entity, each variable represents a FreeBase relation, and samples contain 7.36 relations on average.The data come from FB15K-237, containing approximately 15k entities and 237 relations.
  • B. Structure Learning over KB Relations: Table 4 reports high-confidence extracted edges by ranking relations within the same domain for each target relation.Relations common to associated entity types are omitted from the left-hand side.
  • B. Structure Learning over KB Relations: A dotted relation path R1.R2 denotes a compound FreeBase value path through a virtual node to a real entity.
Loading 1904.10098v1…