Technische Samenvatting: Patch-MoE Mamba

Probleemstelling

Medische beeldsegmentatie is sterk afhankelijk van Convolutional Neural Networks (CNN's) en Transformers. Echter, CNN's zijn inherent beperkt door lokale receptieve velden en worstelen met het modelleren van lange-afstand afhankelijkheden, terwijl Transformers vaak een kwadratische reken- en geheugencomplexiteit met zich meebrengen. Recent zijn State Space Models (SSM's), met name Mamba-gebaseerde architecturen, naar voren gekomen als efficiënte alternatieven met lineaire sequentiecomplexiteit.

Ondanks hun veelbelovende aard, staan bestaande Mamba-gebaseerde segmentatiemodellen voor twee kritieke beperkingen:

1. Verstoring van de lokale ruimtelijke structuur: De meeste modellen converteren 2D-kenmerkkaarten naar 1D-sequenties met behulp van pixel-gebaseerde raster-scanning. Dit proces kan ruimtelijk aangrenzende pixels in de sequentie scheiden, waardoor de lokale 2D ruimtelijke coherentie verzwakt. Dit is nadelig voor dichte voorspellingstaken die nauwkeurige grenslocalisatie en behoud van lokale structuur vereisen.
2. Ongewenste directionele fusie: Huidige methoden fuseren doorgaans outputs van meerdere scanrichtingen via eenvoudige sommatie. Deze vaste aggregatie gaat ervan uit dat alle richtingen op elke ruimtelijke locatie even belangrijk zijn, waardoor ze niet kunnen aanpassen aan de diverse maten, vormen, texturen en grenscomplexiteiten van medische doelen (bijvoorbeeld kleine, laag-contrast laesies versus grote, onregelmatige structuren).

Methodologie

De auteurs stellen Patch-MoE Mamba voor, een patch-geordende Mixture-of-Experts (MoE) state space-architectuur die is ontworpen om de bovengenoemde beperkingen aan te pakken. De architectuur volgt een U-Net-stijl ontwerp, waarbij de standaard Visual State Space (VSS) blokken in VM-UNetV2 worden vervangen door een nieuw Patch-MoE VSS blok.

1. Hiërarchische patch-geordende scanning

In plaats van individuele pixels te scannen via raster, partitioneert de methode kenmerkkaarten in lokale patches en voert directionele scanning uit op een patch-geordende manier.

• Mechanisme: Gegeven een kenmerkkaart wordt het ruimtelijke rooster verdeeld in niet-overlappende patches van grootte $p \times p$. Binnen elke patch worden pixels in rij-majore volgorde genummerd voordat naar de volgende patch wordt overgegaan.
• Voordeel: Dit zorgt ervoor dat alle pixels binnen een patch opeenvolgend blijven in de sequentie, waardoor lokale ruimtelijke buurten behouden blijven, terwijl toch modellering van lange-afstand afhankelijkheden over patches heen mogelijk blijft.
• Hiërarchische strategie: Om multi-schaal context te vangen, hanteert het model een set patchgroottes $\{p^{(1)}, p^{(2)}, \dots\}$. Kleinere patches vangen fijnkorrelige grensdetails, terwijl grotere patches grovere structuren modelleren. Verschillende scanrichtingen (voorwaarts, achterwaarts, breedte-hoogte voorwaarts, breedte-hoogte achterwaarts) kunnen gebruikmaken van verschillende patchgroottes om anisotrope structuren te hanteren.

2. MoE-gebaseerde directionele fusie

Om vaste sommatie te vervangen, introduceren de auteurs een adaptieve fusiemodule die directionele outputs behandelt als "experts".

• Expertconstructie: De module maakt gebruik van vijf experts:
  • Vier richtingspecifieke experts ($E_1$ tot $E_4$): De genormaliseerde outputs van de vier directionele Mamba-scanners.
  • Een concatenatie-expert ($E_{concat}$): Gevormd door de vier genormaliseerde kaarten langs de kanaal-dimensie te concateneren en ze terug te projecteren via een $1\times1$ convolutie. Deze expert vangt cross-directionele correlaties en biedt een holistisch, richtingsonafhankelijk perspectief.
• Ruimtelijk-bewuste gating: Een router-netwerk berekent locatie-afhankelijke gewichten om deze experts adaptief te fuseren. De router wordt aangedreven door een mengsel van lokale context (via diepte-convolutie) en globale context (via adaptieve gemiddelde pooling), gecontroleerd door een leerbaar scalair $\alpha$.
• Residustabilisatie: Een residu-bypass voegt de ruwe directionele outputs terug toe aan het gefuseerde resultaat ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). Dit zorgt voor robuuste basisresponsen en behoudt sterke directionele signalen, wat de training stabiliseert.

3. Semantische en detailinfusie (SDI)

De encoder-outputs worden verwerkt door een SDI-module (aangepast van U-Net v2) om kenmerkkaarten te versterken door semantische informatie van hogere niveaus te infuseren en fijnere details van lagere niveaus te integreren via het Hadamard-product.

Belangrijkste bijdragen

1. Patch-geordende scanning: Een nieuw scansmechanisme dat pixel-gebaseerde raster-scanning vervangt door patch-geordende traversie, wat effectief de lokale 2D ruimtelijke coherentie behoudt terwijl de efficiëntie van state space-modellen behouden blijft.
2. MoE-gebaseerde directionele fusie: Een adaptieve fusiemodule die eenvoudige sommatie vervangt door een leerbaar gating-netwerk. Het weegt dynamisch meerdere directionele experts en een concatenatie-expert op basis van ruimtelijke context, waardoor het model zich kan aanpassen aan variërende laesieschalen en complexiteiten.
3. Hiërarchische multi-schaal modellering: De integratie van hiërarchische patchgroottes stelt het netwerk in staat om gelijktijdig fijnkorrelige grensdetails en grove semantische structuren te vangen.
4. Residustabilisatie: De opname van een residu-pad voor ruwe directionele outputs zorgt voor trainingsstabiliteit en voorkomt het verlies van kritieke directionele signalen.

Experimentele resultaten

De auteurs evalueerden Patch-MoE Mamba op vijf publieke polyp-segmentatiedatasets (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) en twee huidlaesie-segmentatiedatasets (ISIC 2017, ISIC 2018).

• Prestaties: Patch-MoE Mamba bereikte state-of-the-art prestaties op alle datasets, consistent presterend boven CNN-gebaseerde (U-Net, U-Net v2) en Mamba-gebaseerde baselines (VM-UNet, VM-UNetV2).
  • Op de uitdagende ETIS-dataset (bekend om laag contrast en kleine poliepen), verbeterde het model de Dice-score van 72,56% (VM-UNetV2) naar 74,04%.
  • Op ISIC 2018 bereikte het een Dice-score van 89,34%, waarmee het de vorige beste score van 88,36% overtrof.
• Ablatiestudies:
  • Het vervangen van raster-scanning door patch-geordende scanning alleen verbeterde de gemiddelde Dice van 83,46% naar 84,02%.
  • Het toevoegen van de MoE-fusiemodule verhief het gemiddelde verder naar 84,30%.
  • Hiërarchische patchgroottes (grof in vroege stadia, fijn in diepere stadia) leverden de beste resultaten op.
• Rekencomplexiteit: Hoewel de patch-geordende scanning op zich de rekenkosten niet verhoogt (het token-aantal blijft ongewijzigd), verhoogt de opname van de concatenatie-expert het aantal parameters tot ~70M en de FLOPs tot ~28G. De auteurs merken op dat deze overhead te verwachten is vanwege de dichtere kanaalinteracties die nodig zijn voor het MoE-mechanisme.

Betekenis en claims

Het artikel claimt dat Patch-MoE Mamba een algemene en effectieve oplossing biedt voor medische beeldsegmentatie. Door de specifieke beperkingen van bestaande Mamba-modellen aan te pakken—namelijk het verlies van lokale ruimtelijke structuur en de rigiditeit van directionele fusie—demonstreert de voorgestelde architectuur superieure aanpasbaarheid aan diverse medische doelen met variërende maten, vormen en grenscomplexiteiten. De resultaten suggereren dat het behoud van lokale ruimtelijke buurten via patch-geordende scanning, gecombineerd met adaptieve expert-fusie, de grenslocalisatie en segmentatie-nauwkeurigheid aanzienlijk verbetert, met name in uitdagende scenario's met laag-contrast of kleine laesies.