TAPER: Time-Aware Patient EHR Representation

TL;DR

TAPER is a transformer-based model that creates unified, time-aware representations of heterogeneous electronic health record data, improving performance on clinical prediction tasks.

Contribution

The paper introduces TAPER, a novel transformer-based approach that effectively encodes irregular, multi-modal EHR data into unified patient representations for predictive modeling.

Findings

Superior performance on mortality, readmission, and length of stay tasks

Effective encoding of irregular, multi-modal EHR data

Generalizes well across different clinical prediction tasks

Abstract

Effective representation learning of electronic health records is a challenging task and is becoming more important as the availability of such data is becoming pervasive. The data contained in these records are irregular and contain multiple modalities such as notes, and medical codes. They are preempted by medical conditions the patient may have, and are typically jotted down by medical staff. Accompanying codes are notes containing valuable information about patients beyond the structured information contained in electronic health records. We use transformer networks and the recently proposed BERT language model to embed these data streams into a unified vector representation. The presented approach effectively encodes a patient's visit data into a single distributed representation, which can be used for downstream tasks. Our model demonstrates superior performance and generalization on mortality, readmission and length of stay tasks using the publicly available MIMIC-III ICU dataset. Code avaialble at https://github.com/sajaddarabi/TAPER-EHR

Tables

Table 1. TABLE I: Code Recall

Method	Diagnosis@k				Procedure@k
	10	20	30	40	10	20	30	40
Ours	${47.92\%}_{\pm \left(0.6\right)}$	${65.11\%}_{\pm \left(0.4\right)}$	${75.25\%}_{\pm \left(0.3\right)}$	${82.00\%}_{\pm \left(0.4\right)}$	${50.99\%}_{\pm \left(1.2\right)}$	${62.25\%}_{\pm \left(0.9\right)}$	${67.97\%}_{\pm \left(1.1\right)}$	${71.52\%}_{\pm \left(1.0\right)}$
Joint-Skipgram [22]	${41.16\%}_{\pm \left(0.6\right)}$	${59.49\%}_{\pm \left(0.7\right)}$	${71.26\%}_{\pm \left(0.7\right)}$	${79.73\%}_{\pm \left(0.5\right)}$	${50.88\%}_{\pm \left(0.9\right)}$	${61.33\%}_{\pm \left(0.7\right)}$	${67.32\%}_{\pm \left(0.9\right)}$	${70.95\%}_{\pm \left(0.7\right)}$
Med2Vec [12]	${42.24\%}_{\pm \left(0.2\right)}$	${60.18\%}_{\pm \left(0.3\right)}$	${70.89\%}_{\pm \left(0.4\right)}$	${78.42\%}_{\pm \left(0.4\right)}$	${47.44\%}_{\pm \left(0.7\right)}$	${58.10\%}_{\pm \left(0.6\right)}$	${64.19\%}_{\pm \left(0.6\right)}$	${68.36\%}_{\pm \left(0.7\right)}$
MCE [6]	${45.70\%}_{\pm \left(0.4\right)}$	${63.43\%}_{\pm \left(0.4\right)}$	${74.44\%}_{\pm \left(0.4\right)}$	${81.14\%}_{\pm \left(0.4\right)}$	${51.87\%}_{\pm \left(0.8\right)}$	${61.62\%}_{\pm \left(1.0\right)}$	${68.42\%}_{\pm \left(0.8\right)}$	${71.97\%}_{\pm \left(0.8\right)}$
Deepr/P2V [13, 15]	${35.30\%}_{\pm \left(0.3\right)}$	${52.40\%}_{\pm \left(0.4\right)}$	${65.04\%}_{\pm \left(0.4\right)}$	${74.14\%}_{\pm \left(0.3\right)}$	${42.47\%}_{\pm \left(0.7\right)}$	${55.59\%}_{\pm \left(1.1\right)}$	${62.42\%}_{\pm \left(0.8\right)}$	${66.49\%}_{\pm \left(0.9\right)}$
Skipgram	${35.23\%}_{\pm \left(0.4\right)}$	${52.30\%}_{\pm \left(0.4\right)}$	${65.08\%}_{\pm \left(0.4\right)}$	${74.13\%}_{\pm \left(0.2\right)}$	${42.56\%}_{\pm \left(0.7\right)}$	${55.81\%}_{\pm \left(1.0\right)}$	${62.61\%}_{\pm \left(0.9\right)}$	${66.70\%}_{\pm \left(0.9\right)}$

