Amortized Counterfactual Prediction with Graph Neural Networks

1 Introduction

• Counterfactual questions ask what would have happened in an alternative scenario given what occurred in the factual world.They support individualized outcome and explanation queries beyond population-level effects.Source: §1 · pp. 1–2
• Traditional approaches often require a known or estimated causal graph, although such structure may be unavailable or uncertain in domains such as medicine, biology, and social science.Estimating a graph first can add a super-exponential search space and errors in the graph may propagate to counterfactual predictions.Source: §1 · p. 2
• CFGNP is an end-to-end framework that learns counterfactual distributions without prior knowledge of the causal graph.It conditions predictions on observational samples and uses dual-graph structural constraints.Source: §1 · p. 2
• CFGNP outperforms existing counterfactual inference methods on both synthetic and real-world data.The paper also reports strong performance on queries involving features unseen during training.Source: §1 · p. 2

2 Related Work

• Causal discovery estimates a causal graph from associational data, but challenges include unobserved confounding, nonlinearities, and computational scalability.The cited methods range from constraint- and score-based approaches to continuous optimization and deep learning.Source: §2 · pp. 2–3
• Most counterfactual methods rely on a known, estimated, or otherwise fixed causal structure, while end-to-end learning under uncertain structure had not been established for counterfactual inference.The paper identifies dual-graph structural constraints as an opportunity for this setting.Source: §2 · p. 3
• Image-counterfactual methods commonly assume a known causal graph, while VCI avoids an explicit graph but assumes a prescribed treatment–covariate–outcome structure.The passage describes DeepBC, diffusion, and optimal-transport approaches as graph-requiring examples.Source: §2 · p. 3
• Neural Processes map observational data to predictive distributions, and CFGNP extends this approach to counterfactual targets using GNNs over the dual-graph structure.Earlier Neural Process work in causal settings addressed related tasks rather than counterfactual inference.Source: §2 · pp. 3–4

3.1 Structural Causal Models and Causal Inference

• An SCM represents endogenous variables through structural equations, with each variable determined by its direct causes and an exogenous noise variable.The causal graph is directed and acyclic, with edges encoding direct causal relations.Source: §3.1 · p. 4
• Backtracking counterfactuals ask what target variables would have been under an alternative condition given observed facts, while keeping structural equations intact and varying exogenous conditions.The dual graph contains factual and counterfactual copies of the SCM connected through exogenous nodes; inference proceeds by abduction, marginalisation, then prediction.Source: §3.1 · pp. 4–5
• Counterfactual queries are not necessarily identifiable from observational data alone when the underlying causal model is only partially known.Identifiability means that all models agreeing with the available information yield the same query answer.Source: §3.1 · p. 5

3.2 Neural Processes and Graph Neural Networks

• Neural Processes learn conditional predictive distributions from context observations and target inputs, amortizing inference across related tasks.Graph Neural Networks extend message passing to connected nodes; for directed causal graphs, incoming and outgoing neighborhoods can be handled separately.Source: §3.2 · p. 5

4 A GNN Model for Counterfactual Inference without Causal Graph Knowledge

• CFGNP targets the model-averaged counterfactual predictive distribution, accounting for uncertainty over causal models without explicitly inferring or committing to a causal graph.If the distribution is identifiable from observational data, it coincides with the ground-truth model's counterfactual distribution; otherwise, it averages over compatible SCMs.Source: §4 · p. 6
• Training minimizes KL divergence between the target and predicted counterfactual distributions, using counterfactual pairs during training.The training objective is expressed as an expected KL divergence over tasks, and CFGNP's training objective is negative log-likelihood.Source: §4 · p. 6, §4.1 · p. 7

4.1 Approach

• The three-stage architecture summarizes observational context, combines it with the factual observation, then propagates the representation to the counterfactual unit to predict an outcome distribution.The overview depicts this flow from node-wise context representations through factual–counterfactual message passing.Source: §4.1 · p. 6, Figure 3 · p. 6
• For each variable, CFGNP aggregates information from corresponding nodes across observational samples, whose nodes are fully connected because the causal graph is unknown.A K-layer GAT processes the sample subgraphs and aggregation nodes to produce compressed context representations.Source: §4.1 · p. 7
• CFGNP integrates context with the factual unit, then connects corresponding factual and counterfactual variables; a unidirectional GNN can approximate the corresponding bidirectional GNN arbitrarily well.The design removes explicit exogenous nodes, with message-passing depth increasing with graph size but capped at five layers.Source: §4.1 · p. 7
• GNN design makes predictions equivariant to target-node permutations, while applicability to varying graph and context sizes supports generalization beyond training sizes.With infinite data and model capacity, minimizing KL recovers the conditional predictive distribution.Source: §4 · p. 6, §4.1 · p. 7

5 Results

