Heterogeneous Graph Embeddings of Electronic Health Records Improve Critical Care Disease Predictions

Abstract

Electronic Health Record (EHR) data is a rich source for powerful biomedical discovery but it consists of a wide variety of data types that are traditionally difficult to model. Furthermore, many machine learning frameworks that utilize these data for predictive tasks do not fully leverage the inter-connectivity structure and therefore may not be fully optimized. In this work, we propose a relational, deep heterogeneous network learning method that operates on EHR data and addresses these limitations. In this model, we used three different node types: patient, lab, and diagnosis. We show that relational graph learning naturally encodes structured relationships in the EHR and outperforms traditional multilayer perceptron models in the prediction of thousands of diseases. We evaluated our model on EHR data derived from MIMIC-III, a public critical care data set, and show that our model has improved prediction of numerous disease diagnoses.

A Appendix

Definition 1 (Heterogeneous Network)

A heterogeneous network is defined as a graph $\mathbf{G} =(V,E,T)$, where each node $\mathbf{v}$ and each link $\mathbf{e}$ are represented by their mapping functions to a specific node and relation type $\phi (v):V\rightarrow T_{V}$ and $\phi (e):E\rightarrow T_{E}$. Where $T_{V}$ and $T_{E}$ denote the sets of node and relation types, and $|T_{V}|+|T_{E}|>2$.

Definition 2 (Heterogeneous Graph Learning)

Given a heterogeneous network $\mathbf{G}$, the task of heterogeneous graph learning is to learn a function mapping $f:V\rightarrow {R^{d}}$, that connects disparate type of nodes into a $d-dimensional$ uniform latent representation $X\in R^{|V|\times d}$, and $d\ll |V|$, that are able to capture the structural and semantic relations between them.

Definition 3 (One-hop Connectivity)

One-hop connectivity in a heterogeneous network is the local pairwise connection between two consecutive vertices, which directly linked by an edge belongs to a relational type.

1.1 A.1 Skip-Gram Model

The skip-gram model  [11] seeks to maximize the probability of observing the context neighborhood nodes given the center node:

$$\begin{aligned} \mathrm {max}_{f}\sum _{u\in V}logPr(N_c(u)|f(u)) \end{aligned}$$

(5)

Where $N_c(u)$ is the neighborhood context nodes of the center node u, and f(u) is the latent representation of u.

1.2 A.2 Heterogeneous Skip-Gram Model

EHR data is heterogeneous, including varies type of vertices, such as lab tests, diagnoses, prescriptions, and patient demographics. Each of these vertices encodes different information. Heterogeneous Skip-gram model  [5] learns the latent expression of these different type of nodes by maximizing the probability of observing heterogeneous neighborhood given a center node:

$$\begin{aligned} \mathrm {max}\sum _{u\in V}\sum _{t\in T_{V}}logPr(N_{t}(u)|f(u)) \end{aligned}$$

(6)

Where $N_{t}(u)$ is the heterogeneous neighborhood vertices of center node u, and $t\in T_{V}$ is the node type.

1.3 A.3 TransE

TransE model  [1] aims to relate different type of nodes by their relationship type. Specifically, two different types of nodes are connected by a relation type would be represented as a triple (head, relation, tail), denoted as (h, l, t). For example, one triple from EHR data could be (patient, diagnosed, ICD), where patient is the head node, ICD is the specific diagnosis code attributed to the patient, and the relation between these two vertices is diagnosed.

This TransE model leverages the procedure by first projecting different type of node with different initial representation dimension into a same latent dimension space (where the dimension of this latent space can be customized), and these two different type projected nodes are linked by a relation type which is represented as a translation vector in that latent space. Both the projection matrix and the relational translation vector are learnable parameters in the deep learning system.