ATTAIN: Attention-based Time-Aware LSTM Networks for Disease Progression Modeling

1 Introduction

Electronic Health Records (EHRs) are a large-scale and systematic collection of temporal health information of patients. The broad adoption of EHRs in medical systems has promoted the development of various computational methods for understanding the medical history of patients and predicting risks [Marlin et al., 2012; Choi et al., 2016a; Zhou et al., 2013; Choi et al., 2016b]. In this work, we focus on a task named Disease Progression Modeling (DPM), which monitors the disease developing process and predicts future risks based on patients' historical information. DPM is crucial for making clinical decisions and providing prompt medications. A large amount of recent works have been developed for this task [Choi et al., 2016a; Esteban et al., 2016; Lipton et al., 2015; Zhou et al., 2013]. Among them, Recurrent Neural Network (RNN) is one of the most extensively researched deep neural networks to handle the sequential data. As an extension of RNN, the Long-Short Term Memory (LSTM) is specifically designed to capture long-term patterns that commonly exist over a long period of patients'

records [Sundermeyer et al., 2012]. LSTM-based approaches have found a huge success in a variety of tasks involving sequential data, such as video processing, climate changes detecting, and EHRs representation learning [Jia et al., 2019; Jia et al., 2017; Lipton et al., 2015; Esteban et al., 2016].

Despite its great success, significant barriers remain when applying the standard LSTM for modeling EHRs. First, it cannot interpret prediction results. Second, it does not consider irregular time intervals between consecutive events.

Interpretability of computational models is extremely critical in healthcare-related domains. In real hospital settings, it is generally more important to learn about the discriminative interpretable patterns which capture informative progression of a disease than to induce an accurate predictive computational model. Various attention-based neural networks are widely developed to generate interpretations for EHRs. [Choi et al., 2016b; Ma et al., 2017; Sha and Wang, 2017]. As frontier work, RETAIN [Choi et al., 2016b] applied a two-level attention mechanism to identify meaningful visits and specific features that contribute to the prediction.

Measurements in EHRs are commonly acquired with irregular intervals. For example, when a patient is in a severe condition, events are likely to be recorded more frequently than when a patient is in a relatively "healthier" condition. Hence such varying time intervals can reveal patient's health status on certain impending conditions, and it is important to consider the time intervals between temporal events to capture latent progressive patterns of a disease. There have been several previous works on handling the time irregularity [Baytas et al., 2017; Pham et al., 2016; Choi et al., 2016a; Che et al., 2017], e.g, Time-aware LSTM (T-LSTM) [Baytas et al., 2017] transforms time intervals into weights and uses them to adjust the memory passed from previous moments.

In this work, we propose ATTAIN, an attention-based time-aware disease progression model, that incorporates the attention mechanism and models the time irregularity between events. Specifically, we adjust the memory of LSTM when accumulating previous information. Instead of adding memory from one previous event, we retrospect memories of all/several previous events and discount them by weights generated from attention mechanism and the time intervals between those events and current event. The overall weights represent how important each previous event is for the current event to identify the progressing condition. Three attention mechanisms are explored: global (g), local (l), and flexible (f), to generate the attention weights. On the other hand, the time intervals are transformed to decay weights through a decay function so that the outdated events are more likely to play a less important role than recent events for predicting the outcome of the current event. We believe the obtained attention weights together with the decay weights would not only lead to an effective predictive model but also result in an interpretable and clinically reasonable model.

