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

2026-07-13

AI Summary

Causal discovery from observational tabular data faces significant challenges due to heterogeneous causal mechanisms and the combinatorial complexity of high-dimensional DAG search spaces. This work proposes DAG-FM, a novel foundation model that decomposes causal structure learning into two autoregressive stages—leaf node prediction and parent node prediction—and recovers a valid causal ordering through an iterative inference algorithm to construct a fully directed acyclic graph (DAG). To dynamically adapt to unknown families of functional causal models, DAG-FM incorporates a Mixture-of-Leaf-Experts (MoLE) mechanism, complemented by specialized Transformer submodules and tabular interaction blocks to enhance learning efficiency. Extensive experiments on both synthetic and real-world datasets demonstrate that DAG-FM substantially outperforms existing methods in terms of accuracy and scalability.

Abstract

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 \textbf{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.

Problem

• causal discovery
• heterogeneous causal mechanisms
• Directed Acyclic Graphs
• observational data
• Functional Causal Model

Innovation

• DAG-FM
• causal discovery
• Mixture-of-Leaf-Experts
• Transformer-based architecture
• heterogeneous causal mechanisms