🤖 AI

Foundation Models for Partial Causal Identification

This paper proposes a causal foundation model framework that defines a canonical prior over structural causal models to translate the problem of bounding partially-identifiable causal effects under unobserved confounding into learning distributions over functions mapping data and assumptions to causal queries.

Original authors: Alexis Bellot, Anish Dhir

Published 2026-08-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexis Bellot, Anish Dhir

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

In the world of data science, researchers often face a frustrating gap between what they can see and what they want to know. Imagine a doctor who has records of thousands of patients, noting their symptoms, treatments, and outcomes, but who never ran a controlled experiment where they assigned treatments randomly. The doctor wants to know: if a specific patient had received a different medicine, would they have recovered? This is a question about cause and effect, but the data only shows what happened, not what would have happened under different circumstances. Usually, to answer such questions with certainty, scientists need to know the exact hidden rules that govern how the data was created. Without knowing these rules, the answer is often a range of possibilities rather than a single number. This uncertainty is known as partial identification, and it has long forced researchers to build custom, one-off solutions for every new problem, a slow and fragmented process that leaves many important questions unanswered.

A team of researchers has now developed a new approach that acts as a universal solver for these tricky questions. Instead of crafting a unique mathematical formula for each new dataset, they trained a single, powerful computer model to learn the landscape of all possible cause-and-effect scenarios. This model, built on a foundation of "structural causal models," which are essentially detailed maps of how variables influence one another, was taught using a vast library of synthetic examples. The researchers created a specific type of mathematical starting point, or prior, that ensures the model considers every conceivable way the data could have been generated, provided the variables are discrete, like categories or counts. By training on millions of these synthetic scenarios, the model learned to look at a real-world dataset and instantly output a probability distribution that captures the full range of plausible answers.

The result is a system that can take an observational dataset and a specific question about a hypothetical intervention, such as "what would happen if we changed this variable," and return a tight, mathematically guaranteed bound on the answer. In their experiments, the researchers tested this system on simple two-variable systems where the correct answers were already known through traditional, slow mathematical methods. The new model produced results that were just as accurate as the old methods but did so in a fraction of the time. While the traditional approach took over a thousand milliseconds to compute a single answer, the new model delivered its result in less than ten milliseconds, even as the amount of data grew. More importantly, the model's output was not just a guess; as the amount of data increased, the range of answers it provided shrank and converged precisely on the true, mathematically correct limits of what could be known.

This work represents a significant shift in how we handle uncertainty in causal reasoning. Previously, if a researcher wanted to know the bounds of a causal effect in a complex system with hidden confounders, they had to derive a new optimization problem from scratch, a task that required deep expertise and was prone to error. The new method replaces this bespoke derivation with a learned capability. The researchers demonstrated that by training on a "canonical" set of possibilities that covers the entire space of valid causal structures, the model learns to recognize the patterns that define the limits of knowledge. When faced with new data, it does not just predict a single value; it maps out the entire set of values that are consistent with the data and the laws of causality. This means that for any observational dataset, the model can provide a valid, tight interval that contains the true answer, effectively turning a difficult, custom mathematical puzzle into a fast, standard prediction task.

The implications of this approach extend beyond speed. By providing a universal tool that works across different types of causal queries and datasets, the researchers have created a foundation model for partial identification. This means that the same underlying system can be applied to questions in economics, epidemiology, or artificial intelligence without needing to be re-engineered for each specific field. The model's ability to learn the distribution of possible answers, rather than just a single point estimate, respects the inherent uncertainty of the real world. It acknowledges that sometimes, data alone cannot tell us exactly what would have happened, but it can tell us exactly what is possible. In doing so, it offers a robust, principled way to navigate the fog of unobserved factors, giving decision-makers a clear, mathematically sound boundary within which the truth must lie.

Technical Summary: Foundation Models for Partial Causal Identification

Problem Statement

Causal inference often faces the challenge of identifiability: the causal query of interest (e.g., a counterfactual probability or interventional effect) may not be uniquely computable from observational data, even in the large-sample limit, particularly in the presence of unobserved confounding. In such cases, the goal shifts from point estimation to partial identification, where the objective is to determine the tightest possible bounds (the "identified set") containing all values consistent with the observed data and structural assumptions.

Existing methods for partial identification are highly fragmented. Historically, deriving bounds for a specific causal structure, query, and dataset type requires bespoke derivations and separate optimization runs. While recent Causal Foundation Models (CFMs) have successfully learned distributions over causal quantities using Prior-Data Fitted Networks (PFNs), they have largely operated in point-identifiable regimes. When applied to partially identified settings, existing approaches often fail to provide principled bounds because their priors do not guarantee full support over all Structural Causal Models (SCMs) consistent with the data, potentially excluding valid counterfactual values.

Methodology

