📊 statistics

CDFM: Towards a General-Purpose Causal Discovery Foundation Model

This paper introduces the Causal Discovery Foundation Model (CDFM), a unified, general-purpose framework that leverages a principled variational approach and massive pretraining on synthetic causal models to achieve zero-shot structural inference, thereby overcoming the scalability limitations of traditional dataset-specific causal discovery algorithms.

Original authors: Jie Qiao, Ruichu Cai, Zijian Li, Weilin Chen, Pengfei Hua, Boyan Xu, Zhengming Chen, Zhifeng Hao, Peng Cui

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jie Qiao, Ruichu Cai, Zijian Li, Weilin Chen, Pengfei Hua, Boyan Xu, Zhengming Chen, Zhifeng Hao, Peng Cui

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 caused what? You have a pile of clues (data) showing how different things happened together, but you don't have a witness telling you the story. For decades, detectives had to pick a specific "rulebook" for every single case. One rulebook said, "Assume everything is a straight line!" Another said, "Assume the noise is weird and bumpy!" If you picked the wrong rulebook for the wrong case, your solution was garbage.

This paper introduces a new kind of detective: CDFM (Causal Discovery Foundation Model). Instead of picking a rulebook, CDFM is like a super-smart apprentice who has read every possible rulebook in the universe before ever seeing a real case.

The Old Way vs. The New Way

The Old Way (The "Test-Driven" Struggle):
Imagine you have a messy room (your data). The old method says: "Is the mess caused by a tornado? Let's check if the wind was non-Gaussian! Is it caused by a cat? Let's check if the noise is additive!" You have to run a dozen different tests, guess which one fits, and if you guess wrong, you start over. It's slow, fragmented, and breaks when the room is too messy or weird.

The New Way (CDFM):
CDFM skips the guessing game. It looks at the messy room and says, "I've seen millions of rooms like this in my training." It doesn't need you to tell it if the data is linear or non-linear. It just looks at the patterns and draws the map of who caused what. It's a "zero-shot" detective, meaning it can solve a brand-new type of mystery without needing to be retrained first.

How Did It Learn So Much? (The "Synthetic Playground")

You might wonder, "How can a computer know everything about cause and effect?" The authors didn't feed it real-world data (which is messy and often has missing answers). Instead, they built a massive, infinite video game.

In this game, they generated 15 different families of fake worlds. Some worlds had straight-line rules, some had wild, curved rules, some had hidden "ghost" variables pulling strings behind the scenes, and some had missing pieces of data. They created 10 different sizes of these worlds, ranging from tiny puzzles with 10 variables to massive ones with 100 variables. They also varied the number of clues, from 500 up to 4,000 observations.

CDFM played in this video game for a long time, learning to recognize the "footprints" of different causes. It learned that if a pattern looks like a specific type of noise, it's probably one kind of cause; if it looks like another, it's a different kind.

The Magic Trick: "Variational Inference"

Here is the clever part. The paper argues that you can't just guess the answer; you have to mathematically prove that your guess is the only one that fits. The authors used a math trick called variational inference.

Think of it like this: CDFM doesn't just say, "I think A caused B." It says, "I'm going to imagine every possible way A could have caused B, weigh them all, and pick the one that makes the most sense of the data." It treats the "rules" of the world as hidden variables that it tries to figure out while it figures out the structure. This ensures that if the math says a direction is impossible, CDFM won't guess it.

Did It Actually Work? (The Proof)

The authors didn't just say, "It looks cool." They put it to the test in three ways:

The Synthetic Gauntlet: They threw CDFM at the 15 different types of fake worlds it hadn't seen during the final testing phase.
- The Result: CDFM won. In a test with 4,000 samples, it scored an AUROC of 0.888 (a measure of accuracy where 1.0 is perfect). The next best method, TabCausal, only got 0.786. Even when the puzzle got huge (100 variables), CDFM stayed strong with an AUROC above 0.87, while others crashed.
- The Takeaway: It suggests that CDFM can handle a wide variety of messy, real-world-like situations without needing a specific rulebook.
The Real World Test: They tried it on real data, like a physical system called Causal Chamber (where light tunnels and sensors create real physics) and the Tübingen benchmark (real-world pairs like "height causes weight").
- The Result: On the Causal Chamber Task A1, CDFM got an AUROC of 0.952, beating the next best (TabCausal at 0.930). On the Tübingen benchmark, it got 67.1% accuracy, beating the previous best of 63.8%.
- The Takeaway: The skills learned in the video game actually transferred to real life.
The "Honesty" Check: The authors wanted to make sure CDFM wasn't just guessing blindly. They tested it on situations where, mathematically, no one can know the answer (like when two things are perfectly symmetrical).
- The Result: When the math says "it's impossible to tell," CDFM's accuracy dropped to 0.5 (which is just guessing like a coin flip). When the math says "it is possible," its accuracy shot up to 1.0.
- The Takeaway: This suggests CDFM respects the laws of logic. It knows when it shouldn't know the answer.

