KGDAL: Knowledge Graph Guided Double Attention LSTM for Rolling Mortality Prediction for AKI-D Patients

Abstract

With the rapid accumulation of electronic health record (EHR) data, deep learning (DL) models have exhibited promising performance on patient risk prediction. Recent advances have also demonstrated the effectiveness of knowledge graphs (KG) in providing valuable prior knowledge for further improving DL model performance. However, it is still unclear how KG can be utilized to encode high-order relations among clinical concepts and how DL models can make full use of the encoded concept relations to solve real-world healthcare problems and to interpret the outcomes. We propose a novel knowledge graph guided double attention LSTM model named KGDAL for rolling mortality prediction for critically ill patients with acute kidney injury requiring dialysis (AKI-D). KGDAL constructs a KG-based two-dimension attention in both time and feature spaces. In the experiment with two large healthcare datasets, we compared KGDAL with a variety of rolling mortality prediction models and conducted an ablation study to test the effectiveness, efficacy, and contribution of different attention mechanisms. The results showed that KGDAL clearly outperformed all the compared models. Also, KGDAL-derived patient risk trajectories may assist healthcare providers to make timely decisions and actions. The source code, sample data, and manual of KGDAL are available at https://github.com/lucasliu0928/KGDAL.

1. INTRODUCTION

Acute kidney injury (AKI) is a common complication of hospitalized patients and the incidence increase in patients admitted to the intensive care unit (ICU) [33, 34]. AKI that results in the need for dialysis (AKI-D) is associated with a high risk of hospital mortality [13], and for survivors a risk of incident or progressive chronic kidney disease (CKD) [5, 6, 12, 35], cardiovascular disease [11, 26, 39] or end-stage renal disease (ESRD) [15, 16, 23]. By identifying mortality risk factors from patient individual and population data, providers can implement early intervention strategies leading to better health care and substantially reducing the cost of care.

Numerous factors may influence in-hospital mortality including acute anemia, respiratory failure, electrolytes disarrangements, hemodynamic instability, and demographic information. There is a critical need to identify and correlate these patients and dialysis-specific parameters with inpatient mortality in this specific population. Moreover, accurate prediction of mortality over time (i.e., rolling prediction) in real-world healthcare settings for critically ill patients with AKI-D is needed for better utilization of hospital resources, such as intensifying therapies when is needed, or transitioning patients with a high risk of mortality to comfort care [25, 30].

Two general approaches have been used for mortality predictions in the ICU. The first approach uses clinical scores, including Acute Physiology and Chronic Health Evaluation (APACHE) and Sequential Organ Failure Assessment (SOFA), to identify at-risk patients at any time point [29]. The second approach employs machine learning (ML) methods, such as random forest [22] and SVM [37], to predict mortality risks using electronic health record (EHR) data. With the rapid development of ML techniques, ML-based mortality prediction attracts much attention recently. Nevertheless, ML-based methods mainly focus on mortality prediction at the end of the treatment [21], and the clinical needs for rolling mortality prediction is often overlooked [17, 19]. Traditional ML models rely heavily on feature engineering, requiring not only a deep understanding of the domain knowledge but also tremendous efforts on manual feature extraction and model tuning [32].

In recent years, deep learning (DL) models, including Transformer [36], Long-Short Term Memory (LSTM) [14] and gated recurrent neural networks [8], have shown promising performance on end-to-end patient risk prediction using large-scale EHR data [31]. Doctor AI [9] applies RNN on visit-based medical codes to predict the diagnosis and medication categories for the subsequent visit. However, critical long-term dependency could be lost during RNN training [36]. To address the long-term dependency problem, an attention mechanism was introduced [1] on sequence-to-sequence modeling tasks, where an additive attention mechanism is used in an encoder-decoder model to learn where to pay attention in the input sequence. In the following years, multiple types of attention mechanisms have been proposed [4, 24, 36]. Zhou et al. proposed an attention mechanism along the time-step dimension for relation classification [41], and Transformer [36] was developed to learn attentions without using recurrent network units.

Applications of ML models in clinical settings require a certain level of interpretability. To this end, medical ontologies have been incorporated into DL models using graph-based attention mechanisms. GRAM [10] and DG-RNN [40] are graph-based attention models that are designed to obtain attention from medical knowledge graph (KG) and then use the attention to adjust vector representations of medical codes. However, since attention is purely obtained from KG, the concept completeness and correctness are not considered. In fact, KG is rarely completed and does not have 100% correctness [3]. Furthermore, the graph-based attention model is only feasible on discrete data (e.g., diagnosis codes) but not on continuous temporal data (e.g., blood pressure measurements).

In this study, we propose a Knowledge-Graph Guided Double Attention LSTM (KGDAL) model, aiming to make precise rolling mortality predictions in a real-world healthcare setting for critically ill patients with AKI-D. To our knowledge, KGDAL is the first KG-guided model that extracts both time and feature attention on continuous temporal data. KGDAL has the following advantages:

