Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Causal Discovery via Bayesian Optimization

Bao Duong, Sunil Gupta, Thin Nguyen

TL;DR

DrBO tackles the efficiency gap in score-based observational causal discovery by reframing DAG search as unconstrained Bayesian optimization over a low-rank DAG space. It replaces cubic-scale GPs with dropout neural networks as scalable surrogates, and uses node-wise local scores to form an indirect surrogate that couples with BIC-type objectives such as $S_{ ext{BIC-NV}}$ and $S_{ ext{BIC-EV}}$. A linear-in-$d$ search space is achieved via a Vec2DAG-inspired mapping $ au(oldsymbol{p},oldsymbol{R})$, with continual training and targeted pruning enhancing both speed and accuracy. Empirical results on synthetic linear and nonlinear data, as well as real Sachs and bnLearn structures, demonstrate superior sample efficiency and scalability, with ablations validating the contribution of each design choice. The work provides a practical, high-performance framework for causal discovery from observational data and is accompanied by public code.

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.

Causal Discovery via Bayesian Optimization

TL;DR

DrBO tackles the efficiency gap in score-based observational causal discovery by reframing DAG search as unconstrained Bayesian optimization over a low-rank DAG space. It replaces cubic-scale GPs with dropout neural networks as scalable surrogates, and uses node-wise local scores to form an indirect surrogate that couples with BIC-type objectives such as and . A linear-in- search space is achieved via a Vec2DAG-inspired mapping , with continual training and targeted pruning enhancing both speed and accuracy. Empirical results on synthetic linear and nonlinear data, as well as real Sachs and bnLearn structures, demonstrate superior sample efficiency and scalability, with ablations validating the contribution of each design choice. The work provides a practical, high-performance framework for causal discovery from observational data and is accompanied by public code.

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.
Paper Structure (60 sections, 2 theorems, 21 equations, 15 figures, 9 tables, 1 algorithm)

This paper contains 60 sections, 2 theorems, 21 equations, 15 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Bayesian Optimization
  5. Structural Causal Model
  6. Score-based Causal Discovery
  7. Parametrized DAG generation
  8. $\mathbf{DrBO}$: DAG Recovery via Bayesian Optimization
  9. Search Space
  10. Acquisition Function Optimization
  11. Surrogate Modeling with Dropout Networks
  12. Indirect Surrogate Modeling
  13. Continual Model Training
  14. Finalizing the Result
  15. Experiments
  16. ...and 45 more sections

Key Result

Lemma 1

For all $d,k\in\mathbb{N}^{+}$, $\mathbf{p}\in\mathbb{R}^{d}$ and $\mathbf{R}\in\mathbb{R}^{d\times k}$, let $\tau:\mathbb{R}^{d}\times\mathbb{R}^{d\times k}\rightarrow\left\{ 0,1\right\} ^{d\times d}$ be defined as in Eq. (eq:vec2dag-lr). Then, $\tau\left(\mathbf{p},\mathbf{R}\right)$ represents a

Figures (15)

  • Figure 1: DAG learning results on Synthetic data.First column: overall performance in terms of True Positive Rate (TPR, higher is better), Precision, Recall, and $F_{1}$ score (higher is better), as well as False Discovery Rate (FDR, lower is better). Second column: we track the best Bayesian Information Criterion (BIC, higher is better) so far at every optimization step. Third and Fourth columns: we monitor the Structural Hamming Distance (SHD, lower is better) of the DAG whose best BIC so far at every optimization step. Shaded areas in the line plots indicate 95% confidence intervals over 5 random datasets. NOTEARS+TMPI usually stops early before the time limit.
  • Figure 2: Causal Discovery Performance on the Benchmark Sachs DatasetSachs_etall_05Causal. NOTEARS+TMPI stops early before the max no. of evaluations is reached.
  • Figure 3: Ablating our design choices. All configurations are evaluated on 5 linear-Gaussian datasets of $1,\!000$ samples on 30ER8 graphs. Shaded areas indicate 95% confidence intervals.
  • Figure 4: Causal Discovery Performance with Varying Sample Sizes. We apply our $\mathbf{DrBO}$ method on linear-Gaussian data with 20ER4 graphs. Shaded areas represent 95% confidence interval over 5 runs.
  • Figure 5: BGe for Markov Equivalence Class Discovery. We compare the BGe score of ground truth DAGs and the estimations from $\mathbf{DrBO}$ with two popular baselines PC Spirtes_etal_00Causation and GES Chickering_02Optimal. Each point corresponds to one of 50 random datasets with linear Gaussian data on ER graphs of 5 nodes and 5 edges on average. The Pearson correlation coefficient $r$ between the scores of the estimated and ground truth DAGs are included. In addition, we also report the SHD-C metric, which measures the structural distance between MECs.
  • ...and 10 more figures

Theorems & Definitions (8)

  • Example
  • Lemma 1
  • Lemma 2
  • proof
  • proof
  • Remark 1
  • Remark 2
  • Remark 3