Résumé Technique : CLIN-JEPA

Énoncé du Problème

La modélisation des Dossiers Médicaux Électroniques (DME) en Unité de Soins Intensifs (USI) fait face à une limitation fondamentale : les approches existantes traitent soit les données patient comme des séquences de texte statiques (modèles autorégressifs à jetons) qui échouent à modéliser les états physiologiques continus, soit comme des modèles de base par tâche nécessitant un ingénierie de caractéristiques spécifique et un ajustement fin pour chaque tâche en aval. Aucune de ces approches ne fournit une représentation latente unifiée capable à la fois de prévoir les trajectoires des patients et de soutenir diverses tâches de prédiction de risque sans adaptation par tâche.

Bien que les Architectures Prédictives à Encodage Joint (JEPA) aient permis avec succès la planification dans l'espace latent en robotique et en vision, l'adaptation de ce paradigme aux DME cliniques reste un défi ouvert. Les conceptions JEPA antérieures souffrent de lacunes architecturales spécifiques lorsqu'elles sont appliquées à la simulation clinique :

• I-JEPA et V-JEPA : Rejettent le prédicteur après le pré-entraînement, rendant impossible la simulation de trajectoire en aval.
• V-JEPA 2-AC : Conserve le prédicteur mais l'entraîne sur un encodeur figé. Par conséquent, l'encodeur ignore le "signal de déploiement" nécessaire à l'inférence, ce qui entraîne une déconnexion entre l'espace de représentation et la dynamique de prédiction.

Un co-entraînement naïf d'un encodeur et d'un prédicteur sous un objectif JEPA partagé est instable, sujet à un effondrement de la représentation (où l'encodeur dégénère en sorties constantes) et à une dérive en ligne/cible (où le prédicteur diverge alors qu'il accumule des erreurs dans un espace latent en mouvement).

Méthodologie : Cadre CLIN-JEPA

Les auteurs proposent CLIN-JEPA, un cadre de co-entraînement multi-phasique conçu pour entraîner de manière stable un seul socle capable à la fois de prévoir les trajectoires des patients et de servir les tâches de prédiction de risque en aval.

1. Représentation des Données et Encodeur

