KEIM: Knowledge Graph Empowered Interpretable Model for Diagnosis Prediction

Electronic Health Records (EHR) include various sources of healthcare data collected from patients in hospitals. These data are typically stored in structured formats and are widely used in various big data healthcare analysis applications, particularly diagnosis prediction. Deep learning methods have achieved record-breaking results in various real-world prediction tasks. However, deep learning methods usually require a large amount of data for training, and the medical features that rarely appear in the data also pose great challenges for deep learning models. Besides, while deep learning models often achieve high accuracy, the lack of interpretation remains a problem for healthcare applications, which are naturally high-stakes. Existing works utilize medical ontology knowledge to enhance prediction performance and provide interpretable prediction results. Nevertheless, the ontology knowledge is coarse-grained, where many medical concepts and relationships are not included. In this paper, we propose to incorporate large-scale medical knowledge graphs (KGs) into our designed model, called KEIM (Knowledge graph Empowered Interpretable Model), for diagnosis prediction. Specifically, the KGs are first integrated into the time-series module of the model via a laplacian regularization to take advantage of the complex relationships among medical features. Subsequently, we construct a personalized KG for each visit and design a relation-aware attentive graph neural network based on the personalized KG to augment the time-series module for interpretable predictions. Extensive experiments on two benchmark healthcare datasets, namely, MIMIC-III and MIMIC-IV, show that our proposed KEIM not only achieves significant improvement in terms of AUC but also provides interpretability for diagnosis prediction with KGs.