Table 2. TABLE II: Downstream Tasks: Readmission

Method	AUC-ROC	PR-AUC
Text+Code+Demo	$67.42\%$	$68.03\%$
Text+Code	$65.74\%$	$65.43\%$
Text+Demo	$61.44\%$	$62.81\%$
Code+Demo	$64.53\%$	$67.68\%$
Text	$56.44\%$	$56.55\%$
Code	$60.74\%$	$57.89\%$
Demo	$54.76\%$	$59.01\%$

Table 3. TABLE III: Downstream Tasks: Mortality

Model	AUC-ROC	PR-AUC
Text+Code+Demo	$63.42\%$	$65.65\%$
Text+Code	$59.75\%$	$60.33\%$
Text+Demo	$61.54\%$	$64.07\%$
Code+Demo	$60.41\%$	$64.41\%$
Text	$57.14\%$	$57.72\%$
Code	$56.91\%$	$53.71\%$
Demo	$60.10\%$	$62.02\%$

Table 4. TABLE IV: Downstream Tasks: Length of Stay Top-1 avg over 5-fold cross validation

Method	Top-1
Text+Code+Demo	$25.57\%$
Text+Code	$23.38\%$
Text+Demo	$21.10\%$
Code+Demo	$22.22\%$
Text	$18.82\%$
Code	$20.54\%$
Demo	$20.13\%$

Table 5. TABLE V: Comparison with other work on downstream tasks.

Method	Readmisison		Mortality		LOS
	ROC	PR-AUC	ROC	PR-AUC	ACC
Ours	$67.42\%\pm \left(2.1\right)$	$68.03\%\pm \left(1.5\right)$	$63.42\%\pm \left(0.6\right)$	$65.65\%\pm \left(0.9\right)$	$25.57\%\pm \left(2.7\right)$
ClinicalBert [21]	$64.41\%\pm \left(2.3\right)$	$67.71\%\pm \left(1.3\right)$	$61.23\%\pm \left(1.9\right)$	$64.83\%\pm \left(1.1\right)$	$22.65\%\pm \left(1.6\right)$
Joint-Skipgram∗ [22]	$64.59\%\pm \left(2.9\right)$	$65.27\%\pm \left(2.7\right)$	$61.73\%\pm \left(1.6\right)$	$65.03\%\pm \left(1.1\right)$	$23.27\%\pm \left(1.6\right)$
SGcode + SGtext∗	$65.79\%\pm \left(1.9\right)$	$66.72\%\pm \left(0.6\right)$	$62.23\%\pm \left(0.6\right)$	$64.77\%\pm \left(1.1\right)$	$24.54\%\pm \left(0.4\right)$
MCE∗ [6]	$64.61\%\pm \left(2.1\right)$	$63.96\%\pm \left(2.1\right)$	$62.72\%\pm \left(2.1\right)$	$64.78\%\pm \left(2.2\right)$	$22.47\%\pm \left(2.3\right)$
Deepr∗ [13]	$62.63\%\pm \left(1.9\right)$	$62.37\%\pm \left(1.5\right)$	$60.49\%\pm \left(1.2\right)$	$61.45\%\pm \left(1.5\right)$	$22.81\%\pm \left(0.5\right)$
Patient2Vec ∗ [15]	$56.13\%\pm \left(3.3\right)$	$65.12\%\pm \left(1.1\right)$	$61.19\%\pm \left(1.4\right)$	$62.51\%\pm \left(1.6\right)$	$21.46\%\pm \left(1.2\right)$
Med2Vec [12]	$63.44\%\pm \left(2.8\right)$	$62.77\%\pm \left(2.2\right)$	$61.67\%\pm \left(1.6\right)$	$63.22\%\pm \left(0.8\right)$	$24.15\%\pm \left(0.4\right)$
Ours - Supervised	$65.38\%\pm \left(3.4\right)$	$66.07\%\pm \left(2.4\right)$	$62.26\%\pm \left(1.7\right)$	$64.32\%\pm \left(2.3\right)$	$25.12\%\pm \left(3.2\right)$

Equations

$$\frac{1}{T}\sum _{t=1}^T\sum _{-w\le j\le w,j\ne 0}{c}_{t+j}^{\top }\text{log}(\widehat{c_t})+(1-c_{t+j}^{\top })\text{log}(1-{\hat{c}}_t)$$

