Multi-TA: Multilevel Temporal Augmentation for Robust Septic Shock Early Prediction

1 Introduction

A disease progression model (DPM) aims to estimate how a target disease would progress over time based on historical data such as multivariate time-series electronic health records (EHRs) [Mould, 2012]. One important purpose of DPM is to detect diseases as early and accurately as possible since robust and accurate early prediction can assist timely clinical interventions [Zhou et al., 2013; Li et al., 2015; Che et al., 2015; Choi et al., 2016; Choi et al., 2017a] and reduce the risk of mortality as well as the burden on patients and the healthcare system [Kumar et al., 2006; Wang et al., 2014]. While recent deep learning-based early prediction models have demonstrated great success by utilizing RNNs, CNNs, or Transformers [Lipton et al., 2015; Birkhead et al., 2015; Choi et al., 2016; Esteban et al., 2016; Luo et al., 2020; Rasmy et al., 2021; Yang et al., 2021], one of the major shortcomings of the existing approaches is that they do not perform consistently well on EHRs as n increases (i.e. tradeoff between earliness and accuracy) [Achenchabe et al., 2021; Bondu et al., 2022]. In this work, we refer to such phenomena as the lack of temporal robustness.

Typically, EHRs contain patient trajectories, which are represented by a single hospital visit by an inpatient: a visit is a series of events recorded at irregular times; an event consists of a set of data entries (e.g. measurements of vital signs) collected during a certain period of time. As shown in Figure 1 (left), our goal is to predict whether a patient will develop a disease n hours later by leveraging the patient's EHRs in the observation windows until n hours before the onset of the disease or the end of the sequence. Figure 1 (right) shows that as we increase n from 6 to 48 hours, the early prediction performance (red line) deteriorates as observed in other previous works [Lin et al., 2019; Zhang et al., 2019; Khoshnevisan and Chi, 2021]. It is primarily caused by the fact that as n increases, data scarcity problems become more prominent, due to the patient behavior of visiting a hospital when their symptoms become noticeable. More specifically, we identified two levels of data scarcity: the event-level scarcity, decreasing number of events per trajectory (green line in Figure 1) make it more difficult for models to predict the following behavior of diseases or relevant medical signals within the hold-off n hours window. The trajectory-level scarcity refers to as n increases, the number of patient trajectories decreases (blue line) and causes insufficient diversity in the training population to learn possible variations of disease progression across patients.

In this work, we propose a temporal EHRs augmentation (Multi-TA) framework. Figure 2 shows that Multi-TA consists of two key components, one per source of scarcity: (1) Representation learning module (left) is designed to address event-level scarcity by acquiring knowledge about the temporal dynamics of events or patient health status within the holdoff window and fill in the missing pieces; (2) Data augmentation module (right) is devised to address trajectory-level scarcity. By leveraging the representation learning module (G) and optimizing the level of data transformation (θ) within a regular training process, the module generates synthetic trajectories that are realistic but extreme for the prediction model in a way that the model can learn sufficient variations of disease progression and become temporally robust.

We validate Multi-TA on the task of septic shock early prediction due to its critical significance in healthcare and its difficulty. However, our framework is generalizable to diverse tasks involving multivariate irregular time-series data, particularly those facing challenges similar to EHRs. Sepsis, a life-threatening medical condition resulting from a dysregulated body response to infection, can lead to the most severe complication known as septic shock, characterized by high mortality rates and prolonged hospitalization [Singer et al., 2016]. Timely treatment is crucial, as each hour's delay in antibiotic administration increases the mortality risk by 8% [Kumar et al., 2006]. Additionally, sepsis, with diverse etiologies like cancer, presents a wide range of syndromes, and different patient groups may exhibit distinct symptoms [Tintinalli et al., 2011]. Our study, utilizing realworld data from two US medical systems, demonstrates that Multi-TA can enhance the temporal robustness of early prediction models. There have been a few lines of research that aim to improve the robustness of prediction models. One utilizes EHR representation learning models trained with largescale unlabeled data [Li et al., 2020; Rasmy et al., 2021; Pang et al., 2021] while the other employs data augmentation to address the data scarcity problem from the root by generating extra labeled data [Esteban et al., 2017; Che et al., 2017; Baowaly et al., 2019; Poulain et al., 2022]. However, to our knowledge, no prior work has attempted to unify both worlds for one goal - that is, combining temporal knowledge learned from representation learning with synthetic samples generated from data augmentation to improve temporal robustness.

