Collaborative Graph Learning with Auxiliary Text for Temporal Event Prediction in Healthcare

1 Introduction

Electronic health records (EHR) consist of patients' temporal visit information in health facilities, such as medical history and doctors' diagnoses. The usage and analysis of EHR not only improves the quality and efficiency of in-hospital patient care but also provides valuable data sources for researchers to predict health events, including diagnoses, medications, and mortality rates, etc. A key research problem is improving prediction performance by learning better representations of patients and diseases so that improved risk control and treatments can be provided. There have been many works on this problem using deep learning models, such as recurrent neural networks (RNN) [Choi et al., 2016a], convolutional neural networks (CNN) [Nguyen et al., 2017], and attention-based mechanisms [Ma et al., 2017]. However, several challenges remain in utilizing EHR data and interpreting models:

1. Effectively utilizing the domain knowledge of diseases. Recently, graph structures are being adopted [Choi et al., 2017] using disease hierarchies, where diseases are classified into various types at different levels. For example, Figure 1 shows a classification of two forms of hypertension and one form of heart failure. One problem is that existing works [Choi et al., 2017; Shang et al., 2019] only consider the vertical relationship between a disease and its ancestors (hierarchical link). However, they ignore horizontal disease links that can reflect disease complications and help to predict future diagnoses.
2. Collaboratively learning patient-disease interactions. Patients with the same diagnoses may have other similar diseases (patient diagnosis in Figure 1). Existing approaches [Choi et al., 2017; Ma et al., 2017] treat patients as independent samples by using diagnoses to represent patients, but they fail to capture patient similarities, which help in predicting new-onset diseases from other patients.
3. Incorporating unstructured text. Unstructured data in EHR including clinical notes contain indicative features such as physical conditions and medical history. For example, a note: "The patient was intubated for respiratory distress and increased work of breathing. He was also hypertensive with systolic in the 70s" indicates that this patient has a history of respiratory problems and hypertension. Most models [Choi et al., 2016b; Bai et al., 2018] do not fully utilize such data. This often leads to unsatisfactory prediction performance and lack of interpretability.

To address these problems, we first present a hierarchical embedding method for diseases to utilize medical domain knowledge. Then, we design a collaborative graph neural network to learn hidden representations from two graphs: a patient-disease observation graph and a disease ontology graph. In the observation graph, if a patient is diagnosed with a disease, we create an edge between this patient and the disease. The ontology graph uses weighted ontology edges to describe horizontal disease interactions. Moreover, to learn the contributions of keywords for predictions, we design a TF-IDF-rectified attention mechanism for clinical notes which takes visit temporal features as context information. Finally, combining disease and text features, the proposed model is evaluated on two tasks: predicting patients' future diagnoses and heart failure events. The main contributions of this work are summarized as follows:

• We propose to collaboratively learn the representations of patients and diseases on the observation and ontology graphs. We also utilize the hierarchical structure of medical domain knowledge and introduce an ontology weight to capture hidden disease correlations.
• We integrate structured information of patients' previous diagnoses and unstructured information of clinical notes with a TF-IDF-rectified attention method. It allows us to regulate attention scores without any manual intervention and alleviates the issue of using attention as a tool to audit a model [Jain and Wallace, 2019].
• We conduct extensive experiments and illustrate that the proposed model outperforms state-of-the-art models for prediction tasks on MIMIC-III dataset. We also provide detailed analysis for model predictions.

2 RelatedWork

Deep learning models, especially RNN models, have been applied to predict health events and learn representations of medical concepts. DoctorAI [Choi et al., 2016a] uses RNN to predict diagnoses in patients' next visits and the time duration between patients' current and next visits. RETAIN [Choi et al., 2016b] improves the prediction accuracy through a sophisticated attention process on RNN. Dipole [Ma et al., 2017] uses a bi-directional RNN and attention to predict diagnoses of patients' next visits. Both Timeline [Bai et al., 2018] and ConCare [Ma et al., 2020b] utilize time-aware attention mechanisms in RNN for health event predictions. However, RNN-based models regard patients as independent samples and ignore relationships between diseases and patients which help to predict diagnoses for similar patients.

