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

cs.AIOct 23, 2026

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

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

Organizations: Cornell University, New York City, NY · University of Washington, Seattle, WA · Columbia University, New York City, NY · Recursion Pharmaceuticals, Inc., Salt Lake City, UT ·

Abstract

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.

Figures & tables

Capability	RMSN	CRN	TE-CDE	CT	G-Net	IGC-Net	DeepACE	DeepLTMLE	D 3 -Net	PICLE-PFN
Dynamic treatment policy	✗	✗	✗	✗	✗	✓	✗	✓	✓	✓
Multidimensional treatment	✓	✓	✓	✓	✓	✓	✗	✗	✗	✓
No retraining across horizons/policies	✓	✓	✓	✓	✓	✗	✗	✗	✗	✓

Table 1: Flexibility of existing methods: RMSN ( Lim, 2018 ) ; CRN ( Bica et al., 2020 ) ; TE-CDE ( Seedat et al., 2022 ) ; CT ( Melnychuk et al., 2022 ) ; G-Net ( Li et al., 2021 ) ; IGC-Net ( Hess et al., 2026 ) ; DeepACE ( Frauen et al., 2023 ) ; DeepLTMLE ( Shirakawa et al., 2024 ) ; D 3 -Net ( Chen et al., 2026a ) . ✓ indicates that the capability is supported by the method’s released implementation.

Figure 1: PICLE-PFN Architecture . Feature pooling encodes longitudinal covariates and treatments into embeddings processed by a shared backbone with causal temporal attention and row-attention. Window-specific readouts combine window indices, initial policy demonstrations, and factual outcomes (masked for queries). They cross-attend to trajectory-specific backbone representations under window-specific masks: query readouts access only pre-intervention history, while context readouts access factual data through $e$ and additionally attend to target policy demonstrations. Row-attention uses only context keys and values and operates separately within each readout window. Query readouts yield CAPO predictions across admissible windows in one forward pass.

Dataset	$(d_{A},d_{L},d_{B})$	Policy type
cancer	$(2,1,1)$	Static
mimic_y1	$(3,27,3)$	Static
mimic_y2	$(3,27,3)$	Static
mimic_1t	$(1,27,3)$	Static
mimic_1t_hdp0.5	$(1,27,3)$	Dynamic
mimic_1t_hdp1	$(1,27,3)$	Dynamic

Table 2: Dataset characteristics.

Method	Cancer	MIMIC
Method	Cancer	Y1	Y2	1T	1T-HDP0.5	1T-HDP1	1T-HDP2
PICLE-PFN	0.190 $\pm$ 0.03	0.958 $\pm$ 0.50	0.603 $\pm$ 0.23	0.583 $\pm$ 0.20	0.527 $\pm$ 0.13	0.534 $\pm$ 0.13	0.555 $\pm$ 0.15
IGC-Net	0.297 $\pm$ 0.02	0.901 $\pm$ 0.49	0.604 $\pm$ 0.18	0.699 $\pm$ 0.11	0.546 $\pm$ 0.13	0.580 $\pm$ 0.13	0.622 $\pm$ 0.17
CT	0.217 $\pm$ 0.06	0.851 $\pm$ 0.29	0.886 $\pm$ 0.82	0.660 $\pm$ 0.40	—	—	—
RMSN	0.357 $\pm$ 0.04	1.242 $\pm$ 0.70	0.830 $\pm$ 0.13	0.786 $\pm$ 0.21	—	—	—
G-Net	0.661 $\pm$ 0.04	1.333 $\pm$ 0.86	0.984 $\pm$ 0.14	1.127 $\pm$ 0.20	—	—	—

Table 3: Longitudinal CAPO benchmark across datasets. Each cell reports the mean

\pm

SD normalized RMSE over 20 runs, with RMSE averaged across horizons within each repeat. Bold and underlined values denote the best and second-best methods in each column, respectively.

Method	$\tau=2$	$\tau=3$	$\tau=4$	$\tau=5$
PICLE-PFN	0.730 $\pm$ 0.030	0.737 $\pm$ 0.032	0.754 $\pm$ 0.026	0.767 $\pm$ 0.037
IGC-Net	0.736 $\pm$ 0.040	0.740 $\pm$ 0.038	0.762 $\pm$ 0.031	0.779 $\pm$ 0.028
PICLE-PFN wins	7/10	5/10	8/10	8/10

