Temporal-spatial Correlation Attention Network for Clinical Data Analysis in Intensive Care Unit

Traditional machine learning methods for building predictive models based on medical time-series data have been developed over decades, including conditional random field (CRF) [7], Expectation-Maximization algorithm (EM) [8], Bayesian networks [9] and others. Caruana et al. [10] showed that backpropagation can be effective in medical data analysis, and Cooper et al. [11] demonstrated the application of eight machine learning algorithms, including K-nearest neighbor (KNN), in predicting patients’ survival. Awad et al. [12] used a random forest algorithm to predict early mortality of patients, and Bayesian networks [13, 14] were used to develop multi-label classifiers, which were applied to medical multi-classification problems. Logistic Regression (LR) was widely used in the field of predicting mortality risk among hospitalized patients [15, 16, 17, 11] and was chosen as one of the baselines.

In recent years, deep learning technology with deep neural networks as the core has been gradually emerging in the field of medical applications with the increase of computers’ computing power. Lasko et al. [18] used hierarchical autoencoders (AEs) based on time-series data of uric acid measurements to predict the risk of gout and acute leukaemia. Chen et al. [19] proposed a model based on CNN to extract local temporal associations in time-series data for the prediction of futural risk of congestive heart failure (CHF) and chronic obstructive pulmonary disease (COPD). Mobley et al. [20] applied ANN to length of stay prediction, Grigsby et al. [21] considered length of stay prediction as a binary classification task in order to identify patients at risk of long-term hospitalisation in early stage, and Emma et al. [3] proposed the Temporal Pointwise Convolution (TPC) achieving relatively good results in this task.

Lipton et al. [6] extracted 13 variables of laboratory outcome from paediatric ICU patients and used Long Short-Term Memory networks (LSTM) to learn temporal associations and long-distance dependencies in time-series data. Choi et al. [22] used a variety of RNN models such as Gated Recurrent Unit (GRU) and LSTM networks to process clinical time-series data for the prediction of medical events. It is worth mentioning the extensive application of LSTM in clinical prediction tasks, including prediction of cardiac arrest [23], acute kidney injury [24], missing data inference [25], prediction of drugs [26, 27] and length-of-stay prediction [28].

In the task of mortality prediction, [29, 30] illustrated that artificial neural networks (ANN) can achieve better results than Logistic Regression. For decompensation prediction tasks, the Recurrent Attentive and Intensive Model (RAIM) [31] and Simply Attend and Diagnose (SAnD) [32] have applied attention mechanisms to predict physiological decompensation of critical patients in ICU. Phenotype classification is also a promising task and temporal convolutional networks [33] and feedforward networks [34, 18] have already be used in predictions of diagnostic codes based on medical time-series data. There were researches indicating that deep learning models represented by LSTM performs well in predicting mortality [35, 36], length of stay [28] and diagnostic classification [37]. So LSTM [4] was chosen as another baseline.

We use the Medical Information Mart for Intensive Care (MIMIC-IV v0.4) dataset [38]. The MIMIC-IV database is a public database of clinical data that researchers from all over the world can use for free. The database has clinical information on more than 380,000 patients who were admitted to Beth Israel Deaconess Medical Center in Boston, Massachusetts, USA, from 2008 to 2019. Like other EHR databases, it keeps detailed information on patients’ demographics, lab tests, medication administration, vital signs, surgical operations, disease diagnosis, medication management, survival status, and more.

MIMIC-IV uses a modular approach to organize data. It has three modules that are made up of 27 tables that make it easy to use data from different sources separately or together. Another advantage is the protection of patients’ privacy, which is de-identified in two steps: the first is to replace the patient identifier, hospital identifier, and so on with a random code, and the second is to add a random number of days to the date data fixed for each patient.

The data pre-processing workflow is illustrated in Fig.1. The MIMIC-IV critical care database contains 76540 ICU stays from 53150 patients admitted.

In the first step, the important data are taken straight from the original MIMIC-IV tables and put in order by subject, which is another word for patient. This step has two parts: keeping only adults who were in the ICU (over the age of 18), who were in the ICU only once, and who did not move between ICU wards or wards during the same hospital stay. In this step, differences in how children’s and adults’ bodies operate and other unclear factors are taken out. This leaves 47,046 unique patients with a total of 59,372 ICU stays.

