🧬 biology

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

Clin-JEPA introduces a stable, five-phase co-training framework that successfully adapts Joint-Embedding Predictive Architecture (JEPA) to EHR data by jointly training an encoder and predictor, thereby enabling a single backbone to achieve superior latent-space trajectory forecasting and multi-task risk prediction without per-task fine-tuning.

Original authors: Yixuan Yang, Mehak Arora, Ryan Zhang, Baraa Abed, Junseob Kim, Tilendra Choudhary, Md Hassanuzzaman, Kevin Zhu, Ayman Ali, Chengkun Yang, Alasdair Edward Gent, Victor Moas, Rishikesan Kamaleswaran

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yixuan Yang, Mehak Arora, Ryan Zhang, Baraa Abed, Junseob Kim, Tilendra Choudhary, Md Hassanuzzaman, Kevin Zhu, Ayman Ali, Chengkun Yang, Alasdair Edward Gent, Victor Moas, Rishikesan Kamaleswaran

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

The Big Picture: Teaching a Doctor to "Imagine" the Future

Imagine you are trying to teach a student doctor how to predict how a patient will get better or worse in the Intensive Care Unit (ICU).

Most current AI models are like photographers. They take a picture of the patient right now (their vitals, blood work, and treatments) and try to guess a single outcome, like "Will they survive?" or "Will they get sepsis?" They are great at answering specific questions, but they don't really understand how the patient's body changes over time.

Other models are like novel writers. They read the patient's history and try to write the next sentence of the story. But they often just guess the next word based on patterns in text, without truly understanding the underlying physics of the human body.

CLIN-JEPA is a new approach that tries to be a simulation engine. It doesn't just guess a single outcome; it learns to "roll out" a movie of the patient's future. It asks: "If we give this patient this medicine right now, what will their vital signs look like in 1 hour? 10 hours? 48 hours?"

The Problem: The "Frozen Brain" vs. The "Moving Target"

The researchers tried to build this simulation using a technique called JEPA (Joint-Embedding Predictive Architecture), which is famous in robotics and vision. However, they hit a wall when applying it to medical data.

Think of the AI as having two parts:

The Encoder (The Brain): It reads the messy, complex medical text (vitals, drugs, notes) and turns it into a clean, mathematical "thought" (a latent embedding).
The Predictor (The Crystal Ball): It looks at those "thoughts" and tries to guess what the next "thought" will be.

The Failure Modes:

The "Frozen Brain" Problem: In previous methods, they trained the Brain first, then froze it, and then trained the Crystal Ball. The problem? The Brain didn't know the Crystal Ball existed. It learned to think in a way that was good for reading, but bad for predicting the future. The Crystal Ball struggled because the Brain's "thoughts" weren't organized for prediction.
The "Moving Target" Problem: If you try to train both at the same time (co-training), it gets chaotic. The Brain changes its mind constantly, and the Crystal Ball gets confused, chasing a target that keeps moving. Eventually, the whole system collapses, and the AI starts outputting nonsense or "hallucinating" a flat, boring future where nothing ever changes.

The Solution: The Five-Phase "Training Camp"

The authors created CLIN-JEPA, which is a specific, five-step training schedule designed to fix these problems. Think of it like training a new athlete and their coach together, rather than just throwing them into a game.

Phase 1: Warmup (The Solo Drill).
First, they train the Crystal Ball (Predictor) alone while the Brain (Encoder) stays frozen. This lets the Crystal Ball learn the basics of the game without the Brain changing the rules on it.
Phase 2: Joint Refinement (The Partnership).
Now, they unlock the Brain and train both together. This is the magic step. The Brain learns to organize its thoughts specifically so the Crystal Ball can predict the future. They are "grounded" in each other.
Phase 3: Alignment (The Soft Catch-up).
The Brain is frozen again, but the "Target" (a reference version of the Brain) slowly catches up to it. This prevents a sudden shock that would confuse the Crystal Ball.
Phase 4: Hard Sync (The Snap).
They instantly copy the Brain's current state to the Target. This ensures they are perfectly identical, removing any lingering confusion.
Phase 5: Finalize (The Simulation).
Finally, they train the Crystal Ball using the actual future it predicts (rolling out its own guesses) rather than using the "correct" answers from the past. This teaches it to be stable over long periods, just like a real simulation.

