Prototype-Guided Contrastive Knowledge Graph Representation Learning for Diagnosis Prediction

Wooyeol Park, Ahmad Wisnu Mulyadi, Eunsong Kang, Heung-Il Suk

Abstract: The medical knowledge graph (KG) constructed from electronic health records (EHR) offers a comprehensive understanding of patients as it facilitates interconnections among medical codes. While KG has been proven to improve diagnostic accuracy, it inevitably contains noisy and incomplete attributes that could compromise the verdicts. Hence, direct utilization of medical KG for downstream clinical tasks may entail inadequacy, even by resorting to proper machine learning methods. Several prior studies have initiated noise filtering methods upon medical KG to relieve such issues by masking its task-irrelevant nodes and relations. However, they focused solely on a patient-wise perspective that feasibly overlooked essential shared traits from a (sub-)population viewpoint. This study proposes a novel prototype-guided predictive model called ProtoCare to predict diagnoses in EHR by unifying prototype and contrastive learning. ProtoCare leverages a set of prototypes to discover latent shared characteristics among groups of patients, while contrastive learning plays the role of noise filtering and KG refinement accounting for both learned patient and selected prototype features. Experimental results on the real-world EHR cohort exhibited that our proposed ProtoCare outperformed several baselines on the diagnosis prediction task.