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

Riepilogo Tecnico: MHMamba per la Segmentazione di Tumori Cerebrali 3D

Enunciato del Problema
La segmentazione dei tumori cerebrali nelle risonanze magnetiche (MRI) 3D affronta sfide significative a causa dell'alta eterogeneità della morfologia tumorale e del contrasto multimodale. Mentre le Reti Neurali Convoluzionali (CNN) eccellono nella modellazione di texture e confini locali, faticano con le dipendenze a lungo raggio a causa di campi ricettivi limitati. Al contrario, i modelli basati su Transformer possono catturare il contesto globale ma soffrono di complessità computazionale e di memoria quadratica ($O(N^2)$), rendendoli difficili da applicare a volumi 3D ad alta risoluzione senza un forte downsampling o finestramento, che possono portare a incoerenza del contesto e incongruenze ai confini. Inoltre, gli approcci ibridi esistenti soffrono spesso di disallineamenti statistici tra i rami convoluzionali e di attenzione (ad esempio, Normalizzazione Batch/Instance vs. Normalizzazione Livello), portando a un addestramento instabile e a problemi di fusione delle caratteristiche.

Metodologia
Il documento propone MHMamba (Multi-Head Mamba), un'architettura encoder-decoder a forma di U progettata per modellare le caratteristiche globali con complessità lineare preservando al contempo i dettagli ad alta risoluzione. Il framework è composto da tre componenti principali:

1. Modulo Multi-Head Mamba (MHM):

   • Sostituisce i blocchi standard nell'encoder con un modello a spazio di stati (SSM) multi-testa.
   • La dimensione del canale viene suddivisa in più teste SSM parallele ($N_h$), permettendo a diverse teste di apprendere pattern di dipendenza a lungo raggio complementari da input multimodali senza assegnazioni fisse di modalità.
   • Ogni testa serializza le caratteristiche spaziali 3D in una sequenza, applica aggiornamenti di stato selettivi (Mamba) e fonde gli output tramite concatenazione e proiezione.
   • Questo design mantiene una complessità lineare potenziando al contempo le capacità di modellazione globale e la stabilità dell'addestramento in condizioni multimodali.

2. Modulo di Attenzione di Calibrazione Spazio-Canale (CSCA):

   • Progettato per allineare le statistiche delle caratteristiche e potenziare le risposte alle lesioni dopo l'elaborazione Multi-Head Mamba.
   • Impiega un meccanismo a doppio percorso:
     • Percorso Canale: Utilizza pooling medio e massimo globale seguito da un MLP per generare pesi di canale.
     • Percorso Spaziale: Integra in modo innovativo quattro statistiche spaziali (media, deviazione standard, massimo e minimo) per generare una mappa di pesi spaziali, catturando l'eterogeneità e i picchi salienti.
   • Un meccanismo di gating fonde adattivamente questi rami, seguito da una connessione residua per preservare le risposte originali delle caratteristiche.

3. Modulo di Fusione Gated Adattiva (AGF):

   • Posizionato alle connessioni di salto tra encoder e decoder.
   • Affronta il conflitto tra semantica di alto livello e dettagli di basso livello fondendo dinamicamente le caratteristiche multiscala.
   • Le caratteristiche di input vengono suddivise in sottogruppi lungo la dimensione del canale. Per ogni sottogruppo, il modulo impara pesi di fusione indipendenti per bilanciare le caratteristiche dell'encoder e del decoder, prevenendo conflitti di caratteristiche che spesso portano a previsioni di confini incoerenti.

Contributi Chiave

• Architettura Multi-Head Mamba: Il documento introduce un approccio SSM multi-testa per le MRI 3D che scompone i canali in teste parallele. Questo potenzia la modellazione a lungo raggio e l'adattabilità multimodale, mantenendo al contempo il vantaggio della complessità lineare di Mamba.
• Calibrazione e Fusione delle Caratteristiche: L'integrazione del modulo CSCA per unificare le scale e le statistiche delle caratteristiche, e del modulo AGF per fondere dinamicamente le caratteristiche multiscala alle connessioni di salto, migliora significativamente la coerenza dei confini e il rilevamento di lesioni di piccolo volume.
• Validazione Sistematica: Il metodo è stato valutato rigorosamente sui dataset BraTS2021 e BraTS2023, dimostrando miglioramenti stabili rispetto a baseline solide (inclusi nnU-Net, SegMamba e modelli basati su Transformer) in termini di punteggi Dice e Distanza di Hausdorff (HD95).

Risultati Sperimentali
Gli esperimenti condotti sui dataset BraTS2021 e BraTS2023 mostrano che MHMamba raggiunge prestazioni allo stato dell'arte o vicine allo stato dell'arte:

• BraTS2021: Ha raggiunto un punteggio Dice medio del 91,02% e un HD95 medio di 3,38, superando baseline come SegMamba (90,32% Dice, 4,18 HD95) e nnU-Net (89,40% Dice, 5,58 HD95).
• BraTS2023: Ha raggiunto un punteggio Dice medio del 90,23% e un HD95 medio di 3,67.
• Studi di Ablazione:
  • L'aggiunta individuale di CSCA, AGF e del modulo Multi-Head Mamba ha contribuito a guadagni di prestazioni in tutte le sotto-regioni tumorali (Tumore Intero, Nucleo del Tumore e Tumore Enhancing).
  • La configurazione ottimale è risultata essere 4 teste, che ha bilanciato la capacità rappresentativa e l'efficienza computazionale; l'aumento a 8 teste ha portato a un degrado delle prestazioni a causa di una eccessiva partizione e rumore di ottimizzazione nell'addestramento 3D con dimensioni di batch ridotte.
  • I risultati qualitativi indicano che MHMamba produce previsioni meno frammentate con confini più stretti, in particolare nelle regioni tumorali eterogenee e nelle piccole aree enhancing.

Significato e Affermazioni
Il documento afferma che MHMamba affronta efficacemente il compromesso tra la modellazione del contesto globale e la preservazione dei dettagli locali nella segmentazione dei tumori cerebrali 3D. Sfruttando la complessità lineare dei Modelli a Spazio di Stati attraverso un'architettura multi-testa, il metodo evita l'overhead di memoria quadratica dei Transformer superando al contempo le limitazioni del campo ricettivo locale delle CNN. I moduli proposti (CSCA e AGF) mirano specificamente ai disallineamenti statistici e all'incoerenza del contesto spesso riscontrati nella segmentazione 3D multimodale, risultando in un framework robusto, efficiente e accurato adatto per applicazioni cliniche come la pianificazione della radioterapia e la valutazione preoperatoria. Gli autori sottolineano che il metodo mantiene i vantaggi di efficienza della modellazione basata su Mamba migliorando significativamente la sensibilità ai nuclei tumorali e alle lesioni di piccolo volume.