PaperlayerView exact revision on arXiv

Source-linked AI summary

A Foundation Model for Policy In-Context Longitudinal Causal Effect Estimation

Wenxin Chen, Kyra Gan, Mingyu Lu, Zilin Jing, Fei Wang, Hoifung Poon, Tristan Naumann, Peniel Argaw

arXiv:2610.32863v1cs.AI

TL;DR

Time-varying confounding and treatment-confounder feedback complicate longitudinal causal estimation, while existing approaches often require data-specific fitting or handle limited policy and horizon queries. PICLE-PFN uses policy demonstrations in context and synthetic longitudinal SCM pretraining to estimate CAPOs across datasets, policies, and horizons. It matches or outperforms baselines on synthetic and semi-synthetic benchmarks with faster deployment and improved uncertainty calibration, within the scope of its stated assumptions and modeling setup.

  • Problem

    Time-varying confounding and treatment-confounder feedback complicate observational estimation, while many methods require data-specific nuisance fitting or are tailored to fixed treatment sequences and horizons.

    Source: §1 · pp. 1–2

  • Method

    PICLE-PFN encodes policies as history-decision demonstrations and pretrained synthetic longitudinal SCM tasks to predict CAPOs in context across datasets, policies, and horizons.

    Source: §1 · p. 2, §3.1 · p. 4, §4.1 · p. 5

  • Results

    PICLE-PFN matches or outperforms baselines on synthetic and semi-synthetic benchmarks, with the lowest horizon-averaged RMSE in six of seven settings and substantially closer nominal coverage than IGC-Net.

    Source: §5.2 · pp. 7–8, §6 · p. 9

  • Takeaways & Limitations

    The results support a single in-context estimator for static and dynamic policies across multi-step windows, with competitive accuracy and substantially better calibrated uncertainty than the compared method.

    Source: §1 · p. 2, §5.2 · p. 8

    The model is limited to binary treatments and discrete-time trajectories without right censoring, and its identification relies on assumptions such as sequential ignorability.

    Source: §6 · p. 10

Abstract

from arXiv · show

Longitudinal causal effect estimation is complicated by time-varying confounding and treatment-confounder feedback loops. Existing methods generally require data-specific nuisance-model fitting, and many are tailored to prespecified treatment sequences or prediction horizons, leading to time-consuming computation at deployment. We introduce *Policy In-Context Longitudinal causal Effect estimation with Prior-data Fitted Networks* (**PICLE-PFN**), the first foundation model for longitudinal causal effect estimation that supports *dynamic* treatment policies. PICLE-PFN represents target policies through in-context history-decision demonstrations, providing a unified interface for a broad class of treatment policies. Pretrained over synthetic data sampled from longitudinal structural causal models, PICLE-PFN amortizes causal effect estimation across datasets, policies, and horizons. Given a new observational dataset and target policy, it directly predicts conditional average potential outcomes (CAPOs) in one forward pass without task-specific training or hyperparameter tuning. Across synthetic and semi-synthetic longitudinal benchmarks, PICLE-PFN achieves competitive RMSE and substantially improved uncertainty calibration with a fraction of deployment runtime, compared to state-of-the-art conventional estimators trained and tuned separately for each dataset.

1 INTRODUCTION

Longitudinal treatment effects are difficult to estimate from observational data because time-varying confounding creates treatment–confounder feedback. PICLE-PFN uses in-context policy demonstrations and synthetic pretraining to estimate effects across policies and horizons, with competitive accuracy, improved calibration, and faster deployment.

1 INTRODUCTION

  • Time-varying covariates can be consequences of earlier treatments and determinants of later treatments and outcomes, complicating causal estimation.Existing methods also face accumulated estimation error and a trade-off between query flexibility and variance or repeated fitting.Source: §1 · p. 1
  • Many existing methods target fixed treatment sequences, while clinically relevant regimes may adapt decisions to changing patient characteristics.The paper identifies a need for one estimator that handles static and dynamic policies across horizons without retraining for each target.Source: §1 · p. 2
  • PICLE-PFN represents policies with in-context demonstrations and directly learns CAPO from synthetic supervision, avoiding local nuisance-model composition at deployment.Its demonstrations pair decisions with factual histories, enabling a unified interface for static and dynamic policies.Source: §1 · p. 2
  • PICLE-PFN has the lowest RMSE in six of seven benchmark settings, up to 12.2% lower than the strongest baseline, and deploys about 100× to 5000× faster than a comparably accurate method.It also provides substantially better calibrated uncertainty.Source: §1 · p. 2