To summarize, our contributions are: (1) By tackling two levels of data scarcity inherent in EHRs, Multi-TA can build effective early prediction models that are more temporally robust; (2) By integrating a pretrained EHR representation model into the data augmentation process, Multi-TA generates realistic but challenging synthetic time-series data that directly improve temporal robustness; (3) Multi-TA outperforms the baselines on two real-world EHR datasets with various settings for an extremely challenging task.

2 ProposedMethod

Problem Description Our dataset, denoted as Dtrain = {(x1, y1), ..., (xN, yN)}, captures inpatient visits from N different patients, where each visit corresponds to a patient trajectory. The data comprises multivariate irregular time series, with each visit xk being a sequence of events: xk = {x1k , ..., xTk k }, where xtk denotes patient's records at timestamp t and Tk is the number of events in k-th visit, which varies across different visits. Each xtk ∈ RS represents medical data entries collected from clinical measurements at each event, with S being the number of entries. For each visit xk, we have an associated output label yk = {1, 0}, indicating septic shock or non-septic shock, respectively. In addition to the labeled dataset, we leverage an unlabeled dataset Dunlabeled = {x1, ..., xM} for representation learning. This dataset consists of hospital visits that have not been annotated due to budget constraints. The primary objective in early prediction is to learn a prediction function F over Dtrain that best approximates an unknown function f : X → Y , where X represents the true distribution over the entire population. However, when there exists insufficient amount of patient trajectories (small N) in Dtrain or scarce number of events (Tk) that represents each trajectory, it would be difficult to approximate f and temporal robustness cannot be guaranteed.

2.1 Multi-TA: Temporal EHRs Augmentation

We propose Multi-TA, a framework that combines representation learning and data augmentation to address two types of data scarcity. Specifically, as described in Figure 2, MultiTA consists of two stages: (1) Learning EHR representations, which learns the temporal interaction between events within both observation and hold-off window to address event-level data scarcity; (2) Augmenting patient trajectories, which generates sufficiently different synthetic timeseries sequences using constrained worst-case optimization to address trajectory-level data scarcity.

Learning EHR Representations

The goal of this stage is to fill the missing knowledge in the training dataset Dtrain caused by the event-level scarcity. Given that the events within a hold-off window are unknown (refer to Figure 1), the training dataset cannot hold any information from that specific time period, potentially leading to uncertain predictions. In this stage, we utilize the unlabeled dataset Dunlabeled, which spans both observation and hold-off windows, to broaden our knowledge and tackle the eventlevel scarcity.

Step 1: Temporal Abstraction Inspired by [Sohn et al., 2020], we first convert multivariate irregular EHRs into sequences of high-level interval-based concept representations using temporal abstraction [Shahar, 1997], to reduce the random noise inherent in EHRs and implicitly alleviate the effect of irregular time intervals and missing values while preserving temporal information. Specifically, clinical measurements included in the events within each 60 minutes interval are aggregated, and as a result, a visit x = {x1 , ..., xt } composed of t events is transformed into x = {x′1 , ...x′k } in which k < t and x′ = {w1, ...wl} is a list of temporally ordered l abstracted tokens representing a patient health status within each time interval. A token w can be seen as a word concatenating a feature F (e.g., temperature, blood pressure, white blood cell count, etc.) and its state, namely, discretized value V (e.g., low, normal, high, etc.). For example, a token "SystolicBP=H" indicates a certain symptom where systolic blood pressure maintained as high during the time interval. In sum, a visit x can be seen as a document consisting of k events, and an event x′ can be considered as a sentence with l abstracted tokens. The difference between a document in natural language and our visit in EHRs is that ours involves temporal information: (1) temporally ordered token based on the measurement start and end times of each token and (2) time gaps between two consecutive events (x′k-1 , t-hours, x′k ), which could infer hidden knowledge with regard to the progression of a target disease.

