Research map / Probabilistic and scientific modeling / Treatment effects
KG-TREAT: Pre-training for Treatment Effect Estimation by Synergizing Patient Data with Knowledge Graphs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes KG-TREAT, a pre-training and fine-tuning framework that synergizes MarketScan patient records with UMLS knowledge graphs via dual-focus personalized KGs and a bi-level attention module (DIVE), evaluated on four CAD treatment-effect tasks.
Problem
Treatment effect estimation from sparse, high-dimensional observational patient data is limited by small labeled datasets and by the inability of existing methods to exploit medical knowledge while distinguishing treatment, covariate, and outcome relationships.
Approach
Construct two personalized KGs per patient (treatment-covariate and outcome-covariate) from UMLS, fuse them with patient visit sequences through a deep bi-level attention synergy (DIVE) in multiple layers, and pre-train with masked code prediction plus KG link prediction before fine-tuning for joint treatment and outcome prediction.
Key results
- KG-TREAT reaches AUC 0.828/0.858/0.851/0.843 on the four downstream tasks, versus best baseline TransTEE 0.751/0.753/0.770/0.803.
- Averaged across tasks, the paper reports about 7% AUC and 9% IF-PEHE improvement over the best baseline.
- Estimated effects match the published RCT conclusions on all four CAD trial comparisons, while several baselines do not.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| pre-train by synergizing patient data with KGs, then fine-tune for TEE | KG-TREAT · four downstream TEE datasets | AUC: 7% | Helped |
| pre-train by synergizing patient data with KGs, then fine-tune for TEE | KG-TREAT · four downstream TEE datasets | IF-PEHE: 9% | Helped |
| add DIVE instead of final-stage patient/graph concatenation (ablation) | KG-TREAT · four downstream TEE datasets | AUC | Helped |
| add KGs to pre-training (ablation) | KG-TREAT · four downstream TEE datasets | AUC | Helped |
| add pre-training (ablation) | KG-TREAT · four downstream TEE datasets | AUC | Helped |
| add both MCP and LP instead of MCP only (ablation) | KG-TREAT · four downstream TEE datasets | AUC: 3% of AUC | Helped |
| add both MCP and LP instead of LP only (ablation) | KG-TREAT · four downstream TEE datasets | AUC: 1% of AUC | Helped |
| replace DistMult with RotatE for link prediction (ablation) | KG-TREAT · four downstream TEE datasets | AUC | No clear effect |
| replace DistMult with TransE for link prediction (ablation) | KG-TREAT · four downstream TEE datasets | AUC | No clear effect |
Surprising findings
- Removing pre-training hurts more than removing KGs (w/o pre-train vs w/o KGs), indicating pre-training contributes more than the knowledge graph itself.
- Masked code prediction matters more than link prediction in the pre-training ablation.
Limitations
- Downstream evaluation is limited to four CAD datasets derived from RCTs, all with stroke or myocardial infarction outcomes.
- Patient data is proprietary MarketScan claims, limiting reproducibility of the pre-training corpus.
- Ground-truth treatment effects are unavailable, so validation relies on alignment with RCT conclusions and an IF-PEHE approximation.
Experimental setup
Datasets: MarketScan CCAE UMLS CAD-related RCT-derived downstream datasets
Benchmarks: Rivaroxaban v.s. Aspirin Valsartan v.s. Ramipril Ticagrelor v.s. Aspirin Apixaban v.s. Warfarin
Metrics: AUC AUPR IF-PEHE
Baselines: TARNet DragonNet DR-CFR TNet SNet FlexTENet TransTEE
Builds on: BERT TARNet DragonNet DistMult BioLinkBERT
Closest papers by meaning
- Continuous Treatment Effect Estimation Using Gradient Interpolation and Kernel Smoothing (AAAI 2024 · similarity 0.87)
- PairNet: Training with Observed Pairs to Estimate Individual Treatment Effect (ICML 2024 · similarity 0.86)
- LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding (NeurIPS 2025 · similarity 0.86)
- AutoCATE: End-to-End, Automated Treatment Effect Estimation (ICML 2025 · similarity 0.85)
- Enhancing Treatment Effect Estimation via Active Learning: A Counterfactual Covering Perspective (ICML 2025 · similarity 0.85)
- Defining Expertise: Applications to Treatment Effect Estimation (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.