What Did They Find? (The Results)

The team tested this on MIMIC-IV, a massive database of real ICU patient records. Here is what happened:

It Doesn't Crash: When other methods tried to predict 48 hours into the future, their predictions drifted wildly off course (like a GPS losing signal). CLIN-JEPA's predictions actually got more accurate the further out they looked, converging on a stable reality.
It Understands "Sick" vs. "Stable": The researchers looked at the "shape" of the data the AI learned.
- Analogy: Imagine a map where "sick" patients and "healthy" patients are dots.
- Old AI: The dots for sick and healthy patients were mixed up, or they all drifted around randomly.
- CLIN-JEPA: The "sick" patients moved far away from the "healthy" patients in the map. The AI learned a clear, distinct geometry where a deteriorating patient looks very different from a stable one.
One Model Does It All: Usually, you need a different AI model for every specific task (one for predicting death, one for predicting sepsis, one for predicting length of stay). CLIN-JEPA uses one single backbone (one set of learned "thoughts") to beat all the specialized models on 15 different tasks without needing to be retrained for each one.

The Bottom Line

CLIN-JEPA is a new framework that successfully teaches an AI to simulate patient trajectories in the ICU. By using a clever five-step training camp, it solved the problem of making the AI's "brain" and "prediction engine" work together without crashing.

The result is a single, robust model that can look at a patient's history and reliably "imagine" their future state, separating those who are getting worse from those who are stable, and outperforming existing methods at predicting various health risks.

Technical Summary: CLIN-JEPA

Problem Statement

Electronic Health Record (EHR) modeling in the Intensive Care Unit (ICU) faces a fundamental limitation: existing approaches either treat patient data as static text sequences (token-autoregressive models) that fail to model continuous physiological states, or as per-task foundation models that require specific feature engineering and fine-tuning for each downstream task. Neither approach provides a unified latent representation capable of both forecasting patient trajectories and supporting diverse risk-prediction tasks without per-task adaptation.

While Joint-Embedding Predictive Architectures (JEPA) have successfully enabled latent-space planning in robotics and vision, adapting this paradigm to clinical EHR remains an open challenge. Prior JEPA designs suffer from specific architectural gaps when applied to clinical simulation:

I-JEPA and V-JEPA: Discard the predictor after pretraining, rendering trajectory simulation impossible downstream.
V-JEPA 2-AC: Retains the predictor but trains it on a frozen encoder. Consequently, the encoder is unaware of the "rollout signal" required for inference, leading to a disconnect between the representation space and the prediction dynamics.

Naïve co-training of an encoder and predictor under a shared JEPA objective is unstable, prone to representation collapse (where the encoder degenerates to constant outputs) and online/target drift (where the predictor diverges as it accumulates errors in a moving latent space).

Methodology: CLIN-JEPA Framework

The authors propose CLIN-JEPA, a multi-phase co-training framework designed to stably train a single backbone that simultaneously forecasts patient trajectories and serves downstream risk-prediction tasks.

1. Data Representation and Encoder

Input Format: Patient stays are modeled as a sequence of state-action pairs over 1-hour bins (up to 72 hours), preceded by static demographics. Unlike numerical vectorization, the model uses structured natural language text to represent states (vitals, labs, severity scores) and actions (interventions, doses). This preserves human readability, handles missing values naturally (by omission), and avoids modality-specific feature engineering.
Encoder Architecture: The system utilizes a Qwen3-8B language model as the base encoder. The base weights are frozen, and lightweight LoRA adapters are trained to map text inputs (demographics $d$ , state $s_t$ , action $a_t$ ) into a 4096-dimensional latent space.

2. Latent Trajectory Predictor

