USEMA: a Scalable Efficient Mamba Like Attention for Medical Image Segmentation

Riepilogo Tecnico: USEMA – Un'attenzione scalabile ed efficiente simile a Mamba per la segmentazione di immagini mediche

Enunciato del Problema
La segmentazione accurata delle immagini mediche è fondamentale per le pipeline diagnostiche, ma affronta sfide uniche, tra cui dati di addestramento limitati e geometrie di oggetti complesse e sovrapposte. Sebbene i Vision Transformers (ViT) offrano la capacità di catturare dipendenze globali attraverso l'attenzione self-attention, la loro applicazione nell'imaging medico è ostacolata da due problemi principali:

1. Complessità Computazionale Quadratica: L'attenzione self-attention standard comporta un costo di $O(n^2)$ rispetto alla lunghezza della sequenza, rendendola computazionalmente proibitiva per immagini mediche ad alta risoluzione.
2. Dispersione dell'Attenzione: All'aumentare della lunghezza della sequenza, i punteggi nella matrice di attenzione (output softmax) tendono a disperdersi uniformemente verso lo zero. Questo fenomeno, teoricamente fondato sulla legge debole dei grandi numeri, ostacola la capacità del modello di selezionare chiavi e valori importanti, riducendo efficacemente il meccanismo di attenzione a un processo di media uniforme.

Metodologia
Gli autori propongono USEMA, un'architettura ibrida UNet che integra un nuovo meccanismo di attenzione chiamato Scalable Efficient Mamba-like Attention (SEMA). La metodologia affronta le limitazioni sia dei Transformer puri che dei modelli basati su Mamba esistenti attraverso i seguenti componenti:

• Meccanismo SEMA: SEMA combina approssimazioni locali e globali per imitare il comportamento della full self-attention mantenendo una complessità lineare.
  • Componente Locale: Utilizza window attention per mantenere la localizzazione e il focus dei token, evitando il problema della dispersione intrinseco nelle sequenze lunghe.
  • Componente Globale: Per catturare il contesto globale senza costi quadratici, SEMA impiega una media aritmetica teoricamente coerente. Questo termine approssima l'aspetto globale dell'attenzione, riflettendo i limiti $O(1/n)$ osservati nella disuguaglianza di dispersione. Il termine globale è definito come la media di tutti i vettori di valore, trasmessa (broadcasted) lungo la sequenza.
  • Integrazione: Il blocco SEMA è inserito in una struttura macro Mamba, utilizzando positional embedding condizionali, normalizzazione di livello e meccanismi di gating simili a Mamba (ad esempio, attivazione SiLU e convoluzioni depthwise).
• Architettura USEMA: Il modello adotta una struttura simmetrica a forma di U (UNet).
  • Encoder: È composto da blocchi residuali seguiti da blocchi SEMA. Le caratteristiche vengono ridisegnate da $(B, C, H, W)$ a $(B, C, HW)$ prima di entrare nel meccanismo di attenzione. Vengono utilizzati Rotary Positional Embedding e codifiche di posizione localmente potenziate per aumentare la consapevolezza spaziale.
  • Collo di Bottiglia: A causa della compressione spaziale, il collo di bottiglia utilizza full self-attention invece di attention a finestra, poiché la lunghezza della sequenza è sufficientemente piccola da essere gestibile dal punto di vista computazionale.
  • Decoder: Composto da blocchi residuali e convoluzioni trasposte, utilizza connessioni skip per fondere le caratteristiche gerarchiche dall'encoder.

Contributi Chiave

1. Design Ibrido dell'Attenzione: Il paper introduce SEMA, che giustifica teoricamente l'uso della media aritmetica come approssimazione globale per l'attenzione self-attention, complementando l'attenzione a finestra locale per prevenire la dispersione dei punteggi.
2. Integrazione Architettonica: USEMA fonde con successo l'efficienza nell'estrazione di caratteristiche locali dei CNN con le capacità di modellazione globale dell'attenzione simile a Mamba all'interno di un framework UNet.
3. Benchmarking Completo: Il modello è stato valutato su tre diverse modalità di imaging medico (Risonanza Magnetica Addominale, Endoscopia e Microscopia) con risoluzioni variabili e obiettivi di segmentazione (segmentazione di organi, strumenti e cellule).

Risultati Sperimentali
USEMA è stato confrontato con modelli basati su Transformer all'avanguardia (UNETR, Swin-UNETR, nnFormer) e modelli basati su/inspirati a Mamba (U-Mamba, Swin-UMamba, Mamba UNet, MLLA-UNet) su tre dataset:

• Risonanza Magnetica Addominale (MICCAI 2022 AMOS): USEMA ha raggiunto un Coefficiente di Similarità Dice (DSC) di 0.7704 e una Distanza Superficiale Normalizzata (NSD) di 0.8345. Questo ha superato la migliore baseline Transformer (nnFormer) del 5,83% in DSC e la migliore baseline Mamba (U-Mamba Enc) del 10,53% in DSC. Si noti che USEMA ha ottenuto questi risultati con 52M di parametri, ovvero il 13% in meno rispetto al modello Transformer migliore.
• Endoscopia (MICCAI 2017): USEMA ha raggiunto un DSC di 0.6463 e un NSD di 0.6621, superando le baseline Transformer di circa il 5,6% e le baseline Mamba di circa il 13,5% in DSC.
• Microscopia (NeurIPS 2022): In questo compito di segmentazione di istanze, USEMA ha raggiunto un punteggio F1 di 0.5791, superando la migliore baseline Transformer dell'8,61% e il MLLA-UNet ispirato a Mamba del 19,23%. Ha inoltre utilizzato il 43% di parametri in meno rispetto a U-Mamba Enc.

Studio di Ablazione
Uno studio di ablazione ha confermato la necessità dell'approssimazione dell'attenzione globale. La rimozione del termine di media aritmetica (utilizzando solo attention a finestra) ha portato a cali di prestazioni su tutti i dataset (ad esempio, un calo di 0,013 nel DSC sul dataset Risonanza Magnetica Addominale), validando il contributo del termine globale.

Significato e Affermazioni
Gli autori affermano che USEMA bilancia efficacemente l'efficienza computazionale e l'accuratezza della segmentazione. Sfruttando la proprietà di dispersione dell'attenzione self-attention standard e la complessità lineare di Mamba, USEMA fornisce una soluzione scalabile che:

• Supera i modelli puramente convoluzionali, basati su Transformer e basati su Mamba su diversi benchmark di imaging medico 2D.
• Riduce i costi computazionali rispetto ai Transformer con full self-attention.
• Dimostra che un'approssimazione globale efficiente e teoricamente fondata (media aritmetica) può complementare con successo i meccanismi di attenzione locale nell'analisi di immagini mediche.

Il paper conclude con i piani per estendere il meccanismo di media a una media appresa o ponderata e per applicare USEMA a immagini di patologia su larga scala, ma non afferma successi attuali in questi specifici domini futuri.