# Learning Causal Graphs at Scale: A Foundation Model Approach

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 42 references that are in the map.
- First posted on arXiv 23 Jun 2025.

## Builds on (most cited first, 28 of 28)

- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [Learning Sparse Nonparametric DAGs](https://mlsanity.com/paper/1909.13189/) (2019)
- [On the Role of Sparsity and DAG Constraints for Learning Linear DAGs](https://mlsanity.com/paper/2006.10201/) (2020)
- [DAGMA: Learning DAGs via M-matrices and a Log-Determinant Acyclicity Characterization](https://mlsanity.com/paper/2209.08037/) (2022)
- [Causal Autoregressive Flows](https://mlsanity.com/paper/2011.02268/) (2020)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)
- [Integrating Large Language Model for Improved Causal Discovery](https://mlsanity.com/paper/2306.16902/) (2023)
- [DAGs with No Curl: An Efficient DAG Structure Learning Approach](https://mlsanity.com/paper/2106.07197/) (2021)
- [ViTO: Vision Transformer-Operator](https://mlsanity.com/paper/2303.08891/) (2023)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [Causality for Large Language Models](https://mlsanity.com/paper/2410.15319/) (2024)
- [Truncated Matrix Power Iteration for Differentiable DAG Learning](https://mlsanity.com/paper/2208.14571/) (2022)
- [Transformer Meets Boundary Value Inverse Problems](https://mlsanity.com/paper/2209.14977/) (2022)
- [Supervised Whole DAG Causal Discovery](https://mlsanity.com/paper/2006.04697/) (2020)
- [Learning Regularization Parameters of Inverse Problems via Deep Neural Networks](https://mlsanity.com/paper/2104.06594/) (2021)
- [Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery](https://mlsanity.com/paper/2408.07307/) (2024)
- [Causal Discovery with Heterogeneous Observational Data](https://mlsanity.com/paper/2201.12392/) (2022)
- [Multi-task Learning of Order-Consistent Causal Graphs](https://mlsanity.com/paper/2111.02545/) (2021)
- [Effective Causal Discovery under Identifiable Heteroscedastic Noise Model](https://mlsanity.com/paper/2312.12844/) (2023)
- [Deformable Cross-Attention Transformer for Medical Image Registration](https://mlsanity.com/paper/2303.06179/) (2023)
- [To be or not to be stable, that is the question: understanding neural networks for inverse problems](https://mlsanity.com/paper/2211.13692/) (2022)
- [Ambiguity in solving imaging inverse problems with deep learning based operators](https://mlsanity.com/paper/2305.19774/) (2023)
- [Transformer learns the cross-task prior and regularization for in-context learning](https://mlsanity.com/paper/2505.12138/) (2025)
- [Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery](https://mlsanity.com/paper/2505.23106/) (2025)

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

- [Arrow: A Foundation Model for Causal Discovery](https://mlsanity.com/paper/2605.07204/) (2026)
- [CDFM: Towards a General-Purpose Causal Discovery Foundation Model](https://mlsanity.com/paper/2607.11508/) (2026)
- [Causal Foundation Models](https://mlsanity.com/paper/2609.03003/) (2026)
- [Demistifying Data and Simulator Assumptions in Supervised Causal Discovery](https://mlsanity.com/paper/2609.37446/) (2026)

## Abstract (arXiv)

Due to its human-interpretability and invariance properties, Directed Acyclic Graph (DAG) has been a foundational tool across various areas of AI research, leading to significant advancements. However, DAG learning remains highly challenging, due to its super-exponential growth in computational cost and identifiability issues, particularly in small-sample regimes. To address these two challenges, in this work we leverage the recent success of linear transformers and develop a foundation model approach for discovering multiple order-consistent DAGs across tasks. In particular, we propose Attention-DAG (ADAG), a novel attention-mechanism-based architecture for learning multiple linear Structural Equation Models (SEMs). ADAG learns the mapping from observed data to both graph structure and parameters via a nonlinear attention-based kernel, enabling efficient multi-task estimation of the underlying linear SEMs. By formulating the learning process across multiple tasks as a continuous optimization problem, the pre-trained ADAG model captures the common structural properties as a shared low-dimensional prior, thereby reducing the ill-posedness of downstream DAG learning tasks in small-sample regimes. We evaluate our proposed approach on benchmark synthetic datasets and find that ADAG achieves substantial improvements in both DAG learning accuracy and zero-shot inference efficiency. To the best of our knowledge, this is the first practical approach for pre-training a foundation model specifically designed for DAG learning, representing a step toward more efficient and generalizable down-stream applications in causal discovery.