MHMamba: Multi-Head Mamba for 3D Brain Tumor Segmentation

Technische Samenvatting: MHMamba voor 3D Hersentumorsegmentatie

Probleemstelling
Hersentumorsegmentatie in 3D-MRI staat voor aanzienlijke uitdagingen door de hoge heterogeniteit van tumorvorm en multimodale contrast. Hoewel Convolutionele Neuronale Netwerken (CNN's) uitblinken in het modelleren van lokale texturen en grenzen, kampen ze met beperkingen op lange afstand vanwege hun beperkte receptieve velden. Omgekeerd kunnen op Transformer gebaseerde modellen globale context vastleggen, maar lijden ze onder kwadratische reken- en geheugencomplexiteit ($O(N^2)$), waardoor ze moeilijk toepasbaar zijn op hoogresolutie 3D-volumes zonder agressieve downsampling of venstering, wat kan leiden tot contextuele oncoherentie en inconsistenties aan de randen. Bovendien lijden bestaande hybride benaderingen vaak aan statistische mismatches tussen convolutie- en attentie-branches (bijvoorbeeld Batch/Instance Normalisatie versus Layer Normalisatie), wat leidt tot onstabiel training en problemen bij het fuseren van kenmerken.

Methodologie
Het artikel stelt MHMamba (Multi-Head Mamba) voor, een U-vormige encoder-decoder architectuur ontworpen om globale kenmerken te modelleren met lineaire complexiteit terwijl hoogresolutie details behouden blijven. Het kader bestaat uit drie hoofdcomponenten:

1. Multi-Head Mamba (MHM) Module:

   • Vervangt standaard blokken in de encoder door een multi-head State Space Model (SSM).
   • De kanaal-dimensie wordt opgesplitst in parallelle SSM-heads ($N_h$), waardoor verschillende heads complementaire patronen van lange-afstandsafhankelijkheid kunnen leren van multimodale invoer zonder vaste toewijzingen aan modaliteiten.
   • Elke head serialiseert 3D-ruimtelijke kenmerken tot een sequentie, past selectieve staat-updates toe (Mamba) en fuseert de outputs via concatenatie en projectie.
   • Dit ontwerp behoudt lineaire complexiteit terwijl het de globale modellering en trainingsstabiliteit onder multimodale condities verbetert.

2. Channel-Spatial Calibration Attention (CSCA) Module:

   • Ontworpen om kenmerkstatistieken af te stemmen en laesie-responsen te versterken na de multi-head Mamba-verwerking.
   • Het maakt gebruik van een dubbel-pad mechanisme:
     • Kanaal-pad: Gebruikt globale gemiddelde en max-pooling gevolgd door een MLP om kanaalgewichten te genereren.
     • Ruimtelijk-pad: Integreert innovatief vier ruimtelijke statistieken (gemiddelde, standaardafwijking, maximum en minimum) om een ruimtelijke gewichtskarte te genereren, waarbij heterogeniteit en opvallende pieken worden vastgelegd.
   • Een gating-mechanisme fuseert deze branches adaptief, gevolgd door een residuale verbinding om oorspronkelijke kenmerkresponsen te behouden.

3. Adaptive Gated Fusion (AGF) Module:

   • Geplaatst bij de skip-verbindingen tussen de encoder en decoder.
   • Adresst het conflict tussen hoog-niveau semantiek en laag-niveau details door dynamisch multi-schaal kenmerken te fuseren.
   • Invoerkenmerken worden opgesplitst in subgroepen langs de kanaal-dimensie. Voor elke subgroep leert de module onafhankelijke fusiegewichten om encoder- en decoder-kenmerken in evenwicht te brengen, waardoor kenmerkconflicten worden voorkomen die vaak leiden tot inconsistente randvoorspellingen.

Belangrijkste Bijdragen

• Multi-Head Mamba Architectuur: Het artikel introduceert een multi-head SSM-benadering voor 3D-MRI die kanalen decomposeert in parallelle heads. Dit verbetert de modellering van lange-afstandsafhankelijkheid en multimodale aanpasbaarheid, terwijl het het voordeel van lineaire complexiteit van Mamba behoudt.
• Kenmerkcalibratie en Fusie: De integratie van de CSCA-module om kenmerkschalen en statistieken te uniformeren, en de AGF-module om multi-schaal kenmerken dynamisch te fuseren bij skip-verbindingen, verbetert significant de consistentie van randen en de detectie van laesies met klein volume.
• Systematische Validatie: De methode is rigoureus geëvalueerd op de BraTS2021 en BraTS2023 datasets, waarbij stabiele verbeteringen werden aangetoond ten opzichte van sterke baselines (waaronder nnU-Net, SegMamba en op Transformer gebaseerde modellen) in termen van Dice-scores en Hausdorff Afstand (HD95).

Experimentele Resultaten
Experimenten uitgevoerd op de BraTS2021 en BraTS2023 datasets tonen aan dat MHMamba state-of-the-art of near-state-of-the-art prestaties behaalt:

• BraTS2021: Bereikte een gemiddelde Dice-score van 91,02% en een gemiddelde HD95 van 3,38, beter dan baselines zoals SegMamba (90,32% Dice, 4,18 HD95) en nnU-Net (89,40% Dice, 5,58 HD95).
• BraTS2023: Bereikte een gemiddelde Dice-score van 90,23% en een gemiddelde HD95 van 3,67.
• Ablatiestudies:
  • De individuele toevoeging van CSCA, AGF en de Multi-Head Mamba-module droeg bij aan prestatiewinsten over alle tumorsubregio's (Hele Tumor, Tumor Core en Versterkende Tumor).
  • De optimale configuratie bleek 4 heads te zijn, wat de representatieve capaciteit en rekenefficiëntie in evenwicht bracht; een verhoging naar 8 heads leidde tot prestatiedegradeering door over-partitionering en optimalisatieruis bij 3D-training met kleine batchgroottes.
  • Kwalitatieve resultaten gaven aan dat MHMamba minder gefragmenteerde voorspellingen produceert met strakkere randen, met name in heterogene tumorregio's en kleine versterkende gebieden.

Betekenis en Beweringen
Het artikel beweert dat MHMamba effectief de afweging tussen globale contextmodellering en het behoud van lokale details in 3D-hersentumorsegmentatie aanpakt. Door gebruik te maken van de lineaire complexiteit van State Space Models via een multi-head architectuur, vermijdt de methode de kwadratische geheugenoverhead van Transformers, terwijl de beperkingen van lokale receptieve velden van CNN's worden overwonnen. De voorgestelde modules (CSCA en AGF) richten zich specifiek op de statistische mismatches en contextuele oncoherentie die vaak voorkomen in multimodale 3D-segmentatie, resulterend in een robuust, efficiënt en nauwkeurig kader dat geschikt is voor klinische toepassingen zoals bestralingsplanning en preoperatieve beoordeling. De auteurs benadrukken dat de methode de efficiëntievoordelen van op Mamba gebaseerde modellering behoudt, terwijl de gevoeligheid voor tumorcores en laesies met klein volume significant wordt verbeterd.