📊 statistics

FoundCause: Causal Discovery with Latent Confounders from Observational Data

FoundCause is a novel amortized causal discovery model that leverages a transformer architecture with specialized inductive biases to directly infer causal graphs, including latent confounders, from observational data in a single forward pass, outperforming both classical and existing amortized methods across diverse real-world datasets.

Original authors: Patrick Blöbaum, Krishnakumar Balasubramanian, Shiva Prasad Kasiviswanathan

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Patrick Blöbaum, Krishnakumar Balasubramanian, Shiva Prasad Kasiviswanathan

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 figure out who caused what in a complex crime scene, but you only have a pile of blurry surveillance photos (observational data). You can't go back and rerun the crime (interventions), and you suspect there are invisible masterminds (latent confounders) pulling strings behind the scenes that you can't see.

This is the challenge of Causal Discovery. The paper introduces a new detective tool called FoundCause.

Here is how FoundCause works, explained through everyday analogies:

1. The "Training Montage" (Amortized Inference)

Most old-school detectives (traditional algorithms) try to solve every single new case from scratch. They spend hours analyzing the specific photos, running complex math tests, and searching for patterns. This is slow and often fails if the case is messy.

FoundCause is different. It's like a detective who has watched thousands of training movies (synthetic data) before ever seeing a real case.

The Analogy: Instead of solving a puzzle one piece at a time for every new puzzle, FoundCause has already practiced on millions of fake puzzles. It learned the general rules of how causes and effects look.
The Result: When you give it a real dataset, it doesn't need to "think" or calculate for hours. It just looks at the data and instantly (in a single forward pass) draws the map of who caused what. It's like recognizing a face in a crowd instantly because you've seen that face type a million times before.

2. The "Two-Channel Vision" (Handling Missing Data)

Real-world data is often messy. Some surveillance cameras are broken, or some witnesses didn't show up (missing data).

The Analogy: Traditional methods often try to "guess" the missing pieces by filling in the blanks with averages (like guessing a missing puzzle piece is blue because the sky is blue). This often leads to wrong conclusions.
FoundCause's Trick: It has a special "mask" channel. It doesn't just look at the numbers; it looks at where the numbers are missing. It treats the "missingness" itself as a clue. It's like a detective who knows that the absence of a witness in a specific room is just as important as the witness who was there.

3. The "Hidden Puppet Master" (Latent Confounders)

Sometimes, two things happen together not because one caused the other, but because a third, invisible thing caused both.

The Analogy: Imagine you see that "Ice Cream Sales" and "Shark Attacks" go up together. A bad detective might say "Ice cream causes shark attacks." A smart detective knows there's a hidden factor: Summer Heat.
FoundCause's Trick: It has a special module dedicated to finding these invisible "Summer Heats." It uses learnable "tokens" (like invisible sticky notes) to represent hidden common causes. It explicitly asks: "Is there a hidden puppet master pulling the strings for these two variables?" This is a first for this type of AI model.

4. The "Specialized Toolkit" (Inductive Biases)

FoundCause isn't just a generic brain; it's built with specific tools designed for causality.

The "Asymmetry" Sensor: Causes and effects are rarely symmetrical. If you push a ball, it moves; if the ball moves, it doesn't necessarily push you. FoundCause uses classical statistical "asymmetry" measures (like checking if the noise in the data looks different depending on which way you look) to figure out direction.
The "Triangle" Refiner: Causality often happens in chains (A causes B, B causes C) or V-shapes (A and B both cause C). FoundCause has a "triangular refinement" step. It looks at groups of three variables at a time to see if they form a chain or a fork, helping it distinguish between similar-looking patterns that confuse other methods.

5. The "Factored Decoder" (Separating "If" from "Which Way")

When deciding if two variables are connected, FoundCause splits the decision into two steps:

Existence: "Are these two variables connected at all?" (Symmetric check).
Direction: "If they are connected, who is the boss?" (Asymmetric check).
This is like first asking, "Did these two people talk?" and then asking, "Who started the conversation?" rather than trying to answer both at once.