• KGDAL obtains two-dimensional attention in both the time and feature spaces for improved prediction power and enhanced model interpretability.

• The attention mechanism in the feature space is automatically derived based on the KG rather than manual curation.

• KGDAL can model both continuous and discrete temporal EHR data types.

• KGDAL can make precise rolling mortality predictions for AKI-D patients on two independent clinical datasets.

2. METHOD

The goal of this study is to conduct rolling mortality prediction to assist clinicians in making timely decisions and actions [21]. Mathematically, for each particular subsequence of a patient’s EHR data, KGDAL predicts the patient’s outcome in the next K hours, where K is the time granularity (e.g., 72 hours). The outcome could be mortality or survival. KGDAL’s overall architecture, as shown in Figure 1, contains three phases:

Figure 1:

The three phases of the Knowledge-Graph Guided Double Attention (KGDAL) model for rolling mortality prediction for critically ill patients with AKI-D in real-world healthcare settings.

Phase 1. EHR Data Extraction.

Patient clinical features are extracted from patient’s EHR. Each feature is matched to a set of concepts in a Knowledge-Graph (KG), followed by manual validation by clinicians. Patient subsequences with random starting time points and varying lengths are generated. The label of each subsequence is “mortality” or “survival” in the next K hours. Static features of each patient are consistent for all the subsequences of the same patient.

Phase 2. Knowledge-Graph Extraction.

The entire KG is used to learn the concept embeddings for every concept identified in Phase 1. The concept embeddings are grouped based on the KG’s hierarchical structure, resulting in multiple concept embedding groups. Subsequently, the KG-embedding distances (including both pairwise distances between concept embedding groups and the distances from a concept-of-interest to every concept embedding group) are computed.

Phase 3. Knowledge-Graph Guided Double Attention.

All the temporal features are grouped based on their corresponding concept embedding groups. An LSTM is assigned to each temporal feature group. All the LSTM models are trained simultaneously to minimize the overall loss. Both feature and time attentions are learned using fully connected layers followed by softmax. Double attention is formed using both feature and time attention. The double attention is adopted to adjust feature embeddings for the final prediction and to regularize the prediction loss that minimizes the discrepancy between the attention-based distance and the KG-embedding distance.

2.1. Phase 1. Data Extraction

Given a patient’s EHR data $\left\{S_1,S_2,\dots ,S_t,\dots ,S_{N_t}\right\}$, where $S_t\in {ℝ}^{N_f}$ is a set of features (e.g. clinical measurements) at time point t ∈ {1, 2, … , N_t}, N_f is the total number of features, and N_t is the total number of time points. For each patient, we generate N_s subsequences with randomly starting time and random length between 24 and 72 hours (see Figure 1 Phase 1). The outcome label of each subsequence is whether the patient dies or is alive in the next K hours.

Patient clinical features are matched to a set of corresponding concepts in a KG. For features that can be directly matched to a concept in a KG (e.g., a diagnosis code), the corresponding concept will be used. Otherwise, all the associated concepts in the same KG are extracted and then filtered by clinicians. For example, the corresponding concepts of “systolic blood pressure” in HPO [20] are “elevated systolic blood pressure (HP:0004421)” and “decreased systolic blood pressure (HP:0500105)”, both of which are valid and kept for later use.

2.2. Phase 2. Knowledge-Graph Extraction

In phase 2, KGDAL generates concept embeddings and computes the group-wise KG-embedding distances in a KG.

2.2.1. Concept Embedding.

Let E_c be the set of concepts and E_r be all possible relationship types in a KG, a concept relation O can be represented using a triplet denoted as (C_h, r_i, C_t), where C_h, C_t ∈ E_c are the head and tail concepts respectively, and r_i ∈ E_r represents the relationship from C_h to C_t. For example, triplet (“Diabetes mellitus type 1”, “is-a”, “diabetes mellitus”) in SNOMED-CT represents that the Diabetes mellitus type 1 is a subtype of diabetes mellitus.

All the triplets in a KG are used to learn the concept embeddings. The basic idea is to make the learned embeddings of tail concept C_t be close to the sum of the embeddings of head concept C_h and the embeddings of relation r_i. Here, we use TransE [2], one of the most representative translational distance model, to learn the concept embeddings by formulating the problem as follows: given a concept triplet O = (C_h, r_i, C_t) in a KG, we learn the embedding triplets denoted as G = (h, l, t), where h, $t\in {ℝ}^{d_e}$ represents the head and tail concepts embeddings respectively, $1\in {ℝ}^{d_l}$ is the relation vector between h and t. TransE is trained with negative sampling to learn the embeddings that minimize a margin-based ranking loss function:

$$L_{KG}=\sum _{(h,l,t)\in O}\sum _{\left(h^{\prime },l,t^{\prime }\right)\in O^{\prime }}\text{max}\left(0,\gamma +d(h+l,t)-d\left(h^{\prime }+l,t^{\prime }\right)\right)$$ (1)