We validate the proposed ATTAIN on the task of early prediction of septic shock. Sepsis is a life-threatening organ dysfunction [Singer et al., 2016] and a leading cause of death in the United States. Septic shock, the most severe complication of sepsis, leads to a mortality rate as high as 50% and an increasing annualized incidence [Dellinger et al., 2008]. In fact, as many as 80% of sepsis deaths could be prevented with timely diagnosis and treatment [Kumar et al., 2006]. One major challenge associated with such early prediction is the subtle but fast progression at early stage. For example, only minor changes are reflected on white blood cells and body temperature at early stage [Kumar et al., 2006]. Besides, the indicators of sepsis are non-specific, such as infection or fast heart rate. Hence, patients with such symptoms are highly likely to progress to other disease. Because of such delicate progressions, variables in the before-shock stage may either be measured infrequently or not measured at all. Therefore, it is quite critical and challenging to identify previous indicative moments and give accurate predictions of patients' health status.

Our experimental results on real-world EHRs show that ATTAIN outperforms the state-of-the-art models such as RETAIN and T-LSTM . Also, the generated interpretative timeaware attention weights shed some light on the progression behaviors of septic shock.

2 Method

2.1 ProblemDefinition

Our dataset can be represented as X = {x1, x2, ..., xN}, where N is the total number of hospital visits. It is composed of multi-variate irregular time series data and each visit xk consists of a sequence of events: xk = {x1k , ..., xTk k }, where xtk represents the patient's records at time step t in xk. We have xtk ∈ RD , where D is the number of features or measurements recorded at each event and Tk is number of events in the visit k which varies with different visits.

For each xk, we are provided with the event-level label yk = {y1 k , ..., yTk k } for the sequence of events. ytk = 1 indicates a patient is in septic shock state at a given time t in the visit xk, otherwise ytk = 0. The goal of this work is to predict the (t + 1)-th event-level label yt+1 k given the clinical events from time 1 to t: x1k , x2k , ..., xtk in each visit k. For simplicity, we omit index k hereinafter when it does not cause ambiguity.

2.2 Long Short-TermMemory (LSTM)

In the standard LSTM cell unit, the cell state ct serves as an internal memory and controls the information flow. It is generated by forgetting information through a forget gate f t and most recent cell state ct-1 , and adding new information through an input gate it and a candidate cell state ˜ct as listed in Eq. 1.

where ht-1 is a hidden state output by ct-1 , {Wh ∈ RH×H ,Wx ∈ RH×D } denote network parameters to be trained and H is the number of hidden nodes. The new cell state can be obtained as follows:

where ⊗ denotes entry-wise product.

Finally, we generate the hidden states by filtering the new cell state through an output gate layer ot , and produce the probability of each event t at risk of septic shock using a sigmoid function with parameter U :

2.3 ATTAINNetworks

Fig. 1 shows the architecture of proposed ATTAIN framework. In the original LSTM, the cell state of current event ct obtains memory from its most recent event i.e., ct-1 as shown in Eq. 2. Due to the well-known vanishing gradient problem [Baytas et al., 2017], the learning of each cell state still heavily depends on most recent events. However, in disease progression, relationships of events are non-trivial in that the current condition is often not impacted only by recent moments. Single memory-reading also brings that it cannot interpret how critical of each past event. Conversely in practice, doctors usually review the patient's records to identify the critical previous events.

Here we propose to use attention mechanism to collect multiple previous memories and explore the relationships of events in patient trajectories. We expect it can help doctors make prompt decisions and achieve a better prediction.

The standard LSTM implicitly assumes that intervals between time steps, or events, in a sequence are uniformly distributed, whereas the frequency of collecting patients' records often vary greatly based on many factors. The interval between two consecutive events during a patient's visit can span from minutes to several weeks. In general, the events that occur long time ago tend to have less impact to the current event and thus we should properly reduce their contributions. Therefore, it is important to consider the elapsed time when predicting the current event's output. In addition, these irregular time intervals might be indicative for patients' health status on certain impending conditions. For example, when a patient is in a more severe condition, events are likely to be recorded more frequently than when a patient is in a relatively "healthier" condition.

To incorporate the elapsed time into the standard LSTM, we feed the intervals into a monotonic decreasing function. In specific, for each event, the memories collected from its previous events are discounted by a decay factor, which is estimated based on the time intervals between previous events and the current event. In this way, if the previous event occurs a long time ago, the dependency from the previous moment would not play an active role for predicting the current event's output.

