Sample, estimate, aggregate: A recipe for causal discovery foundation models
Menghua Wu, Yujia Bao, Regina Barzilay, Tommi Jaakkola
TL;DR
This work introduces Sea, a foundation-model-inspired framework for causal discovery that predicts a global graph $G=(V,E)$ from a mix of subset-level marginal estimates and global statistics, trained on large-scale synthetic data to generalize across misspecifications and data regimes. The method leverages axial-attention networks to reconcile partial, local causal information into a coherent global structure, with theoretical results supporting its well-specified capacity and skeleton/v-structure recovery capabilities. Empirically, Sea achieves state-of-the-art performance on synthetic and real datasets, scales to graphs with hundreds of nodes, and enables rapid inference with zero or minimal finetuning when adapting to new data assumptions. The work demonstrates robust generalization, adaptability, and efficiency, while outlining practical considerations and avenues for extending causal-discovery foundation models to more diverse assumptions and domains.
Abstract
Causal discovery, the task of inferring causal structure from data, has the potential to uncover mechanistic insights from biological experiments, especially those involving perturbations. However, causal discovery algorithms over larger sets of variables tend to be brittle against misspecification or when data are limited. For example, single-cell transcriptomics measures thousands of genes, but the nature of their relationships is not known, and there may be as few as tens of cells per intervention setting. To mitigate these challenges, we propose a foundation model-inspired approach: a supervised model trained on large-scale, synthetic data to predict causal graphs from summary statistics -- like the outputs of classical causal discovery algorithms run over subsets of variables and other statistical hints like inverse covariance. Our approach is enabled by the observation that typical errors in the outputs of a discovery algorithm remain comparable across datasets. Theoretically, we show that the model architecture is well-specified, in the sense that it can recover a causal graph consistent with graphs over subsets. Empirically, we train the model to be robust to misspecification and distribution shift using diverse datasets. Experiments on biological and synthetic data confirm that this model generalizes well beyond its training set, runs on graphs with hundreds of variables in seconds, and can be easily adapted to different underlying data assumptions.