Riepilogo Tecnico: Patch-MoE Mamba

Enunciato del Problema

La segmentazione di immagini mediche si basa fortemente sulle Reti Neurali Convoluzionali (CNN) e sui Transformer. Tuttavia, le CNN sono intrinsecamente limitate dai campi ricettivi locali, faticando a modellare le dipendenze a lungo raggio, mentre i Transformer spesso comportano una complessità computazionale e di memoria quadratica. Recentemente, i Modelli di Spazio di Stato (SSM), in particolare le architetture basate su Mamba, sono emersi come alternative efficienti con complessità sequenziale lineare.

Nonostante le loro promesse, i modelli di segmentazione esistenti basati su Mamba affrontano due limitazioni critiche:

1. Interruzione della Struttura Spaziale Locale: La maggior parte dei modelli converte le mappe di caratteristiche 2D in sequenze 1D utilizzando una scansione raster pixel per pixel. Questo processo può separare i pixel adiacenti spazialmente nella sequenza, indebolendo la coerenza spaziale 2D locale. Ciò è dannoso per i compiti di previsione densa che richiedono una localizzazione accurata dei confini e la preservazione della struttura locale.
2. Fusione Direzionale Inflessibile: I metodi attuali fondono tipicamente gli output da più direzioni di scansione tramite una semplice somma. Questa aggregazione fissa assume che tutte le direzioni siano ugualmente importanti in ogni posizione spaziale, fallendo nell'adattarsi alle diverse dimensioni, forme, texture e complessità dei confini dei target medici (ad esempio, lesioni piccole e a basso contrasto rispetto a strutture grandi e irregolari).

Metodologia

Gli autori propongono Patch-MoE Mamba, un'architettura di spazio di stato Mixture-of-Experts (MoE) ordinata per patch, progettata per affrontare le limitazioni sopra menzionate. L'architettura segue un design in stile U-Net, sostituendo i blocchi standard Visual State Space (VSS) in VM-UNetV2 con un nuovo blocco Patch-MoE VSS.

1. Scansione Gerarchica Ordinata per Patch

Invece di scansionare i singoli pixel tramite raster, il metodo suddivide le mappe di caratteristiche in patch locali ed esegue la scansione direzionale in modo ordinato per patch.

• Meccanismo: Data una mappa di caratteristiche, la griglia spaziale è divisa in patch non sovrapposte di dimensione $p \times p$. All'interno di ogni patch, i pixel sono enumerati in ordine riga principale prima di passare alla patch successiva.
• Vantaggio: Questo garantisce che tutti i pixel all'interno di una patch rimangano consecutivi nella sequenza, preservando i vicini spaziali locali pur consentendo la modellazione delle dipendenze a lungo raggio tra le patch.
• Strategia Gerarchica: Per catturare il contesto multi-scala, il modello impiega un insieme di dimensioni di patch $\{p^{(1)}, p^{(2)}, \dots\}$. Patch più piccole catturano dettagli fini dei confini, mentre patch più grandi modellano strutture più grossolane. Diverse direzioni di scansione (avanti, inversa, avanti larghezza-altezza, inversa larghezza-altezza) possono utilizzare dimensioni di patch distinte per gestire strutture anisotrope.

2. Fusione Direzionale Basata su MoE

Per sostituire la somma fissa, gli autori introducono un modulo di fusione adattivo che tratta gli output direzionali come "esperti".

• Costruzione degli Esperti: Il modulo utilizza cinque esperti:
  • Quattro esperti specifici per direzione ($E_1$ fino a $E_4$): Gli output normalizzati dei quattro scanner Mamba direzionali.
  • Un esperto di concatenazione ($E_{concat}$): Formato concatenando le quattro mappe normalizzate lungo la dimensione del canale e proiettandole nuovamente tramite una convoluzione $1\times1$. Questo esperto cattura le correlazioni tra direzioni e fornisce una visione olistica e agnostica rispetto alla direzione.
