USEMA: a Scalable Efficient Mamba Like Attention for Medical Image Segmentation

Technische Samenvatting: USEMA – Een Schaalbare Efficiënte Mamba-achtige Aandacht voor Medische Beeldsegmentatie

Probleemstelling
Nauwkeurige medische beeldsegmentatie is cruciaal voor diagnostische pijplijnen, maar staat voor unieke uitdagingen, waaronder beperkte trainingsdata en complexe, overlappende objectgeometrieën. Hoewel Vision Transformers (ViTs) het vermogen bieden om globale afhankelijkheden via zelfaandacht te vangen, wordt hun toepassing in medische beeldvorming gehinderd door twee hoofdzaken:

1. Kwadratische Rekencomplexiteit: Gewone zelfaandacht brengt een $O(n^2)$-kost met zich mee ten opzichte van de sequentielengte, wat rekenkundig onuitvoerbaar maakt voor hoogresolutie medische beelden.
2. Aandachtsspreiding: Naarmate de sequentielengte toeneemt, hebben de scores in de aandachtsmatrix (softmax-uitvoer) de neiging om uniform naar nul te spreiden. Dit fenomeen, theoretisch onderbouwd door de zwakke wet van de grote getallen, belemmert het vermogen van het model om belangrijke sleutels en waarden te selecteren, waardoor het aandachtsmechanisme effectief wordt gereduceerd tot een uniform gemiddelde proces.

Methodologie
De auteurs stellen USEMA voor, een hybride UNet-architectuur die een nieuw aandachtsmechanisme integreert genaamd Scalable Efficient Mamba-like Attention (SEMA). De methodologie adresseert de beperkingen van zowel pure Transformers als bestaande Mamba-gebaseerde modellen via de volgende componenten:

• SEMA-mechanisme: SEMA combineert lokale en globale benaderingen om het gedrag van volledige zelfaandacht na te bootsen terwijl lineaire complexiteit behouden blijft.
  • Lokaal component: Het maakt gebruik van vensteraandacht om tokenlokalisatie en focus te behouden, waardoor het spreidingsprobleem dat inherent is aan lange sequenties wordt vermeden.
  • Globaal component: Om globale context te vangen zonder kwadratische kost, hanteert SEMA theoretisch consistente rekenkundige middeling. Deze term benadert het globale aspect van aandacht, wat de $O(1/n)$-grenzen weerspiegelt die worden waargenomen in de spreidingsongelijkheid. De globale term wordt gedefinieerd als het gemiddelde van alle waardenvectoren, uitgezonden over de sequentie.
  • Integratie: Het SEMA-blok wordt ingevoegd in een Mamba-macrostructuur, waarbij gebruik wordt gemaakt van conditionele positionele embedding, lagennormalisatie en sluitmechanismen vergelijkbaar met Mamba (zoals SiLU-activering en dieptewijze convoluties).
• USEMA-architectuur: Het model adopteert een symmetrische U-vormige structuur (UNet).
  • Encoder: Bestaat uit residublokken gevolgd door SEMA-blokken. Kenmerken worden herschikt van $(B, C, H, W)$ naar $(B, C, HW)$ voordat ze het aandachtsmechanisme binnengaan. Rotatoire positionele embedding en lokaal-versterkte positie-coderingen worden gebruikt om ruimtelijk bewustzijn te versterken.
  • Bottleneck: Vanwege ruimtelijke compressie maakt de bottleneck gebruik van volledige zelfaandacht in plaats van vensteraandacht, aangezien de sequentielengte voldoende klein is om rekenkundig hanteerbaar te zijn.
  • Decoder: Samengesteld uit residublokken en getransponeerde convoluties, waarbij skip-verbindingen worden gebruikt om hiërarchische kenmerken van de encoder te fuseren.

Belangrijkste Bijdragen

1. Hybride Aandachtsontwerp: Het paper introduceert SEMA, wat theoretisch het gebruik van rekenkundige middeling als globale benadering voor zelfaandacht rechtvaardigt, en dit aanvult met lokale vensteraandacht om score-spreiding te voorkomen.
2. Architecturale Integratie: USEMA slaagt erin de efficiëntie van lokale kenmerkextractie van CNN's te verenigen met de globale modelleringcapaciteiten van Mamba-achtige aandacht binnen een UNet-raamwerk.
3. Uitgebreide Benchmarking: Het model wordt geëvalueerd over drie diverse medische beeldvormingsmodaliteiten (Abdominale MRI, Endoscopie en Microscopie) met variërende resoluties en segmentatiedoelen (orgaan-, instrument- en celsegmentatie).

Experimentele Resultaten
USEMA werd gebenchmarkt tegen state-of-the-art Transformer-gebaseerde modellen (UNETR, Swin-UNETR, nnFormer) en Mamba-gebaseerde/geïnspireerde modellen (U-Mamba, Swin-UMamba, Mamba UNet, MLLA-UNet) op drie datasets:

• Abdominale MRI (MICCAI 2022 AMOS): USEMA behaalde een Dice Similarity Coefficient (DSC) van 0,7704 en een Normalized Surface Distance (NSD) van 0,8345. Dit presteerde beter dan de beste Transformer-baseline (nnFormer) met 5,83% in DSC en de beste Mamba-baseline (U-Mamba Enc) met 10,53% in DSC. Opmerkelijk is dat USEMA deze resultaten behaalde met 52M parameters, wat 13% minder is dan het best presterende Transformer-model.
• Endoscopie (MICCAI 2017): USEMA behaalde een DSC van 0,6463 en NSD van 0,6621, en presteerde beter dan Transformer-baselines met ongeveer 5,6% en Mamba-baselines met ongeveer 13,5% in DSC.
• Microscopie (NeurIPS 2022): In deze taak voor instantiesegmentatie behaalde USEMA een F1-score van 0,5791, en overtrof de beste Transformer-baseline met 8,61% en de Mamba-geïnspireerde MLLA-UNet met 19,23%. Het gebruikte ook 43% minder parameters dan U-Mamba Enc.

Ablatiestudie
Een ablatiestudie bevestigde de noodzaak van de globale aandachtsbenadering. Het verwijderen van de rekenkundige middelingsterm (alleen vensteraandacht gebruiken) resulteerde in prestatiedalingen over alle datasets (bijvoorbeeld een daling van 0,013 in DSC op de Abdominale MRI-dataset), wat de bijdrage van de globale term valideert.

Betekenis en Claims
De auteurs claimen dat USEMA effectief een balans vindt tussen rekenkundige efficiëntie en segmentatie-nauwkeurigheid. Door gebruik te maken van de spreidingseigenschap van gewone zelfaandacht en de lineaire complexiteit van Mamba, biedt USEMA een schaalbare oplossing die:

• Pure convolutie-, Transformer-gebaseerde en Mamba-gebaseerde modellen overtreft over diverse 2D-medische beeldvormingsbenchmarks.
• Rekenkosten verlaagt in vergelijking met volledige zelfaandacht-Transformers.
• Aantoont dat een theoretisch onderbouwde, efficiënte globale benadering (rekenkundige middeling) lokale aandachtsmechanismen in medische beeldanalyse succesvol kan aanvullen.

Het paper sluit af met plannen om het middelingmechanisme uit te breiden naar geleerde of gewogen middeling en om USEMA toe te passen op grootschalige pathologie-afbeeldingen, maar claimt geen huidige succes in deze specifieke toekomstige domeinen.