Table 4: Factual Outcome Prediction on MIMIC. Normalized RMSE of the factual diastolic blood pressure on the held-out test patients of 10 sampled cohorts.

Figure 2: RMSE across horizons of the top-3 methods in each setting. Points are the normalized RMSE averaged over the 20 tasks, error bars are $\pm 1$ SD across tasks

Figure 3: RMSE rank (y-axis) versus inference runtime (x-axis; log scale) across datasets.

Figure 4: Uncertainty Calibration Analysis. Left: empirical coverage against its nominal level, averaged across horizons and 20 tasks; dashed diagonal = perfect calibration. Middle: coverage of the 95% interval per horizon (dashed line = nominal). Right: 95% interval width per horizon.

Figure 5: Ablation studies. (a) Unobserved confounding. (b) Context sample size. (c) Our-of-prior target policy: Static alternating treatments. Compared with the in-prior always-treat policy.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure S.1: Ablation study on the target policy demonstration channel. Supplying the differences between demonstrated and factual treatments accelerates convergence.

Component	Value
Transformer layers	8
Embedding width	128
Attention heads	4
MLP hidden width	512
Activation	GELU
Dropout	0

Appendix

Table S.1: PICLE-PFN Architecture Hyperparameters.

Benchmark	Method	$\tau=2$	$\tau=4$	$\tau=6$	$\tau=8$	$\tau=10$
Cancer	PICLE-PFN	0.31 $\pm$ 0.03 (2.3 $\pm$ 0.6)	0.26 $\pm$ 0.04 (1.6 $\pm$ 0.8)	0.18 $\pm$ 0.03 (1.1 $\pm$ 0.3)	0.12 $\pm$ 0.03 (1.1 $\pm$ 0.3)	0.08 $\pm$ 0.02 (1.1 $\pm$ 0.3)
	IGC-Net	0.32 $\pm$ 0.02 (2.6 $\pm$ 0.5)	0.29 $\pm$ 0.04 (2.5 $\pm$ 0.8)	0.27 $\pm$ 0.03 (3.3 $\pm$ 0.7)	0.44 $\pm$ 0.08 (4.8 $\pm$ 0.5)	0.16 $\pm$ 0.07 (3.4 $\pm$ 1.2)
	CT	0.24 $\pm$ 0.02 (1.1 $\pm$ 0.2)	0.27 $\pm$ 0.03 (1.9 $\pm$ 0.7)	0.24 $\pm$ 0.06 (2.4 $\pm$ 1.0)	0.19 $\pm$ 0.09 (2.5 $\pm$ 1.1)	0.15 $\pm$ 0.12 (2.8 $\pm$ 1.3)
	RMSN	0.55 $\pm$ 0.03 (4.7 $\pm$ 0.5)	0.50 $\pm$ 0.07 (5.0 $\pm$ 0.2)	0.35 $\pm$ 0.06 (4.5 $\pm$ 0.8)	0.23 $\pm$ 0.05 (3.2 $\pm$ 0.7)	0.16 $\pm$ 0.04 (3.7 $\pm$ 0.9)
	G-Net	0.75 $\pm$ 0.03 (6.0 $\pm$ 0.2)	0.68 $\pm$ 0.04 (6.0 $\pm$ 0.0)	0.63 $\pm$ 0.05 (6.0 $\pm$ 0.0)	0.62 $\pm$ 0.05 (6.0 $\pm$ 0.0)	0.62 $\pm$ 0.06 (6.0 $\pm$ 0.0)
	CRN	0.50 $\pm$ 0.12 (4.4 $\pm$ 0.8)	0.40 $\pm$ 0.06 (4.0 $\pm$ 0.2)	0.31 $\pm$ 0.05 (3.7 $\pm$ 1.0)	0.23 $\pm$ 0.05 (3.4 $\pm$ 0.9)	0.18 $\pm$ 0.05 (4.0 $\pm$ 0.8)

Appendix

Table S.2: Longitudinal CAPO benchmark, per horizon. Each cell reports the normalized RMSE, mean

\pm

SD across the 20 repeats; values in parentheses are the mean

\pm

SD of the method’s per-task RMSE ranking (

\downarrow

better). The best and second best methods per benchmark and horizon are highlighted.

