论文

Clin-JEPA:针对 EHR 患者轨迹联合嵌入预测预训练的多阶段协同训练框架

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

模型训练预训练

摘要

联合嵌入预测架构(JEPA)通过在潜在空间中进行预测来学习表示,就像在计算机视觉中一样;在推理时保留动作条件预测器会将它们转变为潜在的世界模型,从而实现机器人规划(V-JEPA 2-AC)。尚未探索将此设计应用于 EHR 患者轨迹(模拟潜在空间中患者轨迹的预测器)。我们使用 LLM 作为编码器,将每小时记录读取为文本,避免特征工程和词汇协调。但是,由受监督的 微调 改编的 LLM 不会围绕生理动力学组织其潜在空间,并且将其冻结以训练预测器(如 V-JEPA 2-AC 中那样),使编码器不知道采样轨迹信号:预测器在采样轨迹下会降级。相反,我们在一个潜在预测目标下共同训练编码器和预测器,将编码器置于其预测器必须遵循的动态中。然而,朴素的协同训练是不稳定的:未经训练的预测器将编码器拖向崩溃,并且预测器的采样轨迹随着目标空间的移动而发散。我们提出了 Clin-JEPA,这是一个五阶段课程,可在 MIMIC-IV 上稳定地协同训练 LLM 编码器和潜在轨迹预测器。三个评估支持该设计:(1)在 48 小时自回归采样轨迹下,共同训练的预测器降级最少(预测器降级为 $\times$1.06,两阶段设计为 $\times$1.23--1.36,课程消融为 $\times$6.3--66),而共同训练的编码器则最清晰地解决患者状态的进展(最大状态位移); (2) 共同训练的编码器在其潜在空间中将病情恶化的患者与稳定的患者分开,Cohen 的 $d{=}1.59$,而两级编码器的 $\leq$0.50; (3) 一组嵌入服务于三个基准的 34 个下游任务,其性能优于强大的每任务调整基线和预训练的 EHR 基础模型。