2 RELATED WORK

Prior PFN work addresses time-series forecasting and causal inference, but the cited causal methods focus on static, single-assignment settings. PICLE-PFN instead targets longitudinal sequential treatments and time-varying confounders, addressing limited flexibility across policies and horizons.

2 RELATED WORK

  • PFN time-series approaches either treat time as additional features or explicitly model temporal structure; PICLE-PFN follows the latter approach.The cited PFN causal-inference line focuses on static, single-treatment-assignment settings, unlike this paper’s longitudinal sequential assignments and time-varying confounders.Source: §2 · p. 2
  • Longitudinal estimators model treatment, outcome, or both; treatment-based methods include weighting and balanced representations, while outcome-based methods use G-computation.Doubly robust methods combine treatment and outcome models to reduce sensitivity to misspecification.Source: §2 · pp. 2–3
  • A concurrent approach uses a one-step mean rollout that does not generally recover CAPO and supports only static treatment sequences.Source: §2 · p. 2
  • Existing methods have limited flexibility across causal queries: many use prespecified treatment sequences or require retraining or retuning when policies or horizons change.The paper contrasts this with PICLE-PFN’s single pretrained model supporting these capabilities.Source: §2 · p. 3

3 PROBLEM SETUP AND IN-CONTEXT FORMULATION

The paper formulates longitudinal CAPO estimation as an in-context task: the dataset and policy demonstrations inform estimates for a new query history. Under standard identification assumptions, demonstrations on factual histories retain the policy information needed for CAPO, while synthetic counterfactual targets provide multi-step supervision.

3 PROBLEM SETUP AND IN-CONTEXT FORMULATION

  • The setup considers i.i.d. trajectories over discrete time, with baseline and time-varying covariates, binary treatments, and outcomes after each treatment.Source: §3 · p. 3
  • A target policy may be a static treatment sequence or a history-dependent deterministic rule, and the estimand is the conditional average potential outcome under that policy and horizon.The history before treatment includes baseline and time-varying covariates, prior treatments, and prior outcomes.Source: §3 · p. 3
  • Consistency, sequential ignorability, and positivity are the stated assumptions under which the target policy’s CAPO is identified.Source: §3 · p. 3

3.1 IN-CONTEXT FORMULATION OF LONGITUDINAL CAUSAL EFFECT ESTIMATION

  • A target policy demonstration records its decisions on factual histories rather than encoding the policy’s underlying rules or parameters.Under the stated assumptions, the population demonstrations suffice to identify CAPO.Source: §3.1 · p. 4
  • In practice, the empirical dataset approximates population demonstrations by applying the target policy to observed histories.Source: §3.1 · p. 4
  • Policy information lies in the relationship between histories and decisions across demonstrations, so estimation must aggregate across them rather than rely on one pair.The observed trajectories supply additional context about the data distribution, with corresponding target actions for each patient’s history.Source: §3.1 · p. 4
  • The amortized estimator uses the combined context and a query history to estimate CAPO for new datasets and policies without dataset- or policy-specific parameter updates.Source: §3.1 · p. 4
  • Observed multi-step outcomes do not generally provide valid labels for a target policy because later treatments follow the behavior policy; synthetic SCM rollouts supply CAPO supervision instead.The pretraining tasks vary SCM dimensions, graph structure, transition mechanisms, and target policies, including static and dynamic rules.Source: §3.1 · p. 4

4 METHOD

PICLE-PFN is pretrained on longitudinal structural causal models and target policies, then uses observational trajectories and policy demonstrations to estimate CAPOs jointly across intervention windows. Its training uses counterfactual rollouts as outcome supervision, while a shared transformer backbone and window-specific readouts produce predictions in one forward pass.

4 METHOD

  • PICLE-PFN is pretrained on tasks sampled from priors over longitudinal SCMs and target policies, integrating observational data and policy demonstrations to predict across intervention windows.The training procedure samples an SCM and target policy, generates factual trajectories and policy demonstrations, branches counterfactual rollouts, and minimizes negative log-likelihood.Source: §4 · pp. 4–5

