From Observations to Causations: A GNN-based Probabilistic Prediction Framework for Causal Discovery

TL;DR

This paper introduces a GNN-based probabilistic framework for causal discovery that learns a distribution over causal graphs, improving accuracy and scalability over traditional methods.

Contribution

A novel GNN-based approach that models the entire space of causal graphs probabilistically, capturing complex structures directly from data.

Findings

Outperforms traditional causal discovery methods in accuracy.

Demonstrates scalability on large datasets.

Effective on both synthetic and real-world data.

Abstract

Causal discovery from observational data is challenging, especially with large datasets and complex relationships. Traditional methods often struggle with scalability and capturing global structural information. To overcome these limitations, we introduce a novel graph neural network (GNN)-based probabilistic framework that learns a probability distribution over the entire space of causal graphs, unlike methods that output a single deterministic graph. Our framework leverages a GNN that encodes both node and edge attributes into a unified graph representation, enabling the model to learn complex causal structures directly from data. The GNN model is trained on a diverse set of synthetic datasets augmented with statistical and information-theoretic measures, such as mutual information and conditional entropy, capturing both local and global data properties. We frame causal discovery as a supervised learning problem, directly predicting the entire graph structure. Our approach demonstrates superior performance, outperforming both traditional and recent non-GNN-based methods, as well as a GNN-based approach, in terms of accuracy and scalability on synthetic and real-world datasets without further training. This probabilistic framework significantly improves causal structure learning, with broad implications for decision-making and scientific discovery across various fields.