Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

L'article présente CausalLongPFN, un réseau pré-entraîné par ajustement de prior (prior-fitted network) gelé, pré-entraîné sur des modèles causaux temporels synthétiques, qui prédit les résultats longitudinaux contrefactuels sans entraînement spécifique au domaine, démontrant une performance compétitive sur des benchmarks tels que le cancer, le VIH et MIMIC-III tout en offrant une alternative rentable à la modélisation traditionnelle spécifique à une cohorte.

L'essentiel

Imaginez que vous êtes un médecin essayant de prédire comment un patient réagira à un plan de traitement spécifique au cours des prochaines semaines. Le problème est que les patients sont complexes : leur état de santé change de jour en jour, leurs traitements passés affectent leur état actuel, et d'autres facteurs (comme l'alimentation ou le stress) interfèrent avec les résultats. Habituellement, pour faire ces prédictions, vous devez construire un tout nouveau modèle informatique hautement spécialisé pour chaque nouveau groupe de patients que vous rencontrez. C'est comme embaucher un nouvel architecte pour concevoir une maison à chaque fois que vous déménagez dans un nouveau quartier. C'est lent, coûteux et cela nécessite beaucoup de données.

Ce document présente un nouvel outil appelé CAUSALLONGPFN. Considérez-le comme un « Moteur d'Intuition Médicale Universelle » qui a déjà appris les règles du jeu avant même de voir un vrai patient.

Voici comment il fonctionne, décomposé en concepts simples :

1. Le « Camp d'Entraînement » (Pré-entraînement synthétique)

Au lieu d'attendre l'arrivée de vrais patients, les créateurs ont construit un immense « camp d'entraînement » virtuel. Dans ce camp, ils ont simulé des millions de faux patients avec des millions de types de corps, de maladies et de réactions aux traitements différents. Ils ont programmé ces faux patients pour qu'ils présentent des comportements complexes :

• Certains guérissent lentement ; d'autres s'effondrent rapidement.
• Certains traitements fonctionnent immédiatement ; d'autres ont un effet différé.
• Certains patients réagissent différemment en fonction de leur génétique cachée.

L'IA a passé tout son temps dans ce camp, apprenant à prédire les résultats pour ces faux patients. Elle n'a pas seulement mémorisé des réponses ; elle a appris la logique sous-jacente de la manière dont les traitements, le temps et la biologie interagissent.

2. L'« Expert Gelé » (Aucun nouvel entraînement nécessaire)

Voici le tour de magie : une fois que le modèle a terminé son camp d'entraînement, ils l'ont gelé. Ils ont verrouillé son cerveau. Il ne peut rien apprendre de nouveau ni modifier ses paramètres internes.

Lorsqu'un nouveau groupe de vrais patients arrive (comme des patients cancéreux dans un hôpital), le modèle ne repart pas de zéro. Il ne nécessite pas de réentraînement. Au lieu de cela, il agit comme un super-stagiaire qui lit le dossier médical.

• Les Trajectoires de Support : Vous montrez au modèle quelques exemples de patients de l'hôpital actuel (le « support »). Ce sont comme des dossiers de cas que l'on donne au stagiaire pour qu'il comprenne le style spécifique de cet hôpital.
• La Requête : Vous demandez : « Que se passera-t-il pour ce patient spécifique si nous lui administrons le Traitement A pendant les 5 prochains jours ? »
• La Réponse : Le modèle utilise instantanément ce qu'il a appris dans le camp d'entraînement, combiné au contexte des dossiers de cas que vous venez de lui montrer, pour prédire le résultat. Il fait cela sans effectuer une seule étape de « descente de gradient » (un terme technique pour le processus habituel de réentraînement).

3. Le « Simulateur de Voyage dans le Temps »

Le modèle est conçu pour gérer des données longitudinales, ce qui signifie qu'il comprend le temps. Il ne se contente pas de deviner le résultat de demain ; il simule le futur étape par étape.

• Il prédit le Jour 1.
• Il prend cette prédiction et l'utilise comme point de départ pour le Jour 2.
• Il répète ce processus pour voir ce qui se passe au Jour 5, 6 ou 7.