4.1 PRIOR SAMPLING

  • Synthetic tasks vary longitudinal SCM dimensions, graph structures, transition mechanisms, and behavior policies, and include both static and dynamic target policies.Factual trajectories and target-policy counterfactual rollouts provide supervision whose conditional mean is the CAPO.Source: §4.1 · p. 5

4.2 TRAINING OBJECTIVE

  • Single-draw counterfactual outcomes provide computationally efficient supervision because their conditional expectation given the pre-intervention history equals the CAPO.The model minimizes outcome negative log-likelihood over sampled tasks, query trajectories, and admissible intervention windows; its predictive mean estimates the CAPO at inference.Source: §4.2 · p. 5

4.3 MODEL ARCHITECTURE

  • The architecture uses shared observed-data representations and window-specific readouts, with query readouts restricted to pre-intervention history and context readouts accessing factual data through the endpoint.Context readouts also incorporate target-policy demonstrations, and readout tokens attend across trajectories using only context readouts as keys and values.Source: §4.3 · p. 6, Figure 1 · p. 6

5 EXPERIMENTS

Experiments evaluate longitudinal CAPO estimation across synthetic and semi-synthetic settings, plus a real-world prediction sanity check. PICLE-PFN leads most benchmark settings, runs faster than strong baselines, and shows stronger uncertainty calibration.

5.1 EXPERIMENTAL SETUP

  • The evaluation spans seven settings, 20 datasets per setting, 1,000 context and 1,000 query trajectories, and five horizons; baselines are tuned with Optuna for two hours.Settings vary treatment and covariate dimensions and include static and dynamic policies.Source: §5.1 · pp. 6–7

5.2 RESULTS

  • PICLE-PFN achieves the lowest horizon-averaged RMSE in six of seven settings, reducing RMSE by 12.2% on cancer and up to 10.8% on the remaining MIMIC benchmarks.On mimic y1, it ranks third, with RMSE 12.6% higher than the best baseline.Source: §5.2 · p. 7, Table 3 · p. 7
  • PICLE-PFN is about 100× faster than CT and 5000× faster than IGC-Net, while combining top-ranking accuracy across diverse settings.Figure 3 aggregates RMSE ranking and deployment runtime across horizons and datasets.Source: §5.2 · p. 7, Figure 3 · p. 8
  • On real-world MIMIC data, PICLE-PFN achieves slightly lower factual diastolic blood-pressure prediction error than IGC-Net across all horizons.This is presented as a sanity check, complementary to benchmarks with counterfactual ground truth.Source: §5.2 · p. 7, Table 4 · p. 7
  • PICLE-PFN tracks nominal coverage more closely than IGC-Net, maintaining coverage near 0.9 at the nominal 95% level as intervals widen with horizon; these are not credible intervals.IGC-Net intervals do not widen with longer horizons and even narrow in one setting.Source: §5.2 · p. 8, Figure 4 · p. 9
  • Across tested settings, PICLE-PFN is stable under stronger unobserved confounding, more robust to smaller contexts, and comparable to strong baselines on two unseen static policies.The unseen policies alternate treatment or treat only at the first and last steps.Source: §5.2 · pp. 8–9, Figure 5 · p. 9

6 CONCLUSION, LIMITATIONS, AND FUTURE WORK

PICLE-PFN reformulates longitudinal CAPO estimation as in-context learning and matches or outperforms baselines on synthetic and semi-synthetic benchmarks at a fraction of their deployment runtime. Further analyses show strong uncertainty calibration, resilience to small samples and unobserved confounders, and generalization to policies unseen during pretraining; the authors also identify limits in scalability, pretraining scope, and causal guarantees.

6 CONCLUSION, LIMITATIONS, AND FUTURE WORK

  • PICLE-PFN matches or outperforms baselines on synthetic and semi-synthetic benchmarks at a fraction of their deployment runtime.The model represents target policies through decisions on factual histories across the observed dataset.Source: §6 · p. 9
  • Further analyses demonstrate strong uncertainty calibration, greater resilience to small samples and unobserved confounders, and generalization to policies unseen during pretraining.Source: §6 · p. 10
  • Quadratic costs in sample size N and trajectory length T limit scalability, while current pretraining covers binary treatments, discrete-time trajectories, and no right censoring.Identification requires assumptions such as sequential ignorability; empirical robustness to certain violations does not establish theoretical guarantees under unobserved confounding, and accuracy may deteriorate when deployment differs from the pretraining distribution.Source: §6 · p. 10

A PROOF

