Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

Sintesi Tecnica: Ibrido Transformer-Mamba per la Segmentazione Volumetrica Medica Debolmente Supervisionata

Problematica
La segmentazione di immagini mediche volumetriche (ad es. CT, MRI) richiede tipicamente annotazioni a livello di voxel estremamente laboriose. La segmentazione debolmente supervisionata offre una soluzione, addestrando i modelli utilizzando solo etichette a livello di fetta (plane-level), che però mancano di informazioni sulla localizzazione. Sebbene la Class Activation Mapping (CAM) possa estrarre una localizzazione grossolana dai modelli di classificazione, gli approcci esistenti affrontano sfide significative a causa della natura volumetrica dei dati medici. La maggior parte degli approcci attuali si basa su encoder 2D che elaborano le singole fette indipendentemente, trascurando il contesto 3D. Inoltre, tentare di catturare il contesto volumetrico tramite la standard auto-attenzione (SA) cross-plane è computazionalmente proibitivo, poiché scala quadraticamente con il numero di piani ($N$).

Metodologia: TranSamba
Gli autori propongono TranSamba, un'architettura ibrida che combina i Vision Transformer (ViT) e l'architettura Mamba (Structured State Space Model) per affrontare tali limitazioni. La filosofia di design centrale è quella di sfruttare l'efficienza di Mamba per le interazioni cross-plane, mantenendo al contempo l'efficacia dell'auto-attenzione in-plane dei Transformer.

• Panoramica dell'Architettura: L'encoder consiste in $L$ strati ibridi. Ogni strato contiene un blocco Cross-Plane Mamba (CPM) seguito da un blocco Transformer.
• Modellazione Cross-Plane (CPM): Per catturare il contesto 3D in modo efficiente, i token dei patch da $N$ piani vicini vengono impilati e rimodellati in una singola sequenza di lunghezza $M \times N$ (dove $M$ è il numero di patch per piano). Il blocco CPM utilizza un modulo Mamba per modellare le interazioni tra questi token. Fondamentalmente, i class token vengono separati ed esclusi dal blocco Mamba per preservare la rappresentazione globale del piano. Ciò consente al modello di apprendere le dipendenze inter-piano in tempo lineare.
• Modellazione In-Plane (Transformer): L'output del blocco CPM viene rimodellato in $N$ sequenze parallele. Queste sono elaborate da standard blocchi Transformer utilizzando la Class-to-Patch (C2P) self-attention. Le rappresentazioni dei token arricchite dal passaggio cross-plane migliorano le mappe di attenzione in-plane, portando a una migliore localizzazione degli oggetti.
• Addestramento e Inferenza: Il modello è addestrato con etichette a livello di piano. La funzione di perdita combina termini di cross-entropy derivati da un global average pooling (GAP) del class token e un global weighted ranking pooling (GWRP) dei patch token. Durante l'inferenza, le mappe di attenzione vengono generate aggregando i pesi di attenzione C2P di tutti i blocchi Transformer.

Contributi Chiave

1. Architettura Ibrida: L'introduzione di TranSamba, che integra blocchi CPM con un backbone ViT per consentire una modellazione cross-plane efficiente senza sacrificare i punti di forza dell'auto-attenzione in-plane.
2. Efficienza della Complessità: L'architettura raggiunge una complessità temporale lineare rispetto al numero di piani ($N$) e mantiene una complessità di spazio costante per l'elaborazione dei batch, in contrasto con la scalabilità quadratica della standard auto-attenzione cross-plane.
3. Prestazioni State-of-the-Art: Il metodo dimostra che la modellazione cross-plane è una strategia generalizzabile per la segmentazione medica volumetrica, raggiungendo risultati superiori in diverse modalità e patologie senza fare affidamento su complessi aggregatori multi-scala o strategie di regolarizzazione auto-supervisionata spesso utilizzate nei lavori precedenti.

Risultati Sperimentali
TranSamba è stato valutato su tre dataset distinti:

• BraTS: Segmentazione del tumore cerebrale (T2 FLAIR).

• KiTS: Segmentazione del tumore renale.

• LASC: Segmentazione della cavità dell'atrio sinistro.

• Prestazioni: TranSamba ha raggiunto risultati State-of-the-Art (SOTA) su tutti e tre i dataset, superando i secondi migliori metodi con margini medi del 5,3% nel Dice Similarity Coefficient (DSC) e del 2,4% nell'Intersection-over-Union (IoU).

• Studi di Ablazione:

  • Cross-Plane vs. In-Plane: Sostituire il blocco CPM con blocchi Mamba in-plane (V2) ha prodotto prestazioni significativamente inferiori, confermando che il miglioramento deriva specificamente dalla modellazione cross-plane piuttosto che dal semplice uso di Mamba.
  • Ambito di Modellazione: La modellazione attraverso 16 piani ha fornito i migliori risultati, suggerendo che le interazioni cross-plane locali sono più benefiche rispetto alla modellazione a lungo raggio per le strutture continue.
  • Design dello Strato: Posizionare il blocco CPM prima del blocco Transformer (design CrossIn) si è rivelato ottimale, garantendo che l'estrazione dell'attenzione benefici del contesto cross-plane arricchito.

• Analisi della Dimensione dell'Oggetto: Le prestazioni hanno mostrato una correlazione positiva con la dimensione dell'oggetto. Mentre TranSamba eccelleva su oggetti grandi e enormi, ha incontrato difficoltà con oggetti piccoli (un problema noto nella segmentazione medica), in particolare nel dataset KiTS dove gli oggetti erano spesso piccoli.

Significato e Rivendicazioni
Il paper afferma che TranSamba colma efficacemente il divario tra la necessità di contesto 3D e i vincoli computazionali dell'apprendimento debolmente supervisionato. Decoupling (disaccoppiando) la modellazione cross-plane e in-plane, gli autori dimostrano che la modellazione cross-plane efficiente in tempo lineare è una componente critica e generalizzabile per migliorare la localizzazione degli oggetti nella segmentazione medica volumetrica. Il lavoro suggerisce che le future applicazioni cliniche potrebbero richiedere approcci di post-processing o multi-step, ma l'innovazione architettonica centrale fornisce una base robusta per compiti debolmente supervisionati in diversi scenari di imaging medico.