DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

Deyi Li; Zijun Yao; Qi Xu; Muxuan Liang; Lingyao Li; Zijian Xu; Mei Liu

cs.LGMar 10, 2026

DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

Authors: Deyi Li, Zijun Yao, Qi Xu, Muxuan Liang, Lingyao Li, Zijian Xu, Mei Liu

Organizations: University of Florida, Gainesville, Florida, USA · University of Kansas, Lawrence, Kansas, USA · MD Anderson Cancer Center, Houston, Texas, USA · University of South Florida, Tampa, Florida, USA

Abstract

The growing adoption of electronic health record (EHR) systems has provided unprecedented opportunities for predictive modeling to guide clinical decision making. Structured EHRs contain longitudinal observations of patients across hospital visits, where each visit is represented by a set of medical codes. While sequence-based, graph-based, and graph-enhanced sequence approaches have been developed to capture rich code interactions over time or within the same visits, they often overlook the inherent heterogeneous roles of medical codes arising from distinct clinical characteristics and contexts. To this end, in this study we propose the Disease Trajectory-aware Transformer for EHR (DT-BEHRT), a graph-enhanced sequential architecture that disentangles disease trajectories by explicitly modeling diagnosis-centric interactions within organ systems and capturing asynchronous progression patterns. To further enhance the representation robustness, we design a tailored pre-training methodology that combines trajectory-level code masking with ontology-informed ancestor prediction, promoting semantic alignment across multiple modeling modules. Extensive experiments on multiple benchmark datasets demonstrate that DT-BEHRT achieves strong predictive performance and provides interpretable patient representations that align with clinicians' disease-centered reasoning. The source code is publicly accessible at https://github.com/GatorAIM/DT-BEHRT.git.

Figures & tables

Figure 1. The architecture of DT-BEHRT. Each layer includes a Sequence Representation (SR), Disease Aggregation (DA), and Disease Progression (DP) module. The Patient Representation (PR) is derived via sequence-guided attention. The pre-training framework includes Global Code Masking Prediction (GCMP) and Ancestor Code Prediction (ACP).

Models			G-BERT	BEHRT	Med-BERT	HypEHR	ExBEHRT	HEART	DT-BEHRT
MIMIC-III	Mortality	F1	$59.24{\pm}0.46$	$68.60{\pm}0.43$	$67.91{\pm}1.08$	$70.04{\pm}0.70$	$73.66{\pm}1.09$	$\underline{74.77{\pm}1.26}$	$\mathbf{76.03{\pm}0.28}$
		AUROC	$86.25{\pm}0.82$	$87.23{\pm}0.27$	$87.94{\pm}0.53$	$88.55{\pm}0.39$	$90.72{\pm}0.30$	$\mathbf{92.13{\pm}0.36}$	$\underline{92.09{\pm}0.15}$
		AUPRC	$72.13{\pm}1.51$	$75.33{\pm}0.33$	$75.28{\pm}1.01$	$76.39{\pm}1.06$	$81.44{\pm}0.62$	$\underline{82.76{\pm}0.63}$	$\mathbf{84.50{\pm}0.19}$
	PLOS	F1	$69.62{\pm}1.42$	$70.38{\pm}0.74$	$72.02{\pm}1.43$	$72.73{\pm}0.09$	$74.73{\pm}0.70$	$\underline{75.44{\pm}1.47}$	$\mathbf{76.37{\pm}0.49}$
		AUROC	$72.96{\pm}1.20$	$73.71{\pm}0.98$	$77.25{\pm}0.53$	$78.89{\pm}0.33$	$82.11{\pm}0.71$	$\underline{82.99{\pm}0.40}$	$\mathbf{84.13{\pm}0.26}$
		AUPRC	$72.48{\pm}1.27$	$72.49{\pm}1.47$	$76.98{\pm}0.56$	$78.70{\pm}0.28$	$83.52{\pm}0.79$	$\underline{83.83{\pm}0.59}$	$\mathbf{85.00{\pm}0.22}$

Table 1. Results of general outcome prediction tasks.

	Prevalence	All patients			Patients with $\geq$ 3 visits
		ExBEHRT	HEART	DT-BEHRT	ExBEHRT	HEART	DT-BEHRT