Under consistency, sequential ignorability, and positivity, population policy demonstrations on factual histories identify deterministic policies almost everywhere under the observed history distribution. The proof then uses backward induction on the ICE recursion to establish that the target functional is uniquely determined.

A PROOF

  • Under consistency, sequential ignorability, and positivity, the target functional and observed distribution together with population policy demonstrations uniquely determine the target value almost everywhere under P0.Source: A PROOF · pp. 14–15
  • The population demonstration distribution recovers each deterministic policy for P0-almost every history, and the ICE representation connects those policies to the longitudinal g-formula.Source: A PROOF · p. 14
  • Backward induction establishes agreement of the ICE recursions, using terminal conditional expectations and positivity to align policy-weighted sums across observed histories.The proof also shows that disagreement on an observed-history null set cannot affect the recursion.Source: A PROOF · pp. 14–15

B.1 MODEL ARCHITECTURE

PICLE-PFN separates factual-trajectory representations from policy-conditioned readouts, using window-specific access to predict potential outcomes. Its synthetic-task prior spans longitudinal mechanisms and static or history-dependent target policies, with counterfactual outcomes providing training labels.

B.1 MODEL ARCHITECTURE

  • The shared eight-layer transformer uses an observed-data backbone that remains independent of the target policy and intervention window, while readout tokens combine factual representations with policy demonstrations to predict potential outcomes.Source: §B.1 · p. 15
  • Each trajectory is tokenized as pooled embeddings of interleaved covariate and treatment vectors, with learned feature, position, and type embeddings.Baseline covariates are encoded similarly and their pooled embedding is added to each time-varying covariate embedding.Source: §B.1 · p. 15
  • Window-specific readouts access query histories or context trajectories through the intervention window, attend across context readouts, and map to outcome distributions; supplying demonstrated-minus-factual treatment differences accelerates convergence without changing available information.Readouts for different windows do not interact and share parameters.Source: §B.1 · p. 16, Figure S.1 · p. 16

B.2 PRIOR SAMPLING

  • Pretraining tasks use longitudinal SCMs with varied dimensions, variable types, acyclic temporal dependencies, and behavioral mechanisms drawn from seven function families.Covariate and outcome mechanisms incorporate persistence and smooth time trends, while treatment effects can vary with covariates.Source: §B.2 · pp. 16–17
  • The target-policy prior mixes static treatment sequences and history-dependent policies, then rolls factual prefixes forward under the target policy with shared exogenous noise to generate counterfactual training labels.For a fixed SCM and policy, the conditional expectation of these outcomes given Hs is the CAPO.Source: §B.2 · p. 17

B.3 HYPERPARAMETERS AND TRAINING

The model has 5.4M parameters and selects a checkpoint using held-out tasks drawn from the prior. At inference, it processes each task in one forward pass without per-task fitting; the benchmark table reports normalized RMSE by horizon.

B.3 HYPERPARAMETERS AND TRAINING

  • At inference, each task is processed in a single forward pass with no per-task fitting.Source: §B.3 · p. 17
  • The model has 5.4M parameters, and checkpoint selection uses the lowest validation loss across a fixed set of 128 held-out tasks drawn from the prior.Training uses AdamW with a learning rate warmed up to 0.001 and halved when validation loss does not improve for 40 epochs.Source: §B.3 · p. 17
  • The benchmark table reports normalized RMSE and per-task RMSE ranking by horizon, with mean and standard deviation across 20 repeats and best and second-best methods highlighted.Source: Table S.2 · p. 18

C ADDITIONAL RESULTS

This section describes the benchmark settings and baseline implementations, with full CAPO RMSE and deployment-runtime results reported in Tables S.2 and S.3.

C ADDITIONAL RESULTS

  • Tables S.2 and S.3 provide full CAPO RMSE results across methods and horizons, and deployment runtime results across horizons.The supplied passages identify these tables but do not give their cell values.Source: C ADDITIONAL RESULTS · p. 18

D.1 BENCHMARK DATASETS

  • Benchmarks include one synthetic and six semi-synthetic MIMIC settings, with 20 datasets of 1,000 training and 1,000 test trajectories per setting.Trajectories have length 15; CAPO is estimated at horizons 2, 4, 6, 8, and 10, with start time fixed at 2.Source: §D.1 · p. 19

