Amortized Counterfactual Prediction with Graph Neural Networks

Amortized Counterfactual Prediction with Graph Neural Networks Hendrik Stumpf†1 and Francesco Quinzan†2 1

Department of Computer Science, University of Oxford 2 Department of Engineering, University of Oxford † Corresponding autohrs. Emails: hendrik.stumpf@cs.ox.ac.edu, francesco.quinzan@eng.ox.ac.uk

arXiv:2610.13724v1 [cs.LG] 9 Oct 2026

Abstract Counterfactual inference asks how an observed outcome might have differed under different circumstances, such as whether a patient would have recovered had they received a different treatment. Such questions arise whenever decisions depend on the outcomes of alternatives that were not observed. Answering counterfactual queries generally requires knowledge of the underlying causal structure, which is often unavailable in practice. A common approach is therefore to first estimate the causal graph from observational data and then perform counterfactual inference. However, errors in the estimated causal graph can propagate to downstream counterfactual predictions. To address this challenge, we introduce the Counterfactual Graph Neural Process (CFGNP), an end-to-end approach that answers counterfactual queries directly from observational data without first estimating a causal graph. CFGNP uses Graph Neural Networks (GNNs) to exploit the dual-graph structure by restricting cross-world message passing to corresponding factual– counterfactual variables. Across synthetic and image-based benchmarks, CFGNP achieves stronger predictive performance than existing approaches, including on the German Loan benchmark, where it reduces MSE from 4.51 × 10−2 to 0.19 × 10−2 . Furthermore, on synthetic causal graphs, a model trained only on 10-node graphs maintains strong predictive performance on graphs with up to 30 nodes, without retraining.

1

Introduction

Understanding cause-and-effect relationships is a central goal in many scientific disciplines. In medicine, researchers may wish to determine whether a treatment improves patient outcomes [Burotto et al., 2014, Stovitz and Shrier, 2019]. In economics, policymakers often aim to understand how changes in policy affect employment or inflation [Heckman and Pinto, 2022]. Similar questions arise in biology, social sciences, and machine learning [Morgan and Winship, 2007, Sanderson et al., 2022, Schölkopf et al., 2021]. Especially with the increased deployment of machine learning systems in settings where interpretability is crucial, understanding causeand-effect relationships becomes more important [Li and Han, 2023]. In all of these settings, the objective is not merely to observe statistical associations, but to understand how outcomes would change if different actions or conditions were imposed. Causal inference [Pearl, 2009] provides a framework for studying such questions. Unlike conventional statistical learning, causal inference seeks to uncover cause-effect relationships between variables, rather than statistical associations. The framework formalises this with the structural causal model (SCM), which specifies how variables are generated, and how they causally affect one another. Within the SCM framework, causal questions can be distinguished by the type of reasoning they require. Pearl [Pearl, 2009] differentiates between three levels: associational, interventional, and counterfactual queries. Associational queries concern relationships observed in data, while interventional queries ask how outcomes change when a variable is actively manipulated [Pearl, 1996]. Counterfactual queries go one step further by asking what would have happened in an alternative scenario, given what was observed in 1

the factual world [Shpitser and Pearl, 2012]. Counterfactual queries are particularly important because they allow reasoning about individualized outcomes and explanations, rather than only population-level effects [Nguyen et al., 2020]. This makes them harder to solve than the other two query types as we additionally have to take the alternative hypothesis into account. Traditional approaches to counterfactual inference typically rely on substantial prior knowledge about the underlying cause-effect relationships between variables [Nasr-Esfahany et al., 2023, Rasal et al., 2025, Ribeiro et al., 2025, Xia et al., 2022]. In practice, this often means assuming access to the causal graph, which describes these relationships. Sometimes, less information, such as the topological node ordering of the causal graph, is required [Le et al., 2024]. However, in many real-world settings, such information is unavailable or highly uncertain. This is particularly the case in domains such as medicine, biology, and the social sciences, where causal relationships may be obscured by latent confounders, incomplete observations, and complex dependencies among variables [Loh, 2025, Ruiz-Perez et al., 2024, Weberpals et al., 2024]. Hence, if there is no prior knowledge about the causal graph, we must first use observational samples to estimate the causal graph [Spirtes et al., 2001] and then use the resulting causal graph with the existing methods for counterfactual inference. However, this has two drawbacks; first, the search space of causal graphs scales super-exponentially in the number of variables [Zheng et al., 2018] and second, choosing the wrong causal graph amplifies errors and may cause the counterfactual to be incorrect [Karimi et al., 2020b]. To avoid explicitly estimating the causal graph, we instead condition the counterfactual prediction directly on observational samples. This allows the model to extract information about the underlying causal structure that is relevant to the counterfactual query, without committing to a single estimated graph. We approach this problem through meta-learning using Neural Processes (NPs) [Garnelo et al., 2018], which learn to map a set of observations to a predictive distribution. Building on Graph Neural Processes [Carr and Wingate, 2019], we further incorporate the structural constraints imposed by the dual graph into the prediction. Based on this approach, we introduce the Counterfactual Graph Neural Process (CFGNP), an end-to-end framework for counterfactual inference under unknown causal structure. An overview of our setting is shown in Figure 1. Our contributions. • We introduce the Counterfactual Graph Neural Process (CFGNP), an end-to-end framework for learning counterfactual distributions without prior knowledge of the underlying causal graph (Section 4). • We develop a Graph Neural Process architecture that conditions counterfactual predictions on observational samples (Section 4.1). • We demonstrate that CFGNP outperforms existing counterfactual inference methods on both synthetic and real-world data (Sections 5.1 and 5.2). • We evaluate CFGNP on counterfactual queries involving features that are unseen during training and observe strong performance in this setting (Section 5.2).

2

Related Work

Causal Discovery. For causal inference tasks such as evaluating counterfactual queries, most existing methods rely on explicit knowledge of the causal graph. Estimating this graph from associational data is known as causal discovery [Verma and Pearl, 2022]. The primary goal of causal discovery is to uncover true cause-and-effect relationships by learning the underlying structural dependencies rather than relying on statistical correlations alone [Pearl, 2009, Peters 2

Figure 1: Overview of the task solved by CFGNP. Unknown variables are in light grey and known variables in dark grey. Blue edges represent missing knowledge about the causal graph. The counterfactual is predicted based on the counterfactual query and additional observational samples. et al., 2017]. Classical methodologies encompass constraint-based approaches, such as the PC algorithm [Spirtes et al., 2001], and score-based frameworks, such as Greedy Equivalence Search [Chickering, 2002]. More recently, modern techniques including LiNGAM [Shimizu et al., 2006], NOTEARS [Zheng et al., 2018], and DAG-GNN [Yu et al., 2019] have advanced the field by introducing functional form assumptions, continuous optimization, and deep learning architectures. Nevertheless, key challenges remain, such as unobserved confounding [Diepen et al., 2023, Kaltenpoth and Vreeken, 2023], complex nonlinearities [Li et al., 2024], and computational scalability [Amin and Wilson, 2024, Cheng et al., 2024]. Addressing these challenges continues to make causal inference an active area of research. Counterfactual Inference. Existing approaches attempt to learn counterfactual distributions under varying degrees of prior structural knowledge [Caron et al., 2022, Karimi et al., 2020a, Khemakhem et al., 2021]. For instance, Le et al. [2024] and Javaloy et al. [2023] use normalizing flows that require only a topological ordering of the graph nodes, whereas most alternatives assume the full causal graph. Within deep learning for counterfactual inference, Pawlowski et al. [2020] use normalizing flows to parameterize structural causal models, while recent flow-matching approaches similarly require full graph specification [Ribeiro et al., 2025]. Other methods explicitly exploit relational structure through Graph Neural Networks [Sanchez-Martin et al., 2021, Zečević et al., 2021]. This reliance on a known, estimated, or otherwise fixed causal structure can be restrictive [Caron et al., 2022, Wu et al., 2025]: the space of DAGs grows super-exponentially with the number of variables [Zheng et al., 2018], and errors in the assumed causal structure can propagate to downstream counterfactual predictions [Karimi et al., 2020b]. While end-to-end learning under uncertain causal structure has been explored for interventional inference [Dhir et al., 2025, Geffner et al., 2022], an analogous approach has not yet been established for counterfactual inference, where additional structural constraints from the dual graph can be exploited [von Kügelgen et al., 2022]. Image Counterfactuals. Several works focus specifically on image-based counterfactual generation and analysis [Luedtke and Fukumizu, 2025, Ma et al., 2026, Pawlowski et al., 2020, Rasal et al., 2025]. Among these, DeepBC [Kladny et al., 2024] performs counterfactual inference through latent exogenous variables under a known causal graph. Other recent approaches based on diffusion models [Rasal et al., 2025] and optimal transport [Ribeiro et al., 2025] similarly assume that the causal graph is known. In contrast, VCI [Wu et al., 2025] does not take an explicit causal graph as input, but instead assumes a prescribed treatment– covariate–outcome structure and uses a variational formulation to generate counterfactuals without requiring counterfactual training pairs. Neural Processes. To address these limitations, our approach builds on Neural Processes 3

(NPs) [Garnelo et al., 2018], which have been successfully applied across a range of domains including Bayesian optimization and other scientific applications [Bruinsma et al., 2023, Foong et al., 2020, Kim et al., 2019]. NPs have recently been explored in causal settings, including treatment-effect estimation and Bayesian causal inference [Bynum et al., 2026, Dhir et al., 2025]. However, these works address related causal inference tasks rather than counterfactual inference. In contrast, our approach extends this framework to counterfactual inference, for which we additional leveraging GNNs to capture the structural relationships between factual and counterfactual variables [von Kügelgen et al., 2022]. Graph Neural Processes extend NPs to relational and graph-structured data [Carr and Wingate, 2019, Hu et al., 2023, Liang and Gao, 2021]. By combining these ideas, we directly map observational data to the target counterfactual distribution using GNNs over the dual-graph structure, without first estimating a causal graph.

3

Background

We now introduce notions of causal inference, neural processes, and graph neural networks. Notation. Unstarred and starred variables denote the factual and counterfactual worlds, respectively. Let z ∈ Rn×ddata be the factual observed unit, T ∗ = t∗ a counterfactual condition, obs and Y ∗ ⊆ V the target variables. The observational context is O = {Oi }N i=1 , consisting of i.i.d. samples from the same SCM M, where n and Nobs denote the number of endogenous variables and context samples, respectively. We write p(Y ∗ | O, z, T ∗ = t∗ ) when M is informed by the observational context O, and p(Y ∗ | M, z, T ∗ = t∗ ) when emphasizing the model. If M is unspecified, we assume the true model.

3.1

Structural Causal Models and Causal Inference

