Unimodal Contrastive Pretraining Of Clinical Timeseries

2 minute read

EHR Modeling with Contrastive Learning: A New Approach to Pretraining Clinical Time Series Data

By Ryan King, Shivesh Kodali, Conrad Krueger, Tianbao Yang, and Bobak J. Mortazavi
Published in Machine Learning for Health

In the fast-evolving world of machine learning, the challenge of harnessing Electronic Health Record (EHR) data remains monumental. These datasets, rich with time-series measurements from ICU patients, hold immense potential for advancing clinical decision-making, personalized medicine, and predictive modeling. However, the roadblock has always been the reliance on large amounts of labeled data—expensive and time-consuming to curate.

To address this, our research introduces an efficient contrastive unimodal pretraining method for EHR time series data. This innovative approach not only reduces computational demands but also sets a new benchmark in learning meaningful representations from unlabeled data.


The Problem with Existing Approaches

State-of-the-art methods like supervised deep neural networks (DNNs) demand vast amounts of labeled data. While self-supervised learning alleviates this dependency, most techniques require large batch sizes to ensure performance, creating a computational bottleneck when processing long ICU time-series data. Furthermore, irregular sampling and missing values in EHRs make modeling even more complex.


A Novel Solution: Combining Contrastive and Masked Pretraining

Our method innovatively marries two complementary self-supervised tasks:

  1. Masked Token Imputation
    Inspired by NLP, we treat each EHR measurement as a token in a sequence. By masking only the values (leaving type and time information intact), the model learns to predict missing measurements, addressing data sparsity and irregularity.

  2. Contrastive Sequence Learning
    By generating augmented “global” and “local” views of the same sequence, the model aligns broader patterns with fine-grained details. A variance-reducing estimator enables efficient learning, even with smaller batch sizes.


Key Advantages of Our Method

  • Scalability: Handles long sequences and diverse measurement sets without needing imputation or retraining.
  • Flexibility: Learns from any subset of features, adapting seamlessly to new clinical environments.
  • Efficiency: Reduces computational overhead by leveraging a smaller, GPU-friendly architecture.

Results That Matter

We evaluated our approach on two open-source ICU datasets, MIMIC-III and eICU. Highlights include:

  • In-Hospital Mortality Prediction: Achieved an AUC-ROC of 0.854 and AUC-PR of 0.468, outperforming both traditional and state-of-the-art methods.
  • Phenotyping and Imputation: Demonstrated robust performance on phenotyping tasks and precise imputation of missing measurements.
  • Transferability: Models trained on MIMIC-III successfully transferred knowledge to eICU, proving their generalizability across clinics.

Our results show that pretraining with our combined method yields superior representations, even in semi-supervised settings with limited labeled data.


Real-World Implications

This research bridges a critical gap in EHR modeling, making advanced analytics more accessible to clinics with limited resources. By reducing the reliance on labeled data and enabling efficient computation, our method paves the way for real-time clinical insights, more equitable healthcare, and scalable AI solutions.


Final Thoughts

Our work is a step forward in the journey toward universal, robust foundation models for healthcare. As machine learning continues to integrate into clinical workflows, approaches like ours will play a pivotal role in unlocking the true potential of EHR data.

Explore the full paper on GitHub or on Arxiv

Updated: