[Paper Review] Self-Supervised Graph Learning with Hyperbolic Embedding for Temporal Health Event Prediction

Chang Lü, Chandan K. Reddy+1 more|arXiv (Cornell University)|Jun 8, 2021

Machine Learning in HealthcareComputer Science47 references46 citations

TL;DR

Sherbet proposes a self-supervised graph learning framework with hyperbolic embeddings to improve temporal health event prediction using electronic health records (EHR). It leverages disease hierarchies via a novel hyperbolic embedding method with information flow, constructs a weighted, directed disease interaction graph, and employs multi-level attention to enable both generic and personalized interpretability, significantly improving prediction performance on public EHR datasets while fully utilizing unlabeled data.

ABSTRACT

Electronic Health Records (EHR) have been heavily used in modern healthcare systems for recording patients' admission information to hospitals. Many data-driven approaches employ temporal features in EHR for predicting specific diseases, readmission times, or diagnoses of patients. However, most existing predictive models cannot fully utilize EHR data, due to an inherent lack of labels in supervised training for some temporal events. Moreover, it is hard for existing works to simultaneously provide generic and personalized interpretability. To address these challenges, we first propose a hyperbolic embedding method with information flow to pre-train medical code representations in a hierarchical structure. We incorporate these pre-trained representations into a graph neural network to detect disease complications, and design a multi-level attention method to compute the contributions of particular diseases and admissions, thus enhancing personalized interpretability. We present a new hierarchy-enhanced historical prediction proxy task in our self-supervised learning framework to fully utilize EHR data and exploit medical domain knowledge. We conduct a comprehensive set of experiments and case studies on widely used publicly available EHR datasets to verify the effectiveness of our model. The results demonstrate our model's strengths in both predictive tasks and interpretable abilities.

Motivation & Objective

• To address the underutilization of EHR data in supervised models due to label scarcity in temporal event prediction.
• To exploit hierarchical medical code structures (e.g., ICD-9-CM) for better disease representation learning.
• To model disease complications and interactions through a weighted, directed disease interaction graph.
• To simultaneously provide generic interpretability (disease-level patterns) and personalized interpretability (patient-level contributions) in predictions.
• To design a self-supervised proxy task that leverages all admissions—including single-admission and final admissions—by hierarchically predicting historical diagnoses.

Proposed method

• Proposes a hyperbolic embedding method with information flow to pre-train medical code representations using the hierarchical structure of ICD-9-CM codes.
• Constructs a weighted and directed disease interaction graph based on co-occurrence patterns in patient admission records.
• Employs a graph neural network (GNN) on the disease interaction graph to learn disease complication patterns and hidden representations.
• Introduces a multi-level attention mechanism to compute code-level (disease) and admission-level contributions to predictions, enabling personalized interpretability.
• Designs a hierarchy-enhanced historical prediction proxy task in self-supervised learning to generate labels for all admissions, including single-admission and final admissions.
• Integrates the self-supervised pre-training with a fine-tuning module for downstream temporal prediction tasks such as diagnosis prediction.

Experimental results

Research questions

• RQ1 Can a self-supervised proxy task that hierarchically predicts historical diagnoses improve the utilization of unlabeled EHR data?
• RQ2 Can hyperbolic embeddings effectively model the hierarchical structure of medical codes while preserving semantic similarity?
• RQ3 Can a weighted, directed disease interaction graph capture clinically relevant disease complications?
• RQ4 Can multi-level attention mechanisms simultaneously provide generic and personalized interpretability in EHR-based prediction?
• RQ5 Does the integration of hierarchical structure, disease interactions, and self-supervised pre-training lead to improved performance on temporal health event prediction?

Key findings

• Sherbet achieves state-of-the-art performance on multiple EHR benchmark datasets for diagnosis prediction, outperforming existing supervised and self-supervised models.
• The model effectively captures disease complications, as shown by t-SNE visualizations placing related diseases (e.g., hypertension and heart failure) in close proximity in the hyperbolic space.
• The model distinguishes between different disease subtypes, such as diabetes type I and type II, based on distinct complications, indicating meaningful representation learning.
• Multi-level attention visualization confirms that the model quantifies contributions of specific diseases and admissions to predictions, with higher attention weights aligning with clinically plausible causal pathways.
• The self-supervised proxy task enables full utilization of EHR data, including single-admission and final admissions, by generating synthetic labels for pre-training.
• Case studies demonstrate that Sherbet provides both generic interpretability (e.g., disease comorbidity patterns) and personalized interpretability (e.g., contribution of prior admissions to future diagnoses), enhancing clinical trust.