$$e_{c_t}=W^T{c}_t$$

$$E=\{e_{c_1},e_{c_2},\dots e_{c_T}\}$$

$$Attention(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{softmax}(\frac{\mathbf{Q}{\mathbf{K}}^T}{\sqrt{d}})\mathbf{V}$$

$$h_{1:m}=GR{U}_{enc}(U_m,h_0)$$

$$E_{U_t}=\text{softmax}(\frac{{\mathbf{h}}_{1:\mathbf{m}}{{\mathbf{h}}_{1:\mathbf{m}}}^T}{\sqrt{d^{text}}}){\mathbf{h}}_{1:\mathbf{m}}$$

$$\hat{U}=GR{U}_{dec}(U_m,h_m)$$

$$L(U_{1:m,t},{\hat{U}}_{1:m,t})=\sum _i^m(u_i-\widehat{u_i}{)}^T(u_i-\widehat{u_i})$$

$$\text{recall}@k=\frac{\text{\# of relevant codes in top k}}{\text{\# number of relevant code}}$$

Full text

I Introduction

Electronic health records (EHR) are commonly adopted in hospitals to improve patient care. In an intensive care unit (ICU), various data sources are collected on a daily basis as preempted by medical staff as the patient undergoes care in the unit. The collected data consists of data from different modalities: medical codes such as diagnosis which are standardized by well-organized ontology’s like the International Classification of Disease (ICD)111http://www.who.int/classification/icd/en and medication codes standardized using National Drug Codes (NDC)222http://www.fda.gov. Similarly, at various stages of the patient’s care physicians input text noting relevant events to the patient’s prognosis. Additionally, lab tests and bedside monitoring devices are used to collect signals each of which are collected at varying frequencies for a quantitative measure of the patient care. There is a wealth of information contained within EHRs that has a significant potential to be used to improve care. Examples of inference tasks using such data include estimating the length of stay, mortality, and readmission of patients [1, 2].

The traditional approach for healthcare analysis has mainly focused on classical methods for extracting hand-engineered features and designing rule-based systems. More recently, deep learning has demonstrated state of the art results on a varying set of tasks, in which learning intermediate representation is at the heart of all these analysis [3]. This representation can be obtained without domain-specific expertise by leveraging available EHR data. Although such methods have demonstrated great performance on image, audio datasets, leveraging deep learning techniques on healthcare data present new challenges as the data entered are sparse and contain different modalities.

As is common in natural language processing tasks, the typical method for embedding medical codes and text could be through the use of one-hot vectors; though these are naturally high-dimensional and sparse resulting in poor performance. To alleviate this, the idea of learning a distributed representations as applied to natural language processing [4] has been also applied on medical data [5]. Such methods share a common intuition that similar medical codes should share a similar context. Additionally, codes have varying temporal context, as such patients may have multiple visits with a similar set of codes. As an example, flu is short-lived whereas a diagnosis code for a more terminal disease such as cancer has a longer scope and hence will be present on all of the patient’s visits. Due to the varying temporal context, it is also important to take into account the temporal scope of codes and texts assigned [6]. This demands for a model which takes the sequential dependencies of the patient’s visits into account.

To capture the sequential dependencies present in medical data recurrent neural networks (RNNs), and Long Short-Term Memory (LSTM) have become the go-to model. RNN auto-encoder models are commonly augmented with attention mechanisms allowing the model to attend to specific time steps either through soft/hard attentions resulting in improvement and interpretability in the final representation obtained [7, 8]. In NLP tasks such attention mechanisms are not required to be causal in time and hence can attend to both past representation as well as future representations to generate the current representation. However, in a healthcare setting, it is desirable to have the representation be causal in time as clinical decisions are made sequentially. Recently, transformer models [9] were proposed for natural language processing tasks, and have shown impressive results. It uses self-attention and as the model creates intermediate representations of the input it attends to its representation at previous and future timesteps when considering the present representation.

Majority of patient representation work has solely focused on embedding medical codes or text as a patient representation for downstream tasks but not both. To address this, we study the use of transformer networks to embed structured medical code data as well as a language model to embed the text portion of visits. In this work, we propose to combine the medical representation from text and medical codes into a unified representation which can then be used for downstream prediction tasks. Lastly, the presented study takes into account the temporal context of a visit and embeds subsequent visits given the patient’s history. In the following sections, we briefly go over related works and present our method followed by experimental setup and results.

We have made our code and preprocessing steps available333https://github.com/sajaddarabi/TAPER-EHR.

II Related Works

The idea of learning embeddings for sparse one-hot vector data types using back-propagation was presented in [10]. Follow up work learned these embedding using neural networks Bengio et al.[11]. Since its success in NLP Mikolov et al.[4] for language modeling, similar approaches have been used in the health domain. The two intuitive methods that are commonly used for word embedding are (1) skip-gram where a current word is said to predict surrounding words, (2) continuous bag of words (CBOW) where a set of words are made to predict a center word. In *Choi et al. * [12] medical codes (diagnosis, procedure, and medications) are concatenated as one-hot vectors and embedded using the skip-gram model. The intuition behind skip-gram model is: codes in a visit should be predictive of its surrounding immediate patient visit codes as well. They also present an additional code loss term as regularization to the objective. It follows the intuition that codes in a visit should also predict one another. Similar to [12], Nguyen et al. [13] use the concatenation of code representation and apply 1D convolutional network to obtain a visit representation. Follow up work augment these methods with external ontology’s and attention on such external data sources when learning the representation [14, 15].

More recent work takes advantage of the hierarchical structure present in medical codes as they are assigned to a patient. For example in [16], the final visit representation is created hierarchically: first codes are embedded at a treatment level where a set of medication/procedures codes predict diagnosis codes, followed by diagnosis level where the representations at this level are made to predict next visits codes. Empirically the method can learn from a small set of samples and outperform earlier methods as presented on their proprietary datasets. Although these methods achieve reasonable results they do not explicitly model temporal context. This is important as certain clinical codes are short-lived whereas others could be long-lived or be permanent. As a result, certain codes should not be regarded as the context of one another although they may occur in the same visit or subsequent visit. This has been studied in Cai et al.[6] where they train a CBOW model with temporal attention on the code representations. Most of these methods are concerned with medical code embedding and disregard physician notes which can potentially attribute to an improved representation.

Medical notes contain a vast amount of information but have not been studied adequately, especially for downstream tasks in the medical setting. Most works have focused on clinical concept embedding instead. In a clinical setting, nurses and doctors document patients progress. As notes are typically recorded using medical jargon, they do not necessarily follow the common grammatical structure found in English text. As such building, a representation from medical text using hand-engineered features is a challenge [17]. For example in [rajkomar2018scalable] the authors use unstructured EHR data and learn semi-supervised patient representation which are then evaluated on downstream tasks such readmission, mortality and length of stay. Along this line of work, in [18] text from EHR are used to embed patient text by predicting billing codes and averaged for downstream tasks using neural networks. Similarly, in [19], the authors evaluate different models for embedding clinical notes such as CNNs, LSTMs and evaluate them on chronic disease prediction. Although they showed good results their models are not expressive enough to capture all of the salience present in clinical text. To this end, a recent model namely Bidirectional Encoder Representation Transformers (BERT) language model presented by Devlin et al.[20] have recently outperformed previous methods on many benchmarks. This model was used in [21] to embed medical notes of patients. They evaluate their representations predictive performance on downstream tasks showing state of the art results compared to other methods.

Few works have studied the combination of both text and clinical codes. Previous work [22], trains skip-gram and word2vec models to jointly embed clinical concepts and clinical text into a unified vector. Similarly, in [23] they use clinical text to predict clinical concepts. Although both text and code are taken into account, they are different from transferring the joint representation to downstream tasks that is the focus of our work.

III Method

The objective of this work is to create a distributed representation for a patient based on text and medical codes. This representation is then fed to classifier for predictive analytics tasks such as mortality, length of stay, and readmission (Fig. 2). We split the training into two steps, (1) Skip-gram model using transformer networks to learn medical code representation, (2) a BERT model is trained on medical notes and the resulting representations at a time step are summarized using auto-encoder architectures [24, 25, 26]. The final representation for a patient is a concatenation of these two. We discuss the approach in more detail in the following subsections.

To present the problem setting, the sequence of EHR data under consideration consists of a finite set of medical concepts $\mathcal{C}=\mathcal{M}\cup\mathcal{D}\cup\mathcal{P}$, where $\mathcal{M}$ is the set of medication codes, $\mathcal{D}$ is the set of diagnosis codes, and $\mathcal{P}$ is the set of procedure codes. Accompanying the codes are medical notes $\mathcal{T}$. We denote a patients longitudinal data as $\mathcal{D}^{T}=\{(c_{0},t_{0}),\cdots(c_{T},t_{T})\}$ with $T$ visits where $c_{i}$ and $t_{i}$ correspond to the codes and texts assigned respectively within the same visit window.

III-A Medical Code Embedding

We use skip-gram model to learn code embeddings as it is able to capture relationships and co-occurence between codes. We briefly review skip-gram model presented by Mikolov et al.[4]. Given a sequence of codes $\{c_{1},c_{2},\dotsc,c_{T}\}$, where each code vector is a binary vector $c_{t}\in\{0,1\}^{|\mathcal{C}|}$, the model is tasked to predict the neighbouring codes given a code $c_{t}$. The objective can be written as

[TABLE]

Here $w$ is the context window, and the softmax function is used to model the distribution $p(c_{t+j}|c_{t})$. We use multiple transformer encoder layers with self-attention mechanism as the model for skip-gram. This model is then trained on medical code sequence $S=\{c_{0},\cdots c_{T}\}$ by stacking code vectors into a matrix $K\in\{0,1\}^{T\times|\mathcal{C}|}$. The resulting set of codes are then converted into a set of embedding codes $e_{c_{t}}\in\Re^{d}$ using an embedding matrix $W\in\Re^{|\mathcal{C}|\times d}$. The embedding for the set of codes $c_{t}$ at visit $t$ is obtained as

[TABLE]

As the model does not contain any recurrence or convolution, to enable the model to make use of the ordering we need to inject information about the relative positioning of each embedding. This is done by adding to each embedding position a sinusoidal with frequency as a function of its timestamp $t$ as suggested by the original transformer network. This signal acts as positional-dependent information which the model could use to incorporate time. The model is summarized in Fig. 3. We stack multiple transformer layers following on top of the embedding matrix. By transformer layer, we mean a block containing the multi-head self-attention sub-layer followed by feed-forward and residual connections. For more details on this refer to [9] and the tensor2tensor library444https://github.com/tensorflow/tensor2tensor. As multi-head attention can attend to future time steps, to ensure that the model’s predictions are only conditioned on past visits, that is embedding at time step $t$ can only attend to previous time steps $t-1,t-2\dotsc$, we mask the attention layers with a causal triangular mask. This is the same ”masked attention” in the decoder component of the original transformer network. This mask is applied to the set of embedding

[TABLE]

in the encoder block to ensure causality. Where the query $Q$, key $K$, and value $V$ are set to the sequence of embeddings $E$, and $d$ is the embedding dimension. To obtain the final code representation for timestep $t$, the $t^{th}$ output of the self-attention output is used. We call this code representation $E_{c_{t}}$.

III-B Medical Text Embedding

An overview of the text module is shown in Fig. 4. The embedding for the medical text sequence of a patient for time sequence $\{t_{1},t_{2},\dotsc,t_{T}\}$ is obtained by using a pre-trained BERT model initialized from BioBERT [27] followed by a bidirectional GRU as a text summarizer. This is done, as the pre-trained BERT model has a fixed maximum sequence length of $n$ limiting the sequential scope of the text. Further, medical notes could get very lengthy during a visit and they contain different types of notes such as nurse notes, pharmacy notes, discharge notes, etc. the aggregate length of these at a time step $t$ could surpass the fixed length size of $n$. As the aim is to obtain a single visit representation, the aggregate notes at until time step $T$ are batched into a set of sentences $(u_{1},u_{2},\cdots,u_{m})$, where $u_{i}\in\mathcal{Z}^{n}$ and $m$ is the maximum occurring length in the corpus after batching each visits text into sentences of $n$ words. The resulting set of sentences for a visit is embedded using the BERT model which results in a matrix $U\in\Re^{m\times d_{BERT}}$.

[TABLE]

Subsequently, the set of sentence representations are summarized into a single patient text representation using a text-summarizer module. This module follows an auto-encoder architecture with GRU’s as the building block. The input is a set of sentence representations $\in\Re^{d_{text}}$ obtained by the BERT module applied on the aggregate text until time $t$ followed by a self-attention head on the hidden representations, where the decoder is tasked to output a sentence representations $\in\Re^{d_{text}}$ following the bottleneck. The objective of the summarizer is to reduce the MSE loss objective between the input sequence of the text embeddings and the models predicted representation at the correspondingly same time-step. Finally, the patients visit text representation is obtained by summarizing the set of sentence representations $\{U_{1},U_{2},...,U_{m}\}$, where the output of the attention head applied on the encoders hidden representations is used as the representation.

III-C Patient Representation

The final patient representation $Z^{t}$ at time $t$ is obtained by concatenating the code, and text representations. Additionally, the demographics $d_{t}$ of the patient recorded in visit at time $t$ is concatenated to the resulting vector. The demographics of patient contain information such as age, gender, race, etc, where categorical values are coded as one hot vectors. The final representation is denoted as $Z^{t}=[E_{c_{t}};E_{U_{t}};d_{t}]$, where the size of this vector is the sum of the components $d_{embedding}+d_{enc}+d_{demographics}$. This representation is used for downstream tasks.

We provide specific values for the dimensions of each component in our implementation details in the experiments section.

IV Experiments

IV-A Dataset

We evaluate our model on the publicly available MIMIC-III clinical Database [28]. It consists of EHR records of 58,976 hospital admission consisting of 38,597 ICU patients from 2001 to 2012. On average, each patient has 1.26 visits. The database contains tables associated with different data, where we extracted demographics, medical codes, and medical notes.

IV-A1 Readmission task

The first task is to predict 30-day unplanned readmission to the ICU after being discharged. In this task, we formulate it as a binary problem, that is to predict whether a patient will be readmitted within 30-days after being discharged. Text entered into the MIMIC database, contains different reports, such as nurse notes, lab results, discharge summaries etc. We limit the text for each visit to contain the discharge summaries or text entered within the last 48h before the patient is discharged in the absence of discharge summaries.

IV-A2 Mortality task

The second task is to predict mortality of patient, whether they passed away after being discharged or within the ICU. Similar to readmission it is formulated as a binary task. In this task, mortality related codes are discarded from the dataset and patients who were admitted for organ donations are removed. Additionally, the input text for each visit is limited to the first 24h of the admission.

IV-A3 Length of stay task

: The third task is to forecast the length of stay (LOS) for patients. In this task, longer LOS is an indication of more severe illness and complex conditions. We formulate this problem as a multi-class classification problem by bucketing the length of stay into 9 classes: 1-7 correspond to one to seven days respectively, 8 corresponds to more than 1 week but less than 2, 9 corresponds to more than 2 weeks. The model is tasked to predict $P(Y=L|Z_{t})$ where $L\in{1,2,3,4,5,6,7,8,9}$ denoting the different time intervals defined previously.

IV-A4 Code prediction task

: In this task, clinical codes are predicted for new admissions of patients given past clinical codes and historical data of the patient. The predicted vector is high-dimensional equal to the size of unique codes.

IV-B Data preprocessing

We extracted procedure, and diagnosis codes for each patient visit. These codes are defined by the International Classification of Disease (ICD9) and medications using the National Drug Code (NDC) standard. The total number of ICD9 codes in MIMIC-III is 6984, the number of drug codes is 3389, and the number of procedure codes is 1783. Codes whose frequency are less than $5$ are removed. We used the Clinical Classification Software for ICD9-CM555https://www.hcup-us.ahrq.gov/toolssoftware/ccs/ccs.jsp to group the ICD9 diagnosis codes into 231 categories. The Clinical Classification Software for Services and Procedures666https://www.hcup-us.ahrq.gov/toolssoftware /ccs_svcsproc/ccssvcproc.jsp was used to group the procedure codes into 704 categories. Additionally, patients of age under 18 were removed from the cohort. As medical notes contain many errors, we correct grammatical errors and remove non-alphanumerical characters. The text preprocessing closely follows [21]. After preprocessing, the average recorded number of codes per visit is 20.52, the average number of words in medical notes is $7898$ and the average number of visits per patient is 1.29. The statistics of the compiled cohort are depicted in Fig.5

IV-C Experimental Setup

IV-C1 Model Configuration & training details

The training follows the method discussed in Section III. A medical concept model is trained independently on clinical codes in an unsupervised manner and similarly, the text summarizer is trained on the text portion.

To train the transformer encoder we explored different values for hyperparameters, namely number of layers, number of multi-head attention heads $n_{head}$, dimension for each head $d_{head}$ and the final model representations $d_{code}$. We found $2$ layers perform well. We set the models representation dimension ($d_{code}$) to $128$. Further, the self-attention module contains 8 head ($n_{head}$), each with dimension 64 ($d_{head}$), which are a common configuration used in transformer networks. The network is trained using Adam [29] with a cosine annealing schedule and with a period of 50 epochs. The initial learning rate is set to $0.00025$. The window size for the skipgram objective is set to $2$.

We initialize the BERT model with the pre-trained weights on medical nodes as presented in [30]. In this work, the BERT language model is initialized with BioBERT which is a model that has been trained on a large corpus of public medical data such as PubMed, medical abstracts, etc. Then the model was fine-tuned on the MIMIC-III clinical notes.

To train the text summarizer we use a 2-layer bidirectional GRU autoencoder with the intermediate representation set to $d_{enc}=128$. A teacher-forcing ratio of $0.5$ is used with a step learning rate schedule decay of $0.1$ every $50$ epochs with initial lr set to $10^{-3}$ [31]. We have also tried using cosine annealing schedule, though this did not result in improvements.

Lastly, the classifier for downstream tasks is a simple 2-layer fully connected network. The first layer contains $\frac{d_{code}+d_{text}+d_{demographics}}{2}$ neurons with ReLU activation followed by a layer which maps to number of classes in the downstream task. When training on downstream tasks only the classifier weights are trained for 30 epochs with a step learning rate schedule decay of $0.1$ every $10$ epochs. This setting is used for all downstream tasks.

IV-C2 Implementation details

We implemented all the models with Pytorch 1.0 [32]. For training the models we use the Adam optimizer [29]. In all experiments, the batch size is set to 32 on a machine equipped with 1 NIVIDIA 1080TI CUDA 9.0, 32GB Memory & 8 CPU cores.

IV-D Evaluation Metrics

The compiled cohort consists of patientuids as keys that are unique to each patient which are used to create the test/train split. This is done using k-fold with $k=7$ on the patient keys, which the unsupervised models are trained on the train portion and validated on test ($\approx 15\%$ of total data). The output for a particular time step is evaluated using the patient representation $Z_{t}$ at time $t$.

IV-D1 Area under the precision-recall (AU-PR)

this metric is the cumulative area under the curve by plotting precision and recall while varying the outputs $P(y_{t}=1|Z_{t})$ true/false threshold from 0 to 1.

IV-D2 Receiver operating characteristic curve (AU-ROC)

this metric is the area under the plot of the true positive rate against false positive rate while varying outputs $P(y_{t}=1|Z_{t})$ true/false threshold from 0 to 1.

IV-E Baselines

We compare our model with the following baselines

• •

Med2Vec [12]:

A multi-layer perceptron is trained on medical codes using the skip-gram objective function on a visit basis. An additional loss term is used for the co-occurrence of codes within the same visit as a regularization. The resulting output is a set of code representations in $\Re^{d}$.

• •

ClinicalBERT [21]:

A BERT model is pre-trained on public medical data, which is then fine-tuned on clinical text. Following this pre-training a BERT classifier is initialized with pre-trained weights and further fine-tuned on downstream tasks. The input to this network is text.

• •

Time-aware Embedding [6]:

A multi-layer perceptron is trained on medical codes with an additional attention layer to take into account the temporal context of medical codes. The resulting model is trained using either skip-gram/CBOW.

• •

Patient2Vec [15]:

In this work a sequence of medical codes are embedded using word2vec model. The sequence of visits with irregular time intervals is then binned into a set of subsequences with standard intervals. Subsequently, the embedded vectors are stacked into a matrix where convolution stacked with GRU and attention models are applied to obtain the final patient representation.

• •

Joint-Skipgram[22]:

The embeddings are trained using both text and code as the vocabulary. In addition to the traditional skipgram loss, i.e. codes in the same visit predicting surrounding codes or text predicting surrounding text; the skipgram objective is modified such that text in a visit predict codes in the same visit and vice versa.

• •

Deepr[13]:

A set of clinical codes are embedded using skip-gram model. As visits contain multiple codes, the vectors corresponding to each code is stacked into a matrix, then the set of matrices for each visit is fed to a convolutional neural network and max-pooling layers to extract the final patient representation.

• •

$\bf Sg_{code}+Sg_{text}$:

Embeddings for both code and text are learnt using the skipgram objective independently. Subsequently for downstream tasks a patient representation is obtained by concatenating the code and text embeddings.

• •

Supervised:

The BERT model and summarizer takes as input raw text and transformer model raw codes, which is trained jointly on downstream tasks without pretraining.

We do not compare with more traditional text embeddings such bag of words (BOW) as other work have shown the benefits of using BERT as text representation in NLP tasks.

We study the effect of different components presented by adding/removing text/code/demographics representation to our final patient visit representation.

V Results

To show the expressiveness of our representation, we evaluate its performance compared to baseline methods on downstream tasks and unsupervised learning tasks presented in the experiment section. We present the results obtained by embedding both text and code as patient representation on three downstream tasks: (1) 30-day readmission, (2) mortality, and (3) length of stay (LOS). Note that in MIMIC-III database clinical codes are entered into the database upon discharge of a patient, as consequently the presented results may not be immediately clinically actionable in this case. Although, this may not be the case in other datasets where codes are updated throughout a visit. To this end the codes on the same visit are not fed to the model, but rather the previous set of codes are used.

V-A Code pre-training

The encoder network is trained on clinical concepts and as most patients have 3 hospital visits after filtering to atleast 2 visits the window size for the loss term is set to 2. The performance of our network is compared with baseline methods using recall@k. This metric is evaluated by computing

[TABLE]

This metric mimics a practitioners method of arriving at a diagnosis or prescribing medications where they generally have several sets of candidates as a presumed cause for the underlying condition of the patient. As shown in Fig. 6 our code embedding consistently outperforms the baselines. All baselines are fine-tuned on the same corpora by exploring different architectural hyper-parameters, except the embedding size which is fixed to $d_{code}=128$ for all models. A comparison on recall for different codes (diagnosis, procedure) is also reported in Table. I. From the results the time-aware code representations outperform other baselines.

V-B Ablation Study

To evaluate the different components of the proposed method we conduct an ablation study on the inclusion/exclusion of the components in the final representation for downstream tasks. The complete representation using demographics, text representation and clinical code representation is the concatenation of these on a visit $Z^{t}=[E_{c_{t}};E_{U_{t}};d_{t}]$.

V-B1 Readmission

To better evaluate the effect of the different embedding components, we run an ablation study. The results are reported in table II. ≈

From the results, it can be seen both text and code are informative for classifying readmission. The complete combination of text, code, demographics outperforms others in this case.

V-B2 Mortality

Mortality task is concerned with predicting whether a patient will pass away within a pre-defined window. We predict mortality on a visit basis i.e. does the patient pass away in the current visit to the ICU. An ablation is done in table III. Similar, to previous ablations, the combination of text, code, and demographics outperforms other combinations.

V-B3 Length of Stay

In general length of stay is a much more challenging task compared to readmission binary task. In this classification task, we limit the medical text to the first 24-hours of the current patient visit in which length of stay is being predicted. Limiting the note context window is done as medical text could include information on date patient has been discharged. In this task the network is trained on imbalanced class data split and tested on balanced data, this is done as the majority of classes are discharged within 24h-48h.

As shown in table IV, the combination of text, code and demo outperforms others. We conclude in this set of experiments, solely using text and diagnosis codes is not predictive enough to forcast the length of stay of patients.

V-B4 Comparison with Other Work

We ran all baseline models on three downstream tasks using the same pre-processing steps. The final results are reported in table V. To compare our method, we use the combination of text, code, and demographics. As demographics has proven to have predictive power for the downstream tasks, other methods are augmented to use this as input/concatenated to the output representation prior to prediction. From Table. V the presented method outperforms others by a 1-2% margin, on all tasks demonstrating the usefulness of unsupervised pre-training.

VI Conclusion

Effective representation learning for EHR data is an essential step to improving care. We study embedding both medical codes and notes into a unified vector representation for downstream task prediction. The presented method effectively takes the temporal context of these two data streams and provides a patient visit representation. The proposed method was evaluated on three tasks namely, readmission, mortality, and length of stay outperforming other methods. An ablation study was also done showing the usefulness of both text and code when modeling patient visits. Future work could focus on adding additional data streams to the pipeline by taking into account real-time vitals and measurements taken from a patient as they undergo care.