Step 2: EHRRepresentation Learning Based on the similarity between our data and texts in natural language, we adopt the original BERT [Devlin et al., 2019] architecture to pretrain EHR representations G. We leverage its ability to understand the context within given sequences, but further, we inject our task-specific temporality into the model. First, we utilize the order of temporally sorted tokens as input to position embeddings to better understand how tokens temporally interact in terms of disease progression. Note that the measurement start and end times can be additionally incorporated in the embeddings, but in this work, we only capture the temporal order. Second, motivated by previous works [Nguyen et al., 2017; Pang et al., 2021] and to fully incorporate temporal aspects of our sequences, we introduce an additional special token ([t-HR]) which represents time intervals between consecutive events. Features (e.g., vital signs) in EHRs are measured irregularly based on their needs and clinicians make such decisions to effectively diagnose patients' conditions and keep track of disease trajectory. That being said, by incorporating time intervals into model training, learned representations could embrace useful information regarding clinicians' practice of measuring patients' conditions. We expect that this would play an important role in understanding the temporal relations between tokens within a visit and benefit the self-attention process inside the BERT framework. In sum, our temporality-injected BERT (t-BERT) utilizes visitlevel input sequences (see Figure 2) that contain (1) the temporal order of each token, which helps the model understand temporal interactions relevant to disease progression, and (2) time gap information between consecutive events, which provides implicit information about patient health status to learn effective EHR representations. Lastly, we inherit "masked language modeling" procedure from the original BERT paper [Devlin et al., 2019] in which the authors mask some percentage of the input tokens at random and then predict those masked tokens based on the encoder output.

Augmenting Patient Trajectories

The second stage of Multi-TA aims to provide sufficient variations of disease progression to a prediction model by augmenting patient trajectories in Dtrain. Motivated by a textual data augmentation work [Sohn and Park, 2022], MultiTA generates synthetic visit sequences that can overcome trajectory-level scarcity with two key steps (see Algorithm 1): (1) constrained worst-case data transformation (lines 5-8), which generates sufficiently different synthetic samples by controlling local changes to a set of selective tokens within a trajectory; (2) robust risk minimization (line 9-11), which interchangeably incorporates newly generated synthetic samples into training.

Step 1: Constrained Worst-Case Data Transformation To provide sufficient variations to a prediction model, we leverage the masked language modeling (MLM) objective [Devlin et al., 2019] for data augmentation, which generates variants of a given sentence by masking and predicting a subset of tokens. However, Multi-TA further optimizes this

• Algorithm 1: Trajectory-level Augmentation Input: Training data Dtrain = {xi, yi}ni=1, pretrained EHRs model G, prediction model F, model weights W, loss function L, learning rate η Parameters : Sampling probability p, transformation level θ, upper bound U Output: Prediction model F with trained weights W 1 Randomly initialize model weights W; 2 while termination criterion not met do 3 Select a random sample pair (x, y) fromDtrain; 4 With probability p: 5 Initialize θ0 ← 0; 6 for i ∈ {1, ..., N} do 7 max θ L(y,F(G(x, θi))), θ ≤ U 8 xθ∗ ← G(x, θ∗ ) 9 W ←W- η∇WL(y,F(xθ∗)) 10 With probability 1-p: 11 W ←W- η∇WL(y,F(x))