Method	Type	Cancer			MIMIC_Y1			MIMIC_1T
Method	Type	train	infer	total	train	infer	total	train	infer	total
PICLE-PFN	amortised	—	0.13 $\pm$ 0.00	0.13 $\pm$ 0.00	—	0.14 $\pm$ 0.00	0.14 $\pm$ 0.00	—	0.14 $\pm$ 0.00	0.14 $\pm$ 0.00
IGC-Net	fit per horizon	525 $\pm$ 11	0.08 $\pm$ 0.00	525 $\pm$ 11	582 $\pm$ 5	0.08 $\pm$ 0.01	582 $\pm$ 5	725 $\pm$ 10	0.09 $\pm$ 0.02	725 $\pm$ 10
CT	fit once	23.84 $\pm$ 0.57	0.25 $\pm$ 0.03	24.09 $\pm$ 0.57	78.82 $\pm$ 0.47	1.08 $\pm$ 0.02	79.90 $\pm$ 0.46	66.46 $\pm$ 2.04	1.26 $\pm$ 0.11	67.72 $\pm$ 1.96
RMSN	fit once	65.62 $\pm$ 1.36	0.10 $\pm$ 0.01	65.71 $\pm$ 1.35	33.96 $\pm$ 0.55	0.07 $\pm$ 0.00	34.03 $\pm$ 0.55	6.26 $\pm$ 0.49	0.08 $\pm$ 0.00	6.34 $\pm$ 0.50
G-Net	fit once	4.31 $\pm$ 0.42	6.58 $\pm$ 0.73	10.88 $\pm$ 0.47	12.00 $\pm$ 0.45	6.62 $\pm$ 0.18	18.62 $\pm$ 0.40	4.87 $\pm$ 0.37	8.45 $\pm$ 0.56	13.32 $\pm$ 0.41

Appendix

Table S.3: Runtime(seconds) at deployment. One evaluation task per setting (1000 context, 1000 queries), one A100 GPU, mean

\pm

SD over 5 timing repeats. Baselines are measured on their tuned HPs. CT, RMSN, G-Net, CRN, and TE-CDE is trained once then evaluated on all five horizons; runtime is computed as total time of one training and five inference calls, divided by five); IGC-net, D3, DLTMLE, and DeepACE train one model per horizon; runtime is measured as one training and one inference call. PICLE-PFN is pretrained and only runs inference.

Method	Hyperparameter	Search space
CRN, TE-CDE, G-Net	learning rate	$\log\mathcal{U}[10^{-4},\,10^{-2}]$
	number of layers	$\{1,2,3\}$
	batch size	$\{32,64,128,256\}$
	sequence hidden units	$\{0.5,1,2,3,4\}\times$ input dim
	balanced representation size	$\{0.5,1,2,3,4\}\times$ input dim
	FC hidden units	$\{0.5,1,2,3,4\}\times$ input dim

Appendix

Table S.4: Hyperparameter search spaces for the baselines. Each space is searched with Optuna for 2 h on one A100, selecting by held-out factual RMSE.

\mathcal{U}

and

\log\mathcal{U}

denote uniform and log-uniform ranges; braces denote categorical choices.

Panel of Figure 15 — Figure S.2: Coverage against overlap Overlap (x-axis) is measured as the proportion of observed trajectories that exactly match the target policy, and coverage (y-axis) is evaluated using nominal 95% intervals. Hollow points denote tasks with zero empirical overlap. Tasks with positive overlap generally achieve coverage above 0.9, whereas zero-overlap tasks exhibit the largest deviation from nominal coverage, in either direction.

		mimic_extract_1t				mimic_extract_1t_hdp1
Horizon	$n$ context	PICLE-PFN	CT	IGC-Net	D3	PICLE-PFN	IGC-Net	D3
$h=2$	100	0.456 $\pm$ 0.145	0.631 $\pm$ 0.195	0.447 $\pm$ 0.079	0.596 $\pm$ 0.301	0.403 $\pm$ 0.132	0.406 $\pm$ 0.108	0.486 $\pm$ 0.131
	300	0.382 $\pm$ 0.091	0.385 $\pm$ 0.082	0.387 $\pm$ 0.077	0.498 $\pm$ 0.124	0.347 $\pm$ 0.107	0.355 $\pm$ 0.104	0.399 $\pm$ 0.111
	1000	0.377 $\pm$ 0.125	0.374 $\pm$ 0.198	0.435 $\pm$ 0.309	0.409 $\pm$ 0.086	0.324 $\pm$ 0.100	0.337 $\pm$ 0.096	0.355 $\pm$ 0.101
