Technische Samenvatting: MultiMAE voor Brain MRI

Probleemstelling

In de medische beeldvorming, met name bij hersen-MRI, worden deep learning-modellen vaak geconfronteerd met de uitdaging van ontbrekende inputsequenties. Klinische omgevingen missen vaak volledige multi-modale data (bijv. ontbrekende T1, T1c, T2 of FLAIR-sequenties) vanwege acquisitieproblemen of patiëntbeperkingen. De gebruikelijke benaderingen maken gebruik van een early-fusion strategie, waarbij alle beschikbare modaliteiten langs de kanaaldimensie worden gestapeld (analoog aan RGB-kanalen). Hoewel deze methode eenvoudig is, vereist deze benadering dat alle modaliteiten aanwezig zijn tijdens zowel de training als de inferentie. Bijgevolg lijden modellen die getraind zijn met early fusion onder scherpe prestatieverlies wanneer inputs ontbreken, en slagen ze er niet in om onderscheidende, modaliteit-specifieke representaties te leren die kunnen worden ingezet voor cross-modale redenering.

Methodologie

De auteurs stellen een Multi-Modal Masked Autoencoder (MultiMAE) framework voor, aangepast voor 3D hersen-MRI. Geïnspireerd door de MultiMAE-architectuur voor natuurlijke afbeeldingen, verschuift deze methode van early fusion naar een modaliteit-specifieke encoding strategie gecombineerd met late-fusion via een transformer encoder.

Architectuur en Pretraining

• Backbone: Het model maakt gebruik van een aangepaste Vision Transformer (ViT-B/16) als encoder.
• Modality-Specific Adapters: In plaats van kanalen te stapelen, wordt elke MRI-sequentie (t1, t1c, t2, fla) verwerkt door zijn eigen lineaire adapter om niet-overlappende $16 \times 16 \times 16$ patches naar een gedeelde tokenruimte te projecteren.
• Masking Strategie: Een globale masking ratio van 75% wordt toegepast. Cruciaal is dat masking per-modaliteit wordt gesampled uit een Dirichlet-distributie ($\alpha=1$), waardoor tot 100% van een specifieke modaliteit gemaskeerd kan worden terwijl de globale ratio behouden blijft. Dit dwingt het model om cross-modale afhankelijkheden te leren.
• Encoder-Decoder Design:
  • Encoder: Een standaard ViT-B/16 (12 lagen) verwerkt de niet-gemaskeerde tokens van alle modaliteiten samen met een [cls] token.
  • Decoders: Vier lichtgewicht, modaliteit-specifieke decoders (elk 3 lagen) worden ingezet. Deze maken gebruik van een cross-attention mechanisme waarbij queries (leerbare [mask] tokens voor ontbrekende patches) aandacht besteden aan de gedeelde context die door de encoder wordt geleverd.
• Doelstelling: Het model is getraind om gemaskeerde patches te reconstrueren met behulp van een Mean Squared Error (MSE) loss, wat het model stimuleert om rijke representaties te leren die ontbrekende sequenties kunnen afleiden uit beschikbare sequenties.

Downstream Adaptatie

De getrainde encoder wordt aangepast aan twee downstream taken:

1. Segmentatie: Gebruikmakend van de UNETR architectuur. Om de variabele sequentielengte te hanteren die wordt veroorzaakt door de aparte tokenisatie per modaliteit, worden tokens op elke ruimtelijke locatie gemiddeld om een consistente sequentielengte te behouden.
2. Classificatie: Een enkele lineaire laag vervangt de decoders. Patch-tokens en de [cls] token worden gemiddeld om een feature vector te vormen voor glioma-subtype classificatie (GBM, Astro, Oligo).

Belangrijkste Bijdragen