Recently, graph structures are adopted to explore medical domain knowledge and relations of medical concepts. GRAM [Choi et al., 2017] constructs a disease graph from medical knowledge. MiME [Choi et al., 2018] utilizes connections between diagnoses and treatments in each visit to construct a graph. GBERT [Shang et al., 2019] jointly learns two graph structures of diseases and medications to recommend medications. It uses a bi-directional transformer to learn visit embeddings. MedGCN [Mao et al., 2019] combines patients, visits, lab results, and medicines to construct a heterogeneous graph for medication recommendations. GCT [Choi et al., 2020] also builds graph structures of diagnoses, treatments, and lab results. However, these models only consider disease hierarchical structures, while neglecting disease horizontal links that reflect hidden disease complications. As a result, prediction performance is limited.

In addition, CNN and Autoencoder are also adopted to predict health events. DeepPatient [Miotto et al., 2016] uses an MLP as an autoencoder to rebuild features in EHR. Deepr [Nguyen et al., 2017] treats diagnoses in a visit as words to predict future risks such as readmissions in three months. AdaCare [Ma et al., 2020a] uses multi-scale dilated convolution to capture dynamic variations of biomarkers over time. However, these models neither consider medical domain knowledge nor explore patient similarities as discussed.

In this paper, we explore disease horizontal connections using a disease ontology graph. We collaboratively learn representations of both patients and diseases in their associated networks. We also design an attention regulation strategy on unstructured text features to provide quantified contributions of clinical notes and interpretations of prediction results.

3 Methodology

3.1 ProblemFormulation

An EHR dataset is a collection of patient visit records. Let C = {c1, c2, . . . , c|C|} be the entire set of diseases represented by medical codes in an EHR dataset, where |C| is the medical code number. Let N = {ω1, ω2, . . . , ω|N|} be the dictionary of clinical notes, where |N| is the word number.

EHR dataset. An EHR dataset is given by D = {γu|u ∈ U} where U is the collection of patients in D and γu = (V u 1 , V u 2 , . . . , V u T ) is a visit sequence of patient u. Each visit V u t = {C u t ,N u t } is recorded with a subset of medical codes C u t ⊂ C, and a paragraph of clinical notes N u t ⊂ N containing a sequence of |N u t | words.

Diagnosis prediction. Given a patient u's previous T visits, this task predicts a binary vector ˆy ∈ {0, 1}|C| which represents the possible diagnoses in (T + 1)-th visit. ˆyi = 1 denotes ci is predicted in C u T+1 .

Heart failure prediction. Given a patient u's previous T visits, this task predicts a binary value ˆ y ∈ {0, 1}. ˆ y = 1 denotes that u is predicted with heart failure 1 in (T +1)-th visit.

In the rest of this paper, we drop the superscript u in V u t ,C u t , and N u t for convenience unless otherwise stated.

3.2 The ProposedModel

In this section, we propose a Collaborative Graph Learning model, CGL. An overview of CGL is shown in Figure 2.

Hierarchical Embedding for Medical Codes

ICD-9-CM is an official system of assigning codes to diseases. It hierarchically classifies medical codes into different types of diseases in K levels. This forms a tree structure where each node has only one parent. Note that most medical codes in patients' visits from EHR data are leaf nodes. However, a patient can be diagnosed with a higher level disease, i.e., non-leaf node. Therefore, we recursively create virtual child nodes for each non-leaf node to pad them into virtual leaf nodes. We assume there are nk nodes at each level k (smaller k means higher level in the hierarchical structure).

1 The medical codes of heart failure start with 428 in ICD-9-CM.

