Source-linked AI summary
Gradient-Based Neural DAG Learning
Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, Simon Lacoste-Julien
TL;DR
Learning causal DAGs from observational data is difficult because interventions are often unavailable and nonlinear relationships complicate structure search. The paper introduces GraN-DAG, a neural-network extension of continuous constrained DAG optimization, and finds it outperforms gradient-based continuous approaches on synthetic data while remaining competitive with greedy methods across considered tasks.
Problem
Causal graph learning often must rely on observational data because interventions are costly or impossible, while nonlinear relationships remain insufficiently handled by continuous optimization methods.
Method
GraN-DAG uses neural networks to parameterize nonlinear conditional distributions and enforces acyclicity through a differentiable weighted adjacency matrix within continuous constrained optimization.
Results
GraN-DAG outperforms NOTEARS and DAG-GNN on the considered synthetic data sets and is competitive with the best methods across the considered real and pseudo-real tasks.
Takeaways & Limitations
Continuous constrained optimization can support nonlinear DAG structure learning while providing a competitive alternative to nonlinear greedy-search procedures.
Takeaways & Limitations
The theoretical recovery guarantee is population-level and depends on model compatibility and identifiability, while practical optimization is non-convex and reaches only a stationary point.
Abstract
from arXiv · showhide
We propose a novel score-based approach to learning a directed acyclic graph (DAG) from observational data. We adapt a recently proposed continuous constrained optimization formulation to allow for nonlinear relationships between variables using neural networks. This extension allows to model complex interactions while avoiding the combinatorial nature of the problem. In addition to comparing our method to existing continuous optimization methods, we provide missing empirical comparisons to nonlinear greedy search methods. On both synthetic and real-world data sets, this new method outperforms current continuous methods on most tasks, while being competitive with existing greedy search methods on important metrics for causal inference.
1 INTRODUCTION
The paper addresses causal graph learning from observational data, where interventions are costly or impossible, by extending continuous DAG optimization to nonlinear relationships with neural networks. GraN-DAG is also evaluated against both continuous and nonlinear greedy-search methods.
- Observational data are used because interventions needed for causal-effect prediction are often costly or impossible.
- GraN-DAG extends the continuous constrained optimization framework of NOTEARS to learn nonlinear relationships using neural networks.
- The method focuses experimentally on nonlinear Gaussian additive noise models, which provide theoretical guarantees under specific assumptions.
- The paper adds empirical comparisons with CAM and GSF, nonlinear methods that use discrete greedy search.
- GraN-DAG is reported as competitive across the considered tasks while using preprocessing and postprocessing steps similar to CAM.
2 BACKGROUND
The background frames structure learning as DAG optimization and explains how NOTEARS replaces combinatorial search with continuous constrained optimization. Identifiability from observational data requires assumptions, such as nonlinear additive-noise structure with independent noises.
- 2.1 CAUSAL GRAPHICAL MODELS: The paper assumes no hidden variables and represents each node as one variable in a random vector generated by an underlying DAG.
- 2.1 CAUSAL GRAPHICAL MODELS: Causal graphical models extend Bayesian networks by assigning causal meaning to directed edges and supporting interventional-distribution queries.
- 2.2 STRUCTURE IDENTIFIABILITY: Recovering a DAG from observational samples alone is generally impossible, so assumptions are used to make the structure fully or partially identifiable.
- 2.2 STRUCTURE IDENTIFIABILITY: Under nonlinear structural equations with mild regularity conditions and mutually independent noises, the graph is identifiable from the observational distribution.
- 2.3 NOTEARS: CONTINUOUS OPTIMIZATION FOR STRUCTURE LEARNING: Structure learning maximizes a regularized score over DAGs, but the super-exponential number of DAGs typically motivates heuristic greedy search.
- 2.3 NOTEARS: CONTINUOUS OPTIMIZATION FOR STRUCTURE LEARNING: NOTEARS encodes a graph as a weighted adjacency matrix and replaces combinatorial optimization with a continuous, non-convex constrained problem solved using numerical optimization.
3 GRAN-DAG: GRADIENT-BASED NEURAL DAG LEARNING
GraN-DAG parameterizes conditional distributions with neural networks and constructs a differentiable weighted adjacency matrix whose acyclicity constraint defines a valid joint model. It optimizes likelihood approximately while using several measures to limit overfitting, with recovery guarantees only under identifiable population assumptions.
- Neural-network model: GraN-DAG learns one fully connected neural network per variable, mapping masked inputs to parameters of that variable’s conditional distribution.
- Neural-network model: The product of unconstrained neural-network conditionals is not generally a joint density because the implied graph need not be acyclic.
- Acyclicity constraint: The adapted acyclicity constraint makes the neural-network conditionals correspond to an acyclic graph and therefore to a valid joint density.
- Acyclicity constraint: Neural-network path products form connectivity matrices that quantify input-to-output dependence and support construction of the weighted adjacency matrix Aφ.
- Optimization: GraN-DAG approximately solves maximum likelihood with augmented-Lagrangian subproblems and stochastic gradient optimization.
- Guarantees and limitations: Under model compatibility and identifiable assumptions, the population optimum recovers the ground-truth graph, but empirical non-convex optimization reaches only a stationary point.
- 3.5 OVERFITTING: Overfitting is addressed through held-out early stopping, stochastic-gradient regularization, CAM-style pruning, and preliminary neighbor selection for graphs with 50 nodes or more.
4 EXPERIMENTS
Experiments evaluate GraN-DAG across synthetic, misspecified, real, and pseudo-real settings using SHD and SID, comparing continuous and greedy-search baselines. GraN-DAG is strongest among continuous methods on synthetic data and remains competitive overall, but performance depends on identifiability and model assumptions.
- Evaluation setup: Experiments compare GraN-DAG with continuous baselines DAG-GNN and NOTEARS and combinatorial methods CAM, GSF, GES, and PC using SHD and SID.SHD counts missing, false, or reversed edges; SID measures incorrectly calculated interventional distributions.
- Synthetic data: Synthetic datasets vary data-generating process, node count, edge sparsity, and graph type, using ER and scale-free graphs with 10 datasets of 1000 examples per type.Experiments include graphs with 10, 20, 50, and 100 nodes and edge counts of d or 4d.
- Synthetic data: GraN-DAG and CAM are best across synthetic settings on both SHD and SID, while GSF is usually close behind.Among continuous approaches, GraN-DAG performs best; NOTEARS struggles with nonlinear functions, and DAG-GNN is often near RANDOM in SID.
- Synthetic data: CAM can perform similarly to GraN-DAG because its restricted capacity may reduce variance despite additive-function bias, whereas GraN-DAG can improve with larger sample sizes.The paper presents this as a possible bias-variance explanation rather than a definitive causal account.
- Synthetic data: Under post nonlinear misspecification, GraN-DAG and CAM outperform DAG-GNN and NOTEARS in most comparisons, while their relative performance is generally similar.GSF is worse than GraN-DAG on PNL-GP but better in SID on PNL-MULT.
- Real and pseudo-real data: On real and pseudo-real data, all methods perform worse than on comparable synthetic graphs, while GraN-DAG remains competitive with the best method on each task.GraN-DAG beats CAM in SID on the protein data and obtains the best SHD on SynTReN, but not the best SID.
5 RELATED WORK
GraN-DAG belongs to score-based structure-learning methods and addresses the gap in nonlinear comparisons by using continuous optimization alongside nonlinear greedy baselines.
- Score-based methods optimize over DAG or CPDAG structures, often replacing combinatorial search with greedy procedures.
- GSF extends GES with a generalized score that models nonlinear relationships, while CAM separates node-order search from parent selection.
- Comparisons with nonlinear greedy methods such as CAM and GSF were previously missing, while GraN-DAG significantly outperformed DAG-GNN on the reported benchmarks.
- GraN-DAG and DAG-GNN were not designed for high-dimensional settings where samples are significantly fewer than nodes, unlike NOTEARS and CAM.
- Existing neural-network causal-discovery methods include SAM, which lacks acyclicity enforcement, and CGNN, which requires an initial skeleton for multivariate data.
6 CONCLUSION
GraN-DAG applies continuous constrained optimization to nonlinear structure learning through a neural-network characterization of acyclicity. It outperforms gradient-based approaches on synthetic data and remains competitive with greedy methods across the considered datasets.
- GraN-DAG uses continuous constrained optimization to support nonlinear relationships while avoiding task-specific greedy search procedures.
- Its acyclicity characterization extends the framework of Zheng et al. to neural networks.
- GraN-DAG outperforms NOTEARS and DAG-GNN on the considered synthetic datasets and is competitive on real and pseudo-real datasets.
- GraN-DAG is competitive with greedy approaches across all considered datasets.
A.1 OPTIMIZATION
GraN-DAG approximately solves augmented-Lagrangian subproblems with stochastic gradient optimization, using held-out convergence checks and fewer iterations than NOTEARS on Gauss-ANM data.
- The augmented Lagrangian is initialized with λ0 = 0 and µ0 = 10^-3 for all experiments.
- Each non-convex subproblem is approximately optimized with RMSprop, a stochastic-gradient variant used for neural-network training.
- Minibatch gradients are computed using standard deep-learning libraries, and convergence is assessed on held-out data when its objective stops increasing.
- The Lagrangian and penalty coefficients are updated after each subproblem, with η = 10 and γ = 0.9, and the next subproblem starts from the previous solution.
- GraN-DAG requires fewer total iterations than NOTEARS before augmented-Lagrangian convergence on Gauss-ANM data.
A.2 THRESHOLDING TO ENSURE ACYCLICITY
GraN-DAG uses Jacobian-based thresholding to remove effectively inactive edges and obtain an acyclic graph, because raw adjacency weights can misrepresent functional dependence.
- The learned weighted adjacency matrix may remain cyclic after thresholding, so additional edges must be removed to obtain a DAG.
- Removing edges by raw adjacency magnitude can retain edges with high weights even when the corresponding output is nearly independent of an input.
- High adjacency weights can reflect cancelling neural-network paths or saturated neurons rather than active functional influence.
- A zero adjacency entry is sufficient but not necessary for functional independence.
- The method instead ranks edges by the expected Jacobian, removes them from lowest to highest strength, and stops when acyclicity is achieved.
- Empirically, using the Jacobian matrix rather than the learned adjacency matrix yields better results.
A.3 PRELIMINARY NEIGHBORHOOD SELECTION AND DAG PRUNING
GraN-DAG uses preliminary neighborhood selection and pruning to restrict candidate parents and remove spurious edges. These steps reduce overfitting and substantially improve structural accuracy, while having little effect on SID.
- Preliminary Neighborhood Selection: PNS selects possible parents using extremely randomized trees before GraN-DAG optimization, especially for graphs with 50 nodes or more.The procedure fits one tree-based variable-selection model for each variable against all others.
- DAG Pruning: Pruning fits generalized additive models to current parent sets and removes parents with covariate-test p-values above 0.001.The pruning implementation follows CAM.
- Ablation Results: PNS and pruning substantially improve SHD while barely changing SID because they mainly remove spurious edges rather than alter node ordering.Without these steps, GraN-DAG can find a good ordering while retaining many false-positive edges.
- Comparison Methods: For NOTEARS and DAG-GNN, PNS and pruning yield smaller, statistically weak SHD gains and no significant SID improvement.The authors attribute this contrast partly to weaker overfitting in those methods and to their existing regularization or underfitting.
- Thresholding: Thresholding GraN-DAG edges below 0.33 yields SHD 29.4±15.9 and SID 85.6±45.7, improving over NOTEARS and DAG-GNN without PNS or pruning.The result supports stronger thresholding as an alternative explanation for GraN-DAG’s high unregularized SHD.
A.4 LARGE SAMPLE SIZE EXPERIMENT
The large-sample experiment tests whether GraN-DAG’s greater model capacity becomes advantageous as sample size increases. GraN-DAG eventually surpasses CAM on SID while remaining comparable on SHD, though larger-sample evaluation is left for future work.
- Results: As sample size grows, GraN-DAG outperforms CAM on SID while remaining on par on SHD.The comparison uses Gauss-ANM data with 50-node ER4 graphs, averaged over 10 datasets.
- Interpretation: The authors explain the pattern as reduced variance allowing GraN-DAG to exploit greater capacity, whereas CAM remains limited by modeling bias.Both methods use their respective default hyperparameter combinations.
- Scope: The experiment suggests GraN-DAG may be appealing when sample size is substantial.The paper focuses on sample sizes typically encountered in structure and causal learning and leaves broader sample-size analysis for future work.
A.5 DETAILS ON DATA SETS GENERATION
The supplementary material describes the synthetic graph and data-generating settings, evaluation measures, implementation choices, and additional experiments. Results remain favorable for GraN-DAG and CAM in several nonlinear or misspecified settings, with some computational and modeling boundaries.
- Dataset Generation: Synthetic datasets use randomly sampled ground-truth DAGs, typically with 10 datasets and 1000 examples per dataset.Unless otherwise stated, root variables are sampled from U[−1, 1].
- Data Models: The supplementary benchmarks cover Gauss-ANM, linear Gaussian, additive-function, post-nonlinear, SynTReN, and related synthetic settings.Gauss-ANM uses nonlinear functions with mutually independent noises, while other datasets vary the functional or noise assumptions.
- Supplementary Results: Across larger Gauss-ANM experiments, conclusions remain similar to the 10- and 50-node results, while GES and PC compare unfavorably with GraN-DAG and CAM.Figure 1 tracks weighted adjacency entries during training and shows most true edges recovered with few spurious edges.
- Computational Scope: GSF results are missing for two dataset types because its search exceeded 12 hours, while all other methods finished within 6 hours.This computational boundary occurred even with maximum in-degree limited to 5.
- Evaluation: SHD measures graph-edge discrepancies, whereas SID measures interventional distributions that would be miscalculated using the predicted graph.CPDAG outputs require special handling for SID, and SHD can compare DAGs with DAGs or CPDAGs.
- Implementation: GraN-DAG uses neural networks with RMSprop, while baselines generally use authors’ default hyperparameters and package defaults.Neural-network architecture and learning-rate choices are among GraN-DAG’s selected hyperparameters.
- Hyperparameter Selection: Hyperparameter selection benefits GraN-DAG and DAG-GNN most, although gains are rarely significant and no method improves in every setting.GES and GSF sometimes worsen, potentially because held-out scores poorly proxy graph quality or because tuning reduces training data.