Foundation Models for Causal Inference via Prior-Data Fitted Networks
Authors: Yuchen Ma
Dennis Frauen
Emil Javurek
Stefan Feuerriegel
Venue: ICLR 2026 Poster (2026)
Keywords: Causal Inference
Treatment Effect Estimation
Foundation Models
Primary area: causal reasoning

Abstract: Prior-data fitted networks (PFNs) have recently been proposed as a promising way to train tabular foundation models. PFNs are transformers that are pre-trained on synthetic data generated from a prespecified prior distribution and that enable Bayesian inference through in-context learning. In this paper, we introduce CausalFM, a comprehensive framework for training PFN-based foundation models in various causal inference settings. First, we formalize the construction of Bayesian priors for causal inference based on structural causal models (SCMs) in a principled way and derive necessary criteria for the validity of such priors. Building on this, we propose a novel family of prior distributions using causality-inspired Bayesian neural networks that enable CausalFM to perform Bayesian causal inference in various settings, including for back-door, front-door, and instrumental variable adjustment. Finally, we instantiate CausalFM and explicitly train models to perform in-context learning in these settings. We show that CausalFM achieves competitive in-context learning performance even when compared to baselines that are specifically trained for the task at hand. In sum, our framework can be used as a general recipe to train foundation models for various causal inference settings. In contrast to the current state-of-the-art in causal inference, CausalFM offers a novel paradigm with the potential to fundamentally change how practitioners perform causal inference in medicine, economics, and other disciplines.

Published as a conference paper at ICLR 2026
FOUNDATION MODELS FOR CAUSAL INFERENCE VIA TRIOR-DATA FITTED NETWORKS
Yuchen Ma
Dennis Frauen" LMU Munich & MCML frauen Frauen@lmu.de de
Emil Javurek
LMU Munich & MCML
LMU Munich & MCML
yuchen. ma lmu.de
emi 1 javu rek @ 1u .ee
Stefan Feuerriegel
LMU Munich & MCML
feuerriege1 1mu de
ABSTRACT
Prior-data fitted networks (PFNs) have recently been proposed promising way to train tabular foundation models. PFNs are transformers that are pre-trained on synthetic data generated from prespecified prior distribution and that en- able Bayesian inference through in-context learning. In this paper, we introduce CausalFM, a comprehensive framework for training PFN-based foundation mod- els in various causal inference settings. First, we formalize the construction of Bayesian priors for causal inference based on structural causal models (SCMs) in a principled way and derive necessary criteria for the validity of such priors. Building on this, we propose a novel family of prior distributions using causality-inspired Bayesian neural networks that enable CausalFM to perform Bayesian causal in- ference in various settings, including for back-door, front-door, and instrumental variable adjustment. Finally, instantiate CausalFM and explicitly train models to perform in-context learning in these settings. We show that CausalFM achieves competitive in-context learning performance even when compared to baselines that are specifically trained for the task at hand. In sum, our framework can be used as general recipe to train foundation models for various causal inference settings. In contrast to the current state-of-the in causal inference, CausalFM offers novel paradigm with the potential fundamentally change how practitioners perform causal inference in medicine, economics, and other disciplines.
INTRODUCTION
Causal inference is a cornerstone of empirical research in disciplines such as economics (Angrist 1990, Imbens Angrist 1994), medicine (Feuerriegel et al. 2024 Weberpals et al. (025), and mar- keting /Varian 2016). It enables the estimation of causal effects from observational and randomized data, which is essential for reliable decision decision-making (Kern et al., 2025). In personalized medicine, for instance, it supports identifying the most effective treatment by predicting patient outcomes under different therapeutic options.
In recent years, machine learning, and especially deep learning methods, have gained significant traction in causal inference (Curth & van der Schaar 2021 Ma et al. 2025 2024 Melnychuk et al. 2022; Schheisttal al. 2023] Shalit al. [2017a] Shi al. (2099) These methods offer several advantages for causal effect estimation in practice, including the ability to handle large-scale and nigh-dimensiooal datasets with complex confounding structures and model heterogeneity of causal effects (Feuerriegel et al. 2025). However, most existing approaches require retraining model for each new dataset. To this end, existing approaches lack the flexibility to perform inference for new datasets without additional retraining, which limits their practicality in real-world settings.
"Equal contribution.
1 "orresponning author

---

Published as a conference paper at ICLR 2026
Meanwhile, foundation models have emerged transformative paradigm in machine learning (Devlin 2018 Lahat et al. 2024] Touvron et al.; 2023b/a) a). which offer a key advantage in that they allow for flexible, test-time inference without retraining. These models are pre-trained on large datasets and can generalize across tasks and domains. Examples include large language models (LLMM)) in natural language processing and vision transformers in computer vision. However, this paradigm shift toward test-time inference has not yet had a comparable impact on causal inference. Most current approaches in causal machine learning still rely on specialized models tailored to specific tasks, requiring practitioners to manually select, train, and validate appropriate estimation methods for each new dataset.
In this paper, we propose a change to the paradigm for causal inference based on the idea of foundation models trained for tabular causal inference. For this, we build on the recently proposed prior-data fitted networks (PFNs) (Müller et al. 2022, Hollmann et al. 2023), which are transformers pre-trained on purely synthetic datasets generated from a prespecified prior distribution. PFNs enable Bayesian inference purely through in-context learning, allowing for flexible and efficient predictions without requiring additional training for new tasks (Nagler 2023). While recent works have demonstrated the effectiveness of tabular foundation models based on PFNs for various tasks, only two concurrent works have proposed PFNs tailored for causal inference (Balazadeh et al.) 2025 Robertson et al. 2025). However, these are either restricted specific causal inference settings (namely, only back-door adjustment) or do not offer identifiability guarantees.
We introduce CausalFM, a comprehensive framework for training PFN-based foundation models for various causal inference settings, For this purpose, we introduce CausalFM priors: a novel family of prior distributions based on structural causal models respecting the underlying causal inference problem at hand. We first formalize and derive necessary criteria on how to construct such SCM-based priors for causal inference in principle. Then, we propose a concrete instantiation using Bayesian neural networks and provide learning algorithm that leverages the SCM's 's ability to simulate interventional data to perform Bayesian inference in various causal inference settings.
Compared to classical causal inference methods, models trained based on our CausalFM offer the following advantages: (i) There is no need for additional training for new datasets as our CausalFM performs inference entirely through in-context learning, enabling fast and flexible deployment across new datasets. (ii) The Bayesian nature of our CausalFM provides principled uncertainty quantification, which is critical for downstream decision- poor treatment overlap. (iii) The model automatically learns to "select" an identifiability formula based on the data distribution and task hand. (iv) Our CausalFM builds upon rigorous identifiability guarantees to ensure valid causal inference.
and for detecting situations with
Our contribution arre: (1) We formalize the constructions of priors based on structural causal models (SCMs) for Bayesian causal inference and derive necessary conditions for their validity. (2) We propose an explicit CausalFM prior based on Bayesian neural networks that are compatible with the structure of the causal inference problem at hand. We also propose a learning algorithm to train PFNs for causal inference problems that leverages our CausalFM prior simulate counterfactuals to mitigate the fundamental problem of causal inference. (3) We propose concrete instantiations of our framework by training PFNs for estimating conditional average treatment effects (CATEs) in different causal inference settings. We show empirically that CausalFM performs competitively and outperforms current state-of-the-art CATE estimators on a variety of benchmarks.
2 RELATED WORK
We provide an overview of related literature streams. Additional related work is in Appendix A
Amortized causal inference. Several recent papers pre-train large neural networks on synthetic
data so that they can solve causal tasks via in-context learning. Examples include causal discovery
Maaajaan al. 2025), ATE estimation under unconfoundedness (Zhang al. 2024). zero-shot- and
few-shot learning Nilforoshan et al. [2033] Iwata & Chikahara 2023). and reinforcement-leariin
(Lee et al. 2023 These methods validate the Ieasibility of treating causal inference as an in-context
'Our CausalFM Toolkit available at w com http/////////aaaaattttoo..
Toolkit
docs are at httpss:/ccausa foc wewthiteos.rrrrceccccocooo. io; Project code is available at httpa
7/github. http:h/wwb.coo//cc////csss/l/a/a
2

---

Published as a conference paper at ICLR 2026
learning problem but remain restricted to specific causal inference settings, which typically do not allow accommodating unobserved confounding.
Black-box causal inference (BBCI) (Bynum et al. 2025) proposes maatheicaally-prerainne models to perform causal inference in a variety of settings. However, their approach is different: (i) BBCI does not build upon a Bayesian framework. In contrast, building upon PFNs allows us to perform approximate Bayesian causal inference and thus provide rigorous uncertainty quantification. (ii) The proposed data-generating processes in BBCI are not tailored for high-dimensional causal inference settings (as the authors mention in their Sec. 7). In contrast, our CausalFM prior leverages Bayesian neural networks inspired by TabPFN (Hollmann al., 2023) to create SCM- based prior distributions. (iii) Beyond proposing a new method, we provide novel formalizations and theoretical results of constructing valid SCM-based priors for Bayesian causal inference.
Finally, Dhir et al. 2025 propose framework for amortized estimation of interventional distributions using neural processes Garrelo et al.) 2018). Our work differs that we focus on inference on identified causal quantities without knowledge of the full causal graph (e.g., using clusters of confounders).
PFNs for causal inference: We are aware of only two concurrent works that propose PFN-based mod- els for causal inference, but each with clear limi- tations (see Figure [] (i) (Balazadeh et al. 2025) proposes PFN similar ours, but it is restricted to only back-doot adjustment, i.e., imposes the uncon- foundedness assumption throughout their paper. In
Table 1: Overview of identifiability of PFN- based frameworks for causal inference.
Pramework
Backdoor
Frontdoor
IV
CausaPPFN aal. 2025
x
Do-PEN Robertson al 2025
X
x
Ours (CausalMM)
contrast, we propose a framework for constructing PFN-based foundation models for a large class of causal inference problems, including both front-doo adjustment and instrumental variable settings with unobserved confounding. (ii) Robertson et al. (2025) proposes to train single PFN on various different causal inference settings without providing identifiability assumptions the model. We will show later that the approach of Robertson et al. (2025) has crucial drawback: because the causal quantity of interest is not identified, the PFN learns a posterior that may never concentrate around the true causal quantity, thus leading to asymptotically non-informative estimators. In contrast, we propose to infuse our PFNs with identifiability assumptions required for informative causal inference. As such, we follow established philosophy in causal inference that separates identifiability and estimation steps (Kern et al.) 2025 Pearl (209):: the identifiability step should be established by the practitioner using domain knowledge (e.g.. establishing whether a certain variable is valid instrument), while the estimation step can be treated as a purely statistical learning problem.
3 PROBLEM SETUP
3.1 BACKGROUND ON PFNS
In tabular prediction problems, one considers a population (X, (X - P P. Given a finite sample Dn ~ P of size the goal is to estimate the conditional distribution P(Y = y X =x). PFNs formulate this task in a Bayesian non- parrmmemttriic way by placing prior distribution on i.e. prior over tattu-generating distributions (Müller et al. 2022, Nagler 2023). Sampling proceeds hierarchically via IP
and i.i.d. data (X,, Y.) ~ P. Then, Bayes rule yields the posterior distribution II(P) D,) x II(,, P) ((P). where II(D, P) is the likelihood of the sample Dn under and x denotes proportionality up to a multiplicative constant. The corresponding posterior predicttve distribution is the probability of Y given test point and observed data D,, i.c.,
/ P(Y x = x) [[P Da) dP.
II(Y Da, =
PFNs are neural networks qo (Y Dn, 2) that parameterize the family of predictive posterior distri- butions with trainable parameters 0. That is, PFNs map the entire dataset D, and query to transformers (Ashish Vaswani al. (017) they allow for scalable training and leverage the attention mechanism to effectively extract information from D... PFNs are trained by minimizing the negative log- likelihoddodd loss L()) = Ey-ny [Br-n]] log (Y X, Dx)]]] where TT is prior the sample sizes. In
distribution over y. In terms of architecture, PFNs are
3

---

Published as a conference paper at ICLR 2026
practice, we sample sample size N, II,, probability distribution Py II, dataset Div, p. and test points (2z, UU) ~ P, and then approximate the PFN loss via
(()) = -[-log log qa (vr (,,,,
(2)
which is consistent for the exact cesterior-predictin under regularity conditions (Nagler 2023 Note that all training data are synthetic, i.e., sampled from the prior II. Furthermore, the trained PFN can be deployed on arbitrary real datasets without further training.
3.2 TASK: CAUSAL INFERENCE
In this paper, we aim extend PFNs causal inference. Here, the main challenge that the object of interest is an interventiona P distribution Pint. yet only observe data D, Pobs from potentially different observational distribution (Pearl 2009).
Motivation. As an illustrative example, consider a standard causal inference setting, called backdoor adjustment, where the data comprise (X, AY ~ Pobs. where X are patient covariates, A is treatment, and Y is an outcome of interest (van der Laan & Rubin 2006). For example, in medicine, X may contain treatment history or demographic attributes, A may be a medical treatment, and Y a health outcome. Following the potential outcome framework (Rubin 1974), let Y(a) denote the outcome that would be realized under the treatment A = The interventional distribution is thus over (X, A, (1) (0)) Pint, and common target functional the conditional average treatment effect (CATE) Q(x) E (1) (() X x: (Wager & Athey 2018). The CATE quantifies the expected benefit of providing treatment given the patient's covariates.
Identifiability. To estimate CATE from observational data, need impose identifibilltyy as- sumptions, which link the observational and the interventional distributions and allow us to express Q as a functional of the observational distribution (Rosenbaum & Rubin 1983). These are (i) con- sistency: (AA) = Y, (ii) positivity: Pobe X (1), (0) 11 A X in P Pint.
O. and (111 Unconfoundedness:
Generalized causal inference setting. In the following, we provide a generalized definition of causal inference setting, that allows us to reason about arbitrary causal inference settings and provide generalized statements beyond the standard example above.
Definition 3.1. We define a causal inference setting is a tuple C = (0, x Pint: Q). where O collects the observed variables (and contains at least A and Y); (Ponss Pint)) E Potss x Pint are paired observationa//nnervettinnn distributions over O that correspond to an intervention on A; and (PPm)) a causal query that identifiable, i.c. there exists a measurable functional Q such that (Pim)) (te) for all (Pobs Pint Pint) E Pobs X Pint-
3.2.1 RUNNING EXAMPLES
Example 1 (back-door adjustment). Here, we continue the example from above and define O (X, A, Y) ~ Pobs with bi- nary A above. Pobs Pint contains all observa- tional and interventional distributions that satisfy consistency, positivity, and unconfoundedness. The causal query the CATE (()P)n))())) = (1Y - (0) x = x). which is identified as
A
Figure I: C-DAGs compatible with the three example causal in- ference settings. Yellow variables are observed, blue variables are unobserved, and red variables are clusters of variables.
((Pob))))) EP.... [Y] A - 1, = E.... Y A 0, X z).
*Causal literature often distinguishes between interventional and counterfactual distributions. This is not. relevant for the methods of our paper, and thus interventional distribution as an umbrella term.
4

