Riepilogo Tecnico: CLIN-JEPA

Enunciato del Problema

La modellazione delle Cartelle Cliniche Elettroniche (EHR) in Terapia Intensiva (ICU) incontra una limitazione fondamentale: gli approcci esistenti trattano i dati dei pazienti o come sequenze di testo statiche (modelli autoregressivi token-per-token) che non riescono a modellare stati fisiologici continui, o come modelli fondazione per-task che richiedono ingegnerizzazione specifica delle caratteristiche e fine-tuning per ogni compito downstream. Nessuno di questi approcci fornisce una rappresentazione latente unificata capace sia di prevedere le traiettorie dei pazienti sia di supportare compiti diversificati di previsione del rischio senza adattamento per-task.

Sebbene le Architetture Predittive a Incorporamento Congiunto (JEPA) abbiano permesso con successo la pianificazione nello spazio latente nella robotica e nella visione, l'adattamento di questo paradigma alle EHR cliniche rimane una sfida aperta. I precedenti design JEPA soffrono di lacune architetturali specifiche quando applicati alla simulazione clinica:

• I-JEPA e V-JEPA: Scartano il predittore dopo il pre-addestramento, rendendo impossibile la simulazione della traiettoria a valle.
• V-JEPA 2-AC: Mantiene il predittore ma lo addestra su un encoder congelato. Di conseguenza, l'encoder è ignaro del "segnale di rollout" richiesto per l'inferenza, portando a una disconnessione tra lo spazio di rappresentazione e la dinamica di previsione.

Il co-training ingenuo di un encoder e di un predittore sotto un obiettivo JEPA condiviso è instabile, soggetto a collasso della rappresentazione (dove l'encoder degenera in output costanti) e deriva online/del target (dove il predittore diverge accumulando errori in uno spazio latente in movimento).

Metodologia: Framework CLIN-JEPA

Gli autori propongono CLIN-JEPA, un framework di co-training multi-fase progettato per addestrare stabilmente un singolo backbone che prevede simultaneamente le traiettorie dei pazienti e serve compiti di previsione del rischio downstream.

1. Rappresentazione dei Dati e Encoder

• Formato di Input: I ricoveri dei pazienti sono modellati come una sequenza di coppie stato-azione su finestre di 1 ora (fino a 72 ore), precedute da dati demografici statici. A differenza della vettorizzazione numerica, il modello utilizza testo strutturato in linguaggio naturale per rappresentare stati (segni vitali, esami di laboratorio, punteggi di gravità) e azioni (interventi, dosaggi). Questo preserva la leggibilità umana, gestisce i valori mancanti in modo naturale (per omissione) ed evita l'ingegnerizzazione delle caratteristiche specifica per modalità.
• Architettura dell'Encoder: Il sistema utilizza un modello linguistico Qwen3-8B come encoder di base. I pesi di base sono congelati e vengono addestrati adattatori LoRA leggeri per mappare gli input testuali (demografici $d$, stato $s_t$, azione $a_t$) in uno spazio latente a 4096 dimensioni.

2. Predittore di Traiettoria Latente

• Architettura: Un encoder Transformer da 92 milioni di parametri opera interamente nello spazio latente a 4096 dimensioni.
• Meccanismo: Consuma una sequenza intrecciata di dati demografici, stati e azioni. Utilizzando l'attenzione causale a blocchi (bidirezionale all'interno di un timestep, strettamente causale tra timestep), prevede il prossimo stato latente $\hat{z}_{t+1}$.
• Inferenza: Il predittore viene mantenuto al momento dell'inferenza. Esegue un rollout autoregressivo, alimentando le proprie previsioni continue di embedding come input per i passaggi successivi per simulare le traiettorie dei pazienti su un orizzonte temporale di $H$ ore.

3. Curriculum di Co-Training a Cinque Fasi

Per affrontare l'instabilità del co-training ingenuo, gli autori introducono un programma di pre-addestramento a cinque fasi:

1. Warmup: Il predittore viene addestrato da solo contro un encoder congelato e Fine-Tuned Supervisivamente (SFT). Questo impedisce al predittore "freddo" di trascinare immediatamente l'encoder verso il collasso della rappresentazione.
2. Co-training: Gli adattatori LoRA dell'encoder vengono sbloccati e addestrati congiuntamente al predittore sotto un obiettivo JEPA condiviso. Viene utilizzato un rollout forzato dal docente per garantire che l'encoder apprenda rappresentazioni dinamicamente fondate per il predittore.
3. Allineamento: L'encoder online viene re-congelato e l'encoder target a Media Mobile Esponenziale (EMA) si allinea gradualmente ai parametri dell'encoder online.
4. Hard Sync: Una copia istantanea dei parametri allinea gli encoder target e online, eliminando le discrepanze residue.
5. Finalizzazione: Il predittore viene addestrato sull'encoder stabilizzato utilizzando il rollout autoregressivo nativo (imitando le condizioni di inferenza) per far convergere completamente il predittore sulle rappresentazioni raffinate.

L'obiettivo di addestramento minimizza la distanza $\ell_1$ tra l'embedding del prossimo stato previsto e l'embedding del prossimo stato effettivo dell'encoder target, combinata con una perdita di rollout a breve orizzonte per incoraggiare la stabilità multi-step.

Contributi Chiave

1. Curriculum di Co-Training Stabile: Un programma a cinque fasi innovativo che addestra con successo un encoder basato su un grande modello linguistico e un predittore di traiettoria latente, prevenendo il collasso della rappresentazione e la deriva.
2. Rappresentazione Unificata Basata su Testo: Un framework che consuma testo clinico grezzo senza imputazione, normalizzazione o caratteristiche create a mano, servendo tutti i compiti downstream da un unico insieme di embedding latenti.
3. Primo JEPA Clinico con Predittore Mantenuto: A differenza di lavori precedenti che scartano il predittore, CLIN-JEPA lo mantiene per abilitare una simulazione stabile e autoregressiva dello stato del paziente.

Risultati Sperimentali

Il framework è stato valutato sul dataset MIMIC-IV ICU (84.497 ricoveri) su tre assi:

1. Stabilità e Convergenza dell'Addestramento

• Deriva del Rollout: Su un orizzonte autoregressivo di 48 ore, l'errore di previsione di CLIN-JEPA ha convergito (diminuito del 15,7% rispetto alla baseline a 1 ora). Al contrario, le baseline e le ablazioni (rimozione del warmup o dell'allineamento) hanno diverso catastroficamente, con aumenti di deriva che vanno dal +3% al +4951%.
• Collasso della Rappresentazione: Solo il curriculum completo di CLIN-JEPA ha mantenuto una deviazione standard sana negli output dell'encoder ($z_{std}$), evitando il collasso osservato nelle varianti prive della fase di warmup.

2. Diagnosi della Geometria Latente

• Discriminazione Clinica: L'encoder ha appreso una geometria in cui i pazienti in peggioramento (aumento del punteggio SOFA $\ge$ 3) si sono spostati 4,83 volte più lontano nello spazio latente rispetto ai pazienti stabili. Gli encoder baseline hanno mostrato rapporti di separazione significativamente inferiori (2,62× e 1,03×).
• Dinamiche Temporali: Il potere discriminatorio di CLIN-JEPA è cresciuto con l'orizzonte di osservazione, raggiungendo il picco alla 66ª ora, riflettendo un'internalizzazione della scala temporale lenta e integrativa della progressione in ICU.

3. Prestazioni Multi-Task Downstream

Valutato su due benchmark (ICareFM Early Event Prediction e 8 compiti di rischio binari) senza fine-tuning per-task:

• ICareFM EEP: Ha raggiunto un AUROC medio di 0,851, superando la baseline più forte (LightGBM, 0,827) e le ablazioni dell'encoder.
• Compiti di Rischio Binari: Ha raggiunto un AUROC medio di 0,883, superando la baseline LSTM (0,865).
• Valore del Predittore: L'inclusione del rollout autoregressivo (previsioni future) ha costantemente migliorato le prestazioni rispetto agli input basati solo sulla storia, in particolare per compiti che richiedono la composizione di informazioni eterogenee sulla traiettoria (es. danno renale, sepsi, scompenso).

Significato

Il documento afferma che CLIN-JEPA è il primo framework a colmare con successo il divario tra pre-addestramento JEPA e inferenza clinica mantenendo il predittore. Dimostra che un singolo backbone clinico può fornire rollout autoregressivo stabile, geometria latente clinicamente discriminativa e prestazioni multi-task competitive simultaneamente. Il lavoro stabilisce che il co-training dell'encoder e del predittore sotto un obiettivo condiviso, guidato da un curriculum specifico, è necessario per fondare le rappresentazioni per la simulazione a lungo orizzonte in sistemi dinamici ad alta dimensionalità come l'ICU.