The Results: Why It Matters

The paper tested FoundCause against 15 different real-world datasets (ranging from medical records to microservice logs) and compared it to 11 classical methods and 4 other AI methods.

Speed: It solves the problem in under 2 seconds, while classical methods can take hours or even days.
Accuracy: It got the "direction" of the causes right more often than any other method (improving F1 score by nearly 10%).
Robustness: It didn't crash or fail on difficult, messy datasets where other methods gave up.
Generalization: Even when tested on datasets with more variables than it was ever trained on (like going from a 50-variable puzzle to a 100-variable puzzle), it didn't fall apart; it just got slightly less precise, but still worked.

In summary: FoundCause is a "super-detective" that learned from millions of fake cases, knows how to spot invisible puppet masters, can handle broken cameras, and solves new mysteries instantly by applying the patterns it learned, rather than reinventing the wheel for every new case.

Technical Summary: FoundCause

Problem Statement

Causal discovery from observational data remains a fundamental challenge in machine learning and statistics. The core difficulty lies in recovering directed causal structures and identifying latent confounding without access to controlled interventions. While classical methods (e.g., constraint-based PC, score-based GES, and continuous optimization approaches like NOTEARS) are well-founded, they suffer from significant limitations: they require expensive per-dataset optimization, are sensitive to hyperparameters and distributional assumptions, and often fail to scale to high-dimensional or heterogeneous data regimes. Furthermore, causal discovery is generally not identifiable without strong assumptions (e.g., linearity, additive noise, or faithfulness), which are frequently violated in real-world scenarios.

Recent amortized approaches have attempted to address scalability by training models on synthetic data to map datasets directly to graphs. However, existing neural methods often lack explicit incorporation of domain-specific causal inductive biases, such as asymmetry in cause-effect relationships or higher-order structural patterns (chains, colliders), and typically fail to explicitly model latent confounding.

Methodology: FoundCause

FoundCause is an amortized causal discovery model designed to map observational datasets directly to causal graphs (both directed edges and latent confounding matrices) in a single forward pass. The model is trained entirely on synthetic Structural Causal Models (SCMs) and generalizes to real-world data without per-dataset tuning.

Architecture

The architecture is a hybrid design combining a permutation-invariant transformer with explicit statistical inductive biases:

Input Representation: The model accepts a data matrix $X \in \mathbb{R}^{N \times D}$ and an optional binary observation mask $M$ . Each entry is represented as a 2-channel vector containing the normalized value and the observation status, ensuring the model handles missing data natively without imputation.
Axis-Factorized Encoder: The core encoder utilizes a transformer with alternating attention blocks over samples ( $N$ $N$ ) and variables ( $D$ $D$ ).
- Variable-Attention: Incorporates learnable global tokens and a stat-conditioned bias. This bias injects pairwise statistical features derived from classical asymmetry measures (e.g., correlation, residual variance asymmetry, V-structure indicators) directly into the attention mechanism, guiding the model toward known causal signals.
- Sample-Attention: Aggregates information across samples to capture distributional properties.
Pooling (PMA): A Pooling by Multihead Attention (PMA) mechanism aggregates sample-level embeddings into per-variable representations, combined with max-pooling to capture higher-order distributional properties like multimodality.
Factored Edge Decoder: The prediction head separates edge existence from edge direction:
- Existence Head: Uses symmetric features (sum of embeddings, element-wise product, and confounding scores) to predict if an edge exists.
- Direction Head: Uses antisymmetric features (difference of embeddings and gated asymmetric statistics) to predict the orientation.
Triangular Refinement: Inspired by AlphaFold2, a refinement module performs $R$ rounds of updates over variable triples $(i, k, j)$ . This enables reasoning over higher-order causal motifs (chains, forks, colliders) to resolve Markov equivalence classes.
Confounder Module: A dedicated module uses learnable latent tokens to model hidden common causes. It computes pairwise confounding probabilities via a noisy-OR aggregation of variable-token loadings, producing a symmetric confounding matrix $\hat{C}$ .