As shown in Fig. 1, in the proposed ATTAIN, instead of reading the information from just one previous cell state, we collect all/some previous cell states and regularize these old memories by weights generated by attention mechanism and time intervals as follows:

where αti is the attention weight from i-th event to the current event t, ∆tti is the time interval between i-th event to the current event, g(·) is a decay function, and m stands for the number of events to look backwards. Then we replace ct-1 in Eq. 2 with Ct-1 for generating the new cell state.

Three Attention Mechanisms

To determining m, three ways are explored:

• Global Attention generates attention weights αti from all the previous events for the current time step t, , i.e, m = t - 1 in Eq. 4.

The advantage of the global attention is that it considers all past events and its disadvantage is that it is computationally expensive and attention weights are sparsely distributed. To address these problems, we also explore the following two options.

• Local Attention only considers a fixed number of previous events, m in Eq. 4. The hyperparameter m can be

optimized by grid search or suggested by clinicians.

• Flexible Attention learns the optimal number of previous events adaptively for different given events, i.e, m varies for different time step t. The intuition is that different number of previous events should be considered for different given events.

where Wm and vm are parameters to be learned, and m(t) ∈ [1, t - 1].

Previously RETAIN [Choi et al., 2016b] introduced attention mechanism in learning health record representation. RETAIN generated attention weights from the hidden states of RNN: eti = W ⊤α hi , i ∈ [1, t - 1]. However, it does not capture the relation between the targeting event t and past events is. Different with RETAIN, the aforementioned three attention mechanisms incorporate this relation when estimating attention weights. Additionally, when aggregating previous information (see Eq. 4), αti measures the weight of event i for the current event t and hence its computation should solely depend on the data at events i and t. Therefore we use the original inputs x (i.e. xt and xi ) instead of hidden states h (i.e. ht and hi ) since the pair [xt, xi] extracts immediate relations among the medical variables of two events while [ht, hi] encodes the information from multiple events.

Time-aware Decay Function

Time intervals might range from minutes, hours and even to days in healthcare domain. We integrate the intervals into the model and make LSTM predict in a time-aware way. As a guideline, the function in Eq. 8 is validated and suggested by several prior works for long time periods in EHR data [Pham et al., 2016; Baytas et al., 2017].

Previously T-LSTM [Baytas et al., 2017] also integrated the elapsed time to adjust the cell memory. It divided the previous cell memory as short-term (ct-1 S ) and long-term (ct-1 L ), then discounted only the short-term memory through the same decay function, i.e., ct = f t ⊗ (ct-1 L + ct-1 S · g(∆t)) + it ⊗ ˜ct . T-LSTM depends on only one previous cell state to adjust the memory and is unable to interpret how important each event is for predicting the current outcome.

In summary, through the aforementioned memory adjustment, ATTAIN would identify the critical previous moments and take the effect of elapsed time into account. The two sources of weights, based on attention mechanism and time intervals, are simultaneously tuned together. The intuitions behind ATTAIN can be seen as mimicking doctors reading records for diagnosing in the real world in that they generally pay attention to both very important indicative previous events even though they happened long time ago but also the most recent events. As a result, we expect ATTAIN can lead to more reasonable predictions than the standard LSTM.

3 Experiments

3.1 Data Description

Our EHR data was collected from Christiana Care Health System Health System (CCHS) from July, 2013 to December, 2015. Each data point is a visit with a series of events(see 2.1). In total, there are 210,289 visits, and 10,412,729 medical events. The study population is patients with suspected infection, which is identified by the presence of any type of antibiotic, antiviral, or antifungal administration, or a positive test result of Point of Care Rapid, and it consists of 52,919 visits and 4,224,567 medical events. The definition of study population and the following data prepossessing were determined by leading clinicians from CCHS and Mayo Clinic.