type of local change with a special parameter, transformation level θ that controls the number of tokens and which tokens to transform. By controlling such components, Multi-TA can generate sufficiently different samples, which can mitigate the data scarcity problem. As per line 7 in Algorithm 1, we aim to estimate an optimal θ that can maximize the given prediction model's loss under the upper bound U. This is to generate specific synthetic samples that can maximize the diversity of training dataset at a maximum capacity without any harm (e.g., not flipping its class label). Given the importance of each token in predicting a target disease is different (e.g., a serious condition "WBC=VeryHigh" primarily contributes to the septic shock diagnosis, while other tokens are less important factors), we carefully select tokens to transform aligning with our objective. Specifically, as our goal is to generate a sufficiently different samples with a minimum number of transformations, we select tokens that contribute the most to the target class prediction to quickly maximize the loss value. Note that by selecting the most contributing tokens, instead of the least contributing ones, we can achieve our goal by transforming a minimum number of tokens, and this can prevent the model to generate clinically infeasible samples. Since the MLM predicts a masked token based on its context, if there exist a large number of masked tokens and only a few real tokens are left in context, it would be difficult to reconstruct and form a realistic sequence. However, it is non-trivial to directly calculate token-level contribution scores as EHRs are a multilevel structure similar to documents and we apply sequential model such as LSTM on top (one cell per event) and tBERT on the second level. Figure 3 illustrates our prediction model and the way to calculate token-level contributions in two steps. First, we determine visit-level contribution scores using gradient attribution [Simonyan et al., 2014], that is, calculating the gradient per event and normalize across events to identify top θ ∗ n events in terms of the score, where n is the number of event in a visit. Second, once target events are selected, we use the attention scores of tokens with respect to a special token ([CLS]) as contribution scores, within each selected event. As the [CLS] token contains all information collected from other tokens in a sequence and the token is often used as input for various classification tasks, we assume that the attention scores would represent the contribution level of tokens. Based on the per event token-level scores, we select top θ ∗ l contributing tokens per event where l is the number of tokens in an event. Once the tokens in the chosen events are replaced with [MASK] tokens and all the events in a visit are concatenated to form a visit sequence, a pretrained representation model G predicts the tokens based on the context to generate a realistic visit sequence. At the end of the process, the transformed sequence with the estimated theta is added to the training dataset to update the model weights.

Step 2: Robust RiskMinimization Multi-TA actively and interchangeably generates informative samples and uses them to strengthen a prediction model during each epoch of training. Specifically, this step is analogous to the regular stochastic gradient descent of any differentiable prediction models (lines 9 and 11). Given either the original training samples or transformed samples and their corresponding class labels, Multi-TA updates the model parameters W based on the gradient computed and the learning rate η.

3 Experiment Setup

3.1 Two EHRDatasets

We utilize two real-world EHR datasets collected from 210,289 and 106,844 adult patient visits to Christiana Care Health System (CCHS) and Mayo Clinic (Mayo), respectively. Both datasets comprise 2.5 years (07/2013-12/2015) of anonymized and institutional review board (IRB)-approved EHRs. From the total population, 52,919 visits and 4,224,567 events are selected as our study population to only include the patients with suspected infection, based on the rules designed by the two leading clinicians with extensive experience. Time irregularity and data sparsity are the two major challenges inherent in our datasets. In the study population, time intervals vary from 0.94 seconds to 64.38 hours due to the different measuring frequencies of features, and this causes data sparsity where on average more than 80% of the values are missing. The EHRs from two hospital systems differ in several ways, but the distinct measuring frequencies make the most difference in learning temporal dynamics of EHRs. Compared to Mayo, where the majority of measurements (82%) are carried out within every 1 hour, CCHS measures less frequently - i.e. more distributed time intervals: 1-hour (39%), 2-hour (20%). In addition, to ensure a balanced dataset and avoid bias, we truncated samples from the negative class (non-septic shock) at each prediction hour. We opted not to utilize ICU benchmark datasets such as MIMIC-III and PhysioNet based on our preliminary experiments, which revealed no performance deterioration over time (AUC: 0.961 for 24-hour prediction, 0.980 for 48-hour) and demonstrated near-perfect performance. We believe that this is because of MIMIC-III's higher-quality continuous monitoring, typical of ICUs, which introduces a different type of problem.

3.2 Model Evaluation

We validate our framework in two stages: comparing t-BERT with three representation models for multivariate time-series EHRs, and assessing Multi-TA against mixup and GAN-based augmentation methods. We utilize LSTM for prediction, except with CEHR-GAN-BERT, due to its ability to capture long-term dependencies.