where O represents the positive samples and O′ represents the negative samples that were randomly generated by replacing the head or tail concepts of positive samples. d represents any distance metrics, γ > 0 is a margin hyper-parameter. Equation 1 shows that the distances in positive samples are minimized where the distances in negative samples are maximized. In this step, the concept embedding denoted as E_concept are obtained for each concept associated with the patient clinical features.

2.2.2. Embedding Grouping.

The hierarchical level in a KG represents classes of concepts holding similar characteristics. To capture the commonality within a class and the difference between classes, embedding group E is formed by taking the sum of concept embeddings E_concept in each group based on the hierarchical structure of a KG. The number of groups N_g is determined by the number of classes in the user-specified KG level. In general, using a higher level will form more general concept groups, and using a lower level will form more specific concept groups.

2.2.3. Embedding Group-wise Distance.

In order to measure the difference between concept embedding groups as well as how much each group are related to the outcome (i.e., mortality) and to let a DL model pay more attention to the closely related concept embedding groups, two KG-embedding distances are computed. The pairwise distance between concept embedding groups is computed using Eq 2.

$$K{G}_{\mathit{\text{dist}}}(i,j)=\mathit{\text{dist}}\left(E_i,E_j\right)$$ (2)

where E represents a set of concept embedding groups, the subscript represents the group indexes, and dist denotes a distance metric, such as the Euclidean distance.

The distance between each concept embedding group and the concept-of-interest are computed using Eq 3.

$$K{G}_{-}\mathit{\text{Targe}}{t}_{\mathit{\text{dist}}}\left(E_i\right)=\mathit{\text{dist}}\left(E_i,E_{\mathit{\text{target}}}\right)$$ (3)

where KG_Target_dist (i) represents the distance from the target embedding E_target to the ith concept embedding group.

2.3. Phase 3. Knowledge-Graph Guided Double Attention

2.3.1. LSTM Embedding.

By assigning each temporal feature to its corresponding concept embedding group in a KG, a patient’s subsequence can be denoted as $\left\{{\text{x}}_1^{\text{i}},{\text{x}}_2^{\text{i}},\dots ,{\text{x}}_{\text{t}}^{\text{i}}\right\}$, where ${\text{x}}_{\text{t}}^{\text{i}}\in {ℝ}^{n_i}$ is a list of feature values in the corresponding concept embedding group i at time step t, where i is the group index (0 ≤ i < N_g), t ∈ {1, 2, … ,T} is the time step, and n_i is the number of features at each time step for the ith group.

KGDAL consists of N_g LSTM [14] models, each for a feature group. For simplicity, we assume all the feature groups have the equal number of features, the subscript i of n is omitted in the following text. LSTM has three gates, namely the forget gate f_t, the input gate i_t, and the output gate o_t, where f_t, i_t, $o_t\in {ℝ}^m$, and m is the dimension of the hidden vectors. Using c_t and h_t to represent the cell state vector and the hidden state vector, (c_t, $h_t\in {ℝ}^m$), the updated LSTM cell in KGDAL can be represented as follows:

$$f_t=\sigma \left(W_f{h}_{t-1}+U_f{x}_t+b_f\right)$$ (4)

$$i_t=\sigma \left(W_i{h}_{t-1}+U_i{x}_t+b_i\right)$$ (5)

$$o_t=\sigma \left(W_o{h}_{t-1}+U_o{x}_t+b_o\right)$$ (6)

$${\tilde{c}}_t=\text{tanh}\left(W_c{h}_{t-1}+U_c{x}_t+b_c\right)$$ (7)

$$c_t=f_t\odot c_{t-1}+i_t\odot {\tilde{c}}_t$$ (8)

$$h_t=o_t\odot \text{tanh}\left(c_t\right)$$ (9)

where W_f, W_i, W_o, $W_c\in {ℝ}^{m\times m}$, U_f, U_i, U_o, $U_c\in {ℝ}^{m\times n}$, b_f, b_i, b_o, $b_c\in {ℝ}^m$ are learnable parameters, σ is a sigmoid function, and ⊙ is the Hadamard product.

As a result, each LSTM layer outputs a hidden state matrix for each feature group denoted as $H^{(i)}=\left[h_1^{(i)},h_2^{(i)},\dots ,h_T^{(i)}\right]$, where i is the group index. The hidden state metrics for each feature group is called the feature embedding matrix. Then the feature embedding matrix learned using multiple LSTM layers can be denoted as $\left\{H^{(0)},H^{(2)},\dots ,H^{\left(N_g-1\right)}\right\}$, where $H^{(i)}\in {ℝ}^{m_i\times T}$ represents the feature embedding matrix of group i, and m_i is the dimension of hidden state vector from the ith LSTM for group i.

2.3.2. KG-Guided Double Attention.

To model the latent dependencies between different feature groups and at different time steps, KGDAL learns the attentions in both time and feature spaces guided by a KG. The detailed architecture of KGDAL is in Figure 2.