After preprocessing, our final dataset contains 2,100 visits (1,869 positives and 231 negatives) and 209,346 events (22,430 positives and 186,916 negatives).

Tagging. International Classification of Diseases, Ninth Revision (ICD-9) are widely used for clinical labeling (i.e., septic shock or not). However, as visit-level labels, ICD-9 codes cannot tell when septic shock occurs at event level, and they are not persistently reliable [Ho et al., 2014]. Therefore based on Third International Consensus Definitions [Singer et al., 2016], our clinicians identified septic shock at each event as having received vasopressor(s) or persistent hypotension (systolic blood pressure<90 mmHg or mean arterial pressure<65 mmHg for more than 1 hour). When applying both ICD-9 and our tagging rule, we identify 1,869 shock positive visits and 23,901 negative ones.

Sampling. We further conduct stratified random sampling on negative visits while keeping their underlying distribution of age, gender, race and length of stay the same as positive visits, and then get 1,869 negative visits. Given the number of negative events are dominant in both positive and negative visits, we randomly sample 231 out of 1,869 negative visits to maintain the ratio of positive events around 15%.

Data selection. We focus on the visits in range of 3-90 days since short visits do not hold sufficient information for analysis and long visits introduce data sparsity. Importantly, we exclude the visits that developed septic shock within 8 hours after admission, because our clinicians suggest that such patients generally already show septic-related symptoms on admission and doctors can easily catch such cases.

Feature selection. We exclude features with a missing rate more than 90% and our features can be divided into four categories: 1) vital signs: heart rate, temperature, etc; 2) lab results: BUN, creatinine, white blood cell count (WBC), 18 culture tests, etc; 3) interventions: FIO2, oxygen flow, etc; 4) locations (e.g, emergency or nurse), descriptions, identifiers.

Data aggregation and imputation. We aggregate the events within an hour to an event and take the mean measurement within 1-hour window. For frequently measured features (e.g. vital signs), we add statistical features (e.g. min and max) measured in the 1-hour window. Suggested by our clinicians, we impute missing values with the last value within the fixed length of forward time window (8 hours for vital signs and 24 hours for lab tests) and impute the remaining ones with feature-wise mean values.

3.2 Experimental Setup

Baseline Approaches

We compare ATTAIN with three baselines: 1) LSTM without either attention or time-aware mechanisms; 2) RETAIN [Choi et al., 2016b] with a two-level attention mechanism at both the event-level and the variable level; 3) T-LSTM [Baytas et al., 2017] with time-aware mechanism only.

Our Approaches

We integrate three attention mechanisms with the standard LSTM: LSTMg, LSTMl, LSTMf . The attention weights are generated by Eq. 5 - 7 respectively. Old memory is obtained by accumulating all or selective previous cell states:

Specifically for local attention, we find the optimal setting of m by grid search on a validation set (described later). Fig. 2(a) shows the model achieves competitive performance using tenth-interval m ∈ {20, 30, 40}. We hence further explore the performance for m ∈ [20, 40] incremented by 1 and Fig. 2(b) shows that the optimal values are m ∈ {22, 26}. Here we use m = 24.

Furthermore, we add the time-aware mechanism to the three models above resulted in three ATTAIN models named ATTAINg, ATTAINl, ATTAINf respectively. The Ct-1 is calculated based on Eq. 4 in a way that the old memory is regularized by both attention weights and decay weights.

Evaluation and Setting

Our evaluation metrics include sensitivity/recall, specificity, precision/positive predictive value (PPV), F1-score, and area under the ROC (receiver operator characteristic) curve (AUC) [Nachimuthu and Haug, 2012]. Precision, recall, F1 and AUC are widely used for evaluating machine learning approaches. In healthcare domain, researchers commonly refer to sensitivity, specificity and PPV for the annotation performance.

In our implementation, all the models are updated using mini-batch stochastic optimizer and the batch size is 50. The training epochs is 50 with early stopping, the learning rate is 0.01, and the number of hidden units for LSTM is 72. In training process, we randomly divide the data sets into the training, validation and testing set with the ratio of 70%, 15%, and 15%. Each experiment is repeated 10 times with random model initialization and we report the average values with standard deviation for each evaluation metric.

• The best values of each metric within the session are in bold, and the best values of each metric across all the models are labeled with ⋆.

Table 1: Performance (±standard deviation) of baselines and our approaches on septic shock overall prediction.

4 Results

4.1 Results of Overall Prediction

Table 1 shows the performance of all the models and from top to bottom, we have three baseline models, our proposed three attention mechanisms with LSTM, and the three corresponding ATTAIN models. Columns 'A' and 'T' indicate whether attention or time-aware mechanisms are used in the models. For each sub-session, the best results are marked in bold and the best model across all methods are labeled with ⋆.

Among the three baselines, both RETAIN and T-LSTM outperform the standard LSTM on almost all the evaluations, which shows that either the attention mechanism or timeaware mechanism can lead to a better prediction of the disease progression. As in the original work [Choi et al., 2016b], RETAIN does not boost from LSTM much and this might come from two reasons. First, their attention mechanism did not incorporate the relation between the previous events and the current event. Second, due to high missing rate of features, the variable-level attention did not take into full effect.

All three attention-based models (in the middle subsession) outperform three baselines across all evaluations with only one exception that LSTMg has worse recall than T-LSTM. When comparing the three attention mechanisms, LSTMl performs more stably than LSTMg and LSTMf , i.e., 0.695 vs. 0.682 and 0.695 (larger std) for F1 score and 0.763 vs. 0.748 and 0.755 for AUC. We found that the attention weights of LSTMg are sparsely distributed in long sequences and cannot effectively detect critical moments. The result of LSTMl validates that m = 24 is clinically reasonable. When only using attention mechanism, though close, LSTMf , deciding how many previous events to pay attention in a datadriven manner, is not as competitive as LSTMl.

Except that ATTAINg has a worse recall than T-LSTM, all three ATTAIN models (in the bottom sub-session) outperform three baselines with obvious improvement. Moreover, when comparing between ATTAIN and our three attention-based models, ATTAINg outperforms LSTMg, and the same observation on ATTAINl over LSTMl and ATTAINf over LSTMf . Thus, when integrating both attention and time-aware mechanisms , the two sources of weights allow the model to comprehensively identify the meaningful previous events in terms of medical relation and time effect. For example, if attention weights indicate two previous events are equally important, we pay attention to the recent event first, exactly as doctors perform in real world.

Overall, ATTAINf achieves the best F1-score and AUC. ATTAINl is close to ATTAINf , i.e. 0.718 vs. 0.720 on F1 and 0.804 vs. 0.811 on AUC and gets the best recall. Since our data are aggregated at hourly rate, 24-event retrospect guarantees attending the last 24 hours of a patient status. In practice the clinicians prefer this local-based model because previous information can be sufficiently obtained with certainty.

4.2 Results of Early Prediction

Our best attention-based model LSTMl and our best approach ATTAINf are compared against three baselines: LSTM, RETAIN and T-LSTM for early prediction, i.e. to detect the patients who are developing septic shock in advance. To do so, we utilize the future diagnostic condition as labels for each current event in training process. For example, in previous setting, we use {xik }ti=1 to predict the (t + 1)-th event-level label yt+1 k , one event in advance. In early prediction setting, we will use {xik }ti=1 to predict yt+η , η events in advance. Since our data is aggregated on an hourly basis, η-event early guarantees at least η-hour ahead. We track the performance of each model on F1-score and AUC from one-event earlier up to 24 incremented by 1.