---

Published as a conference paper at ICLR 2026
Pobs, where X, A, and Y are defined as above and M a mediator between A and Y. Interventional distributions are defined using potential outcomes, i.c. (X, A, (O),, Y(1, (I),, (O, ~ Pint and the causal query of interest again the CATE (x) Er...
Example 2 (front-door adjustment). Let O
AX, (X,
- (0, (O)) X x]. (L, [
Identifiability assumptions. We restrict to pairs (P Pmm)) ) that satisfy (i) consistency: Y = ((A, M) and M MM (A); (ii) positivity: Pobs (A
A a X = x) O; and (iii) front-door criterion (a) 1T A X = and Y(a', m) TT M A = d', X
X x) > 0 and Foba (M
Under these assumptions, the CATE is identified and Q is given via the conditional version of Pearl's front-door adjustment formula (Pearl 2009).
Example 3 (Instrumental variables). Let O = Z, Pobs, where Z is an instrumental variable that causes the treatment A but does not directly cause the outcome Y. The interventional Fine for fixed treatment intervention
distribution is defined on (X, Z, A, (1), Y (0))
a.
We are again interested in the CATE Q(TP inc (x))
(1 Y(0) X = xx]
Identifiability assumptions. We restrict pairs (Pobs Pint) that satisfy the following conditions (Newey & Powell 2003): (i): Additive structural equation: Y = -(X, A) g)+g((X, U), with (unknown) functions and g and unobserved confounder U, implying that Y does not directly depend on Z; (ii) Independence: U TL Z X; (iii Relevance: Pobs (A X = I, Z 2) > 0 is non-constant in 2; and (iv) Completeness: For every measurable g. if 18g (xx A)
I, Z = 2] 0 for all z, then fær A) = 0 almost surely in P... Pobs. Then, the CATE can be shown to be identified via an integral equation (Newey Powell, 2003 Harrffrd et al. 2017).
Research question: PFNs have shown to be an effective way to construct tabular foundation models. However, a causal inference setting C comes with additional challenges, such as the distinction between observational and interventional distribution as well as identifiability assumptions.
Research question
How can we train PFNs for a causal inference setting C that provides a Bayesian estimator of (Iim given an observational dataset D,, Pobs and some context (e.g., values a)?
In the following, we introduce CausalFM consisting of (i) appropriate prior distributions that allow for approximating interventional predictive posterior distributions as in Eq. (Section and (ii) training algorithm for the underlying PNFS (see Section 5).
4 CAUSALFM: PRIORS
In this section, we construct prior distributions for CausalFM which are based upon identifiable structural causal models (SCMs). We motivate and formalize our approach (Sec. 4.1, provide necessary criteria for valid causal inference (Sec. A.2), and finally provide a method for constructing such priors in practice (Sec. 4.3). We also provide complete toy example in Appendix
4.1 INTRODUCING SCM-BASED PRIORS
Naïve approach. A naive approach for causal inference would construct prior 11 directly for the observational distribution P. Pobe. If the posterior II(PPons Dx)) - IP's converges to the ground-truth theorem), we can obtain
observational distribution Pr (i.e., satisfying consistent Bayesian estimator of our causal query via Q(T ohs 'D,)). Accordingly, we could train a PFN go Dn, with the loss in Fq. 2 and estimate the CAPO via Q(1o x)).
Bemsseiin-von-Mise
However, the above approach has drawbacks: (i) It requires knowledge of the identification formula Q which must be determined on a case-by-case basis depending on the causal inference setting at hand. This can be tedious or even hard to compute in practice. For example, the IV setting from Example 3 requires solving integral equations to compute Q (Newey & Powell 2003). (ii) Constructing a prior for P Pote makes it harder control the distribution of the causal query Q directly. It has been shown in the literature that this can lead to prior misspecification for Bayesian causal inference or slowly converging posterior distributions (Linero & Antonelli 2022).
5

---

Published as a conference paper at ICLR 2026
Modeling the interventional distribution. Motivated by these drawbacks of constructing priors for only Pobe, we propose to construct priors for observational- iiservational-iiteventtonn distribution pairs (ooss: ), resulting in priors defined on Pohs X Pint- This addresses both drawbacks by (i) inducing an interventional posterior distribution, thus only requiring knowledge of Q (not Q); and (ii) we will that priors Pobss X Pint often allow specify the prior distribution of (P:in)) directly. A natural way to define distributions on Pobs X Pint is via structural causal models (SCMs).
structural causal model (SCM) S is tuple (2 U, f. P), where = (Z,,.., Zx) are observable endogenous variables, U collects latent exogenous variables, f = {fzz..... fz, contains structural assignments Z, = fzz, (pa((i))))) (Z)) with parents pa(Z, S U, and IP is joint distribution on U.
Definition 4.1 (SCMs (Pearl 2009)). A
Every SCM induces a unique directed acyclic graph (DAG), g by defining mapping of the parents (4) to Z; with directed edges. We distinguish two types of latent variables U; in g U; is an unobserved confounder if it is the parent of both A and Y, otherwise, we call it a noise variable. Intuitively, an SCM is a simulator: we can draw latent variables U ~ P and pass them through structural functions/ f. resulting in an induced observational distribution Ih over Z. At the same time, can modify the SCM by intervening then sampling from the SCM mechanism. This induces a corresponding interventional distribution int-
variable via do(AA a), i.c., fixing the variable and
P.S. We call SCM S compatible with causal inference setting C, if (TP IP Pohs Pint-
Definition 4.2 (C-SCM-Priors). A C-SCM-Prior is any probability measure II(S) that puts all its mass on SCMs compatible with C. Via the map S + (IP obss P every such prior induces distribution 11( (Fobs, Pmmt)) on Pobs x Pint.
Sampling from II therefore amounts to sampling a random latent distribution F over U as well as random functional assignments f. These can then be used to internally sample an observational dataset Dn, i.e., there is a well-defined likelihood [T, (S) induced by S. a consequence, we can define the posterior distribution over SCMs via IT(S Dn)) x II(,, ()))(S), where a denotes proportionality up to a normalization constant.
Cluster-DAGs Because an SCM-prior induces distribution over possibly many DAGs, we compress them into a shared structure. Given variables (Z, U), a Cluster-DAG (C-DAG) (Anand et al. ]023)] is a DAG on clusters CI...., C2 which are disjoint subsets of (Z,U). Each C-SCM- Prior induces a unique C-DAG via the following algorithm: (i) draw an edge C₁ C whenever any SCMs S with (1())= > 0 contains some arrow from any node of C2 to any node of C₁ and no SCM S with (S) > 0 contains some arrow from any node of C; to any node of C;; (ii) merge C; and Cj whenever both directions occur across SCMs S with II(S) > O.
4.2 WELL-SPECIFED PRIORS
The question is now how we should design our prior II such that the induced posterior IT(S) | D,) allows for valid Bayesian causal inference. We now define key desirable property of such priors. For this, we call a prior II(S) well-specified for C if, for any true pair (Pr P and every sample D, ~ F.x abs it holds that
of IPS I1(S) D.) ds)
Q(IPL.):
n 00.
In other words, well-specified prior ensures that the causal query Q evaluated on the posterior- predictive interventional distribution (PPID) SPS in II(S D,,) dS is a consistent estimator of the causal target. If we were able to train PFN to approximate the PPID of a well-specified prior, we are sure that we can apply Q on this distribution and obtain consistent estimator.
Identifiability. At this point, one may wonder why we only foeus on priors for identifiable causal inference settings. Indeed, recently proposed method, called do-PFN (Robertson et al.] (025). does not restrict its PFN priors to identifiable settings. The following result shows that, under weak assumptions, such priors cannot be well-specified, leading to asymptotic inconsistency.
Theorem 4.3 (Identifiability rules out prior mass on violating SCMs). Assume the causal inference setting C is identifiable for the causal query Q with corresponding identified functional Q. Let
z == { {S: IPS. E Pots and Q( ((P())) # Q ( obss) }

---

Published as a conference paper at ICLR 2026
denote the set of iddntffiabliity-villatng SCMs, and for any Pobs Pobs define the observational equivalence class
E(Pobs) := {S F = Pabst.
Assume Q is a linear functional of int and that non-identifiability does not cancel out in prior expectation, i.e., for every Pobe E Pobs with [ n ((oba)) > 0,
J (ng
A(P int II(dS) + Q(Poba) 11( E
Then any prior II(S) that is well-specified for C must satisfy (2) =
Proof. See Appendix C
4.3 CONSTRUCTING SCM- SCM-BASED PRIORS
C-DAG design. Our method for constructing priors assumes the knowledge of a well-specified C-DAG G. for C, meaning that Ge is induced by some well-specified C-SCM-Prior. Such C-DAGs are usually known for most causal inference settings (see Fig. T for C-DAGs compatible with the settings in Examples 1-3).
One point of ambiguity is the modeling of noise variables in C-DAGs. Here, we propose a practical design rule: if G. contains an unobserved confounder between A and Y. only add one additional noise variable to either A or Y. Conversely, if 9. is unconfounded, we add noise parents to both A and Y (see Fig. D. The reasoning as follows: if Ge is unconfounded, to both A and Y in order to ensure not restrict ourselves to degenerate observational distributions. Conversely, any unobserved confounder induces noise into both A and Y. thus removing the need to add noise to both. However, it is still necessary to add one additional noise variable to either A or Y since, otherwise, any unconfounded SCM compatible with Ge would need to be degenerate in either A or Y. We provide a concrete toy example in Appendix B to illustrate this.
need add noise
Prior construction. We now propose a practical algorithm to construct C-SCM-priors. We assume that we have access to pair (G., I), where G. is well-specified C-DAG for C and I is set of constraints on SCMs S compatible with Ge ensuring that also compatible with C.
Example I: back-door adjustment. The observable variables are (X, A. together with noise variables. compatible C-DAG is in Fig. (left). The constraint set I(S) {IPS (A
X z) > 0), ensuring that all SCMs satisty the positivity assumption.
A, Example 2: Front-door adjustment. Here, the observed variables are AX M, with noise
variables and unobserved confounder U between A and Y. A compatible C-DAG is in Fig. (middle). The constraint set is I(S) = { (A = a X = x) > 0, pS (M = m X = I,A
a) 0}, ensuring positivity for both treatments and mediators.
Example 3: Instrumental variables. The observed variables are (X, 2, Y), augmented by noise
variables and unobserved confounder U that is joint parent of A and Y has no edge to the instrument Z; see the compatible C-DAG in Fig. (right). The constraint set is I(S) = {PS. (Z = =
x| =) > 0, f (X,
+ g (X,U)).
Overall algorithm. Given II, we propose construct prior distribution II over SCMs as follows: First, we order the clusters (Ci,..., Ck) according to their hierarchy in the DAG (i.e., Ci has no parents). Then, iterate over each cluster C; follows: if C, only contains latent variables, fix their distribution to standard normal distribution via (+) A((O, I). If C. is cluster of observed and latent variables, we assign a clustered Bayesian neural network (BNN) prior to C, (see below). If C2 only contains observed variables, we assign observational BNN prior.
Clustered BNN prior. For clusters that contain both observed and latent variables, we leverage a BNN- prior inspired by TabPFN (Hollmann al. 2023). This prior allows us effectively sample potentially high-dimensional clusters of variables for which the internal causal structure is irrelevant to infer the causal query of interest. The prior is defined via
(6) 9g
pa(C.) RR"
0 0~ TT:
99 (t)
satisfying I(So).

---

Published as a conference paper at ICLR 2026
(6) We then sample random nodes from Do that coincide with observed nodes in C. while the remaining nodes serve as latent noise within the cluster. This corresponds to applying the approach taking in TabPFN (Hollmann et al., 2023) to clusters C2 in the C-DAG in which the causal structure does not matter for estimating our causal query.
Observational BNN prior. If C, contains only observed nodes, define another BNN via
f/
pa(C:)
RIC. P~I. subject to f/ satisfying I(So),
(6)
and set C₁ f (pa(C.)). The observed nodes within C, thus correspond to the output of the neural network and are not randomly subsampled neurons.
Example back-door adjustment. Here, the data distribution IP can be separated as follows: (X,Ux) ~ Px with Ux denoting noise variables withing the cluster X, UA ~ PU Uy ~ Puy- A = fA(XXUA)) and Y = fa(X, A, Uy). Our algorithm proceeds as follows: PU, and Fty are noise variables and are set to standard normal distributions. The cluster (X, Ux) contains both noise and observed variables, meaning that Pxx is sampled from an clustered BNN prior. Finally, and Y are observed variables meaning that fA and fy are sampled from observational BNN priors. We refer to Appendix D.I full implementation details, including for Example 2 and 3.
Notes on identifiability. Our framework follows established causal inference philosophy and separates identifiability from estimation (Pearl 2009]: the identifiability step (=choosing the causal setting) requires careful modeling and usage of domain knowledge, while the estimation step be handed over to our CausalFM. If practitioners suspect identifiability assumptions may be violated, recommend performing causal sensitivity analysis (Dorn & Guo 2022; Frauen al. 2023 to assess the extent of potential violations.
CAUSALFM: TRAINING
5.1 TRAINING ALGORITHM
We look at the case where the causal query QI' X)) is function of the conditional interven- tional distribution Pint(X X)) for some contextual observed variables X. This includes, e.g., the CATE (1) - (0) X| and CAPO EE( X] from our running examples.
Our goal is to train a PFN qY x)) to approximate the conditional PPID (posterior predictive
interventional distribution) Im (Y 'D,, x) =f I (Y Xx= x)) II(S D,) dS. Given an SCM prior and prior 1Iy over sample sizes, we propose the following modified PFN loss
(())= = (ons-N/[Esnn [I [Exx, [ED-PS
[- log (Y go
X, Dx)]]]]
int ols
Importantly, the dataset D sampled from the observational distribution, while the pair (X, is sampled from the interventional distribution induced by random SCM. This ensures that the PFN will aim to predict the interventional outcome Y based on data following the observational distribution. A similar loss has been proposed by Bynum al. (2025), which, however, is only based on the mean-squared error instead of the negative log-likel interpretation for approximating the PPID in Bayesian setting. In particular, modeling the entire PPID allows us not only to provide point estimators of our causal query, but also to account for uncertainty.
and thus does not allow an
In practice, we sample the sample size N2 ~ Nv. an SCM Sj ~ II, and an observational dataset i by sampling from the SCM. Then, we modify the SCM by performing the intervention
Di,, PS, of interest (e.g., do(A = a)) and sample test points (zj,);) The approximated PFN-loss is then
Sp from the interventional SCM.
(()) E[- log go (yy (,,,,1)))
(8)
Finally, once go (Y 1z) is trained, we can obtain an estimator for the causal query via Q(qo(Y [X).
i... by applying the causal query the approximated PPID by the PFN.
8

---

while the full Published as a conference paper at ICLR 2026
Example: back-door adjustment. Here, we sample a sample size N, II,, an SCM S ~ II from our constructed prior distribution II and an observational dataset Di,, ~ F obs" Then, we
perform two interventions do(A A = 1) and do(A = 0) to obtain test points (I,, (1;)) ~ p.d. The PFN loss becomes
int'
(())= Σ[ log (1, (0) #y)]-
Implementation details. Each observation is tokenized during embedding, with separate encoders applied to observational variables. The resulting tokens are processed by a transformer- -based PFN to obtain representations, which are subsequently passed to Gaussian mixture model (GMM) head. Our implementation of qo ar) is based on the TabPFN architecture (Hollmann et al. 2023). We train the model with learning rate of 1e-3, weight decay 1e-5, batch size 16, and sequence length 1024 for up to 150 epochs. Training CausalFM on a single NVIDIA A100 GPU takes about 24 hours. Details on the data prior and generation details are provided in Appendix D.I D.I implementation is given in Appendix D.2
6 EXPERIMENTS
Table 2: Standard CATE estimation over 10 synthetic datasets and Jobs dataset.
We evaluate our method across three causal inference settings: standard CATE estimation, instrumental vari-
Method
Syathetic
Jobs
BASELINES (A): STANDARD CATE ESTIMATORS
S-learner
Künzel et al.
0.734 ams
0.697 1:
ables (IV), and front-door adjustment.
T-leamer Kunzel et
2019
0.661 ar
0.822
+ c.
TARNet (Shait et al
701
0.854 am
0.864 1 ::4
Evaluation metrics. We report the pre- cision in estimating heterogeneous ef-
DR learner (Keanedy 202
0.765 o.
0.959 1: 5:s
RA-leamer (Curth & Van der Schaar 2021
0.609 am
0.652
+ 5.15
fects (PEHE) (Curth & van der Schaar 202T Hill 2011), deñned as the root
X-leamer (Kunzer <q
2019
0.563 0.
0,802 :
BASELINES (B): FOUNDATION MODELL--BSSEE METHODS
mean squared deviation between pre- dicted and ground-truth CATE, to evalu- ate the model performance on the CATE estimation task.
CausalPEN (Balazadeh et al.
2025
0.557 a.
0.528 4 1-
DoPFN (Robertson et al. 2025
0.586 am 0,482 + c:
CausalFM (ours)
0.515 +a
0.478 + a.
Lower = better. Reported: PEHE (mean + std). Top-three per column blue. purple, orange.
6.1 EVALUATION FOR STANDARD
CATE SETTING
Baselines for standard CATE estimation. We consider a broad range of state-of-the-ar methods for the conditional treatment effect estimation from the literature: (1) S-learner (Künzel et al. 2019):: the S-learner is model-agnostic learner that trains single regression model by concatenating the covariate and the treatment as input; (2) T-learner (Künzel et al. 2019): the T-learner a model- agnostic learner that trains separate regression models for treated and control groups; (3) X-learner (Künzel et al. 2019): builds upon the T-learner by first imputing individual treatment effects in each group and then fitting models to these eseudd-effects;; (4) TARNet (Shalit et al. 2017b): using representation learning to extract features of covariates and train separate branches for treated and control groups with regularization; (5) DR-learner (Kennedy, ( (2036): generates pseudo-outcomes based on the doubly-robust AIPW estimator; (6) RA-learner (Curth & van der Schaar 202T]: uses a regression-adjusted neeudo-outcome in the second stage. We also include two PFN-based foundation
models for treatment effect estimation: (7) CausalPFN (Balazadeh et al. 2025) and (8) DoPFN
(Robertson al. 2025). Further implementation details are in Appendix D.3 D.3
Results on standard CATE estimation. We benchmark our model on ten synthetic datasets generated under diverse mechanisms, with implementation details in Appendix D In addition, we evaluate on a semi-synthetic version of the Jobs dataset (Smith & Todd, 2005). derived from the widely used Lal Londe study .onde 1986). Here, we generate outcomes to create a semi-synthetic dataset and allow for evaluation against ground-truth.
Table repports the averaged PEHE across the synthetic datasets (full results in the Appendix) and the Jobs dataset. Our experiments show that CausalFM achieves competitive CATE estimation performance across all benchmarks, without requiring model retraining.

---

Published as a conference paper at ICLR 2026
Baselines for IV setting. We benchmark against a broad set of state-of-the the-art IV methods for treatment effect estimation:
Table 3: IV setting for CATE estimation with binary and continuous instrument variables.
(1) KIV (Singh et al.] 2019): a nonlinear extension of two-stage least squares using
Method
Binary IV
Continuous
BASELINES (A): STANDARD IV ESTIMATORS
kernel ridge regression with feature maps;
KIV Singh et al. 2019
0.454 0.16
0.577 = 0.0
(2) DFIV (Xu et al., 202I): extends KIV
DRIV
Syypkaams
2019
0.531 1
0.693 1:
DeeplV Hartiord et al
2017
0.427 + 6.5
0.516 + 0:13
by parameterizing feature maps with neu-
DeepGMM [senner
0.503
0.588
ral networks trained iteratively; (3) DeepIV
DMLIV S
Syrpkanis el ai.
2019
0.479 - 0.5
0.618 1: 0.30
(Hartford et al. 2017): two-stage neural approach, first estimating the treatment dis-
DFIV Xu al 2024
0.709 +
0.583 + 0.30
MRIV Fraueh & Feuerriggl 2022
0.688 0.1
0.641 =
BASEL INES (B): FOUNDATION MODELS-BASED METHODS
tribution and then solving a counterfactual prediction task; (4) DeepGMM (Bennett
DoPFN (Robertson al. 2025
0.523 0023:00 0.20
0.675 = 0.37
et al., 2019): formulates IV estimation a minimax game based on the generalized method of moments, solved via adversar- ial training; (5) DMLIV (Syrgkanis et al.]
CausalFM (ours)
0.422 0.16
0.579 = 021
Lower better. Reported: PETTE (mean standard deviation). Top three per column are in blue, purpe, orange.
2019): double machine learning framework that estimates nuisance functions and learns the CATE by orthogonalized regression; (6) DRIV (Syrgkamis et al. 2019): a meta-learner combining DMLIV with doubly robust peeudo-outcome for improved stability: and (7) MRIV (Frauen & Feuerriegel 2022): multiply robust framework for binary IVs that directly estimates CATE via ascudo-outcome regression. For foundation model baselines, as CausalPFN (Balazadeh et al., 2025) only for back-door adjustment, we include DoPEN (Robertson et al. 2025
6.2 EVALUATION FOR IV SETTING
Results on IV setting, We evaluate our models on datasets with varying confounding strengths. Table 3 reports the averaged PEHE for binary and continuous IVs. Note that CausalPFN is not designed for IV settings. In contrast, we find that our CausalFM consistently achieves comparable performance relative standard IV estimators and outperforms biased alternatives. Importantly, in contrast the the standard base- lines, these results hold without requiring model retraining. Hence, this confirms the flexibility of our approach to IV settings.
Table 4: Front-door adjustment setting for CATE estimation.
BASELINES (A): STANDARD FRONT DOOR ADJUSTMENT
Plug in front dooe learner (Linear) Pearl 2009
1.124 1 u22
Plug-in front-door learner (RF) Peart 2010)
1.364
Plug in front door learner (NN) Pearl 20019
0.889
BASELINES (B): FOUNDATION MODELS-BASED METHODS
DoPFN (Robertson et al. 2025
1.274 = 024
CausalFM (ours)
0.847 us
Lower better. Reported: PEHE (mean + standard deviation). Top-three per column blue, purple, orange.
6.3 FRONT-DOOR ADJUSTMENT
We additionally evaluate our model under the front-door adjustment setting Table Due to space constraints, details provided in Appendix H.l The experiments show the flexibility of our method to perform causal inference in the front-door adjustment setting.
6.4 DISCUSSION
Limitations and future work. The current evaluation is limited synthetic and semi-synthetic data due to the fundamental problem of missing potential outcomes on real-world data. For future work, it will be interesting to investigate the performance of CausalFM in applied A/B experimental setups to assess its empirical performance and robustness under real-world conditions. Additionally, an important research direction will be to incorporate interpretability or faimess constraints into Ca which is crucial for reliable deployment in practice.
10

---

Published as a conference paper at ICLR 2026
Ethics statement.
Human subjects and IRB. This work does not involve experiments with human subjects. Our training data synthetically generated from prespecified SCM-based priors. For empirical evaluation, additionally use publicly available benchmark data (c.g., Jobs) where outcomes are generated in semi-synthetic manner following common practice; identifiable personal information is introduced by us. Accordingly, no IRB review was required for this study.
Data, privacy, and security. We do not collect, store, release sensitive personal data. Public datasets are used under their respective licenses, and our semi-synthetic outcome generation avoids re- identification risks. We will document preprocessing and generation steps to support reproducibility.
Bias and potential harms. Causal estimators can be misused if applied outside the assumed identifica- tion regime (e.g., back-door, front-door, IV) or under severe violations (e.g., weak instruments, lack of overlap). To mitigate harm: (i) we make assumptions explicit and provide uncertainty quantifica- tion; (ii) we advocate domain- domain-expert validation and sensitivity checks before deployment; (iii) discourage high-stakes automated decision- decision-making without human oversight.
Use of large language models (LLMs). We used LLM- -aM-ddsed tools to assist with writing (clarity, grammar) and for literature research. All claims were authored and verified by the authors; citations were cross-checked against primary sources. No sensitive data were provided to LLM tools.
Reproducibility statement. We ensure reproducibility of our results by providing the full implemen- tation and training scripts. Our CausalFM Toolkit is available / com/ yccm/ Causa 1aMsalFM--a FM too 1 kit, Toolkit does are at https://causa fm- readtheocss heccca o Project code available thtttp : / ctt.c////////// cww./ccmmCausaaa///cccccpaaa///a FM The repository con- tains the necessary code to reproduce our experiments, along with instructions for dataset preparation, model training and evaluation procedures. This setup allows independent researchers to replicate the reported results and extend our work with minimal effort.

---

Published as a conference paper at ICLR 2026
Acknowledgments. This paper is supported by the DAAD program "Konrad Zuse Schools of Excellence in Artificial Intelligence", sponsored by the Federal Ministry of Research, Technology. and Space. Additionally, this work has been supported by the German Federal Ministry of Education and Research (Grant: 01IS24082).
REFERENCES
Ahmed M. Alaa and Mihacla van der Schaar. Bayesian inference of individualized treatment effects using multi-task Gaussian processes. In NeurIPS, 2017.
Tara V. Anand, Adele H. Ribeiro, Jin Tian, and Elias Bareinboim. Causal effect identification in cluster dags. In AAAI, 2023.
Joshua Angrist. Lifetime earnings and the vietnam era draft lotter: Evidence from social security administrative records. The American Economic Review, 80(3)3113-336, 1990.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Vahid Balazadeh, Hamidreza Kamkari, Valentin Thomas, Benson Li, Junwei Ma, Jesse C. Cresswell, and Rahul G. Krishnan. Causalpfn: Amortized causal effect estimation via in context learning, arXiv preprint, arXiv:2506.07918, 2025. BRL/atpp//aaxxlvoorr//t//////
07918v1
Andrew Bennett, Nathan Kallus, and Tobias Schnabel. Deep generalized methodd method of moments for instrumental variable analysis. In NewrIPS, 2019. URL http://axxxv...or/////////
12495v2
Lucius Bynum, Aahlad Manas Puli, Diego Neerero-Quevedee Nhi Nguyen, Carlos Fernandez- Granda, Kyunghyun Cho, and Rajesh Ranganath. Black box causal inference: Effect estimation via meta prediction. arXiv preprint, arXiv:2503.0598. 055985 2025. URL http:://t/x//ii.////// wttpaxpwwwooogo/////////p//
2503 05985v1
Victor Chernozhukov, Denis Chetverikov. Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James M. Robins. Doublc/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1):C1-668, 2018. ISSN 1368-4221. doi: 10.111 1111///ctj. 12097.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In AISTATS, 2021.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics, 2021.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv 2018
preprint
Anish Dhir, Cristiana Diaconu, Valentinian Mihai Lungu, James Requeima, Richard E Turner, and Mark van der Wilk. Estimating interventional distributions with uncertain causal graphs through meta-learning. In NeurIPS, 2025.
Jacob Dorn and Kevin Guo. Sharp sensitivity analysis for inverse propensity weighting via quantile balancing. Journal of the American Statistical Association, 2022. URL http://arxiv..rrt
pdf/2102. 04543v2
Stefan Feuerriegel, Dennis Frauen, Valentyn Melnychuk, Jonas Schweisthal, Konstantin Hess, Alicia Curth, Stefan Bauer, Niki Kilbertus, Isaac S. Kohane, and Mihaela van der Schaar. Causal machine learning for predicting treatment outcomes. Nature Medicine, 2024.
Stefan Feuerriegel, Yash Raj Shrestha, and Georg Von Krogh. A new machine learning approach answers what-if questions. MIT Sloan Management Review, 66(3):65-69, 2025.
Dylan J. Foster and Vasilis Syrgkanis. Orthogonal statistical learning. The Annals of Statistics, 53(3): 879-908, 2023. ISSN 0090-5364. BBR.hhtpp./aarxiv.orr//////////
09036v
12

---

Published as a conference paper at ICLR 2026
Dennis Frauen and Stefan Feuerriegel. Estimating individual treatment effects under unobserved confounding using binary instruments. In ICLR, 2022.
Dennis Frauen, Valentyn Melnychuk, and Stefan Feuerriegel. Sharp bounds for generalized causal sensitivity analysis. In NeurlPS, 2023.
Marta Gamelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.01622 2018.
Richard Hahn, Jared S. Murray, and Carlos Carvalho. Bayesian regression models for causal inference: regularization, confounding, and heterogeneous effects. Bayesian Analysis, 15(3): 965-1056, 2020. URL http:://aaxx..oog/pof///////////
09523v4
Jason Hartford, Greg Lewis, Kevin Leyton-Brown and Matt Taddy. Deep IV: A flexible approach for counterfactual prediction. In JCML, 2017.
Jennifer L. Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: A transformer that solves small tabular classification problems in second. In ICLR, 2023. URL http:
7 1/arxiv. www. 0012000008100000 01848v6
Noah Hollmann, Samuel Müller, Lennart Purucker, Arjun Krishnakumar, Max Körfer, Shi Bin Hoo, Robin Tibor Schirmeisterr, and Frank Hutter. Accurate predictions on small data with a tabular foundation model. Nature, 637(8045):319-33) 2025, doi: 10. 10.1008/413386--2440333444
Shi Bin Hoo, Samuel Müller, David Salinas, and Frank Hutter. The tabular foundation model tabpfn outperforms specialized time series forecasting models based on simple features. arXiv preprint, D.Xiv:2501.02945 2025. Uttptttp: /tartt:///rxiv...o wwwpp/wwwwwwpp////////////
02945v2
Guido W. Imbens and Joshua D. Angrist. Identification and estimation of local average treatment effects. Econometrica, 62(2):467-475 1994.
Tomoharu Iwata and Yoichi Chikahara. Meta-learning for heterogeneous treatment effect estimation with foosed-form solvers. arXiv preprint, S..Ii::2035113555 2023. URL http:::/aax....org
Fredrik Johansson, Uri Shalit, and David Sonntag. Learning representations for counterfactual inference. In ICML, 2016.
Edward H. Kennedy. Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2):3008-3049, 2023a.
Edward H. Kennedy. Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, (112:3308-3049) 2023b.
Christoph Kern, Unai Fischer-Abaigar, Jonas Schweisthal, Dennis Frauen, Rayid Ghani, Stefan Feuerriegel, Mihaela van der Schaar, and Frauke Kreuter. Algorithms for reliable decision-making need causal reasoning. Nature Computational Science, 2025.
Sören R. Künzel, Jasjeet S. Sekhon, Peter J. Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, (16(10):4156-4)
2019.
Adi Lahat, Kassem Sharif, Narmin Zoabi, Yonatan Shneor Patt, Yousra Sharif, Lior Fisher, Uria Shani, Mohamad Arow, Roni Levin, and Eyal Klang. Assessing generative pretrained transformers (gpt) in clinical decision-malee makinn: comparative analysis of gpt-3.5 and gpt-4. Journal of medical Internet research, 26:e54571, 2024.
Robert LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pp. 604-620, 1986.
13

---

Published as a conference paper at ICLR 2026
Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsesa Finn, Ofrr Nachum, and Emma Brunskill. Supervised pretraining can learn in context reinforcement learning. In NeurIPS, 2023.
Antonio R. Linero and Joseph L. Antonelli. The how and why of bayesian nonparametric causal inference. Wiley Interdisciplinury Reviews: Computational Statistics, 15(1):e1583, 2022. URL
http:://axxat
owg. wttp://orggppd/////11/ htt.///2.............22 03897v2
Yuchen Ma, Valentyn Melnychuk, Jonas Schweisthal, and Stefan Feuerriegel. DiffPO: A causal diffusion model for learning distributions of potential outcomes. In NeurIPS, 2024.
Yuchen Ma, Jonas Schweisthal, Hengrui Zhang, and Stefan Feuerriegel. A diffusion-based method for learning the multi-outcome distribution of medical treatments. In KDD, 2025.
Divyat Mahajan, Jannes Gladrow, Agrin Hilmkil, Cheng Zhang, and Meyer Scetbon. Zero-shot learning of causal models. arXiv preprint, arXiv:2410.06122 2025. URL http://aaxx: org/pdf/241s 24t.f///1100661282
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Causal transformer for estimating counterfactual outcomes. In ICML, 2022. URL wwtppxwwwwpp/////2/2/222 22044. 072
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. ICLR, 2022. URL httpp://axxiv..rg/ppp/..
2112. 11110v0
Thomas Nagler. Statistical foundations of prior-data fitted networks. In ICML, 2023.
Whitney K. Newey and James L., Powell. Instrumental variable estimation of nonparametric models. Econometrica, 71 71(5):1565-15777 2003.
Xinkun Nic and Stefan Wager. Quasi-oracle estimation of heterogennouss treatment effects. Biometrika, 108(2):299-33 2199999999999999 2021. ISSN 0006-3444.
Hamed Nilforoshan, Michael Moor, Yusuf Roohani, Yining Chen, Anja Surina, Michihiro Yasunaga, Sara Oblak, and Jure Leskovec. Zero-shot causal learning. In NeurIPS, 2023.
Judea Pearl. Causality. Cambridge University Press, New York City, 2009. ISBN 9780521895640
Jake Robertson, Arik Reuter, Siyuan Guo, Noah Hollmann, Frank Hutter, and Bernhard Schölkopf. Do-pfn: In-context learning for causal effect estimation. arXiv preprint, arXiv:2506.0666 06009, 2025.
htt:httpp///xiiii
/atzg//fff// 2500 06039v1
James M. Robins. new approach to causal inference in mortality studies with sustained exposure period: Application to control of the healthy worker survivor effect. Mathematical Modelling, 1393-1512, 1986.
James M. Robins. Correcting for non-compliance randomized trials using structural nested mean models. Communications in Statistics Theory and Methods, 23(8):2379-2411 1994. ISSN 0361-0926. doi: 10.1080003611999940888133999
James M. Robins. Robust estimation in sequentially ignorable missing data and causal inference models. Proceedings of the American Statistical Association on Bayesian Statistical Science, pp- 6-10, 1999.
James M. Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of reression coefficients when some regressors are not always observed. Journal of the American Statistical Association, 89(427): 846-688, 1994.
Paul R. Rosenbaum and Donald B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald B. Rubin. Estimating causal effeets of treatments in randomized and nonrandomized studies. Journal of Educational Psychology, 66(5):688-701, 1974.
14

---

Published as a conference paper at ICLR 2026
Jonas Schweisthal, Dennis Frauen, Valentyn Melnychuk, and Stefan Feuerriegel. Reliable off off-policy learning for dosage combinations. In NeurIPS, 2023.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: General- ization bounds and algorithms. In ICML, 2017a.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: General- ization bounds and algorithms. In International Conference on Machine Learning, 2017b.
Claudia Shi, David M. Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. NeurIPS, 2019.
Rahul Singh, Maneesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. In NeurlPS, 2019.
Jeffrey A Smith and Petra E Todd. Does matching overcome LaLonde's critique of nonexperimental estimators? Journal of Econometries, 125(1-2):305-33) 333335,,, 2005.
Vasilis Syrgkanis, Victor Lei, Miruna Oprescu, Maggie Hei, Keith Battocchi, and Greg Lewis. Machine learning estimation of heterogeneous treatment effects with instruments. In NeurIPS, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
Boris van Breugel and Mihaela van der Schaar. Why tabular foundation models should be a research priority. In ICML., 2024.
Mark van der Laan. Statistical inference for variable importance. The International Journal of Biostatistics, 2(1):1-31, 2006.
Mark J. van der Laan and Donald B. Rubin. Targeted maximum likelihood learning. The International Journal of Biostatisticss 2(1), 2006.
Hal R. Varian. Causal inference in economics and marketing. Proceedings of the National Academy of Sciences (PNAS), (13(2))7311--735) 7315, 2016. doi: 10. 10.10733pnas.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1136622) 2018. doi: 11011/0000662144.2217.....
Janick Weberpals, Stefan Feuerriegel, Mihaela van der Schaar, and Kenneth I.. Kehl. Opportunities for causal machine learning in precision oncology. NEJM AI, 2025.
Liyuan Xu, Yutian Chen, Siddarth Srinivasan, Nando de Freitas, Amaud Doucet, and Arthur Gretton. Learning deep features in instrumental variable regression. In ICLR, 2021.
Jiaqi Zhang, Joel Jennings, Agril Hilmkril, Nick Pawlowski, Cheng Zhang, and Chao Ma. Towards causal foundation model: On duality between optimal balancing and attention. arXiv preprini, arXiv:2310.00809, 2024.
15

---

Published as a conference paper at ICLR 2026
A EXTENDED RELATED WORK
Prior-data-fitted networks (PFNs) tabular foundation models. Foundation models an emerging paradigm that has revolutionized machine learning for various data modalities, particularly language and vision tasks (Devlin 2018 Lahat al., 2024 Touvron et al., 2023b a). The same paradigm is now being explored for tabular ar data a modality that underpins the large majority of networks (PFNs) (Müller et al. 2022) constitute powerful approach to training tabular foundation models. PFNs are large transformers trained on synthetic data to perform Bayesian inference through in- context learning. TabPFN (Hollmann et al.) [023] (025) scaled this idea by pairing the transformer with a Bayesian neural network prior over structural causal models (SCMs) and demonstrating state-of-the-art performance on various tabular benchmarks. Subsequent work extended PFNs to time-series forecasting (Hoo et al. 2025) and analyzed their in-context learning abilities theoretically (Nagler) 2023). Critically, all existing PFNs are trained only for predictive tasks and do not target causal estimands; they therefore are not designed for causal inference of treatment effects, which is the goal of our paper.
analyses in science and business vaan Breugel van der Schaar 2024 Prior-data-
Treatment effect estimation. Causal inference, such as the estimation of average treatment effects. originates from fields like econometries (Imbens & Angrist 1994
Angrist 1990), statistics (van der Laan & Rubin 2006), and epidemiology (Robins) [986] 1994). Machine learning methods have been proposed to estimate heterogeneous effects support personalized decision-making. One line of work are frequentist methods, which often build on semiparametric theory (Robins et al. 1994]
Robins| 1999). yielding model-agnostic estimators that are doubly robust and Neyman-orthogome
(van der Laan 2006; Chernozhukov et al. 2018 Nie & Wager, 2021: Foster & Syrgkanis 2023
Kennedy 2023a). Another line of work builds upon specific machine learning methods/ architectures
such as regression trees (Wager & Athey 2018) or neural networks (Johansson et al. 2016, Shalit
et al. 2017a; (Shi et al. 2019) and adopts them causal inference. Bayesian alternatives include Bayesian additive regression trees (Hahn et al. 2020) or Gaussian- counterfactual regression 2017). However, all of the existing estimators above must be retrained for every new dataset. In contrast, our CausalFM allows for pre-trained models to approximate Bayesian causal inference.
(Alaa & van der Schaar
A.1 DIFFERENCES BETWEEN CAUSALFM, CAUSALPEN, AND DO-PFN
Identifiability. CausalFM separates identifiability from estimation. The central motivation of CausalFM is that causal identification (choosing an identifiable setting such back-door, IV, or front-door) must be handled before estimation. This mirrors classical causal inference practice and ensures that the PFN only learns within an identifiable causal setting. Our reasoning for identifiability is as follows:
(1) Asymptotically unbiased causal inference: As we show in Theorem 4.3
incorporating identifi- ability assumptions into the prior is necessary for asymptotically unbiased causal inference. As consequence, methods that ignore identifiability assumptions yield biased causal effect estimates, even if we collect large amounts of data. This is highly undesirable in practice.
(2) Informative predictive-posterit distributions: If we do not impose any assumption on the DGP, it is well known that causal effect estimation is not just fundamentally biased, but also that this bias can be of arbitrary size. For example, the sackdoor--djjstmees bias due to omitted unobserved confounding can be written in closed form depending on confounding strength. Thus, if the PFN- prior assigns positive probability mass for DGPs with arbitrary confounding strength, the predictive- posterior must respect the possibility of arbitrarily biased treatment effects, thus rendering PaN-dased inference completely noninformative.
(3) Clear separation between domain knowledge and statistical inference: One might argue that possible remedy would be to restrict the PFN prior only to DGPs with somewhat "weak" identifiability violations (e.g., weak unobserved confounding). However, to assumptions/ domain knowledge on the DGP, similar to those in our paper, that must be made transparent for practitioners and could also possibly be violated.
argue that this would correspond
16

---

Published as a conference paper at ICLR 2026
In short, our paper follows established causal inference philosophy and separates identifiability from estimation: the identifiability step (choosing the causal setting) requires careful modeling and usage of domain knowledge, while the estimation step can be handed over to CausalFM. If practitioners suspect identifiability assumptions may be violated, we recommend performing causal sensitivity analysis to assess the extent of potential violations.
In contrast, CausalPFN implicitly assumes only back-door adjustment and therefore cannot handle IV front-door, resulting in bias under unobserved confounding. Do-PEN mixes many causal graphs in a single prior without conditioning on which setting is identifiable, which, as our Theorem 4.3 shows, can lead to asymptotically biased estimates and non-informative posteriors.
Prior construction. This philosophy requires fundamentally different prior construction. Because identifiability is encoded at the level of causal structure, CausalFM introduces C-SCM priors and C-DAGs that enforce the assumptions required by each identifiability strategy. Do-PEN does nor encode identifiability constraints into the prior family for their prior construction. CausalPFN is restricted solely to back-door adjustment.
Theoretical guarantees. Beyond this framework, we contribute new theoretical results showing that identifiability must be incorporated into PFN priors. Theorem + proves that if PFN prior places nonzero mass on SCMs that violate the identifiability conditions of the chosen setting, then the resulting posterior predictive interventional distribution is necessarily misspecified and cannot yield consistent causal effect estimates, even with infinite data. This explains the empirical behavior of Do-PEN, which may return non-informative posteriors when its prior includes SCMs with strong unobserved confounding and no valid instruments. Our theoretical results show that this issue is structural, merely implementation detail.
Empirical performance. Empirically, our model outperforms others in different settings. Besides, also have experiments showing the necessity to have the correct identifiability assumption in the prior specification.
17

---

Published as a conference paper at ICLR 2026
B EXAMPLE FOR SCM-PRIORS
Here, we consider the IV setting from Example 3 with additional normality assumption and empty X = 0, i.e., observational distribution (Z, A, Y) Pobs. Let us consider the following class of SCMs:
U ~~~ V,,, <z ~ (/ 1), ca NO,, 1), cy ~ (O 1), Z = aez,+ + A = 32 geΛ yU, Y = nU Dey,
(10)
(11)
where U is an unobserved confounder between A and Y, and Ez, EA, ey are noise variables, and 3. 7, 8, Ç, 17, and 0 are scalars describing the functional dependences between observed and noise variables. Our causal query is Q(Pint)) = EE Ç.
General approach. The class of SCMs above is compatible with the linear IV setting whenever it holds that 0 (independence assumption from Example 3). Hence, we can specify prior distribution over this class of SCMs by specifying a distribution over (a, 7, 7; 6, Ç, 7), 0) and setting O. Note that this automatically specifies a distribution over Fobe (by sampling from the SCM) and Pint (by intervening and setting A = 1 in the SCM). Interestingly, this addresses the two drawbacks of observational priors from above as follows: (i) During the PFN training we can sample D, ~ Fobs and y1) ~ Pint and thus fit go (v)) D., for the interventional outcome (see Sec. f details). For estimating the causal query we can thus use Q(gn(v(1) Dn)) and do not need access to the potentially unknown Q. (ii) We can directly control the marginal prior distribution of Ç. thus remedying the above drawbacks and allowing us more control to incorporate prior information of our causal query.
Adding identifiability assumptions to the prior. A key question is whether we should actually impose the identifiability assumption ε = 0 when constructing a prior. A different approach would be to also put a prior on K, thus taking account the possibility of identifiability violations in the prior. Such an approach has been proposed by (Robertson al.) (025), where the authors construct a prior over many possible causal inference settings simultaneously. However, as we show in the following, this would make consistent Bayesian estimation of the causal query of interest impossible, confirming Theorem 4.3
Lemma B.1. Let Let = ("""""""""""""" "n",""""" 00" x²= = 0) be an identified ground-trutt truth SCM. Then for any causal larget { (* there exists another SCM S = (,,8,7,5, 5,4,0, x)) with # that induces the same observational distribution as S
Proof. See Appendix C
Lemma B. has an important consequence: if our prior II puts positive probability mass on all possible combinations of (a, 7, S, 11, 0, s).. the corresponding posterior II(- (D.)) will even for + 2o put positive probability mass on any E R, thus being completely non-informative about the causal target quantity. As a consequence, any Bayesian point estimator using such a prior (c.g.. as in under the approach (Robertson et al.) 2025)) will be asymptotically biased.
In contrast, we present different approach to circumvent the above problems:
namely, we propose to construct PFN-priors that incorporate assumptions that allow for identifiability of the causal target quantity (e.g., setting k 0 in the above example). As such, we follow established philosophy in causal inference that separates identifiability and estimation steps (Pearl 2009): the identifiabilityy
step should be established by the practitioner using domain knowledge (e.g., establishing whether a certain variable is valid instrument). Once identifiability has been established. we can use Bayesian modeling and PFN-based models for the estimation step.
Which noise variables to model? A key question that remains is what classes of SCMs we can use to specify priors for the causal inference setting C at hand. Indeed, the class of SCMs is non-unique: as suggested in the main paper, it is not necessary specify both noise variables CY and ey.
(x*, ,",,"",,,0 0², S", (='"" 0*) be a fixed SCM from the above class with Var" (A) 2) > O and Var (Y a) > for all 2, a. Then, there exist unique SCMs S = (a1, $1, n. 81 = 0, 0. 91) and S. = (az, $2, 72, 82. (2, 72, 02 0) that induce the same observational distribution
Lemma B.2. Let
as S* and thus the same causal query Si C Ç* However, whenever holds that both = 0
and 0 = 0, there exists an SCM S* for which ( # C
18

---

Published as a conference paper at ICLR 2026
Proof. See Appendix
C
Lemma B.2 implies that it suffices to specify priors over SCMs without either treatment noise A or outcome noise ey. However, if remove both, there exist interventional distributions for which the prior will never put probability mass on the ground-trut -tuth causal query, rendering Bayesian inference inconsistent. In the following, we generalize this result to arbitrary SCMs and causal inference settings.
19

---

Published as a conference paper at ICLR 2026
C PROOFS
C.I PROOF THEOREM
Proof of Theorem
4.3 Assume, for contradiction, that 11(2) 0. Then, there exists observational
distribution Pabs
P abs such that
(1 TE E 0,
{(Pbss) (S IP obs Pots)
(12)
Let W := E Z denote the non-violating subset of the same observational equivalence class. Pick any Sth E W with II(5") > 0 (such an SCM exists well-specified for C on Pabe). and draw data Dn ~ ppst Pots.
Consider the posterior-predictive interventional functional
Tn = of p.e. II(S) D.) ds).
(13)
Because Q is linear, we have
T, / (P II(S D,) ds.
(14)
Step 1: Split the integral.
Write R == {(obss Then
Q(P II(S D,) dS + / Q II(S) Da) dS. (15)
Tn
E(Poba))
Step 2: Posterior concentrates the observational equivalence class. Under standard Bayesian consistency conditions for the observational model (e.g., KL support of II at Pobs), we have [ D..) O,
DC,
(16)
and, hence, the last integral over R vanishes asymptotically (assuming the usual integrability for Q)-
Step 3: Inside E(Pabs), the posterior is proportional to the prior. For every S E {(Pots),
the observational likelihood is identical because PPS = Pobs- Therefore, for any measurable A C {(Poss) and all n,, yield
[1(A)
II(dS)
II(A Dx) (T()))))s) Combining with Step yields
and IIdSS Dn, 8 E E(Pohs)
(17)
((tt(())))
(((E()))s) I SECP Q (m)) I(ds),
T,
+ Co. 00.
(18)
Step 4: The limit cannot equal the identifiable target. Split the integral over {(Pbs) into W and Z (n (Pore):
[ Je(Pomm)
( IIds) = =/1, Q(PS I[dd)
/
QPL II(dS).
(19)
ZNEE(Pon))
By definition W and identifiability, for every E W we have Q(P int) = ((Pots), and thus
I, (p (I(S)) II(dS) Q(Pors) (I(B)))((())))))
(20)
Hence,
lim T, Q(Pans) IT(W) (Een(PPn))
APs int I(dS)
(21)
I1(WW) + 11(2 n ((Pob))
By the non-cancellation assumption (applied to this Pota) and w > 0, this limit is not equal to ((Pata).
Finally, since D, ~ " and C is identifiable,
(P Q(Pots).
(22)
Thus lim,- T, * Q(PS") ), contradicting that II is well-specified for Therefore, it must be that IT 0.
20

---

Published as a conference paper at ICLR 2026
C.2 PROOF OF LEMMA I.I (LINEAR IV)
Proof of Lemma B.7 We prove that for any Ç * S", there exists an SCM S = (a, (,, 0) that induces the same observational distribution S
Step 1:
Observational distribution of S"
The observational distribution is characterized by the covariance matrix E" of (Z",'", Y*):
Var(Z*) (a')²
(23)
Cov(Z", (A") (a")88"
(24)
Var(A*) (a' (a*3*)² + - (7*)²
(25)
Cov(Z", y*) = 5"(a")
(26)
Cov(A", Y*) - <"((a" + (8*) + (1*) +"> Var(Y*) = =<""<(a²>>>>> 3-)² (
(27)
(0))²
(28)
Step 2: Construction of alternative SCM S
The covariance matrix Σ has elements:
Var Z) = a² + x²
(29)
Cov(Z, A) = a²3 + n(133 + 7)
(30)
Var(A) = (a.)² + (MB + 7)²+ +
(31)
Cov(Z, Y) =<(a²8 + + + +)) +7p
(32)
Cov(A, <((a:)) + (xB 5*) +7)
(33)
Var(Y)
+ (NB
2 +7) +1² +
(34)
Step Parameter matching
To achieve Σ = 2', we need:
a² + (a²)2
(35)
a² +7) (a*)²8*
(36)
8² (a" 30 (8*)² (7*)²
(37)
+1)) +
qp (a²))
(38)
<[(a33) + (k8 + +) +8 + n +7) = =-'((a" 8012 + (8*)2 + (7*)) 17"
77"
(39)
(a)(3) + (kB + 7) 5$1 + 25n(KB + +7) + 1² 0² (Var(Y*))
(40)
Step Solution construction
We choose k / 0 such that |x| < la*1. We set
a = 2² k²²
(41)
223)
0 a² +x² = 8
(from 35 and (6)).
(42)
8
(43)
K,3 + = + ( (a* 8-)2 + (a.3)
(from 37
(44)
Since a.3 = =a3* a 3*, we have.
a = (a+jz (a (== = (a²)² x²a
a
(a*)² (a 3
(45)
21

---

Published as a conference paper at ICLR 2026
Therefore, have
R²
N3 + - 1090//008
+
(a+)2 (a )3)
(46)
From Eq. (38) and Eq.
39
we can solve for 7) via
- C(a2 3"a"))²0-"
N(NB + +7))
(47)
K
Finally, 0 is determined from Eq. (40)
Step 5: Existence Verification
The system has 8 parameters (a, 3, 7, 6, C, k.) and 6 constraints (the 6 unique entries of the covariance matrix). Since * < is fixed and & chosen, we have 6 remaining parameters for 6 constraints. The key observation is that the introduction of confounding (k * 0) creates additional correlation structures that can compensate for the change in the causal effect Ç, allowing the observational distribution to remain unchanged.
Proof of Lemma B.2 Let S (a, 6, 0) be any SCM from the linear IV class in Eq. (T0).
Then, the following coefficients are identified via observational data alone:
Var((2) Z),
(48a)
E[y = =1],
(48b)
3 E|S(B + dEA)] E((((Ee)]
E| Z = 1] B Z
Ç B
(48c)
E = 1] B[A = 0]
as well as the combination of coefficients
8² 7- Var(A Z), and
172 + 82 Var(Y AA)
(49)
and the back-door adjustment
B| A = E
A = 0] = ( + (BIU A = 1| - EIU A= 0))
(50)
117
= = 7.2 + 52 + (Ba))
(51)
Note that the back-door adjustment is biased for ( due unobserved confounding.
Noiseless treatment case. Let now S = (a", 7" 8t, S ;,"""0 0*) denote an arbitrary fixed SCM
from the linear IV class. We start by constructing S (01, 31, 71, 81 = 0, S1, 71, 01) such that pSi P
obs. Because of Eq. 48) we set
a", B1 3", S1 = S*
(52)
Furthermore, setting 8, implics due to Eq. 44 that
フィ = 8²2__ + 7x
(53)
Due to Eq. 50 it must holds that
7171
8-2 + 7x2 7/
(54)
71 + (B1r))
(Baaa)3''
which implies that
/p 11 22
(55)
8+2 7,22"
Finally, due to Eq. 49 we yield
2
0 = 17 + 8
(56)
8-2 200
22

---

Published as a conference paper at ICLR 2026
which means that every parameter of S₁ has unique solution in terms of parameters of S under the constraints preserving the observational distribution.
Noiseless outcome case. We now construct S2 (ar2, B2, 72, 82, C2, 12, 02 = 0) such that Pos p.s obs Again, Eq. (48) implies that
a2 a², B2 B <2 5
(57)
and setting 02 implies due to Eq.
49 that
n = 1) g0²
(58)
Due to Eq. (50). it must holds that 3- 45 + [5o]²
17
which implies that
8224 + (0
25 11 702 x2
(59)
11 0=2
Finally, due to Eq. 49 we have
17"
n7
8= 17
t,²
+
2
(60)
17"
G=
which means that every parameter of S has unique solution in terms of parameters of S under the constraints of preserving the observational distribution.
23

---

Published as a conference paper at ICLR 2026
D IMPLEMENTATION DETAILS
D.1 IMPLEMENTATION DETAILS OF DATA PRIOR
Data Prior. (i) For each covariate cluster C, containing latent nodes, we sample a random MLP-style graph over pa (C,) by drawing biases and edge-weights from Lc, and then pruning edges at random to ensure acyclicity. We evaluate this graph with tanh activations and noise (from normal, uniform, Laplace, or logistic distribution) to produce continuous features, then apply randomized thresholds to discretize or binarize a subset, yielding mixed-type covariates via our unstructured BNN prior. (ii) For treatment (and outcome) clusters C of purely observed nodes, we instantiate second BNN f()) over pa (C₃) (with A ~ Iec, and the same acyclicity constraint). We forward proopaate the
covariates through f6d with injected noise to compute a scalar propensity score, then threshold to assign a binary treatment. We forward-propagate both covariates and treatment to obtain potential outcomes. The resulting treatment (and outcome) are sampled from our structured BNN prior.
We sample covariates from DAG-structured SCM by drawing random MLP-like directed graph and assigning each node a bias, edge weights sampled from prior distributions. The resulting ML.P-lik graph is transformed into DAG by randomly dropping edges, and structural equations with tanh activations and heterogeneous noise distributions (normal, uniform, Laplace, or logistic) generate continuous features. Then we apply randomized feature transformation that discretizes some features and binarizes others, yielding mixed-type covariates. Next, we assign binary treatments via separate randomly instantiated MLP and forward-propagate each covariate with injected noise to compute a propensity score.
Input format. CausalFM operates as in-context learner like other foundation models, meaning it approximates Bayesian inference by conditioning on a dataset provided in its context window. Therefore, it requires an entire dataset make predictions for specific query samples. (1) Input structure: The model accepts dataset D, = {(T;, ag, yo) acting as the context (or support set) and query point Tquery (or batch of query points). (2) Mechanism: The transformer processes the entire sequence of observed data D, using self-attention to extract context-dependee representations. then outputs the posterior predictive distribution for the causal quantity (e.g., CATE given the context D, and the specific query Tiquery (3) Comparison to fine-tuning baselines: In practice, standard baselines require an explicit training phase on a training set before evaluation. In contrast, C takes the "training" data the input context (support set) and directly generates predictions for the test data (query set) in a single forward pass.
D.2 IMPLEMENTATION DETAILS OF OUR METHOD
We encode observational data tokens, and the embedded tokens are then processed through transformer where attention is applied between the observations. We use transformer- an encoder to extract task- or context-dependent representation from input data. This representation is then passed to a Gaussian mixture model (GMM) head, which predicts the parameters of a GMM, including mixture weights, means, and standard deviations. The model outputs mixture distribution over the target variable, and is trained end-to-end using the negative log-likelihood (NLL) of the observed targets under the predicted GMM. This enables uncertainty-aware and multi-modal predictions while leveraging the few-shot generalization capabilities of our model.
PFN as
We instantiate a featfeeture transformer tailored to CATE estimation. For a mini-batch with sequence length S = Suup + Squary (query set followed by support set). Confounders X E R treatment A E S and factual outcomes Y E Sx are encoded as tokens. To prevent label leakage, we split at Ssupp = [0.8 S] and set A, = NaN and Y, = NaN for t 2 Saugo (on the query set). The model thus observes (X, A. on support steps and learn to infer CATES for the query set from X only.
The X stream uses a feature encoder, while A and Y pass through NaN-indicator handler followed by feature projections. We concatenate the three streams along the token axis to obtain Ho E @xx add a feature-toker positional embedding, and process H, with L transformer encoder blocks (self-attention only). We pool over tokens to produce feature Z R BxSxF After lightweight MLP maps Z to a scalar, a 1D K Kc-oomoneet GMM head outputs mixture parameters (IT, p1, o) via ooftmax(W=2/))) v/T=T)),, W Wpz,0 a softplus(W,,.. (W.=) + e, for each E E Our
24

---

Published as a conference paper at ICLR 2026
training loss is the Gauxsian-mixture mixture negative log-likelihood (GMM-NLL). We thus obtain the distribution
K
p(T x))
Σ
Tx(x) (Ma(2),
(61)
And the CATE can be computed through
K
+(x) = BlT 1
Σ k=1
TA (x))xxx)))
(62)
We use embedding size E = 128, Phends = 4, feed-forward dimension 4E, L = 10 encoder layers. GELU activations, and feature grouping size
K = 5, lemperature 7 1.0, and variance floor pper-feature tokens). For the GMM head we set 10-3 We train with Adam (learning rate 10-3 weight decay 10-5), batch size 16, and up to 150 epochs. We use early stopping on validation loss. Empirically, the total training time for causalFM model is about 24 hours
NVIDIA A100 GPU.
We implement our CausalFM using PyTorch. Our model implementation builds upon the TabPFN architecture (Hollmann et al. 2023) from htttp:://gtttbb.cm///// wtwp:ttuub..omm///..iaaass.. TTbPFN/ tree/main
D.3
IMPLEMENTATION DETAILS OF BASELINES
For the standard CATE setting baselines, we follow the implementation from httpt:://tttttttt
Som/AliciaCurta CATENets/Iree wwarre//maii main for most of the CATE estimators, including S-
learner (Künzel al. 2019), T-leamer (Künzel al. 2019), TARNet (Shalit et al. 2017b). X-
leaner (Künzel et al.) 2019). DR-leamer (Kennedy (023b), RA-learner (Curth & van der Schaar
2021}. For the foundation model baselines, we follow the author implementation from ht
//github. com/ vdb 1m/ Caw.aaaauauaaaaaaaaaaaa /m in for CausalPFN (Balazadeh al. 2025 ) we
follow the author implementation from http::///gttu...c.oooooo/// htt.t/uww..co///211///------- (Robertson et al. 2025).
for DoPFN
For the IV setting, we follow the implementation from https: htttp:://gitub.coo com/ Denni rawueen/ MRIV Net /teettee/// pare/mmann.. www.mmaimaaeaaaa / for the most of the IV methods, in- cluding KIV (Singh et al. 2019) DFIV (Xu et al. 2021), DeepIV (Hartford et al. 2017). DeepGMM (Bennett et al. 2019), DMLIV (Syrgkanis et 2019)) For each dataset and method, evaluated repetitions, each with different random seed. All methods used the same train-test split.
25

---

Published as a conference paper at ICLR 2026
E SYNTHETIC DATA GENERATION FOR THE STANDARD CATE ESTIMATION SETTING
We construct the standard CATE estimation datasets by sampling covariates X, treatment A, and continuous outcomes Y. The design induces rich nonlinearity while preserving strong ignorability (A TT {Y (O)),, X).
E.1 COVARIATES VIA A DAG- SAG-STRUCTURED SCM
We first sample layered directed graph (an MLP-like DAG), then evaluate structural causal model (SCM) on its nodes and expose a random subset as observed features.
Graph. Sample number of layers Lx and hidden size Hx from simple discrete priors (see "Hyper- parameters" below). Build layered graph with H nodes per layer and fully connect layer & to e+1. Randomly drop fraction Phoo of inter-layer edges to sparsify while keeping acyclicity.
Node equations and noise. For each node j. sample weights {wy+ kEpo(j)) bias by. and an exogenous noise distribution Ej Dj. where D is drawn from meta-prior over (Normal, Uniform, Laplace, Logistic] with random scale. Nodes are evaluated in topological order:
Wikk T.k
bj + Ej.
Tj = Lanh(s)),
(63)
kEPa())
with the convention Exeo (.) for roots. Let Ux {Ei} denote the collection of all node noises.
Observed features. Sample feature index set F C V with [F] = uniformly from all graph nodes. A single observation E Rd is obtained by re- sampling Ux evaluating (63 over the DAG, and reading out X = (1y))EF. Each sample uses independent Ux.
Feature typing and transformations (Optional). Each selected feature zj is assigned a random type from (continuous, binary, categorical). Continuous features are kept in their raw form zj e (-1,1). Binary features are obtained by mapping Vy through a logistic function and drawing a Bernoulli sample. For categorical features, we first sample base distribution x E K-1 over K categories from Dirichlet prior. To make the distribution depend the DAG value Zj. we introduce fixed direction vector IR* (normalized) and scale 0, and form
Tιæ;) = sofumax(log n + 0r v)-
(64)
The observed categorical feature is then sampled as X; ~ Categorical(T((;)))
E.2 TREATMENT ASSIGNMENT
Given X, we compute stochastic logit via a feed-forward network with layer-wise exogenous noise and then sample a Bernoulli treatment A fa(X, UA).
Network. Sample depth T. 3 and hidden width HA. Let (0) X E IRd be the input layer. For hidden layers .,.... IA 1,
( = (e0n(e-1) + b(0) + E00 h(0) = (s0).
(65)
and the (scalar) output logit
8A = T (1a-1) + b e(ha).
(66)
We define the propensity p p= (Ss)) and sample
A ~ Bernoullip).
(67)
Let UA = (E( used for the Bernoulli sampling.
- Un) collect all exogenous noises of the network and the random variable Un
26

---

Published as a conference paper at ICLR 2026
E.3
CONTINUOUS OUTCOME
For each unit, we compute the potential outcomes (0) and (1) using the same exogenous noise Uy.
Network. Sample depth Ly 3 and width Hy; optionally drop a fraction Pang of hidden edges to induce sparsity. For given treatment level {0, 1}, the input is X and A, then for hidden layers
(a) = v(o))(-))(a)) + e + g(
(a) tanh (anh(t!9)p (t
(68)
with (a) [X, a|. and the scalar output logit
Y(a) v h(hr (a)
glvr)
(69)
The factual outcome
== (1-
(70)
Let Uy = 3 denote outcome-network noises; the same Uy is reused when constructing (()
and for the same unit.
E,4 INDEPENDENCE AND IDENTIFICATION
All exogenous noises sampled independently across mechanisms and samples: Ux 1 UA 1 Uy and i.i.d. across units. Hence strong ignorability holds:
A IL (O))) X, 0 Pr(A=1 X) < 1,
(71)
with overlap ensured by the sigmoid in (66) to 67
E.5 HYPERPARAMETERS AND PRIORS (AS USED OUR CODE)
We simple, reproducible priors for architecture, weights, and noises:
Covariate DAG: Lx ~ (3,4,5,, Hx ~ (15, 40), edge-drop Parop -0.5.
Treatment net: LA Unif 4J. HA Unif(8..... 20}.
Outcome net: Ly ~ (3 4,5,, Hl Unif 25), edge-drop parop =0.4.
Weights/biases: i.i.d. w, (0, with task-specific Tw.
Node noises: for each node, draw type in (Normal, Uniform, Laplace. Logistic) and a scale from a wide range; sample fresh noises per unit and layer as in (3). (65) (65)-(66)
68 69
Activation: tanh for all hidden layers; output layers are linear (logits).
Features observed: choose F uniformly at random from all DAG nodes, |F| = d.
E.6 GENERATION PIPELINE
For each dataset:
1. Sample the covariate DAG, parameters, and noises; for each unit, evaluate the DAG in topological order to obtain by reading nodes in
2. Given X, construct the treatment network with UA get p and sample A ~ (Brnouuli(p).
3. For outcomes, sample Uy once per unit and use it to compute Y(0) and Y(1) via (69
E.7 SYNTHETIC DATASETS SIZE
We sample 10000 synthetic training datasets from data prior with different data generation mechanism. Each training datasets contain 1024 data samples. The feature dimensions are also different across the datasets, ranging from 10 to 100. The features are mixed data type with continuous, binary and categorical.
27

---

Published as a conference paper at ICLR 2026
F SYNTHETIC DATA GENERATION FOR THE INSTRUMENTAL VARIABLES (IV) SETTING
We aim at estimating CATES from observational data under unobserved confounding using IVs. In contrast to the standard CATE setting, where strong unconfoundedness holds, our IV datasets intentionally violate unconfoundedness by introducing an unobserved confounder U that affects both treatment A and outcome Y. Identification is instead driven by an instrument Z that (i) is relevant for A, (ii) has no direct path to Y beyond A (exclusion), and (iii) is conditionally independent of U given X.
Key differences standard CATE. (i) Ignorability is broken: A (Y(0), x due to + A and U + Y. (ii) We introduce an instrument with Z 1 U X, A X, and no Z - Y edge (exclusion). (iii) Outcomes generated via additive structural form Y = f(X, A) + 9(X, U) + Ły with and deterministic neural networks; the same By is reused across (()) and (1) for unit ensure counterfactual consistency.
F.1 COVARIATES AND LATENT CONFOUNDERS VIA A DAG-STRUCTURED SCM
We reuse the DAG-SCM from the standard setting to produce a wide set of base variables W, then split it into observed covariates X and unobserved confounders U. Thus we have different strength of the unobserved confounders from weak to sufficiently strong.
Graph and node equations. Sample number of layers Lx and hidden size Hx. build a layered DAG (fully connect layer l to E+1), and drop fraction Parop of inter-layer edges sparsify. For each node j. sample weights {wys hepa(f)) bias bj. and node-specifts exogenous noise Ej Dj (type and scale drawn once per node). Evaluate in topological order
8y =
Σ kεpa())
Wikk Up + by + Eg.
By = tanh(s)).
(72)
Draw a feature index set for W (uy) with |W| = dx + dj and then sample the actual confounder dimension du E {2..... 5} uniformly. Split U E Rd from the first dy coordinates of W, X e gdx from the next dx coordinates. Node noises {5,} are drawn independently per unit.
F.2 INSTRUMENT VARIABLE
We generate Z from X only, ensuring T. U X by construction and precluding any direct - Z path. Let Φz be feed-forward network with input X and no layer-wise exogenous noise; the network parameters are sampled once per dataset and then fixed. For a unit with covariates X,
(ernoull((((sz))
binary instrument,
= oz(x),
(73)
82,
continuous instrument.
We randomly choose between the binary and continuous variants when creating datasets. Relevance induced via the Z path in the treatment mechanism below.
Note that the instrument variable Z has a direct influence on the treatment A, but does not have direct effect on the outcome Y.
F.3 TREATMENT VARIABLE
Given (X, Z, U), treatment is generated via a deterministic network A followed by a Bernoulli draw. There is no layer-wise noise inside d the only randomness is the terminal Bernoulli. For unit,
SA A ([x Z;U),,
p a(sA), A Bernoullip).
(74)
This introduces U + and hence breaks ignorability, while maintaining $ U X and Z + A relevance.
28

---

Published as a conference paper at ICLR 2026
F.4 OUTCOME VARIABLES
The instrument variable Z has no direct effect on the outcomes. Outcomes are generated additively from a treatment channel and confounding channel g. both deterministic MLPs with inputs [X; and [X; UU) respectively. Let Ey ~ (N be an i.i.d. scalar noise drawn once per unit,
(() a(X,a) + 9(X EY,
a [0,1),
(75)
Y (1 (1- A)--A)YC (0).
(76)
By construction there is no Z + Y edge (exclusion), since Z influences Y only through A.
F.5 INDEPENDENCE AND IDENTIFICATION (IV)
All exogenous noises are sampled independently across units and mechanisms. The IV conditions hold by construction,
(Independence) Z L U J (Exclusion) (() depends on a, and ey only (no 2),
(77)
(78)
(Relevance) Z M. A A| X.
(79)
F.6 HYPERPARAMETERS AND PRIORS
We use simple priors mirroring our implementation:
DAG-SCM for (X,U): Lx ~ 3,,442 Hx ( .0)], edge-drop Purop 00.4; node noises Ej draw type in (Normal, Uniform, Laplace, Logistic) with random scale.
Instrument net oz: depth I. 2 3, width Hz ~ Unif (a..... :0]: output either Bernoulli with (8z) (binary 7) or real-valued sz (continuous Z); layer-wise noise.
Treatment net @A: depth LA 2 3, width HA Unif ... 30}; layer-wise noise;
Bernoulli (s))).
Outcome nets f.g: depths Ty,, T.g
(...,, 6), widths Hy, Hg Unif 25}: Sy NN(0, o.-) with ay 0.5 by default.
Weights/biases: i.idd w, b ~ NO, 1) sampled once per dataset; tanh activations.
Strength: dy ~ (nif (2..(2.....3)
F.7 GENERATION PIPELINE (IV)
For each dataset, we execute:
Sample the covariate DAG and parameters; for each unit, evaluate 72 to obtain wide matrix then split it into (U, X).
2. Given X, compute the instrument Z via 73
(binary or continuous and mixed).
3. Given (X, Z, U). compute the treatment propensity p = x(sA) via 74 and sample ~ Bernoullipp)
4. Draw a single per unit and compute Y()), ((1) using 75 set the factual outcome by (76).
This yields datasets matching the classical IV graph and enabling evaluation of IV estimators.
29

---

Published as a conference paper at ICLR 2026
G SYNTHETIC DATA GENERATION FOR THE FRONT-DOOR-ADIUSTER SETTING
G.1 FRONT-DOOR ADJUSTMENT DATASETS
We next construct datasets satisfying the front-door criterion. Besides covariates X, treatment A, and continuous outcomes Y, we introduce mediator M. The design ensures that A affects only through M (no direct A-+Y path), U (unobserved) confounds A and Y but does not affect M.
Covariates via a DAG-structured SCM. DAG, draw node-wir wise weights/biases/noise, evaluate in topological order as in (63), and expose d node values as observed features E Rd. Independent exogenous noises Ux x (e;) are re-sampled unit.
Identical the standard setting: we sample a layered
Latent confounders. From the same SCM evaluation also retain g additional node values as unobserved confounders E 129 (not revealed to learners). These induce confounding between A and Y.
G.1.I TREATMENT ASSIGNMENT WITH LATENT CONFOUNDING
Given (X, U), sample ford-forward network and generate treatment. Let LA and HA be the depth and width, respectively. With (0)) = =[X
((9)((11) + b1 h(
tanh(s f= 1=....., LA-
(80)
and scalar logit
A = wT + ba
p = o(AA),
Bernoull(pp)
(81)
G.1.2 MEDIATOR MECHANISM
The mediator is generated from (X, A) only, thereby enforcing the front-door exclusion U Let L 3,, HM be depth and width, with input g( A).
M.
(t) g +
(2) M M
g(0) = tanh (r(0)), 8= 1...., Las - 1,
(82)
and scalar output
M (Tag( + bas + eM (Las)
(83)
We denote UM = (8) LM Te-1
G.1.3 OUTCOME VARIABLE
Outcomes are constructed to satisfy + M - Y as the only causal path from A to Y, while allowing U -> Y and X Y. We decompose Y into M-path component and confounding
component:
Mediator path:
1 ((0 + (0) + gl0) 1 = [X, MJ. 10 = tanh(ry)) (8) ) TY R(X, = (br + c + g(br))
(84)
Confounding path:
G(X,U) = (T + + m(عا) 7000 = [x, UJU 7,00 = tanh(.)
(85)
and define the potential outcomes
(a) = R(X, M(a)) + G(X,U) + ey,
M(a) computed from 82
83 with A=a. (86)
The factual outcome is Y = (1-+) (0). By construction there is no direct A 4 edge; A influences Y solely via M.
30

---

Published as a conference paper at ICLR 2026
G.1.4 HYPERPARAMETERS AND PRIORS
Covariate DAG: Lx ~ Unif{3, 4, 5,5,6 Hx ~ , edge-drop P
node noises drawn per-node from (Normal, Uniform, Laplace, Logistic) with random scale.
Treatment net (Fq. 80
8T
La [3, 4), HA (nif . 203.
Mediator net (Eq. 8
LM [3, 4). HM Unif(8..... 203.
Outcome nets (Eq.
186 y: Ly, L. ~ Unif 4,5}, widths ~ Unif (10
additive Gaussian ey with task-specific scale.
Weights/biases: i.i.d. (0, d tanh nonlinearity.
G.1.5 GENERATION PIPELINE
For each dataset:
1. Sample the covariate DAG and evaluate obtain (X,U) (observed X, hidden U).
2. Compute p(A=1 XU) via 80
81 and sample A.
3. Evaluate the mediator M from (X, A) using
83
4. Sample Uy once per unit and compute (0) and Y (1) via 86 by first obtaining M (0) and M (1) from the mediator net; set Y (1-A) (0).
31

---

Published as a conference paper at ICLR 2026
H ADDITIONAL EXPERIMENTS
H.1 EVALUATION IN THE FRONT-DOOR ADJUSTMENT SETTING
H.1.1 BASELINES FOR FRONT-DOOR ADJUSTMENT SETTING
In contrast to the standard CATE IV settings, there are few established baselines for the front- door case. Identification in this setting is enabled through Pearl's front-door formula (Pearl 2009). The natural baseline therefore the plug-in front-door learner, which estimates the necessary nuisance components, i.c. P(M X),, P(A X)), and E|Y X] and substitutes them into the identification formula to recover causal quantities. To assess the role model flexibility in estimating these nuisance functions, implement the plug-in leamnr with different regression methods, including linear regression, Random Forests, and neural networks.
H.1.2 RESULTS FOR FRONT-DOOR ADJUSTMENT SETTING
Table reportss the averaged PEHE across datasets. We observe that CausalFM achieves competitive CATE estimation. Importantly, these results hold without requiring model retraining for our model, demonstrating the adaptability of our approach to the front-door setting.
H.2 ADDITIONAL RESULTS ON THE STANDARD CATE ESTIMATION
We report the detailed standard CATE estimation 10 synthetic datasets in Table We show our method gives the best estimation on most of the datasets.
Table Standard CATE estimation on 10 synthetic datasets.
Method
Dj
D2
Ds
D,
D,
Ds
Dy
Ds
Da
Dio
RASELINES (A): STANDARD CATE ESTIMATORS
S-learner Künzel al.
2019
0.725
0.583
0.752
0.829
0.614
0.892
0.858
0.421
0.680
0.985
Telearner
Kunzel al
2019
0.652
0.496
0.666
0.746
0.552
0.849
0.761
0.357
0.608
0.931
TARNet Shalit
201
0.769
0.779
0.817
0.984
0.640
0.938
1.405
0.505
0.736
0.968
RA learner
(C
(Curlh & Van der Schaar 2021
0.620
0.421
0.644
0.706
0.523
0.808
0.646
0.353
0.613
0.759
X-learner Künzel et al 9
0.574
0.400
0.614
0.634
0.381
0.713
0.686
0.302
0.549
0.779
DR leamer Kennedy 20/55
0.783
0.533
0.767
0.947
0.867
0.882
0.791
0.4230
0.653
0.998
BASELINES (B): FOUNDATION MODELS-BASED METHODS
NausalPEN (Balazadeh al 2025
0.493
0.489
0.585
0.743
0.413
0.615
0.950
0.288
0.453
0.544
DoPFN Robertson et al. 2025
0.417
0.313
0.228
0.679
0.591
0.475
0.497
0.551
0.610
0.827
M (ours)
0.454
0.487
0.515
0.677
0.204
0.618
0.950
0.278
0.442
0.532
Reported: PEHE (Lower = better, best in bold).
H.3 RESULTS ON OTHER DATASETS
In the following, we present detailed results of the experiments with ACIC 2016 datasets.
Table 6: Standard CATE estimation on ACIC2016 datasets. Reported: PEHE (mean + std.)
We follow CausalPFN Balazadeh et al. (025))
Method
obtaining data from https://githhs http://conedssAA//caasaaaaaaaaa
BASELINES (A): STANDARD CATE ESTIMATORS
S-Jearner
Könzel et al. 2019
1.191
master sestter/cauallli/daaaeeeee
T-learnr
1.143
TARNei Stg.it
701770
0.934
Srilc_challenge.. 2 to evaluate 10 different datasets with various data generation mechanism. The treatment and outcome were simulated from real-world data corresponding 4802 individuals and 58 covariates. Table 6 shows the results of the CATE estimation.
RA-laamer
Van ucr Schaar 2021
0.762 2: 6.4
X-leamer
Runzel
2019
0.519
K2nizuy 6372507
1.485
BASELINES (B): FOUNDATION MODELS-BASED BASED METHODS
CausaIPNN (Balazadeh 2025
0.239
DaPEN (Rabensoo e-mu 201a)
0.857 ems
CausalFM (ours)
0.638
Lower beteer (best in bold)
32

---

Published as a conference paper at ICLR 2026
H.4 ADDITIONAL RESULTS FOR THE IV SETTING
Table 7: IV setting for CATE estimation with binary instrument variable reported with PEHE. Results for benchmarking model performance across 10 different datasets under various confounding strength.
Method
D1
Dg
D2
Do
Dg
Dg
Dio
BASELINES (A): STANDARD CATE ESTIMATORS
TARNes (Shalit et al
7h
0.789
0.790
0.799
0.789
0.831
0.673
0.582
0.978
0.735
0.642
DR-lcamer Keenedy
20236
1.517
1.071
1.022
0.901
0.754
0.676
0.646
1,009
0.664
0.781
BASELINES (B); STANDARD IV ESTIMATORS
KIV (inggh et all 2019
0,660
0.344
0.340
0.394
0.544
0.460
0.299
0.731
0.532
0.241
DFTV cuu A
0.654
0.245
1,022
0.459
1.145
0.770
0.741
0.366
0.971
0.717
DeepIV (orator et 2017
0.614
0.300
0.310
0.372
0.514
0.404
0.309
0.706
0.510
0.235
DeepGMMI [Seennel
20119
0.704
0.403
0.440
0.599
0.569
0.486
0.292
0.737
0.566
0.232
DMLIV (Syigkanis et al.
2019
0.712
0.379
0.361
0.433
0.548
0.450
0.293
0.722
0.549
0.344
DRIV (Syyecans et
0.869
0.470
0.353
0.368
0.565
0.272
0.715
0.587
0.667
MRIV Fraieen & Pouerrigge 2022
0.759
0.632
0.698
1.011
0.348
0.860
0.929
0.707
0.562
0,380
BASELINES (C): FOUNDATION MODEL BASED
DoPEN Robeetsonn al. 2025
0.776
0.265
0.370
0.382
0.552
0.819
0.499
0.794
0.534
0.242
CausaIFM (ours)
0,586
0.224
0.374
0.310
0.543
0.464
0.250
0.701
0.553
0,217
Reported: PETTE (mean standard deviation.) Lower = better (best in hold)..
Table 8: IV setting for CATE estimation with continuous instrument variable reported with PEHE. Results for benchmarking model performance across 10 different datasets under various confounding strength.
Method
Dj
D2
Ds
Da
D;
D6
D7
D3
Dy
D1u
BASELINES (A): STANDARD CATE ESTIMATORS
TARNet (Shalit et al.
2017b
0.943
0.825
1.025
0.458
1.007
1.316
0.848
1.004
0.825
0.884
DR RR-Jeemner Keemady
20250
1.038
1.055
0.946
0.533
0,955
1.071
1.109
1.502
1.258
0.888
BASELINES (B); STANDARD IV ESTIMATORS
KIV (inch al. 9019
0.509
01.677
0.178
01533
0.948
0.420
0.506
DFIV Aur 2011
0.526
0.574
0.691
0.171
0.532
0.991
0.428
0.800
0.609
0.506
DeepIV (Harford c
2017
0.484
0.539
0.664
0.169
0.506
0.901
0.399
0.770
0.572
0.481
DeepCiMMt [Hennel
2019
0.543
01.811
0.165
1.035
0.437
0.789
0.615
0.505
DMLIY (Syrpkams al
2019
0.518
0.642
0,701
0.181
0.600
1.009
0.574
0.813
0.611
0.537
DRIV (yrpkams et al.
2019
0.633
0.705
0.870
0.279
0.663
0.873
0.523
1.009
0.749
0.630
MRIV Franen Feuernger 2022
0.579
0.631
0.760
0.189
0.586
1.091
0.471
0.880
0.669
0.556
BASELINES (C): FOUNDATION MODEL-BASED
DoPFN Robertson al. 2025
0.471
0.528
0.787
0.322
0.649
1.723
0.416
0.588
0.722
0.545
C (ours)
0.515
0.600
0.704
0.152
0.538
0.934
0.414
0.826
0.600
0.509
Reported: PEHE (mean standard deviation.)
(best bold).
H.5 ADDITIONAL RESULTS FOR MISSPECIFICATION AND PRIOR DESIGN CHOICES
Misspecification of causal settings. We conduct experiments study the sensitivity to using an incorrect identifiability strat- egy. Specifically, we generate data from (i) an IV SCM and (ii) front-door SCM (as in our main experiments), and compare model trained under the correct identifiability design (IV or front-door, respectively) with model trained under an incor- rect back-door design. The results are in Table g As expected, using a misspecified identifiability strategy consistently wors- the PEHE. This highlights the importance of including the correct identifiability assumption into the prior specification, which we propose for CausalFM.
Figure 2: Robustness of our model to difference choices of the prior.
Prior design choices. We also analyze the robustness of our model to different choices of the prior. In particular, we vary the strength of unobserved confounding in the data-ged mentaaggneaatngg process, controlled by a parameter a E [0, 1]. The results in Fig. show that model remains robust increases. This again confirms the strong performance CausalFM.
33

---

Published as a conference paper at ICLR 2026
H.6 COMPUTATIONAL TIME
We report computation time in this section. For method and other foundation-modellbasse approaches, we show inference time since these models do need fine-tuning after pretraining For standard baselines, which must be trained for each dataset separately, we report the average total time per dataset, including both training and inference. As shown in Tables I0 and our model is highly efficient.
Table 9: Analysis of the effect of misspecified identifiability strategies.
Data Generating SCM
Stratogy
Identifiaillity Used
PEHE
Correct
IV
0.422
SCM
doer
Correct
Front-dae Back-doce
0.847
door
Incorrect
0.876
Table 10: Overall time comparison for standard CATE setting.
Method
Time (4)
BASEI INFS (A): STANDARD CATE ESTIMATORS
S Selamerr
Klinzel et al.
2019
2.76 x 10"
T-leamer
Künzel al
2019
x 10
TARNet
Shait
2017
3.98 x 10P
DR-leamer
Kenneuy 20236
1.78 x
RA-Icamer
der Schaar
1.21 10'
X- X-leamer
Kunzel et al.
2019
1.93 x 10²
BASEL INTS (By: FOUNDATION MODELS BASED METHODS
CausalPFNN (Balazadeh 2025
10"
DoPFN Roberison et at.
2125
2.31 x 102
Ca (ours)
4.90 x 10-1
Lower better. Reported: Time in seconds.
Table 11: Overall time comparison for IV setting.
Method
Time
BASELINES (A): STANDARD IV ESTIMATORS
st
3.24
DRIV (Syrekanis al
2019
3.87 10"
DeepIV [Hardord et ml
2010
1.27
Dec
201
1,85 10"
DMLIV Kyrg et
2019
1.85 x 10*
10'
MRIV miteen & Feuerriegel
2022
1.56 x 10"
BASEL INES (B): FOUNDATION MODELS BASED METHODS
DoPIN Rohertson et al. 2025
6.53 x 10"
CausalFM (ours)
4.72 x 10
Lower better, Reported; Time in scconds.
34


Reviews:
[{"review_id": "Jg07ZbKsDj", "reviewer": "Reviewer_r9r3", "summary": "This paper proposes CausalFM, a new framework for training foundation models (transformers) to perform causal inference via prior-data fitting. The authors formalize how to construct Bayesian prior distributions based on structural causal models (SCMs) for various causal inference tasks, and derive necessary criteria to ensure these priors lead to valid, identifiable causal inferences. Using this theory, they propose a novel family of SCM-based priors implemented through Bayesian neural networks, which enables the PFN (Prior-Data Fitted Network) to carry out Bayesian causal effect estimation in-context for settings including back-door (unconfounded observational studies), front-door (mediator adjustment), and instrumental variable (IV) scenarios. The paper then instantiates CausalFM by training transformer models on synthetic data generated from those SCM-based priors, effectively teaching the model to infer causal effects (e.g. Conditional Average Treatment Effect, CATE) without retraining on new datasets. Empirically, CausalFM achieves competitive or superior performance compared to both standard specialized estimators and prior foundation-model baselines, across diverse benchmarks for CATE (unconfounded), IV, and front-door settings. In many cases, the CausalFM model’s in-context predictions match or outperform state-of-the-art alternatives even though those baselines are trained specifically for each task. Overall, the paper’s main result is a general recipe to train a single tabular foundation model for causal inference tasks, offering a new paradigm that could improve flexibility (test-time inference without retraining) in fields like medicine and economics.", "questions": "- Robustness to Prior Mis-specification: How sensitive is CausalFM to the correctness of its prior assumptions in practice? In a scenario where the true data-generating process deviates from the assumed SCM prior (for instance, an unobserved confounder is present when the model assumed none, or the functional form of relationships is different), what would happen to the model’s performance? Would the in-context learner recognize the model mismatch (perhaps via larger predictive uncertainty), or could it yield biased estimates? It would be helpful if the authors could clarify whether any experiments were done to assess robustness when the test data falls slightly outside the support of the training prior, and if not, how might one ensure reliability of CausalFM in such cases (e.g., by broadening the prior or diagnosing posterior outputs).\n\n- Interpretability and Practical Use: While the model automatically selects an identification strategy based on the data, it operates as a black box from the user’s perspective. Do the authors have insights or plans for making CausalFM’s reasoning more interpretable? For example, is it possible to extract which causal adjustment the model is implicitly using for a given dataset (such as detecting that it’s performing an IV-like computation versus a back-door adjustment)? Understanding this could be important for practitioners to trust the results. Moreover, how would the authors recommend validating CausalFM’s outputs on a real-world problem where the true causal effect is unknown? Are there diagnostic checks or visualizations one can use (perhaps analyzing the learned attention weights or the model’s posterior over causal effects) to ensure the model’s inference aligns with domain knowledge? Clarifying these points would help users apply CausalFM more confidently and transparently in practice.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper introduces a novel approach by leveraging prior-data fitted networks for causal inference, which might be the first to provide a comprehensive foundation model (CausalFM) covering multiple settings in one framework. Unlike past works that focus on a single identification strategy (e.g. only back-door criteria) or require task-specific models, CausalFM uses SCM-based priors to train a transformer that can flexibly handle back-door, front-door, and IV adjustments within one model. This is one step closer to a real causal FM: the model effectively learns to “select” the appropriate causal estimation formula from data context, enabling zero-shot inference on new datasets without retraining – a clear paradigm shift from traditional retrain-per-dataset approaches. The introduction of SCM-guided synthetic training data (via a Bayesian neural network prior) is especially novel, as it injects domain causal knowledge into the foundation model’s pretraining process. This combination of foundation models + causal SCM priors is an original contribution that extends the foundation model concept into the causal inference domain in a principled way.\n\n- The paper is grounded in strong theory. The authors formally define conditions under which a causal effect is identifiable and incorporate those into the prior design. In particular, they derive necessary criteria for a valid causal prior, proving that any well-specified prior must assign zero probability to SCMs that violate the identifiability of the target causal query. This result (Theorem 4.3) ensures that the model’s Bayesian posterior cannot mislead us with non-identifiable alternatives, thereby guaranteeing asymptotic consistency of the causal estimates under the given assumptions. Furthermore, by building on a Bayesian PFN framework, CausalFM naturally provides uncertainty quantification and connects to Bayesian consistency results (e.g. invoking a Bernstein–von Mises argument for the posterior on the observational distribution). The paper’s separation of identifiability (ensured via SCM prior design) and statistical estimation (learned by the PFN) follows established causal inference principles (Pearl’s identification-before-estimation philosophy), lending the approach a sound theoretical foundation. Overall, the inclusion of formal definitions, theorems, and proofs about prior construction, consistency, and identifiability demonstrates a high level of theoretical rigor that bolsters confidence in the method’s validity.\n\n- Strong Empirical Performance: The experimental evaluation is thorough and shows competitive or superior performance of CausalFM on a wide range of benchmarks. The authors test their model in standard back-door (observational CATE estimation) scenarios (including synthetic datasets and a semi-synthetic Jobs dataset), in instrumental variable settings (with both binary and continuous instruments), and in front-door mediation scenarios. Across these experiments, CausalFM’s one-model approach achieves error rates (PEHE) that are on par with or better than specialized models. For example, in CATE estimation, CausalFM slightly outperforms several state-of-the-art estimators (e.g. S-learner, T-learner, TARNet) and also improves over prior foundation model baselines: CausalFM attained a lower PEHE (≈0.51) than CausalPFN (≈0.56) or DoPFN (≈0.59) on synthetic CATE benchmarks. In the IV setting, CausalFM similarly achieved the best PEHE among foundation models (0.42 vs. 0.52 for DoPFN) and was comparable to the top specialized IV methods. Notably, in a challenging front-door adjustment task with confounded treatment and mediator, the CausalFM model obtained a PEHE around 0.90, outperforming a baseline foundation model (DoPFN at 1.27) and all but the very best plug-in front-door estimator. These results indicate that CausalFM not only generalizes across different causal inference problems but often matches or exceeds the accuracy of dedicated models, all while requiring no retraining per new dataset. The consistency of CausalFM’s performance across diverse settings and its edge over previous PFN approaches (which were limited to back-door or had no identifiability guarantees) stand out as a major strength. This empirical evidence convinces the reader that the proposed foundation model approach is viable and competitive in practice.", "weaknesses": "- While CausalFM is conceptually flexible, there may be practical scalability challenges. Training a PFN of this sort involves generating and learning from a very large number of synthetic datasets drawn from complex SCM priors, which is computationally intensive. The transformer model itself has a fixed context length and model size – this could limit the scale of datasets it can handle at test time (e.g. number of samples or covariates) unless the architecture is scaled up. The paper does not extensively discuss memory or runtime implications; however, foundation models for tabular data (like TabPFN) are known to handle only modest-sized datasets due to the need to input the dataset as context. If one were to apply CausalFM to truly large real-world datasets (say, thousands of units or very high-dimensional covariates), it is unclear if the current model would remain efficient or accurate. In short, questions of scalability (both in training cost and inference on large data) remain a concern. The method’s impressive performance is demonstrated on reasonably sized benchmarks, but its feasibility on substantially bigger or more complex tasks has not been shown, potentially limiting its immediate practical adoption for “big data” causal inference.\n\n- The work would benefit from deeper analysis of why CausalFM works so well. Currently, there is little in the way of ablation studies or interpretability insights. For example, the framework involves several design choices – using a Bayesian neural network prior, simulating counterfactuals during training, particular choices in the SCM parameterization – yet the paper does not present ablation experiments to quantify the contribution of each component (e.g., what if the prior did not enforce certain assumptions, or if the counterfactual simulation was removed?). Without such ablations, it is hard to assess which aspects of the CausalFM design are most critical. Additionally, the interpretability of the PFN’s in-context inference mechanism is not explored. CausalFM essentially acts as a black-box that implicitly decides which causal adjustment to apply (back-door vs front-door vs IV) based on the input data. However, the readers are not shown how the model makes this decision or whether it aligns with known causal formulas in each case. There is no analysis, for instance, of attention weights or internal representations to reveal if the model focuses on certain variables (e.g. an instrument or mediator) to choose the appropriate formula. This lack of interpretability means practitioners might find it hard to trust or validate the model’s reasoning on a given dataset. In summary, the paper’s experiments focus on accuracy, but omit diagnostic analyses: one misses ablation studies to justify design choices and interpretability checks to illuminate the model’s decision-making. This is a limitation because understanding the model’s behavior would greatly enhance confidence and provide insights for future improvements.\n\n- The success of CausalFM hinges on the assumption that the chosen prior distribution (the family of SCMs used in training) accurately captures the relevant causal structure of the target domain. This could make the approach sensitive to mis-specification. If a real-world scenario violates the assumptions baked into the prior (for example, there is an unobserved confounder not accounted for, or the functional relationships differ substantially), the model’s inferences may become unreliable. The authors explicitly restrict the priors to settings where identifiability holds, which is sensible, but it also means the model is not trained to handle scenarios outside those assumptions (e.g. it wouldn’t know how to behave if confronted with a non-identifiable case, except to potentially output high uncertainty). It would strengthen the work to understand how robust the model is to slight violations of its assumptions – this is not addressed in the current evaluation. Moreover, real-world validation is limited in the experiments. While the paper includes a semi-synthetic example (Jobs dataset with simulated outcomes) and numerous fully synthetic benchmarks, it does not demonstrate CausalFM on purely real observational data where ground-truth causal effects are unknown. Consequently, it remains to be seen how the model performs in practice on real datasets, and how one would verify its estimates when the true effect is not available for comparison. The lack of real-world case studies means the approach’s practical utility is still somewhat speculative. In essence, the framework’s generalizability to real, messy data and its robustness to deviations from assumed priors are open questions not fully answered by the paper. These limitations suggest that further work is needed to test CausalFM under less ideal conditions and to guide users in setting an appropriate prior for their specific application.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SxxUbQtRRV", "reviewer": "Reviewer_Zj58", "summary": "In this paper, the authors propose CausalFM for training foundation models for causal inference based on existing Prior-data Fitted Networks (PFNs). The core idea is to pre-train a transformer model on synthetic data generated from a family of Structural Causal Model (SCM), enabling zero-shot causal estimation in various settings, including back-door, front-door, IV, via in-context learning without retraining. The approach addresses a significant limitation of traditionally causal inference methods. The empirical results show competitive performance.", "questions": "In fact, some researchers have begun leveraging PFNs to pre-train causal foundation models, such as the following work. The pretraining method used in this paper is very simimlar to the existing work, although the authors have discussed these methods in the paper. In addition, the authors critique the following work for being restricted to back-door adjustment. However, this work can actually handle back-door, front-door, and IV settings via in-context learning without retraining. Since the pre-training methodology in this paper appears largely similar, its claimed novelty over these prior works remains unclear.\n\nJake Robertson, Arik Reuter, Siyuan Guo, Noah Hollmann, Frank Hutter, and Bernhard Scholkopf. Do-pfn: In-context learning for causal effect estimation, arXiv preprint, arXiv:2506.06039, 2025.\n\nIn the back-door training loss (Eq. 9), the model is trained to predict the individual treatment effect directly. What is the rationale behind this and whether you experimented with predicting potential outcomes Y(a) separately and then differencing?\n\nThe prior in Sec. 4.3 is based on a well-specified C-DAG provided by the practitioner. How sensitive is the method to misspecification of this C-DAG?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Beyond dataset-specific causal inference to in-context causal inference is a promising and important research direction.  In the paper, the authors propose a framework which is not restricted to a single causal setting. It supports back-door, front-door, and instrumental variable scenarios within a single model.\n\nThe introduction of the \"well-specified prior\" concept and the argument for incorporating identifiability assumptions into the prior are valuable theoretical insights.\n\nThe results on synthetic, semi-synthetic (Jobs), and other datasets demonstrate the model's versatility and strong performance.", "weaknesses": "The pre-training methodology in this paper appears largely similar to exsiting work in this area \n\nThe proof sketches are quite dense and lacks intuition, then are hard to follow.\n\nThe 24-hour training time on an A100 GPU is mentioned, but there is no discussion of inference speed or comparative training costs of the baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "psV9jUhVCV", "reviewer": "Reviewer_V8FG", "summary": "This paper introduces CausalFM, a framework for training transformer-based Prior-data Fitted Networks (PFNs) as foundation models for causal inference. The approach centers on \"C-SCM-Priors,\" novel prior distributions over Structural Causal Models (SCMs) constructed to explicitly satisfy identifiability assumptions for specific causal settings (e.g., back-door, front-door, IV). The paper presents a theoretical argument that restricting the prior to identifiable SCMs is necessary for the resulting PFN to be \"well-specified\" and yield consistent causal effect estimates. CausalFM is pre-trained on synthetic data from these priors, enabling it to perform Bayesian causal inference (specifically, CATE estimation) via in-context learning on new observational datasets without retraining. Experiments demonstrate competitive performance against specialized baselines in the targeted settings.\n[Note: I have used LLMs to improve my writing and help me answer paper questions]", "questions": "- How sensitive is CausalFM's performance if the user specifies a causal setting (e.g., assumes back-door identifiability) but the underlying data violates this assumption mildly (e.g., weak unobserved confounding)? Does the model fail catastrophically, or does it offer graceful degradation?\n- How critical are the specific choices made in the BNN architecture and sampling process for the C-SCM-Prior (Sec 4.3)? Were alternatives explored, and how much does performance vary with changes to the prior generation mechanism?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides a clear theoretical argument for why PFN priors in this Bayesian approximation context should enforce identifiability for consistent estimation (I am a little skeptical if we should make this assumption, see below - but realizing and demonstrating it is very valuable).\n- Introduces a structured way to build priors over SCMs using BNNs that respect the assumed causal graph structure and identifiability conditions.\n- CausalFM is designed to handle back-door, front-door, and IV settings, offering broader applicability than methods restricted to unconfoundedness, like CausalPFN.\n- Achieves strong average empirical results across multiple benchmarks and settings compared to specialized estimators that require per-dataset training and tuning.", "weaknesses": "- Reliance on Correct Identifiability Assumptions: The framework's core premise requires the user to correctly identify the true causal structure and select the appropriate identifiability strategy (back-door, front-door, IV) before applying the model. This is a strong assumption, as determining the correct causal graph and valid adjustment strategy from domain knowledge alone is often a major challenge in real-world applications. CausalFM automates estimation given these assumptions but offers no mechanism to validate them or handle uncertainty about the true causal structure. Potential Mismatch with Reality vs. Uncertainty Modeling: By strictly requiring identifiable priors for consistency (Thm 4.3), CausalFM might struggle or produce overly confident (but potentially wrong) estimates if the real-world DGP is non-identifiable or if the chosen identifiable structure is incorrect. CausalFM prioritizes potential consistency under strong assumptions over explicitly modeling uncertainty arising from potential non-identifiability.\n\n- The effectiveness depends on the specific design of the BNN-based C-SCM-Prior. The paper lacks sensitivity analyses regarding the prior's construction (e.g., BNN architecture, diversity of sampled SCMs) and how well this synthetic prior generalizes to the complexities of real-world data generation processes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)