🤖 machine learning

A Causal Foundation Model for Structure and Outcome Prediction

The paper introduces TabPFN-CFM, a causal foundation model trained on synthetic data that generalizes to real-world datasets to simultaneously predict causal structures and outcomes while supporting queries across all levels of Pearl's Causal Hierarchy.

Original authors: Max Zhu, Martino Mansoldo, Ching-Hao Wang, Stefan Groha

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Max Zhu, Martino Mansoldo, Ching-Hao Wang, Stefan Groha

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, but you only have a pile of old, unsorted photos (observational data) and no witness testimony. You need to figure out two things:

The Story: Who caused what? (Did the rain cause the wet grass, or did someone spill a bucket?)
The "What If": If you had done something different, what would have happened? (If you hadn't spilled the bucket, would the grass still be wet?)

Usually, solving this requires a human expert to guess the rules of the world, or a computer to try millions of guesses one by one. This paper introduces TabPFN-CFM, a new kind of "super-detective" AI that has already read a library of millions of made-up mystery stories. Because it has seen so many different scenarios, it can look at your messy photos and instantly guess the story and answer your "what if" questions.

Here is how it works, broken down into simple concepts:

1. The "Causal Foundation Model" (The Super-Detective)

Think of this model as a student who has spent their entire life studying a massive library of synthetic (fake) worlds. In these fake worlds, the authors created thousands of different rules: sometimes gravity works differently, sometimes people act randomly, and sometimes hidden factors (like a secret agent) mess things up.

Because this student has seen so many variations, when you show them a real dataset (like sales numbers or school grades), they don't need to learn the rules from scratch. They just say, "Ah, this looks like Scenario #4,592 from my library," and they instantly know the likely cause-and-effect relationships.

2. Solving Three Levels of Mystery

The paper claims this model can handle three different types of questions, which the authors call the "Causal Hierarchy":

Level 1: The Observer (What is happening?)
- The Question: "Looking at the data, if someone buys a red shirt, do they also buy a hat?"
- The Model's Job: It predicts the outcome based on what it sees.
Level 2: The Intervener (What if I change something?)
- The Question: "If I force everyone to buy a red shirt (even if they wouldn't have), will they buy a hat?"
- The Model's Job: This is tricky because the real world didn't do this. The model simulates the change to predict the result.
Level 3: The Time Traveler (What if things were different in the past?)
- The Question: "John bought a hat and a red shirt. But what if he hadn't bought the red shirt? Would he still have bought the hat?"
- The Model's Job: This is the hardest. It requires rewinding time in the simulation while keeping other facts fixed. The paper says this model is one of the few that can do this well.

3. The "Hidden Villain" (Unobserved Confounders)

In many real-world mysteries, there is a "hidden villain" you can't see. For example, maybe a hidden factor (like "wealth") causes both buying red shirts and buying hats. If you don't know about "wealth," you might wrongly think the shirt causes the hat.

The paper introduces a special tool called an ADMG (Acyclic Directed Mixed Graph). Imagine this as a map that draws two types of lines:

Arrows: Showing direct cause-and-effect (A causes B).
Double-headed squiggly lines: Showing "hidden villains" (A and B are both being influenced by something you can't see).

TabPFN-CFM is unique because it can draw these squiggly lines automatically, guessing where the hidden villains might be, even if you didn't tell it they existed.

4. Using a "Cheat Sheet" (Known Graphs)

Sometimes, you do know part of the story. Maybe an expert told you, "We know for sure that A causes B."
The paper shows that if you give this information to the model (like handing the detective a cheat sheet), the model gets even better at its predictions. It uses this known structure to refine its guesses about the rest of the mystery.

5. How It Learned (The Training)

The model wasn't trained on real human data initially. Instead, the researchers wrote a computer program to generate 100,000s of fake datasets with random rules, random noise, and hidden variables.

They taught the model to look at a fake dataset and guess the structure (the map) and the outcome (the result).
They trained it to do this for all three levels of questions (Observer, Intervener, Time Traveler) at the same time.
The Result: When they tested it on real data (like Amazon sales or Law School admissions), it performed better than older methods, especially for the "What If" questions.

6. The "Speed Upgrade"

The authors also mention they tweaked the model's "engine" (the architecture). They removed some unnecessary parts and added new training tricks (like a better optimizer called "Muon" and a new activation function called "ReLU2").

The Analogy: Imagine upgrading a car engine. They didn't just make the car go faster; they made it reach top speed 4 times faster while using less fuel (computational power). This means the model learns much more efficiently.

Summary of Claims

It predicts structure: It draws the map of cause-and-effect, including hidden factors.
It predicts outcomes: It answers "What happens if I change X?"
It handles counterfactuals: It answers "What would have happened if X were different?"
It uses known maps: If you give it a partial map, it uses it to improve accuracy.
It generalizes: It works well on real-world data even though it was trained on fake data.
It is fast: New training tricks make it 3-4x more efficient to train.

The paper does not claim this is ready for medical diagnosis or legal courtrooms yet. It strictly claims that on the datasets they tested (synthetic math problems, Amazon sales, and Law School admissions), this new "Super-Detective" outperforms the current best tools for figuring out cause-and-effect.

Technical Summary: TabPFN-CFM

Problem Statement

Determining causal structure and predicting causal outcomes from observational data remains a significant challenge in machine learning and statistics. Existing methods typically address these tasks in isolation: some focus on structural predictions (learning the causal graph), while others focus on outcome predictions (estimating the effect of interventions). Furthermore, many approaches struggle to handle unobserved confounding variables, which can lead to unidentifiable graphs and biased outcome estimates. Additionally, current models often lack the ability to support the full spectrum of causal queries defined by Pearl's Causal Hierarchy: observational, interventional, and counterfactual.

Methodology

The authors introduce TabPFN-CFM, a causal foundation model designed to jointly predict causal structure and outcomes from observational data. The approach is grounded in the Bayesian Prior-Over-Functions (PFN) framework.

Core Framework

The model assumes the underlying system follows a Structural Causal Model (SCM) $\psi = \{U, V, F\}$ , where $U$ represents unobserved variables, $V$ observed variables, and $F$ structural equations. The model is trained to approximate the posterior distribution $P(\psi | D_{obs}, G_{est})$ , where $D_{obs}$ is the observational dataset and $G_{est}$ is optional prior knowledge of the graph.

Data Generation and Training

Synthetic Priors: The model is trained on synthetic datasets generated from a prior distribution over SCMs. This prior includes random graph structures (DAGs), missing nodes, diverse noise distributions, and structural equations defined by neural networks with random weights and nonlinearities.
Task Formulation: Each training sample consists of a fit set ( $D_{fit}$ $D_{f i t}$ ), a prediction query ( $D_{pred}$ $D_{p r e d}$ ), and a counterfactual sample ( $D_{causal}$ $D_{c a u s a l}$ ). The model is trained to handle three query types simultaneously:
1. Observational Query: Predicting $P(y^* | x^*, T=t^*)$ .
2. Interventional Query: Predicting $P(y^* | x^*, do(T=1-t^*))$ .
3. Counterfactual Query: Predicting $P(y^* | x^*, do(T=1-t^*), y=y^*)$ .
Graph Handling: The model utilizes Acyclic Directed Mixed Graphs (ADMGs) to represent unobserved confounding via bidirected edges. This allows the model to predict confounding structures rather than assuming all variables are observed.

Architecture

TabPFN-CFM extends the Do-PFN architecture with explicit graph conditioning:

Encoder: Embeds the fit set, prediction query, and the prior ADMG structure (adjacency matrix $A$ , bidirected matrix $C$ , and ancestral matrix $R$ ). The graph encoding provides direct access to parent, child, ancestor, and descendant relationships.
Transformer Backbone: Utilizes row-wise and column-wise attention mechanisms. Crucially, the prediction module attends to the graph embeddings, allowing predictions to incorporate the graph prior.
Outputs:
- Outcome Prediction: A head produces logits over discretized outcome buckets.
- Structure Prediction: A decoder generates element-wise predictions for the directed adjacency matrix ( $\hat{A}$ ), bidirected correlation matrix ( $\hat{C}$ ), and ancestral matrix ( $\hat{R}$ ).
Training Efficiency: The authors implemented several architectural and training optimizations (e.g., Muon optimizer, ReLU2 activation, pre-layer normalization) that improved training speed by approximately 3-4x compared to baseline configurations.

Loss Function

The model is optimized using a joint loss function:
$L = L_{pred} + \lambda \cdot L_{graph}$
Where $L_{pred}$ is the cross-entropy loss for outcome prediction, and $L_{graph}$ is the binary cross-entropy loss for structural predictions (adjacency, ancestral, and confounding matrices). Theoretical proofs in the appendix demonstrate that including graph input as a condition can only improve posterior accuracy (or leave it unchanged) and that minimizing the log-probability is equivalent to minimizing the KL divergence to the true posterior.

Key Contributions

Unified Causal Foundation Model: TabPFN-CFM is the first model to jointly predict both causal structure and outcomes across all three levels of Pearl's Causal Hierarchy within a single foundation model framework.
Handling Unobserved Confounding: By employing ADMGs, the model explicitly predicts unobserved confounding variables, addressing a limitation in existing deep learning causal structure learning methods.
Graph-Conditioned Inference: The model leverages known graph structures when available to improve prediction accuracy, while remaining robust when the graph is unknown.
Efficiency Improvements: The refined training procedure and architectural changes result in nearly 4x faster training convergence compared to previous baselines.

Experimental Results

The model was evaluated on synthetic datasets (both in-distribution and out-of-distribution) and real-world datasets (Amazon Sales and Law School Admissions).

Outcome Prediction:
- On synthetic Out-of-Distribution (OOD) SEMs, TabPFN-CFM achieved lower Mean Squared Error (MSE) than established meta-learners (S-, T-, X-, DR-learners) and Do-PFN in interventional and counterfactual settings.
- In the Law School Admissions dataset, the model significantly improved accuracy in counterfactual settings by utilizing counterfactual information, outperforming baselines that cannot make such predictions.
Structure Prediction:
- TabPFN-CFM outperformed strong baselines (FCI, GES, LiNGAM, PC, AVICI) in predicting adjacency, ancestral, and confounding matrices.
- Specifically, on OOD synthetic data, the model achieved an AUROC of 0.889 for adjacency and 0.828 for confounding matrices, significantly higher than the next best performer (AVICI at 0.623 and 0.563 respectively).
Graph Conditioning: Providing the true graph structure as input improved performance in interventional and counterfactual settings but had negligible impact on observational predictions, consistent with causal theory.

Significance

The paper posits that TabPFN-CFM represents a step forward in unifying causal inference tasks. By training a single model on a diverse prior of causal systems, the approach demonstrates that foundation models can generalize to real-world causal problems without overfitting to specific prior distributions. The ability to handle unobserved confounding and support counterfactual queries suggests a pathway toward more robust causal reasoning in machine learning. The authors note that while the current work focuses on acyclic structures, future extensions could address time series and cyclic causal settings.