Figure 2:

The detailed architecture of the Knowledge-Graph Guided Double Attention LSTM (KGDAL) model.

Time Attention.

All the N_g feature embedding matrices are concatenated into one matrix denoted as H^C with the dimension of ${ℝ}^{\left(m_1+m_2+\cdots +m_{N_g}\right)\times T}$. For simplicity, we assume all LSTM layers have equal dimensions of hidden vectors. Hence, the superscript or subscript i of m is omitted in the following text. and the dimension of H^C is now ${ℝ}^{\left(N_{g}m\right)\times T}$. The time attention is computed as follows:

$$M_{\alpha }=\text{tanh}\left(H^C\right)$$ (10)

$$\mathit{\alpha }=\mathit{\text{softmax}}\left(M_{\alpha }^T{w}_{\alpha }\right)$$ (11)

where $\mathit{\alpha }\in {ℝ}^T$ is the time attention, $w_{\alpha }\in {ℝ}^{N_{g}m}$ is the learnable parameter, and $M_{\alpha }^T$ is the transpose of $M_{\alpha }\in {ℝ}^{\left(N_{g}m\right)\times T}$.

The time attention mechanism is inspired from and is similar to the work by Zhou et al. [41], but there are two key differences: 1) the input to LSTM is the grouped feature sequences and 2) the time attention mechanism is applied on the output of multiple LSTMs.

Feature Attention.

A similar attention mechanism is used to compute the feature attentions. However, there are two changes. Firstly, we transpose H^C and use it as the input so that the attention mechanism will be applied on the feature space instead of the time space. Secondly, the KG-embedding distances (KG_Target_dist) between the concept embedding groups and the concept-of-interest are used to weight the raw feature attentions. Mathematically, the raw feature attentions is computed using:

$$M_{\beta }=\text{tanh}\left({\left(H^C\right)}^T\right)$$ (12)

$$\mathit{\beta }=\mathit{\text{softmax}}\left(M_{\beta }^T{w}_{\beta }\right)$$ (13)

where $\beta \in {ℝ}^{N_{g}m}$ is the raw feature attention, $w_{\beta }\in {ℝ}^T$ is the learnable parameter, and $M_{\beta }^T$ is the transpose of $M_{\beta }\in {ℝ}^{(T)\times N_{g}m}$.

Then the KG-adjusted feature attention ϵ is equals to β weighted by the KG-embedding distances at corresponding position, which is computed using:

$${\mathit{ϵ}}_{(pos,i)}={\mathit{\beta }}_{(pos,i)}\otimes K{G}_{-}\mathit{\text{Targe}}{t}_{\mathit{\text{dist}}}(i)$$ (14)

where $\mathit{ϵ}\in {ℝ}^{N_{g}m}$, ⊗ represents the outer product. pos represents the corresponding positions of feature attentions for each group, where pos = [im : (i + 1)m − 1], i is the group index, m is the dimension of the hidden vectors. For example, if the first m raw feature attentions β are obtained for the first feature group, then these m values are weighted by the distance from the first KG-embedding group to the target embedding.

2.3.3. Double Attended Representations.

We combine the attention on both the time and feature spaces by taking the outer product of the time attentions α and a scaled version of KG-adjusted feature attention ϵ to obtain the double attention. In addition, the feature embeddings are then adjusted by the obtained double attention. The double attention Θ and the adjusted feature embeddings R are computed as follows:

$$\mathit{\Theta }={\mathit{ϵ}}_S\otimes \mathit{\alpha }$$ (15)

$$R=H^C\odot \mathit{\Theta }$$ (16)

where Θ and $R\in {ℝ}^{N_{g}m\times T}.{\mathit{ϵ}}_s\in {ℝ}^{N_{g}m}$ is the scaled feature attention which is computed by taking the ratio between each KG-adjusted feature attention to the first KG-adjusted feature attention ϵ₀ so that the proportion of attentions are maintained, as shown in Equation 17.

$${\mathit{ϵ}}_s=\mathit{ϵ}/{\mathit{ϵ}}_0$$ (17)

Finally, the double-attention adjusted feature embeddings R are passed into the pooling layer for taking the sum/max over each time step, followed with a dense layer with sigmoid function for final predictions.

2.3.4. Loss Function.

To consolidate the concept relations from KG in attentions, a regularization term is added to the original prediction loss function. The new regularization term minimizes the discrepancy between the pairwise attention-based distances and the pairwise KG-embedding distances.

Let the ground truth label be y and the predicted label be $\widehat{y}$, we use the cross-entropy for the original prediction loss denoted as L_pred, and the regularization term is denoted as L_reg, the final loss L is defined as:

$$L=L_{\mathit{\text{pred}}}+L_{reg}$$ (18)

