Clin-JEPA

An intensive-care chart never lines up on a grid: a blood gas at 03:12, three blood pressures in the next ten minutes, a norepinephrine change at 03:40, then nothing from the lab for six hours. Most electronic-health-record models flatten that into a fixed numeric vector — impute the missing lactate, normalize the units, hand-build a rolling mean — discarding the within-hour ordering and the clinically meaningful fact that a measurement was never taken. Clin-JEPA instead serializes each hour of an ICU stay as plain clinical text (t=12h | Heart rate: 88 bpm. MAP: 62 mmHg. Lactate: 2.1 mmol/L. ...) and hands it to a language-model encoder, so absent readings are simply absent from the string and no imputation or featurization step is needed.

What happens to those embeddings is the second idea. Joint-embedding predictive architectures predict future or masked content in latent space rather than reconstructing raw input, but the vision and robotics members of the family treat the predictor as scaffolding: I-JEPA and V-JEPA discard it after pretraining, and V-JEPA 2-AC trains it on a frozen encoder. Clin-JEPA retains the predictor to simulate patient trajectories at inference, so the encoder must be grounded in the rollout signal it consumes. Co-training both under one objective supplies that grounding but destabilizes training: the encoder degrades toward constant outputs, and the predictor forecasts inside a latent space still drifting underneath it.

Rishikesan Kamaleswaran's lab at Duke University answers with a five-phase curriculum that unlocks and refreezes the encoder in a fixed order, neutralizing each failure mode in its own phase. Where CEHR-GPT treats a patient record as a token sequence to generate, Clin-JEPA treats it as a continuous state to roll forward.

#Key Features

• Text-native patient state: State, action and demographics fragments become 4096-dimensional embeddings, unifying vitals, labs, drugs and procedures without a fixed feature schema.
• Retained latent predictor: The predictor survives pretraining and rolls out autoregressively at inference, feeding its own continuous prediction back as the next state rather than sampling a token.
• Five-phase curriculum: Predictor warmup, joint refinement, EMA target alignment, hard sync and finalization; ablating warmup or alignment reproduces one failure mode each.
• One backbone, many tasks: A single set of latent embeddings serves 15 downstream clinical tasks through shallow probes, with no per-task fine-tuning of the encoder.
• Action conditioning: Interventions enter the sequence as their own tokens, so a rollout follows a proposed treatment plan, not observations alone.

#Technical Details

The encoder is a frozen Qwen3-8B language model adapted with LoRA (rank 16, alpha 32, on the attention projections — about 30.7M trainable parameters against 8.2B frozen); the last-token hidden state is the embedding. The predictor is a separate six-layer pre-norm Transformer of roughly 92M parameters (1024 hidden, 8 heads) operating directly in the encoder's 4096-dimensional space, with block-causal attention over the interleaved demographics–state–action sequence. Training minimizes an ℓ1 teacher-forcing loss against an EMA target encoder (τ = 0.996) plus a two-step rollout term. Pretraining used MIMIC-IV v3.1 ICU data — 84,497 stays from 64,874 patients, split 70/15/15 by patient and windowed to 72 hours at one-hour resolution for roughly 197,000 windows — on eight NVIDIA H200 GPUs in about 54 wall-clock hours (≈430 GPU-hours). Over 48 hours Clin-JEPA's mean rollout drift falls 15.7%, while a V-JEPA 2-AC-style baseline stays flat (+3.4%) and the no-alignment ablation diverges by 4951%; deteriorating cohorts displace 4.83× further in latent space than stable ones, against 2.62× for V-JEPA 2-AC. Downstream it reaches mean AUROC 0.851 on the seven-task ICareFM early-event-prediction suite (LightGBM: 0.827) and 0.883 across eight stay-level binary outcomes (LSTM: 0.865), with the largest gains where temporal composition matters — decompensation +0.189 and acute kidney injury +0.111.

#Applications

The intended use is one ICU representation serving both simulation and risk scoring: forecasting where a patient's physiology is heading over the next two days, comparing rollouts under different proposed interventions, and feeding frozen embeddings to lightweight probes for mortality, sepsis, prolonged-stay and organ-failure endpoints. Because the inputs are ordinary clinical text, a newly charted variable needs no schema change. The practical audience is critical-care informatics groups working on MIMIC-IV, which requires PhysioNet credentialed access.

#Impact

Clin-JEPA shows the JEPA recipe transfers to clinical time series when the predictor is part of the deployed model rather than a training aid, and its curriculum answers a failure mode any encoder–predictor co-training run will hit. The caveats are real: it is an unreviewed preprint trained and evaluated entirely on one US academic center's ICU records, so cross-institution transfer is untested, and its margins over a well-tuned LightGBM are modest on threshold-detection tasks that hand-engineered vitals features already handle. The code is MIT-licensed, but no pretrained checkpoints have been released, so reproducing it means re-running roughly 430 GPU-hours on your own MIMIC-IV extract. The same Duke group also released CardioFM, an ECG and PPG waveform encoder sharing several authors but separate work on a different modality.