[Paper Review] Towards Scalable Bayesian Learning of Causal DAGs

Jussi Viinikka, Antti Hyttinen+2 more|arXiv (Cornell University)|Sep 30, 2020

Bayesian Modeling and Causal Inference22 references4 citations

TL;DR

This paper presents Gadget, a scalable MCMC method for Bayesian structure learning of causal DAGs, and Beeps, a novel Bayesian approach for causal effect estimation in linear Gaussian models. By optimizing candidate parent selection and reducing memory/time costs via algorithmic improvements, the methods enable efficient posterior sampling over DAGs with up to 100+ variables and outperform prior approaches in causal effect estimation, especially with limited data.

ABSTRACT

We give methods for Bayesian inference of directed acyclic graphs, DAGs, and the induced causal effects from passively observed complete data. Our methods build on a recent Markov chain Monte Carlo scheme for learning Bayesian networks, which enables efficient approximate sampling from the graph posterior, provided that each node is assigned a small number $K$ of candidate parents. We present algorithmic techniques to significantly reduce the space and time requirements, which make the use of substantially larger values of $K$ feasible. Furthermore, we investigate the problem of selecting the candidate parents per node so as to maximize the covered posterior mass. Finally, we combine our sampling method with a novel Bayesian approach for estimating causal effects in linear Gaussian DAG models. Numerical experiments demonstrate the performance of our methods in detecting ancestor-descendant relations, and in causal effect estimation our Bayesian method is shown to outperform previous approaches.

Motivation & Objective

• Address the computational intractability of Bayesian model averaging over the combinatorial space of DAGs in large-scale causal inference.
• Overcome the limitations of existing MCMC samplers for DAGs, including high memory usage and biased sampling due to over-representation of certain orderings.
• Enable efficient and accurate posterior sampling over DAG structures with hundreds of variables by reducing time and space complexity.
• Develop a Bayesian framework for estimating causal effects that leverages full DAG structure and uncertainty, improving upon non-Bayesian IDA-style methods.
• Investigate optimal candidate parent selection strategies to maximize posterior coverage while maintaining computational feasibility.

Proposed method

• Use a partition-based MCMC scheme that samples over ordered node partitions instead of full DAGs or node orderings, reducing sampling bias.
• Restrict each node’s parent set to a small candidate set of size $ K $, precomputing score sums in a lookup table to drastically reduce per-step computation.
• Introduce algorithmic optimizations—such as efficient data structures and pruning strategies—to reduce memory and time complexity, enabling larger $ K $ values (e.g., $ K=15 $) with 2–3 orders of magnitude savings.
• Propose a three-stage sampling-based method (Beeps) to approximate the posterior of causal effects: (1) sample DAGs via Gadget, (2) conditionally sample model parameters per DAG, (3) compute causal effects via matrix inversion in linear Gaussian models.
• Use matrix inversion to map structural parameters to interventional distributions, enabling exact causal effect computation under linear Gaussian assumptions.
• Evaluate candidate parent selection via exact optimization for small graphs and scalable heuristics for larger ones, aiming to maximize posterior coverage.

Experimental results

Research questions

• RQ1 How can we scale Bayesian DAG learning to larger networks (e.g., 100+ variables) while maintaining computational feasibility and sampling accuracy?
• RQ2 What is the optimal strategy for selecting candidate parents per node to maximize posterior coverage of the true DAG structure?
• RQ3 Can we design a Bayesian method for causal effect estimation that fully leverages structural uncertainty and outperforms non-Bayesian approaches like IDA?
• RQ4 How do memory and time requirements scale with increasing $ K $, and what algorithmic optimizations can reduce these costs by orders of magnitude?
• RQ5 To what extent do heuristic parent selection methods approximate the performance of optimal selection, and when do they fail?

Key findings

• The proposed algorithmic optimizations reduce memory and time requirements by 2–3 orders of magnitude, enabling feasible use of $ K=15 $, which covers a significantly larger portion of the posterior DAG space.
• Beeps outperforms previous IDA-based methods in causal effect estimation, particularly with small to moderate sample sizes, demonstrating improved accuracy and uncertainty quantification.
• On a real-world gene expression dataset with $ n=107 $ variables, Beeps achieved better accuracy than BiDAG and Gadget, especially with fewer data points (e.g., 200–400), where BiDAG often failed to converge.
• For datasets with 800 and 1600 data points, BiDAG failed to complete 100 iterations in 4/10 and 8/10 cases, respectively, while Gadget with $ K=15 $ remained stable and scalable.
• Optimal candidate parent sets yield strong posterior coverage even with moderate $ K $, and simple heuristics often achieve near-optimal performance, though not always—highlighting the need for further research on efficient high-scoring parent set discovery.
• The Bayesian causal effect estimator (Beeps) provides a full posterior distribution over causal effects, enabling principled uncertainty-aware decision-making, and outperforms non-Bayesian alternatives in terms of estimation accuracy.