$$L_{\mathit{\text{pred}}}=\sum _{k=1}^{N_s}-\left(y_k\text{log}\left(\widehat{y_k}\right)+\left(1-y_k\right)\text{log}\left(1-\widehat{y_k}\right)\right)$$ (19)

$$L_{reg}=\sum _{i=1}^{N_g-1}\sum _{j=i+1}^{N_g}\left(\mathit{\text{dist}}\left({\mathit{\Theta }}_i,{\mathit{\Theta }}_j\right)-K{G}_{\mathit{\text{dist}}}(i,j)\right)$$ (20)

where Θ represents the double attentions, KG_dist is the KG-embedding distance discussed in 2.2.3, and i, j are the group indexes, N_s is the number of samples, k is the sample index.

3. EXPERIMENTS

The performance of KGDAL was evaluated using two (proprietary and public) AKI-D datasets.

3.1. Data Preprocessing

3.1.1. Proprietary EHR Data.

The proprietary EHR data include 608 AKI-D patients who were admitted to the University of Kentucky (UK) HealthCare from January 2009 to October 2019. Among them, 247 (41%) died in the hospital and 361 (59%) survived. This cohort excluded patients who were less than 18 years old, or were diagnosed with end-stage kidney disease (ESKD) at the time of index hospital admission, or were recipients of kidney transplant. The EHR records during renal replacement therapy (RRT), including both haemodialysis (HD) and continuous renal replacement therapy (CRRT), were extracted. The duration of RRT was limited from 72 hours to 2,000 hours. Any records beyond this range were excluded.

Twelve types of temporal features were collected from EHR, which were systolic blood pressure, diastolic blood pressure, creatinine, bicarbonate, hematocrit, potassium, bilirubin, sodium, temperature, white blood cells (WBC) count, heart rate, and respiratory rate. Six types of static features were also collected, which were demographics (age, race, and gender), admission weight, body mass index (BMI), and Charlson comorbidity score. Three status flags were constructed, which indicated the on or off of CRRT or HD, and the status of being in the ICU or not. In total, 21 features were included in the UK data and the average missing rate of the temporal features was 58.7%.

For each feature, outlier values greater than 97.5 percentile or below 2.5 percentile were both excluded. The temporal granularity of the temporal features was set to one value (median) per hour. Linear interpolation was employed to fill the gaps between two actual measurements if needed. The only exception is creatinine, for which we only kept one value every six hours to maintain the in-practice frequency.

3.1.2. Public EHR Data.

The public data were extracted from the MIMIC-III [18]. We first identified all the AKI patients by the presence of ICD-9 codes of 584.5 to 584.9, then we identified AKI-D patients with the additional presence of ICD-9 procedure codes of 3995 as well as diagnosis codes of V45.11 and V561 [38]. Applying the same cohort exclusion criteria resulted in the MIMIC-III data with 170 AKI-D patients. Among them, 66 (39%) died in the hospital and 104 (61%) survived. The temporal features of MIMIC-III were the same as the UK data except for WBC and temperature, since neither of them was available in the RRT duration. The average missing rate of the temporal features was 49.4%. The same static features and status flags as those in the UK data were included in the MIMIC-III data. In total, 19 features were included in the MIMIC-III data. The same data extraction and outlier detection procedures were conducted.

3.1.3. Knowledge Map.

The Human phenotype ontology (HPO) was used as the clinical knowledge map to learn the concept embeddings. HPO, a widely used biomedical ontology, provides a standardized vocabulary of phenotypic abnormalities encountered in human disease [27]. Concepts in HPO are organized in hierarchies. Among the six sub-hierarchies on the top level, we focused on the “Phenotypic abnormality” sub-hierarchy since it includes most of the concepts of abnormalities related to the selected features in this experiment. To represent the strength of the relations between any two concepts, we counted the number of the common ancestors of the two concepts. Due to the large number of unique descendants in the “Phenotypic abnormality” sub-hierarchy (15,560), we computed the relations between every concept and “Acute Kidney Injury” (HP:0001919) and 100 random selected concepts. Finally, we obtained 1,566,363 identical concept triplets (see definition in Section 2.2) from HPO, where the number of common ancestors was considered as the relation strength between two concepts.

3.1.4. Experimental Data Generation.

From each patient’s EHR sequence during RRT, we randomly generated at most 30 subsequences with their lengths varying from 48 to 96 hours. For each subsequence in the UK data, the possible class labels are whether the patient died (positive) or survived (negative) 24, 48, or 72 hours after the end of that subsequence. For the MIMIC-III data, the possible class labels of each subsequence are whether the patient died (positive) or survived (negative) 48 or 72 hours after the end of that subsequence. Note that due to the small study cohort size, the “24 hours” label was absent from the MIMIC-III data. In summary, the UK data include 14,757, 15,468, and 16,660 subsequences labeled as negative (alive) and 3,455, 2,744, and 1,552 subsequences labeled as positive (mortality) in the next 72, 48 and 24 hours respectively. For the MIMIC-III data, 4,391 and 4,574 subsequences were labeled as negative (alive); 727 and 544 subsequences were labeled as positive (mortality) in the next 72 and 48, respectively. The ratio between positive and negative ranged between 10% and 25%.