Structural Causal Model. A structural causal model (SCM) describes the causal relationships among a set of endogenous variables V = {V1 , . . . , Vn } through a collection of structural equations f = {f1 , . . . , fn }. Each variable Vk is determined by its direct causes P ak ⊆ V \ Vk and an exogenous noise variable Uk according to Vk = fk (P ak , Uk ). The exogenous variables U = U1 , . . . , Un capture factors outside the model and are assumed to be mutually independent. The causal relationships encoded by the structural equations can be represented by a directed and acyclic causal graph G, where an edge Vj → Vk indicates that Vj ∈ P ak . In particular, with U coming from a probability distribution we can define associational queries (L1 ) to concern distributions p(Y ) with Y ⊂ V . Observational samples Oi are sampled from p(V ). Interventional queries (L2 ) concern questions p(Y | do(T = t)) where we actively set nodes T ⊂ V to a specific value and consider the outcome Y ⊂ V . (Backtracking) counterfactual queries (L3 ) are probability distributions, with T ∗ , Y ∗ , Z ⊂ V and T ∗ ∩ Y ∗ = ∅, of the form P (Y ∗ = y ∗ | T ∗ = t∗ , Z = z), concerning questions such as: given that Z = z happened, what would have happened to variables Y ∗ if alternatively T ∗ = t∗ . In particular, we use notation X ∗ = x∗ instead of the do-operator to make clear that we are using backtracking counterfactuals [von Kügelgen et al., 2022] and not interventionist counterfactuals [Pearl, 2009]. While interventionist counterfactuals artificially modify the structural equations in the counterfactual world, they thereby locally violate the laws of nature. In contrast, backtracking counterfactuals keep the structural equations intact and instead consider alternative exogenous conditions that could have produced X ∗ = x∗ . 4

Dual Graph. To answer backtracking counterfactuals, we need the dual graph. Figure 2 shows an example of the dual graph for backtracking counterfactuals. The dual graph consists of two copies of the same SCM. The left part, the observed factual unit, represents what actually happened, whereas the copy on the right, the counterfactual unit, represents the alternative circumstances. Both worlds are connected via the exogenous nodes, which share a probability distribution p(U, U ∗ ). We can solve the counterfactual query via the abduction– marginalisation–prediction framework [von Kügelgen et al., 2022], consisting of the following steps: 1.) Abduction: We infer the joint posterior over the exogenous nodes P (U, U ∗ | Z = z, T ∗ = t∗ ). 2.) Marginalisation: We marginalise out U to obtain P (U ∗ | Z = z, T ∗ = t∗ ). 3.) Prediction: We use the SCM to predict P (Y ∗ | Z = z, T ∗ = t∗ ).

Figure 2: Example of the dual graph for a counterfactual query with a causal graph of four nodes. Light grey nodes are unknown.

Identifiability. When the underlying causal model is only partially known, a causal query may not be uniquely determined by the available information. An Li query is identifiable from Lj information if all models in the considered model class that agree on the available Lj information yield the same answer [Pearl, 2009]. This is relevant to our setting because counterfactual queries are not necessarily identifiable from observational data alone.

3.2

Neural Processes and Graph Neural Networks

Neural Processes. Neural Processes (NPs) [Garnelo et al., 2018] are probabilistic models that learn distributions over functions from a collection of related tasks. Each task is represented by a NG C context set C = {(xi , yi )}N i=1 of observed input-output pairs and a target set G = {(xj , yj )}j=1 at which predictions are made. NPs learn the conditional predictive distribution p(yG | xG , C), thereby amortizing inference across tasks by mapping context observations directly to predictive distributions. Extensions of NPs to graph-structured domains replace neural architectures with graph-based models to capture relational structure [Ashman et al., 2024, Carr and Wingate, 2019]. Graph Neural Networks. Graph Neural Networks (GNNs) learn representations of graph(k) structured data by iteratively propagating information between connected nodes. Let hi (0) denote the representation of node vi at layer k, with hi = xi . Following the message-passing framework [Gilmer et al., 2017], a GNN layer aggregates messages from the neighbourhood N (i) of each node, n o (k) (k) (k) (k+1) (k) (k) mi = AGG(k) ϕ(k) (hi , hj , eji ) vj ∈ N (i) , hi = ψ (k) (hi , mi ), where ϕ(k) and ψ (k) are learnable message and update functions, respectively, and AGG(k) is permutation invariant. For directed graphs, the neighbourhood can be partitioned into incoming and outgoing neighbours, N − (i) and N + (i), allowing messages to be propagated separately along each direction. This distinction is particularly relevant for causal graphs, where incoming and outgoing edges represent causal relations in opposite directions.

4

A GNN Model for Counterfactual Inference without Causal Graph Knowledge

This work focuses on predicting counterfactual distributions directly from observational data without prior knowledge about the causal graph. For each task, we observe a context set 5

Figure 3: Overview of CFGNP. Observational samples are first aggregated into node-wise context representations, which are combined with the observed factual unit. Message passing between factual and counterfactual unit then produces parameters of the counterfactual predictive distribution. obs O = {zi }N i=1 of i.i.d. observational samples from an unknown SCM M = (G, f ). Given a factual observation z and a counterfactual condition T ∗ = t∗ , our goal is to predict the counterfactual distribution p(Y ∗ | O, z, T ∗ = t∗ ) over a set of query variables Y ∗ . The observational context O provides evidence about M. We aim to directly learn the model-averaged counterfactual predictive distribution

p(Y ∗ | O, z, T ∗ = t∗ ) = EM∼p(M|O,z) [p(Y ∗ | M, z, T ∗ = t∗ )] ,

(1)

thereby accounting for uncertainty over causal models without explicitly inferring or committing to a particular causal graph. Existing methods on the other hand try to estimate the causal graph first. Since we intend to approximate the counterfactual probability distribution, we use the KullbackLeibler (KL) divergence [Kullback and Leibler, 1951] as the minimization objective. Suppose we have a list L of samples (τ, y ∗ ) ∼ q, where τ = ((z, T ∗ , t∗ , Y ∗ ), O) contains the factual representation z, counterfactual query, and observation pattern O, while y ∗ contains the corresponding counterfactual target values. Let pθ (y ∗ | τ ) denote the distribution predicted by the model with weights θ, where y ∗ may contain multiple target variables. We then want to obtain the weights that minimize the KL divergence, hence θ∗ = arg min Eτ ∼q [KL(p(Y ∗ = y ∗ |z, T ∗ = t∗ ) || pθ (Y ∗ = y ∗ |z, T ∗ = t∗ , O))] θ h i = arg max Eτ ∼q Ey∗ |τ [log pθ (Y ∗ = y ∗ |z, T ∗ = t∗ , O)] + C θ

for some constant C independent of the model parameters. During training we assume access to counterfactual pairs, as we sample y ∗ | τ ∼ p(Y ∗ = y ∗ |z, T ∗ = t∗ , O). Recovery of the Counterfactual Distribution. In the limit of infinite data and model capacity, minimizing the KL divergence recovers the conditional predictive distribution [Bruinsma, 2022] p(Y ∗ | O, z, T ∗ = t∗ ), with latent quantities, such as the exogenous variables, implicitly marginalized out [Gordon et al., 2019]. If the counterfactual distribution is identifiable from observational data, this predictive distribution coincides with the counterfactual distribution of the ground-truth causal model. Otherwise, the model instead recovers the model-averaged counterfactual distribution over SCMs compatible with the observations.

4.1

Approach

Figure 3 illustrates the architecture of the Counterfactual Graph Neural Process (CFGNP). Given observational samples O, a factual observation z, and a counterfactual intervention x∗ , CFGNP predicts the corresponding counterfactual outcome distribution. The architecture consists of three stages: summarizing the observational context (middle), combining it with the factual observation (top right), and propagating the resulting representation to the counterfactual unit (bottom right). We describe each component below. 6

Embedding Input. We project all nodes to the hidden dimension dhidden using node-wise MLPs. For O, we apply the same MLP hO across all nodes. For z and x∗ , we distinguish intervention nodes T ∗ from the remaining nodes using separate MLPs: hz,int and hz,common for z, and hx,int and hx,common for x∗ . In x∗ , all entries outside T ∗ are masked. Observational Context Encoding. We next incorporate the observational samples O by allowing nodes within each sample Oj to communicate and compressing their information into q. For each variable i, we introduce an aggregation node qi , initialized to zero, and connect the corresponding node i from every observational sample to qi . Since we have no prior knowledge of the causal graph, we additionally fully connect the nodes within each observational sample. The resulting graph consists of Nobs disjoint fully connected n-node subgraphs and n aggregation nodes. A K-layer GAT [Brody et al., 2022] performs message passing, [O′ , q ′ ] = (1) mθ ([O, q]), such that q ′ captures the compressed observational context. We then integrate this context with the factual unit z. To keep the graph compact, we connect each qi′ only to its corresponding zi and add a self-loop to each zi . A 3-layer GAT propagates (2) the contextual information to the factual unit, z ′ = mθ (z, q ′ ). Predicting Counterfactual Distribution. The final stage aims to approximate the abduction–marginalisation–prediction steps using the context-informed factual representation z ′ and counterfactual condition x∗ . The dual graph enables these steps by coupling the factual and counterfactual worlds through connected exogenous nodes. Approximating these steps generally requires bidirectional message passing. Because we do not have prior knowledge of the causal graph, factual and counterfactual units are fully connected. Thus, each endogenous node has direct access to the information needed to reproduce its corresponding exogenous state, allowing us to remove explicit exogenous nodes and connect only corresponding factual and counterfactual variables across worlds. In this case, a unidirectional GNN can approximate the corresponding bidirectional GNN arbitrarily well (App. 1). As a design choice, we apply min{⌊n/2⌋, 5} GAT layers, scaling depth with graph size while (3) capping it at five to limit over-smoothing [Li et al., 2018]. This yields [z ′′ , h∗ ] = mθ (z ′ , x∗ ), ∗ ∗ where x encodes t at conditioned nodes and masks the rest. The resulting h∗ represents the counterfactual unit. To obtain the counterfactual predictive distribution, a node-wise MLP maps h∗ to its distributional parameters. We use a mixture of Nc Gaussians, allowing CFGNP to represent uncertainty arising from ambiguity over the underlying SCM. Training Objective. During training CFGNP minimizes the negative-log-likelihood objective, as this is equivalent to minimizing the KL divergence [Murphy, 2022]: Lθ = −

|L| d X X i=1 j=1

log

Nc X

∗ wk (τi )j N yi,j |µk (τi )j , σk (τi )j



.

k=1

Model Properties. The structure of CFGNP reflects two desirable properties. First, predictions should be independent of the ordering of variables. Permuting target nodes Y ∗ should induce the same permutation in the corresponding predictions, while permutations of the remaining nodes (Y ∗ )c should leave predictions for Y ∗ unchanged. This is provided by design of GNNs (App. .3). Second, a trained model should be applicable to graphs of different sizes and arbitrary numbers Nobs of observational samples (App. .4). This enables CFGNP to generalize beyond the graph sizes and context sizes encountered during training. 7

We note for completeness that, for n nodes, data/hidden dimensions ddata /dhidden , Nobs observations, K GNN iterations, and NC mixture components, one forward pass has complexity (.2):   Θ n ddata dhidden Nobs + KNobs ndhidden (dhidden + n) + nd2hidden + n2 dhidden + ndhidden NC .

5

Results