We create an embedding tensor {Ek}k∈[1,2,...,K] for nodes in the tree. Ek ∈ Rnk×dc is the embedding matrix for nodes in level k, and dc is the embedding size. For a medical code ci as a leaf node, we first identify its ancestors in each level k = [1, 2, . . . ,K - 1] in the tree and select corresponding embedding vectors from {Ek}. Then, the hierarchical embedding ei ∈ RKdc of ci is calculated by concatenating the embeddings in each level: ei = Ei1 ⊕ Ei2 ⊕ , . . . ,⊕ EiK , where ⊕ denotes the concatenation. We use E ∈ R|C|×Kdc to represent medical codes after hierarchical embedding.

Graph Representation

In visit records, specific diagnosis co-occurrences could reveal hidden similarities of patients and diseases. We explore such relationship by making the following hypotheses:

1. Diagnostic similarity of patients. If two patients get diagnosed with the same diseases, they tend to have diagnostic similarities and get similar diagnoses in the future.
2. Medical similarity of diseases. If two diseases belong to the same higher-level disease, they might have medical similarities such as symptoms, causes, and complications.

Based on these hypotheses, we construct a collaborative graph G = {GUC, GCC} for patients and medical codes. GUC is the patient-disease observation graph built from EHR data. Its nodes are patients and medical codes. We use a patientcode adjacency matrix AUC ∈ {0, 1}|U|×|C| to represent GUC. Given patient u, if u is diagnosed with a code ci in a previous visit, we add an edge (u, ci) and set AUC[u][i] = 1. GCC is the ontology graph. Its nodes are medical codes. To model horizontal links of two medical codes (leaf nodes), we create a code-code adjacency matrix A′CC ∈ N|C|×|C| . If two medical codes ci and cj have their lowest common ancestor in level k, we add an ontology edge (ci, cj)k and set A′CC [i][j] = k. This process is based on the idea that two medical codes with a common ancestor in lower levels of the hierarchial graph of ICD-9-CM should be similar diseases. Finally, we set A′CC [i][i] = 0 for all diagonal elements. Although A′CC can reflect the hierarchical structure of medical codes, it is a dense matrix and generates a nearly complete ontology graph, which will cause a high complexity for graph learning. We further propose a disease co-occurrence indicator matrix BCC initialized with all zeros. If two medical codes ci and cj appear in a patient's visit record, we set BCC[i][j] and BCC[j][i] as 1. Then, we let ACC = A′CC ⊙ BCC be a new adjacency matrix for GCC to neglect disease pairs in A′CC which never co-occur in EHR data. Here ⊙ denotes elementwise multiplication. Finally, we not only create a sparse ontology graph for computational efficiency, but also focus on more common and reasonable disease connections in the ontology graph.

Collaborative Graph Learning

To learn hidden features of medical codes and patients, we design a collaborative graph learning method on the fact and ontology graphs. Instead of calculating patient embeddings with medical codes like DeepPatient [Miotto et al., 2016], we assign each patient an initial embedding. P ∈ R|U|×dp is the embedding matrix of all patients with the size of dp. Let H(0) p = P,H(0) c = E and H(l) p ∈ R|U|×d(l) p ,H(l) c ∈ R|C|×d(l) c be the hidden features of patients and medical codes (i.e., inputs of l-th graph layer). We design a graph aggregation method to calculate the hidden features of patients and medical codes in the next layer. First, we map the medical code features H(l) c into the patient dimension and aggregate adjacent medical codes from the observation graph (AUC) for each patient:

HereWCU (l) ∈ Rd(l) c ×d(l) p maps code embeddings to patient embeddings. For the ontology graph, if ci, cj are connected in level k, we assign an ontology weight φj to cj when aggregating cj into ci:

Here σ is the sigmoid function. µj , θj ∈ R are trainable variables for cj. φj(k) is a monotonic function w.r.t. level k. This function enables the model to describe the horizontal influence of a disease on other diseases via assigning increasing or decreasing weights by levels. Let Φ ∈ R|C|×|C| be the ontology weight matrix and M,Θ ∈ R|C| be the collection of µ, θ. H(l) p is mapped into themedical code dimension and aggregated with adjacent patients from the observation graph:

Here WUC ∈ Rd(l) p ×d(l) c maps patient embeddings to code embeddings. Given that ACC stores the level where two diseases are connected, we use ACC to compute Φ. Finally, H(l) p and H(l) c of the next layer are calculated as follows:

where W(l) {p,c} maps Z(l) {p,c} to the (l+1)-th layer, and we use batch normalization to normalize features. In the L-th graph layers, we do not calculate H(L) p and only calculate H(L) c as the graph output, since the medical codes are required for further calculation. We let Hc = H(L) c ∈ R|C|×d(L) c be the final embedding for medical codes.

Temporal Learning for Visits

Given a patient u, we first compute a embedding vt for visit t:

After the collaborative graph learning, Hic contains the information of its multi-hop neighbor diseases by the connection of patient nodes. Hence, different fromGRAM, it enables the model to effectively predict diseases that have never been diagnosed on a patient before. We then employ GRU on vt to learn visit temporal features and get a hidden representation R = {r1, r2, . . . , rT } where the size of the RNN cell is h:

Then we apply a location-based attention [Luong et al., 2015] to calculate the final hidden representation ov of all visits:

where wα ∈ Rh is a context vector for attention and α is the attention weight for each visit.

Guiding Attention on Clinical Notes

We incorporate the clinical notes NT from the latest visit VT , since NT generally contains the medical history and future plan for a patient. We propose an attention regulation strategy that automatically highlights key words, considering traditional attention mechanisms in NLP have raised concerns as a tool to audit a model [Jain and Wallace, 2019; Serrano and Smith, 2019]. Pruthi et al. [Pruthi et al., 2020] present a manipulating strategy using a set of pre-defined impermissible tokens and penalizing the attention weights on these impermissible tokens. To implement the regulation strategy, we propose a TF-IDF-rectified attention method on clinical notes. Regarding all patients' notes as a corpus and each patient's note as a document, for a patient u, we first calculate the TF-IDF weight βi for each word ωi in u's note NT and normalize the weights into [0, 1]. Then, we select the embedding qi ∈ Rdw from a randomly initialized word embedding matrix Q ∈ R|N|×dw . For attention in Eq. (8), the context vector wα is randomly initialized, while clinical notes are correlated with diagnoses. Therefore, we adopt ov as the context vector. Firstly, we project word embeddings Q into the dimension of visits to multiply the context vector ov:

Then, let N be the embedding matrix selected from Q′ for words in NT , we calculate the attention weight α′ as well as the output on for clinical notes:

For a word with a high TF-IDF weight in a clinical note, we expect the model to focus on this word with a high attention weight. Therefore, we introduce a TF-IDF-rectified attention penalty L0 for the attention weights of words:

The attention weights that mismatch the TF-IDF weights will be penalized. We believe that irrelevant (impermissible) words such as "patient" and "doctor" tend to have low TFIDF weights. Finally, we concatenate on and ov as the output O ∈ R2h for patient u: O = ov ⊕ on.

Prediction and Inference

Diagnosis prediction is a multi-label classification task, while heart failure prediction is a binary classification task. We both use a dense layer with a sigmoid activation function on the model output O to calculate the predicted probability ˆy. The loss function of classification for both tasks is cross-entropy loss Lc. Then, we combine the TF-IDF-rectified penalty L0 and cross-entropy loss as the final loss L to train the model:

Here, y is the ground-truth label of medical codes or heart failure, and λ is a coefficient to adjust L0. In the inference phase, we freeze the trained model and retrieve the embeddings Hc of medical codes at the output of heterogeneous graph learning. Then, given a new patient for inference, we continue from Eq. (6) and make predictions.

