Diffusion MRI Transformer with a Diffusion Space Rotary Positional Embedding (D-RoPE)

1. Problématique et Contexte

L'Imagerie par Résonance Magnétique de Diffusion (dMRI) est une technique non invasive cruciale pour étudier la microstructure du cerveau, notamment l'intégrité de la matière blanche. Cependant, l'apprentissage de représentations généralistes à partir de données dMRI reste limité par rapport à d'autres modalités d'imagerie médicale.

Les défis majeurs identifiés sont :

• Nature complexe des données : Les données dMRI ne sont pas de simples volumes 3D statiques. Elles sont composées de plusieurs volumes correspondant à différentes directions de gradient de diffusion et à différentes valeurs de "b" (b-values). Ces données résident dans un espace de diffusion sphérique (défini par des coordonnées sphériques) et non dans un espace cartésien linéaire.
• Dépendances multiples : Il est nécessaire de modéliser conjointement les dépendances spatiales (x, y, z), directionnelles et les variations d'intensité liées aux valeurs de b.
• Hétérogénéité des protocoles : Les acquisitions cliniques varient considérablement (nombre de directions, valeurs de b), ce qui rend difficile l'application de modèles standards qui nécessitent une normalisation rigide ou un rééchantillonnage.
• Limites des approches existantes : Les modèles de deep learning actuels (comme les CNN ou les Transformers standards) ne capturent pas bien les propriétés géométriques uniques de l'espace de diffusion et peinent à généraliser entre différents protocoles d'acquisition.

2. Méthodologie

Les auteurs proposent une architecture basée sur le Transformer, adaptée spécifiquement aux données dMRI, et entraînée via une tâche de Masked Autoencoding (MAE).

Architecture du Modèle

• Entrée : Les volumes dMRI sont découpés en patches 3D. Chaque patch est linéarisé en un token.
• Encodage de Position Absolu : Contrairement aux Transformers classiques, l'encodage de position intègre deux composantes :
  1. Espace spatial (x, y, z) : Encodé via un encodage sinusoïdal 3D.
  2. Espace de diffusion (ρ, θ, φ) : Encodé via une couche linéaire apprenable, représentant la direction du gradient et la valeur de b.
• Mécanisme d'Attention : Inspiré des Transformers vidéo, les blocs d'attention alternent entre les dimensions spatiales et les dimensions de diffusion pour capturer les interactions dans les deux espaces.

Innovation Clé : D-RoPE (Diffusion Space Rotary Positional Embedding)

C'est la contribution centrale de l'article. Le Rotary Positional Embedding (RoPE) standard, conçu pour des séquences 1D, est inadapté à la géométrie sphérique de l'espace de diffusion.

• Concept : Les auteurs généralisent RoPE en remplaçant la distance linéaire simple $(m-n)$ par une distance géométrique $D(m, n)$ dans l'espace de diffusion.
• Formulation de la distance : La distance entre deux tokens $m$ et $n$ (correspondant à des acquisitions $(b_m, v_m)$ et $(b_n, v_n)$) est définie comme une combinaison pondérée de la différence de valeurs de b et de l'angle entre les vecteurs de gradient :
  $$D(m, n) = \sqrt{\gamma(b_m - b_n)^2 + \arccos^2(|v_m \cdot v_n|)}$$
  où $\gamma$ est un paramètre de pondération. Le terme $\arccos$ assure une symétrie de 180° (les directions opposées $v$ et $-v$ sont traitées comme identiques car elles produisent un encodage de diffusion équivalent).
• Avantage : Cela permet au modèle de comprendre les relations relatives entre les directions de diffusion, indépendamment du nombre total de directions acquises.

Entraînement (Pretraining)

Le modèle est pré-entraîné avec un objectif MAE (Masked Autoencoder) :

• Une partie des tokens (patches) est masquée.
• Le modèle doit reconstruire les volumes originaux à partir des tokens visibles.
• Stratégies de masquage : Les auteurs testent trois stratégies : masquage spatial (75%), masquage directionnel (50% des directions entières masquées), et une alternance des deux. La stratégie alternée s'avère la plus performante.