We evaluate CFGNP in terms of counterfactual prediction performance, generalization across causal graph sizes and unseen features, and computational efficiency, with additional experiments assessing its ability to recover analytically available counterfactual distributions with increasing context size (App. C.4). We also include an ablation study (App. C.5) where a normal transformer is run on the concatenated inputs. The results show that CFGNP performs much better than this. Baselines. We compare against the following baselines: • VCI [Wu et al., 2025]: uses a conditional VAE with a counterfactual variational objective under a restricted causal structure. • DeepBC [Kladny et al., 2024]: performs counterfactual inference through exogenous variables in a learned SCM, assuming a known causal graph (CG). • VACA [Sanchez-Martin et al., 2021]: uses a variational graph autoencoder on a known CG. • BGM [Nasr-Esfahany et al., 2023]: represents structural equations as invertible mappings between exogenous and observed variables given the causal structure. • CFM [Le et al., 2024]: learns flow-based causal mechanisms assuming a given causal ordering. Unlike the baseline methods, CFGNP uses counterfactual pairs during training, but no causal graph. Metrics. We evaluate counterfactual predictions using the following metrics: • MSE/MAE: Mean squared/absolute error (MSE/MAE) for numerical variables. • Accuracy/ROC-AUC: Accuracy and ROC-AUC for categorical variables. For datasets containing both numerical and categorical variables, we report metrics separately for each variable. For synthetic experiments with multiple variables, we follow Sanchez-Martin et al. [2021] and report MSE aggregated across all variables and samples. Synthetic Datasets. For the synthetic experiments, we use the Triangle, M-shape, and German Loan benchmarks from Sanchez-Martin et al. [2021], together with randomly generated SCMs for the graph-size experiments. Triangle and M-shape have 3 and 5 nodes respectively and are evaluated with linear (LIN), nonlinear additive (NLIN), and nonlinear non-additive (NADD) mechanisms, while German Loan is a 7-node semi-synthetic seven-variable SCM inspired by the German Credit domain. Detailed definitions can be found in B.3. Further, we conduct experiments on causal graphs of increasing size where each sample comes from its own SCM. We follow Soleymani et al. [2020]. We generate the SCMs by randomly permuting the nodes into a topological order and adding each admissible edge Xi → Xj with probability 12 . Each variable X is generated from its direct causes using one of seven mechanisms f per sample, followed by posterior additive noise (App. B.4). We consider causal graphs of sizes S = 10, 15, 20, 25, 30. Real-world Datasets. For the real-world experiments, we use the CheXpert dataset [Irvin et al., 2019], consisting of chest X-rays (App. B.2). We augment each image with four explicit 8

variables: age, sex, disease, and AP/PA, where AP/PA denotes the radiographic viewing direction. Following Ribeiro et al. [2025], we supply the baselines with a causal structure in which age affects disease, while age, sex, disease, and AP/PA determine the image. Since ground-truth counterfactuals are unavailable [Ma et al., 2026], we evaluate whether generated images exhibit the intended counterfactual attributes using classifiers. Details on the classifiers and image-generating CFGNP variant are provided in App. B.2. We assess whether images resemble real-world images qualitatively.

5.1

Synthetic Experiments

On synthetic experiments, we compare against VACA [Sanchez-Martin et al., 2021], BGM [Nasr-Esfahany et al., 2023], and CFM [Le et al., 2024], which have demonstrated strong performance in prior work. VACA and BGM require the full causal graph, whereas CFM only requires its topological ordering, allowing comparison across different levels of causal knowledge. For the graph-size experiments, we omit CFM as it was not computationally viable on larger causal graphs. Counterfactual Prediction. Table 1 shows the results for the synthetic datasets. CFGNP improves across almost all variants of the datasets and in many by a substantial amount. For instance, for German loan, CFGNP has an MSE of 0.19 × 10−2 compared to 4.51 × 10−2 of CFM. However, on the LIN variant of M-Shape, the MSE of 3.60 × 10−5 is higher than BGM’s with 1.46 × 10−5 . Dataset

Variant

BGM

Triangle

LIN NLIN NADD

M-shape German Loan

VACA

(3.71 ± 0.08) × 10−3 (4.78 ± 0.97) × 10−6 (1.80 ± 1.80) × 10−6

CFM MSE ↓ (23.52 ± 1.83) × 10−2 (3.32 ± 0.05) × 10−3 (13.24 ± 1.93) × 10−2 (1.35 ± 0.24) × 10−5 (22.50 ± 2.61) × 10−2 (1.54 ± 0.40) × 10−6

CFGNP (Ours) (2.59 ± 2.13) × 10−7 (2.14 ± 0.34) × 10−7 (2.81 ± 1.11) × 10−7

LIN NLIN NADD

(1.46 ± 0.40) × 10−5 (1.97 ± 0.50) × 10−5 (1.97 ± 0.07) × 10−3

(12.83 ± 0.69) × 10−2 (16.26 ± 0.95) × 10−2 (41.06 ± 0.31) × 10−2

(2.73 ± 0.17) × 10−5 (4.67 ± 1.14) × 10−5 (1.89 ± 0.03) × 10−3

(3.60 ± 7.00) × 10−5 (2.50 ± 0.47) × 10−6 (1.08 ± 0.35) × 10−4

–

(46.81 ± 4.20) × 10−2

(10.30 ± 0.40) × 10−2

(4.51 ± 0.78) × 10−2

(0.19 ± 0.09) × 10−2

Table 1: Mean MSE ± standard deviation achieved by BGM, VACA, CFM, and CFGNP (ours) across the Triangle, M-shape, and German Loan datasets. Lower values indicate better performance. Best scores are marked in bold. Graph-Size Experiments. For experiments with increasing graph size we use two variants of CFGNP. In one variant, CFGNP is retrained for each s ∈ S, while in the other, CFGNP is trained only on graphs with s = 10 nodes and applied to the other graph sizes without training. We refer to this model as 10-CFGNP. Additionally, as the benchmarks require a fixed causal graph, we provide the fullest DAG, connecting each node to all topological successors. This highlights an advantage of CFGNP: when causal graphs vary across the dataset, baselines must still assume a fixed graph. Figure 4 shows MSE with increasing graph size (note the non-linear y-axis). 10-CFGNP performs similarly to retrained CFGNP for Linear and Sum-sine, but worse for Log-sum-exp and Sqrt-sum. MSE increases for all methods, but both CFGNP variants outperform the baselines and degrade more slowly. Other mechanisms are reported in App. C.2.

5.2

CheXpert Experiments

We compare against VCI [Wu et al., 2025] and DeepBC [Kladny et al., 2024] which both have previously demonstrated strong performance. Again, VCI assumes treatment–covariate–outcome structure while DeepBC needs full graph knowledge. Counterfactual Prediction. Table 2 shows that CFGNP performs particularly well on age, 9

Figure 4: Performance (MSE) of models with regard to size of the causal graph. Models are compared on causal graphs with Linear (1.), Sum-sine (2.), Log-sum-exp (3.), Sqrt-sum (4.), and Sum-sqrt (5.) structural equations. MSE mean over distinct causal graphs. Vertical lines show 25th and 75th quantiles. reducing MAE from 10.88 for DeepBC to 2.71, and consistently achieves high ROC-AUC on categorical variables. For disease, for example, accuracies are around 0.87, while CFGNP achieves a ROC-AUC of 0.77, compared with 0.75 for DeepBC and 0.45 for VCI. Generated counterfactual images are provided in Appendix C.2. Feature

Metric

CFGNP (Ours)

DeepBC

VCI

Sex

Accuracy ↑ ROC-AUC ↑

0.8452 ± 0.0608 0.9155 ± 0.0540

0.8222 ± 0.1220 0.9204 ± 0.0752

0.5173 ± 0.0296 0.5238 ± 0.0122

Age

MAE ↓

2.7136 ± 0.7326

10.8750 ± 0.8732

66.8172 ± 62.1769

AP/PA

Accuracy ↑ ROC-AUC ↑

0.9664 ± 0.0140 0.9787 ± 0.0136

0.9572 ± 0.0552 0.9658 ± 0.0584

0.8597 ± 0.0000 0.5741 ± 0.0342

Disease

Accuracy ↑ ROC-AUC ↑

0.8830 ± 0.0368 0.7710 ± 0.0572

0.8720 ± 0.0240 0.7517 ± 0.0327

0.8548 ± 0.0667 0.4516 ± 0.0202

Table 2: Performance of CFGNP, VCI, and DeepBC on the CheXpert dataset. MAE is reported for the continuous feature, while accuracy and ROC-AUC are reported for categorical features. Arrows indicate whether higher or lower values are preferred. Best values marked in bold. Out-of-Distribution Variable-Set Experiment. We evaluate CFGNP’s generalization to unseen features on CheXpert. For each OOD setting, we exclude one feature kOOD during training and evaluate on the full test set, leveraging CFGNP’s ability to handle graphs of different sizes. Table 3 shows strong generalization for sex and AP/PA, with ROC-AUC above 0.99, while age and disease are more challenging, with an MAE of up to 11.53 and ROC-AUC between 0.61 and 0.64, respectively. Feature

Metric

OOD Sex

OOD Age

OOD AP/PA

OOD Disease

Sex

Accuracy ↑ ROC-AUC ↑

0.9516 ± 0.0108 0.9927 ± 0.0037

0.9634 ± 0.0157 0.9941 ± 0.0039

0.9692 ± 0.0084 0.9965 ± 0.0017

0.9627 ± 0.0029 0.9964 ± 0.0005

Age

MAE ↓

10.7037 ± 1.5283

11.5291 ± 0.0950

9.2965 ± 1.9205

8.1875 ± 0.4218

AP/PA

Accuracy ↑ ROC-AUC ↑

0.9917 ± 0.0069 0.9998 ± 0.0004

0.9904 ± 0.0085 0.9998 ± 0.0003

0.9995 ± 0.0008 1.0000 ± 0.0000

0.9996 ± 0.0005 1.0000 ± 0.0000

Disease

Accuracy ↑ ROC-AUC ↑

0.6560 ± 0.0176 0.6378 ± 0.0068

0.6678 ± 0.0119 0.6296 ± 0.0163

0.6690 ± 0.0073 0.6058 ± 0.0264

0.6579 ± 0.0212 0.6155 ± 0.0197

Table 3: OOD generalization on CheXpert. Each column indicates the feature excluded during training; evaluation is performed on the full test set containing all features. Best results in bold.

10

6

Conclusion

We propose CFGNP, an end-to-end architecture designed to learn counterfactual distributions directly from observational data without knowledge of the causal graph. CFGNP applies GNNs to enforce structural constraints on the dual graph implied by the notion of backtracking counterfactuals. By leveraging observational context samples alongside counterfactual queries, CFGNP directly approximates the Bayesian model-averaged counterfactual distribution. Evaluations on image and synthetic counterfactual queries have shown strong performance compared to baselines. Further experiments on synthetic datasets show that CFGNP retains strong performance on larger causal graphs compared with retraining from scratch, while outperforming baselines. Finally, we empirically validate that CFGNP performs adequately even on out-of-distribution features. In terms of limitations, the preceding capabilities come with a computational trade-off. The required observational context increases both training and inference costs. Moreover, if the counterfactual distribution is non-identifiable, CFGNP cannot recover the true distribution. However, this also applies to the baseline methods. In the future, one could incorporate priors about the causal graph. While the entire causal graph may not be known, an expert may have knowledge about the existence and absence of certain causal relationships. Taking this even further, one could provide simple distributions over edges, for instance obtained by amortized causal structure learning methods [Lorch et al., 2022].

AI use statement In this work, we used generative AI tools for refining writing, formulating mathematical claims, assisting in the writing of proofs, and feedback on research methodology or experiments. We have not used generative AI tools for generating synthetic data sets, helping develop theoretical models or conceptual frameworks, providing critical ingredients for proving mathematical claims, implementing methods, assisting with translation, cleaning and reformatting datasets, and interpreting results. Support for qualitative and thematic data analysis and proposing or refining hypotheses are not applicable to this work. We have reviewed all AI-assisted work. We checked that the generated formulation of the claim matches the intended claim. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI.