Training and Loss

The model is trained on thousands of synthetic SCMs generated via DoWhy and custom generators, spanning diverse graph topologies, nonlinear mechanisms, and noise distributions. Key training strategies include:

Continuous Regeneration: Synthetic data is regenerated every epoch to prevent overfitting to a fixed corpus.
Anti-Shortcut Augmentation: Random permutation of variable indices and subpopulation mixing are used to prevent the model from relying on spurious correlations (e.g., variable ordering).
Loss Function: The objective combines a class-balanced binary cross-entropy for directed edges with several auxiliary terms:
- Direction Calibration ( $L_{cal}$ ): Limits extreme logit gaps between opposite directions.
- Skeleton Supervision ( $L_{skel}$ ): Provides direction-agnostic supervision for adjacency.
- Confounder Loss ( $L_{conf}$ ): Supervises the latent confounding module.
- Sparsity and Regularization: Penalties for overly dense graphs and regularization for pooling/gating mechanisms.

Key Contributions

Explicit Latent Confounding Modeling: FoundCause is the first amortized causal discovery approach to explicitly model latent confounding using dedicated learnable tokens, outputting both a directed acyclic graph (DAG) and a symmetric confounding matrix.
Statistical Inductive Biases: Unlike generic neural architectures, FoundCause injects classical pairwise statistical features (asymmetry, V-structures) directly into the attention mechanism and decoding process, bridging classical causal discovery principles with deep learning.
Triangular Refinement: The introduction of a triangle-based refinement module allows the model to reason over higher-order causal motifs, improving the resolution of Markov-equivalent structures.
Native Missing Data Handling: The architecture natively processes masked inputs, avoiding the bias introduced by imputation methods used in classical baselines.
Zero-Shot Generalization: The model demonstrates robust generalization to variable counts ( $D$ ) and sample sizes ( $N$ ) beyond its training range ( $D \in [2, 50]$ ), maintaining performance on datasets with up to $D=100$ .

Experimental Results

FoundCause was evaluated on 15 real-world and semi-realistic benchmarks (including ASIA, CHILD, SACHS, and PETSHOP) and the Tübingen cause-effect pairs benchmark.

Performance: FoundCause outperforms 11 classical non-amortized methods (e.g., PC, GES, NOTEARS) and 4 amortized methods (e.g., AVICI, CSIvA, SEA).
- It achieves a +9.6% improvement in F1 and +1.2% in AUROC over the strongest non-amortized methods.
- It reduces the Structural Hamming Distance (SHD/d) by 18.9% relative to the best classical baseline.
- It is the only method to achieve top rank on F1, AUROC, and SHD/d while successfully completing inference on all 15/15 datasets (some classical methods time out or fail on dense graphs).
Bivariate Direction: On the Tübingen benchmark, FoundCause approaches the performance of specialized bivariate methods (SLOPPY, HECI) and significantly outperforms classical baselines like ICA-LiNGAM.
Robustness: Under 30% Missing At Random (MAR) corruption, FoundCause retains 92% of its clean-data F1, whereas classical methods relying on mean imputation degrade significantly.
Efficiency: Inference is performed in under 2 seconds on average, a single forward pass, compared to hours required by optimization-based methods.

Significance and Claims

The paper claims that FoundCause represents a significant step toward scalable, practical causal discovery. By combining amortized inference with structured inductive biases, it addresses the trade-off between the theoretical guarantees of classical methods and the scalability of neural approaches.

The authors emphasize that FoundCause does not rely on new identifiability guarantees but rather learns practical inference procedures that perform well across a broad spectrum of data-generating processes. The explicit modeling of latent confounding and the ability to handle missing data natively are highlighted as critical advancements for real-world applicability. The work positions FoundCause as both a practical tool for applied causal analysis and a foundation for future models that natively incorporate causal reasoning. The authors plan to open-source the model to support reproducibility.