Task Setup Our task involves predicting the onset of septic shock in a patient m hours after receiving their last n hours of Notes: Best results for section are in bold and the best overall results have asterisk. (RL: representation learning, DA: data augmentation)

medical records. Here, m ranges from 6 to 48 hours, based on input from clinicians. Shock visits are aligned to their initial onset, while non-shock visits are truncated accordingly. Nonshock visits are further balanced to match the length distribution of shock visits. Given sepsis's rapid progression, our analysis focuses on the initial five days of patient records.

RepresentationModels (1) EHR, raw (continuous) EHRs with the expert rule-based imputation [Kim and Chi, 2018]; (2) MuLan [Sohn et al., 2020], Skipgram-based static model which takes temporally abstracted visit sequences without time gap information; (3) BERT [Li et al., 2020], BERTbased contextualized model which takes MuLan's input; (4) t-BERT (Ours), BERT-based model that incorporates taskspecific temporal information into training.

Data Augmentation Models (1) mixup [Zhang et al., 2018], linear interpolation-based augmentation; (2) Manifold Mixup (M.Mixup) [Verma et al., 2019], variation of mixup, which combines intermediate vector representations; (3) CEHR-GAN-BERT (GAN) [Poulain et al., 2022], GANbased data augmentation approach which utilizes EHR representation model CEHR-BERT [Pang et al., 2021]. For a fair comparison, we substituted t-BERT for CEHR-BERT; (4) Multi-TA (Ours), our proposed EHRs augmentation algorithm which controls transformation level to generate sufficiently different but informative samples to predictions.

Evaluation Strategy We measure model temporal robustness using two evaluation methods: (1) Raw prediction performance is evaluated using five metrics (accuracy, precision, recall, F1-score, and AUC) with a focus on F1 and AUC due to their balanced measure; (2) For lower-bound performance, we identify the lowest performance across all prediction hours, denoted as -LB. A more temporally robust model exhibits higher lower-bound performance. In addition, we report the mean and standard deviation of model performance from six experimental trials (three repetitions for two-fold stratified cross-validation) with critical differences [Benavoli et al., 2016] between models. Hyperparameters, such as p and U, are chosen based on the validation performance.

Evaluation Sets (Regular vs. Robust) We use two test sets to gauge the model's sensitivity to changes in training data size. "Regular" is a commonly used test set for early prediction, in which the numbers of patients in training, validation, and test sets decrease as prediction hour n increases. "Robust" uses the same group of patients across varying prediction hours for model evaluation, in which only the number of training and validation data changes.

4 Result and Discussion

We present two experimental results. 48 hours early prediction reveals the model's static robustness while 6-48 hours prediction uncovers the temporal robustness of the model.

4.1 48 Hours Early Prediction

Representation Models Table 1 (RL sections) shows the performance for four representation models. Firstly, BERT consistently outperforms EHR and MuLan on every metric across both datasets, with approximately a 10% increase in F1 and AUC for both (9.3% of AUC increase in CCHS). This demonstrates BERT's effectiveness in learning contextualized and subtle information inherent in EHRs. Moreover, our proposed t-BERT outperforms BERT and exhibits lower variance across metrics and datasets, except for recall in CCHS. This suggests that incorporating time intervals into training enhances understanding of temporal dynamics in patient trajectories. Notably, the performance improvement fromBERT to t-BERT is more pronounced in CCHS (AUC increase: 3% in CCHS, 1.7% in Mayo), indicating that t-BERT may be more effective on more irregularly measured EHRs.

