Estimating Interventional Distributions with Uncertain Causal Graphs through Meta-Learning
Authors: Anish Dhir
Cristiana Diaconu
Valentinian Mihai Lungu
James Requeima
Richard E. Turner
Mark van der Wilk
Venue: NeurIPS 2025 poster ()
Keywords: Bayesian causal inference
Meta learning
neural processes
causal inference
Primary area: probabilistic_methods
TL;DR: We directly estimate interventional distributions from observational data with meta-learning

Abstract: In scientific domains---from biology to the social sciences---many questions boil down to \textit{What effect will we observe if we intervene on a particular variable?} If the causal relationships (e.g.~a causal graph) are known, its possible to estimate the intervention distributions. In the absence of this domain knowledge, the causal structure must be discovered from the available observational data. However, observational data are often compatible with multiple causal graphs, making methods that commit to a single structure prone to overconfidence. A principled way to manage this structural uncertainty is via Bayesian inference, which averages over a posterior distribution on possible causal structures and functional mechanisms. Unfortunately, the number of causal structures grows super-exponentially with the number of nodes in the graph, making computations intractable. We propose to circumvent these challenges by using meta-learning to create an end-to-end model: the Model-Averaged Causal Estimation Transformer Neural Process (MACE-TNP). The model is trained to predict the Bayesian model-averaged interventional posterior distribution, and its end-to-end nature bypasses the need for expensive calculations. Empirically, we demonstrate that MACE-TNP outperforms strong Bayesian baselines. Our work established meta-learning as a flexible and scalable paradigm for approximating complex Bayesian causal inference, that can be scaled to increasingly challenging settings in the future.

Estimating Interventional Distributions with Uncertain Causal Graphs through Meta-Learning
Anish Dhir* Imperial College London
Cristiana Diaconu" University of Cambridge
Valentinian Mihai Lungu University of Cambridge
James Requeima
Richard E. Turner
Mark van der Wilk University of Oxford
University of Toronto Vector Institute
University of Cambridge Alan Turing Institute
Abstract
In scientific domains-fron biology the social sciences- questions boil down to What effect will we observe if we intervene on a particular variable? If the causal relationships (e.g. causal graph) are known, it is possible to estimate the intervention distributions, In the absence of this domain knowledge, the causal structure must be discovered from the available observational data. However, observational data are often compatible with multiple causal graphs, making meth- ods that commit a single structure prone to overconfidence. principled way to manage this structural uncertainty is via Bayesian inference, which averages over posterior distribution on possible causal structures and functional mecha- nisms. Unfortunately, the number of causal structures grows super-exponentially with the number of nodes in the graph, making computations intractable. We propose to circumvent these challenges by using meta-learning to create an end-to- end model: the Podel-Averaged Causal Estimation Transformer Neural Process MIACE-TNPP) The model is trained to predict the Bayesian model-averagee in- terventional posterior distribution, and end-to-end nature bypasses the need for expensive calculations. Empirically, we demonstrate that MACE-TND outperforms strong Bayesian baselines. Our work establishes meta-learnigg as a flexible and scalable paradigm for approximating complex Bayesian causal inference, that scaled to increasingly challenging settings in the future.
Introduction
Answering interventional questions such as: "What happens to Y when we change X?" central to areas such healthcare (37] and economies [57]. One can estimate such interventional distributions by actively intervening on the variable interest and observing the effects (obtaining interventional data), but this can costly, difficult, unethical, or even impossible in practice (35). Causal inference offers an alternative by leveraging readily available observational data alongside knowledge of the underlying causal relationships in the form of causal graph (50). A causal graph can be manually specified when domain knowledge is available. In the absence of this, causal discovery techniques attempt to learn the causal structure from data [45]. However, causal discovery from purely observational data is notoriously difficult. Identifying the true graph requires strong assumptions, such as the use of certain restricted model classes [54 1 Ch. 4] and the acquisition of infinite data, that are rarely met in practice. It is therefore often the case that data provides plausible evidence for a set of causal graphs, even though each of these graphs may imply drastically different causal
Equua contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

effects. Picking a single graph can thus result in poor downstream decisions [5] 4.. In this work, we address the challenge of tractably estimating interventional distributions when the true causal graph is uncertain, a common scenario in real world applications.
Instead of using single graph drive decisions, the Bayesian framework provides a principled way to manage the uncertainty over the causal models by maintaining distribution over possible causal structures and functions relating variables. With access to both, model uncertainty can be accounted for by marginalising the interventional distributions over both posteriors (fig. this procedure has two main challenges. First, the space of causal graphs grows super-exponie expeeeexponentially with the number of variables, making exact posterior inference intractable, and sampling difficult to scale. Second, even with a posterior over causal graphs, estimating an interventional distribution in each plausible causal graph necessitates computing the posterior over functional mechanisms, which analytically intractable except in simple models. Poor approximations at any point in this pipeline can result in inaccurate interventional distributions. Consequently, most works restrict themselves to simple functional mechanisms and constrain the allowable structures, limiting their applicability.
40 However,
To overcome these bottlenecks, we turn to recent advances in meta-learning. Neural processes (NPs) 20 are family of meta-learning models that approximate the Bayesian posterior with guarantees 25 by directly mapping from datasets to the predictive distribution of interest, thus bypassing the
19
intractable explicit modelling of intermediate posteriors. They underpin several successful methods with strong empirical performance across a range of real-life domains, ranging from tabular data classification (8 to weather modelling (33 Recently, NPs have also been applied to causal discovery
14]. with Dhir et al. [15] showing that they can accurately recover posterior distributions over causal graphs. However, these existing approaches cannot estimate interventional distributions.
30 39
In this work, we apply NPs to the problem of causal inference directly from data by developing a meta-learning framework that targets Bayesian posteriors for interventional queries-the Model- Averaged Causal Estimation Transformer Neural Process (MACE-TNP). As shown in fig. D our method amortises the full Bayesian causal inference pipeline (learning the posterior over the causal structure and functions, and marginalising over them), all within a single model. By directly estimating the interventional distribution, our approach avoids compounding errors from intermediate approximations of posteriors and marginalisation, enabling more accurate and computationally- efficient inference under model uncertainty. Our contributions are threefold. First, we propose an end-to-end model trained on synthetic datasets to directly approximate the Bayesian posterior interventional distribution. Second, we empirically show that when the analytical posterior available in closed-form form, MACE-TNP converges to it. Third, we demonstrate that MACE-TNP outperforms a range of Bayesian and non-Bayesian baselines across diverse experimental settings of increasing complexity, highlighting the method's 's potential to scale to high-dimensional datasets. Our framework paves the way for gettaleerrning-asee foundation models for causal interventional estimation.
Classical Approaches
Step Interventional distributions
pla. 22), Gr, D.w.)
Gi
plnittols j). G.biin,
Step Causal discovery
Prest P ob
pla.dd((z Daa.) pla, d. Daw
Debs
G2
piz: do 2z), Pats)
pla, GN. Dahe)
[0] grows superexponentially
8 Gx
Generally intractable
Modeae-Aeergee Causal Estimation Transformer Neural Process
Figure 1: Overview of MACE-TNP. Unlike classical approaches, that usually require two-step pro- cedure which 1) first involves posterior inference over the graph structure, followed by 2) complicated inference over the functional mechanism, MACE-TNP amortises the full causal inference pipeline.

---