C'est comme un simulateur de vol qui ne montre pas seulement le décollage, mais qui simule tout le parcours de vol en fonction des décisions du pilote, même si la météo change en cours de vol.

4. Pourquoi cela importe (Les Résultats)

Les auteurs ont testé cet « expert gelé » par rapport à l'ancienne méthode (construire un nouveau modèle pour chaque hôpital).

• Le Test : Ils ont utilisé des données sur le cancer, le VIH, les anticoagulants (warfarine) et de vrais dossiers de soins intensifs (réanimation).
• Le Résultat : Le modèle gelé a obtenu des performances aussi bonnes, voire meilleures, que les modèles qui ont été spécialement entraînés pour chaque maladie spécifique.
• La Grande Victoire : Il a particulièrement bien fonctionné sur les données réelles de soins intensifs où ils ne pouvaient pas tester de scénarios « et si » (car on ne peut pas éthiquement tester différents traitements sur de vrais patients dans une simulation). Le modèle pouvait prédire ce qui allait probablement se passer ensuite en se basant uniquement sur les données observées.

L'Essentiel

Le document affirme que vous n'avez pas toujours besoin de construire un nouveau modèle personnalisé pour chaque nouvel ensemble de données médicales. Au lieu de cela, vous pouvez entraîner un seul modèle massif sur une grande variété de scénarios « et si » (données synthétiques) et l'utiliser comme un outil gelé et prêt à l'emploi.

C'est comme avoir un chef cuisinier de haut niveau qui s'est exercé à cuisiner tous les types de cuisine dans une cuisine virtuelle. Lorsque vous lui apportez un nouvel ensemble d'ingrédients (un nouveau groupe de patients), il n'a pas besoin de réapprendre à cuisiner ; il regarde simplement vos ingrédients et sait immédiatement comment préparer le plat, en utilisant sa vaste intuition pré-apprise.

Note Importante : Les auteurs précisent que cet outil est destiné à la prédiction et à la recherche, et non à la prise de décisions de vie ou de mort en clinique. Il aide les médecins à comprendre les résultats potentiels, mais repose toujours sur les mêmes hypothèses médicales standards que n'importe quel autre modèle causal. C'est un puissant outil de recherche, pas un remplacement du jugement d'un médecin.

Pour ceux qui souhaitent explorer l'implémentation ou utiliser le modèle, le code source est disponible sur GitHub : https://github.com/Amirhossein-Zare/causal-long-pfn et les poids du modèle pré-entraîné sont hébergés sur Hugging Face : https://huggingface.co/Amirhossein-Zare/causal-long-pfn .

Résumé technique

Résumé technique : Réseaux de pré-ajustement de prior longitudinaux causaux pour la prédiction de résultats contrefactuels

Formulation du problème

L'article traite du défi de la prédiction des résultats potentiels sous des séquences de traitements futures dans des données observationnelles longitudinales. La tâche centrale consiste à estimer le résultat potentiel conditionnel à l'historique $E[Y(\bar{a})_{t+\tau} | H_t]$, où $H_t$ représente l'historique observé jusqu'au temps $t$, et $\bar{a}$ est une séquence de traitement planifiée.

Ce problème est compliqué par trois facteurs primaires :

1. Confusion temporelle variable : Les assignations de traitement à chaque étape dépendent de covariables qui sont elles-mêmes des résultats de traitements antérieurs.
2. Dynamiques de patients hétérogènes : Les trajectoires individuelles présentent une évolution d'état complexe et non linéaire ainsi qu'une hétérogénéité latente.
3. Limitations des données : Les cohortes observationnelles sont souvent trop petites pour entraîner de manière fiable des modèles de séquences profonds à partir de zéro pour chaque nouveau domaine ou simulateur.

Les estimateurs causaux longitudinaux existants (par exemple, RMSN, CRN, G-Net, Causal Transformer) nécessitent généralement un entraînement supervisé distinct, incluant la sélection d'hyperparamètres et la modélisation de la propension, pour chaque nouvelle cohorte. Ce pipeline est coûteux et peu pratique lorsqu'un entraînement répété spécifique au domaine est requis.