MIMIC-III
Acute and unspecified renal failure	16.00%	$\mathbf{49.96{\pm}0.96}$	$\underline{47.29{\pm}1.81}$	$44.62{\pm}1.91$	$46.98{\pm}4.75$	$46.52{\pm}2.52$	$\mathbf{54.26{\pm}4.55}$
Acute cerebrovascular disease	0.90%	$3.95{\pm}2.30$	$3.21{\pm}0.91$	$\mathbf{7.53{\pm}2.11}$	$2.11{\pm}0.96$	$\underline{3.82{\pm}2.99}$	$\mathbf{17.74{\pm}33.72}$
Acute myocardial infarction	3.70%	$\mathbf{18.49{\pm}1.40}$	$\underline{16.61{\pm}1.41}$	$17.27{\pm}1.18$	$\mathbf{17.81{\pm}7.67}$	$\underline{16.21{\pm}4.25}$	$11.28{\pm}3.64$
Cardiac dysrhythmias	20.10%	$74.68{\pm}1.48$	$\mathbf{74.98{\pm}1.47}$	$\underline{72.81{\pm}1.87}$	$71.35{\pm}6.60$	$\mathbf{77.02{\pm}6.09}$	$\underline{70.47{\pm}2.39}$

Table 2. Results of phenotyping prediction tasks.

	Performance
Variant	Architectures	DA ${}^{\text{w/o cov}}$	$\times$	$\checkmark$	$\times$	$\times$	$\times$	$\times$	$\times$	$\times$
		DA	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$
		DP	$\times$	$\times$	$\times$	$\checkmark$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$
	Pre-training Tasks	GCMP	$\times$	$\times$	$\times$	$\times$	$\times$	$\checkmark$	$\checkmark$	$\checkmark$
	Pre-training Tasks	ACP	$\times$	$\times$	$\times$	$\times$	$\times$	$\times$	$\times$	$\checkmark$
MIMIC-III	Mortality	F1	$71.59{\pm}1.29$	$72.01{\pm}1.76$	$73.51{\pm}1.69$	$72.06{\pm}0.82$	$73.48{\pm}0.47$	$74.17{\pm}0.09$	$\underline{75.40{\pm}0.49}$	$\mathbf{76.03{\pm}0.28}$

Table 3. Ablation study for general outcome prediction tasks using MIMIC-III and MIMIC-IV.

Figure 2. Subgroup performance radar plots for a mortality, b PLOS, and c readmission prediction tasks in comorbidity groups.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
$c,{\mathcal{C}}$	A medical code; the medical code vocabulary
${\mathcal{D}},{\mathcal{M}},{\mathcal{L}},{\mathcal{P}}$	Sets of diagnosis, medication, laboratory test, and procedure codes
$T$	Total number of hospital visits
$v_{t},{\mathcal{V}}$	Set of codes at visit $t$ ; the entire sequence of visits / patient trajectory
$N_{v_{t}},N_{V}$	Number of codes in visit $v_{t}$ and total number of codes in trajectory ${\mathcal{V}}$ , i.e., $N_{V}=\sum_{t=1}^{T}N_{v_{t}}$
${\bm{e}}_{c},{\bm{e}}_{{type}(c)},{\bm{e}}_{{visit}(c)}$	Embedding vector of code $c$ , its type, and its visit index

Appendix

Table 4. Notations used in this paper.

Figure 3. A sample attention mask with a DA token triggering threshold of $k=3$ . Each DA token is restricted to attend only to diagnosis codes within its corresponding ICD-9 chapter and to itself. In this case, ICD-9 code 291.0 (Alcohol withdrawal delirium) does not trigger a DA token.

Characteristic	MIMIC-III	MIMIC-IV	eICU
Number of patients	33,067	60,709	85,839
Diagnosis vocabulary size	1,998	1,983	838
Medication vocabulary size	145	140	2,042
Procedure vocabulary size	801	801	–
Laboratory test vocabulary size	1,500	1,281	755
Average visits per patient	1.21	1.39	1.16

Appendix

Table 5. Descriptive statistics of the MIMIC-III, MIMIC-IV, and eICU datasets.

