Foundation model predicts effects of changing treatment over time
A Foundation Model for Policy In-Context Longitudinal Causal Effect Estimation
Artificial Intelligence
Summary
Estimating the impact of treatments over time is hard because treatments and patient conditions can influence each other in complex ways. The authors developed a model called PICLE-PFN that learns from simulated data and can quickly predict how different treatment plans affect outcomes without needing to be retrained for each new situation. This model works with changing treatment policies and different time lengths, making it faster and more reliable than older methods tested on similar problems. Essentially, it offers a flexible tool to understand what might happen if treatments change over time.
What this means in practice
- •For clinical decision support teams: Quickly predict patient outcomes under changing treatment rules without retraining models for every new case.
- •For health policy analysts: Evaluate the effects of different health intervention policies over time using a unified prediction model.
Authors
Wenxin Chen, Kyra Gan, Mingyu Lu, Zilin Jing, Fei Wang, Hoifung Poon, Tristan Naumann, Peniel Argaw
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.