Table 1 showed the data used for training, validation, and testing. All the subsequences of 50 randomly selected patients (25 died vs. 25 alive) were used for validation. All the rest patient data were randomly split into training (90%) and testing (10%). The data split was patient-wise so that the subsequences from the same patient only appeared in one of the three datasets.

Table 1:

Training, validation and testing data at the subsequence level.

	Proprietary EHR dataset (UK data)						Public EHR data (MIMIC-III)
	In next 72 hours		In next 48 hours		In next 24 hours		In next 72 hours		In next 48 hours
	Alive	Death	Alive	Death	Alive	Death	Alive	Death	Alive	Death
Train	12275	2643	12817	2101	13717	1201	2725	391	2831	285
Validation	1077	419	1147	349	1305	191	1281	211	1337	155
Test	1405	393	1504	294	1638	160	385	125	406	104

3.2. Experiment Results

3.2.1. Concepts Embedding and Concepts Grouping.

Figure 3 shows the HPO concepts matched to the temporal features in both patient cohorts. These concepts were clustered into four groups based on the structure of the “phenotypic abnormality” sub-hierarchy in HPO, which were “Abnormality of the cardiovascular system (Cardiovascular)”, “Abnormality of metabolism/homeostasis (Metabolism)”, “Abnormality of blood and blood-forming tissues (Blood)”, and “Abnormality of the respiratory system (Respiratory)”. Figure 3A shows the partial hierarchical structure of HPO where colors indicate concept groups.

Figure 3:

A: The partial hierarchical structure of the Human Phenotype Ontology (HPO) that includes the following concepts. Colors indicate different concept groups (Red: Acute Kidney injury (AKI); Orange: “Cardiovascular”; Green: “Metabolism”; Blue: “Blood”; Pink: “Respiratory”. B: The similarities of the same selected features in a projected space generated using t-SNE.

The concept embeddings for all the selected concepts were obtained by training the TransE model using all the HPO concept triplets. The resulting concept embeddings were used to compute two types of pairwise distances, i.e., distances between any two concept embedding groups and distances from the concept-of-interest “AKI” to each concept embedding group. The concept-wise distances were visualized using the t-SNE plot in Figure 3B. “Respiratory” is mostly related to “AKI” (averaged distance 0.02), while “Cardiovascular” is the most distant from “AKI” (averaged distance 0.98).

3.2.2. Performance Comparison.

We compared KGDAL with various mortality rolling prediction models on both the UK data and the MIMIC-III data. In addition, an ablation study was conducted to test whether KGDAL’s KG-adjusted feature attentions were critical in mortality rolling prediction by 1) only using the time attention mechanism, or 2) removing from the loss the KG adjustment that minimizes the discrepancy between the pairwise attention based distance and the pairwise KG-embedding distances. We compared all the models’ performance on both the balanced and imbalanced test sets, the ratio of positive samples (died) to negative samples (survived) are 1:1 and 1: 2, respectively. By randomly sampling five times for each case, we reported the averaged performance on all the evaluation metrics. All the compared models and their inputs are described as following:

• Random Forest: The input to this model is the un-grouped temporal features where each temporal feature at each time step is appended as a column, and each static feature is a column. Column-wise mean imputation is used to fill the missing values.

• Boosted Tree: We use Extreme Gradient Boosting (XG-Boost) [7] as the second baseline model. The input to this model is the same as Random Forest.

• LSTM: The input to a LSTM is the un-grouped temporal features and static features concatenated at each time point.

• Transformer: The input to a Transformer is the same as LSTM. We use the encoder part of the original transformer with a dense layer for the prediction task.

• KGDAL_α: This KGDAL model only uses the time attention mechanism. By removing feature attention, it explores the usefulness of the feature attention mechanism. The inputs are the grouped (KG-guided) features.

• KGDAL_αβ: This KGDAL model uses both time attention and feature attention, but it removes the KG-adjusted attention and the KG-adjusted loss from KGDAL. It explores the usefulness of the KG-adjusted attention mechanism. The inputs are the grouped (KG-guided) features.

The performance of mortality rolling prediction in the next 72, 48, and 24 hours on the UK data are listed in Table 2, Table 3, and Table 4 respectively. The performance of mortality rolling prediction in the next 72 and 48 hours on MIMIC-III data are listed in Table 5 and Table 6 respectively. In all the tables, precision (PREC), recall (REC), and F1 scores are for the positive (died) class.

Table 2:

Performance of morality prediction in the next 72 hours during RRT (UK data)

	Npos : Nneg = 1 : 1					Npos : Nneg = 1 : 2
	ROCAUC	ACC	PREC	REC	F1	ROCAUC	ACC	PREC	REC	F1