Data AugmentationModels Table 1 (DA sections) reveals the experiment result for four data augmentation models. Data augmentation models shows mixed results when compared to their counterparts, either EHR or t-BERT. While mixup applied to EHR slightly improves the model performance (AUC increase: 1.9% in CCHS, 0.2% in MAYO), M.Mixup and GAN applied to t-BERT rather degrade the performance in most cases. We hypothesize that this underperformance is derived from their architectural design in which (1) they generate intermediate latent vector representations for augmentation, which may only contain condensed information instead of actual visit sequences; (2) they generate representations without conditioning on class labels and regardless of prediction objective. Furthermore, GAN does not capture temporal dependencies between events due to the lack of LSTM layer and rather depends on its BERT model to capture the information. On the other hand, our proposed Multi-TA consistently outperforms the other models across all metrics and datasets except for recall in CCHS, where GAN performs the best with 88.8%. Further, Multi-TA shows smaller variance than others, demonstrating its robustness to the data changes. This validates the efficacy of optimizing transformation levels in relation to a target prediction task.

4.2 6-48 Hours Early Prediction

Representation Models Figure 4 shows the model performance (F1) in the robust setting across varying hold-off window sizes, ranging from 6 to 48 hours at 6-hour intervals. The plots exhibit consistent patterns across both datasets: (1) t-BERT outperforms baseline models across different hours, except for 18 hours with Mayo data; (2) t-BERT maintains stable performance from 6 to 48 hours, showing minimal performance decline and indicating robustness to training data Notes: Best results for section are in bold and the best overall results have asterisk. (RL: representation learning, DA: data augmentation)

scarcity. Similar to 48-hour prediction, the biggest performance gap is observed between EHR/MuLan and BERT, emphasizing the importance of capturing contextual information in EHRs. The consistent improvement from BERT to t-BERT reveals the significance of incorporating time gap signals to comprehend temporal dynamics of disease trajectories. Table 2 (RL sections) presents detailed model performance for both settings and lower-bound metrics that measure the worst performance over time. Firstly, compared to the regular setting, the robust setting exhibits slightly lower performance. We hypothesize that the patient group in the robust setting may pose greater challenges possibly due to vague symptoms. Secondly, t-BERT consistently performs better with CCHS across all metrics, as observed in Figure 5. This underscores the significance of modeling time irregularity in disease early prediction and highlights the effectiveness of our proposed t-BERT. Lastly, t-BERT outperforms other models in all settings including robust and lower-bound metrics, establishing it as the most temporally-robust among RL models.

Notes: Models with no statistical difference (P≥.05) are linked.

Data Augmentation Models Figure 5 compares F1-scores for data augmentation models in the robust setting, where Multi-TA demonstrates more stable and superior performance over time compared to others. Specifically, GAN exhibits unstable performance, degrading even with more training data (as hour n decreases), highlighting its ineffectiveness in ensuring temporal robustness. M.Mixup produces relatively stable predictions and performs better than t-BERT on Mayo, but its inconsistent performance across varying prediction hours and datasets suggests sensitivity to changes. In contrast, Multi-TA consistently improves performance over time with its target-oriented sample generation, even when n is large and training data is the scarcest. Table 2 (DA sections) underscores the temporal robustness of the four data augmentation models. Across eight prediction hours, Multi-TA outperforms other models consistently, with higher lower-bound scores (F1-LB and AUC-LB), indicating greater reliability, especially for early prediction tasks. Finally, Figure 6 demonstrates that Multi-TA significantly outperforms other models.

4.3 Analysis of Data Augmentation

Quantity of Data Multi-TA determines optimal data augmentation levels. Unlike methods that assumes more random data enhances model performance, our approach augments training data specifically tailored to the current model by adding fewer samples. Specifically, a sampling probability p in Algorithm 1 controls the augmented dataset size efficiently. Figure 7 addresses the question of "how many new samples are sufficient?" and demonstrates that excessive data addition (≥ 0.25) rather harms performance.

Quality of Data Multi-TA optimizes the quality of augmented data by determining the optimal transformation level θ to maximize diversity. For example, larger θ values indicate more extensive token transformation, increasing distance from the original sample. Figure 8 (left) displays the distributions of estimated θ values for both datasets, which are neither uniform nor skewed, showing Multi-TA's ability to adaptively set transformation levels for target datasets and prediction models, unlike fixed or random approaches.