The paper proposes a universal foundation model for partial identification that outputs a posterior distribution over a causal query, where the support of this distribution provably encodes the identified set. The methodology rests on three pillars:

1. Canonical Priors with Full Support

To ensure the model considers all SCMs compatible with the data, the authors define a canonical prior ( $\Pi_0$ ) over the space of SCMs with discrete observables.

Canonical SCMs: Leveraging the result that any SCM with discrete observables can be equivalently represented by a canonical SCM parameterized by categorical probability mass functions over exogenous variables (Zhang et al., 2022), the authors define a parameter space that is representation-complete.
Prior Construction: The prior $\Pi_0$ $Π_{0}$ is defined over the joint space of causal diagrams ( $G$ $G$ ) and canonical parameters ( $\theta$ $θ$ ).
- The graph structure $G$ is sampled from an Erdős-Rényi model (Bernoulli edge probabilities).
- The exogenous parameters $\theta$ are sampled from Dirichlet distributions.
Full Support Property: A key theoretical result (Proposition 1) proves that this canonical prior has full support over the space of all SCMs ( $\mathcal{M}$ ). This ensures that no SCM consistent with the observed distribution is assigned zero probability a priori, a critical requirement for deriving valid bounds in partial identification.

2. Prior-Data Fitted Networks (PFNs)

The authors employ a Prior-Data Fitted Network to approximate the posterior distribution directly from synthetic data.

Architecture: A transformer-based architecture is used, featuring input embeddings that preserve variable values, types (intervention vs. observation), and correlations. The model utilizes alternating sample-wise and feature-wise attention mechanisms.
Training Regime: The model is trained on synthetic batches of $(D_n, \Psi)$ $(D_{n}, Ψ)$ pairs.
- $D_n$ : Observational datasets sampled from SCMs drawn from the canonical prior $\Pi_0$ .
- $\Psi$ : Ground-truth counterfactual queries computed from the generating SCM.
Objective: The network learns to predict a discrete distribution (binned over $[0, 1]$ ) for the query $\Psi$ given the dataset $D_n$ , effectively amortizing the inference of the posterior support.

3. Handling Structural Assumptions

The framework is extensible to settings with partial or full knowledge of the causal diagram. By training the PFN on tuples $(G, D_n, \Psi)$ , the model can condition on specific causal structures or equivalence classes, leveraging qualitative knowledge to tighten bounds.

Key Contributions

Canonical Prior Definition: The definition of a canonical prior with full support over SCMs with discrete observables. This enables the derivation of bounds that are guaranteed to be valid and asymptotically tight.
Universal Solver for Partial Identification: The development of a foundation model that, given an observational dataset, outputs a posterior distribution whose support captures the full identified set. This contrasts with previous methods that require problem-specific derivations.
Theoretical Guarantees:
- Consistency (Theorem 1): As the dataset size $n \to \infty$ , the PFN's predictive distribution converges weakly to the conditional prior $\Pi_0(\Psi | P_M)$ , which corresponds to the exact Bayesian posterior over the identified set.
- Frequentist Coverage (Theorem 2): The Bayesian credible sets derived from the PFN achieve asymptotic frequentist coverage for the identified set.
- Universality (Corollary 1): The 100% Bayesian credible interval induced by the model converges asymptotically to the true identified set.

Results

The authors evaluate the proposed Causal Foundation Model (CFM) on binary two-variable systems where analytical bounds are known (e.g., Manski bounds, Tian & Pearl bounds).

Accuracy: The predicted posterior distributions concentrate on the analytically derived identified sets.
Coverage: In experiments with varying context sizes ( $n$ ), the CFM achieves coverage rates close to 1.00 (e.g., 0.99–1.00), comparable to or exceeding the performance of a Gibbs sampler (Zhang et al., 2022) adapted for this task.
Efficiency: The CFM demonstrates a significant reduction in inference time compared to the Gibbs sampler. For example, with $n=400$ , the CFM inference time is ~25ms, whereas the Gibbs sampler takes ~1521ms.
Extension: Additional experiments on a three-variable instrumental variable graph confirm the model's ability to handle more complex structures while maintaining high coverage and tight interval widths.

Significance and Claims

The paper claims to extend the paradigm of causal foundation models to the estimation of partially-identifiable causal effects. By translating the problem of bounding counterfactuals into learning distributions over functions mapping data to causal queries, the authors enable a single model to handle broad classes of queries, datasets, and structural assumptions without deriving problem-specific bounds for each instance.

The significance lies in the amortization of partial identification. Instead of solving a new optimization problem for every new dataset and query, the proposed method learns a universal solver that provides valid, tight bounds asymptotically. The use of a canonical prior with full support is presented as the crucial mechanism that distinguishes this approach from previous foundation models, ensuring that the resulting bounds are not only data-driven but also theoretically grounded in the space of all possible causal mechanisms consistent with the observations.