W3PHIAI-26 · AAAI Workshop · 2026
Foundation Models for Heterogeneous Treatment Effect Estimation: Adapting Pretrained EHR Embeddings in Small Clinical Cohorts
Overview
Can pretrained patient representations help estimate how treatment effects vary when clinical cohorts are small? This study combines CLMBR embeddings with outcome-specific principal component analysis and a regularized R-learner. It compares GLP-1 receptor agonists with insulin or sulfonylureas in Stanford EHR data across four outcomes, with 46–151 patients per outcome. The manuscript reports lower cross-validated R-loss than the evaluated feature baselines and explores event-level model attributions. The study is retrospective and single-institution; broader clinical generalization remains to be established.
Method
- Represent pre-treatment patient histories with pretrained CLMBR embeddings.
- Use outcome-specific PCA and strong regularization to adapt high-dimensional embeddings to small cohorts.
- Estimate heterogeneous effects with an R-learner and cross-fitting; evaluate with cross-validated R-loss.
- Explore event-level model attributions connecting representation dimensions to estimated effects.
My contribution
I worked on the adaptation of pretrained EHR embeddings and the treatment-effect estimation pipeline with Bin Yang and Prof. Yuanjia Wang.
Full paper
Open PDF ↗
