Technical Summary: Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Problem Formulation

The paper addresses the challenge of predicting potential outcomes under future treatment sequences in longitudinal observational data. The core task is to estimate the history-conditional potential outcome $E[Y(\bar{a})_{t+\tau} | H_t]$, where $H_t$ represents the observed history up to time $t$, and $\bar{a}$ is a planned future treatment sequence.

This problem is complicated by three primary factors:

1. Time-varying confounding: Treatment assignments at each step depend on covariates that are themselves outcomes of prior treatments.
2. Heterogeneous patient dynamics: Individual trajectories exhibit complex, nonlinear state evolution and latent heterogeneity.
3. Data limitations: Observational cohorts are often too small to train reliable deep sequence models from scratch for every new domain or simulator.

Existing longitudinal causal estimators (e.g., RMSN, CRN, G-Net, Causal Transformer) typically require a separate supervised training run, including hyperparameter selection and propensity modeling, for each new cohort. This pipeline is costly and impractical when repeated domain-specific training is required.

Methodology: CAUSALLONGPFN

The authors introduce Causal Longitudinal Prior-Fitted Networks (CAUSALLONGPFN), a prior-fitted in-context predictor designed to amortize longitudinal causal prediction across a broad prior over temporal structural causal models (TSCMs).

1. Synthetic Pretraining on a Broad Prior

The model is pretrained entirely on synthetic episodes sampled from a broad prior over TSCMs. This prior is designed to span a wide class of longitudinal causal dynamics rather than replicating a single simulator. Key features of the synthetic prior include:

• Causal Temporal Graphs: Sparse contemporaneous and lagged dependencies with acyclic instantaneous graphs.
• Nonlinear Mechanisms: State coordinates follow sparse nonlinear autoregressive updates using diverse elementary nonlinearities (identity, tanh, sinusoidal, ReLU, etc.) and various noise distributions.
• Dynamical Motifs: Structured motifs such as action-memory, saturating, homeostatic, feedback-control, and smoothed-readout channels are overlaid to capture qualitative mechanisms like delayed effects and regulatory feedback.
• Confounded Behavior Policies: Treatments are sampled from state-dependent stochastic policies influenced by latent unit heterogeneity ($Z_i$), creating time-varying treatment–confounder feedback.
• Outcome Models: Scalar outcomes are generated via autoregressive readouts with direct and cumulative treatment effects.

2. Architecture

CAUSALLONGPFN employs a dual-encoder architecture:

• Causal History Encoder: A trajectory-level causal Transformer (using masked self-attention) that maps longitudinal sequences to history representations, ensuring the representation at time $r$ depends only on information available up to that time.
• PFN Context Encoder: A Transformer-based context encoder that performs in-context adaptation. It processes support trajectories (treated as an unordered set) and a query token jointly via self-attention. No positional encoding is assigned to the ordering of support trajectories.
• Gaussian-Mixture Prediction Head: The final query representation parameterizes a 5-component Gaussian mixture distribution for the normalized next outcome, providing both point predictions and uncertainty estimates.

3. In-Context Prediction and Rollout

At test time, the model is frozen. It receives:

• Support trajectories: Examples from the new domain/task.
• Query history: Observed up to time $t_{obs}$.
• Proposed future treatment sequence: The planned interventions.

The model returns a predictive distribution without gradient updates, propensity-model fitting, or adversarial balancing. For multi-step prediction ($\tau > 1$), the model performs an autoregressive plug-in rollout: it predicts the one-step outcome distribution, inserts the mixture mean into the query history, and repeats the process under the specified treatment sequence.

Key Contributions

1. A Prior-Fitted Model for Longitudinal Causal Prediction: CAUSALLONGPFN is the first PFN-style model for history-conditional potential-outcome prediction under planned longitudinal treatment sequences. It operates as a frozen model requiring no test-time adaptation.
2. A Synthetic Prior over Longitudinal Causal Tasks: The authors design a TSCM prior generating diverse tasks with treatment-confounder feedback, latent heterogeneity, nonlinear dynamics, delayed/cumulative effects, and regime changes.
3. Architecture for Longitudinal In-Context Inference: A novel dual-encoder combining a causal Transformer history encoder with a PFN context encoder and a Gaussian-mixture head.
4. Autoregressive Counterfactual Rollout: An extension of the learned one-step predictor to multi-step prediction via deterministic plug-in rollout.
5. Zero-Shot Evaluation: A comprehensive evaluation of a single frozen model against domain-trained baselines (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) on branchable counterfactual benchmarks and factual real-world data.

Results

The model was evaluated on four benchmarks: Cancer tumor growth, Warfarin PK/PD, HIV treatment dynamics, and MIMIC-III ICU trajectories.

• Domain-Balanced Performance: CAUSALLONGPFN achieved the best domain-balanced one-step normalized RMSE (0.222), narrowly outperforming MSM and RMSN. For five-step prediction, it ranked third overall, behind RMSN and G-Net, but outperformed MSM, CRN, and transformer-based baselines.
• Per-Domain Performance:
  • MIMIC-III (Factual): The model ranked first in both one-step and five-step prediction, demonstrating strong transfer to real-world clinical trajectories without domain-specific training.
  • Counterfactual Benchmarks (Cancer, HIV, Warfarin): The model remained competitive, ranking second or third on one-step tasks. However, on longer-horizon counterfactual tasks (e.g., 5-step cancer prediction), specialized domain-trained recurrent models (RMSN, CRN) achieved lower error, suggesting an advantage when sufficient target-domain data exists for specific fitting.
• Uncertainty Calibration: The Gaussian-mixture head provided useful distributional information. Calibration varied by domain, with Warfarin showing the best calibration and MIMIC-III showing wider intervals due to higher heterogeneity.

Significance and Claims

The paper claims that broad synthetic causal pretraining can provide a useful frozen alternative to repeated domain-specific training when:

• Retraining specialized models is costly or impractical.
• Rapid adaptation to a new cohort is required.
• Counterfactual supervision is unavailable (as in real-world factual prediction tasks like MIMIC-III).

The authors emphasize that CAUSALLONGPFN does not remove the standard causal assumptions (consistency, positivity, sequential exchangeability) required for interpreting observational data. Instead, it amortizes the estimation problem. The results suggest that a sufficiently broad synthetic prior can capture reusable structure across treatment-response tasks, making the model a strong general-purpose in-context predictor. However, the paper modestly notes that domain-specific training remains valuable when sufficient target-domain data and validation signals are available, particularly for long-horizon counterfactual predictions in specific domains.

The work positions itself as a research tool for causal sequence modeling and hypothesis generation rather than a standalone clinical decision system, cautioning against over-trust in predictions when causal assumptions or prior support are inadequate.

Code & Model Availability

The implementation is available on GitHub at https://github.com/Amirhossein-Zare/causal-long-pfn and pretrained model weights are available on Hugging Face at https://huggingface.co/Amirhossein-Zare/causal-long-pfn .