• Format d'Entrée : Les séjours hospitaliers sont modélisés comme une séquence de paires état-action sur des intervalles d'une heure (jusqu'à 72 heures), précédés par des données démographiques statiques. Contrairement à la vectorisation numérique, le modèle utilise du texte structuré en langage naturel pour représenter les états (signes vitaux, analyses de laboratoire, scores de sévérité) et les actions (interventions, dosages). Cela préserve la lisibilité humaine, gère naturellement les valeurs manquantes (par omission) et évite l'ingénierie de caractéristiques spécifique à une modalité.
• Architecture de l'Encodeur : Le système utilise un modèle de langage Qwen3-8B comme encodeur de base. Les poids de base sont figés, et des adaptateurs LoRA légers sont entraînés pour mapper les entrées textuelles (démographie $d$, état $s_t$, action $a_t$) dans un espace latent de 4096 dimensions.

2. Prédicteur de Trajectoire Latente

• Architecture : Un encodeur Transformer de 92 millions de paramètres opère entièrement dans l'espace latent de 4096 dimensions.
• Mécanisme : Il consomme une séquence entrelacée de données démographiques, d'états et d'actions. En utilisant une attention causale par blocs (bidirectionnelle au sein d'un pas de temps, strictement causale entre les pas de temps), il prédit le prochain état latent $\hat{z}_{t+1}$.
• Inférence : Le prédicteur est conservé au moment de l'inférence. Il effectue un déploiement autorégressif, réinjectant ses propres prédictions d'embeddings continus comme entrées pour les étapes suivantes afin de simuler les trajectoires des patients sur un horizon de $H$ heures.

3. Curriculum de Co-entraînement en Cinq Phases

Pour remédier à l'instabilité du co-entraînement naïf, les auteurs introduisent un calendrier de pré-entraînement en cinq phases :

1. Échauffement : Le prédicteur est entraîné seul contre un encodeur figé, finement ajusté de manière supervisée (SFT). Cela empêche le prédicteur "froid" de tirer immédiatement l'encodeur vers un effondrement de la représentation.
2. Co-entraînement : Les adaptateurs LoRA de l'encodeur sont déverrouillés et entraînés conjointement avec le prédicteur sous un objectif JEPA partagé. Un déploiement forcé par un enseignant est utilisé pour s'assurer que l'encodeur apprend des représentations dynamiquement ancrées pour le prédicteur.
3. Alignement : L'encodeur en ligne est refigé, et l'encodeur cible à Moyenne Mobile Exponentielle (EMA) rattrape progressivement les paramètres de l'encodeur en ligne.
4. Synchronisation Stricte : Une copie instantanée des paramètres aligne les encodeurs cible et en ligne, éliminant les écarts résiduels.
5. Finalisation : Le prédicteur est entraîné sur l'encodeur stabilisé en utilisant un déploiement autorégressif natif (imitant les conditions d'inférence) pour converger pleinement le prédicteur sur les représentations affinées.

L'objectif d'entraînement minimise la distance $\ell_1$ entre l'embedding de l'état suivant prédit et l'embedding de l'état suivant réel par l'encodeur cible, combiné à une perte de déploiement à court horizon pour encourager la stabilité multi-étapes.

Contributions Clés

1. Curriculum de Co-entraînement Stable : Un calendrier en cinq phases novateur qui permet de co-entraîner avec succès un encodeur basé sur un grand modèle de langage et un prédicteur de trajectoire latente, empêchant l'effondrement de la représentation et la dérive.
2. Représentation Unifiée Basée sur le Texte : Un cadre qui consomme du texte clinique brut sans imputation, normalisation ou caractéristiques conçues à la main, servant toutes les tâches en aval à partir d'un seul ensemble d'embeddings latents.
3. Premier JEPA Clinique avec Prédicteur Conservé : Contrairement aux travaux antérieurs qui rejettent le prédicteur, CLIN-JEPA le conserve pour permettre une simulation stable et autorégressive de l'état du patient.

Résultats Expérimentaux

Le cadre a été évalué sur le jeu de données MIMIC-IV ICU (84 497 séjours) selon trois axes :

1. Stabilité de l'Entraînement et Convergence

• Dérive du Déploiement : Sur un horizon autorégressif de 48 heures, l'erreur de prédiction de CLIN-JEPA a convergé (diminution de 15,7 % par rapport à la base de référence à 1 heure). En revanche, les bases de référence et les ablations (suppression de l'échauffement ou de l'alignement) ont divergé de manière catastrophique, avec des augmentations de dérive allant de +3 % à +4951 %.
• Effondrement de la Représentation : Seul le curriculum complet CLIN-JEPA a maintenu un écart-type sain dans les sorties de l'encodeur ($z_{std}$), évitant l'effondrement observé dans les variantes dépourvues de la phase d'échauffement.

2. Diagnostic de la Géométrie Latente

• Discrimination Clinique : L'encodeur a appris une géométrie où les patients en détérioration (augmentation du score SOFA $\ge$ 3) se sont déplacés 4,83 fois plus loin dans l'espace latent que les patients stables. Les encodeurs de base ont montré des ratios de séparation nettement inférieurs (2,62 fois et 1,03 fois).
• Dynamiques Temporelles : Le pouvoir discriminatif de CLIN-JEPA a augmenté avec l'horizon d'observation, culminant à la 66e heure, reflétant une internalisation de l'échelle de temps lente et intégrative de la progression en USI.

3. Performance Multi-Tâches en Aval

Évalué sur deux benchmarks (Prédiction d'Événements Précoces ICareFM et 8 tâches de risque binaires) sans ajustement fin par tâche :

• ICareFM EEP : A atteint une AUROC moyenne de 0,851, surpassant la base de référence la plus forte (LightGBM, 0,827) et les ablations d'encodeur.
• Tâches de Risque Binaires : A atteint une AUROC moyenne de 0,883, surpassant la base de référence LSTM (0,865).
• Valeur du Prédicteur : L'inclusion du déploiement autorégressif (prédictions futures) a constamment amélioré les performances par rapport aux entrées basées uniquement sur l'historique, en particulier pour les tâches nécessitant la composition d'informations de trajectoire hétérogènes (par exemple, lésion rénale, sepsis, décompensation).

Signification

L'article affirme que CLIN-JEPA est le premier cadre à combler avec succès le fossé entre le pré-entraînement JEPA et l'inférence clinique en conservant le prédicteur. Il démontre qu'un seul socle clinique peut fournir un déploiement autorégressif stable, une géométrie latente discriminative cliniquement et des performances multi-tâches compétitives simultanément. Ce travail établit que le co-entraînement de l'encodeur et du prédicteur sous un objectif partagé, guidé par un curriculum spécifique, est nécessaire pour ancrer les représentations pour la simulation à long horizon dans des systèmes dynamiques de haute dimension comme l'USI.