What It Can't Do (The Rules)

The paper is very clear about what CDFM is not.

It is not a magic wand that removes the need for assumptions. The paper explicitly states that you must have some causal assumptions; CDFM just learns a broad spectrum of them instead of one single one.
It is not perfect at everything. In very specific, simple cases (like a perfectly linear world with non-Gaussian noise), older, specialized tools like DirectLiNGAM can still perform slightly better. CDFM is a "generalist," not a "specialist" for one tiny niche.
The results on real-world data are measured on specific benchmarks (Causal Chamber, Tübingen). The paper does not claim it has solved causal discovery for every problem in the universe, only that it works remarkably well across a huge range of them.

The Missing Piece Bonus

As a side effect of its training, CDFM also got really good at filling in missing puzzle pieces. If you give it a dataset with holes, it can guess what the missing numbers were. In tests, it beat standard methods like Mean or KNN imputation, achieving a Mean Absolute Error (MAE) of 0.452 on continuous data and 0.717 on discrete data.

The Bottom Line

This paper suggests a big shift: instead of building a new, specialized tool for every new type of data, we might be able to build one giant, pre-trained brain that understands the language of cause and effect. It's not a solved problem for every mystery in the world, but it's a massive step toward a detective that can walk into any room, look at the clues, and say, "I've seen this before. Here is the map."

Technical Summary: CDFM: Towards a General-Purpose Causal Discovery Foundation Model

Problem Statement

Causal discovery, the process of recovering underlying causal structures from observational data, has traditionally relied on a fragmented, test-driven paradigm. Existing algorithms are meticulously designed for specific causal mechanisms (e.g., linear non-Gaussianity, additive noise, or specific independence constraints). Consequently, these methods excel only when their underlying assumptions strictly hold. In real-world scenarios characterized by heterogeneous data and unknown generating processes, practitioners face an exhaustive loop of manually selecting algorithms, verifying assumptions, and re-evaluating statistical tests. This dataset-specific approach fails to scale to the demands of modern scientific discovery and struggles to generalize across diverse, unknown causal mechanisms.

Methodology

The authors propose the Causal Discovery Foundation Model (CDFM), a unified framework designed for zero-shot structural inference across heterogeneous and initially unknown causal mechanisms. The methodology is grounded in a principled variational framework that bridges causal identifiability theory with deep learning architecture design.

1. Theoretical Foundation: Variational Decomposition

The core theoretical insight is that causal discovery under unknown mechanisms is a Bayesian model selection problem. Since the true causal mechanism $M^*$ is unknown, the model must marginalize over a vast hypothesis space of mechanisms $\mathcal{M}$ .

Identifiability Boundaries: The paper establishes that without constraints, causal direction is mathematically impossible to identify (Theorem 1). However, if the space of mechanisms $\mathcal{M}$ satisfies a "strict space identifiability" property, the exact marginal likelihood uniquely identifies the true graph (Theorem 3).
Variational Inference: Directly computing the marginal likelihood is intractable. CDFM formulates the problem using variational inference, introducing a parameterized distribution $Q(M|X)$ $Q (M ∣ X)$ to approximate the posterior over mechanisms. This yields an Evidence Lower Bound (ELBO) that decomposes the learning objective into three tractable components:
1. Mechanism Inference: $Q(M|X)$ , extracting statistical footprints (noise properties, non-linearities) to infer latent mechanism encodings.
2. Data Reconstruction: $\log P(X|G, M)$ , ensuring the inferred mechanism and graph explain the observed data (self-supervised signal).
3. Graph Inference: $\log P(G|M)$ , decoding the structural graph conditioned on the encoded mechanism.

2. Architecture Design

CDFM's architecture explicitly mirrors the variational decomposition:

Column-Wise Embedding & Alternating Attention: Input data (handling missing values via a binary indicator channel) is embedded using a Set Transformer to capture marginal distributions. A stack of alternating row- and column-wise attention blocks integrates joint dependencies to parameterize the latent mechanism distribution $Q(M|X)$ .
Data Reconstruction (Imputation): A residual block predicts conditional quantile values for masked entries. This serves as a self-supervised task to satisfy the data reconstruction term. A "soft causal mask" is employed during attention, penalizing information flow from non-causal variables based on current edge weight hypotheses, creating an unsupervised inner loop for structure refinement.
Graph Inference: Induced self-attention extracts "cause" and "effect" summary vectors for each variable. A Mixture-of-Experts (MoE) bilinear head decodes these vectors into edge logits, naturally accommodating directed edges and bidirected edges (for latent confounding). A lightweight graph-adaptive threshold calibrator converts continuous probabilities into a binary graph.

3. Pretraining Strategy

CDFM is pretrained on a massive, highly diverse space of synthetic Structural Causal Models (SCMs). The data generator spans:

Graph Structures: Six random graph models (ER, Scale-Free, WS, SBM, etc.) with varying sparsity and sizes ( $D \in [2, 100]$ ).
Mechanism Diversity: Linear and nonlinear models (Additive Noise, Post-Nonlinear, Random Fourier Features), discrete and mixed-type mechanisms, and physics-inspired models.
Noise & Confounding: Diverse noise distributions (Gaussian, Laplace, Cauchy, etc.), heteroscedasticity, and latent confounding (15% of batches).
Realism: Simulated measurement errors, discretization, and random missing values (MCAR).

Key Contributions

Paradigm Shift: Formulates causal discovery as a general-purpose foundation model task, shifting from specialized algorithm selection to a unified inference approach capable of zero-shot structural inference.
Principled Framework: Unifies foundation models and structural identifiability through a variational framework, treating unknown mechanisms as latent variables to translate the marginalization problem into a tractable learning objective.
Architecture Design: Designs CDFM to align mechanism encoding, structure-aware reconstruction, and graph decoding with the theoretical decomposition, enabling robust inference directly from observational data.
Large-Scale Pretraining: Introduces a comprehensive pretraining regime driven by diverse causal knowledge, synthesizing a massive hypothesis space to internalize complex statistical asymmetries.

Experimental Results

The authors evaluate CDFM across synthetic benchmarks, real-world datasets, and controlled identifiability studies.

Synthetic Benchmarks: CDFM consistently outperforms traditional algorithms (PC, GES, DirectLiNGAM, NOTEARS) and amortized baselines (AVICI, TabCausal, Arrow) across 15 distinct mechanism families and varying graph sizes ( $D \in [10, 100]$ ) and sample sizes ( $N \in [500, 4000]$ ). It achieves superior AUROC and F1 scores, demonstrating stable performance as problem scale increases.
Real-World Benchmarks:
- Causal Chamber: On physical multivariate systems (Tasks A1 and A2), CDFM achieves high AUROC (0.952 and 0.965), outperforming all baselines and demonstrating transferability from synthetic to physical systems.
- Tübingen Benchmark: On 95 bivariate cause-effect pairs, CDFM achieves 67.1% accuracy, surpassing other state-of-the-art methods.
Identifiability Validation: Controlled experiments confirm that CDFM adheres to theoretical identifiability boundaries. It correctly identifies causal directions when functional asymmetry (non-linearity/non-Gaussianity) or V-structures exist, and reverts to chance-level performance (accuracy $\approx$ 0.5) in non-identifiable regimes (e.g., linear Gaussian models).
Missing Value Imputation: As a byproduct of the reconstruction objective, CDFM's imputation head achieves state-of-the-art performance on both continuous and discrete data with missing values, outperforming classical methods like MissForest and MICE.

Significance and Claims

The paper claims that CDFM represents a significant step toward a general-purpose causal discovery engine. By shifting the bottleneck from manual algorithm selection to the scalable generation of diverse, identifiable synthetic data, CDFM internalizes complex causal asymmetries that allow it to generalize across unknown generating processes.

The authors emphasize that the model does not remove the need for causal assumptions but rather learns to adopt them over a broad spectrum of assumption classes. The work suggests a future direction where causal models learn broad representations over families of mechanisms rather than optimizing for single, rigid assumptions. The results demonstrate that such a foundation model can drive a paradigm shift, offering robust, zero-shot causal discovery capabilities that scale with data volume and heterogeneity.