$h=4$	100	0.621 $\pm$ 0.204	0.838 $\pm$ 0.210	0.698 $\pm$ 0.178	0.745 $\pm$ 0.217	0.574 $\pm$ 0.181	0.661 $\pm$ 0.148	0.687 $\pm$ 0.159
	300	0.537 $\pm$ 0.133	0.578 $\pm$ 0.142	0.645 $\pm$ 0.202	0.596 $\pm$ 0.133	0.500 $\pm$ 0.151	0.539 $\pm$ 0.148	0.575 $\pm$ 0.161

Appendix

Table S.5: Context sample size under a static and a history-dependent target policy. RMSE on the query split (size=1000), mean

\pm

SD over 20 tasks (

\downarrow

better).CT cannot handle dynamic policies and is not run on the history-dependent setting.

	Normalized RMSE $\downarrow$		Runtime (s) $\downarrow$
Horizon $\tau$	PICLE-PFN	CausalLongPFN	PICLE-PFN	CausalLongPFN
1	0.275 $\pm$ 0.026	0.195 $\pm$ 0.020	0.039	6.67
2	0.308 $\pm$ 0.030	0.278 $\pm$ 0.025	0.051	12.49
3	0.298 $\pm$ 0.037	0.306 $\pm$ 0.026	0.065	18.80
4	0.264 $\pm$ 0.038	0.298 $\pm$ 0.030	0.085	24.75
5	0.216 $\pm$ 0.032	0.278 $\pm$ 0.035	0.100	31.18

Appendix

Table S.6: Our PFN vs. CausalLongPFN on cancer : normalized RMSE (mean

\pm

std over 20 tasks) and inference runtime per task (1000 context / 1000 query trajectories, one A100).

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 266 papers

Smooth Multi-Policy Causal Effect Estimation in Longitudinal Settings
May 14, 2026Wenxin Chen, Weishen Pan, Kyra Gan +1Causal Effect Estimation Causal Inference
Cornell University, New York, US.
Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction
Jun 4, 2026Amirhossein Zare, Amirhessam Zare, Herlock Rahimi +2Structural Causal Models Counterfactual Prediction
School of Medicine, Tehran University of Medical Sciences, Tehran, Iran. · Department of Electrical and Computer Engineering, Yale University, New Haven, CT, USA. · School of Medicine, Shiraz University of Medical Sciences, Shiraz, Iran. +2
Interventional Time Series Priors for Causal Foundation Models
Mar 11, 2026Dennis Thumm, Ying ChenCausal Foundation Models Synthetic Data Generation
National University of Singapore · Department of Mathematics, Centre for Quantitative Finance, Risk Management Institute, National University of Singapore

May 14, 2026J/K move · Enter open · S save

Smooth Multi-Policy Causal Effect Estimation in Longitudinal Settings

Wenxin Chen, Weishen Pan, Kyra Gan, Fei Wang

Cornell University, New York, US.

Causal Effect Estimation Causal Inference

Open paper AskPDF arXiv

Comparative evaluation of multiple dynamic treatment policies is essential for healthcare and policy decisions, yet conventional longitudinal causal inference methods estimate each in isolation, preventing information sharing across counterfactuals. We demonstrate that this separate estimation paradigm induces a structurally uncontrolled second-order bias, inflating finite-sample variance even after standard debiasing with longitudinal targeted maximum likelihood estimation(LTMLE). To address this, we propose a policy-aware reparameterization of Iterative Conditional Expectation (ICE) Q-functions that enables joint estimation through shared representations. We implement this approach in the Policy-Encoded Q Network (PEQ-Net), an architecture centered on a shared policy encoder. The encoder is trained using kernel mean embeddings, ensuring that the learned representation space reflects population-level policy dissimilarities. After applying an LTMLE correction step, we prove this design imposes a structural constraint on the second-order remainder, thereby stabilizing finite-sample variance. Experiments on semi-synthetic datasets demonstrate that PEQ-Net consistently outperforms existing ICE-based methods, achieving substantial reductions in root-mean-square error, particularly when evaluating closely related policies.

Abstract

Causal Effect Estimation

Explore similar work

Smooth Multi-Policy Causal Effect Estimation in Longitudinal Settings

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Interventional Time Series Priors for Causal Foundation Models