4 Experiments

4.1 Experimental Setup

Dataset Description

We use the MIMIC-III dataset [Johnson et al., 2016] to evaluate CGL. Table 1 shows the basic statistics of MIMICIII. We select patients with multiple visits (# of visits ≥ 2) and select clinical notes except the type of "Discharge summary", since it has a strong indication to predictions and is unfair to be used as features. For each note, we use the first 50,000 words, while the rest are cut off for computational efficiency, given the average word number per note is less than 5,000. We split MIMIC-III randomly according to patients into training/validation/test sets with patient numbers as 6000/125/1000. We use the codes in patients' last visit as labels and other visits as features. For heart failure prediction, we set labels as 1 if patients are diagnosed with heart failure in the last visit. Finally, the observation graph is built based on the training set. A 5-level hierarchical structure and the ontology graph are built according to ICD-9-CM.

Evaluation Metrics

We adopt weighted F1 score (w-F1 [Bai et al., 2018]) and top k recall (R@k [Choi et al., 2016a]) for diagnosis predictions. w-F1 is a weighted sum of F1 for each class. R@k is the ratio of true positive numbers in top k predictions by the total number of positive samples, which measures the prediction performance on a subset of classes. For heart failure predictions, we use F1 and the area under the ROC curve (AUC), since it is a binary classification on imbalanced test data.

Baselines

To compare CGL with state-of-the-art models, we select the following models as baselines: 1) RNN-based models: RETAIN [Choi et al., 2016b], Dipole [Ma et al., 2017], Timeline [Bai et al., 2018]; 2) CNN-based models: Deepr [Nguyen et al., 2017]; 3) Graph-based models: GRAM [Choi et al., 2017], MedGCN [Mao et al., 2019]; and 4) A logistic regression model, LRnotes, on clinical notes using only TF-IDF features of each note (whose dimension is the dictionary size).

Deepr, GRAM, and Timeline use medical code embeddings as inputs, while others use multi-hot vectors of medical codes. We do not consider SMR [Wang et al., 2017] because 1) it does not compare with the above state-of-theart models and 2) it focuses on medication recommendation which is different from our tasks. We also do not compare with MiME [Choi et al., 2018] and GCT [Choi et al., 2020] because we do not use treatments and lab results in our data.

Parameters

We randomly initialize embeddings for diseases, patients, and clinical notes and select the paramters by a grid search. The embedding sizes dc, dp, and dw are 32, 16, and 16. The graph layer number L is 2. The hidden dimensions d(1) p , d(1) c , and d(2) c are 32, 64, and 128, and the GRU unit h is set to 200. The coefficient λ in L0 for diagnosis and heart failure prediction is 0.3 and 0.1. We set the learning rate as 10-3 , optimizer as Adam, and use 200 epochs for training. The source code of CGL is released at https://github.com/LuChang-CS/CGL.

4.2 Experimental Results

Diagnosis and Heart Failure Prediction

Table 2 shows the results of baselines and CGL on diagnosis prediction. We use k = [20, 40] for R@k. Each model is trained for 5 times with different variable initializations. The mean and standard deviation are reported. The proposed CGL model outperforms all the baselines. We think this is mostly because CGL captures hidden connections of patients and diseases and utilizes clinical notes. In addition, the results of LRnotes indicate that only using clinical notes does not improve performance in predicting diagnosis. Table 3 shows the heart failure prediction results. We observe that CGL also achieves the best performance in terms of AUC and F1.

Ablation Study

To study the effectiveness of components, we also compare 3 CGL variants: CGL without hierarchical embedding (CGLh-), CGL without clinical notes as inputs (CGLn-), and CGL without ontology weights (CGLw-). The results are shown in Table 4. We observe that even without clinical notes, CGLn- with hierarchical embeddings and ontology weights still achieves the best performance among all other baselines. This indicates that domain knowledge including hierarchical embeddings and ontology weights also help to learn better representations of medical codes. In addition, from Table 4 we can infer that the complexity of CGL mostly comes from modeling clinical notes, i.e., word embeddings. Therefore, CGL is scalable and can be generalized to other tasks when clinical notes are not accessible.

