CS-MUNet: A Channel-Spatial Dual-Stream Mamba Network for Multi-Organ Segmentation

1. Il Problema

La segmentazione degli organi addominali, in particolare quelli cavi (come stomaco, intestino tenue e crasso), rappresenta una sfida significativa nell'analisi di immagini mediche (TC e Risonanza Magnetica). Le difficoltà principali includono:

• Morfologia variabile e confini indistinti: Gli organi cavi hanno pareti sottili e contrasti di tessuto molle bassi, rendendo difficile la definizione dei bordi.
• Limitazioni degli approcci esistenti basati su Mamba: Sebbene i metodi basati su Mamba (State Space Models - SSM) offrano un'efficienza computazionale lineare $O(n)$ per la modellazione a lungo raggio, le attuali implementazioni presentano tre carenze critiche:
  1. Mancata collaborazione semantica tra canali: Trattano i canali come componenti paralleli indipendenti, ignorando le dipendenze ordinate tra canali che codificano semantiche anatomiche correlate (texture, forma, densità).
  2. Integrazione superficiale della consapevolezza dei bordi: Le informazioni sui bordi vengono spesso trattate come vincoli post-processo o perdite ausiliarie, senza penetrare nei parametri fondamentali dello scanning SSM.
  3. Carenza di ottimizzazione specifica per organi cavi: I metodi esistenti sono validati principalmente su organi solidi, fallendo nell'affrontare le sfide specifiche degli organi cavi.

2. Metodologia: CS-MUNet

L'autori propongono CS-MUNet, una rete a doppio flusso (dual-stream) che modella simultaneamente le dimensioni spaziali e dei canali attraverso due moduli innovativi integrati in un'architettura U-Net basata su Res2Net-50.

A. Modulo Mamba Consapevole dei Confini (BASM - Boundary-Aware State Mamba)

Posizionato nelle connessioni di salto (skip connections) del decoder, il BASM integra la consapevolezza dei bordi direttamente nel meccanismo di transizione di stato dello SSM.

• Generazione di Mappe Posteriori dei Confini: Utilizza un framework di attenzione bayesiana per generare mappe posteriori a livello di pixel ($P_b$) combinando una prior geometrica (distanza dal bordo) e una likelihood di attenzione.
• Modulazione Parametrica Diretta: Invece di usare i bordi come segnale esterno, le mappe posteriori vengono iniettate direttamente nei parametri fondamentali di Mamba:
  • $\Delta$ (Tasso di dimenticanza): Modulato per mantenere gli stati storici più a lungo nelle regioni di confine.
  • $B$ (Peso di scrittura): Aumentato per amplificare l'importanza dell'input corrente nelle zone di confine.
• Fusione a Doppio Ramo: Combina l'output dello SSM (modellazione globale) con un ramo di compensazione strutturale locale (SASF) basato su convoluzioni depthwise multi-scala, fusi tramite un meccanismo di pesi stile SE (Squeeze-and-Excitation).

B. Modulo di Aggregazione dello Stato Mamba per Canale (CMSA - Channel Mamba State Aggregation)

Posizionato nel collo di bottiglia (bottleneck) dell'encoder, il CMSA ridefinisce la dimensione del canale come la dimensione della sequenza per lo SSM.

• Modellazione delle Dipendenze Inter-Canale: Tratta i canali come una sequenza temporale per catturare le dipendenze semantiche anatomiche tra di essi (es. correlazione tra texture e densità).
• Vincoli di Stato Limitati (Lipschitz): Poiché canali adiacenti possono rappresentare concetti anatomici disomogenei, l'SSM standard potrebbe causare un'accumulazione incontrollata di stati. Il CMSA introduce vincoli di clipping espliciti sulla matrice di transizione di stato ($\bar{A}_k$) e sul decadimento cumulativo, garantendo la convergenza teorica e prevenendo la contaminazione semantica tra gruppi di canali eterogenei.
• Aggregazione a Gruppi: I canali sono suddivisi in gruppi per bilanciare la capacità di modellazione delle dipendenze e l'efficienza computazionale.

3. Contributi Chiave

1. Nuovo Paradigma di Modellazione SSM: CS-MUNet è il primo approccio a integrare simultaneamente la collaborazione semantica tra canali e la fusione di caratteristiche consapevoli dei bordi direttamente nell'architettura SSM, superando i limiti dei metodi precedenti.
2. Modulazione Attiva dei Bordi (BASM): Trasforma la consapevolezza dei bordi da un vincolo passivo a un segnale attivo di guida per lo scanning, modificando dinamicamente i parametri di transizione di stato ($\Delta$ e $B$) in base alla probabilità del confine.
3. Aggregazione Semantica dei Canali (CMSA): Risolve il problema della modellazione indipendente dei canali ridefinendo l'asse dei canali come sequenza SSM, con vincoli matematici rigorosi per garantire la stabilità e l'indipendenza semantica.
4. Prestazioni Superiori: Dimostra un miglioramento costante rispetto agli stati dell'arte (SOTA) su dataset pubblici complessi.

4. Risultati Sperimentali

Il modello è stato valutato su due benchmark pubblici: UW-Madison GI Tract (MRI, 3 organi cavi) e WORD (TC, 16 organi).

• Prestazioni Quantitative:
  • UW-Madison: mDice di 86.16% (miglioramento di +2.91% rispetto al baseline più forte).
  • WORD: mDice di 94.47% (miglioramento di +1.03% rispetto al baseline più forte).
  • CS-MUNet supera metodi basati su CNN (U-Net, UNet++), Transformer (TransUNet, Swin-UNet) e altri approcci Mamba recenti.
• Efficienza: Nonostante le prestazioni superiori, il modello mantiene un numero di parametri contenuto (52.2M), significativamente inferiore a modelli basati su Transformer come TransUNet (117M) e Swin-UNet (41M, ma con prestazioni inferiori).
• Analisi di Ablazione:
  • L'aggiunta di BASM e CMSA singolarmente porta a guadagni significativi.
  • La rimozione della modulazione $\Delta/B$ nel BASM causa un calo drastico delle prestazioni, confermando che l'integrazione dei bordi nei parametri SSM è cruciale.
  • La sensibilità ai iperparametri è bassa, indicando la robustezza dell'architettura.

5. Significato e Impatto

CS-MUNet stabilisce un nuovo standard per la segmentazione multi-organo addominale.

• Risoluzione di Limiti Strutturali: Dimostra che la modellazione congiunta delle dipendenze spaziali (con bordi attivi) e delle dipendenze semantiche tra canali è essenziale per la precisione, specialmente per gli organi cavi.
• Efficienza Clinica: Offre un compromesso ottimale tra accuratezza e costo computazionale, rendendolo adatto per il deployment in ambienti clinici reali dove sia la precisione che l'efficienza sono critiche.
• Futuro: Sebbene l'approccio sia promettente, gli autori riconoscono la necessità di estendere il modello a 3D per migliorare la coerenza volumetrica e ottimizzare ulteriormente il peso del modello per dispositivi con risorse limitate.

In sintesi, CS-MUNet rappresenta un avanzamento fondamentale nell'applicazione degli State Space Models alla visione medica, passando da una semplice sostituzione di blocchi convoluzionali a una riprogettazione architetturale che sfrutta le proprietà uniche degli SSM per la semantica anatomica e la precisione dei bordi.