In the second step, we exclude clinical events that cannot be matched with ICU stays. Firstly, it considers events missing admission ID (HADM_ID), and only events with HADM_ID are reserved. Secondly, since the $stays.csv$ table allows the connection of HADM_ID with ICU stays, we exclude events that do not have HADM_ID in $stays.csv$. Then, for events where ICU stay IDs (ICUSTAY_ID) are defaults, they can be recovered by attempting to check the ID. Finally, we only retain events for which the ICUSTAY_ID is present in $stays.csv$.

In the third step, a patient has a single ICU admission, which is also described as an “episode”. In this script, for each episode, a list of events is put together in order of time, with only the variables from a list that has already been set. We use 155 physiologic variables, expanding from 17 variables, which are a subset from the Physionet/CinC Challenge 2012 [39]. During the collaboration, the local hospitals helped choose the 155 physiological variables, which included 5 categorical variables and 150 continuous value variables. The 24 variables of 155 are list in Table.1. Thus far, we have obtained over 92 million events from five tables (icu/inputevents, icu/procedureevents, icu/chartevents, icu/outputevents, hosp/labevents). Finally, we fixed a test set of 15% (7,057) of patients, including 8,906 ICU stays.

For each patient, we need to predict in-hospital mortality, and other tasks in each time interval, such as the length-of-stay prediction. Fig.3 shows the problems. To predict, we only use previous clinical data from each time interval. As shown in Fig.3, we chose the data from $t$ hours before the prediction time point as the reference data for the predicting related tasks in this paper.

For each patient $i$, we can obtain a set of samples $X^{i}=\{x^{i}_{1},x^{i}_{2},....,x^{i}_{t}\}$, where $X^{i}\in R^{t\times d}$. $t$ denotes that we select $t$ times clinical data before prediction time point. $d$ denotes that the $d$ clinical test values are selected in each hour. In Section 3, we select $155$ clinical values based on the advice of Professor Doctor. Here, we apply the one-hot method [40] to handle clinical data and convert these clinical values in to the format of vector. Finally, in each hour, clinical data $1\times 155$ can be convert to $1\times d$ dimension vector for final predict task.

In the prediction task, our model needs to map medical features to clinically meaningful labels such as length of stay, in-hospital mortality, etc. A predictive model for c classification can be represented as a mapping from input to category, i.e.$f(X^{i}):R^{t\times d}\rightarrow\{0,1,\ldots,c\}$ and the entire sample data set is then denoted as $D=\{X^{i},y^{i}\}_{i=1}^{N}$, where $y_{i}$ is the target label, $N$ is the number of samples. For the in-hospital mortality task, the patient target label is $y\in\{0,1\}$, where $0$ denotes survival and $1$ denotes death. The predictive model is derived from the input data $X^{i}$ to obtain the output $y^{*}_{i}=f(X^{i},\theta)$, where $\theta$ is the parameters of the model. The key of our work is to obtain the best parameter $\theta^{*}$ for different tasks, which can minimize the difference between the output of model $y^{*}_{i}$ and the true label $y_{i}$. In the next subsection, we will detail the structure of our model.

In order to consider the temporal and spatial information, we apply the two-branch network as Fig.2. Each branch has the same structure, which includes the initial Encoder $E_{e}$ and the Fusion-Encoder $E_{f}$.

We divided each sample $X^{i}$ into $n$ equal parts based on the time dimension and the model follows these chunks for input:

Temporal Branch. As shown in the upper part of Fig.2, this branch consists of two parts: an Encoder and $n-1$ Fusion-Encoder. $X_{0}^{i}$ is the input of Encoder and the output is denoted as $z_{0}$:

where $\omega$ denotes the network parameters of Encoder. Encoder’s core structure is based on a multi-head self-attention mechanism. The calculation in Encoder can be expressed in detail as follows:

Then we process the following set of data, $X^{i}_{j},j=1,...,n-1$, and merge it with previous data information, $z_{j-1}$. The Fusion-Encoder $E_{f}$ has the same internal structure, but there is a recursive relationship. The j-th Fusion-Encoder is represented by:

where $\delta_{j}$ denotes the network parameters of the j-th Fusion-Encoder. When $j=1$, connect line in model and the second input of the 1st Fusion-Encoder is the output of Encoder, $z_{0}$. For $j=2,3,...,n-1$, the second input of Fusion-Encoder is the output of the last Fusion-Encoder, $z_{(j-1)}$, at which point line is connected. The calculation in the j-th Fusion-Encoder is as follows:

The core is a multi-head cross-attention mechanism in which Q is derived from self-attentive output of the j-th time series $X^{i}_{j}$, K from the previous Fusion-Encoder’s output, and V from a mapping that combines Q and K. In this case, V focuses on a relatively global feature, so the recursive structure can also focus on global information.

Spatial Branch. We transpose time series blocks of $n$ equal parts as the input of the spatial branch to obtain the relevance of variable dimensions by attention mechanism. As shown in Fig.2, $(X_{j}^{i})^{T}\in R^{d\times\frac{t}{n}}$, replaces $X_{j}^{i}$ of temporal branch. There is no additional difference between this branch and the upper one, both are composed of Encoder and Fusion-Encoder.

Finally, we can obtain $z_{n-1}^{t}$ from the temporal branch and $z_{n-1}^{s}$ from the spatial branch. Here, we connect $z_{n-1}^{t}$ and $z_{n-1}^{s}$ as the final feature of patient in the predict time point, and apply the Softmax to train and obtain final predict result. Algorithm 1 shows the overall training process.

Predicting in-hospital mortality is a binary classification task, so the area under the receiver operating characteristic (AUC-ROC) is the main way to measure how well it works, and the area under the Precision-Recall (AUC-PR) is the supplement. The main metric for length-of-stay prediction is Median Absolute Deviation (MAD) which is lower to indicate better performance, and Cohen’s linear weighted kappa score (KAPPA) that the higher the better is usually used as a supplement. Most patients have more than one condition, so it’s clear that phenotype classification is a multi-label classification problem, and we chose macro- and micro-averaged AUC-ROC as the evaluation criteria for this problem. Predicting decompensation is made easier by treating it as a binary classification task like predicting death, which is also mainly judged by the AUC-ROC.

Traditionally, Logistic Regression (LR) was the most accurate model for medical time series. In recent years, LSTM and some deeplearning methods have done well with time series. In this section, Logistic Regression [4] and LSTM [4] were used as a baseline and compared with our approach.

For in-hospital mortality prediction, the prediction time point is defined as 48h after the beginning of the ICU stay, and $t$=48h which means one ICU stay is one sample. As can be seen from the Table.2, TSCAN shows a huge improvement in classification scores relative to the baseline, with feature concatenation improving by 0.011 relative to LR and Max pooling improving by 0.052 relative to LSTM, and we have achieved significant performance benefits of 0.2% compared to other SOTA methods (TPC) in ACU-ROC.

For length-of-stay prediction, we select prediction time point every 12 hours, begining at the fourth hour after the patient is admitted to the ICU and ending with the patient’s discharge or death. $t=320$h, $n$=4 and the task aims to predict the remaining time spent of patients in ICU. As shown in Table.3, TSCAN realize a reduction of three points compared to LSTM and more than forty points compared to logistic regression in MAD. For the metric of KAPPA, TSCAN also improves by 0.009 compared to channel-wise LSTM.

For phenotype classification, $n=4$, the prediction time point is described as the end of the ICU stay (the patient is discharged or dies), and the time-series data of $t=320$h are used to classify phenotype in order to make full use of the data during the patient’s stay in the ICU. In other words, for phenotype classification, those less than 320 in length in the temporal dimension are padded with zero, and those that are redundant are removed. As shown in the Table. 4, it can be seen that TSCAN improves by 0.056 in macro-averaged AUC-ROC and 0.040 in micro-averaged AUC-ROC over logistic regression. Compared to LSTM, the experiment shows an improvement of 0.025 in the macro-averaged AUC-ROC and 0.018 in the micro-averaged AUC-ROC. This shows that the internal Fusion-Encoder module also works well for problems with more than one label.

For decompensation prediction, $n=4$, we select a prediction time point every hour, 4h after the beginning of the ICU stay, and $t=320$h. And the results are shown in the Table 5. It can be seen that TSCAN has an improvement of 0.043 in the AUC-ROC metric and 0.112 in the AUC-PR metric compared to Logic Regression, and 0.021 in the AUC-ROC metric compared to LSTM. And we have achieved significant performance benefits of 0.7% compared to other SOTA methods (channe-wise LSTM) in ACU-ROC. Even when the data is pre-processed to get a time-equal sequence, it is clear that our method can still do a great job.

We take data from the initial 48 hours of ICU stays as the sample for the prediction task. Early prediction of hospital deaths can help identify patients who are at high risk and give medical staff an alert. In the next subsections, we apply this task to define some parameters of a prediction model that will be utilized in other tasks.

This task is to predict remaining time spent in ICU at every 12 hours of stay, which is framed as a classification problem with 10 buckets (one for ICU stays shorter than a day, seven day-long buckets for each day of the first week, one for stays of over one week but less than two, and one for stays of over two weeks). With this work, ICU beds can be used more efficiently and patients can be cared for better.

We also focus on the contribution of the selected 155 indicators to the task. As shown in Fig.7, the horizontal axis shows the 15 variables with the highest proportion of the 155 variables used in the experiment, and the vertical axis indicates their contribution to the prediction of length-of-stay. From these results, we can see that ALT is responsible for the most weights, which is close to 0.03.

Phenotyping has applications in cohort construction for clinical studies, comorbidity detection and risk adjustment, quality improvement and surveillance, and diagnosis [45]. Based on the ICU data, the phenotyping classification task is to figure out which type of acute care the patient is in. There are 25 common conditions selected for this paper, including 8 chronic conditions that are common co-morbidities and risk factors in ICU such as essential hypertension, 12 severe conditions that are relatively more dangerous to life such as pneumonia, and 5 mixed conditions that are recurrent or chronic, with periodic acute episodes, as shown in Table.7.

The attention map of clinical indicators is shown in Fig.8. The weights of the top 15 variables are shown on the horizontal axis, and their contributions to the task of phenotyping classification are shown on the vertical axis. Several indicators such as the esophageal echo, height, PA catheter, and PICC line accounted for a higher weighting of over 0.008, which suggests that medical staff could pay more attention to those factors when judging the phenotype.

The goal of this task is to figure out which patients are likely to get much worse in the next 24 hours. This setting is currently manually designed with early warning scores and thresholds below which alerts are triggered, so most of these scoring systems are based on simple thresholds and a small number of common physiological indicators, such as the National Early Warning Score (NEWS) [46] and the Modified Early Warning Score (MEWS) [47]. Decompensation prediction in the paper aims to detect patients who are physiologically decompensating, or whose conditions are deteriorating rapidly.

Referring to previous work [5], and trying to fit the early warning scoring system as closely as possible, the task of decompensation prediction in our experiment is considered as a prediction of mortality in the next 24 hours for patients in the ICU and is done every hour. So, our experiments can be run on the MIMIC-IV database, which will pull out correct features and labels.

For each hour as a prediction time point, the data is matched to a target label that says whether the patient died within 24 hours of that hour. In contrast to predicting in-hospital mortality and phenotype classification, for a single ICU stay, more than one sample can be made.

The attention map of clinical indicators is shown in Fig.9. The vertical axis shows the contribution ratios of the first 15 variables, which account for the highest weight in the task. The percentages for the top 15 variables are relatively average, ranging from 0.005 to 0.006. The predictions are consistent with clinical reality. For the PICC line, the patient needs intravenous nutrition and rehydration therapy, and for the PA catheter, the patient needs to be checked for complex hemodynamic disturbances. Both of these are treatments for patients who are pretty sick.

On the one hand, the results suggest that patients with IABP on left heart assist devices, arrhythmia drugs such as amiodarone, verapamil, and chest pain might be more likely to deteriorate, while patients with cardiac dysfunction might be less likely to survive. On the other hand, the capillary refill rate and phenylephrine are important predictors of a patient’s prognosis in cases of circulatory failure. This means that stable circulation and enough blood flow to tissues and organs are needed to improve the prognosis of a patient. The above indicates that cardiovascular system disease plays an important role in predicting decompensation.