KAT-GNN: A Knowledge-Augmented Temporal Graph Neural Network for Risk Prediction in Electronic Health Records

Lin, Kun-Wei, Kuo, Yu-Chen, Wang, Hsin-Yao, Tseng, Yi-Ju

arXiv.org Artificial Intelligence 

Abstract-- Clinical risk prediction using electronic health records (EHRs) is vital to facilitate timely interventions and clinical decision support. However, modeling heterogeneous and irregular temporal EHR data presents significant challenges. We propose KAT -GNN (Knowledge-Augmented Temporal Graph Neural Network), a graph-based framework that integrates clinical knowledge and temporal dynamics for risk prediction. These graphs are then augmented using two knowledge sources: (1) ontology-driven edges derived from SNOMED CT and (2) co-occurrence priors extracted from EHRs. Subsequently, a time-aware transformer is employed to capture longitudinal dynamics from the graph-encoded patient representations. KAT -GNN is evaluated on three distinct datasets and tasks: coronary artery disease (CAD) prediction using the Chang Gung Research Database (CGRD) and in-hospital mortality prediction using the MIMIC-III and MIMIC-IV datasets. KAT - GNN achieves state-of-the-art performance in CAD prediction (AUROC: 0.9269 0.0029) and demonstrated strong results in mortality prediction in MIMIC-III (AUROC: 0.9230 0.0070) and MIMIC-IV (AUROC: 0.8849 0.0089), consistently outperforming established baselines such as GRASP and RETAIN. Ablation studies confirm that both knowledge-based augmentation and the temporal modeling component are significant contributors to performance gains. These findings demonstrate that the integration of clinical knowledge into graph representations, coupled with a time-aware attention mechanism, provides an effective and generalizable approach for risk prediction across diverse clinical tasks and datasets. NTRODUCTION This study was supported by grants from the National Science and T echnology Council, T aiwan (NSTC 114-2221-E-A49-061), the Higher Education Sprout Project of the National Y ang Ming Chiao T ung University and MOE, T aiwan (CGMH-NYCU-114-CORPG2P0072), and Chang Gung Memorial Hospital (CMRPG2P0342).