Ethics statement This work uses existing datasets and collects no new human-subject data. While counterfactual models may have applications in healthcare, their predictions can reflect biases and limitations in the data and should not be used for clinical decision-making without appropriate validation.

Reproducibility statement To support reproducibility, we provide the source code and experimental setup. The repository can be found under https://github.com/Zelos1/CFGNP. Model architectures and hyperparameters are detailed in the appendix.

References Amin, A. N. and Wilson, A. G. (2024). Scalable and flexible causal discovery with an efficient test for adjacency. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F., editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 1331–1358. PMLR.

11

Ashman, M., Diaconu, C., Kim, J., Sivaraya, L., Markou, S., Requeima, J., Bruinsma, W. P., and Turner, R. E. (2024). Translation equivariant transformer neural processes. arXiv preprint arXiv:2406.12409. Brody, S., Alon, U., and Yahav, E. (2021). How attentive are graph attention networks? arXiv preprint arXiv:2105.14491. Brody, S., Alon, U., and Yahav, E. (2022). How attentive are graph attention networks? Bruinsma, W. (2022). Convolutional Conditional Neural Processes. PhD thesis. Bruinsma, W., Markou, S., Requeima, J., Foong, A. Y. K., Andersson, T., Vaughan, A., Buonomo, A., Hosking, S., and Turner, R. E. (2023). Autoregressive Conditional Neural Processes. In International Conference on Learning Representations. Burotto, M., Chiou, V. L., Lee, J.-M., and Kohn, E. C. (2014). The mapk pathway across different malignancies: A new perspective. Cancer, 120(22):3446–3456. Bynum, L. E. J., Ranganath, R., and Cho, K. (2026). Computational identifiability. Caron, A., Baio, G., and Manolopoulou, I. (2022). Counterfactual Learning with Multioutput Deep Kernels. arXiv:2211.11119 [cs]. Carr, A. and Wingate, D. (2019). Graph neural processes: Towards bayesian graph neural networks. Cheng, Y., Li, L., Xiao, T., Li, Z., Suo, J., He, K., and Dai, Q. (2024). Cuts+: High-dimensional causal discovery from irregular time-series. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 11525–11533. Chickering, D. M. (2002). Optimal structure identification with greedy search. Journal of machine learning research, 3(Nov):507–554. Dhir, A., Diaconu, C., Lungu, V. M., Requeima, J., Turner, R. E., and Wilk, M. v. d. (2025). Estimating Interventional Distributions with Uncertain Causal Graphs through Meta-Learning. arXiv:2507.05526 [cs]. Diepen, M. M. V., Bucur, I. G., Heskes, T., and Claassen, T. (2023). Beyond the markov equivalence class: Extending causal discovery under latent confounding. In van der Schaar, M., Zhang, C., and Janzing, D., editors, Proceedings of the Second Conference on Causal Learning and Reasoning, volume 213 of Proceedings of Machine Learning Research, pages 707–725. PMLR. Foong, A., Bruinsma, W., Gordon, J., Dubois, Y., Requeima, J., and Turner, R. (2020). Meta-learning stationary stochastic process prediction with convolutional neural processes. Advances in Neural Information Processing Systems, 33:8284–8295. Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F., Rezende, D. J., Eslami, S. M. A., and Teh, Y. W. (2018). Neural Processes. arXiv:1807.01622 [cs]. Geffner, T., Antoran, J., Foster, A., Gong, W., Ma, C., Kiciman, E., Sharma, A., Lamb, A., Kukla, M., Pawlowski, N., Allamanis, M., and Zhang, C. (2022). Deep end-to-end causal inference. Gichoya, J. W., Banerjee, I., Bhimireddy, A. R., Burns, J. L., Celi, L. A., Chen, L.-C., Correa, R., Dullerud, N., Ghassemi, M., Huang, S.-C., et al. (2022). Ai recognition of patient race in medical imaging: a modelling study. The Lancet Digital Health, 4(6):e406–e414.

12

Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. (2017). Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263–1272. Pmlr. Gordon, J., Bronskill, J., Bauer, M., Nowozin, S., and Turner, R. E. (2019). Meta-learning probabilistic inference for prediction. He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778. Heckman, J. J. and Pinto, R. (2022). Causality and econometrics. Working Paper 29787, National Bureau of Economic Research. Hornik, K., Stinchcombe, M., and White, H. (1989). Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359–366. Hu, J., Liang, Y., Fan, Z., Chen, H., Zheng, Y., and Zimmermann, R. (2023). Graph neural processes for spatio-temporal extrapolation. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD ’23, page 752–763. ACM. Irvin, J., Rajpurkar, P., Ko, M., Yu, Y., Ciurea-Ilcus, S., Chute, C., Marklund, H., Haghgoo, B., Ball, R., Shpanskaya, K., Seekins, J., Mong, D. A., Halabi, S. S., Sandberg, J. K., Jones, R., Larson, D. B., Langlotz, C. P., Patel, B. N., Lungren, M. P., and Ng, A. Y. (2019). Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. Javaloy, A., Sánchez-Martín, P., and Valera, I. (2023). Causal normalizing flows: from theory to practice. arXiv:2306.05415 [cs]. Javanmardi, M., Sajjadi, M., Liu, T., and Tasdizen, T. (2018). Unsupervised total variation loss for semi-supervised deep learning of semantic segmentation. Kaltenpoth, D. and Vreeken, J. (2023). Nonlinear causal discovery with latent confounders. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 15639–15654. PMLR. Karimi, A.-H., Kügelgen, J. v., Schölkopf, B., and Valera, I. (2020a). Algorithmic recourse under imperfect causal knowledge: a probabilistic approach. arXiv:2006.06831 [cs]. Karimi, A.-H., von Kügelgen, J., Schölkopf, B., and Valera, I. (2020b). Algorithmic recourse under imperfect causal knowledge: a probabilistic approach. Khemakhem, I., Monti, R. P., Leech, R., and Hyvärinen, A. (2021). Causal Autoregressive Flows. arXiv:2011.02268 [stat]. Kim, H., Mnih, A., Schwarz, J., Garnelo, M., Eslami, A., Rosenbaum, D., Vinyals, O., and Teh, Y. W. (2019). Attentive neural processes. arXiv preprint arXiv:1901.05761. Kladny, K.-R., Kügelgen, J. v., Schölkopf, B., and Muehlebach, M. (2024). Deep Backtracking Counterfactuals for Causally Compliant Explanations. arXiv:2310.07665 [cs]. Kullback, S. and Leibler, R. A. (1951). On information and sufficiency. The annals of mathematical statistics, 22(1):79–86. Le, M. K., Do, K., and Tran, T. (2024). Learning Structural Causal Models from Ordering: Identifiable Flow Models. arXiv:2412.09843 [cs]. Li, C., Shen, X., and Pan, W. (2024). Nonlinear causal discovery with confounders. Journal of the American Statistical Association, 119(546):1205–1214. 13

Li, Q., Han, Z., and Wu, X.-M. (2018). Deeper insights into graph convolutional networks for semi-supervised learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32. Li, T. and Han, L. (2023). Dealing with explainability requirements for machine learning systems. In 2023 IEEE 47th Annual Computers, Software, and Applications Conference (COMPSAC), pages 1203–1208. Liang, H. and Gao, J. (2021). How neural processes improve graph link prediction. Loh, W. W. (2025). Causal inference with unobserved confounding: leveraging negative control outcomes using lavaan. Multivariate Behavioral Research, 60(5):1029–1041. Lorch, L., Sussex, S., Rothfuss, J., Krause, A., and Schölkopf, B. (2022). Amortized inference for causal structure learning. Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. Luedtke, A. and Fukumizu, K. (2025). DoubleGen: Debiased Generative Modeling of Counterfactuals. arXiv:2509.16842 [stat]. Ma, H., Quinzan, F., Willem, T., and Bauer, S. (2026). Ai alignment in medical imaging: Unveiling hidden biases through counterfactual analysis. Morgan, S. L. and Winship, C. (2007). Counterfactuals and causal inference. New York: Cambridge. Murphy, K. P. (2022). Probabilistic machine learning: an introduction. MIT press. Nasr-Esfahany, A., Alizadeh, M., and Shah, D. (2023). Counterfactual Identifiability of Bijective Causal Models. arXiv:2302.02228 [stat]. Nguyen, T.-L., Collins, G. S., Landais, P., and Le Manach, Y. (2020). Counterfactual clinical prediction models could help to infer individualized treatment effects in randomized controlled trials—an illustration with the international stroke trial. Journal of clinical epidemiology, 125:47–56. Pawlowski, N., Castro, D. C., and Glocker, B. (2020). Deep Structural Causal Models for Tractable Counterfactual Inference. arXiv:2006.06485 [stat]. Pearl, J. (1996). A causal calculus for statistical research. In Learning from Data: Artificial Intelligence and Statistics V, pages 23–33. Springer. Pearl, J. (2009). Causality. Cambridge University Press, 2 edition. Peters, J., Janzing, D., and Schölkopf, B. (2017). Elements of causal inference: foundations and learning algorithms. The MIT press. Rasal, R., Kori, A., Ribeiro, F. D. S., Xia, T., and Glocker, B. (2025). Diffusion Counterfactual Generation with Semantic Abduction. arXiv:2506.07883 [cs]. Ribeiro, F. D. S., Santhirasekaram, A., and Glocker, B. (2025). Counterfactual Identifiability via Dynamic Optimal Transport. arXiv:2510.08294 [cs]. Ruiz-Perez, D., Isabella, G., Musfiqur, S., Kalai, M., and Giri, N. (2024). Unfolding and deconfounding: biologically meaningful causal inference from longitudinal multi-omic networks using metalica. Msystems, 9(10). Sanchez-Martin, P., Rateike, M., and Valera, I. (2021). VACA: Design of Variational Graph Autoencoders for Interventional and Counterfactual Queries. arXiv:2110.14690 [stat]. 14

