📊 statistics

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

This paper introduces DAG-FM, a novel foundation model that leverages a two-stage auto-regressive Transformer architecture with a Mixture-of-Leaf-Experts mechanism to achieve state-of-the-art causal discovery performance on heterogeneous and high-dimensional tabular data.

Original authors: Yikang Chen, Zhengkang Guan, Haoyuan Qian, Peng Cui, Yi Yang, Kun Kuang

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yikang Chen, Zhengkang Guan, Haoyuan Qian, Peng Cui, Yi Yang, Kun Kuang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery: who is causing what? You have a giant spreadsheet full of clues (data), but the rules of the game are tricky. Sometimes, the clues follow strict, predictable laws; other times, the laws change from one clue to the next. This is the world of causal discovery, and for a long time, the tools detectives used were like trying to solve a Rubik's Cube while wearing blindfolded gloves. They either got stuck on simple puzzles or broke when the rules got messy.

Enter DAG-FM, a new kind of "super-detective" foundation model. Think of it not as a single, rigid rulebook, but as a master chef with a massive, magical kitchen.

The Old Way vs. The New Chef

Previously, detectives tried to guess the entire map of connections (who causes whom) all at once. It was like trying to draw a whole city map in one go. If the city had weird, mixed-up traffic rules (heterogeneous mechanisms), the map usually ended up a mess. Some old methods were like specialists who only knew how to drive on straight, flat highways (Linear Non-Gaussian models). If the road turned into a bumpy dirt track or a winding mountain pass, those specialists got lost.

DAG-FM changes the game. Instead of drawing the whole map at once, it plays a game of "reverse scavenger hunt." It works in two steps:

Find the Leaves: It looks at the data and asks, "Who is at the very end of the chain? Who isn't causing anyone else?" It peels off the last person in the line.
Find the Parents: Once it knows who the "leaf" is, it asks, "Who was pulling their strings?" It finds the parents of that leaf.

It repeats this process, peeling away layers of the onion until the whole structure is revealed. This is much smarter than guessing the whole picture at once.

The Magic Ingredient: The "Mixture of Experts"

Here is where DAG-FM gets really cool. In the real world, the rules of cause-and-effect aren't the same everywhere. Sometimes a cause acts like a simple math equation; other times, it's a complex, squiggly function.

DAG-FM uses a team of specialists inside its brain, called Mixture-of-Leaf-Experts (MoLE). Imagine a room full of different detectives: one is great at solving linear puzzles, another at spotting noise patterns, another at handling weird curves. When DAG-FM looks at a piece of data, it doesn't just guess; it asks, "Which detective is best for this specific clue?" It routes the problem to the right expert dynamically. This allows it to handle a chaotic mix of rules that would confuse older, single-minded detectives.

The Training Ground: A Synthetic Playground

Before DAG-FM could tackle real-world mysteries, the authors had to teach it. They didn't just feed it real data; they built a massive, virtual playground. They generated millions of fake scenarios using random graphs and random noise, mixing and matching different types of causal rules (like Linear, Additive Noise, and Post-Nonlinear).

Crucially, they designed this playground with a specific safety net. They proved mathematically that if the model is trained on this specific mix of rules, it will be able to find the one true answer, even if the data is messy. It's like training a pilot in a simulator that includes every possible storm and engine failure, so when they fly a real plane, they know exactly what to do.

The Results: Speed and Smarts

When the authors tested DAG-FM, the results were impressive.

On Fake Data: In simulations with 100 different types of mixed-up rules, DAG-FM beat almost every other method. It got the connections right more often (higher precision and recall) and made fewer mistakes (lower error rates) than the old-school algorithms and even other new AI models.
On Real Data: They tested it on real-world datasets, like a protein signaling network with 7,466 samples and 11 variables, and a physical system with 10,000 samples and 38 variables. While many other methods crashed (ran out of memory) or gave up, DAG-FM kept going. It found the best maps in these tests, outperforming traditional tools that struggled with the size or complexity of the data.
Efficiency: It's also incredibly efficient. It can handle datasets with up to 50,000 samples or 500 variables on a standard computer with 24GB of video memory, without needing fancy tricks to slow it down.

What It Doesn't Do (Yet)

It's important to know what this super-detective can't do. The paper explicitly states that DAG-FM assumes it has all the clues (no hidden confounders). If there are secret variables influencing the data that the model can't see, the method might get confused. The authors also note that while the model works great on the rules it was trained on, it is still being tested on completely new, unseen types of chaos.

The Bottom Line

DAG-FM isn't a magic wand that solves every mystery instantly, but it's a massive leap forward. By breaking the problem down into smaller, manageable steps and using a team of specialized experts to handle different types of rules, it manages to find the true cause-and-effect map in situations where previous tools failed. It suggests that with the right training and architecture, AI can finally navigate the messy, mixed-up reality of how the world actually works.

Technical Summary: DAG-FM

Problem Statement

Causal discovery from observational tabular data faces two fundamental challenges: the heterogeneity of underlying causal mechanisms and the high-dimensional combinatorial search space of Directed Acyclic Graphs (DAGs). Traditional methods, including constraint-based, score-based, and optimization-based approaches, often struggle without strong assumptions. Specifically, Functional Causal Model (FCM) based methods typically assume homogeneous causal mechanisms (e.g., all relationships are linear or all are additive noise), limiting their generalizability to real-world scenarios where mechanisms vary across variables. Conversely, emerging amortized (foundation) models often lack theoretical guarantees for outputting strictly identifiable DAGs and frequently fail to handle heterogeneous mechanisms, as their performance is highly sensitive to the design of the prior space used for pre-training.

