# KG-TREAT: Pre-training for Treatment Effect Estimation by Synergizing Patient Data with Knowledge Graphs

arXiv: https://arxiv.org/abs/2403.03791 · HTML: https://mlsanity.com/paper/2403.03791/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 34 references that are in the map.
- First posted on arXiv 6 Mar 2024.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction](https://mlsanity.com/paper/2005.12833/) (2020)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [BEHRT: Transformer for Electronic Health Records](https://mlsanity.com/paper/1907.09538/) (2019)
- [LinkBERT: Pretraining Language Models with Document Links](https://mlsanity.com/paper/2203.15827/) (2022)
- [GRAM: Graph-based Attention Model for Healthcare Representation Learning](https://mlsanity.com/paper/1611.07012/) (2016)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Scalable Multi-Hop Relational Reasoning for Knowledge-Aware Question Answering](https://mlsanity.com/paper/2005.00646/) (2020)
- [Large-scale Pretraining for Visual Dialog: A Simple State-of-the-Art Baseline](https://mlsanity.com/paper/1912.02379/) (2019)
- [CETransformer: Casual Effect Estimation via Transformer Based Representation Learning](https://mlsanity.com/paper/2107.08714/) (2021)

## Built on by (most cited first, 5 of 5)

- [HealthFlow: A Self-Evolving AI Agent with Meta Planning for Autonomous Healthcare Research](https://mlsanity.com/paper/2508.02621/) (2025)
- [Large Language Models for Knowledge Graph Embedding: A Survey](https://mlsanity.com/paper/2501.07766/) (2025)
- [A Comparative Study of Model Adaptation Strategies for Multi-Treatment Uplift Modeling](https://mlsanity.com/paper/2511.01185/) (2025)
- [Knowledge Graphs Meet Graph Neural Networks: A Comprehensive Survey](https://mlsanity.com/paper/2607.09666/) (2026)
- [Estimating Treatment Effects for Depression in Longitudinal Therapy Switching Settings](https://mlsanity.com/paper/2607.27214/) (2026)

## Abstract (arXiv)

Treatment effect estimation (TEE) is the task of determining the impact of various treatments on patient outcomes. Current TEE methods fall short due to reliance on limited labeled data and challenges posed by sparse and high-dimensional observational patient data. To address the challenges, we introduce a novel pre-training and fine-tuning framework, KG-TREAT, which synergizes large-scale observational patient data with biomedical knowledge graphs (KGs) to enhance TEE. Unlike previous approaches, KG-TREAT constructs dual-focus KGs and integrates a deep bi-level attention synergy method for in-depth information fusion, enabling distinct encoding of treatment-covariate and outcome-covariate relationships. KG-TREAT also incorporates two pre-training tasks to ensure a thorough grounding and contextualization of patient data and KGs. Evaluation on four downstream TEE tasks shows KG-TREAT's superiority over existing methods, with an average improvement of 7% in Area under the ROC Curve (AUC) and 9% in Influence Function-based Precision of Estimating Heterogeneous Effects (IF-PEHE). The effectiveness of our estimated treatment effects is further affirmed by alignment with established randomized clinical trial findings.