Amortized Bayesian Causal Discovery of Extended Factor Graphs

2026-07-24

AI Summary

Learning high-dimensional causal graphs—such as gene regulatory networks—poses significant challenges in scalability, unknown intervention targets, uncertainty quantification, and identifiability. This work proposes a Bayesian causal discovery method that leverages an extended factor graph formulation combined with amortized variational inference to enable efficient posterior inference on graphs with up to thousands of nodes, while rigorously preserving the directed acyclic graph constraint and model identifiability. Notably, the approach naturally accommodates settings where intervention targets are unknown. In synthetic benchmarks, the method achieves state-of-the-art accuracy with well-calibrated posterior uncertainties; when applied to single-cell perturbation data, it successfully recovers known growth factor gene targets and identifies novel ones.

Abstract

Learning causal graphs from interventional data is a challenging problem with broad applications. In molecular biology, for example, a central goal is to uncover gene regulatory networks from large-scale perturbation data. An ideal algorithm for this task should scale to thousands of nodes, incorporate interventions even when their targets are unknown, quantify uncertainty, and provide identifiability guarantees. However, existing approaches---e.g. approaches using score-based optimization or approximate Bayesian inference---often fail to meet all of these criteria. To address these limitations, we develop Amortized Bayesian Causal Discovery of Extended Factor Graphs (ABCDEFG). Our method guarantees exact acyclicity, scales to graphs with thousands of nodes, and naturally handles interventions even when their targets are unknown. Additionally, ABCDEFG estimates a posterior distribution whose maximum a posteriori estimate provably identifies the true causal graph up to an equivalence class. On simulated datasets, ABCDEFG achieves state-of-the-art accuracy, producing a well-calibrated posterior distribution while outperforming previous score-based and approximate Bayesian methods. Applied to large-scale single-cell perturbation data, ABCDEFG identifies both established and novel gene targets of growth factors.

Problem

• causal discovery
• interventional data
• gene regulatory networks
• Bayesian inference
• scalability

Innovation

• amortized inference
• Bayesian causal discovery
• interventional data
• acyclicity constraint
• factor graphs