Event-Based Contrastive Learning for Medical Time Series

TL;DR

The paper tackles risk stratification after clinically meaningful index events in heterogeneous medical time series. It proposes EBCL, an event-centric contrastive pretraining method that forms positive pairs from pre- and post-event windows around domain-informed index events and contrasts them against cross-patient negatives, shaping a temporally aware latent space. EBCL demonstrates superior downstream performance on heart failure and MIMIC-IV ICU tasks (mortality, readmission, LOS) and yields embeddings conducive to patient subtyping, including clear clustering of phenotypes. By aligning the pretraining objective with clinically significant events, EBCL improves risk stratification and offers a practical, workflow-compatible approach for health AI applications.

Abstract

In clinical practice, one often needs to identify whether a patient is at high risk of adverse outcomes after some key medical event. For example, quantifying the risk of adverse outcomes after an acute cardiovascular event helps healthcare providers identify those patients at the highest risk of poor outcomes; i.e., patients who benefit from invasive therapies that can lower their risk. Assessing the risk of adverse outcomes, however, is challenging due to the complexity, variability, and heterogeneity of longitudinal medical data, especially for individuals suffering from chronic diseases like heart failure. In this paper, we introduce Event-Based Contrastive Learning (EBCL) - a method for learning embeddings of heterogeneous patient data that preserves temporal information before and after key index events. We demonstrate that EBCL can be used to construct models that yield improved performance on important downstream tasks relative to other pretraining methods. We develop and test the method using a cohort of heart failure patients obtained from a large hospital network and the publicly available MIMIC-IV dataset consisting of patients in an intensive care unit at a large tertiary care center. On both cohorts, EBCL pretraining yields models that are performant with respect to a number of downstream tasks, including mortality, hospital readmission, and length of stay. In addition, unsupervised EBCL embeddings effectively cluster heart failure patients into subgroups with distinct outcomes, thereby providing information that helps identify new heart failure phenotypes. The contrastive framework around the index event can be adapted to a wide array of time-series datasets and provides information that can be used to guide personalized care.

Figures (11)

• Figure 1: Event-based contrastive learning (EBCL). For patient $i$ and event $j$, we denote pre-event data, $L_i^j$, and post-event data, $R_i^j$. For EBCL pretraining, we sample a batch, $\mathcal{B}$, of pre and post-event trajectory pairs from a dataset, $D_{\mathcal{B}}=\{(L_i^j, R_i^j): i,j \in \mathcal{B}\}$. We choose the event of interest to be an inpatient admission. $L_i^j$ and $R_i^j$ are passed separately into a transformer encoder $f_\theta$ to get $\bar{L}_i^j = f_\theta(L_i^j)$ and $\bar{R}_i^j = f_\theta(R_i^j)$ which is pretrained with CLIP contrastive loss. The positive pairs are pre and post-event data of the same event, $(\bar{L}_i^j, \bar{R}_i^j)$. The negative pairs are mismatched pre-event and post-event trajectories from different patients, such as $(\bar{L}_i^j,\bar{R}_l^m)$ where $i \neq l$.
• Figure 2: Clustering and outcome risk subtyping from pretrained EBCL Embedding.(a) Clustering and dimensionality reduction of pretrained embedding showed distinct heart failure clusters with unique outcome prognoses. Clusters are sorted by the prevalence of 1-Year Mortality (cluster 1 is the fatal group with the highest mortality rate, while cluster 6 is the healthiest group). The survival curve of the identified clusters plotting (b) time to mortality and (c) time to readmission
• Figure 3: Contrastive Learning Framework of EBCL, EBCL Ablations (EBCL Censored, EBCL Outpatient) and OCP.
• Figure 4: XGBoost requires a fixed size vector as input. We create this vector from time series data as follows: we restrict to the 128 most prevalent features in the data and the relative time of those observations (relative to the time of the decision). For each of the time windows prior to the time of the decision. For the heart failure cohort, windows end at discharge time for 1-Year Mortality and 30-Day Readmission and admission time for 7-Day LOS. For the MIMIC-IV cohort, windows end at the event time (i.e. hypotension or mechanical ventillation onset time). We perform feature aggregations (mean, count, min value, max value) for each feature over all pre-defined time windows. We then concatenate all of the feature aggregation outputs across all windows to get a final vector that is the input to XGBoost. We have 5 window sizes and 4 aggregations each generating a 129 length vector. After concatenating these, we get a 2,580 size vector that is the input for XGBoost.
• Figure 5: OCP Positive and Negative Pairs. In OCP, positive pairs are sequences with correct order, while negative pairs have the order of the first and second halves of the timeline swapped. The model, $f_\theta$ is pretrained to predict whether the sequence was swapped.