Source-linked AI summary
Are Joint-Embedding Predictive Objectives Effective for Longitudinal EHRs?
Ashvin Gupta, Brendan Delaney, Alessandra Russo
TL;DR
MLM may emphasize local event reconstruction rather than higher-level patient state, while meaningful predictable targets are difficult to define for EHRs. This study tests three JEPA formulations on longitudinal EHRs and finds that they train without collapse but underperform MLM, with smaller targets failing to restore downstream performance.
Problem
For longitudinal EHRs, it remains unclear whether latent prediction can provide useful patient representations when events are irregular and future events may depend on healthcare decisions.
Source: §1 · pp. 1–2Method
The study evaluates span, future, and random JEPA target constructions for longitudinal EHR sequence representation.
Source: §1 · p. 2Results
JEPA models remained non-collapsed but underperformed MLM across downstream tasks, and reducing target size did not consistently improve downstream performance.
Source: §3 · pp. 3–4Takeaways & Limitations
The studied formulations show that stable latent prediction alone is insufficient to outperform token-level reconstruction, highlighting the challenge of defining prediction units that represent clinical patient state.
Source: §4 · p. 4Experiments are limited to MIMIC-IV and do not establish suitability for clinical deployment.
Source: §5.1 · p. 7
Abstract
from arXiv · showhide
Masked language modelling (MLM) is the dominant pretraining objective for transformer encoders applied to longitudinal electronic health records (EHRs). However, reconstructing individual events may overemphasise noisy or weakly informative observations and does not explicitly encourage learning higher-level patient states. Joint-Embedding Predictive Architectures (JEPAs) instead learn by predicting representations in latent space and have shown strong performance in a variety of domains. In this paper, we investigate to what extent JEPA can be adapted to EHR sequence representation. Our empirical investigation using MIMIC-IV shows that although all three of our JEPA formulations train stably without representational collapse, they consistently underperform standard MLM. We find that target size, not target heterogeneity, increases prediction error, yet reducing the target size does not recover downstream performance. This suggests that stable latent prediction alone is insufficient for competitive EHR representation learning and highlights the challenge of defining clinically meaningful latent targets for EHR.
1 Introduction
EHR foundation models depend on useful patient representations, but event-level MLM may prioritize local event identity over patient state. JEPA offers a latent-prediction alternative, though meaningful EHR targets are difficult to define.
1 Introduction
- Downstream performance depends on representation quality because EHR foundation models are adapted to clinical prediction tasks.These models aim to learn general-purpose patient representations through self-supervised learning.Source: §1 · p. 1
- MLM reconstructs masked events equally, potentially emphasizing routine or unpredictable details over higher-level patient state.This may lead representations to capture local event identity rather than patient state.Source: §1 · p. 1
- JEPA predicts latent representations from context rather than reconstructing observations, an approach intended to capture semantic features without unnecessary low-level reconstruction.The paper notes strong representation-learning performance for JEPA in vision, world modelling, and language [10] [11] [13].Source: §1 · p. 1
- Defining predictable EHR targets is difficult because events are irregularly sampled and semantically meaningful, while future events can depend on healthcare decisions.Temporal proximity does not necessarily imply semantic similarity.Source: §1 · p. 2
- The study tests whether I-JEPA-style latent prediction transfers to longitudinal EHRs across target constructions, stability, downstream performance, and training characteristics.The evaluated constructions include span masking, future masking, and random event masking.Source: §1 · p. 2
2 Methodology
The method represents each patient as a sequence of clinical-event embeddings and evaluates three JEPA target constructions against MLM. Each JEPA formulation uses an EMA target encoder and latent prediction from context.
2.1 Medical Event Representation
- Sequence order is encoded with Rotary Position Embeddings (RoPE), and all models share a transformer backbone with 6 layers, 8 attention heads, hidden dimension 768, and maximum context length N = 1024.Source: §2.1 · p. 2
2.2 Objectives
- Span JEPA masks contiguous target spans and predicts their EMA-encoder latents from remaining context using a shallow transformer predictor.Predictions are matched to target representations by mean squared distance.Source: §2.2 · pp. 2–3, Figure 1 · p. 2
- Future JEPA uses events observed before time T as context and the following 24 hours as one target region, while Random JEPA matches MLM's random corruption scheme.Random JEPA differs from MLM in predicting continuous target representations rather than discrete medical codes.Source: §2.2 · p. 3
- All JEPA experiments use a stop-gradient EMA target encoder; preliminary SIGReg experiments were unstable and collapsed, while MLM is the primary baseline.The EMA formulation remained stable and non-collapsed, allowing downstream performance to be assessed without conflating it with optimization failure.Source: §2.2 · p. 3
3 Experiments and Results
Across the evaluated downstream tasks, MLM outperforms all three JEPA formulations even though JEPA training remains non-collapsed. Smaller Span-JEPA targets reduce prediction difficulty but do not consistently improve downstream performance.
3 Experiments and Results
- MLM outperformed every JEPA formulation across downstream tasks, with Future JEPA closest to MLM.The evaluation compares pretraining objectives and masking strategies on MIMIC-IV clinical prediction tasks.Source: §3 · p. 3
- JEPA target-encoder effective rank and feature standard deviation remained stable as training and validation prediction losses decreased, indicating non-collapsed solutions.The models nevertheless learned representations less useful for downstream clinical prediction than MLM.Source: §3 · p. 3
- Smaller Span-JEPA targets reduced prediction difficulty but produced no consistent downstream improvement across the masking-ratio analysis.Target size was the only examined factor consistently positively associated with prediction error across all three formulations.Source: Figure 2 · p. 4, §3 · p. 4
4 Discussion
The studied JEPA formulations avoid representational collapse but underperform MLM, and reducing target size does not recover downstream performance. The discussion points to the prediction unit: event-level targets may not capture clinically meaningful patient states.
4 Discussion
- JEPA trains without collapse yet remains weaker than MLM on evaluated downstream tasks; reducing target size also fails to improve performance.The authors suggest the limitation may lie in how the prediction target is defined, rather than simply in optimization difficulty.Source: §4 · p. 4
- EHR events already carry specific semantic meaning, so predicting event-level latents may reconstruct event meaning rather than encourage a coherent patient-state representation.Unlike image patches, increasing the amount of masked EHR information does not necessarily increase semantic abstraction; MLM directly preserves clinically informative event identity, and recent JEPA-style EHR work retains token-level supervision alongside latent prediction [19].Source: §4 · p. 4
- Future JEPA performs more strongly, but future observations also reflect clinical decisions, and irregular sampling makes either temporal proximity or token contiguity an imperfect proxy for patient state.A fixed number of events can span different clinical intervals, while a fixed temporal window can contain varying amounts of information.Source: §4 · p. 4
- An EHR-specific JEPA may instead predict aggregated patient-state representations and model transitions between them, separating persistent clinical state from individual observations.The authors present this as a possible way to better match the abstraction motivating JEPA.Source: §4 · p. 4
5 Appendix
The appendix limits this study to retrospective prediction experiments on MIMIC-IV and does not establish clinical deployment suitability. It calls for further validation and oversight before clinical application.
5 Appendix
- The study evaluates retrospective EHR prediction and does not establish that JEPA is suitable for clinical deployment.The authors state outputs should not guide patient care without validation and human supervision.Source: §5.1 · p. 7
- Experiments are limited to MIMIC-IV, so performance may not generalize across institutions, patient populations, or care settings.Source: §5.1 · p. 7
- Before clinical application, the authors call for external validation, subgroup fairness assessment, calibration analysis, and prospective evaluation under clinical oversight.Source: §5.1 · p. 7
5.3 Training Dynamics
The appendix reports non-collapsed JEPA training dynamics and describes regression, target-selection, hyperparameter, dataset, and linear-probe materials. Its captions identify the scope of tables but do not provide their unseen values.
5.3 Training Dynamics
- Across 50 epochs, prediction losses decrease while target-embedding effective rank and feature variance remain non-degenerate for all three JEPA formulations.Source: Figure 3 · p. 7
5.4 Regression
- Validation per-target log(1 + MSE) is modeled with standardized OLS covariates, HC3 robust errors, and 95% confidence intervals.Covariates are target size, unique-code fraction, event-type diversity, target–context code overlap, and code rarity.Source: §5.4 · pp. 7–8
5.5 Context and Target selection
- Span targets use sampled contiguous blocks, while future targets use a cut meeting minimum context and target counts within 24 hours, with a last-10%-of-events fallback.Source: §5.5 · p. 8
5.6 Hyperparameters
- The hyperparameter-sweep table covers the three JEPA objectives while holding transformer and predictor settings fixed.Source: Table 3 · p. 9
5.7 MIMIC Dataset Information
- The MIMIC-IV dataset table is titled as covering population and task-specific distributions.Source: Table 4 · p. 9
5.8 Linear Probing
- The downstream linear-probe table reports AUROC and AUPRC with 95% percentile bootstrap confidence intervals from 1,000 stratified resamples.Source: Table 5 · p. 10