Causal Discovery via Bayesian Optimization

2025-01-25

AI Summary

This work addresses causal graph learning from observational data. We propose DrBO, the first framework to integrate Bayesian optimization (BO) into directed acyclic graph (DAG) structure search for efficient approximation of the optimal causal graph. To overcome the scalability limitations of Gaussian processes in high-dimensional graph spaces—traditional BO’s bottleneck—we innovatively employ a continually trained dropout neural network as the surrogate model, jointly achieving reliable uncertainty estimation, strong generalization, and computational efficiency. DrBO jointly optimizes both the DAG constraint and the causal scoring function. Empirical evaluation on synthetic and real-world benchmarks demonstrates that DrBO significantly outperforms state-of-the-art methods: it recovers higher-fidelity causal structures with fewer iterations and substantially reduced runtime. The implementation is publicly available.

Technology Category

• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies

Abstract

Existing score-based methods for directed acyclic graph (DAG) learning from observational data struggle to recover the causal graph accurately and sample-efficiently. To overcome this, in this study, we propose DrBO (DAG recovery via Bayesian Optimization)-a novel DAG learning framework leveraging Bayesian optimization (BO) to find high-scoring DAGs. We show that, by sophisticatedly choosing the promising DAGs to explore, we can find higher-scoring ones much more efficiently. To address the scalability issues of conventional BO in DAG learning, we replace Gaussian Processes commonly employed in BO with dropout neural networks, trained in a continual manner, which allows for (i) flexibly modeling the DAG scores without overfitting, (ii) incorporation of uncertainty into the estimated scores, and (iii) scaling with the number of evaluations. As a result, DrBO is computationally efficient and can find the accurate DAG in fewer trials and less time than existing state-of-the-art methods. This is demonstrated through an extensive set of empirical evaluations on many challenging settings with both synthetic and real data. Our implementation is available at https://github.com/baosws/DrBO.

Problem

• Causal Inference
• Directed Acyclic Graphs (DAGs)
• Machine Learning

Innovation

• Bayesian Optimization
• Causal DAG Learning
• Neural Network Prediction