{'fraction_non_alphanumeric': 0.03262213486351964, 'fraction_numerical': 0.012078749643344006, 'mean_word_length': 4.755154168947272, 'pattern_counts': {'":': 5, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
TRACE: Early Detection of Chronic Kidney Disease Onset with Transformer-Enhanced Feature Embedding

Chronic kidney disease (CKD) has a poor prognosis due to excessive risk factors and comorbidities associated with it. The early detection of CKD faces challenges of insufficient medical histories of positive patients and complicated risk factors. In this paper, we propose the TRACE (Transformer-RNN Autoencoder-enhanced CKD Detector) framework, an end-to-end prediction model using patients' medical history data, to deal with these challenges. TRACE presents a comprehensive medical history representation with a novel key component: a Transformer-RNN autoencoder. The autoencoder jointly learns a medical concept embedding via Transformer for each hospital visit, and a latent representation which summarizes a patient's medical history across all the visits. We compared TRACE with multiple state-of-the-art methods on a dataset derived from real-world medical records. Our model has achieved 0.5708 AUPRC with a 2.31% relative improvement over the best-performing method. We also validated the clinical meaning of the learned embeddings through visualizations and a case study, showing the potential of TRACE to serve as a general disease prediction model.


Introduction
Chronic kidney disease (CKD) is a general term for many heterogeneous diseases that irreversibly alter kidney structure or cause a chronic reduction in kidney function (Levey and Coresh 2012). It is defined by the presence of kidney damage, or decreased kidney function, or both, for a minimum of three months. Diagnosis of CKD is often made after accidental findings from screening laboratory tests, or when the symptom already becomes severe (Webster et al. 2017). According to the US Centers for Disease Control and Prevention, approximately 15% of US adults have CKD, but most people may not feel ill or notice any symptoms until CKD is advanced (CDC 2019). Two major causes of CKD are hypertension and diabetes. Kidney failure is the most serious outcome of CKD and severe conditions can only be treated by dialysis and transplantation, which are indications of end-stage kidney disease. CKD is a well-known risk factor for cardiovascular disease and all-cause mortality (Abboud and Henrich 2010;Tonelli et al. 2006). However, less than 2% of CKD patients finally require renal replacement therapy, because many of them die from cardiovascular causes before end-stage kidney disease can occur (Keith et al. 2004).
CKD has a poor prognosis due to excessive risk factors and comorbidities associated with it (Kronenberg 2009). The early detection of CKD, for health benefits, is an even more challenging task. For the purposes of early detection and control, a promising direction is to regularly monitor risk factors of CKD for high-risk patients. It is also worth trying to focus the screening for CKD on younger, healthier populations, although it is less likely to detect CKD in such cohorts (Tonelli et al. 2006). In short, early detection of CKD faces challenges of insufficient medical histories of positive patients and complicated risk factors from a data standpoint. This calls for more effective machine learning prediction models to address these issues.
To address the first challenge, we need a prediction model that can better extract knowledge from the insufficient medical histories. In recent years, deep learning has emerged as a powerful tool to gain insight into EHR data Ching et al. 2018;Miotto et al. 2018;Shickel et al. 2017). Previous studies predicted heart failure onset with recurrent neural networks (RNNs) (Choi et al. 2017c) and the reverse time attention mechanism (Choi et al. 2016c). These models proposed different ideas for sequence modeling, but they did not perform very well in our experiments. Apart from sequence modeling, we also need a better feature representation for the early detection of CKD onset. This ties in with the second challenge -complicated risk factors.
Given how extensive the risk factors are when assessing a patient's likelihood of having CKD, it is critical to learn an embedding to measure the latent similarity between these risk factors and CKD. An embedding maps discrete medical concepts to a continuous latent space and summarizes the interactions between medical concepts. Several models can learn embeddings of medical concepts (Choi et al. 2016b(Choi et al. , 2020. However, these models have yet to make the most of the sequential nature of EHR data because they only learned medical concept embeddings for individual hospital visits. We can use a sequence aggregator, such as RNN or 1-D convolutional neural networks (CNNs), to link a sequence of visits and encode them to a patient-level representation.
Motivated by the challenges and the previous work, we propose the TRACE (Transformer-RNN Autoencoder-enhanced CKD Detector) framework, which combines ideas of both RNN autoencoder and Transformer (Vaswani et al. 2017;Choi et al. 2020). TRACE jointly learns the Transformer-encoded hidden structure of individual hospital visits while predicting CKD onset with RNN. In this work, we used Cerner Health Facts (DeShazo and Hoffman 2015), a large database derived from EHR systems across the US, as our data source. The database includes comprehensive patient-level details of diagnoses, procedures, medications, and laboratory tests. The EHR data are structured data without clinical notes and images. Our features include both medical codes and non-medical code information. We summarize our contributions as follows: • To the best of our knowledge, this is the first work that builds advanced sequential deep learning models to predict CKD onset.
• We propose a Transformer-RNN autoencoder architecture. This autoencoder jointly learns a medical concept embedding via Transformer for individual hospital visits, as well as a patient embedding via an RNN encoderdecoder structure which summarizes the entire medical history of this patient.
• We adopt two pre-training processes to capture proper latent representations for patients' conditions and hospital visit histories respectively. We also demonstrated that TRACE successfully alleviated the aforementioned challenges by incorporating the pre-trained latent representations.

Related Work Deep Learning in Healthcare Domain
Deep learning algorithms have become popular approaches for modeling disease progression (Choi et al. 2016a;Ma et al. 2017;Choi et al. 2016c;Lipton et al. 2016), patient characterization (Baytas et al. 2017;Che et al. 2015), and generating synthetic EHR data for research purposes (Choi et al. 2017b). The most common application of modeling disease progression is predicting disease outcomes. Deep neural networks have very limited power when learning disease trajectories from scratch, sometimes it is necessary to incorporate prior medical knowledge (Ma et al. 2018;Pham et al. 2017) or supplement EHR data with inherent hierarchical structure of medical ontologies (Choi et al. 2017a). Missing value is also a challenge is modeling EHR data. (Che et al. 2018) has demonstrated that RNNs are able to capture the long-term dependency in time series and improve prediction performance.
Deep learning models anticipate a large volume of data to achieve satisfactory results, which usually exceeds the capacity of most healthcare facilities. A straightforward solution is to combine EHR data from multiple sources, but data harmonization is a labor-intensive process. (Rajkomar et al. 2018) recently proposed a representation of EHRs based on the Fast Healthcare Interoperability Resources (FHIR) format for deep learning models without site-specific data harmonization. When working on an imbalanced dataset with insufficient positive samples, data augmentation techniques can benefit training. CONAN (Cui et al. 2020) incorporated generative adversarial networks (GANs) to create candidate positive and negative samples in rare disease detection. Pre-training and transfer learning (Bengio 2012;Dauphin et al. 2012) can also help to solve this problem. G-BERT (Shang et al. 2019) used the pre-trained hospital visit representations for downstream predictive tasks. (Rios and Kavuluru 2019) trained a CNN on a large global database with biomedical abstracts, and transferred the learned knowledge to predict diagnosis codes for one medical center.

Representation Learning for Medical Concepts
Representation learning algorithms in healthcare domain are mainly borrowed from natural language processing (NLP). The general idea is to encode discrete medical concepts (e.g., medical codes) to one-hot vectors (Bengio, Courville, and Vincent 2013) and then apply Word2Vec algorithms (Mikolov et al. 2013a) to learn embeddings.
For example, Med2Vec (Choi et al. 2016b) utilized skipgram (Mikolov et al. 2013b) to learn intra-visit medical code co-occurrences as well as inter-visit sequential information. The generic skip-gram model is based on the assumption that a word can play different roles at different positions in a sentence. However, this assumption doesn't hold for medical codes given their unordered nature. When we adopt NLP algorithms to model medical concepts, we typically apply the algorithms to the feature dimension instead of the temporal dimension, and the order in which these medical concepts occur is ignored. MiME  leveraged the inherent structure of medical codes to learn a multilevel embedding of EHR data, but this model required the EHR data to contain complete structure information between diagnoses and treatments. Recently, GCT (Choi et al. 2020) has been proposed to solve this problem. GCT learned the graphical structure of EHR data during training and proved that Transformer is a suitable model to learn such structure. Our work was motivated by GCT to use Transformer to encode hospital visits.

Method Problem Statement
We formulate this problem as a binary prediction task. Given a patient whose medical history is in the form of a sequence of hospital visits P = {V 1 , V 2 , · · · , V T } in chronological order, where T ∈ N is the total number of visits that the patient has. Each visit V t (t ∈ {1, 2, · · · , T }) consists of a list of medical codes, clinical observations and other information related to this patient. We want to predict whether the patient will be diagnosed with CKD for the first time in the following visit V T +1 . Figure 1 illustrates the vector representations of a patient's hospital visit history P, which contains both medical codes  Figure 1: Vector representation of a patient's EHR data. This patient had seven hospital visits at time {t 1 , t 2 , · · · , t 7 }. We used the first six visits {t 1 , · · · , t 6 } to construct input vectors to the model and predicted whether there would be a CKD onset at the seventh visit.

Vector Representations of EHRs
and non-medical code information. For simplicity, all notations and algorithms in this paper are presented for a single patient unless otherwise specified.
Medical Code Representations Medical codes include diagnosis codes, procedure codes, and medication codes. Medical codes are the primary features for our prediction task. We denote the set of medical codes in our EHR data by C = {c 1 , c 2 , · · · , c |C| } with size |C|. All medical codes that occur at a hospital visit V t are represented by a multi-hot vector Non-medical Code Information Besides medical codes, we also included the patient's observations (e.g., lab tests, vital signs, etc.), age, race, gender, and the timestamp of visit V t . Observations are the secondary features in our dataset. Let d t denote the vector representation of the non-medical code information, which is a concatenation of multi-hot vector and numeric values. We will provide details of theses non-medical code features in the "Dataset" section.

Model Architecture of TRACE
In this section, we describe TRACE in detail, with the following components: a patient embedding from a pre-trained Transformer-RNN autoencoder, a medical code history encoder, and a joint attention module. The overall architecture is illustrated in Figure 2.

Patient Embedding via Transformer-RNN Autoencoder
We pre-train a Transformer-RNN autoencoder as a joint feature extractor for both feature embedding and patient embedding. This design adopts a Transformer (Choi et al. 2020) for computing self-correlation of all features at individual hospital visit level, and a subsequent RNN autoencoder for learning a patient representation by reconstructing the input sequence in an unsupervised fashion.
The autoencoder ingests a sequence of hospital visits in the form of and n is the total number of features. The sequence {x ′ t } T t=1 then runs through an encode-decoder structure to reconstruct itself.
• Encoder. We first map each x ′ t (t ∈ {1, 2, · · · , T }) to a latent space with a learnable embedding matrix W x by This extends the raw input vector x ′ t to a vector array for the Transformer to process. Since Transformer has a quadratic time and space complexity and is expensive to compute, we downsize the feature dimension from n toñ (ñ ≪ n) via a linear trans-formZ t =W z Z t ∈ Rñ ×dz . This improves scalability by reducing the complexity to O(ñ 2 ). We learn an embedding for the downsized feature space using a Transformer with one encoder block and a single attention head as follows, where X t ∈ Rñ ×d emb . Positional encoding is removed from our framework, since the features are not ordered.
To aggregate the embedding of all features occurred at time t, we average-pool X t by the downsized feature dimension to obtain a single embedding vector v t ∈ R d emb .
Eventually, we feed the sequence {v t } T t=1 in an encoding RNN layer to encode the entire input sequence to a single-vector patient embedding e p .
• Decoder. We propagate e p to a decoding RNN layer to obtain a decoded sequence {h dec t } T t=1 , then reconstruct the input sequence from {h dec t } T t=1 . Specifically, we add separate fully-connected layers on top of {h dec t } T t=1 to reconstruct different types of features, and obtain a reconstructed sequence {x ′ t } T t=1 . The loss between the input sequence {x ′ t } T t=1 and the reconstructed sequence is the sum of multiple losses. For multi-hot medical codes and observations, we use softmax classifiers with cross entropy loss. For race and gender, we apply sigmoid classifiers with cross entropy loss. For age and timestamp, we apply linear transform and minimize mean squared errors.
• Patient Embedding for TRACE. To obtain a patient embedding for our end-to-end prediction task, we feed the input sequence {[x t ; d t ]} T t=1 in the pre-trained encoder and compute the embedding vector e p for this patient (Figure 2).
We also summarize the detailed pre-training process as an algorithm in the appendix.

Medical Code History Encoder
Considering that medical codes summarize other non-medical code features to some extent, we separately encode the patient's medical code history for information gain in our model.
We first map the discrete medical code inputs x t to a continuous latent embedding space m t as follows, where W m ∈ R |C|×dm is a word embedding lookup table pre-trained via Med2Vec (Choi et al. 2016b), and d m is the size of the embedding vector. To encode a sequence of medical codes, we then apply an RNN layer on top of the medical Figure 2: End-to-end structure of TRACE. The model ingests a patient's medical history in the form of two sequences of vectors X = {x 1 , x 2 , · · · , x T } and D = {d 1 , d 2 , · · · , d T }, next propagates them to the encoder module of a pre-trained Transformer-RNN autoencoder for a patient embedding, then combines the patient embedding with the patient's medical code history to compute joint attention, and finally outputs a probability scoreŷ.
code embedding m t by where h m t is the hidden state of the RNN layer at time t. Joint Attention We want to further have the patient embedding interact with the medical code history. Specifically, we compute interactions between e p and each h m t (t ∈ {1, 2, · · · , T }) as follows, where u is a learable weight, and α t is the attention weight assigned to visit V t . Then we obtain a context vector c for this patient by CKD Onset Prediction We use the context vector c to predict the binary label y ∈ {0, 1} as follows, whereŷ is the predicted probability score for this patient. The training objective is to use the predicted scoreŷ and the true label y to minimize the following binary cross entropy loss: where N is the total number of patients in our training set.

Experiments Dataset
We collected our experimental dataset from two healthcare systems in Cerner Health Facts, each healthcare system comprises multiple healthcare facilities. This is a casecontrol study where negative patients were downsampled through a statistical analysis, such that our model was trained to distinguish positive and negative patients who were similar in terms of age, race and gender.

Features and Data Preprocessing
We extracted a patient's diagnosis codes, procedure codes, medication codes, observations, age, race, gender, and admission date for features of each hospital visit. Statistics of the features are available in Table 1. The raw diagnosis codes, procedure codes and medication codes in Cerner Health Facts are respectively International Classification of Diseases (ICD), Current Procedural Terminology (CPT) and generic drug names. We grouped ICD diagnosis codes by the Clinical Classifications Software (CCS) to obtain higher-level diagnosis codes for experiments, which reduced the number of diagnosis codes from over 69,000 to 275. CKD diagnoses were identified by the CCS codes and were excluded from the feature set. We did not consider hospital visits without any diagnosis codes documented and removed medical codes that appeared in less than 50 hospital visit records.
Apart from medical codes, we also included observations, age, race, gender, and admission date to represent a hospital visit. Observations, race and gender were categorical features encoded to a multi-hot vector (Figure 1)  converted it to a numeric timestamp by calculating the duration in days from the patient's first visit to this visit. We took the logarithm of numeric features for model inputs.

Selection of Cases and Controls
We excluded hospital visits made by non-adult patients because we aimed at predicting CKD onset for adults only. The case/control selection criteria are as follows.
• Cases (positives) were patients who had at least one hospital visit prior to CKD onset. The CKD onset was the positive class label for our prediction task.
• Controls (negatives) were non-CKD patients who had at least two hospital visits in our dataset. Controls were identified for each case using the propensity score-matching based on logistic regression (Rosenbaum and Rubin 1983) and the greedy algorithm (Bergstralh and Kosanke 1995). Matching variables include age, gender and race. Class labels of control patients came from their latest hospital visit records and were negative labels.
Six controls were selected for each case to match the prevalence of CKD in US adults (i.e., 1/7 ≈ 14.29%). Eventually, we extracted a total of 147,791 patients for experiments. The dataset was further split into training, validation and test sets in a 75/10/15 ratio. The case/control ratio in each of the training, validation and test sets was the same as the disease prevalence rate in the entire experimental dataset. Table 1 provides details of the study cohort. Since 90% of patients in the dataset had less than 30 hospital visits, we only kept up to 30 most recent hospital visits per patient to improve scalability.

Pre-training for Medical Concepts
There were two types of medical concepts in our dataset: medical codes and observations. We performed two pretraining processes to get proper embeddings for them.
Independent Pre-training We pre-trained embedding weights for medical codes 1 using Med2Vec. The dataset for this pre-training task was extracted from 10 healthcare facilities in Cerner Health Facts, not including the two healthcare systems for our experiments. We trained the Med2Vec model on 1,155,450 patients with 15,115,251 hospital visits and obtained pre-trained embedding weights for 3,884 distinct medical codes (Table 1). In our prediction task, we treated medical codes outside the independent pre-training as out-of-vocabulary (OOV) tokens and initialized embedding weights for OOV medical codes with zeros.
Transformer-Encoded Embedding We pre-trained the Transformer-RNN autoencoder on our training set to encode all features, where only age, race, gender, and timestamp were not medical concepts (Table 1). This means that the Transformer-encoded feature embedding was a good latent representation of medical concepts. The pre-trained feature embedding was built into TRACE as part of the encoder module for fine-tuning (Figure 2).

Baseline Models
For comparison, we implemented the following models with as the input vector for a hospital visit. • Logistic regression (LR). We counted the occurrences of each medical code and each observation for a patient, all the other features were determined by the patient's last hospital visit in the inputs. A LR model was trained on the resulting vectors.
• Multi-layer perceptron (MLP). We used the same approach to construct model inputs as the LR model, but added a fully-connected layer with relu activation between the input layer and the output layer.
• RNN and BiRNN. We used a fully-connected layer with relu activation to encode inputs and then propagated the resulting vectors to a forward/bidirectional RNN layer. Logistic regression was applied to the last hidden state of the RNN layer to predict CKD onset.
• RETAIN (Choi et al. 2016c). RETAIN model was designed to predicts heart failure onset using backward RNN and two levels of attention weights. We used the same architecture as the RNN baseline, but replaced the RNN layer with the RETAIN module.
• Dipole (Ma et al. 2017). Dipole model predicts multiple disease outcomes via a bidirectional RNN layer and three different attention mechanisms. We used the same structure as the RNN baseline, but replaced the RNN layer with the Dipole module and trained it using each of the three attention mechanisms, i.e., Dipole l , Dipole g and Dipole c .
• 1-D CNN. A modification of AlexNet (Krizhevsky, Sutskever, and Hinton 2012). We replaced all 2-D convolutional layers with 1-D convolutional layers, which served as a sequence aggregator of individual hospital visits. We computed the mean of AlexNet's outputs across the temporal dimension and applied logistic regression on top of it to generate predictions. The inputs were encoded in the same way as the RNN baseline.

Evaluation Metrics
We measured the model performance on our test set by area under the precision-recall curve (AUPRC). AUPRC can effectively evaluate the fraction of true positives among positive predictions (Saito and Rehmsmeier 2015), thus it is an appropriate metric when evaluating binary classifiers on imbalanced datasets like ours. In addition to AUPRC, we also calculated negative log likelihood by Eq. 9 to measure the model loss on the test set.

Implementation Details
We implemented all models and calculated all evaluation metrics using TensorFlow 2.2.0 (Abadi et al. 2015). For Med2Vec, we used the code provided by the authors 2 . The dimension of the downsized feature space wasñ = 100. The sizes of all embedding vectors and hidden layers were 128. The dropout rate for the feed-forward layer of Transformer was 0.5. We used the Adadelta optimizer (Zeiler 2012) and set the learning rate as 1.0 to match the exact form in the paper. We trained each model for 50 epochs with 100 patients per batch. All experiments were run on a 16GB NVIDIA Tesla V100 PCIe GPU.

Results
Prediction Performance We evaluated our model and all baselines on the test set. Table 2 shows the AUPRC and negative log likelihood scores of the test set. In general, sequential models (RNNs and 1-D CNN) outperformed nonsequential models (LR and MLP). This is because sequential models were more capable of capturing the underlying patterns in disease trajectories, while non-sequential models only learned from aggregated information of medical histories. In the real-world clinical practice, doctors need to 2 https://github.com/mp2893/med2vec  carefully review a patient's medical histories and monitor the patient for a long time to decide whether a chronic disease like CKD can be diagnosed. Occasional occurrences of some symptoms related to CKD are insufficient to support the doctor's decision.
It is noteworthy that all RNN-based baselines demonstrated comparable performance in terms of both metrics, and increased model complexity failed to surpass the simplest RNN model. This seems to indicate that training RNN models from scratch is not suitable for our task. Both RE-TAIN and Dipole computed attention scores with the outputs of RNN layers. The attention mechanism assigns a score to each hospital visit using the sequential information learned from scratch, which is fine when detecting diseases in their original tasks. However, CKD is quite different because its excessive risk factors could be intertwined. It is hard to determine whether a patient has CKD simply by the existence of several risk factors without extensive prior knowledge.
TRACE successfully alleviated this problem by introducing a pre-trained Transformer-RNN autoencoder. The autoencoder produced a good patient embedding which compressed information in the entire input sequence. With this patient embedding as the prior knowledge, our end-to-end prediction model was able to better discover the correlation between CKD diagnoses and past medical records. TRACE achieved a 2.31% gain in AUPRC compared with the bestperforming baseline (i.e., Dipole g ).
Ablation Study To understand how each major model component contributed to the overall prediction performance, we compared TRACE with its several variants.
• TRACE base. This is TRACE without the medical code history encoder and the joint attention. We directly used the fine-tuned patient embedding to get predictions.
• RACE. This is TRACE without Transformer-encoded feature embedding. In the RNN autoencoder, we got a feature embedding through a fully-connected layer with relu activation instead.
• RACE base. This is RACE without the medical code his- Figure 4: Medical concept attentions produced by TRACE for a CKD patient in the test set. The labels are the medical concepts occurred at each hospital visit. It illustrates how medical concepts on the vertical axis attended to medical concepts on the horizontal axis. "dx-": diagnosis code (medical code), "lab-": lab test (observation), "vs-": vital sign (observation), "other-": observation other than lab test and vital sign, "BP": blood pressure. This patient got a 0.9198 prediction score with four hospital visits as inputs.
tory encoder and the joint attention. The fine-tuned patient embedding were directly used for getting predictions.
We trained the three variants with the same set of hyperparameters as TRACE. We note the AUPRC scores for the analyses here (Table 3). Overall, pre-trained RNN autoencoders provided richer patient-level information than raw input features. Even the worst-performing model in Table 3 (i.e., RACE) achieved a 0.93% relative improvement in AUPRC over the best-performing baseline in Table 2 (i.e., Dipole g ). Evidently, Transformer has demonstrated its superiority over pure fully-connected layers in encoding medical concepts (RACE vs. TRACE and RACE base vs. TRACE base). As we expected, there was slight information gain after adding medical code histories and the joint attention, but the strength was limited (RACE base vs. RACE and TRACE base vs. TRACE). Figure 3 plots patient embeddings produced by TRACE and the RNN baseline respectively. We used the t-SNE (Maaten and Hinton 2008) algorithm for dimensionality reduction. Obviously, TRACE learned a clearer boundary between positive and negative patients, as well as a better clustering of positive patients. Given that this is a case-control study, TRACE met our expectation to better distinguish cases and controls who were similar in terms of age, race and gender. We also provide an illustration of pre-trained patient embeddings in the appendix. Transformer-RNN autoencoder was able to produce a more gathered patient embedding than the generic RNN autoencoder, showing the strength of Transformer in encoding features.

Patient Embedding Visualization
Attention Visualization and Case Study We visualize the attention behavior of Transformer in the course of CKD onset prediction. Since TRACE computed self-attention for the downsized feature spaceZ t , we need to back-propagate to the original feature space by Z t =W z ⊤Z t to get desired attention weights. To improve readability of the visualization, we randomly selected a CKD patient from the test set, who had four hospital visits and at most 20 medical concepts per visit. Figure 4 illustrates the attention behavior of medical concepts occurred at each hospital visit for the selected patient, which also shows the patient's disease trajectory. This patient had hypertension and diabetes -two major causes of CKD that usually intertwine with other risk factors of CKD. No remarkable attention behavior was present at the first hospital visit. At the second visit, we noticed that esophageal disorders, diabetes and low diastolic blood pressure were mutually attended. At the third visit, the high hemoglobin A1c level and headache attended to each other, indicating a poor blood sugar control and a higher risk of diabetes complications. Eventually, at the fourth visit, the patient had a group of CKD risk factors tested, such as blood urea nitrogen, serum potassium and serum creatinine. The abnormal test results all attended to diabetes, which suggested the correlation between diabetes and CKD. The high serum potassium level also attended to the high serum creatinine level. Moreover, this patient got a true positive prediction with a 0.9198 prediction score.
We also visualize the pre-trained attention behavior for the same patient in the appendix, which was produced by the pre-trained Transformer-RNN autoencoder. The pre-trained attention behavior is similar to the one fine-tuned by TRACE, but some CKD risk factors stood out after fine-tuning.

Conclusion
In this work, we proposed the TRACE framework, a novel end-to-end prediction model that incorporated a pre-trained Transformer-RNN autoencoder for early detection of CKD onset. It is hard to predict CKD onset by training a model from scratch due to the excessive risk factors and insufficient medical histories of positive patients. TRACE alleviated this problem by introducing prior knowledge learned by the autoencoder. Experimental analyses showed that TRACE outperformed all baselines and its several variants in predicting CKD onset. We also validated the clinical meaning of the learned embeddings through visualizations and a case study, which demonstrated the potential of TRACE to be generalized to other disease prediction tasks. In the future, we plan to combine data augmentation techniques like GAN to better address the data insufficiency. We will also adopt more advanced NLP algorithms to train embeddings for patients and features.