3. Contributions Principales

1. Architecture Transformer pour dMRI : Proposition d'un modèle capable de traiter des acquisitions dMRI arbitraires (nombre variable de directions et de valeurs de b) sans rééchantillonnage.
2. D-RoPE : Introduction d'un encodage de position rotatif spécifique à l'espace de diffusion, permettant de modéliser les interdépendances géométriques complexes entre les directions de gradient et les valeurs de b.
3. Représentations Transférables : Démonstration que le pré-entraînement auto-supervisé sur de grandes quantités de données (HCP) permet d'obtenir des représentations robustes transférables à des tâches cliniques avec peu de données (ADNI).

4. Résultats Expérimentaux

Les expériences ont été menées sur les ensembles de données HCP (pour le pré-entraînement) et ADNI (pour les tâches en aval).

Reconstruction (MAE)

• L'ajout de D-RoPE combiné à un masquage alterné (spatial + diffusion) améliore significativement la qualité de reconstruction par rapport aux modèles de base (ViT standard).
• Métriques : Réduction de 37-40% du FID (Fréchet Inception Distance) et amélioration du PSNR et SSIM.
• Validité anatomique : Les cartes de Fractional Anisotropy (FA) et de Mean Diffusivity (MD) dérivées des volumes reconstruits présentent une faible erreur par rapport aux vérités terrain, prouvant que la microstructure est préservée.

Tâches en Aval (Downstream Tasks)

Le modèle pré-entraîné a été évalué via fine-tuning partiel ou linear probing sur des tâches de classification et de régression :

1. Prédiction de l'âge et classification du sexe : Le modèle atteint des performances compétitives, souvent supérieures aux méthodes supervisées classiques (3D ResNet) ou aux Transformers standards avec RoPE classique, même avec des représentations "gelées" (frozen).
2. Classification CN vs MCI (Déclin Cognitif Léger) :
   • Le modèle pré-entraîné avec une tête MLP sur des latents gelés atteint une précision équilibrée (BA) de 64,7%, surpassant les caractéristiques manuelles (Hand-Crafted) et les modèles entraînés de zéro.
   • Cela suggère que le modèle a appris des motifs microstructuraux prédictifs sans nécessiter de modèle biophysique explicite.
3. Prédiction du score ADAS-Cog :
   • Amélioration significative du coefficient de corrélation (+0,05 par rapport aux meilleurs baselines) pour la prédiction des scores cognitifs.
   • C'est la première étude à prédire directement un score cognitif à partir de volumes dMRI bruts avec une telle performance.

5. Signification et Conclusion

Ce travail marque une avancée significative dans le domaine de l'apprentissage de représentations pour l'imagerie par résonance magnétique de diffusion :

• Harmonisation des protocoles : Le modèle peut traiter des données provenant de protocoles d'acquisition hétérogènes sans nécessiter de rééchantillonnage vers un schéma fixe, ce qui est crucial pour l'analyse de grandes cohortes cliniques.
• Fondation pour un "Foundation Model" : L'approche ouvre la voie vers un modèle fondamental (Foundation Model) pour la dMRI, capable d'être adapté à diverses tâches cliniques avec peu de données étiquetées.
• Efficacité des données : La capacité à obtenir de meilleures performances avec des données pré-entraînées gelées ou partiellement affinées démontre l'efficacité de l'apprentissage auto-supervisé pour extraire des biomarqueurs microstructuraux pertinents.
• Limites : L'auteur note que le calcul de D-RoPE est légèrement plus coûteux ($O(lh^2)$) que le RoPE standard ($O(lh)$) en raison de la complexité de la distance géométrique, mais ce coût est jugé acceptable au vu des gains de performance.

En résumé, l'introduction de D-RoPE permet de combler le fossé entre la géométrie sphérique des données dMRI et les architectures Transformer modernes, offrant des représentations robustes et généralisables pour le diagnostic et la recherche en neurologie.