XGBoost	0.50	0.50	0.51	0.35	0.42	0.50	0.55	0.33	0.35	0.34
LSTM	0.62	0.61	0.60	0.67	0.63	0.63	0.60	0.44	0.67	0.53
Transformer	0.70	0.64	0.64	0.64	0.64	0.69	0.63	0.46	0.64	0.53
KGDALα	0.75	0.66	0.64	0.77	0.69	0.75	0.64	0.47	0.77	0.59
KGDALαβ	0.70	0.63	0.63	0.63	0.63	0.70	0.63	0.46	0.63	0.53
KGDAL	0.76	0.71	0.66	0.87	0.75	0.74	0.64	0.48	0.87	0.62

Table 3:

Performance of morality prediction in the next 48 hours during RRT (UK Data)

	Npos : Nneg = 1 : 1					Npos : Nneg = 1 : 2
	ROCAUC	ACC	PREC	REC	F1	ROCAUC	ACC	PREC	REC	F1
XGBoost	0.50	0.49	0.49	0.35	0.41	0.50	0.55	0.33	0.35	0.34
LSTM	0.61	0.60	0.59	0.67	0.63	0.62	0.59	0.43	0.67	0.52
Transformer	0.69	0.63	0.62	0.66	0.64	0.70	0.63	0.46	0.66	0.54
KGDALα	0.74	0.67	0.63	0.80	0.70	0.75	0.63	0.47	0.80	0.59
KGDALαβ	0.70	0.64	0.63	0.66	0.64	0.72	0.64	0.47	0.66	0.55
KGDAL	0.73	0.68	0.63	0.88	0.74	0.74	0.63	0.47	0.88	0.61

Table 4:

Performance of morality prediction in the next 24 hours during RRT (UK Data)

	Npos : Nneg = 1 : 1					Npos : Nneg = 1 : 2
	ROCAUC	ACC	PREC	REC	F1	ROCAUC	ACC	PREC	REC	F1
XGBoost	0.50	0.51	0.51	0.37	0.43	0.50	0.55	0.34	0.37	0.35
LSTM	0.62	0.62	0.61	0.69	0.65	0.61	0.58	0.42	0.69	0.53
Transformer	0.72	0.65	0.64	0.71	0.67	0.71	0.63	0.46	0.71	0.56
KGDALα	0.78	0.71	0.66	0.89	0.76	0.78	0.65	0.49	0.89	0.63
KGDALαβ	0.75	0.68	0.66	0.76	0.71	0.75	0.65	0.48	0.76	0.59
KGDAL	0.75	0.70	0.64	0.94	0.76	0.75	0.62	0.47	0.94	0.62

Table 5:

Performance of morality prediction in the next 72 hours during RRT (MIMIC-III)

	Npos : Nneg = 1 : 1					Npos : Nneg = 1 : 2
	ROCAUC	ACC	PREC	REC	F1	ROCAUC	ACC	PREC	REC	F1
XGBoost	0.51	0.49	0.47	0.15	0.23	0.50	0.59	0.28	0.15	0.20
LSTM	0.70	0.64	0.80	0.38	0.52	0.69	0.72	0.65	0.38	0.48
Transformer	0.64	0.60	0.69	0.38	0.49	0.62	0.67	0.50	0.38	0.44
KGDALα	0.57	0.38	0.08	0.02	0.04	0.56	0.49	0.04	0.02	0.03
KGDALαβ	0.69	0.72	0.97	0.46	0.62	0.68	0.81	0.97	0.46	0.62
KGDAL	0.65	0.59	0.58	0.62	0.60	0.63	0.56	0.40	0.62	0.48

Table 6:

Performance of morality prediction in the next 48 hours during RRT (MIMIC-III)

	Npos : Nneg = 1 : 1					Npos : Nneg = 1 : 2
	ROCAUC	ACC	PREC	REC	F1	ROCAUC	ACC	PREC	REC	F1
XGBoost	0.50	0.50	0.49	0.17	0.25	0.50	0.60	0.32	0.17	0.22
LSTM	0.71	0.64	0.79	0.38	0.52	0.70	0.71	0.61	0.38	0.47
Transformer	0.66	0.59	0.67	0.37	0.47	0.65	0.65	0.48	0.37	0.41
KGDALα	0.59	0.39	0.08	0.02	0.03	0.57	0.49	0.03	0.02	0.02
KGDALαβ	0.67	0.70	0.92	0.44	0.60	0.66	0.79	0.84	0.44	0.58
KGDAL	0.62	0.56	0.56	0.59	0.57	0.61	0.53	0.37	0.59	0.46

As a baseline, a random model achieved ROCAUC 0.52, accuracy 0.52, precision 0.52, recall 0.51, and F1 0.52. The random forest model predicted that almost every patient survived in all experiments. Its average performance on all experiments are ROCAUC 0.50, accuracy 0.58, precision 0.25, recall 0.01 and F1 0.02. The XGBoost model performed better than random forest. However, it is only slightly better than the random model. The performance of deep learning models, including LSTM and Transformer, performed significantly better than the compared traditional machine learning models, indicating that for temporal data based rolling prediction, deep learning models can better capture the critical temporal patterns, resulting in better performance than that of traditional machine learning methods.

Table 2 and Table 3 show that KGDAL has the best performance on 72 and 48-hour rolling prediction on the UK data on almost all evaluation metrics on both balanced and imbalanced test data. The second-best model is KGDAL_α followed by KGDAL_αβ and Transformer. However, Table 4 shows that KGDAL_α outperformed KGDAL on both the balanced and imbalanced test data for the 24-hour mortality rolling prediction, even though KGDAL has the best recall. It suggests that for shorter prediction windows, the time attention mechanism has more contribution than the feature attention mechanism with/without KG adjustment, and KG-adjusted feature attention can improve the overall model performance slightly. In summary, all the experiments on UK data show that the attention-based models including Transformer, KGDAL_α, KGDAL_αβ, have better overall performance than LSTM, indicating the attention per se is critical for rolling prediction tasks.

On the MIMIC-III data, KGDAL_αβ has the best overall performance for both 72 and 48-hour rolling mortality prediction. While LSTM has the highest ROCAUC and KGDAL achieves better recall, KGDAL_αβ has higher scores on accuracy, precision, and F1. Surprisingly, KGDAL_α, which is the second-best model on the UK data, has the lowest performance among the three in all the experiments on the MIMIC-III data. The performance of KGDAL_α on this experiment suggests that it is feature attentions rather than time attentions that play the critical role in MIMIC-III rolling mortality prediction.

In summary, we found that KGDAL’s time and feature attentions play an essential role in rolling mortality prediction. While on one data, time attention is more important; feature attention could dominate the model performance on another data. The KG-guided grouping is crucial for all experiments and datasets. Nevertheless, other factors such as sample size and patient distribution may affect model performance as well.

3.2.3. Mortality Risk Trajectory Analysis.

A patient’s mortality risk trajectory is a series of predicted risks of all the subsequences of that patient ordered by time. An example trajectory is shown in Figure 5. Given all the mortality risk trajectories, we computed the pairwise trajectory distances using dynamic time wrapping [28], and used hierarchical clustering to identify similar risk trends among all the correctly predicted patients. Two trajectory clusters shown in Figure 4 revealed multiple episodes of increasing risks for non-survivors and quick decreasing risks for survivors. The trend-based analysis may assist healthcare providers in making early decisions before the risk increases.

Figure 5:

The risk trajectory of a survival patient.

Figure 4:

Two risk trajectory clusters with different endings.

As a case study, we visualized the mortality risk trajectory of an AKI-D patient in Figure 5. In the figure, the x-axis is the days before outcome event (survival (end of follow-up) / mortality), the y-axis is the predicted risk of death, and the risk scores range from 0 (survived) and 1 (died). Every blue point is the predicted mortality risk from a subsequence of the same patient. All the points were fitted to a smooth curve using polynomial regression showing the trajectory of predicted mortality risks. In the case study, the patient finally survived. However, the risk score was not monotonically decreasing. Starting with a high risk, KGDAL predicted that the risk was gradually decreased for 5 days. The risk trajectory then stayed roughly at 50% for 10 days with mild fluctuations. Finally, the trajectory decreased quickly in the last 5 days of RRT.

The corresponding KG-adjusted attentions are shown in Figure 6, where each row represents the attentions of a feature group, each column is a time point (the darker the color, the high the score). The risk trajectory aligned well with the KG-adjusted attention. The attention hit three times the maximum in the early RRT duration, indicating a high mortality risk at that time. The figure also reveals that “Respiratory” and “Blood” had overall higher attentions than the other concept groups. This is well-aligned with the observations in the knowledge graph that “Respiratory”, “Blood”, and “AKI” are closely related concepts. The high agreement between the risk trajectory and the time and feature attentions suggests that the attentions obtained from KGDAL may be useful to explain to clinicians the potential risks and why the risk is high or low so that interventions can be taken in place timely.

Figure 6:

An example of the KG-adjusted 2-D attentions.

4. CONCLUSION

In this work, we developed a novel model called KGDAL for rolling mortality prediction for AKI-D patients. KGDAL uses a knowledge graph to guide the generation of 2D attention in both time and feature spaces. KGDAL and its variations achieved the best performance on both the UK data and the MIMIC-III data. Using a case study, we demonstrated the interpretability of KGDAL and the capability of using KGDAL for assisting timely decisions for clinicians.

For future work, we will further optimize the architecture of KGDAL, adjust the training and testing processes, and visualize the resulting trajectories for clinicians. Also, we will incorporate a customized knowledge graph to better capture the relationships between the temporal features for patients from different institutions and further validate KGDAL on large-scale cross-institutional datasets.

CCS CONCEPTS.

• Computing methodologies → Neural networks; • Applied computing → Health informatics.