让医疗记录预测更准:用多阶段联合训练提升患者轨迹建模能力
Clin-JEPA: A Multi-Phase Co-Training Framework for Joint-Embedding Predictive Pretraining on EHR Patient Trajectories

- 分五阶段联合训练编码器与预测器,稳定解决自回归生成发散问题
- 48小时预测误差下降15.7%,患者恶化状态在隐空间区分度提升4.83倍
- 单一模型可通用多个临床任务,优于传统表格/序列方法
我们提出Clin-JEPA,一种用于电子健康记录(EHR)患者轨迹的多阶段联合嵌入预测(JEPA)预训练框架。尽管JEPA架构已在机器人和视觉领域实现高质量表征学习,但将其拓展至EHR数据以构建一个能同时预测患者轨迹并支持多种下游风险预测任务的统一骨干网络,仍面临挑战。现有框架或在预训练后丢弃预测器(I-JEPA, V-JEPA),或在冻结编码器上训练预测器(V-JEPA 2-AC),导致编码器无法感知推理时预测器所需的滚动信号。为解决此问题,我们设计了五阶段预训练流程——预测器预热、联合精炼、EMA目标对齐、硬同步和预测器终调——稳定地联合训练基于Qwen3-8B的编码器与92M参数的潜在轨迹预测器。在MIMIC-IV ICU数据集上的三项独立评估表明:(1)潜空间ℓ₁滚动漂移在48小时预测中唯一收敛(-15.7%),而基线和消融组发散(+3%至+4951%);(2)编码器学习到具有临床判别性的潜在几何结构(恶化患者群体在潜空间移动距离是稳定患者的4.83倍,而基线≤2.62倍);(3)单一骨干模型在多任务下游评估中超越强基线。Clin-JEPA在ICareFM EEP任务上达到0.851平均AUROC,8个二分类风险任务上达0.883(分别优于基线平均0.038和0.041)。
原文摘要 · Abstract (English)
We present Clin-JEPA, a multi-phase co-training framework for joint-embedding predictive (JEPA) pretraining on EHR patient trajectories. JEPA architectures have enabled latent-space planning in robotics and high-quality representation learning in vision, but extending the paradigm to EHR data -- to obtain a single backbone that simultaneously forecasts patient trajectories and serves diverse downstream risk-prediction tasks without per-task fine-tuning -- remains an open challenge. Existing JEPA frameworks either discard the predictor after pretraining (I-JEPA, V-JEPA) or train it on a frozen pretrained encoder (V-JEPA 2-AC), leaving the encoder unaware of the rollout signal that the retained predictor must use at inference; co-training the encoder and predictor under a shared JEPA prediction objective would supply this grounding, but naïve co-training is unstable, with representation collapse and online/target drift causing autoregressive rollout to diverge. Clin-JEPA's five-phase pretraining curriculum -- predictor warmup, joint refinement, EMA target alignment, hard sync, and predictor finalization -- addresses each failure mode by phase, stably co-training a Qwen3-8B-based encoder and a 92M-parameter latent trajectory predictor. On MIMIC-IV ICU data, three independent evaluations support the framework: (1) latent $\ell_1$ rollout drift uniquely converges ($-$15.7%) over 48-hour horizons while baselines and ablations diverge (+3% to +4951%); (2) the encoder learns a clinically discriminative latent geometry (deteriorating-patient cohorts displace 4.83$\times$ further than stable patients in latent space, vs $\leq$2.62$\times$ for baseline encoders); (3) a single backbone outperforms strong tabular and sequence baselines on multi-task downstream evaluation. Clin-JEPA achieves mean AUROC 0.851 on ICareFM EEP and 0.883 on 8 binary risk tasks (+0.038 and +0.041 vs baseline average).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。