[Paper Review] Diffusion Models for Causal Discovery via Topological Ordering

Pedro L. Sánchez, Xiao Liu+2 more|arXiv (Cornell University)|Oct 12, 2022

Bayesian Modeling and Causal Inference4 citations

TL;DR

This paper proposes DiffAN, a scalable causal discovery method that leverages denoising diffusion probabilistic models (DPMs) to estimate the Hessian of the data log-likelihood for topological ordering in additive noise models. By training a neural score model once and efficiently updating it via a 'deciduous score' mechanism without retraining, DiffAN enables causal discovery on high-dimensional datasets with up to 500 variables and 10⁵ samples, outperforming prior methods in scalability while maintaining state-of-the-art accuracy on small datasets.

ABSTRACT

Discovering causal relations from observational data becomes possible with additional assumptions such as considering the functional relations to be constrained as nonlinear with additive noise (ANM). Even with strong assumptions, causal discovery involves an expensive search problem over the space of directed acyclic graphs (DAGs). \emph{Topological ordering} approaches reduce the optimisation space of causal discovery by searching over a permutation rather than graph space. For ANMs, the \emph{Hessian} of the data log-likelihood can be used for finding leaf nodes in a causal graph, allowing its topological ordering. However, existing computational methods for obtaining the Hessian still do not scale as the number of variables and the number of samples increase. Therefore, inspired by recent innovations in diffusion probabilistic models (DPMs), we propose \emph{DiffAN}\footnote{Implementation is available at \url{https://github.com/vios-s/DiffAN} .}, a topological ordering algorithm that leverages DPMs for learning a Hessian function. We introduce theory for updating the learned Hessian without re-training the neural network, and we show that computing with a subset of samples gives an accurate approximation of the ordering, which allows scaling to datasets with more samples and variables. We show empirically that our method scales exceptionally well to datasets with up to $500$ nodes and up to $10^5$ samples while still performing on par over small datasets with state-of-the-art causal discovery methods. Implementation is available at https://github.com/vios-s/DiffAN .

Motivation & Objective

• To address the scalability bottleneck in causal discovery for high-dimensional datasets with many variables and samples.
• To overcome the computational infeasibility of kernel-based Hessian estimation in large-scale settings.
• To enable efficient topological ordering by learning a score function via denoising diffusion models.
• To develop a method that updates the score function without retraining after each leaf node removal, preserving efficiency.
• To achieve state-of-the-art performance on small datasets while scaling to large-scale problems with up to 500 nodes and 10⁵ samples.

Proposed method

• Training a denoising diffusion probabilistic model (DPM) to learn the score function of the data distribution, enabling efficient Hessian computation via backpropagation.
• Using the DPM's score and its full Jacobian to compute the Hessian of the log-likelihood for identifying leaf nodes in the causal graph.
• Introducing the 'deciduous score' framework, which analytically updates the score function after removing a leaf node without retraining the neural network.
• Applying the updated score to a subsample of data for Hessian evaluation, enabling constant-time complexity in sample size and enabling scalability.
• Iteratively identifying and removing leaf nodes based on the Hessian's diagonal elements, reconstructing the causal DAG through topological ordering.
• Leveraging the fact that DPMs can generalize to out-of-distribution points and support subsampling, unlike kernel-based methods.

Experimental results

Research questions

• RQ1 Can denoising diffusion models be effectively used to estimate the Hessian of the data log-likelihood for causal discovery?
• RQ2 Can the score function be updated efficiently after each leaf node removal without retraining the neural network?
• RQ3 Does the proposed method scale to datasets with up to 500 variables and 10⁵ samples, where prior methods fail?
• RQ4 How does the performance of the method compare to state-of-the-art causal discovery algorithms on small datasets?
• RQ5 Can subsampling during Hessian computation maintain accuracy while drastically reducing computational cost?

Key findings

• DiffAN scales to datasets with up to 500 variables and 10⁵ samples, while existing methods like SCORE fail beyond 2,000 samples on a 64GB RAM machine.
• The method maintains performance on par with state-of-the-art methods on small datasets (e.g., 50 variables), demonstrating competitive accuracy.
• The deciduous score mechanism allows Hessian computation with constant complexity in sample size, enabling efficient updates after each leaf removal.
• Using a subsample for Hessian evaluation yields accurate topological ordering, validating the approximation strategy.
• The DPM-based score estimation enables generalization to out-of-sample points, unlike kernel-based methods that require full kernel matrices.
• The approach avoids architectural constraints on the neural network, unlike prior deep learning-based causal discovery methods.