Graph Neural Networks for Probabilistic Causal Discovery

Rezaur Rashid, Gabriel Terejanu

Keywords: causal graphs, causal pairs, causal structure learning, DAG, graph neural networks, observational data

Abstract: Conventional causal discovery algorithms face significant challenges in dealing with large-scale observational datasets and in capturing global structural information. ⁤⁤To address these limitations, we introduce a novel graph neural network (GNN)--based probabilistic framework for causal structure learning that generates a probability distribution over the entire graph space. By encoding the node and edge attributes into a unified graph representation, our framework enables the GNN to learn the complex causal structure directly from the data augmented with statistical and information-theoretic measures, which exploit the local and global data properties. ⁤⁤Our approach outperforms benchmark methods, both traditional and recent non-GNN-based, in terms of accuracy and scalability on synthetic and real-world datasets. Notably, our framework advances the causal discovery paradigm by generating a probability distribution over the causal graphs, rather than learning a single causal graph.

Submission Number: 28