2 Background
Our goal in this paper is to compute interventional distributions that take uncertainty over the causal model- and functional mechanism-ing account. Like a majority of works that learn causal models from data, we assume no hidden confounders. We set up the problem and provide background here. Throughout the paper random variables are denoted by uppercase letters (e.g., X)), and their realizations by lowercase letters (e.g.. 2). Boldface is used to denote vectors (e.g., x, X).
Causal model: Causal concepts can be formally defined by considering a causal model. Given directed acyclic graph (DAG) g with node set V == (1...., D), functional mechanisms f == .......')) J and independent noise terms U == {U1,.... UD}, Structural Causal Model (SCM) defines variables X, recursively as follows (0
(J1.-.-,Je
: fi(PA,,
for i = l...,, D.
where PA, C (x),,.,,XD))(((X is the set of parents for X; This process induces a joint distri- bution over all the variables. Such construction can then be used to formally define interventions. We focus on hard interventions, denoted do(r,), where X, is set to a fixed value z,, leaving all other mechanisms unchanged [50]. The resulting distribution of any variable X, is known as the interventional distribution p(X; do(z,)), We use Dabs to refer to datasets of independent and identically distributed (i.i.d) observations drawn from the model, and Dim for i.i.d. observations drawn from any interventional distribution.
The task of causal discovery is to reconstruct the data generating graph 9 from an observational dataset Daas from an SCM. However, this typically only identifies a Markov equivalence class (MEC) of graphs that encode the same conditional independences [50]. Unique graph identification requires strong assumptions that may not hold practice, such as hard restrictions the allowable model classes [54] Ch. 4]. Furthermore, the identifiability guarantees of these methods also only hold in the infinite data setting. For a lot of tasks, causal discovery is a means an end-namely, estimating interventional distributions for downstream tasks. With a causal graph and observational data, causal inference allows for estimating an interventional distribution p(z,/do(z;) )), if it is identifiable [50 59 However, inferring the ground truth causal graph is difficult. This has drastic implications for computing interventional distributions. Two graphs, even within the same MEC, may have very different interventional distributions [5I]. Relying on single causal graph to compute interventions can thus lead to incorrect conclusions.
Causal discovery and inference:
Bayesian causal inference: Due to the limits of causal discovery. uncertainty is inherent in causal inference. Finite data issues further compound the problem. The Bayesian framework allows for quantifying the model uncertainty, both in the causal structure and functions, and use it for downstream decision making.
Definition 2.1.
We define a Bayesian causal model (BCM) as the following hierarchical Bayesian model over causal graphs G. functional mechanisms f. and dataset Dobs of Nots samples: G ~ (G), f = {fi)iev ~ PRCM (flg), Don. := (X" J ~ IIN= Ihev PROM (x²*) 9), where x² denotes the i-th node of the n-th observational sample. This implies the joint distribution PBCM(Dos obs f)
As BCMs are defined with a causal graph, they induce a distribution over interventional quantities as well. Analogous to standard causal models, interventions PBCM (z, (d0 fú, 9) can be computed by setting (-) = Ij and leaving all other mechanisms unchanged.
Given an observational dataset, our task is to estimate an interventional distribution of interest. To do this, is necessary to infer the possible graphs and functional mechanisms that generated the dataset. The Bayesian answer to this question is through the posterior
PBCM (f, g(t)) x PBCM Dots f, (f
If the underlying model is identifiable, for example by restricting the function class of [54] Ch. 4), then under suitable condition the posterior over 9 will concentrate on the true graph in the infinite data limit [13 [4 12). However, for finite data, if the causal model is not identifiable, the posterior
2 Note that we assume no hidden confounders which is common for causal discovery. Hence, given the ground truth causal structure, interventional distributions are identifiable. 'The prior has to have positive density over the true underlying data generation process.
3

---

will quantify the uncertainty over causal graphs. To make use of the uncertainty. Bayes prescribes to average the interventions over the models [40], which we call the posterior interventional distribution
p)cm(())dd((=)) Dabs) = Σ/ g
preM (x:))((()))) ())ncc (f|g, Pan.) (GPPaa).
(3)
Computing the above quantity often intractable for two main reasons: 1) computing PBCM (g (aats) is challenging the number of causal graphs increases super-exponentially with the number of variables, 2) (BCM(f]G, Doss) is only tractable for simple models.
Transformer neural processes: To bypass the need to compute the intermediate intractable quantities in eq. (3), we turn towards neural processes (NP) [2 19 ). From a Bayesian perspective. NPs
incorporates prior through the distribution over datasets that it is trained on [46], and, during inference, directly provides estimates of the posterior distribution of interest, side-stepping any explicit approximations of intermediate quantities. In particular, the Transformer Neural Process (TNP; [4 16 ) which builds on the scalability and expressiveness of the transformer architecture (62)
has achieved strong results across diverse domains [2 6 29 14, motivating its use in our model- averaged approach for causal intervention estimation.
3 Related work
Estimating the posterior interventional distribution (eq. (3) is challenging. The dominant paradigm involves a two stage process: 1) obtaining samples from the high dimensional posterior over graphs, and 2) estimating the interventional distribution under each sampled DAG, followed by averaging the result (fig. D. Although principled, this process faces computational challenges in both stages.
The first stage is challenging due to the super-exponential size of the space of DAGs. Early score- based methods addressed this by leveraging score equivalence, allowing search over the small space of MECs instead of individual DAGs. To achieve this, they used restricted the model family to linear Gaussians, with specific priors make the scores analytically tractable 22 22
To accommodate broader model classes, Madigan et al. [4] introduced an MCMC scheme over the space of DAGs. However, the large space of DAGs leads slow mixing and convergence issues, limiting the number of effective posterior samples 17 31 58, 48, 33). A common bottleneck in these approaches is that scoring the proposed structures at each MCMC step requires expensive marginal likelihood estimation. This is often mitigated through reducing the graph space by restricting the in-degree of each node. Variational inference (VT) offers a cheaper alternative [ii] but struggles to capture multi-modal posteriors inherent in causal discovery [61] Sec. 3.11, and can still have demanding computational costs (e.g. SVGD used in (88 scales quadratically with samples). Crucially, any inaccuracies or biases in this stage affects the downstream estimation in the second stage.
The second tional burden. It requires performing inference with a potentially complex functional model for every single DAG sampled from the approximate posterior pG D. As result, previous work has only considered simple functional models where the inference is not too prohibitive. While early work in simple settings like linear Gaussian models allowed for closed-fome form averaging 64 10 recent works often employ Gaussian Process (GP) networks [8] where this is not possible. To tackle this, Giudice et al. 24 use complex MCMC schemes for both hyperparameter posteriors of the GPs and graph sampling, but have to resort to approximating the final interventional posterior distribution with Gaussian for computational tractability. Toth et al. (60 61 also use GP networks but use the cheaper alternative of using MAP estimates for hyperparameters. However, both ultimately rely on the expensive process of estimating interventions by sampling from the GP posterior conditional on each DAG. Hence, despite variations, the core limitation of expensive inference persists across these approaches, especially prohibiting the use of more flexible function model classes.
over the posterior of causal graphs- has its own significant computa-
In contrast to this explicit two-stage procedure, we propose leveraging NPs 20 to directly learn an estimator for the target interventional distribution conditional on the observational data D Our approach aims to learn mapping Dobs + p(y]do(x:) Dats) that does not require approximating potentially problematic intermediate quantities. This effectively amortises the complex inference and averaging procedure over the training of the NP. Our method thus seeks to mitigate the severe computational bottlenecks and avoid the compounding of approximation errors inherent in the standard two-stage pipeline for Bayesian causal inference.

---

There have been recent attempts at end-to-end or meta-learning approaches to estimating inter- ventional distributions. While some methods assume knowledge of the ground truth causal graph others offer similar data driven approach as ours, but only in restricted settings. For example, Geffner et al. [2]] offer an end-to-end approach but restrict to additive noise models, and do not perform functional inference. Sauter et al. (56) also use meta-learning to directly target the interventional distribution. Apart from differences in architecture and the loss used, their method is limited to discrete interventions. In contrast, types functional mechanisms or types of interventions. Further, by viewing meta-learning through a Bayesian lens, we provide insight into the role of the training data as encoding prior distribution [25 46 28 R Tying our method to Bayesian inference also provides an understanding of the behaviour
9 42 69
propose general framework that is not restricted to
of our model under identifiability and non-identifiabilityy of the causal model [ [4 13
4 A transformer model for meta-learning causal inference
Causal inference with neural processes: The focus of this work is causal inference directly from data-preddd pred--ciiiniigg the distribution of a variable of interest X; when we intervene on another variable do(2;) given access only observational data Dobs (eq. 33) Instead of using a two-step approach as in fig. [ which is computationally expensive and approximation error prone, we propose to directly learn the map from the observational dataset the posterior interventional distribution of BCM in an end-to-end fashion with NPs. With a chosen BCM, our aim is to approximate the true posterior interventional distribution defined in eq. 3 To do this, we minimise the expected
Kullback-Lein (KL) divergence over the tasks & == (Dobs. i, j. X;) between the true posterior interventional distribution and the NP model predictions po (z: (d((2,)) Dons)
8 == arg min Eg KL # = arg max Be [x [l, Dots )][ + C,
Pots)) IIpo (X, (d)) Dobs)))
(4)
where { ~~ p(Data, i, J. x, X proM (x:]o((( Dats), are the parameters of the NP, and C is some constant independent of 0.
Hence, our objective requires us to generate tasks, and interventional data from BCMs, in order to find the optimal 0*, To do this we, 1) sample a graph G ~ (9), and 2) a functional mechanism for each of the D variables from the graph f PBCM (f]g). Conditioned functional mechanism f we then 3) draw Nots samples for each variable to construct the observational (Dcm(D..... functions as the observational data, we 4) randomly sample a variable index j to intervene upon and Nin intervention values X, ~ (O, I), set the values of node j to be Xj. and 5) draw Nim samples of each node forming an interventional dataset Dim ~ PBCM (Dm (d((x;)))))))) ), f). Finally we sample an outcome node index i extract samples of PBCM (x| |do(x)) Jú: 9) from Djnt-
the sampled graph & and
data Dobs PBCM (Dob. [f, 9). To construct the interventional data, keeping the same graph and
While at training time we assume access to an explicitly specified Bayesian causal model, at test time we do not need access to such model. In fact, the Bayesian prior is implicitly encoded into the NP through the distribution over its training datasets [46] Inference for a new dataset simply requires a forward pass through the network.
Recovery of exact prediction map: Bruinsma g Proposition 3.26] shows that, in the limit of infinite tasks and model capacity, the global maximum of eq. (4) is achieved if and only if the model exactly learns the map (Dats, xj) + PROM (1r, [do(x;), Dats). Hence the NP learns to implicitly marginalises out any latent variables in eq. (3) 25 While the constraint of infinite tasks is limiting when applying NP to real-world datasets, if the tasks are generated through a known Bayesian causal model, we have in theory access to infinite amount of tasks.
Model architecture and desirable properties: Given we are interested in predicting (BCMa (do(x), D. variables play distinct roles as either the outcome node X,. the intervening node X,, or the nodes that are being marginalised. Thus, properties of this distribution, and the role of the variables, guide our architecture choice. First, the interventional distribution remains invariant when the observational data samples are permuted or the nodes being marginalised over are permuted (permutation inneriaanvvaeeeee with respect to observational samples and to all nodes ex- cept the outcome X, and intervention X,). Second, permuting the interventional queries should permute the samples of the target distribution accordingly (permutationn-eqiiariiace with respect
5

---

to interventional samples). Similarly, permuting any nodes involving the outcome or intervention nodes should yield the corresponding permuted interventional distribution peamantttton-eqvvvriiitn with respect to outcome and intervention nodes). For example, permutting the outcome and inter- vention nodes <+> j should result in the permuted p(x; [do(x;). Pats). Furthermore, we assume no correlations among the interventional samples and as such restrict our attention to the family of con- ditional neural processes (CNPs), where the predictive distributions factorises over the interventional samples pa xx (x; non-Gaussian even in very simple cases, we opt for Mixture of Gaussians (MoG) representation of (o(2:]o((1)) Dos)) [6.
), Dats) = I[>= po(2; (a())) Dons). As interventional distributions can be
An architecture that is flexible enough to satisfy these desiderata is the transformer 662 34]. We
provide a schematic architecture for our proposed model, the Modet-Averaged Causal Estimation Transformer Neural Process (MACE-TNP), in fig. 2 and give detailed explanation of each of its components in Appendix A.2
1. Embed with variable-specitic MLPs
2. MACE Transformer Encoder
3. MACE Transtormer Decoder
Attantion among samples
(Attantion among nodes
dasts
dumbas
Nobe
Naw
Sel-attenfon ameng observatonal samples
Lue, Th Nars
MLP
Cross-attittio behowee intervections and dbservationl samples
Nu
Nus
XL
Figure 2: Overview of MACE-TNP yielding px(x; |do(x,); Dane). Inputs are 1) embedded via variable-specific MIPs, 2) fed into transformer encoder that alternates sample-wise and node-wise attention. The resulting outcome node representation from the unknown interventional distribution is 3) decoded to obtain the parameters of the NP distribution.
The model takes as input a matrix of Nobs observational samples of D nodes and an intervention matrix of Nm queries for a node of interest X,, with the rest of the D nodes masked out (by zeroing them out). Variables play distinct roles in both matrices, requiring different encoding strategies- either as the node we intervene upon (j), outcome node (i), or node we marginalise over. Our input representation also needs to reflect that observational and interventional samples originate from different distributions. As such, we employ six variable-specific MLPs, one for each combination of (node type, sample type). These MLPs produce dember-dimensional embeddings, resulting in representations Zons E IR- x. Dxdenaa and Zim E IR Nm DX Dxdental for the observational and interventional data, respectively.
Embedding:
MACE transformer encoder: To satisfy the afore- mentionet permutation symmetries, we construct an encoder of L layers where we alternate between attention among samples, followed by attention
The two attention mechanisms that we use are multi-head self-attention (MHSA) and multi-head cross-attention (MHCA) bothh defined in eq. (5) and eq. 7 respectively. More specifically, at each layer I {1,.... ... L} we 1) update the observational data representation Z E R Nox dentel through MHSA. This then used 2) modulate the interventional data repre- sentation Z N. x Dedened through MHCA, an operation that assures permutation equivariance with respect to the interventional samples. We then 3) concatenate the two representations to obtain z" E R Niecc Nm)) x D xemnne followed by 4) MHSA among the nodes to yield the output at layer / which acts as input at layer 1+11 1:
among nodes [5]] 39 32
1. MHSA(Z + 2. Z MHCA(Zİ 2 3.Z" [Z] Z
4. obs 2 MHSA(Z")
attantion amang samples
attention ameng nodes
The information required for the target distribution is now encoded in the outcome node (i-th index) of the interventional matrix. Z 12 xeeebodd This is passed through an MLP decoder to obtain the final distribution. To parametrise expressive interventional distributions, we construct the output distribution of the NP as an MoG with Neump components (6]. The NP outputs the mean, standard deviation and weight corresponding to each component for each interventional query z²: {u, o, wxzz)) == {mx(2])) oR (x()), wo( wx(x")) Noue
MACE decoder:

---

Loss: The model is trained to maximise the log posterior interventional distribution according to eq. where, with a MoG parameterisation:
Co(x, wxx,))
log po do(u?), Damn)
Σ log
Σ
wx(zf) -Ma" 1
5 Experiments
We evaluate the performance of our model, MACE-TNP, against Bayesian causal inference baselines, and causal discovery method that selects single graph. With our experiments we aim to answer: 1) When analytically tractable, can we confirm that our model recovers the true posterior interventional distribution under identifiability and non- non-identifiabilityy of the causal graph, 2) How does model compare against baselines when the baselines' assumptions are respected and when they are violated, 3) How does our model perform when the number of nodes scaled, 4) How does our model perform when we do not have knowledge of the data generating process? Code for our experiments is available at: https: httttp:://githb... com/ An
To train MACE-TNP, we randomise the number of observational samples Nots ~ (500, 750), and set Nin = 1000 Nobs. The training loss is evaluated on these Nit samples. For testing, sample 500 observation points and compute the loss against 500 intervention points.
Baselines: We benchmark against methods that infer distributions over causal graphs and sample to marginalise across these graphs when estimating posterior interventional distributions. DiBS-GP [60], ARCO-GP [61], and BCI-GPN (24) all use GP networks, but differ in the inference procedure over graphs. DiBS-GP uses a continuous latent parametrise graph, ARCO GP uses an order parametrisation of DAGs, whereas BCI-GPN uses an MCMC scheme to sample DAGs. We also compare against DECI 21 which assumes additive noise and uses autoregressive neural networks to learn distribution over causal graphs while only learning point estimates for functions. Finally, to show that learning a distribution over graphs is useful for causal inference, we compare against a non-Bayesian baseline that uses NOGAM [4] to infer single DAG, and estimates the interventional distribution by using GPs: NOGAM-GP.
For evaluation, we compare the model's posterior interventional distributions on held out datasets. Unlike graph-based metrics that only assess structural accuracy, this task requires correct inference of both the causal graph as well as the functional mechanisms. The posterior interventional distribution of the data generating model only analytically tractable simple cases (section (5.1). In these instances, we report the KI. divergence between the data generating model's 's posterior interven- tional distribution and that of the NP, averaged over intervention queries. When the analytical solution is not available, we report the negative log-posterior interventional density (NLPID) of the true D)]]
Metries:
intervention outcomes under the model: -Exx~~((01)) (X
23
5.1 Two-node linear Gaussian model
First, we test on synthetic data where the ground-truth posterior interventional distribution is tractable. Specifically, we are interested in the behaviour of our model under identifiability and non-identifiability of the causal structure. For this, we generate from bivariate, single edge, linear model, with Gaussian noise-a model that identifiable when the variances are known [53], but non-identifiable under specific priors [22) (appendix B.T. Besides the KI. between the ground-truth and MACE-TNP, we also report the KL between the interventional distribution conditioned on the true function and graph, and MACE-TNPP output distribution. The latter gauges accuracy in learning the true interventional distribution, which requires identifiability of the causal graph.
The results are shown in fig. 3 for the identifiable (left) and non- identifiallee (right) cases. They confirm that the output of MACE-TNY TNP does indeed converge the Bayesian optimal posterior, as the dark blue lines indicating (L (za dd( Dots) (2: (z) Dona)) )) go to O with increasing sample size in both cases. Moreover, as expected, KL(PECM (x: (d((oo f" G (*)lle Ilpp (x Dats)) )), where (f", g*) characterise the true data-generating mechanism, does not go to 0 in the second case due to the non-identifiability of the causal graph (as indicated by the red line). The flexibility of our architecture also allows for conditional queries, multiple interventions, well as casily incorporating
7

---

interventional data to help identify causal relations. Hence, we investigate here whether providing
small number Mi 5 of true interventional samples, alongside the observational data, resolves
- identifiability challenges in the non-identifiablee case. As shown in fig. (right) with the green line, we Ma
find that this does indeed lower the (L(PBCM (z:)do(z;) f" 9*) (z.]ooo))(2)) Dots: {a; ] I= )),
suggesting that even limited interventions can enhance identifiability, Moreover, we show that the divergence further decreases with more interventional samples Min in appendix B.
Rl. divergenie
0.4
KL (pacual (NN)
KL. (ebssl PAACE TNP)
Kl. divergennee
0.8
-
KL (pucaall Puacsss (uaccss-rr))
Kl (Phccll PMACE TSP)
0.3
KL. (Piccall puscee--nx)) Mix
0.2
0.1
0:22
0.0
10
10
10
10a
Observacionll sample size
Oberrvational sample size
KL divergences as a function of the observational sample size, for the identifiable case (left) and the non non-identifiable one (right). Dark blue denotes PBCM PBCM-the posterior interventional distribution defined in eq. (3). red and green use P the interventional distribution conditioned on { }r}. We additionally provide MACE-TNP with Miat = 5 interventional samples. We indicate the median and the 10-90% quantiles.
Figure 3:
5.2 Three-node linear Gaussian confounder vs. mediator
In the previous section, showed that in the two-node case, when the causal graph is identifiable, MACE-TNP can identify the correct interventional distribution. However, with more variables, finding the correct interventional distribution requires proper adjustment of the variables that are not the treat- ment or the outcome. For example, if we are interested in the distribution p(y do(x)) and another vari- able is confounder, the interventional distribution is p(y] do(x)) = p(y|æ,, 2)p(z)d2. However, if the variable is mediator, the interventional distribution is p(y]do(z)) f p(ylæ, (z. 2 z)p( (222222))22 Here, we show that in the three-node case, with identifiable causal structures, MACE-TI TNP implicitly performs the required adjustment.
We train a MACE-TNP model on fully-connected three-node graphs with data generated from identifiable linear Gaussian structures [53] We then estimate the KI. divergence between the interventional distribution conditioned on the true data-generating mechanism and the model's interventional distribution: KL(PRCM (zz/do(zz;), f", G* plppoozz (z,))((z;) Dons)) for 1) confounder graphs, 2) mediator graphs, and 3) confounder graphs where the confounder is unobserved in the context. Figure 4 (right) shows the results. As expected, as increase the sample size for 1) and 2). MACE-TNNE TNP more accurately identifies the correct interventional distribution, showing that implicitly adjusting the third variable depending on whether it is mediator When the confounder is unobserved, and the interventional distribution and the causal graph are not identifiable, the KL between the true distribution and the MACE-TNP tends to constant value above zero.
confounder.
5.3 Three-node experiments
Next, we compare our model MACE-TNP against the baselines in three-node setting. Here, there are 25 graphs in total, making inference over the graph easier than in higher-node settings. Given that most baselines either neural networks or GPs, we compare MACE-TNP to the baselines under two scenarios: 1) when tested on GP data, and 2) when tested on data generated using neural networks (NN). For each functional mechanism, we train separate MACE-TNY model. Full experimental details are provided in Appendix C.2.2.
When tested in-distribution (on datasets from the same distribution the model was trained on). MACE-TNE TNP consistently outperforms all baselines across both functional mechanisms, as shown in table I MACE-TNP outperforms GP-based methods through its implicit handling of hyperparameter inference, that the GP baselines may struggle with. It also surpasses DECI (an NN-based approach) on both GP and NN data by employing a Bayesian treatment over functions. This highlights key

---

RL (pS (MACE-TXP)
5-
Confomdder
Confounder unaheerved
4-
Confounder
Mediator
KL. divergenos
Medistor
Z
22
1-
0
10'
X
X
Y
Observationa sample
Figure 4: KL divergence (right) of the interventional distribution conditioned on {f" ," and the model's for the confounder (dark blue), mediator (red) and unobserved confounder (light blue) cases. With increasing sample size, MACE-TNP identifies the correct distribution for both the mediator and confounder cases, implicitly carrying out the required adjustment.
advantage of MACE-TNP it can easily incorporate complex Bayesian causal models into its training pipeline by sampling training datasets, whereas traditional Bayesian methods rely on inadequate approximations for complex scenarios.
Table I: Results for MACE-TNP and baselines on the three-node experiments. We show the NLPID (L) and report the mean + the error of the mean over 100 datasets. Each row corresponds to a different functional mechanism used in the test set (GP/ NN).
MACE-TNP
DiBS-GP
ARCO-GP
BCI-GPN
DECI
NOGAM+GP
GP
563.9 23.4
644.2 27.2
630.7 1 22.3
628.5 27.5
632.0 25.6
749.4
43.0
NN
527.9 1 19.8
807.6 1 50.1
851.2 1 55.0
706.8 1 5.0
588.0
23.6
815.6 1 58.4
Out-of-distrie distrbbtitti testing: All previous experiments trained MACE-TNP on the true functional mechanism (GP or NN). However, a natural question is: how does the model perform when tested on out-of-distributtinn (OOD) data? To probe this, evaluate MACE-TNP (GP) NN-generated data and MACE-TNE (NN) on GP-generated data. As expected, performance degrades when the test mechanism differs from training, since model lacks built-in inductive bias for unseen mechanisms: MACE-TNP (GP) tested on NN achieves 608.3 + 17.3, compared to MACE-TNP (NN) with 527.9 = 19.8. Similarly, MACE-TNP (NN) tested on GP achieves 678.0 + 10.0, compared to MACE- TNP (GP) with 563.9 = 23.4. However, NPs trivially support additional training on any data likely be informative. Indeed, training MACE-TNP on the combined GP+NN data nearly recovers in-distribution accuracy, achieving 531.0 19.4 on NN data and 583.9-21.3 on GP data (see table 4).
5.4 Higher dimensional experiment
We next investigate the scalability of our method. This is especially relevant in many modern applications (e.g. genomics, neuroscience, econometrics, and social-network analysis) which naturally involve high-dimensional systems where understanding intervention effects is crucial. We do so by testing single trained model on increasingly higher-dimensional data, scaling up from 20 up to 40 nodes. The functional mechanisms used are a mix of NNs and functions drawn from a GP with an additional latent variable input. More details on data generation are given in appendix B.3
Table 2 shows that MACE-TNP outperforms both the Bayesian, well as the non-Bayesian baselines across all node sizes. Moreover, the underperformance of the non- non-Bayesian baseline NOGAM+GP underscores the importance of capturing uncertainty with higher number of variables. The majority of our baselines involve Ga-based approaches, which can become prohibitively expensive with higher number of variables. For example, we do not report BCI-GPN as its MCMC scheme is too expensive for these node sizes. In contrast, MACE-TNP can readily leverage advancements that have made neural network architectures scale favourably other domains. Inference after training only requires forward pass through the network.

---

Table Results for MACE-TNP and baselines on the higher dimension experiment. We show the NLPID (1) and report the mean + the error of the mean over 100 datasets. Each row corresponds to different number of variables.
MACE-INP
DiBS-GP
ARCO-GP
DECI
NOGAM+GP
20 variables
660.4 + 5.2
701.9 + 4.0
701.9 4.0
686.3 + 6.7
942.7 + 23.8
30 variables
653.3 5.7
7.3.2
713.0 4.7
675.6 6.1
946.9 19.2
40 variables
665.8 + 4.8
711.5 4.6
712.1
683.0 5.1
986.0 + 20.0
5.5 Unknown dataset generation process
Finally, we apply our proposed method on the Sachs proteomics dataset [55 which includes measurements of D Crucially, we do not retrain any model for this task; instead, we reuse the model from section 5.4
proteins from thousands of cells under various molecular interventions.
which was trained exclusively on synthetic data. Following 7 36 66], we retain only samples with interventions directly targeting one of the D 11 proteins, yielding 5846 samples: 1755 observational and 4091 interventional across five single-proted perturbations. The results indicate that our method performs competitively with our strongest baseline, giving an NLPID for MACE-TNY of 998.9 + 104.9 compared to 1000.9 133.5 5 for DECI, when averaged over interventional queries and 10 outcome nodes. We provide additional comparisons the remaining baselines in appendix C.2.3 This shows the potential of tackling interventional queries in real-world settings with a fast, data-driven framework that captures uncertainty in principled manner and leverages flexible and expressive neural network architectures.
Conclusions
We address the challenge of efficiently estimating interventional distributions when the causal graph structure is unknown. Our solution, MACE-TNP, an end-to-end meta-learning framework that directly approximates the Bayesian model-averaged interventional distribution by mapping observational data to posterior interventional distributions. When the true posterior available, we show empirically that the model's predictions converge to it. Moreover, in simple non-identifiable case, we show that interventional data allows for capturing the true underlying mechanism. When employing more complex functional mechanisms, well MÁCE-TNN outperforms strong Bayesian and non- -on-Bayesian baselines, with the only requirement being access samples from prior distribution (implicit explicit) meta-train time. One limitation our model is its reliance on substantial training-time compute and data to effectively capture diverse range of causal mechanisms. Moreover, if the test distribution is not properly covered by the training data distribution, the model may struggle to generalise. However, as our cuitof-distribution experiments demonstrate, integrating additional data better cover the target distribution is straightforward and efficient in improving MACE-TNP's generalisation capabilities. Finally, the attention mechanism scales quadratically with the number of variables and samples, which can be costly. However, MACE-TNE can leverage recent advances in sparse and low-rank attention to mitigate this overhead [67 68 Future work includes a thorough investigation into the interventional sample complexity required by the NP for accurate interventional estimation, as well as how best construct the prior over BCMs to capture the complexities of real-life data.
data (up to 40 nodes),
References
[1] Matthew Ashman, Cristiana Diaconu, Junhyuck Kim, Lakee Sivaraya, Stratis Markou, James Requeima, Wessel Bruinsma, and Richard Turner. Translation Equivariant Transformer Neural Processes. In Forty-first International Conference Machine Learning, 2024.
[2] Matthew Ashman, Cristiana Diaconu, Eric Langezaal, Adrian Weller, and Richard E. Turner. Gridded transformer neural processes for large unstructured spatio-temporal data, 2024. URL https://arxix: org/ /
[3] Albert-I Barabási and Réka Albert. Emergence scaling in random networks. science, 286(5439), 1999.
10

---

[4] Alexis Bellot. Towards bounding causal effects under markov equivalence. In Proceedings of the Fortieth Conference on Uncertainty in Artificial Intelligence, 2024.
[5] José M Bernardo and Adrian FM Smith. Bayesian theory, volume 405. John Wiley & Sons, 2009.
[6] Christopher M Bishop. Mixture density networks. 1994.
[7] P. Brouillard, S. Lachapelle, A. Lacoste, S. Lacoste-Julien, and A. Drouin. Differentiable causal discovery from interventional data. Advances Neural Information Processing Systems, volume 33, pages 21865-21877, 2020.
[8] Wessel Bruinsma. Convolutional Conditional Neural Processes. PhD thesis, Department of Engineering, University of Cambridge, 2022. URL handle/1810/55 /554383
cam ac uk/
[9] Lucius E. J. Bynum, Aahlad Manas Puli, Diego Herrero-Quevedo, Nhi Nguyen, Carlos Fernandez- Granda, Kyunghyun Cho, and Rajesh Ranganath. Black box causal inference: Effect estimation via meta prediction, 2025. https://aaxxii org/ abs /203 11503.0985 05985
[10] Federico Castelletti and Guido Consonni. Bayesian inference of causal effects from observa- tional data in gaussian graphical models. Biometrics, 77(1), 2021.
[11] Bertrand Charpentier, Simon Kibler, and Stephan Gunnemann. Differentiable dag sampling. In International Conference on Learning Representations, 2022.
[12] David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 3(Nov), 2002.
[13] Anish Dhir, Samuel Power, and Mark van der Wilk. Bivariate causal discovery using bayesian model selection. In Forty-first International Conference on Machine Learning. 2024.
[14] Anish Dhir, Ruby Sedgwick, Avinash Kori, Ben Glocker, and Mark van der Wilk. Continuous bayesian model selection for multivariate causal discovery. arXiv preprint 401054241..00154 2024.
[15] Anish Dhir, Matthew Ashman, James Requeima, and Mark der Wilk. A meta-learning approach to bayesian causal discovery. In The Thirteenth International Conference on Learning Representations, 2025.
[16] Lco Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent Bottlenecked Attentive Neural Processes. In The Eleventh International Conference Learning Representations, 2023.
[17] Nir Friedman and Daphne Koller. Being bayesian about network structure. bayesian approach to structure discovery in bayesian networks. Machine learning, 50, 2003.
[18) Nir Friedman and Iftach Nachman. Gaussian process networks. In Proceedings of the Sixteenth conference on Uncertainty artificial intelligence, 2000.
[19] Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional Neural Processes. In International Conference on Machine Learning, pages 1704-1713, PMLR, 2018.
[20] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv: /007v0180. 2018
[21] Tomas Geffner, Javier Antoran, Adam Foster, Wenbo Gong, Chao Ma, Emre Kiciman, Amit Sharma, Angus Lamb, Martin Kukla, Agrin Hilmkil, et al. Deep end-to-end causal inference. In NeurlPS 2022 Workshop on Causality for Real-world Impact, 2022.
[22] Dan Geiger and David Heckerman. Parameter priors for directed acyclic graphical models and the characterization of several probability distributions. The Annals of Statistics, 30(5), 2002.

---

[23] Andrew Gelman, Jessica Hwang, and Aki Vehtari. Understanding predictive information criteria for bayesian models. Statistics and computing, 24, 2014.
[24]
Enrico Giudice, Jack Kuipers, and Giusi Moffa. Bayesian causal inference with gaussian process networks. arXiv preprint arXiv:2402.0063: 2024.
[25] Jonathan Gordon, John Bronskill, Matthias Bauer, Sehastian Nowozin, and Richard Turner. Meta-learning probabilistic inference for prediction. In International Conference on Learning Representations, 2019.
[26] David Heckerman. A bayesian approach to learning causal networks. In Proceedings of the Eleventh conference on Uncertainty in artificial intelligence, 1995.
[27] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, 00s/1606.08415,, 2016. URL http:://axxi, org/ /bs/
1606. 08415
[28] Noah Hollmann. Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfi: A transformer that solves small tabular classification problems in a second. The Eleventh International Conference on Learning Representations, 2023.
[29] Daniel Jenson, Jhonathan Navott, Mengyan Zhang. Makkunda Sharma, Elizaveta Semenova, and Seth Flaxman. Transformer neural processes kernel regression, 2025. URL https
//arxiv. wwt./o/////////ee/eee/oeeeeeee// abs/2411. 12502
[30] Nan Rosemary Ke, Silvia Chiappa, Jane X Wang, Jorg Bornschein, Anirudh Goyal, Melanie Rey, Theophane Weber, Matthew Botvinick, Michael Curtis Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. In International Conference Learning Representations, 2023.
[31] Mikko Koivisto and Kismat Sood. Exact bayesian structure discovery bayesian networks. Journal of Machine Learning Research, 5(May), 2004.
[32] Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal. Self-attention between datapoints: Going beyond individual input-output pairs in deep learning. Advances in Neural Information Processing Systems, 34, 2021.
[33] Jack Kuipers and Giusi Moffa. Partition meme for inference on acyclic digraphs. Journal of the American Statistical Association, 112(517), 2017.
[34] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariin neural networks. In International conference on machine learning. PMLR, 2019.
[35] Haichao Li, Kun Wu, Chenchen Ruan, Jiao Pan, Yujin Wang. and Hongan Long. Cost-reduction strategies in massive genomics experiments. Marine Life Science & Technology, 1, 2019.
[36] P. Lippe, L.. von Rueden, J. Bühler, B. Schölkopf, I. Gurevych, and J.M. Mooij. Amortized inference for causal structure learning. In Advances in Neural Information Processing Systems, volume 35, pages 1737-1751, 2022.
[37] Roderick Little and Donald B Rubin. Causal effects in clinical and epidemiological studies via potential outcomes: concepts and analytical approaches. Annual review of public health, 21
(1), 2000.
[38] Lars Lorch, Jonas Rothfuss, Bernhard Schölkopf, and Andreas Krause. Dibs: Differentiable bayesian structure learning. Advances in Neural Information Processing Systems, 34, 2021.
[39] Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized inference for causal structure learning. Advances in Neural Information Processing Systems, 35, 2022.
[40] David Madigan and Adrian E Raftery. Model selection and accounting for model uncertainty in graphical models using occam's 's window. Journal of the American Statistical Association, 89 (428), 1994.
12

---

[41] David Madigan, Jeremy York, and Denis Allard. Bayesian graphical models for discrete data. International statistical review/revue internationale de statistique, 1995.
[42] Divyat Mahajan, Jannes Gladrow, Agrin Hilmkil, Cheng Zhang, and Meyer Scetbon. Zero-shot learning of causal models. NeurlPS 2024 Causal Representation Learning Workshop, 2024.
[43] Massimo Marinacci. Model uncertainty. Journal of the European Economic Association, 13(6). 2015.
[44] Francesco Montagna, Nicoletta Noceti, Lorenzo Rosasco, Kun Zhang, and Francesco Locatello. Causal discovery with score matching on additive models with arbitrary noise. In Conference on Causal Learning and Reasoning. PMLR, 2023.
[45] Joris M Mooij, Jonas Peters, Dominik Janzing, Jakob Zscheischler, and Bernhard Schölkopf. Distinguishing cause from effect using observational data: methods and benchmarks. Journal of Machine Learning Research, 17(22), 2016.
[46] Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In International Conference on Learning Representa- tions, 2022.
[47] Tung Nguyen and Aditya Grover. Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling. In International Conference on Machine Learning. PMLR, 2022.
[48] Teppo Niiiim, Pekka Parviainen, Mikko Koivisto, et al. Structure discovery in bayesian networks by sampling partial orders. Journal of Machine Learning Research, 17(57), 2016.
[49] Weronika Ormaniec, Scott Sussex, Lars Lorch, Bernhard Schölkopf, and Andreas Krause. Standardizing structural causal models. In The Thirteenth International Conference Learning Representations, 2025.
[50] Judea Pearl. Causality. Cambridge university press, 2009.
[51] Judea Pearl. Interpretation and identification of causal mediation. Psychological methods, 19 (4), 2014.
[52] J. Peters and P. Bühlmann. Identifiability of gaussian structural equation models with equal error variances. Biometrika, 101:219-22 228, November 2013. ISSN 0006-3444 3444.
[53] Jonas Peters and Peter Bühlmann. Identifiability of gaussian structural equation models with equal error variances. Biometrika, 101(1), 2014.
[54] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- tions and learning algorithms. The MIT Press, 2017.
[55] Karen Sachs, Omar Perez, Dana Pe' er, Douglas A Lauffenburger, and Garry P Nolan. Causal protein-signalie networks derived from multiparameter single-cell data. Science, 308(5721): 523-529, 2005. doi: 10.1 126/science. 105809.
[56] Andreas Sauter, Saber Salehkaleybar, Aske Plaat, and Erman Acar. Activa: Amortized causal effect estimation without graphs via transformer-based variational autoencoder, 2025. URL
https://arxiv.org/ abs wttp://ww...g///////./..... 012200 01290
[57] Ran Spiegler. Can agents with causal misperceptions be systematically fooled? Journal of the European Economic Association, 18(2). 2020.
[58] Mare Teyssier and Daphne Koller. Ordering-basee based search: simple and effective algorithm for learning bayesian networks. In Proceedings of the Twenty-First Conference on Uncertainty in Artificial Intelligence, 2005.
[59] Jin Tian and Judea Pearl. A general identification condition for causal effects. In Aaai, 2002.
13

---

[60] Christian Toth, Lars Lorch, Christian Knoll, Andreas Krause, Franz Perkoof, Robert Peharz, and Julius Von Kügelgen. Active bayesian causal inference. Advances in Neural Information Processing Systems, 35, 2022.
[61] Christian Toth, Christian Knoll, Franz Pernkopf, and Robert Peharz. Effective bayesian causal inference via structural marginalisation and autoregressive orders. The 28th International Conference Artificial Intelligence and Statistics, 2025.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. Advances in Neural Information Processing Systems, 30, 2017.
[63] Anna Vaughan, Stratis Markou, Will Tebbutt, James Requeima, Wessel P. Bruinsma, Tom R. Andersson, Michael Herzog, Nicholas D. Lane, Matthew Chantry, J. Scott Hosking, and Richard E. Turner. Aardvark weather: end-to-end data-driven weather forecasting, 2024. URL https://arx// org/ abwp//04..0..00 00411
[64] Jussi Viinikka, Antti Hyttinen, Johan Pensar, and Mikko Koivisto. Towards scalable bayesian learning of causal dags. Advances in Neural Information Processing Systems, 33, 2020.
[65] Stephen G. Walker. Modern bayesian asymptotics. Statistical Science, 19(1:111-1117, 2004.
[66] Y. Wang, L. Solus, K. Yang, and C. Uhler. Permutation-based causal inference algorithms with interventions. In Advances in Neural Information Processing Systems, volume 30, 2017.
[67] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung. Yin Li, and Vikas Singh. Nyströmforner: A nyström-based algorithm for approximating self-attention. CoRR, abs/2102.03902, 2021. URL https:///rai org/ abs/ 22022 03902.
[68] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. CoRR, abs/2007, 14062, 2020. URL https:///axxx
abs/2007 14062
[69] Jiaqi Zhang, Joel Jennings, Agrin Hilmkil, Nick Pawlowski, Cheng Zhang, and Chao Ma. Towards causal foundation model: on duality between optimal balancing and attention. In Forty-first International Conference on Machine Learning, 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We empirically show in section S that MACE-TNP outperforms a range of strong baselines, illustrate its scalability with the higher-dimensional experiment, and pro- vide evidence that its output converges to the Bayesian posterior interventional distribution in the simple bivariate linear Gaussian case.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We provide discussion of the limitations in section 6
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
15

---

Answer: [Yes]
Justification: We provide theorems alongside their complete proofs for the linear Gaussian model in appendix B.1
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed description of the architecture of the model in ap- pendix A provide details about data generation in appendix B and the hyperparameters used in the experiments appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
16

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide anonymised code in the supplementary material. This includes code for data generation for the synthetic experiments, as well as for model implementation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co.
www.ibimmsiiooooooiiiioo.. for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 //nips. www.c/ppwwpppcccccecc htte://wiissseeeeeee www.eueiiisssoonnooooooooo
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide detailed information regarding training and testing the model, including hyperparameters, optimiser details, ablations on model size and attention mecha- nisms, well additional results the ones presented in the main paper in appendix C
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Whenever we report the metric of our choice (NLPID), we also report the eπоо of the mean on the test sets. When reporting the KL divergence in the two-node experiments, we provide the median as well as the 10 90% quantiles.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide hardware details, including type of compute, memory requirements and time of execution, in appendix C.I
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp////.ppppp...ooooo//aaaaaaa
Answer: [Yes]
Justification:
The paper conforms the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiciio))
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Although our method has promising real-life applications, the paper is more heavily focused on fundamental research rather than particular applications.
18

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We include formal citations for each baseline we benchmark-and whose code we use usc-and comply with their published license terms.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
19

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thooe cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide documented code in the supplementary material, as well as detailed description of the model and training procedure in appendix.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involce crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
20

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development our paper does not involve LLMs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy hatt:e///eerrpppccc/oooffeereeee /tt///////oeeeeeee////////.... for what should or should not be described.
21

---

A Architecture
This section provides the definitions of the architectures described in the main paper.
Transformers [62] can viewed as general set functions [34], making them ideally suited for NPs, which must ingest datasets. We begin by briefly overviewing transformers, defining the attention operations and how we construct transformer layer, followed by how we integrate transformers into the MACE-TNE architecture.
A.1
Transformers
MHSA and MHCA Throughout this work we make use of two operations: multi- self- attention (MHSA) and multi-head cross-attention (MHCA). Let Z E N be set of N tokens of dimensionality Dz. Then, for V I.....,N the MHSA operation updates this set of tokens as follows
Zn 5 cat ({E On (Zn, 2m ).m W V.A
Wo:
where WVA RD.xDv and Wo E RDDxxD, are the value and projection weight matrices, H denotes the number of heads, and an is the attention mechanism. We opt for the most widely used softmax formulation
(Zn 2m) = softmax((z,, W WQ
W K.A 4m J 1 m=
(6)
where Wo. E x.DDDDx and WK, K.A E D.Dx are the query and key matrices.
The MHCA operation performs attention between two different sets of tokens Z₁ Nxx, and Za E D.Nx. For V 1, I,...,NI, NI, the following update on 2 is performed:
1in f cat { ===1
N2
"
an Z2.m 22.m)Z2.m 22.m
WWV. v.h
h=1 Wo.
(7)
In order to obtain the attention blocks used within the transformer, these operations are typically combined with residual connections, layer-isations and posnt-wise MLPs.
More specifically, we define the MHSA operation as follows:
Z Z + MHSA(layer-norm (Z))
(8)
Z + Z + MLP(layernormng (Z)).
Similarly, the MHCA operation is defined as:
Ż, + Zı + MHCA(layernoomm, (Zz). layer-norm, (Zz))
Z Z, MT.P(laye--oorm (Z,)).
Masked-MHSA Consider the general case which we want to update N token Z E N.D. There might be some situations where we want make the update of certain token n,, Z independent of some other tokens. In that case, we can specify a set M, C N N containing the indices of the tokens we want to make the update of nn independent of. Then, we can modify the pre-softmax activations within the attention mechanism âr. (Zn,?m)) where an (Zn, 2m) softmax (ār, (Zn, 2m follows:
-00
if m E Mn TT Q.A WT Am otherwise
āi (Z, Zm)
(10)
From the indices of Mn we can construct binary masking matrix M E 10
50 if E Mn
M
[
otherwise
When used the context of MHSA, we refer this operation as masked-MHSS and represent it as Z masked-MHSAL M).
22

---

A.2 Model-Averaged Causal Estimation Transformer Neural Processes (MACE-TNPs)
We refer to Nguyen and Grover 47 Ashman al. for complete description of standard TNP architectures, and focus on describing the architecture of the MACE-TNP in more detail. Our proposed architecture is conceptually similar to the standard TNP architectures, but incorporates specific design choices and inductive biases that make it suitable for causal estimation.
We assume have access Noh. observational samples and want predict the distribution of No interventional samples. The inputs to the MACE-TNP are: the observational dataset Dobs E I Naw x Dxdaee, the values of the node we intervene upon X, E Ig- (implying we intervene on node j), and the outcome node index i. Let Dass. E IR xdea denote the observational data node i. We omit the batch dimension for notational convenience.
Data pre-processing The model takes as input matrix of Nots observational samples of D nodes and an intervention matrix of Nint queries for node of interest Xj, with the rest of the D nodes masked out (by zeroing them out). Let Din,i E R N. Xdae Xdae denote the interventional data at node i. In the following use (ktt((A)) denote nodes in the observational dataset that are being marginalised over.
Embedding To differentiate between the different type variables, employ six different types of encodings, depending on the source of the data (observational (obs) or interventional (int)), and the type of the node (node intervene upon (f), outcome node (i), or node marginalise over). These are all performed using 2-layer MLPs of dimension dembedd
observational, intervention node:
Zoba = MLP obs (Dobs.s)
(1I)
observational, outcome node:
Zoba MLP / (Poss, 2)
observational, marginal nodes:
Zoba. (AA., (AE(DD) = MLPobs (Pobs,
interventional, intervention node:
Zim MLPim., (Din.s)
interventional, outcome node:
Zmn MLPPm, (Pimi,))
interventional, marginal nodes: Zin (A(D)) = ML.Pmm (
where (k E [D] \ {i,j}} represents the set of indices from {1...., D} excluding i and j. The
representations then concatenated back together in the original node order:
Zobs,
if =
Zobs concal ((ED)),)))E;
where De.
Zobs.j Zobs
if j otherwise
Zint,s
if =
Zin concat ((Z(t)))RE(()) where Zu
Zim.ij
if k = j otherwise
Zint,k
After the embedding stage, we obtain the representation of the observational dataset Zobs T.Nas x Dixqueed and the representation of the interventional one Z Nin Dedetedd
MACE Transformer Encoder We utilise transformer-basee architecture composed of L layers, where we alternate between attention among samples, followed by attention among nodes. This choice preserves 1) cermutation-invaiian equivariance with respect to the interventional samples, 3) permutation- invarriacc with respect to the nodes we marginalise over, and 4) interventional nodes. Although we generally omit the batch dimension for convenience, we include in this subsection accurately reflect our inplementation. Thus, the input to the MACE transformer encoder are the observational data representation Zed R x a xDD.DDxxxee Oxdaae and interventional data representation Ziz E B x Nai Dxnand
with respect to the obervational samples, 2) permutation-
with respect to the outcome and
with B the batch size.
23

---

Attention among samples We propose two variants to perform attention among samples. We use the less costly MHSA + MHCA variant for the experiments the main paper and show that performs better in appendix C.2.2
1. MasKed-MHSA among the observational and interventional samples: At each layer I, we first move the node dimension to the batch dimension for efficient batched attention: Z E BB D x BR Bx
(Bx ((BBD)xNNa Xdaated and Z E We then concatenate the two representations [Z Z and construct a mask M R. &a Nit that only allows interventional tokens attend to observational ones.
7 R Min Z' R(BXD)x(NN + X =
M {
if m < Nobs otherwise
We then perform masked-MHSA: z' masked-MHSSA(Z M). This strategy has computational complexity (( N).)
2. MHSA MHCA: An alternative, less costly strategy, is to perform MHSA the observa- tional data, followed by MHCA between the interventional and observational data. More specifically, as in the previous case move the node dimension the batch dimension and then perform:
Z MHSA(Z Z mm MHCA (Z ).
We then concatenate the two representations into Z E ((BxD)x(NNct Nin) Xdented Z ubs Zn This strategy has reduced computational cost of O(N obs + NonsNii and is the strategy we use for the results in the main paper.
Attention among nodes The output of the attention among samples at layer Z D[BxD)x(NNxxt>)))) Xddette is then fed into the next stage: attention among nodes. We first reshape the data Z' E () Na.) dented + Z" E IR (Bx NNu))- Dxdente and then perform MHSA between the nodes:
+ MHSA(Z")
This is then reshaped back into Z R Bx Nui) x Dxdente and then split into the observational and intervational data representations that fed into layer I + Z I BxBxN Dxdend and z!+1 E Bx No. x D2 dentei Int
MACE Decoder We parameterise the output distribution of the NP as Mixture of Gaus- sians (MoG) with Neump components. The NP outputs the mean, standard deviation and weight corresponding to each component for each interventional query {x)) No. (,,a,,)(rr" (xz))) =
These are computed based on the outcome interventional rep- resentation from the final layer of the MACE Transformer Encoder. More specifically, the input to the decoder is ZL E Ni Xdented This is then passed through a two-layer MI.P of hidden size demb. followed by activation function
{Ha (x""), og (x")), ) Wa (x)))
Zout = CitiiiationMMLP(I....
Finally, we use linear layers to project the embedding Zoet E R. Nia Xelented to the parameters of a mixture of Noomp Gaussian components:
Linearmean (Zout) E R Nix Neup
pre-o Linearsa (Zout) R Neou
pre-w Linearweight (Zomt) E R Nix Neeepp
We then apply element-wise transforms to obtain valid parameters:
softplus(pre-o))
W softmax(pre-w)
with the softmax being applicd along the component dimension.

---

Loss The output parameters are then used to evaluate the per-dataset loss of the MACE-TNP, which, the evaluation of the log- posteriotoor interventional distribution of the
as shown MoG. We restate the equation of the loss presented in section 4 for completeness:
Nax
log po( (odo)))(((;)) D)
log
(u(x)) -x(2 (((3)) (12)
where o)) represents the Gaussian distribution with mean and standard deviation
B Data Generation
We provide in fig. 5 diagram showing how we sample training data from a specified Bayesian Causal Model to infer its posterior interventional distribution (see discussion in section
1. Sample causal graph
Sample
Sample interventional data
g rv p)ca(()
observational data
U(L, - (O0, N Dan Pea (,0)
Daue PRcu((Pa.
2. Sample functional mechanism condilioned on graph for each variable
Nata
Nut
<~ (f(0
Figure 5: Overview of the data generation process. We first sample graph 9, and functional mechanism (conditioned on the sampled graph) for each of the D nodes in the dataset. These are then used draw Nobs observational samples. construct the interventional dataset, we first randomly sample a node to intervene upon j, draw Nint intervention values x, (O0 I), and set the values of node j to be Xj. We then drawn Nin samples of each node to form an interventional dataset Dint
B.1 Two-node Linear Gaussian Models
The data generation details for the two-node linear Gaussian experiments from section 5.1 and the derivations of the posterior interventional distribution are explained in this section.
We examine the basic scenario involving n independent and identically distributed (i.i.d.) random vectors, each consisting of two components, defined as X [X XiJ" for E 21.2..... n). Let the observed dataset be denoted by Dobs := {X X2 X"). For the sake of notational simplicity, drop the subscript BCM from PBCM in eq. throughout the subsequent proofs. In this setting, where the random vectors are composed of only two nodes (x), there exist three distinct possible structural SCMs:
0
G1 =
X1 = wX, U1 and X2 U2
(13)
0
G2
Xi Un and X₂ wX1 Up
(14)
0
0
0
G3
X1 = Un and X₂ U2
(15)
0
We consider two models, one where the causal graph is identifiable (appendix B.I.I and one where it is not identifiable (appendix B.1.2).
25

---

B.1.1
Identifiable Case
We begin with the case where the enror terms U1 and U2 are Gaussian distributed and the noise variances of U1 and U2 are equal and known- a setting shown to be identifiable in Peters and Bühlmann [22] Fixing o², o E 1t,, we consider the following hierarchical model:
~ U~~N00,0²)) for 1,2 G G~U 92.93)
~~0,00 if g (91, 92).
which induces the following joint distribution:
p(X, w. 9) = (XXXxx ()p(w(9)(()) i {91. 92) or
(16)
otherwise.
(17)
We show below that the above models can be identified by the posterior.
Theorem B.1. Let Dais == (x(1), x ......
X(m) be i.i.d.
observations generated by The posterior over the graphs
one of the simple models described in eqs. (3) 10 15
(p(Gi(Does) ) 22 (obs), Dobr]]
(o2.52
exp
2x
S12 oZ S2
o²
V(o2
σ
02)
exp
20
02
S σ2
S
12
1
where e is a constant of normalisation and
=1 = X $
Σ
x
S2
Σ x
(18)
The posterior interventional distribution is mixture of Gaussian distributions
p(X, = y|do(X2 = x), Dots) = P(Gi (1) [f + (1- 1e1-ppan
with
S1
22 σ S2 + σ²
11 (1-2(x)=
02 S2 a² x and of-(2x) = 0²
(19)
First, find the full conditional distribution, p( w G Dans) and the posterior distribution over the DAG models, p(g Dots). Following Bayes' rule we have
Proof:
p(w|g1, Dens) ((wll))p(G)) II
T p(X:191, p(w|gi)
II x(Xx w),w)
1=1
==1
w
1
x wxx)
exp
20
2
t=1
and by completing the square we obtain
o2S12 S2
of
p(w|G1, Dote)) -N
(20)
+ of
S2 S2 of
with S12 and S2 being defined in eq. 18))
26

---

Similarly, conditioned on the G2 model, the full conditional distribution is again Gaussian
o² S + 0² o o +
S2
D(w]02. Dabs) = =N Using eq. (20 and eq. (21 the posterior for the G is
(21)
0²
p( x p(w))(G() II 1( w)d
S 51+
02
Si2
x exp
CXP
10
(22)
202
Vo + o²
20
o S2 + o²
p(gg)Doos) x ((w)p(92)
II
p 9(,, w)dw
+1+52
0²
S7
x exp
20²
V +
exp
20 o S + 0²
(23)
n1 P(os)Dots) & p(9a) II
81 + S2
p(X.(9s) x exp 1=1
(24)
20²
Next, the posterior interventional distribution is
p(Xildo(X2 =x), Dobs) p(Xx Ga)o((z)) Dobs)
E
/ G, wldo(x), Dors)
9E(01.9a)
= p({x)do(z), Gs. Doss Dess) +
Σ
o()) /
9, Debs ddw.
9e(9
Conditioned on the model graphs, the interventional distributions are
p(x1 y//ooXX 91, w) (wylwx, p(x (j = 02, = NN(10,
p(X1 y = x)) NN(g0,,
Then, we have
/ |Dxx y|do(X2 = x), G2, w p (92, Dota f N(3)0,00))pte G2, Doma)dw = (oy10,00)
( p(xx g/do(X2 2), Gs, w) w)p(w(as Doba )dw / (()10,00)pp( Gs. Dots)dw N3))0,
and
J v/do(X2 yldo== x)), Gi, w))(wpli, Dama)dw- NollooNNNlluuus (l)/w,,,o)p((wle Dow
-N(1
S1 02 + x, 0²
0222
1+
S2
o²
Hence, the interventional distribution is simply mixture of 2 Gaussian distributions
p(Xi y/2o(X2 x), Donss) = pp( Da W (1ln- + p(9 Do.))
27

---

with (Dobs) calculated in eqs. (22) to (24) and (1-2(z), 2 defined in 34
Similarly, the next result easily follows
p(X, yydo(XX z). Dats) = (
(1- pl9g/Dam))
with ((G2)Danbi calculated in eq. (33 and
S1 122 12-1(x) S o²
and 02- = 0² of( Bus o + +
Remark: can be shown that if Pobs is generated by one of the models presented eqs. to (15). then the posterior distribution p( Dots) asymptotically concentrates around the truc data-ged generaannn structure 9*
12 65 14]. Consequently, in the infinite data limit, the posterior interventional distribution converges to a Gaussian distribution whose mean and variance depend on the intervened node and the true underlying causal mechanism g*.
B.1.2 Non-identifiable Case
Second, we consider the errors' variances to be unknown while keeping the same SCMs described in eqs. (T3) to (I5). Therefore, we place priors on these extra parameters well chosen such that the model is not identifiable [22]. We propose the following hierarchical model for fixed a >> 2 and B,
9~ G2, 9s}
If = G1, then rî InvGammala,, 3), - ~ InnGammaaa 2 3), ~ nr),
1
If g = 92, then r TnoGamma(a
2' T2 IneGamma(a, 3), (0,
3),
1
If Ga, then 7 InuGamma(a
2' 3), r IneGamma(a, 3).
Then, this hierarchical model introduces the following joint distributions
If g 91, then p(X, w, 7, -2. 91) = p(X wu 7 -2 G1)p(w(( TPp(( If G G2, then pl X, w. 7,, 7g: (.: (2) = -(XX G.) 1
If G = Gs. then p(X, -.. -2, Ga) = p( (
For completeness, we show below that the above priors result in the same posterior for graphs in the same Markov equivalence class - Si and G2. We begin by recalling simple result before stating the main theorem of this subsection.
Lemma 1. For any > and A, B,C C E IR such that CA² we have
dx
r(
(A2) -d]
X
(APx2 2Br + C)
(CA2 B
where T(-) is the usual Ganma--uuncion.
PRoOF: Completing the square in the dominator we have
10
dr
A²
dr
(ARx² 2Bx
=
CA22
B
-00
2
A
CA2
B
A2
+
28

---

Next, we recall the probability density function (pdf) of shifted and scaled version of the standard student-t distribution (i.e. a oT, with T t(2))
)( VT 1 ($02 0²
(2=1)
))
fz(z)
Then, matching the terms gives the desired result.
Theorem B.2. Let Dobs == (x(1)x((1) ..... [1) be i.i.d observations generated by
one of the simple models described in eqs. (13 to (15).
The posterior over the graphs
(p(Gi(Dobs) pl92 (ots),
(ot) P
(S 1w-1)/2 11/2
(
(S
(/ 1 S 1/2 (5
10)2 0/22 (0))(-11)/2 (5)
where c is the constant of normalisation, = 2a + n and
S.¹ = S, 23 (1,2), =
(25)
with S1, S2 and S12 defined in eq. 18 shifted and scaled Student-t distributions
The posterior interventional distribution is a mixture of
S
p(1, = y/do(X2 = x), Dobn) = MiG + p(gi)Doou)) v. V v V-
with
(1-2 = xS12 S
and 2--(())= == SS
5 + x 2 5%22
(26)
(5
Similar to the presented in Appendix B.I.1 we start by deriving the posterior over the three models described above and use the same definitions from eq. (8
Proof:
p(Ci)Dotss) x / p(Dooall1 S1,, w, (ow.
d(T2) dw
=/ (V2nT?)
exp - 2- iml
Σ
(xi -xxx))
1
1
S2
n exp
(V2TTz)
2T
v /
exp (
2
320-
(72)
-a-1 exp
[T2)
-a 1/2
exp
3 d(r d(T dw
2
I(a)T
1
TI
1
B2a-1/2 1/2
P(a+ + )
I(a +
dw
VT (2m (2m) 1+122 1aaln/////aaaaaaaaa_
(8+ 82
(u-1)/2
(B+
In -ui--
at
+
2
2n
(2 ir(a + = 4)
(1/n S2) (2)(111//2
1
T( ))
(29 + (u-(11)//
[23 + S1)((11) + S2)
S22 jatn/2
(23) tT(a - [(a 2) (() ))
S (v- 1)/2
(27)
(S S12)
where in the last step we used the result of Lemmal fulfilled as
I
We note that the conditions in the lemma are
S,S² 2 (S12) 2
(23 + S1)(1n S2)
(28)
29

---

where we used the fact that n, > 0 and the Cauchy-Schwarde inequality in the second step. Similarly, we find the posterior for the second model, 92
(23)2 -ara-rra+ + " 5) Vin-T(a))(a- ;) 59
(5
(u-1)/2
(ott)) x and for the third model, 93
(S) ST
512)
p(ggg/Dtaa) x / (2))) d(r2)
=
exp (-
277 S1
1
Sz
exp
(22T2
2T
82a-1/2
B
(T1) 1 exp
(T2)
2 ---
exp
d(
d(ri)
(29)
((a)T(a- 2)
72
20- 1/2T 223))a-1123
+ n-1 ]T + 2) T^T(a)T(a- 2)
2
(30)
(())((11)/2 1//2
Conditioned on the models, the interventional distributions are
p(X1 y)) 9I, w, Tt) =- p(x y)do(X2 2), 91, w, +2) Nyllwx Tf),
p( pxx [1 y)do(r). G2, w. 7ी) p(X, y)do(X2 z), G2, w, Tt) NNy10,
p(x = y)do()), -2) := p(Xi = y|do(X2 = x), Ga, 7 (,)
Then, the posterior interventional distribution, (dX))) Dots), is
y/do(x), Σ D(G)D / p(xi= - ",,,,,,pppp T2. G, Dota) d(r?) d(r2) dw
GE .a)
+ p(GalDota) w/lo())) pp)( 221, Ga, Data) d(r?) d(T2))
First, we find the last term
I p(x2 = =)) pp) 721.
1. Dow.) d(T?) d((T)) &
x),
/
1 [VIT exp (-
122 27
p (D So, 22
d(tz) x
V2
x
(2g + S + 12)aat"""
x (1+ aa
y² S
(31)
which a scaled Student-t distribution with = 2a + n degrees of freedom and o²
S
D-1
Computing the next integral follows the same pattern presented in eq. (30)
/ p(x1 = =y/oo(X2 = =x),
= G2, w. 7) 7i, 92, Data) d(rs) d(T2) du x
/ exp (-
y2²
pl (D w, (3)
d( d( dw x
/2TT
2-
(S1 + 1 1)a+
x (1+ x² 2+ (28 + (Si + -))1(())2
(32)
(28 + S y ) "y"
sa
30

---

and y degrees of freedom. Finally, which is again a scaled Student-t distribution with o² Si using similar steps in eq. (27). the last term --1
/ p(x1 v/do( X2 = x). 91, p/p(w, (i r.. 1221:11 91: Dobs) d(ri) ((1))((1)) dw x
/
exp (- [y (
wx)2
pl (Dbuuluu w. (2, (otpomp)pppp()))))))) d(ri) d(r dw
2TT
2
1
(33)
(S2 x2 (2 + S1 v2) (S12 + mzy) a+
where the denominator is always bounded away from 0 as in eq. (28). By completing the square, we obtain another scaled and shifted student-t distribution with v 2a + n, (1-2(z) and 2 defined in eq. (34. Then, combining eqs. (27) to (33) we obtain the result.
Similarly, we show
p(X2 y)do(x)) Dots) p( M1s
p(9g Dobs)))
B,0. v
with
x.122 S
and 02 (2) 59 - S 1 -S
(2-1(x):=
(34)
(5
We employ asymmetric priors for of and T ensure that the posterior assigns equal probability to the models &1 and G2, which belong to the same Markov equivalence class. In particular, setting 233 = results in SW = S and S = S2 which implies that the posterior distribution over graphs,
Remark:
(p(9) Dots), p(92 Dons); p(Gs Dohx)]
takes the form:
1
(35)
( 8% 1/2 (
ST.) 0/22
where c is a normalising constant.
This setup corresponds to the prior structure proposed by Geiger and Heckerman 22 Equation 12). obtained by setting the precision matrix T the Wishart distribution the identity. As noted in their Geiger and Heckerman [22) Section 4), change of variables transforms the Wishart prior on the covariance X into the prior used here for the weights w and error variances of and T
Assuming a true data-generating mechanism, the posterior concentrates on its Markov equivalence class. If the true graph is G, G2, then p true graph, then p(Ga (oons) + 1 65 [4].
p(G2|Pand 1/2. Conversely, if G3 is the
The degrees of freedom, distributions converge to Gaussians in the large-sample regime.
n. grow with the sample size the corresponding Student-t
B.2 Three-node Experiments
In the three-node experiments (section 5.3) we use two datasets with two different functional mecha- nisms fz( as defined in eq. (T): one sampled from a GP prior, and one based neural networks. In both cases, we sample Frdös-Rényi graphs with graph degree chosen uniformly from 21. Following Ormanicc et al. [49], standardise all variables upon generation.
31

---

GP functional mechanism To model fz(.) we use GP with a squared exponential kernel, with randomly sampled lengthscale for each parent set PA, of size [PA,]. More specifically, we sample the lengthscale from log- distribution (Ap)p=) IP. ~ Log 1), followed by clipping between Ap clip( pp, 0.1,5) 5) to ensure that too long lengthscale does not result in independence of the variable from a parent. This defines the kernel matrix between the n-th and m-th samples as:
K Kmm = exxp(((AA3 PAA")"
(PA ))
with A = Diag( A1 .i...... the variables as follows
A APA)) We then add noise with variance o² ~ Gamma(1, (1, 5) and sample
X; ~ (0,
Neural network- functional mechanism
We sample each variable as follows
T NN 1)
X ((es ul) + oe,
where d² Gamma((1 10), (O, 1). Res residual neural network with randomly
sampled number of blocks Nhlocks ~u 111,..... 8} and randomly sampled hidden dimension dhidden 4122; 26, 27, 2 We the GELU [27 activation function.
B.3
Higher-dimennional experiments
For the higher dimension experiments section 5.4 follows:
we generate the training data for MACE-TNP
We sample number of variables D ~ 40).
We sample type of graph, either an Erdōs-Rént graph or a scale-free graph 3.
The density of the graph (number of edges) sampled from 24 G6].
For each node, we sample functional mechanism randomly from either GP with additional latent variable input, or a Neural network with an additional latent variable input:
GP with latent: We sample latent 1 AO0 1), and lengthscales {Aplp=l IPAd++ Log(-0.5, 1), where PA, denotes the set of parents of node index i. Functions are sampled from squared expoenential kernel with y included as an input and Gaussian noise added with variance o² Gamma(1 (1, 5).
- NN with latent:
7h ((, 1) X ~ NNo m)]) + de,
where o² Gamma(1, 10), c~ (O, 1). NN, denotes randomly initialised neural network with 128 hidden dimensions and one hidden layer.
Using latent as an input ensures that the final distribution is not Gaussian. Following Ormaniec et al. [4 we standardise all variables during the data generation process.
For testing for each variable size in table we only generate Erdōs-Rényi graphs with density 4D. This is to test the performance of the baselines and our method the difficult dense graph case. The rest of the data generation process is the same the training data.
32

---

C Experimental Details
This section provides additional details and results for the experiments presented in Section
C.1 Architecture, training details and hardware
Throughout our experiments we use H = 8 attention heads, each of dimension Do = DRV
dmodet/& /8. The MIPSs used in the encoding use two layers and hidden dimension of dembed moddel- Unless otherwise specified, we use a learning rate of 5 x 10 with a linear warmup of 2% of the total iterations, and batch size of 32.
To train MACE-TNP. we randomise the number of observational samples Nots ~> (4500 750), and set Nm = 1000 Nots- The training loss is evaluated on these Nm samples. For testing, we sample 500 observation points and compute the loss against 500 intervention points.
Two-node linear Gaussian model We use L 2 transformer encoder layers, where each trans- former encoder layer involves the attention over samples, followed by attention over nodes. The model dimension is dinueel 128, and feedforward width dn 128. We train the model for 1 epoch on 50.000 datasets and test on 100 datasets. Training takes roughly 60 minutes on a single NVIDIA GeForce RTX 2080 Ti GPU 11GB, and testing performed in less than 5 seconds.
Three-node experiments For the experiment in the main paper, we use I. = 2 transformer encoder layers, model dimension dinadel 128, and feedforward width dp 128. We train the model for 2 epochs on 50.000 datasets for the GP experiment and 100.000 datasets for the NN one, and test on 100 datasets in both cases. When testing the OOD performance, we train on the union of the two datasets for 2 epochs. Training the models described in the main text required roughly 4 - 6 hours of GPU time; however, because we ran them on shared cluster, actual runtimes may vary with cluster utilization.
Higher dimensional and Sachs experiments For the higher dimensional experiments we use L = 4 encoder layers. The model dimension is dmodel = 256 with feedforward dimension du 1024. We train the model on data generated as listed in appendix [B.]] with 2, 500, 000 datasets in total. The model was trained on an NVIDAI A100 80GB GPU for 2 epochs which took roughly 20 hours. We use the model trained for the higher dimensional experiment for the Sachs experiment.
Hardware For the two- and three-node experiments, we ran both training and inference on single NVIDIA GeForce RTX 2080 Ti (11 GB) with 20 CPU cores on a shared cluster. The only exception was for our largest three-node GP and NN models (with d'model = 1024), where we used a single NVIDIA RTX 6000 Ada Generation (50 GB) paired with 56 CPU cores; those models required roughly 25 GB of GPU memory. For the higher-node experiments, we used a single NVIDIA A100 80GB GPU, as well as an RTX 4090 24GB GPU.
C.2 Additional Results
C.2.1 Two-node Linear Gaussian Model
We study the performance of MACE-TNP in both identifiable and non-identifiable causal settings by generating data according to the models described appendix (B.. and appendix 1.2 respectively. For all experiments, we set = ow 1 in the identifiable case and = 3, 1/ 2,3 = 1 in the non-identifiable case. We investigate at different interventional queries, x, how the NP predicted function of the observational sample size. For simplicity, for the first model B.I.I we consider an NP which outputs a mixture 2 components, while for the second model B. the NP approximates the true analytical distribution using components.
distributions compare with the analytical ones
The flexibility of our architecture also allows for conditional queries, multiple interventions, as well as easily incorporating interventional data to help identify causal relations. Hence, we investigate here whether providing small number Min of true interventional samples, alongside the observational data, resolves identifiability challenges in the non-identifiable case. As already shown in fig. 3 (right) with the green line and discussed in section 5.1
we find that adding extra interventional data
33

---

does indeed lower the (2; f, (=)))o(( D Dobs { } that even limited interventions can enhance identifiability. We also test this with an increasing number of interventional samples in fig. 6 As soon as the interventional information is rich enough (Mint (50, (. the NP recovers the interventional disribution of the true data- aata-generating mechanism even with little to no observational data, as indicated by the near-flat KI. curves. We note that the KL divergence between two Gaussian mixtures (or between a Student-t mixture and a Gaussian mixture) lacks closed form expression. Therefore, we approximate it in our experiments by averaging 1000 Monte Carlo estimates of the log-density ratio.
)), suggesting
Then, we show in fig. 7 two examples where the intervention is made at for the identifiable model and at r 2 for the non-identifiable model. clear distinction is observed between the two settings: for the identifiable case, the analytical posterior interventional distribution is a mixture of two Gaussian distributions, which, at high observational sample sizes, converges to a single Gaussian (i.e. because the observational data gives information regarding the causal structure, the weight corresponding to one mode collapses 0). In contrast, for the non-identifiable case, the posterior places equal mass on both 9 and G2. and therefore, the mixture structure persists across both regimes. In both settings, the NP-predicted distributions closely match the correct interventional distributions, with accuracy improving as the number of observational samples increases. This improvement is due to two factors. First, larger sample sizes provide the NP with more information about the underlying causal model, allowing for enhanced inference. Second, in the non -dnendentifi case, the posterior interventional distribution a mixture of two Student-t distributions with a number of degrees of freedom proportional to the number of observational samples. Thus, in the high sample regime, the mixture distribution converges to a mixture of Gaussians, which is the class that parameterises the output of the NP model.
0.5
KL (PhEM PMACE-
0.4
KL divergence
KL (PhaM PMMACE-TN)), Min 5
0.3
KL (PhcM PMACE-TNU Min 50
0.2
KL (PROM PMACE TNT), Min 300
0.1
0.0
10
102
Observational sample size
Figure 6: Average KL divergence between the interventional distribution of the true generating mechanism, {9*, f*), and the NP-predee NP-predicted distribution shown function of the observational sample size for the non-identifiable setting. Results are shown for various interventional sample sizes. For simplicity, only report the medians.
34

---

Posterior Interenttonal distribution x=1 (Phamples
Posterior interyend distiburiinn (Asempples
weight
weight 88,
weight M:19:M)-0-2; M)-022211300
0.20
RP prediction
Anslytical
6.24
weight D.48: Ni2 0.581
NP prediction
Dengit
0.25
Dersity
Analytical
0.20
6.14
0.20
0.55
0.00
Outcome
Fosterfor interventional distibutton
i#samples
Posterior Interventionl distrbuttonn (#somptes
0.40
weight
1.49
0.35
weight NJC. 87, 931
weight 0:39: MI-0-12, 1.52)
NF peediction
weight D.I6; NJ2 6D. 0.59
0.20
Anslytical
NP prediction
Denoit
Density
Analytical
6.15
0.55
0.00
6.05
Outcome
Interventional distribution
100)
distribution
0.40
woight
1.17)
0.25
weight 0.04 NO 84, 0.54)
weight 0:41: 6.11, 1.12)
MP prediction
weight 0.45: N12.4b 0.451
0.30
Analylical
NP prodiction
Density
0.22
Ansiytical
0.13
0.25
0.11
0.10
0.55
6:35
0.00
0:05
Outcome
Oulcome
Interventioall distibutionn
distributinn
0.60
weight
1.331
0.25
weight 0.07. NO. 0.94)
Weight 0.39: 0.11, 1.13)
MP prediction
Weight D.45: 39, 0.544
Analylical
0.22
NP prodiction
Deroial
0.30
0:05
0.50
0.00
Outcome
Outcome
Figure 7: Fitted NP posterior interventional distributions true posterior interventional distributions for identifiable (left) and non-identifiable models (right) at increasing observational sample sizes (5, 50, 100, 500).
C.2.2 Threc-node Experiments
In this section we provide additional results on the three-node experiments where we aim to address three questions: 1) between the MHSA and MHCA schemes for sample attention introduced in appendix A.2) which one performs better? 2) Does increasing the number of MoG components improve performance, and 3) How does the model performance vary with the size of the architecture?
3 the results for the two functional mechanisms used in the three-node experiments: GP and NN-based. For each model configuration, we present four sets of results: for model trained on GP and tested on (GP/ / GP), a model trained on NN and tested on NN (NN NN), and for model trained on the combination between the two datasets and tested on each of them (GP+NN / GP and GP+NN NN). These results allow us assess whether the influence of model architecture is consistent across the functional mechanisms, Notably, models trained on the combined GP+NN
35

---

dataset are able to match-within error the performance of models trained specifically on either GP or NN data. This highlights the strength of the meta-learning approach: even when trained data generated from diverse functional mechanisms, a single model can generalise effectively across both, achieving performance comparable specialised models while also benefiting from broader prior coverage. We summarise the findings from table 3
1. MHSA + MHCA outperforms the masked-MHSS strategy for attention over samples.
2. Increasing the number of MoG components increases the performance of MACE- -NP. There is larger gap in performance when going from 1 to 3 mixture components, indicating the importance of allowing the model to output non- non-Gaussian marginal predictions. Increasing the number of components to 10 further improves performance, but the gains not as sigmificant.
3. Scaling up the model architecture generally leads decreased NLPID.
4. Training a model on the combination of the two datasets (GP+NN) is able to recover- error- the performance on both datasets.
Table 3: Results of MACE-TNP under different architectural configurations. M-SA stands for Masked-MHSA, while SA+CA indicates the MHSA+MHCA attention mechanism. For each model, the column trained the model on the GP+NN dataset and tested it on the GP one). We report the mean + the error of the mean of the NLPID over 100 datasets.
under NLPID indicates the training set / test set (i.e. GP+NN / GP indicates
NLPID ())
MoG
Attention
dmodel
dR
L
GP/ GP
NN/NN
GP+NN GP
GP+NN / NN
1
M-SA
128
128
4
629.0 20.0
664.1 16.0
640.1 17.3
668.3 17.2
1
SA+CA
128
128
1
617.4 + 20.1
661.9 16.4
629.8 + 17.5
688.0 = 31.5
3
M-SA
128
128
4
581.8 21.8
538.9+197 19.1
+
547.1 + 17.8
3
SA+CA
128
128
4
540.5 17.1
28..1
540.6 18.8
10
M-SA
128
128
21.9 4
533.2 18.3
599.4
531.5 19.6
10
SA+CA
128
128
4
563.9
527.9 19.8
21.
531.0 19.4
10
SA+CA
512
256
555.7
527.0 19.1
6
532.1 18.4
10
SA+CA
1024
558.0 + 23.9
518.2 19.7
365.6 22.3
521. 20.8
Finally, table 4 summarises the results for the out-of-distribution (OOD) evaluation for the configura- tion presented in the main text.
Table 4: Results for the OOD two-node experiment. We show the NLPID (L) and report the mean + the error of the mean over 100 datasets. Each row corresponds different functional mechanism used in the test set (GP/ NN).
Training
Test
GP
NN
GP+NN
GP
563.9 + 23.4
678.0 = 10.0
583.9 +: 21.5
NN
608.3 + 17.3
527.9 + 19.8
531.0 19.4
C.2.3 Sachs Full Results
The full set of results for the Sachs dataset are shown in table 5 The MACE-TNP performs competitively with DECI, and both outperform other methods that use GPs as functional models.
36

---

Table Results for the Sachs dataset [55] We show the NLPID (4) and report the mean the error of the mean across 5 interventions and across all 10 nodes used as the outcome for each intervention. Each row corresponds to different baseline.
Sachs
MACE-TNP
998.9 104.9
DiBS-GP
1417.5 = 186.7
ARCO-GP
1400.7 208.7
DECI
1000.9 + 133.5
NOGAM+GP
1763.7 297.4
37


Reviews:
[{"review_id": "wE7yWmckoX", "reviewer": "Reviewer_vjb8", "summary": "The paper introduces MACE-TNP, a meta-learning framework designed to estimate interventional distributions under uncertain causal graphs by directly mapping observational data to Bayesian model-averaged interventional distributions using a Transformer Neural Process. Unlike traditional methods that rely on computationally expensive graph averaging or two-stage approximations, MACE-TNP leverages end-to-end training on synthetic data from Bayesian Causal Models to amortize inference. The model features variable-specific MLPs, a permutation-invariant transformer encoder, and a Mixture of Gaussians decoder, trained to minimize KL divergence from true interventional distributions. Theoretically sound and empirically validated across settings from small-scale (2–3 nodes) to large-scale (20–40 nodes) systems, MACE-TNP outperforms existing Bayesian and non-Bayesian baselines and shows promise on real-world data (e.g., Sachs dataset). Despite its high training cost and sensitivity to out-of-distribution functions, it offers a scalable and accurate solution for causal inference under uncertainty, marking a significant step forward in the field.", "questions": "See above weakness", "limitations": "Generalization may be a limitation.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n1. **Computational Efficiency**\n   - Avoids expensive two-stage Bayesian inference by amortizing computation into a single forward pass\n   - Scales to moderately high-dimensional systems (up to 40 nodes)\n\n2. **Theoretical Foundations**\n   - Provably converges to true Bayesian posterior in identifiable cases\n   - Properly handles both identifiable and non-identifiable graphs\n   - Incorporates interventional data to resolve non-identifiability\n\n3. **Architectural Innovations**\n   - Transformer-based neural process handles permutation symmetries\n   - Variable-specific embeddings distinguish intervention/outcome roles\n   - Mixture of Gaussians decoder captures complex distributions\n\n4. **Empirical Performance**\n   - Outperforms Bayesian baselines (DiBS-GP, ARCO-GP) on synthetic tasks\n   - Beats non-Bayesian methods (DECI, NOGAM+GP) in accuracy\n   - Shows robustness to functional misspecification\n   - Validated on real-world Sachs dataset\n\n5. **Practical Design**\n   - Includes ablation studies and architecture analysis\n   - Discusses limitations and future improvements\n\n## Weaknesses\n\n1. **Fundamental Limitation: Latent Confounding**\n   - Cannot handle unmeasured confounders (assumes causal sufficiency)\n   - With latent variables, graph space becomes infinite\n   - Theoretical results suggest this case may be undecidable as it may require the model to learn a classifier of infinite number of classes implicitly (https://arxiv.org/pdf/2411.17898)\n\n2. **Training Requirements**\n   - Needs large amounts of synthetic training data (minor)\n   - Computationally expensive meta-training phase (minor)\n   - Performance depends on match between training and test distributions (minor)\n\n3. **Generalization Challenges**\n   - Limited out-of-distribution robustness\n   - May not extrapolate to unseen causal mechanisms\n\n4. **Scalability Constraints**\n   - Attention mechanism has quadratic complexity\n   - Untested on very large graphs (>100 nodes)\n   - Memory intensive for high-dimensional interventions\n\n5. **Methodological Comparisons**\n   - Lacks comparison to recent table-pfn based approaches such as causalPFN(https://arxiv.org/pdf/2506.06039)\n   - Doesn't benchmark against latent variable methods\n   - Limited discussion of failure modes", "was_revised": true, "final_justification": "Several of my concerns has been resolved and thus I will maintain my positive score.", "rebuttal": "We thank the reviewer for their thoughtful and constructive feedback. We’re glad they appreciated, among others, the **theoretical foundations** of our work, the ability of MACE-TNP to **avoid expensive two-stage Bayesian inference** through amortised learning, and its **strong empirical performance** across both synthetic benchmarks and real-world data. \nWe appreciate their comment that our approach **\"offers a scalable and accurate solution for causal inference under uncertainty, marking a significant step forward in the field\"**.\n\n> Fundamental Limitation: Latent Confounding\n\nAs also noted in our response to Reviewer nB8D, this is a valid concern. While we assume causal sufficiency to ensure a fair comparison with existing baselines, our method does not inherently require this assumption.\n\nIn fact, MACE-TNP **can accommodate hidden confounding** by training on data generated from Bayesian Causal Models (BCMs) with unobserved confounders. To demonstrate this, we conducted an experiment using a three-node linear SCM. We trained two MACE-TNP models: one under observed confounding and one under hidden confounding. As expected, the model trained with observed confounding accurately recovered the interventional distribution, while the one trained under hidden confounding appropriately expressed higher uncertainty (similar to Figure 3), reflecting the fundamental identifiability limits.\n\nAlthough we're unable to include figures in the rebuttal, we will include these results in the revised manuscript to illustrate the model’s behaviour under hidden confounding, and thank the reviewer for highlighting this important point.\n\n> Training Requirements\n\nThis is indeed another limitation of our approach that we acknowledge in L371. However, we argue that this strategy is consistent with recent successful approaches in the machine learning community—such as TabPFN [1], a state-of-the-art method for tabular data that relies on large-scale synthetic training.\n\nCrucially, although training requirements are high, inference-time requirements are light, requiring a single forward pass. This makes our approach particularly appealing in settings with frequent inference-time queries, where the upfront training cost is effectively amortised.\n\n[1] Hollmann, N., Müller, S., Purucker, L., Krishnakumar, A., Körfer, M., Hoo, S. B., Schirrmeister, R. T., & Hutter, F. (2025). Accurate predictions on small data with a tabular foundation model. Nature.\n\n> Generalization Challenges: Limited OOD robustness, may not extrapolate\n\nWe also acknowledge this limitation in L373 and thank the reviewer for pointing it out. However, we would like to stress that\n- When tested on real-world data (Section 5.4), our approach is competitive, suggesting that it is possible to generalise to real-life data better than the considered benchmarks.\n- Robustness to prior misspecification can be improved by training on a broader mix of prior classes (see Appendix C.2.2). While selecting appropriate priors for unknown tasks is non-trivial, this is an active and emerging research area [2] which we plan to explore in future work.\n\n[2] Zhang, X., & Maddix Robinson, D. (2025, July 22). Mitra: Mixed synthetic priors for enhancing tabular foundation models. Amazon Science. Retrieved from Amazon Science blog.\n\n> Scalability Constraints: Attention mechanism has quadratic complexity\n\nWhile we currently use full attention, the architecture of MACE-TNP is not inherently constrained to this. It can readily incorporate efficient attention mechanisms---such as linear or kernelized attention [3, 4] and low-rank approximations [5, 6]---to significantly reduce computational cost with a relatively low impact on performance.\n\n[3] Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020, November). Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning (pp. 5156-5165). PMLR.\n\n[4] Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., ... & Weller, A. (2020). Rethinking attention with performers. arXiv preprint arXiv:2009.14794.\n\n[5] Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.\n\n[6] Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., & Singh, V. (2021, May). Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence (Vol. 35, No. 16, pp. 14138-14148).\n\n> Scalability Constraints: Untested on very large graphs (>100 nodes)\n\nTo address this concern, we conducted an additional experiment using the model from Sections 5.3 and 5.4, which was trained on graphs with 5–40 nodes. We evaluated it on a new test dataset consisting of 100-node graphs with approximately 400 expected edges. Importantly, the new dataset was only used at test time--**no samples from the 100-node dataset were included in training**. We obtained the NLPID for our model and for DECI, the most competitive baseline. Whereas our model has a NLPID score of $671.4 \\pm 5.1$, DECI achieved a NLPID score of $691.8 \\pm 5.7$. With this additional experiment we prove that \n- 1) we are able to scale the model to a higher-dimensional setting;\n- 2) when tested on this higher-dimensional setting, our model continues to outperform the strongest benchmark, even if this **testing scenario is out-of-distribution for our model**.\n\n(We used one 24GB memory GPU for this experiment).\n\n> Methodological Comparisons\n> a) Lacks comparison to recent table-pfn based approaches such as causalPFN (https://arxiv.org/pdf/2506.06039)\n> b) Doesn't benchmark against latent variable methods\n\nWe thank the reviewer for highlighting do-PFN, which is indeed a closely related contribution. Its appearance reinforces the growing interest in addressing interventional queries / causal effect estimation using meta-learning approaches such as Neural Processes and Prior-Data Fitted Networks.\n\nAs do-PFN was released on arXiv **after the NeurIPS submission deadline**, we did not include it in our comparisons. According to NeurIPS policy, comparisons to works appearing after March 1, 2025, are not required (see the NeurIPS-FAQ).\n\nWhile do-PFN offers an accessible interface, it is currently limited to binary treatments in the context of CATE, and cannot be directly applied to our broader setting involving continuous treatments and general interventional distributions. Conceptually, both methods aim to learn the same target quantity but differ in architectural choices (e.g., treatment/covariate encoding, model scale, decoder parameterization (MoG in MACE-TNP vs. categorical distribution in do-PFN, etc.).\n\nWe appreciate the connection and will discuss it in the revised version.\n\n> c) Limited discussion of failure modes\n\nWe appreciate this point. In Lines 372–373, we do acknowledge key limitations, including out-of-distribution generalization, training cost, and scalability. Additionally, we plan to include the hidden confounder experiment in the revision to directly address the limitation related to latent confounding. If there are any other failure modes the reviewer thinks we have not properly addressed, we are happy to discuss them.\n\n---\n\nWe would like to thank the reviewer once again for their insightful feedback and, if we sufficiently addressed your concerns, we would appreciate it if you could reconsider your score."}, {"review_id": "SAph0GWm51", "reviewer": "Reviewer_RHFT", "summary": "This paper studies an interesting problem, directly estimating the Bayesian model-averaged interventional distribution from observational data with a meta-learning framework. It addresses the major challenge of computational intractability in averaging over uncertain causal graphs and functional mechanisms, which is critical when the causal structure is not known or is only partially identified. This paper proposes MACE-TNP, a Transformer Neural Process architecture that amortizes the Bayesian inference pipeline, and demonstrates empirical convergence to the true interventional posterior under both identifiable and non-identifiable settings. This paper shows scalability to high-dimensional data (up to 40 variables) while outperforming strong Bayesian baselines.", "questions": "1. How does the model behave when the test data is generated from functional mechanisms not covered during training? Are there known inductive biases that improve generalization?\n2. Could MACE-TNP be extended to support counterfactual queries? What modifications would be needed in the architecture or loss?\n3. What is the sensitivity of the model to initialization, training set size, or hyperparameter choices? How robust is the meta-learning process?\n4. Have you assessed how well-calibrated the estimated interventional distributions are? Could calibration methods be integrated during training?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: 1. The paper correctly identifies a fundamental bottleneck in causal inference, making decisions based on a single inferred graph leads to overconfidence and incorrect interventions. The use of Bayesian averaging over graphs is sound and aligns with established theory. 2. Instead of traditional two-stage Bayesian inference (graph posterior + function posterior), MACE-TNP learns an end-to-end mapping from data to interventional distributions. This bypasses expensive intermediate computations and reduces approximation compounding. 3. The framework is well-grounded in Bayesian theory. Equation (3) formalizes the Bayesian model-averaged interventional distribution, and MACE-TNP aims to learn this via minimizing KL divergence over a synthetic task distribution. 4. Unlike baselines relying on restrictive assumptions (e.g., Gaussian or additive noise models), MACE-TNP handles both neural network-based and GP-based structural equations.\n\nWeaknesses: 1. Although inference is efficient (forward pass), the training phase is expensive: requires many synthetic SCM samples, uses transformer layers with quadratic complexity, needs large-scale simulation of interventional and observational datasets. 2. The model learns from synthetic Bayesian causal models. It implicitly relies on the fidelity of the training task distribution to approximate real-world data complexity. If the real-world mechanism falls outside this distribution, performance may degrade. 3. While the model converges empirically, no theoretical guarantees are given about identifiability under finite data or architectural constraints. This may limit trust in sensitive domains (e.g., policy, medicine). 4. The model estimates interventional distributions, but it does not output interpretable causal graphs or functions, which might be desirable in scientific applications where explainability is key.", "was_revised": true, "final_justification": "I have carefully read the comments by the authors and other reviews by other reviewers and responses by the authors, my concerns have been addressed and I will keep my score unchanged.", "rebuttal": "We thank the reviewer for their thoughtful and detailed feedback. We’re glad they recognised the importance of addressing the challenge of averaging over uncertain causal structures, and appreciated our proposed end-to-end approach with MACE-TNP. We're especially pleased by the reviewer’s positive remarks on the **theoretical grounding** of our framework, its ability to **bypass costly intermediate inference steps**, its **flexibility** across different structural equation models, and its **scalability** to high-dimensional settings.\n\n**Weaknesses and Questions**\n\n> ... training phase is expensive: requires many synthetic SCM samples, ... large-scale simulation of interventional and observational datasets.\n\nWe appreciate the reviewer highlighting this important point, which we also acknowledge in our limitations section (L371). While the training phase does require generating many synthetic SCM samples, this strategy is consistent with recent successful approaches in the machine learning community—such as TabPFN [1], a state-of-the-art method for tabular data that relies on large-scale synthetic training.\n\nCrucially, our method incurs a one-time training cost, after which inference becomes highly efficient, requiring only a single forward pass (as the reviewer also notes). This makes the approach particularly appealing in settings with frequent inference-time queries, where the upfront cost is effectively amortised.\n\n[1] Hollmann, N., Müller, S., Purucker, L., Krishnakumar, A., Körfer, M., Hoo, S. B., Schirrmeister, R. T., & Hutter, F. (2025). Accurate predictions on small data with a tabular foundation model. Nature, 637(8045), 319–326. https://doi.org/10.1038/s41586-024-08328-6\n\n> ... uses transformer layers with quadratic complexity\n\nWhile MACE-TNP currently uses full attention, the architecture is not inherently constrained to this. It can readily incorporate efficient attention mechanisms--such as linear or kernelized attention [2, 3] and low-rank approximations [4, 5]--to significantly reduce computational cost with minimal impact on performance.\n\n[2] Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020, November). Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning (pp. 5156-5165). PMLR.\n\n[3] Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., ... & Weller, A. (2020). Rethinking attention with performers. arXiv preprint arXiv:2009.14794.\n\n[4] Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.\n\n[5] Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., & Singh, V. (2021, May). Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence (Vol. 35, No. 16, pp. 14138-14148).\n\n> How does the model behave when the test data is generated from functional mechanisms not covered during training?\n\nWe agree this is an important limitation and explicitly acknowledge it in L373. Indeed, performance degradation when the real-world mechanism falls outside of the prior support is a general concern for all Bayesian methods.\n\nHowever, a key advantage of neural processes is their flexibility in incorporating a diverse mixture of priors, easily widening the support and thus improving robustness. We demonstrate this empirically in:\n\n- L325 (Out-of-distribution testing) — showing that, although MACE-TNP performs poorly when the prior is misspecified, we can train a MACE-TNP on a larger mix of priors to recover in-distribution performance for each prior class.\n\n- Section 5.4 — showing that MACE-TNP performs competitively with baselines on real-world data, where the true data-generating process is unknown (and might not be covered by our prior). The prior used for this experiment is detailed in Appendix B.3.\n\n> While the model converges empirically, no theoretical guarantees are given about identifiability under finite data or architectural constraints. This may limit trust in sensitive domains (e.g., policy, medicine)\n\nWe're unsure the identifiability of which quantity you are referring to here. \n- If the underlying causal structure (and interventional distribution) is identifiable, the posterior will concentrate on the correct interventional distribution [6,7]. If not, the posterior will place mass on the interventional distributions consistent with the equivalence class of causal structures that is identifiable from observations. This is shown in Section 5.1.\n- While the recovery of the true posterior is only guaranteed in the limit of infinite tasks (L203), we do show empirically that in the finite task case we outperform baseline methods.\n\n[6] Dhir, Anish, et al. \"Continuous Bayesian Model Selection for Multivariate Causal Discovery.\" Forty-second International Conference on Machine Learning.\n\n[7] Lungu, Valentinian, et al. \"Bayesian causal discovery: Posterior concentration and optimal detection.\" arXiv preprint arXiv:2507.16529 (2025).\n\n\n\n> The model estimates interventional distributions, but it does not output interpretable causal graphs or functions, which might be desirable in scientific applications where explainability is key.\n\nWe agree this is a valid concern—our experiments focus on accurately estimating interventional distributions, rather than recovering interpretable causal graphs or functions. However, since the model can output interventional effects for all node pairs, it is in principle possible to reconstruct the underlying causal structure from these outputs.\n\nMoreover, our framework is flexible and could be extended to jointly estimate a posterior over causal graphs—for example, by adding an auxiliary decoder head, as done in [8]. We see this as a promising direction for future work where interpretability is essential, and will mention it in the Conclusions section of the updated manuscript. We thank the reviewer for highlighting this point.\n\n[8] Dhir, Anish, et al. \"A meta-learning approach to Bayesian causal discovery.\" 2024.\n\n>  Are there known inductive biases that improve generalization?\n\nWe incorporate basic inductive biases into our architecture as discussed in the paragraph starting at L212.\nWhile other inductive biases such as translation equivariance [9] have been incorporated into neural processes for spatiotemporal prediction, the suitability of these additional biases is task-dependent.\n\n[9] Ashman, Matthew, et al. \"Translation Equivariant Transformer Neural Processes.\" Forty-first International Conference on Machine Learning.\n\n\n> Could MACE-TNP be extended to support counterfactual queries? What modifications would be needed in the architecture or loss?\n\nWe believe this approach holds promise for efficiently computing counterfactual queries (and their posteriors). However, pursuing this direction would require substantial modifications at both the architectural and data generation levels, likely constituting a research project in its own right.\n\n> What is the sensitivity of the model to initialization, training set size, or hyperparameter choices? How robust is the meta-learning process?\n\nIn essence, MACE-TNP incorporates a transformer architecture, enabling us to draw on the extensive literature regarding effective initialization and training of such models. While we employ standard architectural choices, the method is flexible and can be adapted to any architecture the user prefers.\n\nFurthermore, we provide an ablation of a set of hyperparmeters (i.e. MoG components, attention mechanism, model size, etc.), and some resulting guidance, in Appendix C.2.2. All configurations yield reasonable performance, with slight improvements observed as the architecture size or number of MoG components increases. This suggests that the model is both robust to a range of design choices and able to benefit from scaling.\n\nThe reviewer is right to highlight training set size as an important factor, particularly in the context of meta-learning. While we did not investigate the sample efficiency of our method, we agree that it is an important direction for future research. However, empirically evaluating sample efficiency can be computationally expensive.\n\n> Have you assessed how well-calibrated the estimated interventional distributions are? Could calibration methods be integrated during training?\n\nA lower score on our metric - NLPID (averaged over multiple datasets)- implies that the distribution is closer in KL to the true distribution [10]. \nHence we expect the models with a lower score to be better calibrated under the assumption that datasets have been generated from the prior.\nIf this is not the case, post hoc calibration methods can indeed be used to better calibrate the model [11].\n\n[10] David Madigan and Adrian E Raftery. Model selection and accounting for model uncertainty in graphical models using occam’s window.\n\n[11] Fong, Edwin, and Chris C. Holmes. \"Conformal bayesian computation.\" Advances in Neural Information Processing Systems 34 (2021).\n\n---\n\nWe thank the reviewer once again for their thoughtful and constructive feedback. We are happy to address any remaining questions, and if we have satisfactorily addressed the majority of the concerns, we would be grateful if the reviewer would consider revisiting their score."}, {"review_id": "4Tbwk0kNO5", "reviewer": "Reviewer_nB8D", "summary": "This paper studies the problem of evaluating interventional distributions / causal effects from the observational data without access to a detailed causal graph describing the underlying data-generating mechanisms. The authors focus on a specific family of causal models where there is no unobserved confounder. They also assume access to the Bayesian prior over the possible causal graph and parameters of the underlying causal models. Based on these knowledge, the authors propose a variational inference method to approximate the posterior distribution over the interventional distributions provided with the observational data. The simulation results support the proposed method for linear causal models and causal models with functions parametrized as Gaussian processes.", "questions": "1. How is the principle of meta-learning applied to this problem setting?\n2. What are the parametric model assumptions that enable the proposed method?\n3. In Line 72, the authors state, \"Like a majority of works that learn causal models from data, we assume no hidden confounders.\" This statement is controversial, as unobserved confounding is one of the key challenges to causal inference. A majority of advancements in causal estimation after 2000 have focused on addressing the issues of confounding bias. For example, there is a growing line of work for identifying causal effects provided with a partial description of causal knowledge represented as a Markov equivalence class (Jiji Zhang, 08; Malinsky & Spirtes, 16; Jaber, Zhang & Bareinboim 19).", "limitations": "The authors have not adequately addressed the theoretical limitation of their work. The model assumptions underlying the proposed method are unclear. Regarding the societal impact, the paper is mainly theoretical, and its long-term implications remain unclear.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "This paper studies an exciting problem of causal inference without access to a complete causal graph. The direction of leveraging generative modeling is promising. Overall, this paper is well organized, and the simulation results support the proposed method under a specific family of causal models. However, I do have some questions and concerns, summarized below.\n\n1. The relations to the meta-learning method are somewhat unclear. Typically, meta-learning algorithms refer to the methods that learn to learn. For example, a popular line of meta-learning approach (Finn, Abbeel & Levine, 2017) learns an optimal initialization point for generalizing to any future environment after taking a few steps of gradient descent. It could also be framed in terms of recurrent models that are trained via a meta-learning objective. This paper focuses on learning a posterior distribution over the potential outcomes provided with observational data. It is unclear how this learning objective relates to the existing meta-learning literature.\n\n2. The model assumptions underlying the proposed method are unclear. The proposed method requires access to the detailed parametric knowledge about the underlying data-generating models to properly assign the prior. The authors should list these model assumptions more explicitly, similar to [57]. Otherwise, it would be confusing when compared to other non-parametric methods that attempt to leverage qualitative knowledge, e.g., (Jiji Zhang, 08).\n\n3. The learning objective here is the posterior interventional distribution, averaging over all possible models. It would be appreciated if the authors could provide more insights on this learning objective. Since the posteior could be sensitive to the prior and one may not have access to the correct prior, it is possible to assign a prior distribution that leads to a posterior interventional distribution that is overly concentrated towards the mean. In this case, it is sufficient to learn a neural process model that recovers the first moment of the posterior but underestimates the variance. The learned posterior interventional distribution would not be able to capture the representations if the ground-truth model is an extreme case, located around the margin. This is one of the main reasons why existing causal partial identification focuses on the quantile evaluation over the boundary case (Balke & Pearl, 1994; Zhang, Tian & Bareinboi, 2021). This issue could be clarified by comparing the proposed method with existing bounding literature over a small number of candidate graphs over a discrete domain. In that case, one could clearly plot the upper and lower bound over the target causal effect, and compare how the learned posterior covers the whole feasible region. One such example is Figure 3 in (Zhang, Tian & Bareinboi, 2021).", "was_revised": true, "final_justification": "The authors have addressed my concerns regarding the posterior distribution inference.", "rebuttal": "We thank the reviewer for their thoughtful feedback. We’re pleased that the reviewer described our paper as addressing “an exciting problem” and noted that “the direction of leveraging generative modeling is promising.” We also appreciate their positive remarks that the \"paper is well organized, and the simulation results support the proposed method under a specific family of causal models\".\n\n**Summary**\n\nWe would like to bring some clarifications to the summary provided by the reviewer.\n\n> \"… assume access to the Bayesian prior over the possible causal graph and parameters of the underlying causal models.\"\n\n**We do not assume access to the \"true\" Bayesian prior**, but rather that we can construct an empirical Bayesian prior with enough support that it is informative for the test data (whose performance is the one that is most relevant in the end). \n\nUnlike many other methods, if we have reason to believe the prior is not wide enough (e.g. due to poor prediction on a held out set), **we can easily incorporate synthetic datasets from an even wider variety of Bayesian causal models (BCM) (def. 2.1) at training time and thus widen the support**. \n\nWe illustrate this in section 5.4 where we test our method on a real-life dataset, implying we do not have access to the true data generating process.\n\n> \"… a variational inference method\"\n\n**We do not use variational inference in the training**. For a single task (defined in L186), the task-dependent loss (Eq. 4) involves computing the log-likelihood of the interventional targets (sampled from the BCM) conditioned on the observational data. This is available in closed form. Furthermore, since tasks exhibit statistical similarity, we adopt a meta-learning approach, computing the final objective as an expectation over the task distribution.\n\n> \"...functions parametrized as Gaussian processes\"\n\nIn our experiments, we also include neural network-based functional mechanisms. For a description of how we construct our synthetic tasks, we refer the reviewer to appendix B.2 and B.3.\n\n**Weaknesses and Questions**\n\n> How is the principle of meta-learning applied to this problem setting?\n\nNeural Processes (NPs) are models that learn to map datasets to predictive distributions, where **each dataset is sampled from a distribution over tasks that share common statistical structure**. This two-level hierarchy is captured in our loss function (Eq. 4), which reflects learning both within each task (as in standard supervised learning) and across tasks via the expectation over task distributions (where we denote a task with $\\xi$) (L186) — the latter enables the model to improve its performance on unseen tasks (i.e., meta-learning or learning to learn). For a comprehensive overview of Neural Processes, we refer the reader to Section 1.2 of [1], as well as to the original papers that classify neural processes as a meta-learning paradigm [2, 3].\n\nIn short, **meta-learning is happening at training time, where the loss involves an outer expectation over statistically-similar tasks**.\n\n[1] Bruinsma, W. P. (2022). Convolutional Conditional Neural Processes (PhD thesis, Department of Engineering, University of Cambridge). University of Cambridge.\n\n[2] Garnelo, Marta, et al. \"Conditional neural processes.\" International conference on machine learning. PMLR, 2018.\n\n[3] Garnelo, Marta, et al. \"Neural processes.\" arXiv preprint arXiv:1807.01622 (2018).\n\n> What are the parametric model assumptions that enable the proposed method?\n\nIn our method, at train time, a user can specify any prior (or a mixture of priors) by specifying the Bayesian causal models (BCMs) (def. 2.1) that the training data is sampled from. At test time the posterior interventional distribution is obtained by a forward pass through the network of the test dataset. Thus, **the only model assumptions in our approach are the ones determined by the choice of BCMs we train our model on**. Our method does not rely on any specific choice of BCM, and can be changed depending on the problem. \n\nWe have specified the BCMs that generated the training data (and hence determine the prior) for each experiment in Appendix B. For further clarity, we will include the description of the BCMs that generated the data for each experiment in the main text. \n\n> The learning objective here is the posterior interventional distribution, averaging over all possible models. It would be appreciated if the authors could provide more insights on this learning objective.\n\nLet us unpack the learning objective of MACE-TNP, stated in Eq. 4. As mentioned above, the objective is hierarchical:\n1) At the task level, the quantity we want to learn is the true posterior interventional distribution. Thus, at training time, we seek to minimise the KL divergence between this quantity and the output of our model. Equivalently, this can be framed as maximising the expected log predictive distribution of the model under the interventional distribution (i.e. interventional targets conditioned on the current task $\\mathbb{E}_{X_i | \\xi}$). At training time, this can be evaluated in closed-form.\n2) Because these tasks share statistical similarities, we take a meta-learning approach where we also average over multiple tasks  $\\xi$ (which are defined by the observational data $\\mathcal{D}_{\\text{obs}}$, the interventional node $j$ and value $x_j$, and the outcome node $i$).\n\nCombined, these two lead to Eq. 4.\n\n> ... existing causal partial identification focuses on the quantile evaluation over the boundary case (Balke & Pearl, 1994; Zhang, Tian & Bareinboi, 2021).\n\nWe appreciate the reviewer highlighting this point.\nThe work you shared [4], and others [5], are similar in spirit to ours. In their approach, credible regions (e.g., Figure 3 in [4]) are derived from a posterior distribution obtained via MCMC, using a specified prior (see Section 3: Bayesian Approach for Partial Identification in [4]).\n\nIn contrast, **our method offers an alternative way to obtain a posterior over interventions, without relying on MCMC**. As demonstrated in our experiments, our approach outperforms MCMC-based methods such as BCI-GPN.\nSamples from the posterior of the neural process can be used to estimate the (support) credible region.\nImportantly, our framework does not commit to a specific prior choice: if there are concerns about underestimating credible region widths, one can incorporate broader or even uninformative priors in the training process without modifying the core methodology. \nOur contribution is thus orthogonal to the choice of prior.\n\n[4] Zhang, Junzhe, Jin Tian, and Elias Bareinboim. \"Partial counterfactual identification from observational and experimental data.\" International Conference on Machine Learning. PMLR, 2022.\n\n[5] Chickering, David Maxwell, and Judea Pearl. \"A clinician's tool for analyzing non-compliance.\" Proceedings of the National Conference on Artificial Intelligence. 1996.\n\n> In Line 72, the authors state, \"Like a majority of works that learn causal models from data, we assume no hidden confounders.\"\n\nWe agree with the reviewer that the original statement was overstated and will revise it accordingly. Our assumption of causal sufficiency was made to maintain consistency with the baselines (and a lot of works on causal discovery), which also rely on this assumption. That said, we acknowledge the existence of methods that do not assume causal sufficiency.\n\nImportantly, our framework **can accommodate hidden confounding** by training on data generated from BCMs with unobserved confounders. We conducted an experiment using a simple 3-node linear model, demonstrating that MACE-TNP correctly identifies the interventional distribution in the presence of observed confounding and appropriately reflects uncertainty when a confounder is hidden (and thus the interventional distribution is not identifiable) (similar to Figure 3 in our work). Unfortunately, because we are not allowed to include any images in the rebuttal, we cannot share the results at this stage, but we will include them in the revised version.\n\nWe thank the reviewer for raising this important point, as the additional experiment further highlights the flexibility of MACE-TNP in estimating interventional distributions, even under hidden confounding.\n\n---\n\nIf the reviewer feels we have addressed the majority of their concerns, we would appreciate it if they reconsidered their score. We are also happy to further address any remaining concerns."}, {"review_id": "TlwBMO7q0a", "reviewer": "Reviewer_2FPW", "summary": "This paper tackles the challenge of estimating interventional distributions when the true causal graph is unknown and must be inferred from observational data. Instead of relying on a single estimated causal structure—which can lead to overconfident or incorrect conclusions—the authors propose using meta-learning to approximate the full Bayesian model-averaged interventional distribution.\n\nThey introduce MACE-TNP (Model-Averaged Causal Estimation Transformer Neural Process), an end-to-end framework that bypasses the intractable steps of explicitly modeling posteriors over causal structures and functions. By directly learning to predict the interventional posterior, their approach avoids compounding errors and scales to more complex scenarios.\n\nEmpirical results show that MACE-TNP recovers known analytical posteriors when available, outperforms strong Bayesian and non-Bayesian baselines across diverse experiments, and offers a promising path toward scalable, meta-learning-based causal inference.", "questions": "Please respond to the weakness section.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1) Addresses an important problem: Tackles the realistic and challenging scenario where causal graphs are uncertain, avoiding overconfidence from single-structure methods.\n2) Principled approach: Builds on Bayesian model averaging to handle structural uncertainty in a theoretically sound way.\n3) Innovative use of meta-learning: Leverages neural processes to amortize the entire causal inference pipeline, bypassing otherwise intractable computations.\n4) Strong empirical validation: Demonstrates that MACE-TNP recovers known analytical posteriors and outperforms existing Bayesian and non-Bayesian baselines across increasingly complex settings.\n5) Potential scalability: Shows promise for scaling to high-dimensional problems where traditional Bayesian methods become impractical.\n\nWeaknesses:\n1) Reliance on synthetic data: Most results are on synthetic datasets; it remains unclear how well the method performs on noisy, imperfect real-world data.\n2) Interpretability limitations: While the method predicts interventional distributions, it offers limited insight into the underlying causal structures or mechanisms it learns, which may be important in scientific applications.\n3) Assumes training distributions generalize: The success of the meta-learning approach depends on training data sufficiently representing the range of possible causal scenarios; this may limit applicability in domains with very different or unseen causal patterns.\n4) Potential sensitivity: As a meta-learned model, it may be sensitive to hyperparameters or require large, well-designed training datasets to achieve robust performance.", "was_revised": true, "final_justification": "I am reasonably satisfied with the authors' response. Some of the issues raised are common in the field and largely beyond the authors' control; therefore, I will not penalize them for these.", "rebuttal": "We thank the reviewer for their thoughtful feedback. We're glad they found our approach **principled**, our use of **meta-learning innovative**, and appreciated the **strong empirical** results showing MACE-TNP's ability to recover analytical posteriors and outperform strong baselines. We're also encouraged by their recognition of the method’s **scalability** to complex, high-dimensional settings.\n\n> Reliance on synthetic data\n\nWe agree that real-world evaluation is crucial, **which is why we included experiments on the real-world Sachs dataset in Section 5.4**. However, the lack of established benchmarks with reliable interventional data remains a challenge in the community. Despite this, we are actively exploring additional real-world settings to further assess the strengths and limitations of our approach beyond synthetic benchmarks.\n\n\n> Interpretability limitations ... limited insight into the underlying causal structures or mechanisms\n\nAs mentioned in our reply to reviewer RHFT, we agree this is a valid concern. We focus on accurately estimating interventional distributions, rather than recovering interpretable causal graphs or functions. However, since the model can output interventional effects for all node pairs, it is in principle possible to reconstruct the underlying causal structure from these outputs.\n\nMoreover, our framework is flexible and could be extended to **jointly estimate a posterior over causal graphs**—for example, by adding an auxiliary decoder head, as done in [1]. We see this as a promising direction for future work where interpretability is essential, and will mention it in the Conclusions section of the updated manuscript. We thank the reviewer for highlighting this point.\n\n[1] Dhir, Anish, et al. \"Continuous Bayesian Model Selection for Multivariate Causal Discovery.\" Forty-second International Conference on Machine Learning.\n\n> Assumes training distributions generalize\n\nWe agree that the effectiveness of our model depends on how well the training prior captures the target domain. This is a known limitation, and we discuss it in the paper. A conservative strategy is to use a broad prior during training to encourage generalization. Moreover, our framework is flexible: if performance in a specific application is suboptimal due to prior mismatch, the model can be easily fine-tuned on more domain-aligned priors [2].\n\nWe also demonstrate empirically:\n\n- L325 (Out-of-distribution testing) — showing that, although MACE-TNP performs poorly when the prior is misspecified, we can train a MACE-TNP on a larger mix of priors to recover in-distribution performance for each prior class.\n\n- Section 5.4 — showing that MACE-TNP performs competetively with baselines on real-world data, where the true data-generating process is unknown. The prior used for this experiment is detailed in Appendix B.3.\n\n[2] Scholz, Jonas, et al. \"Sim2real for environmental neural processes.\" arXiv preprint arXiv:2310.19932 (2023).\n\n> Potential sensitivity: As a meta-learned model, it may be sensitive to hyperparameters or require large, well-designed training datasets to achieve robust performance.\n\n\nWe agree that the choice of training distribution (prior) plays a key role in performance. The choice of prior for tasks without any knowledge is an active research area [3], and future iterations of our approach could benefit from advances in identifying realistic synthetic priors.\n\nRegarding architectural hyperparameters, Appendix C.2.2 includes ablations on attention types, number of MoG components, model dimensions, and transformer depth. While some configurations perform slightly better, the differences are modest, suggesting that MACE-TNP is relatively robust to these design choices.\n\n[3] Zhang, X., & Maddix Robinson, D. (2025, July 22). Mitra: Mixed synthetic priors for enhancing tabular foundation models. Amazon Science. Retrieved from Amazon Science blog.\n\n---\n\nWe would like to thank the reviewer once again for their feedback and hope that our responses have satisfactorily addressed the raised concerns."}]

Author rebuttal:

Meta-review: This work proposes a meta-learning and model-averaging approach to infer interventional distributions. Despite limitations (expensive training, scalability, unobserved confounding,...) reviewers overall agree that the work is valuable and notably emphasize the principled approach and good experimental results with strong baselines.
Decision: Accept (poster)