Transformation, Where? By design, Multi-TA transforms most contributing time steps within a visit. Figure 8 (right) shows the distribution of the indices of transformed time steps on its x-axis, and we denote the most recent time step as 0. As expected, the results show that more recent time steps (which are close to 0) were transformed more, indicating that they contributed more to predictions and increasing diversity.

5 RelatedWork

Representation Learning for EHRs Given the significance of contextualized vector representations, [Richardson et al., 2020; Tonelli et al., 2018], BERT [Devlin et al., 2019] has gained widespread adoption for EHRs, including variants like BEHRT [Li et al., 2020] incorporating patient ages to imply temporal code orders, G-BERT [Shang et al., 2019] enhancing clinical context with a graph neural network, Med-BERT [Rasmy et al., 2021] utilizing domain-specific pretraining and a large dataset (20M patients), and CEHRBERT [Pang et al., 2021], akin to our t-BERT, training a temporal BERT with visit type prediction and time-indicative tokens. Despite these advancements, prior works may not effectively capture rapid disease progression like sepsis due to discrete and coarse-grained medical codes [Lee et al., 2020], nor address temporal robustness in early prediction tasks. While recent advances in Large Language Models (LLMs) offer promising capabilities for learning representations from medical corpora [Singhal et al., 2023], their application to specialized tasks like septic shock prediction remains limited.

DataAugmentation for EHRs While time-series data augmentation shows promise [Wen et al., 2021; Iglesias et al.,

2023], deep learning-based methods for healthcare prediction remain limited. Firstly, simple yet effective linear interpolation techniques like mixup [Zhang et al., 2018] and Manifold Mixup [Verma et al., 2019] can augment both continuous and discrete inputs by combining a pair of training samples, and the generated samples act as a regularization that enhances model robustness. Secondly, GAN-based approaches are divided into two groups: (1) generation-focused models for preserving privacy in data [Mogren, 2016; Choi et al., 2017b; Esteban et al., 2017; Yoon et al., 2019; Baowaly et al., 2019; Li et al., 2021]; (2) prediction-focused model with semisupervised learning [Che et al., 2017; Yu et al., 2019; Cui et al., 2020], but both approaches have not shown their effectiveness in improving temporal robustness for early disease prediction. CEHR-GAN-BERT [Poulain et al., 2022] shares similarities with our approach by integrating a EHR representation model with a discriminator for robust predictions. However, it lacks capturing temporal dependencies and generates only unlabeled vector representations.

6 Limitations and FutureWork

Time Complexity Multi-TA's iterative transformation process increases training time. Training for 48-hour early prediction spans around 3-min for EHR, 7-min for t-BERT, and 100-min for Multi-TA. We expect enhanced efficiency with parallel processing like distributed optimization.

BERT as a Transformation Function When the majority of tokens in a visit is masked, it is challenging for BERT to fill in the spots due to the lack of context. However, we expect that causal language models such as GPTs can address this issue and generate more diverse and creative patient trajectories.

Comparison with Existing Approaches Multi-TA has not yet been comprehensively compared to other existing methods for septic shock prediction or handling time gap information [Ma et al., 2020; Wei et al., 2023]. We plan to conduct further experiments to facilitate a comparative analysis.

Binary Prediction We followed established literature precedents [Gao et al., 2022; Yang et al., 2023; Fleuren et al., 2020; Yan et al., 2022], which use binary prediction models with manually designated windows up to 48-hour. Yet, we admit that a regression model capable of predicting the exact onset hour would significantly benefit healthcare providers.

7 Conclusion

Accurate early disease predictions face challenges due to inherent data scarcity at trajectory and event levels. This paper introduces Multi-TA, a unified temporal augmentation framework, which addresses these challenges by integrating EHR representation learning with data augmentation. Multi-TA mitigates event-level scarcity by learning temporal dynamics from unlabeled events and enhances trajectory diversity by generating new samples via constrained worst-case transformations. Experiment results on two real-world EHR datasets reveal the temporal robustness of Multi-TA compared to various state-of-the-art models, affirming its efficacy in improving early prediction for enhanced clinical decision-making.

Acknowledgments

This research was supported by the NSF Grants: #1726550, #1651909, and #2013502.