DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

Riepilogo Tecnico: DAMamba-UNet3D

Problematica
La segmentazione di immagini mediche 3D richiede un equilibrio tra la precisa delineazione dei bordi locali e un contesto spaziale globale coerente, il tutto rispettando rigidi budget computazionali e di parametri. Le architetture esistenti affrontano scambi distinti:

• U-Net Convoluzionali: Offrono efficienza nei parametri e un mescolamento locale $O(n)$, ma mancano di un contesto globale esplicito, richiedendo l'aumento della profondità per catturare dipendenze a lungo raggio.
• Vision Transformers: Forniscono un ragionamento globale tramite self-attention, ma incorrono in costi di memoria e calcolo proibitivi $O(n^2)$ per grandi volumi 3D (ad esempio, un patch di $128^3$ contiene circa 2,1 milioni di token).
• Modelli a Spazio di Stato (SSM/Mamba): Propagano l'informazione globalmente in tempo $O(n)$ e memoria $O(n)$. Tuttavia, gli esistenti segmentatori SSM medici (ad es. SegMamba) si affidano a pattern di scansione fissi e ampi budget di parametri ($\sim 70$ milioni di parametri), mancando di adattabilità dipendente dai dati.

Sebbene la Scansione Adattiva Dinamica (DAS) sia stata introdotta per il riconoscimento di immagini naturali 2D per apprendere un riordinamento dipendente dai dati prima della scansione selettiva, non è stata applicata all'imaging medico né estesa a volumi 3D.

Metodologia
Gli autori propongono DAMamba-UNet3D, un'architettura ibrida encoder-decoder che integra blocchi Tri-Plane 3D-DAS in un framework U-Net.

1. Design dell'Architettura:

   • Struttura Ibrida: Il modello mantiene blocchi convoluzionali per lo stem (E1), il bottleneck e l'intero decoder per preservare i dettagli locali.
   • Integrazione Encoder: I blocchi Tri-plane DAMamba sono inseriti esclusivamente nelle fasi dell'encoder E2–E4.
   • Variante Compatta: Il modello primario utilizza una larghezza di base $C=16$, risultando in circa 5,3 milioni di parametri ($\sim 13\times$ meno di SegMamba).
   • Variante su Larga Scala (DAMamba-L): Una variante più ampia ($\sim 70$ milioni di parametri) con blocchi DAMamba impilati nell'encoder e un bottleneck convoluzionale, progettata per eguagliare la capacità di SegMamba per un confronto diretto.

2. Scansione Adattiva Dinamica Tri-Piano (Tri-Plane DAS):

   • Meccanismo: A differenza delle scansioni fisse (row-major, zigzag o flattening multi-orientamento), la DAS apprende un riordinamento spaziale dipendente dall'input. Una Rete di Predizione degli Offset (OPN) predice gli offset 2D dalla mappa delle caratteristiche, che vengono utilizzati per ricampionare le caratteristiche tramite interpolazione bilineare prima della scansione selettiva.
   • Estensione 3D: L'operatore DAS 2D viene esteso ai volumi applicandolo attraverso tre piani anatomici: Assiale ($H \times W$), Coronale ($D \times W$) e Sagittale ($D \times H$). Gli output vengono mediati per formare il blocco 3D DAS:
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • Scansione Selettiva: La mappa delle caratteristiche riordinata è elaborata da un blocco di scansione selettiva Mamba standard, consentendo la propagazione del contesto globale con costo $O(N)$ senza mappe di attenzione quadratiche.

3. Addestramento e Valutazione:

   • Dataset: Set di addestramento BraTS 2020 (369 soggetti, 4 modalità).
   • Loss: Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$) per concentrarsi sui voxel difficili e sui bordi.
   • Protocollo: Validazione incrociata a cinque pieghe (five-fold) a livello di soggetto con inferenza a volume completo (sliding-window, patch $128^3$). Le metriche includono la media di Dice e HD95 per l'Enhanced Tumor (ET), Tumor Core (TC) e Whole Tumor (WT).

Contributi Chiave

1. Prima Applicazione Medica di DAS: Trasferisce la Scansione Adattiva Dinamica dalla classificazione di immagini naturali 2D alla segmentazione medica densa 3D.
2. Estensione da 2D a 3D: Eleva l'operatore DAS ai dati volumetrici tramite scansione adattiva tri-piano, offrendo un'alternativa apprendibile al flattening multi-orientamento progettato a mano.
3. Efficienza dei Parametri: Dimostra che un'U-Net ibrida con blocchi tri-plane DAS solo nell'encoder raggiunge una capacità quasi paritaria a modelli molto più grandi con $\sim 13\times$ meno parametri.
4. Validazione a Due Scale: Valida l'approccio sia su scala compatta ($\sim 5,3$M) che su larga scala ($\sim 70$M), mostrando che la DAS appresa può superare i pattern di scansione fissi quando il posizionamento architettonico è ottimizzato.

Risultati

• Modello Compatto (DAMamba-UNet3D, $\sim 5,3$M): Ha ottenuto un Dice medio di 0,815 $\pm$ 0,013 su BraTS 2020. Questo è quasi paritario con il SegMamba riaddestrato (0,824 $\pm$ 0,014) nonostante abbia $\sim 13\times$ meno parametri.
• Modello su Larga Scala (DAMamba-L, $\sim 70$M): Ha superato il SegMamba riaddestrato con un Dice medio di 0,829 $\pm$ 0,012 (un miglioramento di +0,5 punti), con il maggior guadagno osservato nella classe Tumor Core (TC) (0,842 vs 0,827).
• Studi di Ablazione:
  • Posizionamento: Posizionare i blocchi DAS solo nell'encoder (E2–E4) ha prodotto i risultati migliori (0,815). Aggiungere la DAS al bottleneck o al decoder ha abbassato le prestazioni (0,807 e 0,805, rispettivamente), indicando che il posizionamento solo nell'encoder è critico.
  • Pattern di Scansione: La DAS tri-piano appresa ha superato gli ordini di scansione fissi (Raster, Snake, Cross, Tri-orientated) quando il posizionamento dell'encoder era corrispondente.

Significato e Rivendicazioni
Il paper sostiene che la DAS tri-piano appresa in una U-Net ibrida è un'alternativa competitiva alla scansione fissa Tri-orientated Mamba (ToM) per la segmentazione del tumore cerebrale 3D.

• Efficienza: Il modello compatto dimostra che è possibile ottenere prestazioni elevate con un numero significativamente ridotto di parametri sfruttando la scansione adattiva piuttosto che la capacità bruta.
• Prestazioni: Alle grandi scale, il meccanismo adattivo appreso (DAS) può migliorare rispetto alle traiettorie di scansione fisse (ToM) sotto identici protocolli di addestramento e valutazione.
• Modestia: Gli autori avvertono che, sebbene i risultati su BraTS 2020 siano promettenti, lo studio è limitato a un singolo dataset. Affermano che "sono necessari la validazione su più dataset e i test statistici prima di poter fare affermazioni generali su DAS rispetto a ToM". Il contributo primario è l'estensione della DAS ai volumi medici 3D e la dimostrazione della sua fattibilità in un'architettura ibrida efficiente nei parametri.