Fig. 3 shows the results of early prediction on five models. As expected, as η increases, the early prediction tasks become more and more challenging in that the performance of all the models has gradually degraded. Despite this, ATTAINf has been decreasing much slower and steadily staying on the top of all other models in terms of both AUC and F1. LSTMl follows ATTAINf to be the second best model and the standard LSTM performs the worst.

4.3 Case Study

To illustrate the interpretability, Fig. 4 shows a case study by comparing the weights learned from RETAIN, LSTMl and ATTAINl for a sample patient's visit. To simplify, we fix m = 10 for all three models and the weights are normalized. The patient is in septic shock at events 11 and 12. For event 11, we expect the models can identify some critical previous event(s). We consider three prior events with the highest weights to be critical events. At event 12, septic shock continues, we expect a reliable model: a) identifies event 11 to be critical for event 12, because the former clearly explains why the patient was still in septic shock at event 12, and b) identifies the same critical events (other than event 11) as those for event 11 for event 12. Fig. 4 shows that RETAIN indeed identifies that event 11 is critical for event 12, but it discovers different sets of critical events for the two events. LSTMl consistently detects the same set of critical events, {4, 5, 8}, for both events 11 and 12, but it fails to identify that event 11 is critical to event 12. Finally, our ATTAIN model successfully identify that event 11 is important and besides event 11, event 12 has the same critical set: {4, 5, 8} as event 11.

5 RelatedWork

EHRs have been a popular research platform with increasing availability to develop predictive models for tasks of disease progression [Zhang et al., 2017; Wang et al., 2014], phenotyping [Che et al., 2015; Baytas et al., 2017; Liu et al., 2015; Li et al., 2015], diagnosis prediction [Choi et al., 2016b], etc. However, EHRs also pose numerous challenges, such as they are noisy, fragmental and high dimensional. To this end, deep learning networks, such as RNN [Che et al., 2018] and LSTM [Lin et al., 2018; Yang et al., 2018] can assist in learning complex relationships among medical events.

In recent years, attention mechanisms are extensively explored to interpret the model output and greatly improve the prediction performance. For example, RETAIN applies a reverse time attention mechanism in an RNN [Choi et al., 2016b] and Dipole [Ma et al., 2017] uses the similar attention networks for diagnosis prediction. Another challenge associated with EHR data, time irregularity, has also been tackled. T-LSTM [Baytas et al., 2017] divides short-term from the previous cell memory, and adjusts it with a time-aware mechanism. In [Pham et al., 2016], the time intervals are used to modify the forget gate of LSTM. In [Che et al., 2018], time gaps are made regular through data imputation methods. Finally, Health-ATM [Ma et al., 2018] extracts patient information patterns with attentive and time-aware models through RNN and Convolutional Neural Networks (CNN). Compared with the prior works, our proposed method explores different attention mechanisms to generate weights for the past events while handling the time irregularity in EHRs. For acute medical conditions such as septic shock, it is extremely significant to identify critical and timely meaningful events.

6 Conclusion

Disease progression modeling is an important task especially for acute medical conditions such as septic shock. In this work, we propose ATTAIN, an attention-based time-aware LSTM networks to effectively improve the interpretability of LSTM while also modeling the irregular time intervals. In specific, ATTAIN employs attention weights combining with time decay function to identify the contributions of the historical events to the current event. The experimental results on a real EHR dataset demonstrate the effectiveness of ATTAIN compared with state-of-the-art baselines. One limitation of this work is not generating feature-level attentions at each event, since the feature-level attentions can vary drastically over multiple events in acute disease conditions, especially given the impact of high missing rate in real-world EHRs. We keep this challenging task as our future work. We will be working on imputing realistic missing data for heterogeneous patient groups through domain adaptation and generative networks. Also, we will validate ATTAIN on other EHRs such as MIMIC-III and for other disease prediction tasks.

Acknowledgements

This research is supported by the NSF Grants: #1522107, #1651909, #1660878 and #1726550.