D.1.1 CANCER

  • The cancer benchmark has tumor volume as its time-varying covariate, two binary treatments, and a static policy that always applies both treatments.Treatment assignment depends on recent tumor progression, introducing time-varying confounding.Source: §D.1.1 · p. 19

D.1.2 MIMIC Y1, MIMIC Y2

  • The MIMIC Y1 and Y2 benchmarks use distinct simulated outcomes and confounders; each has 27 time-varying confounders and three treatments.Both target the policy that applies the second and third treatments for the intervention window.Source: §D.1.2 · p. 20

D.2 BASELINE IMPLEMENTATIONS

  • The nine baselines have differing coverage: several encoder–decoder models require fixed treatment sequences, while some DR-based methods are limited to single binary treatments.IGC-Net is extended beyond its published scope to use per-trajectory policy decisions; G-Net is omitted for dynamic policies because its released inference path does not reevaluate policies on simulated states.Source: §D.2 · pp. 20–21

D.3 HYPERPARAMETER TUNING AND EVALUATION PROTOCOL

The protocol tunes and evaluates baselines on held-out tasks and shared data, then compares deployment cost and uncertainty calibration. Calibration deviations are greatest on tasks with zero empirical policy overlap.

D.3 HYPERPARAMETER TUNING AND EVALUATION PROTOCOL

  • Baselines are tuned with two hours of Optuna search per setting or setting–horizon cell, selected by held-out factual RMSE without counterfactual labels.Selected hyperparameters are fixed before fresh models are trained on each evaluation task.Source: §D.3 · p. 21, Table S.4 · p. 22

D.4 RUNTIME MEASUREMENT

  • Deployment timing excludes hyperparameter tuning: PICLE-PFN is measured by inference alone, while baselines include training and inference.Horizon-independent baselines are timed across one training run and five inferences; horizon-specific methods are timed separately per horizon.Source: §D.4 · p. 21

D.5 REAL-WORLD MIMIC PREDICTION

  • The real-world prediction task uses ten cohorts of 3,000 patients and predicts hourly diastolic blood pressure at horizons 2–5 from hour 5.IGC-Net is the only baseline supporting both multivariate treatments and dynamic policies in this comparison; PICLE-PFN uses factual treatment sequences as policy demonstrations.Source: §D.5 · pp. 22–23

D.7 UNOBSERVED CONFOUNDING

  • The unobserved-confounding experiment adds latent severity to outcomes and behavior propensities, fixing its outcome effect at 0.5 while varying propensity confounding over 1, 2, and 4.The supplied table caption identifies a separate context-sample-size comparison under static and history-dependent policies, but provides no result values.Source: §D.7 · pp. 23–24, Table S.5 · p. 24

D.6 UNCERTAINTY CALIBRATION.

  • On tasks with positive empirical overlap, nominal 95% intervals generally achieve coverage above 0.9; zero-overlap tasks show the largest deviations from nominal coverage.Coverage is evaluated for PICLE-PFN and IGC-Net on CAPO intervals.Source: §D.6 · p. 23, Figure S.2 · p. 24

E COMPARISON WITH CAUSALLONGPFN

The comparison identifies a limitation of CausalLongPFN’s mean plug-in rollout and contrasts its cancer-benchmark results with PICLE-PFN. The appendix also gives PICLE-PFN’s attention-cost complexity.

E COMPARISON WITH CAUSALLONGPFN

  • On cancer, CausalLongPFN has lower RMSE at horizons 1 and 2, while PICLE-PFN has lower RMSE at longer horizons.This comparison is limited to cancer and maximum horizon 5 because the released checkpoint supports at most 10 time-varying covariates; the benchmark has 27.Source: Table S.6 · p. 25, E COMPARISON WITH CAUSALLONGPFN · p. 25
  • Even with exact one-step conditional means, CausalLongPFN’s mean plug-in rollout does not generally recover multi-step CAPO because expectation and nonlinear prediction need not commute.The rollout inserts the first-step predicted mean into the next-step predictor rather than averaging predictions over the intermediate outcome.Source: E COMPARISON WITH CAUSALLONGPFN · pp. 24–25

F COMPLEXITY

  • PICLE-PFN’s stated attention cost is quadratic in trajectory count and linear in sequence length for a fixed horizon cap.The full-triangle case is at most O(NTd^2 + T^2N^2); row-attention cost is linear in query count when context size is fixed.Source: F COMPLEXITY · p. 25

Related papers

Loading 2610.32863v1…