[Paper Review] Differentiable DAG Sampling

Bertrand Charpentier, Simon Kibler+1 more|arXiv (Cornell University)|Mar 16, 2022

Bayesian Modeling and Causal Inference4 citations

TL;DR

This paper introduces DP-DAG, a differentiable probabilistic model for fast, differentiable sampling of directed acyclic graphs (DAGs) by first sampling a topological ordering and then edges consistent with it. It further proposes VI-DP-DAG, a variational inference method that learns DAG structures from observational data without requiring non-differentiable post-processing, achieving state-of-the-art performance and up to 10x faster training than baselines like GraN-DAG and Masked-DAG.

ABSTRACT

We propose a new differentiable probabilistic model over DAGs (DP-DAG). DP-DAG allows fast and differentiable DAG sampling suited to continuous optimization. To this end, DP-DAG samples a DAG by successively (1) sampling a linear ordering of the node and (2) sampling edges consistent with the sampled linear ordering. We further propose VI-DP-DAG, a new method for DAG learning from observational data which combines DP-DAG with variational inference. Hence,VI-DP-DAG approximates the posterior probability over DAG edges given the observed data. VI-DP-DAG is guaranteed to output a valid DAG at any time during training and does not require any complex augmented Lagrangian optimization scheme in contrast to existing differentiable DAG learning approaches. In our extensive experiments, we compare VI-DP-DAG to other differentiable DAG learning baselines on synthetic and real datasets. VI-DP-DAG significantly improves DAG structure and causal mechanism learning while training faster than competitors.

Motivation & Objective

• To address the challenge of differentiable DAG learning from observational data, where existing methods rely on computationally expensive augmented Lagrangian schemes or non-differentiable post-processing.
• To develop a probabilistic model that enables fast, differentiable sampling of valid DAGs through structured sampling of permutations and edge sets.
• To design a variational inference framework (VI-DP-DAG) that guarantees valid DAG outputs at all training steps and supports end-to-end optimization.
• To improve both DAG structure and causal mechanism learning while significantly reducing training time compared to state-of-the-art baselines.

Proposed method

• DP-DAG samples a DAG in two steps: first, it uses Gumbel-Sinkhorn or Gumbel-Softmax to differentiably sample a topological ordering (permutation) of nodes.
• Second, it uses Gumbel-Top-k or Gumbel-Softmax to sample edges that are consistent with the sampled ordering, ensuring acyclicity.
• The method leverages differentiable relaxations of discrete distributions to enable gradient-based optimization over the space of valid DAGs.
• VI-DP-DAG combines DP-DAG with variational inference to approximate the posterior over DAG edges given observed data, using a variational distribution parameterized by neural networks.
• The approach avoids the need for complex dual ascent or non-differentiable pruning steps, ensuring valid DAGs at every training step.
• The model is trained via variational lower bound maximization, with edge probabilities derived from a learned latent variable model conditioned on the topological ordering.

Experimental results

Research questions

• RQ1 Can a differentiable probabilistic model over DAGs be designed to enable fast and reliable sampling without relying on non-differentiable post-processing?
• RQ2 How does VI-DP-DAG compare to existing differentiable DAG learning methods in terms of structure and causal mechanism learning performance?
• RQ3 Does the differentiable sampling mechanism in DP-DAG lead to more robust and confident predictions under data perturbations compared to baselines?
• RQ4 Can VI-DP-DAG achieve faster training speeds while maintaining or improving performance on synthetic and real-world datasets?
• RQ5 To what extent does the choice of threshold in edge discretization affect causal mechanism recovery in VI-DP-DAG?

Key findings

• VI-DP-DAG achieves state-of-the-art performance on DAG structure learning, outperforming GraN-DAG and Masked-DAG on AUC-PR and AUC-ROC metrics across synthetic and real datasets.
• On the ER-100-400 dataset, VI-DP-DAG trains over 10x faster than GraN-DAG and avoids the need for time-consuming non-differentiable processing steps.
• The PNS pre-processing step is ~4x slower than VI-DP-DAG on ER-100-400, while CAM is over 900x slower and did not complete within 2 days.
• The CAM algorithm and DAG pruning post-processing are respectively ~23x and ~4x slower than VI-DP-DAG on large graphs, highlighting its efficiency advantage.
• VI-DP-DAG assigns consistently decreasing confidence scores to perturbed graphs, indicating reliable uncertainty estimation, unlike noisier baselines such as GraN-DAG and Masked-DAG.
• Causal mechanism learning in VI-DP-DAG is robust to threshold selection, with MSE scores remaining stable except near threshold 1, where sparsity causes loss of Granger-causal relationships.