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

Probleemstelling

Diffusie-MRI (dMRI) is een cruciale techniek voor het bestuderen van microstructurele veranderingen in de hersenen, maar het ontwikkelen van algemene, overdraagbare representaties voor deze data is beperkt gebleven. Bestaande deep learning-modellen zijn niet optimaal geschikt voor dMRI vanwege twee hoofduitdagingen:

1. Complexe data-structuur: dMRI bestaat uit meerdere volumina die afhankelijk zijn van de richting en sterkte van diffusie-gevoelige gradiënten (b-vectoren en b-waarden). Modellen moeten zowel ruimtelijke afhankelijkheden als richtingsafhankelijkheden in de diffusieruimte gezamenlijk modelleren.
2. Variatie in acquisitieprotocollen: Klinische studies gebruiken vaak verschillende protocollen met een verschillend aantal diffusierichtingen en b-waarden. Traditionele modellen hebben moeite met deze variabiliteit en vereisen vaak upsampling naar een gemeenschappelijk protocol, wat informatie kan verliezen.

Methodologie

De auteurs stellen een nieuw transformer-architectuur voor, gebaseerd op Masked Autoencoders (MAE), die specifiek is ontworpen voor de unieke geometrie van dMRI-data.

1. Architectuur en Invoer:

• De input bestaat uit een reeks 3D-diffusievolumina. Deze worden opgedeeld in 3D-patches en gemapt naar tokens.
• In tegenstelling tot standaard Vision Transformers (ViT) die alleen ruimtelijke posities coderen, combineert dit model absoluut positionele encoding voor zowel de ruimtelijke coördinaten $(x, y, z)$ als de diffusieruimte-coördinaten (sferische coördinaten $\rho, \theta, \phi$).

2. Diffusion Space Rotary Positional Embedding (D-RoPE):
Dit is de kerninnovatie van het paper. Standaard Rotary Positional Embeddings (RoPE) zijn ontworpen voor lineaire sequenties (zoals tekst of video-tijd). dMRI-data ligt echter in een sferische ruimte.

• De auteurs generaliseren RoPE door de lineaire afstand $(m-n)$ te vervangen door een geometrisch gedefinieerde afstand $D(m, n)$ tussen twee diffusie-acquisities.
• De afstand wordt berekend als een gewogen combinatie van het verschil in b-waarde en de hoek tussen de b-vectoren:
  $$D(m, n) = \sqrt{\gamma(b_m - b_n)^2 + \arccos^2(|v_m \cdot v_n|)}$$
• Hierbij zorgt de absolute waarde in de arc-cosine voor 180° symmetrie (tegenovergestelde gradiëntvectoren worden als identiek behandeld).
• Deze D-RoPE wordt geïntegreerd in de attention-mechanismen, waardoor het model de onderlinge afhankelijkheid tussen verschillende diffusierichtingen en b-waarden kan leren, ongeacht het aantal richtingen.

3. Training (Self-Supervised Learning):

• Het model wordt voorgeprogrammeerd (pretrained) met een Masked Autoencoding (MAE) doelstelling.
• Er worden drie maskeringstrategieën getest: ruimtelijk maskeren, maskeren in de diffusierichting (hele richtingen weglaten), en een afwisselende strategie.
• Het doel is om de gemaskerde patches te reconstrueren, wat het model dwingt om zowel ruimtelijke als microstructurele patronen te leren.

Belangrijkste Bijdragen

1. Nieuwe Architectuur: Een transformer-ontwerp dat willekeurige dMRI-acquisities kan verwerken, ongeacht het aantal diffusierichtingen of b-waarden.
2. D-RoPE: Een nieuwe positie-encoding die de geometrie van de diffusieruimte (b-waarde en b-vector) expliciet modelleert binnen de attention-mechanismen.
3. Uitgebreide Evaluatie: Een grondige test van de geleerde representaties op meerdere downstream-taken, inclusief leeftijdsvoorspelling, geslachtsclassificatie, en klinische beoordeling van cognitieve stoornissen.

Resultaten

De prestaties zijn geëvalueerd op datasets van het Human Connectome Project (HCP) en ADNI.

• Reconstructie (Pretraining): Het model met D-RoPE en afwisselende maskering behaalde de beste reconstructie-kwaliteit (hoogste PSNR/SSIM, laagste FID) en produceerde de meest anatomisch plausibele diffusie-metrieken (Fractional Anisotropy en Mean Diffusivity).
• Downstream Taken:
  • Leeftijd en Geslacht: Het gefinetunede model presteerde vergelijkbaar met of beter dan volledig getrainde baselines (zoals 3D ResNet en ViT), zelfs met beperkte data.
  • Cognitieve Classificatie (MCI vs. CN): Het gebruik van de "frozen" (niet-finetuned) latent features van de pretrained encoder met een MLP-head resulteerde in een balanced accuracy van 64,7%, wat een verbetering is ten opzichte van handgemaakte features (61,9%) en andere baselines.
  • Voorspelling Cognitieve Scores (ADAS): Het model bereikte een correlatiecoëfficiënt van 0,38, een verbetering ten opzichte van de baselines.
• Robuustheid: Het model toonde aan dat het kan werken met willekeurig geselecteerde richtingen per subject, wat aantoont dat het niet afhankelijk is van een vast acquisitieprotocol.

Betekenis en Impact

Dit werk is een belangrijke stap richting een foundation model voor dMRI.

• Generalisatie: Het lost het probleem op van variabele acquisitieprotocollen, waardoor analyses kunnen worden geharmoniseerd over verschillende cohorten zonder resampling.
• Klinische Toepasbaarheid: De resultaten tonen aan dat self-supervised geleerde representaties waardevolle microstructurele informatie bevatten die direct gerelateerd is aan cognitieve achteruitgang en ziekteprogressie (zoals MCI en Alzheimer), zelfs zonder expliciete biophysieke modellering.
• Toekomstperspectief: De methode biedt een basis voor toekomstige multi-modale modellen en kan worden uitgebreid naar andere toepassingen zoals super-resolutie of het ontdekken van nieuwe biomarkers voor witte stof-integriteit.

Kortom, door de unieke geometrie van diffusie-data expliciet te modelleren via D-RoPE, slaagt dit paper erin om robuuste, overdraagbare en klinisch relevante representaties te leren uit ruwe dMRI-data.