Parameter	Search Space
Learning rate	{0.01, 0.001}
Batch size	{32, 64}
Number of layers ( $L$ )	{2, 3}
Number of GAT blocks ( $L_{\mathcal{G}}$ )	{2}
Hidden dimension ( $\mathtt{d}$ )	{64, 128}
DA token threshold ( $k$ )	{3, 4}

Appendix

Table 6. Model parameters and their search space.

Figure 4. Illustration of Case 1 with attention scores of the a PR module, b SR module, c DA module, and d DP module. Only edges with scores $>$ 0.001 are displayed, and self-loops are removed.

Figure 5. Illustration of Case 2 with attention scores of the a PR module, b SR module, c DA module, and d DP module.

Explore similar work

MiGHT-EHR: A Multi-task Graph Transformer for Heterogeneous Temporal Electronic Health Records
Aug 5, 2026Anirudh Rayas, Yuan Wang, Pavan TuragaHeterogeneous GNNs Clinical Decision Support
Arizona State University, Tempe, AZ, USA · University of South Carolina, Columbia, SC, USA
DT-Transformer: A Foundation Model for Disease Trajectory Prediction on a Real-world Health System
May 14, 2026Yunying Zhu, Andrew R Weckstein, Kueiyu Joshua Lin +1Healthcare Clinical Prediction
Division of Pharmacoepidemiology and Pharmacoeconomics, Department of Medicine, Brigham and Women’s Hospital, Harvard Medical School, Boston, MA, USA · Harvard T.H. Chan School of Public Health, Harvard University, Boston, MA, USA · Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University, MA, USA +2
Explainable Transformer Models for Clinical Prediction Tasks on Structured Electronic Health Records
Aug 20, 2026Jun Ni Du, Lukas Adamek, Maxim Kryukov +4Clinical Outcome Prediction Longitudinal EHR Modeling
Data & Computational Science, R&D, Sanofi Toronto, Canada · Data & Computational Science, R&D, Sanofi Barcelona, Spain · Clinical Real World Evidence, R&D, Sanofi Cambridge, USA +2

Aug 5, 2026J/K move · Enter open · S save

MiGHT-EHR: A Multi-task Graph Transformer for Heterogeneous Temporal Electronic Health Records

Anirudh Rayas, Yuan Wang, Pavan Turaga

Arizona State University, Tempe, AZ, USA · University of South Carolina, Columbia, SC, USA

Heterogeneous GNNs Clinical Decision Support Clinical Outcome Prediction Multi-Task Learning Longitudinal EHR Modeling Graph Transformer

Open paper AskPDF arXiv

Learning from Electronic Health Records (EHRs) has gained significant attention due to its potential to improve clinical prediction. However, effective learning remains challenging because EHRs encode heterogeneous, temporally ordered clinical interactions. In particular, EHRs contain: (i) heterogeneous clinical entities, including patients, visits, diagnoses, prescriptions, and procedures, together with their heterogeneous interactions, (ii) longitudinal patient trajectories across hospital visits and (iii) shared statistical dependencies across related clinical prediction tasks. Existing EHR learning methods capture only a subset of these properties. To bridge this gap, we propose Multi-task Graph transformer for Heterogeneous Temporal EHRs (MiGHT-EHR), which jointly models all three within a unified representation learning method. MiGHT-EHR constructs a heterogeneous graph from EHRs in which nodes represent clinical entities and edges connect statistically associated entities identified via normalized point-wise mutual information. Across MIMIC-III and MIMIC-IV datasets, MiGHT-EHR outperforms state-of-the-art methods on average across four tasks: drug recommendation, prediction of length-of-stay, mortality, and readmission, with particularly strong improvements in mortality and readmission prediction. Furthermore, a post-hoc analysis of the learned representations reveals that patient neighborhoods are organized by clinical outcomes, salient medical concepts are recoverable as linear directions in the representation space, and task probabilities are well calibrated. Collectively, these findings demonstrate that MiGHT-EHR representations support diverse prediction tasks while preserving clinically interpretable structure.

Abstract

Explore similar work

MiGHT-EHR: A Multi-task Graph Transformer for Heterogeneous Temporal Electronic Health Records

DT-Transformer: A Foundation Model for Disease Trajectory Prediction on a Real-world Health System

Explainable Transformer Models for Clinical Prediction Tasks on Structured Electronic Health Records