BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Résumé technique de "BrainAnytime : Préentraînement intermodal conscient de l'anatomie pour l'analyse d'images cérébrales avec disponibilité de modalité arbitraire"

Énoncé du problème

Les flux de travail de diagnostic clinique pour les maladies neurodégénératives, en particulier la maladie d'Alzheimer (MA), suivent une « voie d'escalade de modalité ». Les cliniciens commencent généralement par des imageries structurelles de routine (par exemple, IRM pondérée en T1), ajoutent sélectivement des séquences (par exemple, FLAIR, T2) pour affiner les diagnostics différentiels, et réservent l'imagerie moléculaire (par exemple, amyloïde-TEP) aux cas incertains. Par conséquent, les données des patients dans les cohortes réelles se caractérisent par des sous-ensembles de modalités hétérogènes et souvent incomplets.

Les modèles d'IA actuels supposent généralement des modalités d'entrée fixes, nécessitant un changement d'architecture ou un réentraînement lorsque la disponibilité des données change. De plus, les modèles existants de préentraînement à grande échelle sur l'imagerie cérébrale souffrent souvent de trois limitations dans ce contexte :

1. Ils sont généralement réservés à l'IRM, échouant à apprendre des ponts explicites vers l'amyloïde-TEP, un biomarqueur critique.
2. Ils supposent des combinaisons de modalités fixes, ne correspondant pas à l'absence dépendante du stade des données réelles.
3. Le préentraînement standard par Autoencodeur Masqué (MAE) repose sur un masquage spatialement uniforme, ignorant le fait que la MA présente une vulnérabilité régionale sélective (par exemple, l'hippocampe) plutôt qu'une pathologie uniforme.

Méthodologie

Les auteurs proposent BrainAnytime, un cadre de préentraînement unifié conçu pour soutenir l'analyse d'images cérébrales sous une disponibilité de modalité arbitraire. Le cadre est préentraîné sur 34 899 scans cérébraux 3D provenant de cinq ensembles de données (A4, DOD-ADNI, AIBL, BraTS23, NACC) couvrant l'IRM multi-séquences et l'amyloïde-TEP. Il intègre trois composants principaux au sein d'une architecture partagée d'autoencodeur masqué 3D :

1. Multi-MAE3D (Autoencodeur Masqué Intermodal 3D)

• Architecture : Un encodeur Transformer partagé traite n'importe quel sous-ensemble de modalités disponibles (T1, T2, FLAIR, TEP).
• Tokenisation : Chaque modalité utilise une intégration de patchs convolutifs 3D spécifique à la modalité.
• Gestion des données manquantes : Les modalités manquantes sont entièrement masquées. Pour maintenir un nombre de tokens constant indépendamment de la disponibilité de l'entrée, le budget de tokens visibles est redistribué entre les modalités observées à l'aide d'une distribution de Dirichlet (Dir($\alpha$)).
• Décodage : Chaque modalité possède un décodeur léger qui reconstruit tous les patchs. La perte d'entraînement ($L_{MAE}$) est calculée uniquement sur les patchs masqués en utilisant l'erreur quadratique moyenne avec une normalisation par patch.

2. Distillation Intermodale Réciproque (RCMD)

• Objectif : Apprendre les correspondances entre l'IRM structurelle et la TEP moléculaire.
• Mécanisme : Un réseau enseignant à Moyenne Mobile Exponentielle (EMA) partage l'architecture de l'encodeur. Pour les échantillons appariés IRM-TEP, les représentations au niveau du groupe ($z_{IRM}$ et $z_{TEP}$) sont calculées par moyennage.
• Perte : Deux prédicteurs MLP symétriques ($g_{M\to P}$ et $g_{P\to M}$) sont entraînés pour prédire les caractéristiques de l'enseignant intermodales. La perte ($L_{RCMD}$) maximise la similarité cosinus entre les caractéristiques prédites par l'élève et les caractéristiques de l'enseignant dans l'espace de modalité opposé.

3. Masquage de Curriculum Conscient de la Pathologie (PACM)

• Objectif : Prioriser l'anatomie vulnérable à la maladie lors de la reconstruction plutôt que de masquer uniformément.
• Évaluation de l'importance : Un score statique est dérivé des priors neuropathologiques (stade de Braak) mappés via l'atlas AAL116. Un score dynamique est obtenu en agrégeant les cartes d'attention d'un enseignant anatomique EMA. Ceux-ci sont combinés pour générer des scores d'importance des patchs.
• Calendrier du curriculum : Une stratégie en trois phases est employée :
  1. Phase précoce : Masquage aléatoire uniforme.
  2. Phase intermédiaire : Le paramètre de température ($\tau$) recuit par décroissance cosinus, affinant progressivement la distribution du masquage pour favoriser le masquage des régions à haute importance (c'est-à-dire forçant le modèle à reconstruire des régions critiques).
  3. Phase tardive : Température fixe.
• Sélection : Les patchs sont sélectionnés en utilisant l'astuce Gumbel-top-k, garantissant que les patchs à haute importance ont plus de chances d'être masqués et donc reconstruits.

Entraînement et Affinage :
La perte totale est $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[apparié]} \cdot L_{RCMD}$. Pour les tâches en aval, seul l'encodeur élève est utilisé avec une tête de tâche légère. Les modalités manquantes restent à zéro et bloquées par l'attention lors de l'inférence, garantissant que le modèle gère des sous-ensembles arbitraires sans changement d'architecture.

Contributions clés

1. Cadre unifié : Un modèle unique capable d'accepter n'importe quel sous-ensemble de modalités (d'un simple scan T1 à un bilan multimodal complet) sans réentraînement ni changement d'architecture.
2. Apprentissage intermodal : Apprend explicitement les correspondances structurelles-moléculaires entre l'IRM et la TEP par distillation réciproque, comblant le fossé entre les preuves structurelles et pathologiques.
3. Préentraînement conscient de l'anatomie : Introduit une stratégie de masquage de curriculum qui oriente l'effort de reconstruction du modèle vers la neuroanatomie saillante de la MA, répondant aux limitations du masquage uniforme dans les MAE standards.
4. Préentraînement évolutif : Préentraîné avec succès sur un ensemble de données hétérogène à grande échelle (34 899 scans) imitant les schémas d'absence clinique.

Résultats expérimentaux

Le modèle a été évalué sur quatre tâches en aval (CN vs MA, CN vs MCI, régression MMSE, prédiction de l'âge) à travers cinq configurations de modalités motivées cliniquement (allant de T1 seul à T1+T2+FLAIR+TEP).

• Performance : BrainAnytime a surpassé les modèles spécifiques à la modalité (par exemple, 3D ResNet, M3T), les méthodes de fusion bimodale, les modèles de base cérébraux à grande échelle (BrainIAC, SAM-Brain3D) et les approches d'apprentissage multimodal incomplet (Flex-MoE, FuseMoE).
• Métriques clés :
  • CN vs MA : A atteint une précision moyenne de 89,4 % (contre 84,2 % pour la prochaine meilleure référence, Flex-MoE), représentant une amélioration relative de 6,2 % par rapport à la référence la plus forte avec modalité manquante.
  • CN vs MCI : A atteint une précision moyenne de 67,2 % (contre 60,6 % pour MoE-retriever), une amélioration relative de 7,0 %.
  • Régression : A démontré une performance supérieure dans la prédiction du MMSE et de l'âge pour toutes les configurations de modalités.
• Robustesse : Une étude de robustesse intra-sujet a confirmé que la performance s'améliorait constamment à mesure que plus de modalités étaient ajoutées (par exemple, le F1 CN vs MA est passé de 65,0 % avec T1 seul à 80,4 % avec une entrée complète), démontrant la capacité du modèle à exploiter efficacement les données disponibles.
• Ablation : La suppression individuelle de RCMD ou PACM a entraîné des baisses de performance, confirmant la nécessité de la distillation intermodale et du masquage conscient de l'anatomie.

Signification

L'article affirme que BrainAnytime comble un fossé critique entre la réalité clinique (acquisition de données échelonnée et incomplète) et les paradigmes actuels de l'IA (modèles à entrée fixe). En permettant à un modèle unique d'accepter flexiblement toute imagerie disponible, il s'aligne sur la voie d'escalade de modalité clinique. Le cadre démontre que le préentraînement avec distillation intermodale et masquage conscient de la pathologie produit un modèle de base robuste qui surpasse significativement les références existantes, en particulier dans les scénarios où les données sont incomplètes ou où des modalités manquent. Cette approche offre une solution pratique pour le déploiement de l'IA dans des contextes cliniques réels où des données multimodales complètes sont rarement disponibles pour chaque patient.