1. MultiMAE Pretraining voor 3D MRI: Het werk onderzoekt de toepassing van MultiMAE op 3D hersen-MRI, wat modaliteit-specifieke encoding en gemaskeerde modellering over meerdere sequenties mogelijk maakt.
2. Robuustheid tegen Ontbrekende Modaliteiten: De auteurs demonstreren dat deze pretraining strategie de prestaties en robuustheid aanzienlijk verbetert in zowel segmentatie (glioma en metastase) als classificatietaken wanneer input-modaliteiten ontbreken, waarbij het de standaard MAE-ViT baselines overtreft.
3. Synthese van Ontbrekende Modaliteiten: Het framework demonstreert het vermogen om tijdens de inferentie volledig ontbrekende MRI-sequenties te synthetiseren door gebruik te maken van de geleerde cross-modale relaties tijdens de pretraining.

Experimentele Resultaten

De methode werd geëvalueerd op een grote gecombineerde dataset (BraTS 2021/2023, UPENN-GBM, UCSF-PGDM, etc.) en getest op interne en externe datasets.

• Generatie van Ontbrekende Modaliteiten: Het model slaagde erin ontbrekende modaliteiten met hoge getrouwheid te reconstrueren (PSNR ~20–21, SSIM ~0.71–0.76), wat de standaard MAE-ViT baseline (PSNR ~7–10, SSIM ~0.52–0.61) aanzienlijk overtreft.
• Segmentatie Prestaties:
  • De getrainde MultiMAE presteerde consequent beter dan de MAE-ViT baseline op zowel interne als externe testsets voor glioma- en metastase-segmentatie.
  • In scenario's met ontbrekende modaliteiten vertoonde de MultiMAE een superieure veerkracht. Bijvoorbeeld, wanneer de T1c-modaliteit ontbrak, behield de MultiMAE een Dice-score van ~51.55 (frozen encoder) voor de tumorkern, terwijl de ViT-baseline daalde naar ~47.91.
  • Het model behaalde een absolute verbetering van 10.1 in de algemene Dice-score en 0.46 in MCC ten opzichte van de baselines wanneer inputs ontbraken.
• Classificatie Prestaties:
  • Het model behaalde een sterke multi-class nauwkeurigheid en MCC (tot 0.70) voor glioma-subtypering.
  • In tegenstelling tot de segmentatietaak was het bevriezen van de encoder minder effectief voor classificatie; volledige fine-tuning leverde de beste resultaten op.
  • Zelfs met ontbrekende modaliteiten (bijv. geen T1c) behield de MultiMAE een hoge nauwkeurigheid (0.89) vergeleken met de ViT-baseline (0.41), wat duidt op effectieve cross-modale redenering.

Betekenis en Claims

Het artikel claimt dat het MultiMAE-paradigma een flexibele en generaliseerbare encoder biedt voor hersen-MRI die een cruciaal klinisch probleem van incomplete data aanpakt. Door elke sequentie als een aparte modaliteit te behandelen en cross-modale reconstructie af te dwingen, leert het model ontbrekende sequenties af te leiden uit beschikbare inputs.

De auteurs benadrukken dat hoewel de aanpak kleine wijzigingen in de ViT-backbone introduceert, het substantiële winst oplevert in robuustheid. De pretraining strategie stelt het model in staat om "ongeziene inputconfiguraties" te hanteren zonder hertraining, wat het geschikt maakt voor klinische omgevingen met beperkte middelen waar de volledigheid van data niet gegarandeerd kan worden. De auteurs merken echter bescheiden op dat:

• Reconstructies enige onscherpte vertonen door de beperkingen van pixel-wise MSE loss.
• Er een domeingap bestaat tussen interne en externe datasets, wat de bredere uitdaging van cross-site adaptatie weerspiegelt.
• De toegevoegde complexiteit van cross-modale integratie aanzienlijke pretraining compute en data vereist om te convergeren, aangezien willekeurig geïnitialiseerde versies van het model in sommige gevallen slechter presteerden.

Uiteindelijk presenteert dit werk een herformulering van het MAE-pretraining paradigma die een elegantere omgang met multi-modale hersen-MRI data mogelijk maakt, waarbij de nadruk ligt op robuustheid in het licht van ontbrekende inputs.