Clin-JEPA: A Multi-Phase Co-Training Framework for Joint-Embedding Predictive Pretraining on EHR Patient Trajectories

Area of Science:

• Artificial Intelligence
• Machine Learning
• Biomedical Informatics

Background:

• Joint-embedding predictive architectures (JEPA) excel in robotics and vision but face challenges in Electronic Health Record (EHR) data.
• Existing JEPA methods struggle to create a unified model for both trajectory forecasting and diverse risk prediction without task-specific fine-tuning.
• Naïve co-training of JEPA encoder and predictor leads to instability, representation collapse, and divergent rollouts.

Purpose of the Study:

• To develop a stable co-training framework (Clin-JEPA) for joint-embedding predictive pretraining on EHR patient trajectories.
• To enable a single EHR backbone for simultaneous trajectory forecasting and downstream risk prediction.
• To overcome instability issues in co-training JEPA models for EHR data.

Main Methods:

• Introduced Clin-JEPA, a five-phase co-training curriculum: predictor warmup, joint refinement, EMA target alignment, hard sync, and predictor finalization.
• Stably co-trained a Qwen3-8B encoder with a 92M-parameter latent trajectory predictor.
• Evaluated on MIMIC-IV ICU data using latent rollout drift, latent geometry discriminability, and multi-task downstream performance.

Main Results:

• Clin-JEPA demonstrated stable latent $\ell_1$ rollout convergence ($-$15.7%) over 48-hour horizons, unlike diverging baselines.
• The learned latent space showed clinically discriminative geometry, with deteriorating patients displacing significantly further than stable patients.
• The single Clin-JEPA backbone outperformed strong tabular and sequence baselines in multi-task downstream risk prediction tasks.

Conclusions:

• Clin-JEPA provides a stable and effective framework for co-training JEPA models on EHR data.
• The framework enables a unified model for both EHR trajectory forecasting and diverse risk prediction tasks.
• Clin-JEPA significantly advances EHR representation learning, achieving superior performance in clinical forecasting and risk assessment.