# Sample, estimate, aggregate: A recipe for causal discovery foundation models

arXiv: https://arxiv.org/abs/2402.01929 · HTML: https://mlsanity.com/paper/2402.01929/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 95 references that are in the map.
- First posted on arXiv 2 Feb 2024.

## Builds on (most cited first, 30 of 45)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Prompt Programming for Large Language Models: Beyond the Few-Shot Paradigm](https://mlsanity.com/paper/2102.07350/) (2021)
- [Are Emergent Abilities of Large Language Models a Mirage?](https://mlsanity.com/paper/2304.15004/) (2023)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [Neural Autoregressive Flows](https://mlsanity.com/paper/1804.00779/) (2018)
- [CAM: Causal additive models, high-dimensional order search and penalized regression](https://mlsanity.com/paper/1310.1533/) (2013)
- [Causal Inference in the Presence of Latent Variables and Selection Bias](https://mlsanity.com/paper/1302.4983/) (2013)
- [Causal Discovery from Heterogeneous/Nonstationary Data with Independent Changes](https://mlsanity.com/paper/1903.01672/) (2019)
- [Beware of the Simulated DAG! Causal Discovery Benchmarks May Be Easy To Game](https://mlsanity.com/paper/2102.13647/) (2021)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [On the Turing Completeness of Modern Neural Network Architectures](https://mlsanity.com/paper/1901.03429/) (2019)
- [Causal-learn: Causal Discovery in Python](https://mlsanity.com/paper/2307.16405/) (2023)
- [Optimal Approximation Rate of ReLU Networks in terms of Width and Depth](https://mlsanity.com/paper/2103.00502/) (2021)
- [Deep End-to-end Causal Inference](https://mlsanity.com/paper/2202.02195/) (2022)
- [Amortized Inference for Causal Structure Learning](https://mlsanity.com/paper/2205.12934/) (2022)
- [High-dimensional learning of linear causal networks via inverse covariance estimation](https://mlsanity.com/paper/1311.3492/) (2013)
- [Efficient Neural Causal Discovery without Acyclicity Constraints](https://mlsanity.com/paper/2107.10483/) (2021)
- [Permutation-based Causal Inference Algorithms with Interventions](https://mlsanity.com/paper/1705.10220/) (2017)

## Built on by (most cited first, 21 of 21)

- [Black Box Causal Inference: Effect Estimation via Meta Prediction](https://mlsanity.com/paper/2503.05985/) (2025)
- [Large-Scale Targeted Cause Discovery via Learning from Simulated Data](https://mlsanity.com/paper/2408.16218/) (2024)
- [CauScale: Neural Causal Discovery at Scale](https://mlsanity.com/paper/2602.08629/) (2026)
- [Arrow: A Foundation Model for Causal Discovery](https://mlsanity.com/paper/2605.07204/) (2026)
- [Amortized Inference of Causal Models via Conditional Fixed-Point Iterations](https://mlsanity.com/paper/2410.06128/) (2024)
- [CausalMan: A physics-based simulator for large-scale causality](https://mlsanity.com/paper/2502.12707/) (2025)
- [TabCausal: Pretraining Across Causal Environments for Tabular Causal Discovery](https://mlsanity.com/paper/2605.31156/) (2026)
- [Meta-Statistical Learning: Supervised Learning of Statistical Estimators](https://mlsanity.com/paper/2502.12088/) (2025)
- [FoundCause: Causal Discovery with Latent Confounders from Observational Data](https://mlsanity.com/paper/2606.17516/) (2026)
- [Test-Time Learning of Causal Structure from Interventional Data](https://mlsanity.com/paper/2602.19131/) (2026)
- [CDFM: Towards a General-Purpose Causal Discovery Foundation Model](https://mlsanity.com/paper/2607.11508/) (2026)
- [Identifying biological perturbation targets through causal differential networks](https://mlsanity.com/paper/2410.03380/) (2024)
- [Learning to refine domain knowledge for biological network inference](https://mlsanity.com/paper/2410.14436/) (2024)
- [Large Causal Models for Temporal Causal Discovery](https://mlsanity.com/paper/2602.18662/) (2026)
- [CauTion: Knowing When to Trust LLMs for Ensemble Causal Discovery](https://mlsanity.com/paper/2606.03602/) (2026)
- [Causal Foundation Models](https://mlsanity.com/paper/2609.03003/) (2026)
- [Scalable Contrastive Causal Discovery under Unknown Soft Interventions](https://mlsanity.com/paper/2603.03411/) (2026)
- [TCD-Arena: Assessing Robustness of Time Series Causal Discovery Methods Against Assumption Violations](https://mlsanity.com/paper/2605.03045/) (2026)
- [CausalSteward: An Agentic Divide-Conquer-Combine Copilot for Causal Discovery](https://mlsanity.com/paper/2607.01936/) (2026)
- [CausalArena: Benchmarking Causal Discovery in the Foundation Model Era](https://mlsanity.com/paper/2609.11897/) (2026)
- [Demistifying Data and Simulator Assumptions in Supervised Causal Discovery](https://mlsanity.com/paper/2609.37446/) (2026)

## Abstract (arXiv)

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.