Sanderson, E., Glymour, M. M., Holmes, M. V., Kang, H., Morrison, J., Munafò, M. R., Palmer, T., Schooling, C. M., Wallace, C., Zhao, Q., et al. (2022). Mendelian randomization. Nature reviews Methods primers, 2(1):6. Schölkopf, B., Locatello, F., Bauer, S., Ke, N. R., Kalchbrenner, N., Goyal, A., and Bengio, Y. (2021). Toward causal representation learning. Proceedings of the IEEE, 109(5):612–634. Shimizu, S., Hoyer, P. O., Hyvarinen, A., and Kerminen, A. (2006). A linear non-gaussian acyclic model for causal discovery. Journal of Machine Learning Research, 7(72):2003–2030. Shpitser, I. and Pearl, J. (2012). What counterfactuals can be tested. arXiv preprint arXiv:1206.5294. Soleymani, A., Raj, A., Bauer, S., Schölkopf, B., and Besserve, M. (2020). Causal feature selection via orthogonal search. arXiv preprint arXiv:2007.02938. Spirtes, P., Glymour, C., and Scheines, R. (2001). Causation, prediction, and search. The MIT press. Stovitz, S. D. and Shrier, I. (2019). Causal inference for clinicians. BMJ Evidence-Based Medicine, 24(3):109–112. Verma, T. S. and Pearl, J. (2022). Equivalence and synthesis of causal models. In Probabilistic and causal inference: The works of Judea Pearl, pages 221–236. von Kügelgen, J., Mohamed, A., and Beckers, S. (2022). Backtracking Counterfactuals. Version Number: 3. Weberpals, J., Raman, S. R., Shaw, P. A., Lee, H., Russo, M., Hammill, B. G., Toh, S., Connolly, J. G., Dandreo, K. J., Tian, F., et al. (2024). A principled approach to characterize and analyze partially observed confounder data from electronic health records. Clinical Epidemiology, pages 329–343. Wu, Y., McConnell, L., and Iriondo, C. (2025). Counterfactual Generative Modeling with Variational Causal Inference. arXiv:2410.12730 [cs]. Xia, K., Pan, Y., and Bareinboim, E. (2022). Neural Causal Models for Counterfactual Identification and Estimation. arXiv:2210.00035 [cs]. Yu, Y., Chen, J., Gao, T., and Yu, M. (2019). DAG-GNN: DAG structure learning with graph neural networks. In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 7154–7163. PMLR. Zečević, M., Dhami, D. S., Veličković, P., and Kersting, K. (2021). Relating Graph Neural Networks to Structural Causal Models. arXiv:2109.04173 [cs]. Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. (2018). The unreasonable effectiveness of deep features as a perceptual metric. Zheng, X., Aragam, B., Ravikumar, P., and Xing, E. P. (2018). Dags with no tears: Continuous optimization for structure learning.

15

Appendix .1

Equivalence of a Direct MPNN to a Non-directional MPNN on the Dual Graph

This section provides an expressivity result justifying the removal of explicit exogenous nodes from the dual-graph architecture. Proposition 1 (Expressive sufficiency after eliminating exogenous nodes). Let Gexo denote the dual graph containing explicit exogenous nodes, and let Gdir denote the reduced graph obtained by removing the exogenous nodes and directly connecting each factual–counterfactual pair (xi , x∗i ). Assume that the endogenous subgraph of Gdir is fully connected. Then, for every function of the endogenous nodes that can be approximated arbitrarily well by a bidirectional MPNN on Gexo , there exists a direct MPNN on Gdir , with sufficiently large hidden dimension, that approximates the same function arbitrarily well. Proof. A bidirectional MPNN layer on Gexo has the form 



 (ℓ−1) , h(ℓ) v = Φℓ hv

X

Mℓ→ (h(ℓ−1) , h(ℓ−1) ), v w

(w,v)∈E

X

Mℓ← (h(ℓ−1) , h(ℓ−1) ) v w

(v,w)∈E

where Mℓ→ and Mℓ← are neural networks handling incoming and outcoming edges respectively. Since the values of the exogenous variables are not observed, the initial exogenous embeddings are fixed constants, (0)

h(0) ui = cU ,

hu∗ = cU ∗ . i

Consequently, the hidden state of every exogenous node is completely determined by the hidden states of neighbouring endogenous and exogenous nodes. Gdir does not have exogenous nodes, so we model them by doubling the hidden dimension of the endogenous nodes. Instead of explicitly representing ui and u∗i , the direct MPNN stores their hidden states inside the corresponding endogenous representations, (ℓ) (ℓ) (ℓ)
e hx∗ = hx∗ , hu∗ .

(ℓ) (ℓ) e h(ℓ) xi = hxi , hui ,

i

i

i

Assume inductively that these augmented representations reproduce the hidden states of the bidirectional MPNN after layer ℓ − 1. The bidirectional update of an exogenous node is a continuous function   (ℓ−1) (ℓ−1) (ℓ−1) h(ℓ) = F h , h , h , ℓ,i ui ui xi u∗ i

and analogously for u∗i . Since the endogenous graph is fully connected, node xi receives the representation of x∗i through an ordinary directed message. Hence, all arguments of Fℓ,i are available locally in the augmented representations. By the universal approximation theorem [Hornik et al., 1989], an MLP can approximate Fℓ,i arbitrarily well, allowing the direct MPNN to update the virtual exogenous states. Next, consider the update of an endogenous node,

16

  (ℓ−1) h(ℓ) , {h(ℓ−1) : y ∈ Nend (xi )}, h(ℓ−1) . xi = Gℓ,i hxi y ui (ℓ−1) (ℓ−1) The virtual state hui is stored locally in e hxi . Because the endogenous graph is fully connected, node xi receives the hidden state of every endogenous node. Hence, all arguments of Gℓ,i are available to the direct MPNN. Again, the update function can be approximated arbitrarily well by an MLP.

After each layer,
(ℓ) (ℓ) e h(ℓ) xi = hxi , hui , and similarly for x∗i . By induction, the hidden representations of all endogenous nodes produced by the bidirectional MPNN can be reproduced arbitrarily closely by the direct MPNN. Since the readout is performed only on endogenous nodes, reproducing the exogenous hidden states explicitly is unnecessary.

.2

CFGNP Complexity Analysis

We calculate the complexity in each part of the module. We first note that the GAT layer has computational complexity Θ(|V|dd′ + |E|d′ ) for a graph with |V| nodes and |E| edges and the hidden dimensions go from d to d′ [Brody et al., 2021]. Let |A| be the number of edges in the graph with adjacency matrix A. Changing Data Dimension Applying the MLP to each entry in O ∈ R(Nobs ·n)×ddata and z, x∗ ∈ R(Ncount ·n)×ddata takes (Nobs n + 2Ncount n) ddata dhidden ∈ Θ(n ddata dhidden (Ncount + Nobs )) operations. Since we use Ncount = 1 throughout this work we get complexity Θ(n ddata dhidden Nobs ) Propagation of Observational Information Each layer of the GAT uses adjacency matrix   Mn 0n · · · 0n In  .. ..   0n M n . . . . .     . ..  . . . Gc =  . . . . . 0n .   .   0n · · · 0n Mn In  0n · · · · · · 0n 0n With |Mn | = n(n − 1) and |In | = n we get |Gc | = Nobs n(n − 1) + Nobs n. Further, we have |V| = (Nobs + 1)n. For one layer we then get computational complexity Θ((Nobs + 1)nd2hidden + (Nobs n(n − 1) + Nobs n)dhidden ). For all K layers in the block, we have computational complexity Θ(K(Nobs + 1)nd2hidden + K(Nobs n2 )dhidden ) = Θ(KNobs ndhidden (dhidden + n)). 17

Mixing Factual Information Each layer in the module for mixing factual information uses adjacency matrix   0n In ′ GC = . 0n In With |G′C | = 2n and |V| = 2n we get computational complexity Θ(nd2hidden + ndhidden ). All three layers for this module thus have complexity Θ(nd2hidden + ndhidden ). Estimating Counterfactual Graph Each layer in this block uses adjacency matrix   Mn In ′′ GC = . I n Mn We then get |G′′C | = 2n(n − 1) + 2n = 2n2 and |V| = 2n. This gives us computational complexity of Θ(nd2hidden + n2 dhidden ) in each layer. All min{⌊ n2 ⌋, 5} layers together give us computational complexity Θ(nd2hidden + n2 dhidden ) for the entire block. The MLP layer needs Θ(n dhidden NC ) operations. Overall, the computational complexity of CFGNP is  Θ n ddata dhidden Nobs + KNobs ndhidden (dhidden + n)  + nd2hidden + n2 dhidden + ndhidden NC .

.3

Permutation Properties of CFGNP

Proposition 2. Let F (z, T ∗ , t∗ , Y ∗ , O) denote the output of the CFGNP model restricted to the target nodes Y ∗ . Then the CFGNP is 1. permutation equivariant with respect to permutations of the nodes in Y ∗ ; and 2. permutation invariant with respect to permutations of the nodes in (Y ∗ )c . Proof. In the following, let Mn be the n × n matrix that has 0 on the diagonal and everywhere else 1. Let In be the n × n identity matrix. Finally, let 0n be the n × n matrix with only 0. Let P ∈ {0, 1}n×n be the permutation matrix associated with a permutation π of the n variables. We use the convention P ei = eπ(i) , where ei denotes the i-th standard basis vector. Consequently, for every node-indexed matrix H ∈ Rn×d , (P H)j = Hπ−1 (j) . A permutation represents only a relabelling of the variables and must therefore be applied consistently to every node-indexed input. We define zP := P z,

x∗P := P x∗ ,

Ok,P := P Ok

for every observational sample Ok . We further define Tπ∗ := π(T ∗ ),

Yπ∗ := π(Y ∗ ). 18