Méthodologie : CAUSALLONGPFN

Les auteurs introduisent les Causal Longitudinal Prior-Fitted Networks (CAUSALLONGPFN), un prédicteur de contexte pré-ajusté sur un prior conçu pour amortir la prédiction causale longitudinale à travers un large prior de modèles structurels causaux temporels (TSCM).

1. Pré-entraînement synthétique sur un prior large

Le modèle est entièrement pré-entraîné sur des épisodes synthétiques échantillonnés à partir d'un large prior sur des TSCM. Ce prior est conçu pour couvrir une large classe de dynamiques causales longitudinales plutôt que de répliquer un seul simulateur. Les caractéristiques clés du prior synthétique incluent :

• Graphes temporels causaux : Dépendances contemporaines et décalées (lagged) éparses avec des graphes instantanés acycliques.
• Mécanismes non linéaires : Les coordonnées d'état suivent des mises à jour autorégressives non linéaires éparses utilisant diverses non-linéarités élémentaires (identité, tanh, sinusoïdal, ReLU, etc.) et diverses distributions de bruit.
• Motifs dynamiques : Des motifs structurés tels que la mémoire-action, la saturation, l'homéostasie, le contrôle par rétroaction et les canaux de lecture lissés sont superposés pour capturer des mécanismes qualitatifs comme les effets retardés et la rétroaction de régulation.
• Politiques de comportement confondées : Les traitements sont échantillonnés à partir de politiques stochastiques dépendant de l'état et influencées par l'hétérogénéité latente de l'unité ($Z_i$), créant une rétroaction traitement-confondeur temporellement variable.
• Modèles de résultats : Les résultats scalaires sont générés via des lectures autorégressives avec des effets de traitement directs et cumulatifs.

2. Architecture

CAUSALLONGPFN emploie une architecture à double encodeur :

• Encodeur d'historique causal : Un Transformer causal de trajectoire (utilisant l'auto-attention masquée) qui cartographie les séquences longitudinales en représentations d'historique, garantissant que la représentation au temps $r$ ne dépend que des informations disponibles jusqu'à ce temps.
• Encodeur de contexte PFN : Un encodeur de contexte basé sur un Transformer qui réalise l'adaptation en contexte. Il traite conjointement les trajectoires de support (traitées comme un ensemble non ordonné) et un jeton de requête (query token) via l'auto-attention. Aucun encodage positionnel n'est assigné à l'ordre des trajectoires de support.
• Tête de prédiction par mélange gaussien : La représentation de requête finale paramètre une distribution de mélange gaussien à 5 composantes pour le résultat normalisé, fournissant à la fois des prédictions ponctuelles et des estimations d'incertitude.

3. Prédiction en contexte et déploiement (Rollout)

Au moment du test, le modèle est gelé. Il reçoit :

• Trajectoires de support : Exemples provenant du nouveau domaine/tâche.
• Historique de requête : Observé jusqu'au temps $t_{obs}$.
• Séquence de traitement future proposée : La séquence d'interventions planifiées.

Le modèle renvoie une distribution prédictive sans mise à jour de gradient, sans ajustement de modèle de propension, ni équilibrage adversarial. Pour une prédiction multi-étapes ($\tau > 1$), le modèle effectue un déploiement (rollout) par injection autorégressive : il prédit la distribution du résultat d'une étape, insère la moyenne du mélange dans l'historique de requête, et répète le processus sous la séquence de traitement spécifiée.

Contributions clés

1. Un modèle de pré-ajustement de prior pour la prédiction causale longitudinale : CAUSALLONGPFN est le premier modèle de type PFN pour la prédiction de résultats potentiels conditionnels à l'historique sous des séquences de traitements longitudinales planifiées. Il fonctionne comme un modèle gelé ne nécessitant aucune adaptation au moment du test.
2. Un prior synthétique sur des tâches causales longitudinales : Les auteurs conçoivent un prior TSCM générant des tâches diverses avec rétroaction traitement-confondeur, hétérogénéité latente, dynamique non linéaire, effets retardés/cumulatifs et changements de régime.
3. Architecture pour l'inférence en contexte longitudinale : Un double encodeur novateur combinant un Transformer causal d'historique avec un encodeur de contexte PFN et une tête de mélange gaussien.
4. Déploiement contrefactuel autorégressif : Une extension du prédicteur d'une étape appris vers une prédiction multi-étapes via un déploiement par injection déterministe.
5. Évaluation zero-shot : Une évaluation complète d'un modèle gelé unique contre des bases de référence entraînées par domaine (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) sur des benchmarks contrefactuels branchables et des données réelles factuelles.

Résultats

Le modèle a été évalué sur quatre benchmarks : croissance tumorale cancéreuse, Warfarin PK/PD, dynamique de traitement du VIH et trajectoires de l'unité de soins intensifs MIMIC-III.

• Performance équilibrée par domaine : CAUSALLONGPFN a obtenu la meilleure RMSE normalisée d'une étape équilibrée par domaine (0,222), dépassant de peu MSM et RMSN. Pour la prédiction à cinq étapes, il se classe troisième au total, derrière RMSN et G-Net, mais surpasse MSM, CRN et les bases de référence à base de transformer.
• Performance par domaine :
  • MIMIC-III (Factuel) : Le modèle arrive premier tant pour la prédiction d'une étape que de cinq étapes, démontrant un fort transfert aux trajectoires cliniques réelles sans entraînement spécifique au domaine.
  • Benchmarks contrefactuels (Cancer, VIH, Warfarin) : Le modèle reste compétitif, se classant deuxième ou troisième sur les tâches d'une étape. Cependant, sur les tâches contrefactuelles à horizon plus long (ex: prédiction de cancer à 5 étapes), des modèles récurrents spécialisés et entraînés par domaine (RMSN, CRN) obtiennent des erreurs plus faibles, suggérant un avantage lorsqu'un volume suffisant de données du domaine cible est disponible pour l'ajustement spécifique.
• Calibration de l'incertitude : La tête de mélange gaussien fournit des informations distributionnelles utiles. La calibration varie selon le domaine, le Warfarin présentant la meilleure calibration et MIMIC-III des intervalles plus larges en raison d'une plus grande hétérogénéité.

Signification, affirmations et disponibilité

L'article affirme qu'un pré-entraînement sur un prior synthétique large peut fournir une alternative gelée utile au ré-entraînement répété spécifique au domaine lorsque :

• Le ré-entraînement de modèles spécialisés est coûteux ou peu pratique.
• Une adaptation rapide à une nouvelle cohorte est requise.
• La supervision contrefactuelle est indisponible (comme dans les tâches de prédiction factuelle du monde réel telles que MIMIC-III).

Les auteurs soulignent que CAUSALLONGPFN ne supprime pas les hypothèses causales standard (consistance, positivité, échangeabilité séquentielle) nécessaires pour interpréter les données observationnelles. Au lieu de cela, il amortit le problème d'estimation. Les résultats suggèrent qu'un prior synthétique suffisamment large peut capturer des structures réutilisables à travers les tâches de traitement-réponse, faisant du modèle un prédicteur en contexte robuste et généraliste. Cependant, l'article note modestement que l'entraînement spécifique au domaine reste précieux lorsqu'il existe des données et des signaux de validation suffisants pour le domaine cible, particulièrement pour les prédictions contrefactuelles à long horizon.

Ce travail se positionne comme un outil de recherche pour la modélisation de séquences causales et la génération d'hypothèses plutôt que comme un système de décision clinique autonome, mettant en garde contre une confiance excessive dans les prédictions lorsque les hypothèses causales ou le support du prior sont inadéquats.

Code et disponibilité des modèles : L'implémentation est disponible sur GitHub à l'adresse https://github.com/Amirhossein-Zare/causal-long-pfn et les poids du modèle pré-entraîné sont disponibles sur Hugging Face à l'adresse https://huggingface.co/Amirhossein-Zare/causal-long-pfn .