A Foundation Model for Policy In-Context Longitudinal Causal Effect Estimation
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 | ✓ | ✓ | ✓ | ✓ | ✓ | ✗ | ✗ | ✗ | ✗ | ✓ |
| Dataset | Policy type | |
| cancer | Static | |
| mimic_y1 | Static | |
| mimic_y2 | Static | |
| mimic_1t | Static | |
| mimic_1t_hdp0.5 | Dynamic | |
| mimic_1t_hdp1 | Dynamic |
| Method | Cancer | MIMIC | |||||
| Y1 | Y2 | 1T | 1T-HDP0.5 | 1T-HDP1 | 1T-HDP2 | ||
| PICLE-PFN | 0.190 0.03 | 0.958 0.50 | 0.603 0.23 | 0.583 0.20 | 0.527 0.13 | 0.534 0.13 | 0.555 0.15 |
| IGC-Net | 0.297 0.02 | 0.901 0.49 | 0.604 0.18 | 0.699 0.11 | 0.546 0.13 | 0.580 0.13 | 0.622 0.17 |
| CT | 0.217 0.06 | 0.851 0.29 | 0.886 0.82 | 0.660 0.40 | — | — | — |
| RMSN | 0.357 0.04 | 1.242 0.70 | 0.830 0.13 | 0.786 0.21 | — | — | — |
| G-Net | 0.661 0.04 | 1.333 0.86 | 0.984 0.14 | 1.127 0.20 | — | — | — |
| Method | ||||
| PICLE-PFN | 0.730 0.030 | 0.737 0.032 | 0.754 0.026 | 0.767 0.037 |
| IGC-Net | 0.736 0.040 | 0.740 0.038 | 0.762 0.031 | 0.779 0.028 |
| PICLE-PFN wins | 7/10 | 5/10 | 8/10 | 8/10 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Component | Value |
| Transformer layers | 8 |
| Embedding width | 128 |
| Attention heads | 4 |
| MLP hidden width | 512 |
| Activation | GELU |
| Dropout | 0 |
| Benchmark | Method | |||||
| Cancer | PICLE-PFN | 0.31 0.03 (2.3 0.6) | 0.26 0.04 (1.6 0.8) | 0.18 0.03 (1.1 0.3) | 0.12 0.03 (1.1 0.3) | 0.08 0.02 (1.1 0.3) |
| IGC-Net | 0.32 0.02 (2.6 0.5) | 0.29 0.04 (2.5 0.8) | 0.27 0.03 (3.3 0.7) | 0.44 0.08 (4.8 0.5) | 0.16 0.07 (3.4 1.2) | |
| CT | 0.24 0.02 (1.1 0.2) | 0.27 0.03 (1.9 0.7) | 0.24 0.06 (2.4 1.0) | 0.19 0.09 (2.5 1.1) | 0.15 0.12 (2.8 1.3) | |
| RMSN | 0.55 0.03 (4.7 0.5) | 0.50 0.07 (5.0 0.2) | 0.35 0.06 (4.5 0.8) | 0.23 0.05 (3.2 0.7) | 0.16 0.04 (3.7 0.9) | |
| G-Net | 0.75 0.03 (6.0 0.2) | 0.68 0.04 (6.0 0.0) | 0.63 0.05 (6.0 0.0) | 0.62 0.05 (6.0 0.0) | 0.62 0.06 (6.0 0.0) | |
| CRN | 0.50 0.12 (4.4 0.8) | 0.40 0.06 (4.0 0.2) | 0.31 0.05 (3.7 1.0) | 0.23 0.05 (3.4 0.9) | 0.18 0.05 (4.0 0.8) |
| Method | Type | Cancer | MIMIC_Y1 | MIMIC_1T | ||||||
| train | infer | total | train | infer | total | train | infer | total | ||
| PICLE-PFN | amortised | — | 0.13 0.00 | 0.13 0.00 | — | 0.14 0.00 | 0.14 0.00 | — | 0.14 0.00 | 0.14 0.00 |
| IGC-Net | fit per horizon | 525 11 | 0.08 0.00 | 525 11 | 582 5 | 0.08 0.01 | 582 5 | 725 10 | 0.09 0.02 | 725 10 |
| CT | fit once | 23.84 0.57 | 0.25 0.03 | 24.09 0.57 | 78.82 0.47 | 1.08 0.02 | 79.90 0.46 | 66.46 2.04 | 1.26 0.11 | 67.72 1.96 |
| RMSN | fit once | 65.62 1.36 | 0.10 0.01 | 65.71 1.35 | 33.96 0.55 | 0.07 0.00 | 34.03 0.55 | 6.26 0.49 | 0.08 0.00 | 6.34 0.50 |
| G-Net | fit once | 4.31 0.42 | 6.58 0.73 | 10.88 0.47 | 12.00 0.45 | 6.62 0.18 | 18.62 0.40 | 4.87 0.37 | 8.45 0.56 | 13.32 0.41 |
| Method | Hyperparameter | Search space |
| CRN, TE-CDE, G-Net | learning rate | |
| number of layers | ||
| batch size | ||
| sequence hidden units | input dim | |
| balanced representation size | input dim | |
| FC hidden units | input dim |
| mimic_extract_1t | mimic_extract_1t_hdp1 | |||||||
| Horizon | context | PICLE-PFN | CT | IGC-Net | D3 | PICLE-PFN | IGC-Net | D3 |
| 100 | 0.456 0.145 | 0.631 0.195 | 0.447 0.079 | 0.596 0.301 | 0.403 0.132 | 0.406 0.108 | 0.486 0.131 | |
| 300 | 0.382 0.091 | 0.385 0.082 | 0.387 0.077 | 0.498 0.124 | 0.347 0.107 | 0.355 0.104 | 0.399 0.111 | |
| 1000 | 0.377 0.125 | 0.374 0.198 | 0.435 0.309 | 0.409 0.086 | 0.324 0.100 | 0.337 0.096 | 0.355 0.101 | |
| 100 | 0.621 0.204 | 0.838 0.210 | 0.698 0.178 | 0.745 0.217 | 0.574 0.181 | 0.661 0.148 | 0.687 0.159 | |
| 300 | 0.537 0.133 | 0.578 0.142 | 0.645 0.202 | 0.596 0.133 | 0.500 0.151 | 0.539 0.148 | 0.575 0.161 | |
| Normalized RMSE | Runtime (s) | |||
| Horizon | PICLE-PFN | CausalLongPFN | PICLE-PFN | CausalLongPFN |
| 1 | 0.275 0.026 | 0.195 0.020 | 0.039 | 6.67 |
| 2 | 0.308 0.030 | 0.278 0.025 | 0.051 | 12.49 |
| 3 | 0.298 0.037 | 0.306 0.026 | 0.065 | 18.80 |
| 4 | 0.264 0.038 | 0.298 0.030 | 0.085 | 24.75 |
| 5 | 0.216 0.032 | 0.278 0.035 | 0.100 | 31.18 |
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