# DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 57 references that are in the map.
- First posted on arXiv 13 Jul 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [Causal Discovery with Continuous Additive Noise Models](https://mlsanity.com/paper/1309.6779/) (2013)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [On the Role of Sparsity and DAG Constraints for Learning Linear DAGs](https://mlsanity.com/paper/2006.10201/) (2020)
- [CAM: Causal additive models, high-dimensional order search and penalized regression](https://mlsanity.com/paper/1310.1533/) (2013)
- [Beware of the Simulated DAG! Causal Discovery Benchmarks May Be Easy To Game](https://mlsanity.com/paper/2102.13647/) (2021)
- [TabICLv2: A better, faster, scalable, and open tabular foundation model](https://mlsanity.com/paper/2602.11139/) (2026)
- [DAGMA: Learning DAGs via M-matrices and a Log-Determinant Acyclicity Characterization](https://mlsanity.com/paper/2209.08037/) (2022)
- [Score matching enables causal discovery of nonlinear additive noise models](https://mlsanity.com/paper/2203.04413/) (2022)
- [Amortized Inference for Causal Structure Learning](https://mlsanity.com/paper/2205.12934/) (2022)
- [LimiX: Unleashing Structured-Data Modeling Capability for Generalist Intelligence](https://mlsanity.com/paper/2509.03505/) (2025)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [On the Identifiability and Estimation of Causal Location-Scale Noise Models](https://mlsanity.com/paper/2210.09054/) (2022)
- [Scalable Causal Discovery with Score Matching](https://mlsanity.com/paper/2304.03382/) (2023)
- [Causal Discovery with Score Matching on Additive Models with Arbitrary Noise](https://mlsanity.com/paper/2304.03265/) (2023)
- [Identifying Patient-Specific Root Causes with the Heteroscedastic Noise Model](https://mlsanity.com/paper/2205.13085/) (2022)
- [Directed Cyclic Graphical Representations of Feedback Models](https://mlsanity.com/paper/1302.4982/) (2013)
- [A Scale-Invariant Sorting Criterion to Find a Causal Order in Additive Noise Models](https://mlsanity.com/paper/2303.18211/) (2023)
- [A Meta-Learning Approach to Bayesian Causal Discovery](https://mlsanity.com/paper/2412.16577/) (2024)
- [Standardizing Structural Causal Models](https://mlsanity.com/paper/2406.11601/) (2024)
- [Ordering-Based Causal Discovery for Linear and Nonlinear Relations](https://mlsanity.com/paper/2410.05890/) (2024)
- [Effective Causal Discovery under Identifiable Heteroscedastic Noise Model](https://mlsanity.com/paper/2312.12844/) (2023)
- [Heteroscedastic Causal Structure Learning](https://mlsanity.com/paper/2307.07973/) (2023)
- [A Skewness-Based Criterion for Addressing Heteroscedastic Noise in Causal Discovery](https://mlsanity.com/paper/2410.06407/) (2024)
- [Large-Scale Targeted Cause Discovery via Learning from Simulated Data](https://mlsanity.com/paper/2408.16218/) (2024)
- [Distinguishing Cause from Effect with Causal Velocity Models](https://mlsanity.com/paper/2502.05122/) (2025)
- [CauScale: Neural Causal Discovery at Scale](https://mlsanity.com/paper/2602.08629/) (2026)

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

- [Causal Foundation Models](https://mlsanity.com/paper/2609.03003/) (2026)
- [SkillTFM: Gated Skill Evolution for Training-Free Adaptation of Tabular Foundation Models](https://mlsanity.com/paper/2608.06137/) (2026)

## Abstract (arXiv)

Causal discovery from observational tabular data remains fundamentally challenging, primarily due to the heterogeneity of underlying causal mechanisms and the high-dimensional combinatorial search space of Directed Acyclic Graphs (DAGs). In this paper, we propose DAG-FM, a novel foundation model architecture that amortizes causal discovery. Unlike direct matrix prediction, DAG-FM decomposes the causal discovery process into two auto-regressive stages using two specialized Transformer-based sub-modules: a leaf-node predictor and a parent-node predictor. To effectively model complex row-column interactions, we adopt a robust tabular interaction block to output feature-wise representations. Crucially, to handle diverse and unknown Functional Causal Model (FCM) assumptions in real-world scenarios, we introduce Mixture-of-Leaf-Experts (MoLE), allowing the model to dynamically route and adapt to identifiable mechanism families. Through an iterative inference algorithm, DAG-FM seamlessly extracts causal orderings and constructs valid DAGs. Extensive experiments demonstrate that DAG-FM achieves state-of-the-art performance on both synthetic benchmarks and complex real-world datasets, significantly outperforming traditional classical algorithms and recent foundation models in both accuracy and scalability.