A Framework for Early Sepsis Prediction via Self-Supervised (JEPA) and Federated Representation Learning
Machine Learning
Summary
The authors studied different computer models to predict sepsis early using patient health records, which are often messy and incomplete. They tested four main approaches, including self-supervised methods that learn useful patterns without labeled data, a semi-supervised method, and a fully supervised one. Their best model nearly matched top benchmarks while using fewer health markers, showing efficient prediction. They also found that self-supervised pretraining combined with fine-tuning made the models' predictions more stable over time compared to other methods. This means their approach can better predict sepsis earlier and maintain accuracy over different time frames.
Authors
Umair bin Mansoor, Munaf Rashid, Roomi Naqvi
Abstract
Early sepsis prediction from electronic health records is challenged by irregular sampling, high missingness, and class imbalance. We systematically compare four modeling paradigms -- self-supervised Joint Embedding Predictive Architecture (JEPA) via masked latent prediction, self-supervised VICReg (variance-invariance-covariance regularization) with two-view augmentation, semi-supervised fine-tuning of a VICReg-pretrained encoder, and supervised Temporal Convolutional Network (TCN) -- alongside raw-feature baselines. All models share a common preprocessing pipeline of hourly binning with forward-fill imputation applied to 7 biomarkers selected via sparsity analysis from the MIMIC-III dataset. Our best model (JEPA + XGBoost + mean pooling) achieves AUPRC 0.636 at the time of onset (H0), approaching the SupMix benchmark (0.667) while using 83\% fewer biomarkers. The Tier 1 pipeline -- VICReg pretraining followed by semi-supervised fine-tuning and XGBoost -- achieves AUPRC 0.510 at H0, a 3.1$\times$ improvement over the raw-feature baseline (0.165) and a 7.6\% improvement over the end-to-end supervised TCN (0.474). Crucially, the fine-tuned VICReg encoder exhibits the most temporally persistent representations, degrading only 16.8\% from H0 to H10 compared to 47.5\% for supervised TCN and 65.3\% for JEPA, demonstrating that self-supervised pretraining with task-aware fine-tuning yields features that are both sharp near onset and robust across prediction horizons.