• CFGNP uses counterfactual pairs during training but no causal graph; CheXpert image evaluation checks intended counterfactual attributes with classifiers because ground-truth counterfactuals are unavailable.Synthetic comparisons include methods requiring a full causal graph or only its topological ordering.Source: §5 · pp. 8–9, §5.1 · p. 9

5.1 Synthetic Experiments

• CFGNP achieves MSE 0.19 × 10^-2 versus CFM's 4.51 × 10^-2 on German Loan and improves across almost all synthetic dataset variants.On LIN M-Shape, CFGNP's MSE of 3.60 × 10^-5 is higher than BGM's 1.46 × 10^-5.Source: Table 1 · p. 9, §5.1 · p. 9
• Across graph sizes, both CFGNP variants outperform baselines and degrade more slowly; 10-CFGNP matches retrained CFGNP on Linear and Sum-sine but does worse on Log-sum-exp and Sqrt-sum.10-CFGNP is trained only on 10-node graphs and applied to other sizes without retraining.Source: §5.1 · p. 9, Figure 4 · p. 10

5.2 CheXpert Experiments

• On CheXpert, age MAE is 2.71 for CFGNP versus 10.88 for DeepBC; disease ROC-AUC is 0.77 for CFGNP, 0.75 for DeepBC, and 0.45 for VCI.The CheXpert comparison reports MAE for the continuous feature and accuracy and ROC-AUC for categorical features.Source: §5.2 · pp. 9–10, Table 2 · p. 10
• For unseen CheXpert features, sex and AP/PA achieve ROC-AUC above 0.99, while age MAE reaches 11.53 and disease ROC-AUC ranges from 0.61 to 0.64.Each OOD setting excludes one feature during training and evaluates on the full test set containing all features.Source: Table 3 · p. 10, §5.2 · p. 10

6 Conclusion

• CFGNP directly approximates Bayesian model-averaged counterfactual distributions and performs strongly across image and synthetic queries, including on larger graphs and out-of-distribution features.Its observational context increases training and inference costs, and it cannot recover a true distribution when the counterfactual distribution is non-identifiable; future work could incorporate expert knowledge or distributions over causal edges.Source: §6 · p. 11

Ethics statement

• The work uses existing datasets and collects no new human-subject data; counterfactual predictions should not guide clinical decisions without appropriate validation.The authors warn that predictions may reflect biases and limitations in the data.Source: Ethics statement · p. 11

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

• When the endogenous subgraph is fully connected, a direct MPNN with sufficiently large hidden dimension can approximate any endogenous-node function approximable by a bidirectional MPNN with explicit exogenous nodes.This expressivity result justifies eliminating explicit exogenous nodes from the dual-graph architecture.Source: Equivalence of a Direct MPNN to a… · p. 16

Permutation Properties of CFGNP

• CFGNP is permutation equivariant to relabeling target nodes and permutation invariant to relabeling non-target nodes.The permutation must be applied consistently to node-indexed inputs.Source: Permutation Properties of CFGNP · pp. 18, 21

Graph-Size and Observation-Count Transferability of CFGNP

• A CFGNP can transfer to different graph sizes and observational sample counts without changing learned parameters, subject to the stated message-passing parameter-sharing conditions.The architecture’s parameters do not depend on graph size or the number of observational samples.Source: Graph-Size and Observation-Count… · pp. 21–22
• The same CFGNP parameter vector applies across causal graphs with different node counts and observational datasets with different sample counts.Shared node-wise maps and message-passing functions are evaluated over differently sized computational graphs without changing their learned parameters.Source: Graph-Size and Observation-Count… · pp. 22–24
• Architectural transferability does not guarantee identical predictive performance when graph size or sample count differs substantially from training, because performance depends on the training and test instance distributions.This is a limitation on predictive generalization, not on applying the same parameter vector.Source: Graph-Size and Observation-Count… · p. 24

A Extended Model Architecture

• The appendix defines graph-structure matrices for the model’s computational graph.Source: A Extended Model Architecture · p. 24

A.1 Model Input

• CFGNP takes a factual representation, counterfactual query, target variables, and observational context as input.Source: §A.1 · p. 24

A.2 Changing data dimension

• Separate MLPs encode intervened and non-intervened nodes in both factual and counterfactual inputs.Source: §A.2 · p. 24

A.3 Propagation of Observational Information

• Observational representations are reshaped and propagated using a K-layer GAT with attention over neighboring nodes.Source: §A.3 · pp. 24–25

A.4 Mixing Factual Observation

• Factual observations are mixed by constructing a new graph and applying a three-layer GAT.Source: §A.4 · p. 25

A.5 Estimating Counterfactual Distribution

• The model forms per-node Gaussian-mixture parameters after propagating information through its counterfactual graph.The output contains mixture means, standard deviations, and weights for each node.Source: §A.5 · p. 25

B.1 Model Hyperparameters

• Synthetic and semi-synthetic runs use hidden dimension 80, five GAT heads, and batch size 128; image runs use 256, nine heads, and batch size 32.Models train for 200 epochs with early stopping after 25; AdamW uses learning rate 10^-4 and weight decay 10^-5.Source: §B.1 · p. 25

B.2 CheXpert Experiment Setup

• CheXpert counterfactuals are evaluated with ResNet-plus-MLP attribute classifiers and regressors, while image generation adds reconstruction losses to attribute objectives.The reconstruction loss combines MSE, LPIPS, and total variation.Source: §B.2 · pp. 26–27

B.3 Synthetic Experiments Setup

• Synthetic experiments use five random seeds, Nobs = 100, and report mean and standard deviation across runs.Source: §B.3 · p. 27

B.3.1 Triangle

• The Triangle SCM has three variables, with X1 affecting X3 directly and through the mediated path X1 → X2 → X3; its NLIN variant is nonlinear additive.Source: §B.3.1 · pp. 27–28

B.3.2 M-graph

• The M-graph has five endogenous variables, and its exogenous variables are independent standard normals.Source: §B.3.2 · p. 28

B.3.3 German Loan

• The semi-synthetic German Loan SCM has seven endogenous variables and is inspired by the German Credit domain.Its variables include gender, age, education, loan amount and duration, income, and savings.Source: §B.3.3 · p. 29

B.4 Structural Equations for Artificial Graphs

• Artificial-graph variables are generated as X = f(Pa(X)) + ϵ, with f sampled from a set of transformations.Source: §B.4 · p. 30

B.5 Identifiability in a Two-Node SCM

• The paper’s backtracking counterfactual distribution coincides with the corresponding interventional abduction–action–prediction construction.Source: §B.5 · p. 31
• Observational and interventional distributions can leave λ unidentified, so even infinite observational data may leave backtracking-counterfactual uncertainty.The likelihood is independent of λ, which can change the backtracking counterfactual while remaining invisible to observational and interventional distributions.Source: §B.5 · pp. 30, 31, 33

• An example changes the factual age from 80 to 58, with the pixel difference displayed alongside the images.Source: Figure 9 · p. 33
• Another example changes the factual AP/PA value from 1 to 0 and displays the pixel difference.Source: Figure 10 · p. 33
• A further example changes the factual age from 74 to 85.Source: Figure 11 · p. 34

• With sex as the OOD feature, examples intervene on age from 65 to 72 and on AP/PA from 0 to 1.Source: Figure 12 · p. 34, Figure 13 · p. 34
• With age as the OOD feature, examples intervene on age from 30 to 52 and sex from 1 to 0.Source: Figure 14 · p. 35, Figure 15 · p. 35
• With AP/PA as the OOD feature, examples intervene on sex from 0 to 1 and age from 26 to 64.Source: Figure 16 · p. 35, Figure 17 · p. 35
• With disease as the OOD feature, examples intervene on AP/PA from 0 to 1 and age from 67 to 37.Source: Figure 18 · p. 36, Figure 19 · p. 36

C.2 Graph Size Results

• Figures 20–21 plot MSE against causal-graph size for four structural equations, averaging over distinct graphs and showing the 25th and 75th quantiles.Source: Figure 20 · p. 37, Figure 21 · p. 37

C.3 Runtime of Experiments

• Runtime tables compare CFGNP with BGM, CFM, and VACA on synthetic datasets, and with VCI and DeepBC on CheXpert.Experiments use batch sizes of 128 for (semi-)synthetic data and 32 for CheXpert, on an NVIDIA H100 GPU.Source: Table 4 · p. 38, Table 5 · p. 38, §C.3 · p. 38

C.4 Convergence in (Non-)Identifiable Cases

• The convergence experiment compares CFGNP’s predicted distribution with the ground-truth counterfactual and an observation-conditioned alternative in identifiable and non-identifiable two-node SCMs.Figure 22 reports KL divergences across numbers of supplied observations, with 10th and 90th quantiles.Source: §C.4 · p. 38, Figure 22 · p. 39

C.5 Transformer Ablation

• The ablation separately embeds observations, factual units, and counterfactual/intervention inputs, concatenates the vectors, and processes them with a Transformer encoder.Intervened and non-intervened variables retain separate embeddings; the output for x* is handled as in CFGNP.Source: §C.5 · p. 39
• The ablation keeps CFGNP’s listed hyperparameters identical and uses a five-layer Transformer with four-times embedding feed-forward dimension, 0.1 dropout, GELU, and layer normalization.Pre-layer normalization is used in each encoder layer, followed by final layer normalization.Source: §C.5 · p. 39
• Table 6 compares full CFGNP with the Transformer ablation across the Triangle, M-shape, and German Loan datasets; lower MSE indicates better performance.The supplied passage does not report the table’s numeric values or identify a winning model.Source: Table 6 · p. 39, §C.5 · p. 39