Architecture: A 92M-parameter Transformer encoder operates entirely within the 4096-dimensional latent space.
Mechanism: It consumes an interleaved sequence of demographics, states, and actions. Using block-causal attention (bidirectional within a timestep, strictly causal across timesteps), it predicts the next latent state $\hat{z}_{t+1}$ .
Inference: The predictor is retained at inference time. It performs autoregressive rollout, feeding its own continuous embedding predictions back as inputs for subsequent steps to simulate patient trajectories over an $H$ -hour horizon.

3. Five-Phase Co-Training Curriculum

To address the instability of naïve co-training, the authors introduce a five-phase pretraining schedule:

Warmup: The predictor is trained alone against a frozen, Supervised Fine-Tuned (SFT) encoder. This prevents the "cold" predictor from immediately dragging the encoder toward representation collapse.
Co-training: The encoder's LoRA adapters are unlocked and trained jointly with the predictor under a shared JEPA objective. Teacher-forced rollout is used to ensure the encoder learns representations dynamically grounded for the predictor.
Alignment: The online encoder is re-frozen, and the Exponential Moving Average (EMA) target encoder smoothly catches up to the online encoder's parameters.
Hard Sync: An instantaneous parameter copy aligns the target and online encoders, eliminating residual mismatches.
Finalize: The predictor is trained on the stabilized encoder using native autoregressive rollout (mimicking inference conditions) to fully converge the predictor on the refined representations.

The training objective minimizes the $\ell_1$ distance between the predicted next-state embedding and the target encoder's embedding of the actual next state, combined with a short-horizon rollout loss to encourage multi-step stability.

Key Contributions

Stable Co-Training Curriculum: A novel five-phase schedule that successfully co-trains a large language model-based encoder and a latent trajectory predictor, preventing representation collapse and drift.
Unified Text-Based Representation: A framework that consumes raw clinical text without imputation, normalization, or hand-engineered features, serving all downstream tasks from a single set of latent embeddings.
First Clinical JEPA with Retained Predictor: Unlike prior works that discard the predictor, CLIN-JEPA retains it to enable stable, autoregressive patient-state simulation.

Experimental Results

The framework was evaluated on the MIMIC-IV ICU dataset (84,497 stays) across three axes:

1. Training Stability and Convergence

Rollout Drift: Over a 48-hour autoregressive horizon, CLIN-JEPA's prediction error converged (decreased by 15.7% relative to the 1-hour baseline). In contrast, baselines and ablations (removing warmup or alignment) diverged catastrophically, with drift increases ranging from +3% to +4951%.
Representation Collapse: Only the full CLIN-JEPA curriculum maintained a healthy standard deviation in encoder outputs ( $z_{std}$ ), avoiding the collapse observed in variants lacking the warmup phase.

2. Latent Geometry Diagnosis

Clinical Discrimination: The encoder learned a geometry where deteriorating patients (SOFA score increase $\ge$ 3) displaced 4.83× further in latent space than stable patients. Baseline encoders showed significantly lower separation ratios (2.62× and 1.03×).
Temporal Dynamics: The discriminative power of CLIN-JEPA grew with the observation horizon, peaking at hour 66, reflecting an internalization of the slow, integrative timescale of ICU progression.

3. Downstream Multi-Task Performance

Evaluated on two benchmarks (ICareFM Early Event Prediction and 8 binary risk tasks) without per-task fine-tuning:

ICareFM EEP: Achieved a mean AUROC of 0.851, outperforming the strongest baseline (LightGBM, 0.827) and encoder ablations.
Binary Risk Tasks: Achieved a mean AUROC of 0.883, outperforming the LSTM baseline (0.865).
Predictor Value: The inclusion of the autoregressive rollout (future predictions) consistently improved performance over history-only inputs, particularly for tasks requiring the composition of heterogeneous trajectory information (e.g., kidney injury, sepsis, decompensation).

Significance

The paper claims that CLIN-JEPA is the first framework to successfully close the gap between JEPA pretraining and clinical inference by retaining the predictor. It demonstrates that a single clinical backbone can deliver stable autoregressive rollout, clinically discriminative latent geometry, and competitive multi-task performance simultaneously. The work establishes that co-training the encoder and predictor under a shared objective, guided by a specific curriculum, is necessary to ground representations for long-horizon simulation in high-dimensional dynamical systems like the ICU.