• Gating Consapevole dello Spazio: Una rete router calcola pesi dipendenti dalla posizione per fondere adattivamente questi esperti. Il router è guidato da una miscela di contesto locale (tramite convoluzione depth-wise) e contesto globale (tramite pooling medio adattivo), controllato da uno scalare apprendibile $\alpha$.
• Stabilizzazione Residua: Un bypass residuo aggiunge gli output direzionali grezzi al risultato fuso ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). Questo garantisce risposte di base robuste e preserva i segnali direzionali forti, stabilizzando l'addestramento.

3. Infusione di Semantica e Dettagli (SDI)

Gli output del codificatore sono elaborati da un modulo SDI (adattato da U-Net v2) per migliorare le mappe di caratteristiche infondendo informazioni semantiche dai livelli superiori e integrando dettagli più fini dai livelli inferiori tramite il prodotto di Hadamard.

Contributi Chiave

1. Scansione Ordinata per Patch: Un meccanismo di scansione innovativo che sostituisce la scansione raster pixel per pixel con una traversata ordinata per patch, preservando efficacemente la coerenza spaziale 2D locale mantenendo al contempo l'efficienza dei modelli di spazio di stato.
2. Fusione Direzionale Basata su MoE: Un modulo di fusione adattivo che sostituisce la semplice somma con una rete di gating apprendibile. Pesa dinamicamente molteplici esperti direzionali e un esperto di concatenazione in base al contesto spaziale, consentendo al modello di adattarsi a diverse scale e complessità delle lesioni.
3. Modellazione Gerarchica Multi-Scala: L'integrazione di dimensioni di patch gerarchiche permette alla rete di catturare simultaneamente dettagli fini dei confini e strutture semantiche grossolane.
4. Stabilizzazione Residua: L'inclusione di un percorso residuo per gli output direzionali grezzi garantisce la stabilità dell'addestramento e previene la perdita di segnali direzionali critici.

Risultati Sperimentali

Gli autori hanno valutato Patch-MoE Mamba su cinque dataset pubblici di segmentazione di polipi (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) e due dataset di segmentazione di lesioni cutanee (ISIC 2017, ISIC 2018).

• Prestazioni: Patch-MoE Mamba ha raggiunto prestazioni all'avanguardia su tutti i dataset, superando costantemente le basi basate su CNN (U-Net, U-Net v2) e quelle basate su Mamba (VM-UNet, VM-UNetV2).
  • Sul dataset impegnativo ETIS (noto per il basso contrasto e i polipi piccoli), il modello ha migliorato il punteggio Dice dal 72,56% (VM-UNetV2) al 74,04%.
  • Su ISIC 2018, ha raggiunto un punteggio Dice del 89,34%, superando il precedente miglior risultato del 88,36%.
• Studi di Ablazione:
  • Sostituire la scansione raster con la scansione ordinata per patch da sola ha migliorato il Dice medio dall'83,46% all'84,02%.
  • Aggiungere il modulo di fusione MoE ha ulteriormente portato la media all'84,30%.
  • Le dimensioni gerarchiche delle patch (grossolane nelle fasi iniziali, fini nelle fasi più profonde) hanno prodotto i migliori risultati.
• Complessità Computazionale: Sebbene la scansione ordinata per patch di per sé non aumenti il costo computazionale (il numero di token rimane invariato), l'inclusione dell'esperto di concatenazione aumenta i parametri a circa 70M e i FLOPs a circa 28G. Gli autori notano che questo sovraccarico è atteso a causa delle interazioni di canale più dense richieste dal meccanismo MoE.

Significato e Affermazioni

Il paper afferma che Patch-MoE Mamba offre una soluzione generale ed efficace per la segmentazione di immagini mediche. Affrontando le limitazioni specifiche dei modelli Mamba esistenti, ovvero la perdita della struttura spaziale locale e la rigidità della fusione direzionale, l'architettura proposta dimostra una superiore adattabilità a diversi target medici con dimensioni, forme e complessità dei confini variabili. I risultati suggeriscono che preservare i vicini spaziali locali attraverso la scansione ordinata per patch, combinata con una fusione adattiva degli esperti, migliora significativamente la localizzazione dei confini e l'accuratezza della segmentazione, in particolare in scenari impegnativi che coinvolgono lesioni a basso contrasto o di piccole dimensioni.