Prediction Analysis

New-onset diseases. For a patient, new-onset diseases denote new diseases in future visits which have not occurred in previous visits of this patient. We use the ability of predicting new-onset diseases to measure learned diagnostic similarity of patients. It is natural for a model to predict diseases that have occurred in previous visits. With the help of other similar patients' records, the model should be able to predict new diseases for a patient. The idea is similar to collaborative filtering in recommender systems. If two patients are similar, one of them may be diagnosed with new-onset diseases which have occurred in the other patient. We also use R@k (k = [20, 40]) to evaluate the ability of predicting occurred and new-onset diseases. Here, R@k denotes the ratio between the number of correctly predicted occurred (or new)

diseases and the number of ground-truth diseases. We select GRAM and MedGCN which have good performance in diagnosis prediction, and CGLn- without clinical notes, because we want to explore the effectiveness of the proposed observation and ontology graphs. Table 5 shows the results of R@k on test data. We can see that CGLn- has similar results to GRAM on occurred diseases while achieving superior performance on new-onset diseases. This verifies that our proposed collaborative graph learning is able to learn from similar patients and predict new-onset diseases in the future.

Disease embeddings. To show the similarity of diseases, we plot the learned 4795 code embeddings Hc using t-SNE [Maaten and Hinton, 2008]. Figure 3 shows the embeddings learned by GRAM, Timeline, and CGL in 3 levels. Colors denotes different disease types in each level. In Figure 3, disease embeddings learned by GRAM and CGL are basically clustered according to their real categories, while Timeline seems like a random distribution. In the plot of GRAM, we observe the clusters are far away from each other given large inter-cluster distances, while nodes in a cluster are close to each other due to small intra-cluster distances. We can observe that the embeddings learned by GRAM do not capture distinguishable features of low-level diseases as well as the relationships across clusters. Therefore, we can infer that learning proper representations that reflect disease hierarchical structures and correlations is helpful for predictions.

Contribution of notes. We compare the proposed TF-IDF rectified attention weights with regular attention weights to verify if the model focuses on important words. Table 6 demonstrates an example with a part of a note and predicted diagnoses. In this example, the patient is diagnosed with 33 diseases, and CGL predicts 10 of them correctly in top 20 predicted codes. Important words with high α′ values are highlighted in pink. We first observe that pink words are relevant to diagnoses. In addition, we notice the rectified attention weights are more semantically interpretable. For example, "acute" and "HCAP" (Health care-associated pneumonia) get higher weights with the rectified attention loss. Meanwhile, we show the unimportant words with low α′ values in gray. We observe that our model detects unimportant words which have less contributions. For example, "patient" and "diagnosis" are regarded as an unimportant word but not captured in the regular attention mechanism. Therefore, we may conclude that the TF-IDF-rectified attention method improves the accuracy of interpretations using clinical notes.

5 Conclusion

In this paper, we propose CGL, a collaborative graph learning model to jointly learn the representations of patients and diseases, and effectively utilize clinical notes in EHR data. We conducted experiments on real-world EHR data to demonstrate the effectiveness of the learned representations and performance improvements of CGL over state-of-the-art models. We also provide analysis of CGL on multiple aspects, including new onset diseases, disease embeddings, and contribution of clinical notes. In the future, we plan to explore methods to quantify the contributions of certain admissions to each predicted medical code. Usage of single admission records in EHR data will also be considered for further investigation.

Acknowledgments

This work was supported in part by US National Science Foundation under grants 1838730 and 1948432. SK was supported in part by the NLM of the NIH under Award Number R01LM013308.