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

Causal discovery from observational data is often hindered by the heterogeneity of underlying causal mechanisms and the difficulty of searching through the massive space of possible Directed Acyclic Graphs (DAGs). Traditional methods typically assume a single, homogeneous mechanism (e.g., linear relationships), which limits their real-world applicability. This paper asks: can a foundation model be designed to theoretically guarantee the identification of a unique DAG while remaining robust to diverse, heterogeneous causal mechanisms?

The authors propose DAG-FM, an amortized foundation model that treats causal discovery as a two-stage autoregressive process. Instead of predicting the entire adjacency matrix at once, the model first identifies the causal ordering by recursively predicting leaf nodes, and then prunes the graph by identifying the parents of those nodes. To handle heterogeneous mechanisms, the authors introduce a "Mixture-of-Leaf-Experts" (MoLE) architecture. This module dynamically routes data to specialized experts, each optimized for different functional causal model (FCM) assumptions (such as LiNGAM or Post-Nonlinear models). The model is pre-trained on a synthetic prior space designed to satisfy specific identifiability conditions, ensuring that the model converges to a unique DAG as the sample size increases.

DAG-FM demonstrates state-of-the-art performance across both synthetic benchmarks and complex real-world datasets. By decoupling the discovery process into leaf and parent prediction, the model significantly improves scalability, handling up to 500 variables or 50,000 samples within a 24GB VRAM constraint. The authors provide a theoretical framework showing that by restricting the prior space to specific identifiable mechanism families, the model can correctly recover the true causal structure even when the underlying mechanisms are heterogeneous.