Clin-JEPA: A Multi-Phase Co-Training Framework for Joint-Embedding Predictive Pretraining on EHR Patient Trajectories
Abstract
Joint-embedding predictive architectures (JEPA) learn representations by predicting in latent space, as in computer vision; retaining the action-conditioned predictor at inference turns them into latent world models, enabling planning in robotics (V-JEPA 2-AC). Bringing this design to EHR patient trajectories---a predictor that simulates a patient's trajectory in latent space---has not been explored. We use an LLM as the encoder, reading the hourly record as text, avoiding feature engineering and vocabulary harmonisation. But an LLM adapted by supervised fine-tuning does not organise its latent space around physiological dynamics, and freezing it to train the predictor, as in V-JEPA 2-AC, leaves the encoder unaware of the rollout signal: the predictor degrades under rollout. We instead co-train encoder and predictor under one latent-prediction objective, grounding the encoder in the dynamics its predictor must follow. Naïve co-training, however, is unstable: the untrained predictor drags the encoder toward collapse, and the predictor's rollout diverges as its target space moves. We present Clin-JEPA, a five-phase curriculum that stably co-trains an LLM encoder with a latent trajectory predictor on MIMIC-IV. Three evaluations support the design: (1) under 48-hour autoregressive rollout the co-trained predictor degrades least (predictor degradation ×1.06, against ×1.23--1.36 for two-stage designs and ×6.3--66 for curriculum ablations) while the co-trained encoder resolves the progression of patient state most sharply (largest state displacement); (2) the co-trained encoder separates deteriorating from stable patients in its latent space with Cohen's d=1.59, against ≤0.50 for two-stage encoders; (3) one set of embeddings serves 34 downstream tasks across three benchmarks, outperforming strong per-task tuned baselines and a pretrained EHR foundation model.
Related
- FlowMap: Geometry Dynamics Consistent Embedding of RNA Velocity for Interpretable Cellular Trajectories
- GSET: Gene Signals from Embedding Trajectories select compact contexts for mechanistic GRN inference
- RDEmbed: A Sentence-Embedding Framework for Rare Disease Phenotype Representation and Phenotypic Similarity Analysis
- Integrated Predictive and Remote-Sensing Framework for Early Warning and Regulatory Compliance in Environmentally Sensitive Urban Zones
- From mice to humans: A multi-omic predictive framework for translational immunology
BibTeX
@misc{yang2026clinjepa,
title = {{Clin-JEPA: A Multi-Phase Co-Training Framework for Joint-Embedding Predictive Pretraining on EHR Patient Trajectories}},
author = {Yixuan Yang and Mehak Arora and Ryan Zhang and Baraa Abed and Junseob Kim and Tilendra Choudhary and Md Hassanuzzaman and Kevin Zhu and Ayman Ali and Chengkun Yang and Alasdair Edward Gent and Victor Moas and Rishikesan Kamaleswaran},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2605.10840},
eprint = {2605.10840},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.10840},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Clin-JEPA: A Multi-Phase Co-Training Framework for Joint-Embedding Predictive Pretraining on EHR Patient Trajectories
AU - Yixuan Yang
AU - Mehak Arora
AU - Ryan Zhang
AU - Baraa Abed
AU - Junseob Kim
AU - Tilendra Choudhary
AU - Md Hassanuzzaman
AU - Kevin Zhu
AU - Ayman Ali
AU - Chengkun Yang
AU - Alasdair Edward Gent
AU - Victor Moas
AU - Rishikesan Kamaleswaran
PY - 2026
DA - 2026/09/16
PB - arXiv
DO - 10.48550/arxiv.2605.10840
UR - https://arxiv.org/abs/2605.10840
AB - Joint-embedding predictive architectures (JEPA) learn representations by predicting in latent space, as in computer vision; retaining the action-conditioned predictor at inference turns them into latent world models, enabling planning in robotics (V-JEPA 2-AC). Bringing this design to EHR patient trajectories---a predictor that simulates a patient's trajectory in latent space---has not been explored. We use an LLM as the encoder, reading the hourly record as text, avoiding feature engineering and vocabulary harmonisation. But an LLM adapted by supervised fine-tuning does not organise its latent space around physiological dynamics, and freezing it to train the predictor, as in V-JEPA 2-AC, leaves the encoder unaware of the rollout signal: the predictor degrades under rollout. We instead co-train encoder and predictor under one latent-prediction objective, grounding the encoder in the dynamics its predictor must follow. Naïve co-training, however, is unstable: the untrained predictor drags the encoder toward collapse, and the predictor's rollout diverges as its target space moves. We present Clin-JEPA, a five-phase curriculum that stably co-trains an LLM encoder with a latent trajectory predictor on MIMIC-IV. Three evaluations support the design: (1) under 48-hour autoregressive rollout the co-trained predictor degrades least (predictor degradation ×1.06, against ×1.23--1.36 for two-stage designs and ×6.3--66 for curriculum ablations) while the co-trained encoder resolves the progression of patient state most sharply (largest state displacement); (2) the co-trained encoder separates deteriorating from stable patients in its latent space with Cohen's d=1.59, against ≤0.50 for two-stage encoders; (3) one set of embeddings serves 34 downstream tasks across three benchmarks, outperforming strong per-task tuned baselines and a pretrained EHR foundation model.
ER -