Methodology

The paper proposes DAG-FM, a foundation model architecture designed to amortize causal discovery while intrinsically guaranteeing the output of a unique, identifiable DAG from observational data.

1. Theoretical Foundation: Heterogeneous Identifiability

The authors establish a design condition for the prior space that allows for heterogeneous causal mechanisms. Unlike traditional FCM approaches that enforce a single mechanism type, DAG-FM's prior space permits a Structural Causal Model (SCM) to be a composite mixture of four distinct mechanism families:

LiNGAM: Linear relationships with non-Gaussian noise.
ANM: Non-linear relationships with additive noise.
HNM: Heteroscedastic noise models.
PNL: Post-nonlinear models.

Theoretical results (Theorems 4.4 and 4.5) demonstrate that under these conditions (assuming causal sufficiency and faithfulness), the true DAG is almost surely identifiable from the observational distribution. This ensures that the amortized model converges to a unique ground-truth DAG as sample size increases.

2. Architecture: Decomposed Prediction

To address the scalability issues of directly predicting adjacency matrices (which scale quadratically and lack DAG constraints), DAG-FM decomposes the discovery process into two sequential, auto-regressive stages using specialized Transformer-based sub-modules:

Leaf-Node Predictor ( $Q_{\theta_1}$ ): Identifies the causal ordering by recursively predicting leaf nodes (nodes with no children) in the current sub-graph.
Parent-Node Predictor ( $Q_{\theta_2}$ ): Given a confirmed leaf node, this module identifies its parents, effectively pruning the graph to an upper-triangular structure based on the established ordering.

3. Mixture-of-Leaf-Experts (MoLE)

To handle the diverse mechanism families within the prior space, the leaf prediction head integrates a Mixture-of-Leaf-Experts (MoLE) framework.

The model employs $m+1$ bilinear heads, each corresponding to a specific mechanism family (plus a general mechanism expert and a dummy expert).
A top-k router dynamically selects the appropriate expert based on the input data features.
This design allows the model to adapt to identifiable mechanism families dynamically, bridging the gap between theoretical FCM assumptions and practical data heterogeneity.

4. Prior Space Construction

The model is pre-trained on a large-scale synthetic dataset generated from a prior space $\Phi$ that includes:

Random Graphs: Erdős–Rényi, Barabási–Albert, and Watts–Strogatz models with varying densities.
Random Mechanisms: A diverse library of functions (MLPs, GPs, Random Forests, etc.) sampled from the four defined families.
Random Noise: A comprehensive library of distributions (Normal, Laplace, Student's t, etc.) with controlled Signal-to-Noise Ratios (SNR).
Data Augmentation: Non-linear warping and discretization to simulate real-world data characteristics, while enforcing strict distribution normalization to prevent spurious statistical artifacts.

5. Inference

Inference is performed iteratively (Algorithm 1). The model repeatedly identifies a leaf node, determines its parents, removes the node from the set, and updates the causal ordering until all variables are processed. This process guarantees a valid DAG structure without post-hoc cycle correction.

Key Contributions

Theoretical Guarantee: The paper provides a design condition for the prior space that supports heterogeneous causal mechanisms while theoretically ensuring the amortized model converges to a unique, identifiable DAG.
DAG-FM Architecture: The development of a foundation model that decouples DAG identification into leaf-node and parent-node prediction tasks, utilizing Transformer-based sub-modules and the MoLE framework to handle mechanism heterogeneity.
Empirical Performance: Extensive experiments demonstrating that DAG-FM achieves state-of-the-art performance on both synthetic benchmarks with heterogeneous mechanisms and complex real-world datasets, significantly outperforming traditional algorithms and recent foundation models in accuracy and scalability.

Experimental Results

Synthetic Benchmarks: On datasets with $n=1000$ samples and $d=20$ variables featuring heterogeneous mechanisms, DAG-FM achieved the highest F1-score (0.64) and lowest normalized Structural Hamming Distance (nSHD) (0.19), outperforming baselines like Direct-LiNGAM, AVICI, and TabCausal.
Homogeneous Subspaces: Even in strictly homogeneous settings (e.g., pure LiNGAM), DAG-FM remained competitive, achieving an Order Divergence Rate (ODR) of 0.12, comparable to specialized algorithms.
Real-World Benchmarks:
- Sachs Dataset (Protein Signaling): DAG-FM achieved an F1 of 0.53 and nSHD of 0.25, significantly outperforming all other methods (e.g., FoundCause: F1 0.33; GES: F1 0.31).
- Causal Chamber: DAG-FM achieved an F1 of 0.53 and nSHD of 0.14, surpassing all baselines.
Scalability: DAG-FM demonstrated robust scalability, handling up to 50,000 samples or 500 variables within a 24GB VRAM constraint, whereas many baselines (e.g., AVICI, CauScale) encountered out-of-memory errors on large datasets.
Robustness: The model showed resilience to varying graph densities, data warping, and discretization, and maintained strong performance in Out-of-Distribution (OOD) scenarios involving unseen graph structures, mechanisms, and noise distributions.

Significance

The paper claims that DAG-FM represents a significant advancement by integrating the theoretical guarantees of FCM-based identifiability with the scalability of amortized foundation models. By moving away from direct adjacency matrix prediction and instead decomposing the task into sequential, theoretically grounded sub-tasks, DAG-FM overcomes the limitations of both traditional methods (lack of generalizability to heterogeneous mechanisms) and existing amortized approaches (lack of structural guarantees and sensitivity to prior design). The work establishes a new paradigm for causal discovery that is both theoretically rigorous and practically applicable to high-dimensional, real-world data.