Let τT ∗ ∈ {0, 1}n be the indicator vector of the intervention set T ∗ , i.e., ( 1, i ∈ T ∗ , (τT ∗ )i = 0, i ∈ / T ∗. The indicator vector of the permuted intervention set Tπ∗ is τT ∗ ,P := P τT ∗ . Indeed, for every j ∈ {1, . . . , n}, (τT ∗ ,P )j = (P τT ∗ )j = (τT ∗ )π−1 (j) , which is equal to 1 precisely when π −1 (j) ∈ T ∗ , or equivalently when j ∈ π(T ∗ ). Recall that the counterfactual input x∗ contains the intervention values at the indices in T ∗ and is zero elsewhere: ( t∗ , i ∈ T ∗ , ∗ xi = i 0, i ∈ / T ∗. For every j,

(x∗P )j = (P x∗ )j = x∗π−1 (j) .

(x∗P )j =

( t∗π−1 (j) , π −1 (j) ∈ T ∗ , 0,

⇔

π −1 (j) ∈ / T ∗.

( (P t∗ )j , j ∈ Tπ∗ , (x∗P )j = 0, j∈ / Tπ∗ .

Thus, the intervention values and the intervention indices are permuted consistently. We now show the characteristics for each of the model components. Equivariance of nodewise MLPs. We first consider the nodewise MLPs applied at the beginning of the model. For every observational sample Ok , the same MLP hO is applied independently to every node. Therefore, hO (Ok,P ) = P hO (Ok ). For the factual observation z, define the nodewise transformation
hz (z, τT ∗ )i = (τT ∗ )i hz,int (zi ) + 1 − (τT ∗ )i hz,common (zi ). We now verify permutation equivariance coordinatewise. For an arbitrary j ∈ {1, . . . , n}:


hz (zP , τT ∗ ,P )j = (τT ∗ ,P )j hz,int (zP )j + 1 − (τT ∗ ,P )j hz,common (zP )j
= (τT ∗ )π−1 (j) hz,int zπ−1 (j)

+ 1 − (τT ∗ )π−1 (j) hz,common zπ−1 (j) = hz (z, τT ∗ )π−1 (j)
= P hz (z, τT ∗ ) j . Since this equality holds for every j, hz (P z, P τT ∗ ) = P hz (z, τT ∗ ). The same argument applies to the counterfactual input. hx (P x∗ , P τT ∗ ) = P hx (x∗ , τT ∗ ).

19

Equivariance of the GATs. We now consider the GATs. A GAT mθ is permutation equivariant in the sense that, for a node-feature matrix H, adjacency matrix G, and permutation matrix Q, mθ (QH, QGQ⊤ ) = Qmθ (H, G). For the observational graph, the same permutation is applied to every observational sample and to the aggregation nodes q. The corresponding block-diagonal permutation matrix is   QO = diag  P, . . . , P  = INobs +1 ⊗ P. | {z } Nobs +1 times

The adjacency matrix Gc consists only of blocks Mn , In , and 0n . Since P Mn P ⊤ = Mn ,

P In P ⊤ = In ,

P 0n P ⊤ = 0n ,

we have QO Gc Q⊤ O = Gc . This suppresses the unchanged adjacency matrix from the notation, (1) mθ

  ′   ′    e O O O (1) QO = QO mθ = QO ′ . q q q

It follows in particular that qP′ = P q ′ . For the graph containing the factual representation z and the aggregated observational representation q ′ , define the permutation   P 0n Qzq = . 0n P ′ Because Qzq G′C Q⊤ zq = GC , permutation equivariance gives (2) mθ



    z z (2) Qzq ′ = Qzq mθ q q′  ′ z = Qzq ′′ q  ′ Pz = . P q ′′

For the factual–counterfactual graph, define  Qzx∗ = With

 P 0n . 0n P

′′ Qzx∗ G′′C Q⊤ zx∗ = GC ,

we similarly obtain mcf θ

  ′   ′  z z cf Qzx∗ = Qzx∗ mθ ∗ x∗ x  ′′  z = Qzx∗ x∗ ′  ′′  Pz = . P x∗ ′

In particular, the counterfactual representation satisfies x∗ ′P = P x∗ ′ . 20

The final MLP. The final MLP is applied independently to every node of x∗ ′ , with the same parameters at each node. Therefore, (3)

(3)

mθ (P x∗ ′ ) = P mθ (x∗ ′ ). Let

(3)

R := mθ (x∗ ′ )

denote the nodewise output containing the mixture parameters. The preceding arguments imply that the nodewise output of the model transforms as RP = P R. The final nodewise output R ∈ Rn×3Nc contains one row for each node, and the CFGNP returns only the rows indexed by the target set: F (z, T ∗ , t∗ , Y ∗ , O) = R Y ∗ . Moreover, the target set is relabelled according to Yπ∗ = π(Y ∗ ). Since for every i ∈ Y ∗ ,

(P R)π(i) = Ri

restricting P R to the indices in Yπ∗ yields precisely the same nodewise outputs as restricting R to Y ∗ , but indexed according to the permuted target nodes. Hence,
F (zP , Tπ∗ , t∗P , Yπ∗ , OP ) = (P R) Y ∗ = PY ∗ R Y ∗ , π

where PY ∗ denotes the permutation induced by π on the target outputs. Therefore, F (zP , Tπ∗ , t∗P , Yπ∗ , OP ) = PY ∗ F (z, T ∗ , t∗ , Y ∗ , O). This establishes permutation equivariance with respect to the target nodes. If the permutation acts only on nodes in (Y ∗ )c , then every target node is fixed and the induced permutation PY ∗ is the identity. The same identity therefore reduces to F (zP , Tπ∗ , t∗P , Y ∗ , OP ) = F (z, T ∗ , t∗ , Y ∗ , O), which establishes permutation invariance with respect to permutations of the non-target nodes.

.4

Graph-Size and Observation-Count Transferability of CFGNP

The CFGNP architecture can be applied to causal graphs with different numbers of nodes and to observational datasets containing different numbers of samples. In particular, the learnable parameters of the model do not depend on either the graph size n or the number of observational samples Nobs . Consequently, a model trained using graphs with n nodes and Nobs observational samples can be applied, without modifying its learned parameters, to ′ graphs with m nodes and Nobs observational samples. We formalize this property in the following proposition. Proposition 3 (Graph-size and observation-count transferability). Let (n,Nobs )

Fθ

denote the CFGNP applied to a causal query with n endogenous variables and Nobs observational samples. Suppose that the number of message-passing layers is fixed depending on n and 21

independently of Nobs , or, alternatively, repeated message-passing iterations reuse the same parameters. (n,Nobs )

Then the parameter space Θ of Fθ arbitrary

is independent of both n and Nobs . Consequently, for

′ n, m, Nobs , Nobs ∈ N,

the same parameter vector θ ∈ Θ defines both (n,Nobs )

Fθ

and

′ ) (m,Nobs

Fθ

.

Hence, a CFGNP trained on one graph size and one number of observational samples can be evaluated on different graph sizes and different numbers of observational samples without introducing, removing, or resizing any learned parameters. Proof. We show that every learnable component of the CFGNP is independent of both the number of endogenous variables n and the number of observational samples Nobs . First, consider the node-wise input transformations. Let hO : Rddata → Rdhidden denote the MLP applied to the observational data. For O ∈ RNobs ×n×ddata , the same map hO is applied independently to every node in every observational sample, Ok,i 7→ hO (Ok,i ),

k = 1, . . . , Nobs ,

i = 1, . . . , n.

Thus, changing either n or Nobs changes only the number of evaluations of hO , but not its input dimension, output dimension, or learned parameters. The same argument applies to the node-wise maps hz,int , hz,common , hx,int , and hx,common . Next, consider a GAT layer. For a node i, its update is of the form   X h′i = σ  αij W hj  , j∈N (i)

where αij = P and

exp(eij ) , k∈N (i) exp(eik )

  eij = LeakyReLU a⊤ [W hi ∥W hj ] .

The learnable quantities W and a depend only on the feature dimensions. They do not depend on the total number of nodes in the graph, the number of edges, or the number of observational samples represented in the graph. Changing n or Nobs therefore changes only the set of nodes and neighborhoods on which the same local update rule is evaluated. We now consider the observational aggregation graph Gc . For n endogenous variables and Nobs observational samples, its adjacency matrix is   Mn 0n · · · 0n In   ..   0n Mn . . . . I n   (n,Nobs )  . . .. .. Gc = . ..  . . . 0n .    0n · · · 0n Mn In  0n · · · · · · 0n 0n 22

(n,N

)

Thus, Gc obs contains Nobs diagonal blocks Mn , one for each observational sample, and Nobs corresponding blocks In connecting the observational nodes to the aggregation nodes q. Changing Nobs only changes the number of Mn - and In -blocks, while changing n only changes (1) their dimensions. The construction rule itself remains unchanged. Since the GAT mθ applies the same attention and message-passing functions to every node and edge, its learned parameters are independent of both n and Nobs . In particular, if

(1)

[O′ , q ′ ] = mθ ([O′ , q]; Gc ), then changing the number of observational samples only changes the size of O′ and the number of incoming messages available to the aggregation nodes q. It does not change the dimensionality of the representation qi′ ∈ Rdhidden associated with each endogenous variable. Hence, the remainder of the architecture receives the same type of representation regardless of Nobs . Next, consider the graph used to combine the observational representation with the factual observation. For n variables,   0n In ′(n) GC = . 0n In For m variables, the corresponding graph is ′(m) GC =

  0m Im . 0m Im

These graphs have different dimensions but are generated according to the same construction rule. Moreover, the number Nobs no longer appears explicitly at this stage, because the observational samples have already been aggregated into the n representations contained in q ′ . Therefore, the parameters of the GAT used in this stage are independent of both n and Nobs . The same argument applies to the graph used for counterfactual inference. For n variables,   Mn In ′′(n) GC = , I n Mn whereas for m variables, ′′(m) GC =

  Mm Im . I m Mm

Changing the graph size changes only the number of nodes and edges on which the shared message-passing functions are evaluated. The learned parameters of the GAT are unchanged. Furthermore, Nobs does not affect the dimensionality of the node representations entering this stage. Finally, the output network is applied independently and identically to each target node, (3)

mθ : Rdhidden → R3Nc . Changing the number of endogenous variables changes only the number of node-wise predictions, while changing Nobs has no effect on the dimensions of the output map. Hence, neither quantity (3) changes the learned parameters of mθ . Consequently, every learned component of the CFGNP is parameterized only by quantities such as ddata , dhidden , Nc , and not by n or Nobs . The matrices Mn , In , and 0n and the repetition of the observational blocks determine only the structure and size of the computational graph; they contain no learned parameters. Therefore, ′ Θn,Nobs = Θm,Nobs =Θ

23

for arbitrary

′ n, m, Nobs , Nobs ∈ N.

Thus, the same parameter vector θ can be applied to causal graphs with different numbers of nodes and to observational datasets containing different numbers of samples. This establishes architectural transferability across both graph size and observation count. It does not, by itself, guarantee identical predictive performance when n or Nobs differs substantially from the values encountered during training, since such performance additionally depends on the distribution of training and test instances.

A

Extended Model Architecture

In the following, let Mn be the n × n matrix that has 0 on the diagonal and everywhere else 1. Let In be the n × n identity matrix. Finally, let 0n be the n × n matrix with only 0.

A.1

Model Input

Our input is a tuple ((z, T ∗ , t∗ , Y ∗ ), O) with O ∈ RNobs ×n×ddata , z ∈ R1×n×ddata , and x∗ ∈ RNcount ×n×ddata .

A.2

Changing data dimension

We apply MLPs on each node. On O we apply the same MLP hO across all nodes. However, for the other parts we want the model to learn to distinguish between nodes we intervene on and the rest. Therefore, we apply an MLP hz,int on the indices T ∗ in z and an MLP hz,common on the other nodes. Similarly, we apply the MLP hx,int on the indices T ∗ of x∗ and an MLP hx,common on the other indices.

A.3

Propagation of Observational Information

We reshape O into O′ ∈ R(Nobs ·n)×dhidden . Furthermore, we introduce a set of nodes q ∈ Rn×dhidden , initialized as q = 0. For the GAT we use the adjacency matrix   Mn 0n · · · 0n In  .. ..   0n M n . . . . .     . ..  . . . Gc =  . . . . . 0n .   .   0n · · · 0 n M n I n  0n · · · · · · 0n 0n (1)

We then propagate information on Gc using a Graph Attention Network (GAT) mθ consisting of K layers. A single GAT layer updates each node representation by computing attention (t) coefficients over its neighbors. Let hi denote the representation of node i at layer k. The attention coefficients are computed as (k)

(k)

αij = P

exp(eij )

, (k) exp(e ) k∈N (i) ik

  (k) (k) (k) eij = LeakyReLU a⊤ [W hi ∥W hj ] ,

where W and a are learnable parameters. The node representations are then updated according to   X (k+1) (k) (k) hi = σ αij W hj  . j∈N (i)

24

Applying the entire GAT to the initial node representations yields (1)

[O′ , q ′ ] = mθ ([O′ , q]).

A.4

Mixing Factual Observation (2)

We construct a new graph G′C and apply a GAT mθ

consisting of three layers.

The resulting adjacency matrix is G′C =



 0n In . 0n In

Applying the GAT three times yields (2)

[z ′ , q ′′ ] = mθ ([z, q ′ ]).

A.5

Estimating Counterfactual Distribution

We build a new graph G′′C with adjacency matrix   Mn In ′′ GC = . I n Mn To propagate information through the entire graph, we use min{⌊ n2 ⌋, 5} GAT layers, collected (2)

in the MPNN mθ , which yields (2)

[z ′′ , x∗ ′ ] = mθ ([z ′ , x∗ ]). (3)

Finally, we apply an MLP mθ with hidden dimension 3Nc . After reshaping, the network c outputs the parameters {µk , σk , wk }N k=1 for each node.

B

Experimental Setup

B.1

Model Hyperparameters

For the (semi-)synthetic datasets, we use a hidden dimension of 80, 5 attention heads in the graph attention network (GAT) layers [Brody et al., 2022], and a batch size of 128. For the predicted MoG we use NC = 3 Gaussian mixture components. For the image-based counterfactual datasets, which consist of chest X-rays, we use a hidden dimension of 256, 9 attention heads in the GAT layers, and a batch size of 32. We choose a hidden dimension of 256 to match the dimensionality of the encoded image representations, providing a natural choice for the width of the network. We train each model for 200 epochs although we include early stopping after 25 epochs. Further, we use the AdamW [Loshchilov and Hutter, 2019] with initial learning rate 10−4 and a weight decay of 10−5 .

B.2

CheXpert Experiment Setup

We use Nobs = 100 for the CheXpert samples. For each image, the disease information was represented by the 13 CheXpert findings. We combine them into one No Finding attribute. The Disease attribute is just the complement. The resulting dataset has some data imbalances. The Sex feature has 41.24% female values and the rest is male. Only 17.24% of samples have PA as value and only 9.05% of the samples exhibit a disease. Finally, the average age in the samples is 60.66 years. 25

Figure 5: The training curves of the classifiers and regression models used for the classification of the generated CheXpert images. Counterfactual Evaluation. To evaluate which attributes a generated counterfactual image exhibits, we use classifiers/regression models. For the classifiers and regression models of attributes Sex, AP/PA, Age, and Disease we stick closely to Gichoya et al. [2022]. We apply the pretrained ResNet [He et al., 2016] and subsequently apply multiple MLP layers. For regression attributes we have final dimension 1 and for the classification attributes we apply the softmax activation function in the end to normalize the outputs to probabilities. We add here the performance of the classification/regression models of the features. The Sex classifier has a test accuracy of 0.9901 and a ROC-AUC of 0.9995. The Age regression model has MAE 6.1473. Further, the AP/PA classifier has both accuracy and ROC-AUC of 1.0. Finally, the No Finding classifier has accuracy 0.9010 and ROC-AUC 0.9314. Figure 5 shows the training curves for each of the classifiers. CheXpert training objective. To the normal classification/regression losses we also add a reconstruction loss to ensure that the images look like actual images. We use MSE, LPIPS[Zhang et al., 2018], and the total variation [Javanmardi et al., 2018] of x̂cf , to get Lrec = λM MSE (xreal , x̂cf ) + λL LPIPS (xreal , x̂cf ) + λT V TV (x̂cf ) ,

(2)

with λM = 15, λL = 30, and λT V = 30. The resulting total loss combines the classification,

26

regression, and reconstruction objectives: L=

X

CE (yi , ŷi ) +

i∈C

X 1 r∈R

tr

MSE (yr , ŷr ) + Lrec .

(3)

where C and R denote the sets of classification and regression features, respectively.

B.3

Synthetic Experiments Setup

This appendix gives the complete definitions of the SCMs used in the (semi-) synthetic experiments. Their definitions follow the benchmark specification in Appendix E of SanchezMartin et al. [2021]. They are reproduced here so that the experimental data-generating processes can be understood without the original paper. We also keep to Sanchez-Martin et al. [2021] for the computation of the standard deviation of the test MSE. We train the model with 5 different seeds and then compute mean and standard deviation of the runs. We use Nobs = 100 in all synthetic experiments. B.3.1

Triangle

Figure 6: The causal graph of the triangle dataset. The causal graph has three nodes and three edges. The Triangle SCM has the underlying causal graph as depicted in Figure 6. Hence X1 affects X3 through the direct path X1 → X3 and through the mediated path X1 → X2 → X3 . Exogenous distributions. For the LIN and NLIN variants, U1 ∼ 21 N (−2, 1.5) + 21 N (1.5, 1).

(4)

U1 ∼ 12 N (−2.5, 1) + 21 N (2.5, 1).

(5)

For NADD,

The remaining noises are LIN: U2 ∼ N (0, 1),

U3 ∼ N (0, 1),

(6)

NLIN: U2 ∼ N (0, 0.1),

U3 ∼ N (0, 1),

(7)

NADD: U2 ∼ N (0, 0.25),

U3 ∼ N (0, 0.0625).

(8)

Here the second parameter of N (µ, ·) follows the convention used in the benchmark specification. Linear (LIN). X1 = U1 ,

(9)

X2 = −X1 + U2 ,

(10)

X3 = X1 + 0.25X2 + U3 .

(11)

27

Nonlinear additive (NLIN). X1 = U1 ,

(12)

3 X2 = −1 + + U2 , 1 + exp(−2X1 ) X3 = X1 + 0.25X22 + U3 .

(13) (14)

The causal mechanisms are nonlinear, but each exogenous variable enters its structural equation additively. Nonlinear non-additive (NADD). (15)

X1 = U1 , X2 = 0.25 sgn(U2 )

X12 1 + U22

,

X3 = −1 + 0.1 sgn(U3 ) X12 + X22 + U3 .

(16) (17)

The function sgn(·) returns the sign of its argument. B.3.2

M-graph

Figure 7: The causal graph of the M-shape data. It has five nodes and four edges. The M-graph contains five endogenous variables as shown in Figure 7. All exogenous variables are independent standard normal variables: Ui ∼ N (0, 1),

i = 1, . . . , 5.

(18)

Linear (LIN). X1 = U1 ,

(19)

X2 = U2 ,

(20)

X3 = X1 + U3 ,

(21)

X4 = −X2 + 0.5X1 + U4 ,

(22)

X5 = −1.5X2 + U5 .

(23)

X1 = U1 ,

(24)

X2 = U2 ,

(25)

X3 = X1 + 0.5X12 + U3 , X4 = −X2 + 0.5X12 + U4 , X5 = −1.5X22 + U5 .

(26)

Nonlinear additive (NLIN).

28

(27) (28)

Nonlinear non-additive (NADD). X1 = U1 ,

(29)

X2 = U2 ,

(30) (31)

X3 = X1 U3 , X4 = −X2 + 0.5X12 X5 = −1.5X22 U5 . B.3.3

U4 ,

(32) (33)

German Loan

Figure 8: The causal graph of the German loan data. The nodes are real attributes used for German credit rating system. they are connected based on the real-world connections. The Loan benchmark is a semi-synthetic SCM from the causal-recourse literature and is inspired by the German Credit domain. It contains seven endogenous variables as shown in Figure 8. In the figure, G is gender, A age, E education, L loan amount, D loan duration, I income, and S savings. The structural assignments are G = UG ,

(34)

A = −35 + UA ,  h i−1 E = −0.5 + 1 + exp 1 − 0.5G − (1 + exp(−0.1A))−1 − UE ,

(35) (36)

L = 1 + 0.01(A − 5)(5 − A) + G + UL ,

(37)

D = −1 + 0.1A + 2G + L + UD ,

(38)

I = −4 + 0.1(A + 35) + 2G + GE + UI ,

(39)

S = −4 + 1.5 I{I>0} I + US .

(40)

Here I{I>0} is one if I > 0 and zero otherwise. The exogenous distributions are UG ∼ Bernoulli(0.5),

(41)

UA ∼ Gamma(10, 3.5),

(42)

UE ∼ N (0, 0.25),

(43)

UL ∼ N (0, 4),

(44)

UD ∼ N (0, 9),

(45)

UI ∼ N (0, 4),

(46)

US ∼ N (0, 25).

(47)

29

B.4

Structural Equations for Artificial Graphs

Variables are generated as X = f (Pa(X)) + ϵ, where f is sampled from the following set of transformations: X Linear: f1 (Pa(X)) = a X ′ + b, a = 0.5, b = 0, X ′ ∈Pa(X)

Sum-sqrt:

X

f2 (Pa(X)) = a

p |X ′ | + b,

a = 0.5, b = 0,

sin(cX ′ ) + b,

a = 1, b = 0, c = 0.5,

tanh(cX ′ ) + b,

a = 1, b = 0, c = 2,

X ′ ∈Pa(X)

Sum-sine:

X

f3 (Pa(X)) = a

X ′ ∈Pa(X)

Sum-tanh:

X

f4 (Pa(X)) = a

X ′ ∈Pa(X)



 Geometric mean:

f5 (Pa(X)) = a 

Y

1 | Pa(X)|

|X ′ |

+ b,

a = 3, b = 0.1,

X ′ ∈Pa(X)



 Log-sum-exp:

f6 (Pa(X)) = a log 

′

X

eX  + b,

a = 1, b = log 2,

X ′ ∈Pa(X)

Sqrt-sum:

v u u u f7 (Pa(X)) = at

X

X ′ + b,

a = 1, b = 0.

X ′ ∈Pa(X)

B.5

Identifiability in a Two-Node SCM

Setup. Let X = (X0 , X1 ) follow the SCM X0 := U0 ,

X1 := aX0 + ⟨gλ (X0 ), W ⟩,

where U0 ∼ N (0, 1), W ∼ N (0, I2 ) are independent, and p φ(x) := arctan(σv x/σu ), s(x) := σu2 + σv2 x2 ,
gλ (x) := s(x) cos(λφ(x)), sin(λφ(x)) , λ ∈ [0, 1].

(48)

(49) (50)

Here a controls the observational mechanism, whereas λ controls the cross-world coupling. Proposition 4 (Observational invariance). For every λ, X0 ∼ N (0, 1),

X1 | X0 = x ∼ N ax, s(x)2 .

(51)

Hence the observational distribution is independent of λ. The same is true of the intervention Pdo(X0 =α) . Proof. Since

∥gλ (x)∥2 = s(x)2 cos2 (λφ(x)) + sin2 (λφ(x)) = s(x)2 ,

we have ⟨gλ (x), W ⟩ ∼ N (0, s(x)2 ), which gives (51). Under do(X0 = α), X1 = aα + ⟨gλ (α), W ⟩ ∼ N (aα, s(α)2 ), again independently of λ.

30

Backtracking semantics. Let the factual evidence be e = {X0 = x0 , X1 = x1 },

r := x1 − ax0 .

To construct a counterfactual world satisfying X0′ = α, we use the minimum-change exogenous coupling (U0′ , W ′ ) = (α, W ). (52) Equivalently, (52) is the solution of min c0 |u′0 − U0 |2 + cW ∥w′ − W ∥2

u′0 ,w′

s.t. u′0 = α,

for any c0 , cW > 0. Thus backtracking changes only the exogenous realization necessary to produce X0′ = α while preserving the downstream noise W . Theorem 5 (Backtracking counterfactual). Under (48)–(52),
X1∗ (α) | e ∼ N aα + kλ r, τλ2 , where kλ :=

s(α) cos(λ∆φ), s(x0 )

τλ2 := s(α)2 sin2 (λ∆φ),

(53) (54)

and ∆φ := φ(α) − φ(x0 ). Moreover, (53) is identical to the corresponding abduction–action–prediction counterfactual. Proof. The evidence fixes U0 = x0 and imposes ⟨gλ (x0 ), W ⟩ = r. With ĝ0 := gλ (x0 )/s(x0 ), Gaussian conditioning gives   r ⊤ W |e∼N ĝ0 , I − ĝ0 ĝ0 . s(x0 )

(55)

Backtracking changes U0 to α while preserving W , hence X1∗ (α) = aα + ⟨gλ (α), W ⟩. Therefore

E[X1∗ (α) | e] = aα +

r ⟨gλ (α), ĝ0 ⟩, s(x0 )

and Var[X1∗ (α) | e] = s(α)2 −

⟨gλ (α), gλ (x0 )⟩2 . s(x0 )2

Since ⟨gλ (α), gλ (x0 )⟩ = s(α)s(x0 ) cos(λ∆φ), substitution yields (53)–(54). The interventional abduction–action–prediction construction evaluates the same quantity aα + ⟨gλ (α), W ⟩ under the same conditional law (55), so the two counterfactual distributions coincide. Corollary 6 (Cross-world non-identifiability). Although λ is invisible to every observational and interventional distribution, it changes the backtracking counterfactual whenever λ 7→ (kλ , τλ ) is non-constant. In particular, τλ2 = 0 ⇐⇒ λ∆φ ∈ πZ. For λ = 0, X1∗ (α) | e = δaα+ s(α) r , s(x0 )

(56)

whereas for λ > 0 the counterfactual has positive conditional variance for almost every (x0 , α) with α ̸= x0 . 31

For λ = 1, g1 (x) = (σu , σv x), and k1 = Use the prior

σu2 + σv2 x0 α , s(x0 )2

τ12 =

a ∼ N (0, σa2 ),

and let

σu2 σv2 (α − x0 )2 . s(x0 )2

λ ∼ πλ , (m)

(57)

a ⊥ λ,

(m)

Dobs = {(x0 , x1 )}M m=1

be observational context data generated from (51). Writing E := Dobs ∪ {(x0 , x1 )}, we get for the optimal backtracking predictive Z pBCM (y | x0 , x1 , X0′ = α, Dobs ) = pbt (y | e, α, a, λ)p(a, λ | E) da dλ.

(58)

Proposition 7 (Exact posterior and model average). The likelihood is independent of λ, so p(a, λ | E) = p(a | E)πλ (λ). With (i)

wi := s(x0 )−2 , we have

  2 a | E ∼ N µa , σa|E ,

(59)

where −2 σa|E = σa−2 +

X

(i)

X

(i) (i)

wi x 0 x 1 .

(60)

  2 N y; µa βλ + kλ x1 , τλ2 + σa|E βλ2 πλ (dλ),

(61)

wi (x0 )2 ,

2 µa = σa|E

i∈E

i∈E

Consequently, Z pBCM (y | ·) = where βλ := α − kλ x0 . Proof. Proposition 4 implies that the likelihood does not depend on λ, which yields the posterior factorization. (60) is the Gaussian posterior for weighted linear regression. Finally, substituting r = x1 − ax0 into (53) gives aα + kλ r = a(α − kλ x0 ) + kλ x1 = aβλ + kλ x1 . Marginalizing the Gaussian posterior over a gives (61).

32

Identifiable versus non-identifiable regimes. Let p∗BCM := pbt (· | e, α, a∗ , λ∗ ) denote the backtracking counterfactual under the true mechanism. As M → ∞, 2 = Θ(M −1 ) → 0, σa|E

µa → a∗

a.s.,

and therefore Z pBCM =⇒

N y; a∗ α + kλ r∗ , τλ2 πλ (dλ),

r∗ := x1 − a∗ x0 .

(62)

If πλ = δλ∗ , the limit in (62) equals p∗BCM . If instead πλ is non-degenerate, then for almost every (x0 , α) with α ̸= x0 the map λ 7→ (kλ , τλ ) is non-constant, so the limiting Bayes predictor remains a genuine mixture over observationally indistinguishable cross-world mechanisms. Thus even infinite observational data identifies a but not λ: uncertainty about the observational mechanism vanishes, while uncertainty about the backtracking counterfactual need not.

C

Additional Results

C.1

Additional CheXpert Results

C.1.1

CheXpert Experiments

We include some counterfactual images generated by the model.

Figure 9: Example counterfactual of a CheXpert image. Observed factual unit (left) has age 80 and the intervention value is 58 in the counterfactual (middle). The difference between the pixels is on the right.

Figure 10: Example counterfactual of a CheXpert image. Observed factual unit (left) has AP/PA as 1 and the intervention value is 0 in the counterfactual (middle). The difference between the pixels is on the right. 33

Figure 11: An image where the observed factual unit of age is 74 and the counterfactual intervention is 85. C.1.2

CheXpert Out-of-Distribution Experiment

This model was trained with sex chosen to be the out-of-distribution feature. However, during inference the sex feature is included again to test the performance. Out-of-Distribution Feature Sex. We here include the results when the model was trained with OOD feature sex.

Figure 12: An image where the observed factual unit of age is 65 and the counterfactual intervention is 72. The model was trained with OOD index sex.

Figure 13: An image where the observed factual unit of AP/PA is 0 and the counterfactual intervention is 1. The model was trained with OOD index sex. Out-of-Distribution Feature Age. We here include the results when the model was trained with OOD feature age.

34

Figure 14: An image where the observed factual unit is age 30 and the counterfactual intervention is age 52. The model was trained with OOD index age.

Figure 15: An image where the observed factual unit of sex is 1 and the counterfactual intervention is 0. The model was trained with OOD index age. Out-of-Distribution Feature AP/PA. We here include the results when the model was trained with OOD feature AP/PA.

Figure 16: n image where the observed factual unit of sex is 0 and the counterfactual intervention is 1. The model was trained with OOD index AP/PA.

Figure 17: An image where the observed factual unit is age 26 and the counterfactual intervention is age 64. The model was trained with OOD index AP/PA. 35

Out-of-Distribution Feature Disease. We here include the results when the model was trained with OOD feature disease.

Figure 18: An image where the observed factual unit of AP/PA is 0 and the counterfactual intervention is 1. The model was trained with OOD index disease.

Figure 19: An image where the observed factual unit is age 67 and the counterfactual intervention is age 37. The model was trained with OOD index disease.

36

C.2

Graph Size Results

Figure 20: Performance (MSE) of models with regard to the size of the causal graph for sqrt-sum (left) and sum-sqrt (right) structural equations. The MSE is averaged over distinct causal graphs. Vertical lines show the 25th and 75th quantiles.

Figure 21: Performance (MSE) of models with regard to the size of the causal graph for geometric-mean (left) and sum-tanh (right) structural equations. The MSE is averaged over distinct causal graphs. Vertical lines show the 25th and 75th quantiles. 37

C.3

Runtime of Experiments

As before, we use a batch size of 128 for the (semi-)synthetic experiments and a batch size of 32 for the CheXpert dataset. All experiments are run on an NVIDIA H100 GPU. Dataset

CFGNP (Ours)

BGM

CFM

VACA

German Loan Triangle LIN Triangle NLIN Triangle NADD M-Shape LIN M-Shape NLIN M-Shape NADD

31.5719s 8.1853s 7.9175s 7.9084s 12.9088s 12.9274s 12.9784s

5.2537s 2.9878s 3.0323s 3.0364s 3.3865s 3.3941s 3.4396s

115.1188s 26.3716s 26.2195s 26.3470s 41.5352s 41.8969s 41.9043s

5.3015s 3.7258s 3.7365s 3.7963s 3.8361s 3.8385s 3.8209s

Artificial Size 10 Artificial Size 15 Artificial Size 20 Artificial Size 25 Artificial Size 30

18.4545s 27.6381s 39.2469s 54.7471s 74.2878s

3.8820s 5.3534s 9.2543s 8.5072s 13.9569s

– – – – –

2.4226s 2.6127s 2.6588s 2.5591s 2.7384s

Table 4: Runtime comparison of CFGNP (Ours), BGM, CFM, and VACA across (semi-) synthetic datasets. The datasets were used throughout the results section.

Dataset

CFGNP (Ours)

VCI

DeepBC

CheXpert

47.7369s

3.8506s

6032.4057s

Table 5: Runtime comparison of CFGNP (Ours), VCI, and DeepBC on the CheXpert dataset.

C.4

Convergence in (Non-)Identifiable Cases

We empirically evaluate the convergence behaviour of CFGNP. To this end, we evaluate CFGNP on an SCM for which the ground-truth counterfactual distribution can be computed analytically. We consider a causal graph with two nodes and at most one edge. In Appendix B.5, we show that, under certain choices of parameters, such an SCM is identifiable. Furthermore, we analytically compute the ground-truth counterfactual distribution p∗T := pT (Y ∗ = y ∗ | z, T ∗ = t∗ , f , G), given the true causal graph and structural mechanisms. We compare the distribution predicted by CFGNP to both p∗T and pT := pT (Y ∗ = y ∗ | O, z, T ∗ = t∗ ), which conditions only on the observations rather than on knowledge of the SCM (f , G). Subsequently, we repeat the experiment for a non-identifiable parameterization of the two-node SCM. In both cases, we use 40000 samples, split into 21000 train samples, 9000 validation samples, and 10000 test samples.

38

Figure 22: KL divergences for the two-node experiments across different numbers Nobs of observations supplied to each sample. Vertical lines show the 10th and 90th quantiles. Results are shown for the identifiable (left) and non-identifiable (right) settings.

C.5

Transformer Ablation

To assess the importance of the CFGNP architecture, we consider a Transformer-based ablation that replaces the core CFGNP processing mechanism with a standard Transformer encoder while retaining the same input and output representations. The observation, factual unit, and counterfactual/intervention inputs are first embedded separately using the same embedding dimensionality as CFGNP. Again, eparate embeddings are retained for intervened and nonintervened variables. Then, we concatenate all these vectors and run a transformer architecture on them. The output of the transformer for x∗ is handled in the same way as for CFGNP Unless stated otherwise, all hyperparameters are kept identical to those used for CFGNP, including the base hidden dimension, number of attention heads, number of mixture components, and input/output dimensionality. As in CFGNP, the effective embedding dimension is given by the base hidden dimension multiplied by the number of attention heads. The Transformerspecific hyperparameters are five encoder layers, a feed-forward dimension of four times the embedding dimension, a dropout rate of 0.1, and GELU activations. Pre-layer normalization is used within each Transformer encoder layer, followed by a final layer normalization after the encoder stack. Table 6 compares the resulting Transformer ablation against the full CFGNP model. Dataset

Variant

Triangle

LIN NLIN NADD

CFGNP (Ours) Transformer MSE ↓ (2.59 ± 2.13) × 10−7 (8.33 ± 0.25) × 10−4 (2.14 ± 0.34) × 10−7 (8.37 ± 0.46) × 10−4 (2.81 ± 1.11) × 10−7 (5.02 ± 0.03) × 10−4

M-shape

LIN NLIN NADD

(3.60 ± 7.00) × 10−5 (2.50 ± 0.47) × 10−6 (1.08 ± 0.35) × 10−4

(4.79 ± 0.12) × 10−3 (6.04 ± 0.19) × 10−3 (3.04 ± 0.17) × 10−3

German Loan

–

(0.19 ± 0.09) × 10−2

(46.92 ± 7.72) × 10−2

Table 6: Mean MSE ± standard deviation of the full CFGNP model and the Transformer ablation across the Triangle, M-shape, and German Loan datasets. Lower values indicate better performance. Best scores are marked in bold.

39

Related documents

• Amortized Counterfactual Prediction with Graph Neural Networks#1298220
• Convex optimization for actionable & plausible counterfactual explanations#779283
• Counterfactual explanations as interventions in latent space#478134
• P$^2$CE: Model-Agnostic Plausible Pareto-Optimal Counterfactual Explanations#1084257
• Artificial intelligence surrogates for treatment effect estimation with before-and-after data#1063391