Diffusion Models for Causal Discovery via Topological Ordering (2210.06201v2)

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 .

• The paper introduces DiffAN, which uses diffusion probabilistic models and topological ordering to streamline causal discovery.
• It employs denoising diffusion techniques to estimate scores and compute Hessians, enabling efficient analysis of complex, large-scale datasets.
• Empirical results demonstrate that DiffAN scales to 500 nodes and 100,000 samples, offering practical applications in fields like economics and healthcare.

Diffusion Models for Causal Discovery via Topological Ordering

This paper presents a new approach to causal discovery by leveraging diffusion models to identify and construct causal graphs from observational data through topological ordering. Developing causal structures from data is an NP-hard problem often constrained by computational limitations, especially when involving large datasets with numerous variables and samples. Traditional methods, such as those optimizing directed acyclic graphs (DAGs), require heavy computational resources. The authors propose DiffAN, a method rooted in diffusion probabilistic models (DPMs), to overcome these computational barriers.

Methodological Insights

DiffAN introduces a novel approach to causal discovery by addressing the scalability issues faced by conventional methods. It utilizes denoising diffusion techniques to approximate the data's score, thereby facilitating causal structure identification without the need for re-training the neural network repeatedly during topological ordering. The proposed method leverages gradient-based differentiation, enabling it to handle large datasets efficiently.

Key Components:

• Nonlinear Additive Noise Model (ANM): This is the foundational assumption for the causal discovery process. The paper leverages ANM's identifiability to ascertain the causal ordering, using the Hessian of the data log-likelihood.
• Score Estimation via Diffusion Models: The neural networks are trained using diffusion processes that approximate the score of the data distribution. DiffAN computes the Hessian of these scores, a crucial step in identifying leaf nodes iteratively.

Results and Implications

The empirical results underscore DiffAN's scalability, showing its capability to handle up to 500 nodes and 100,000 samples efficiently. Unlike previous methods, which falter at large sample sizes, DiffAN maintains practical runtime without compromising on the accuracy of causal discovery. Its strategic use of a subset of samples to compute the Hessian helps in achieving these results.

The implications of this work are significant:

• For practical applications, it offers a feasible solution to analyze causal structures in large datasets, relevant in fields such as economics and healthcare.
• Theoretically, it expands the capabilities of diffusion models, demonstrating their utility beyond generative tasks to causal inference and graph learning.

Future Directions

The methodology sets the stage for further exploration into scalable causal discovery techniques, particularly using deep learning approaches. The integration of score estimation and diffusion models can be extended to develop more robust, scalable methods for causal inference, potentially incorporating adaptive learning techniques.

In summary, DiffAN not only promises enhanced performance in causal discovery tasks but also heralds a new era of applying diffusion models to complex problems in data science, providing a pathway to addressing longstanding challenges in causal inference and graph-based learning models.