Resumen Técnico: CLIN-JEPA

Enunciado del Problema

El modelado de Registros Electrónicos de Salud (EHR) en la Unidad de Cuidados Intensivos (UCI) enfrenta una limitación fundamental: los enfoques existentes tratan los datos del paciente ya sea como secuencias de texto estáticas (modelos autorregresivos de tokens) que no logran modelar estados fisiológicos continuos, o como modelos base por tarea que requieren ingeniería de características específica y ajuste fino para cada tarea posterior. Ninguno de estos enfoques proporciona una representación latente unificada capaz tanto de pronosticar trayectorias de pacientes como de apoyar diversas tareas de predicción de riesgos sin adaptación por tarea.

Aunque las Arquitecturas Predictivas de Incrustación Conjunta (JEPA) han permitido con éxito la planificación en el espacio latente en robótica y visión, adaptar este paradigma a los EHR clínicos sigue siendo un desafío abierto. Los diseños previos de JEPA sufren de brechas arquitectónicas específicas al aplicarse a la simulación clínica:

• I-JEPA y V-JEPA: Descartan el predictor después del preentrenamiento, haciendo imposible la simulación de trayectorias aguas abajo.
• V-JEPA 2-AC: Mantiene el predictor pero lo entrena sobre un codificador congelado. En consecuencia, el codificador desconoce la señal de "despliegue" (rollout) requerida para la inferencia, lo que genera una desconexión entre el espacio de representación y la dinámica de predicción.

El entrenamiento conjunto ingenuo de un codificador y un predictor bajo un objetivo JEPA compartido es inestable, propenso al colapso de la representación (donde el codificador degenera en salidas constantes) y a la deriva en línea/de objetivo (donde el predictor diverge a medida que acumula errores en un espacio latente en movimiento).

Metodología: Marco CLIN-JEPA

Los autores proponen CLIN-JEPA, un marco de entrenamiento conjunto de múltiples fases diseñado para entrenar de manera estable un único backbone que pronostica simultáneamente trayectorias de pacientes y sirve a tareas de predicción de riesgos aguas abajo.

1. Representación de Datos y Codificador

• Formato de Entrada: Las estancias de los pacientes se modelan como una secuencia de pares estado-acción en intervalos de 1 hora (hasta 72 horas), precedidos por datos demográficos estáticos. A diferencia de la vectorización numérica, el modelo utiliza texto estructurado en lenguaje natural para representar estados (signos vitales, análisis de laboratorio, puntuaciones de gravedad) y acciones (intervenciones, dosis). Esto preserva la legibilidad humana, maneja los valores faltantes de forma natural (mediante omisión) y evita la ingeniería de características específica de la modalidad.
• Arquitectura del Codificador: El sistema utiliza un modelo de lenguaje Qwen3-8B como codificador base. Los pesos base se congelan y se entrenan adaptadores LoRA ligeros para mapear las entradas de texto (demografía $d$, estado $s_t$, acción $a_t$) a un espacio latente de 4096 dimensiones.

2. Predictor de Trayectoria Latente

• Arquitectura: Un codificador Transformer de 92 millones de parámetros opera enteramente dentro del espacio latente de 4096 dimensiones.
• Mecanismo: Consume una secuencia intercalada de datos demográficos, estados y acciones. Utilizando atención causal por bloques (bidireccional dentro de un paso de tiempo, estrictamente causal entre pasos de tiempo), predice el siguiente estado latente $\hat{z}_{t+1}$.
• Inferencia: El predictor se mantiene durante la inferencia. Realiza un despliegue autorregresivo, alimentando sus propias predicciones de incrustación continua de vuelta como entradas para pasos posteriores para simular trayectorias de pacientes durante un horizonte de $H$ horas.

3. Currículo de Entrenamiento Conjunto de Cinco Fases

Para abordar la inestabilidad del entrenamiento conjunto ingenuo, los autores introducen un programa de preentrenamiento de cinco fases:

1. Calentamiento (Warmup): El predictor se entrena solo contra un codificador congelado y ajustado finamente de forma supervisada (SFT). Esto evita que el predictor "frío" arrastre inmediatamente al codificador hacia un colapso de la representación.
2. Entrenamiento Conjunto: Los adaptadores LoRA del codificador se desbloquean y se entrenan conjuntamente con el predictor bajo un objetivo JEPA compartido. Se utiliza un despliegue forzado por el maestro (teacher-forced rollout) para asegurar que el codificador aprenda representaciones dinámicamente fundamentadas para el predictor.
3. Alineación: El codificador en línea se vuelve a congelar, y el codificador objetivo de Promedio Móvil Exponencial (EMA) alcanza suavemente los parámetros del codificador en línea.
4. Sincronización Rápida (Hard Sync): Una copia instantánea de parámetros alinea los codificadores objetivo y en línea, eliminando discrepancias residuales.
5. Finalización: El predictor se entrena sobre el codificador estabilizado utilizando despliegue autorregresivo nativo (imitando condiciones de inferencia) para converger completamente el predictor sobre las representaciones refinadas.

El objetivo de entrenamiento minimiza la distancia $\ell_1$ entre la incrustación del siguiente estado predicho y la incrustación del siguiente estado real del codificador objetivo, combinada con una pérdida de despliegue de corto horizonte para fomentar la estabilidad multi-paso.

Contribuciones Clave

1. Currículo de Entrenamiento Conjunto Estable: Un programa de cinco fases novedoso que entrena conjuntamente con éxito un codificador basado en un modelo de lenguaje grande y un predictor de trayectoria latente, previniendo el colapso de la representación y la deriva.
2. Representación Unificada Basada en Texto: Un marco que consume texto clínico crudo sin imputación, normalización o características diseñadas a mano, sirviendo a todas las tareas posteriores desde un único conjunto de incrustaciones latentes.
3. Primera JEPA Clínica con Predictor Retenido: A diferencia de trabajos anteriores que descartan el predictor, CLIN-JEPA lo mantiene para habilitar una simulación estable y autorregresiva del estado del paciente.

Resultados Experimentales

El marco se evaluó en el conjunto de datos MIMIC-IV ICU (84,497 estancias) a través de tres ejes:

1. Estabilidad y Convergencia del Entrenamiento

• Deriva del Despliegue: Sobre un horizonte autorregresivo de 48 horas, el error de predicción de CLIN-JEPA convergió (disminuyó un 15.7% en relación con la línea base de 1 hora). En contraste, las líneas base y las ablaciones (eliminación del calentamiento o la alineación) divergieron catastróficamente, con aumentos de deriva que oscilaron entre +3% y +4951%.
• Colapso de la Representación: Solo el currículo completo de CLIN-JEPA mantuvo una desviación estándar saludable en las salidas del codificador ($z_{std}$), evitando el colapso observado en variantes que carecían de la fase de calentamiento.

2. Diagnóstico de la Geometría Latente

• Discriminación Clínica: El codificador aprendió una geometría donde los pacientes que se deterioran (aumento de la puntuación SOFA $\ge$ 3) se desplazaron 4.83 veces más en el espacio latente que los pacientes estables. Los codificadores de línea base mostraron ratios de separación significativamente menores (2.62 veces y 1.03 veces).
• Dinámicas Temporales: El poder discriminatorio de CLIN-JEPA creció con el horizonte de observación, alcanzando un pico a la hora 66, reflejando una internalización de la escala de tiempo lenta e integradora de la progresión en la UCI.

3. Rendimiento Multi-tarea Aguas Abajo

Evaluado en dos puntos de referencia (Predicción de Eventos Tempranos de ICareFM y 8 tareas de riesgo binarias) sin ajuste fino por tarea:

• EEP de ICareFM: Logró un AUROC medio de 0.851, superando a la línea base más fuerte (LightGBM, 0.827) y a las ablaciones del codificador.
• Tareas de Riesgo Binarias: Logró un AUROC medio de 0.883, superando a la línea base LSTM (0.865).
• Valor del Predictor: La inclusión del despliegue autorregresivo (predicciones futuras) mejoró consistentemente el rendimiento sobre las entradas de solo historial, particularmente para tareas que requieren la composición de información de trayectoria heterogénea (por ejemplo, lesión renal, sepsis, descompensación).

Significado

El artículo afirma que CLIN-JEPA es el primer marco que cierra con éxito la brecha entre el preentrenamiento JEPA y la inferencia clínica al retener el predictor. Demuestra que un único backbone clínico puede ofrecer despliegue autorregresivo estable, geometría latente clínicamente discriminatoria y rendimiento multi-tarea competitivo simultáneamente. El trabajo establece que es necesario entrenar conjuntamente el codificador y el predictor bajo un objetivo compartido, guiado por un currículo específico, para fundamentar representaciones para la simulación de largo horizonte en sistemas dinámicos de alta dimensión como la UCI.