From Convolution to Transformer: A Comparative Study of U-Net Variants for Brain Tumor and Retinal Vessel Segmentation

Sintesi Tecnica: Dalla Convoluzione al Transformer

Problematica
La segmentazione di immagini mediche è fondamentale per la diagnosi assistita dal computer, la pianificazione del trattamento e il monitoraggio delle malattie. Sebbene l'architettura U-Net sia diventata lo standard per la segmentazione biomedica grazie alla sua struttura encoder-decoder e alle connessioni skip, le varianti convenzionali basate sulla convoluzione affrontano diverse limitazioni. Nello specifico, le classiche U-Net si affidano a operazioni convoluzionali locali, che ne limitano la capacità di catturare dipendenze a lungo raggio e informazioni contestuali globali. Questa limitazione è particolarmente deleteria in compiti complessi come la segmentazione dei tumori cerebrali, dove le regioni presentano forme irregolari, aspetti eterogenei e confini ambigui. Inoltre, le prestazioni delle varie modifiche della U-Net (residua, basata sull'attenzione o basata sul transformer) possono variare significativamente a seconda della modalità di imaging (MRI volumetrica rispetto a fundus 2D) e della struttura target (sottoregioni tumorali rispetto a sottili vasi retinici).

Metodologia
Questo studio conduce una valutazione comparativa di cinque architetture rappresentative basate su U-Net attraverso due distinti dataset di riferimento:

1. BraTS 2023: Un dataset volumetrico e multimodale di risonanza magnetica (MRI) per la segmentazione dei tumori cerebrali (che coinvolge le sottoregioni: tumore totale, tumore captante e nucleo tumorale).
2. DRIVE: Un dataset di immagini del fundus 2D per la segmentazione dei vasi retinici, caratterizzato da strutture sottili, ramificate e a basso contrasto.

I modelli valutati rappresentano uno spettro di filosofie di progettazione:

• 3D U-Net: Estende l'architettura 2D originale ai dati volumetrici utilizzando convoluzioni 3D per catturare la continuità anatomica tra le fette.
• Residual U-Net: Incorpora blocchi residui per migliorare la propagazione del gradiente e la stabilità dell'apprendimento delle caratteristiche, aiutando l'estrazione di strutture complesse ed eterogenee.
• Attention U-Net: Introduce gate di attenzione nelle connessioni skip per sopprimere le informazioni di sfondo irrilevanti ed enfatizzare le regioni rilevanti per il compito.
• UNETR: Sostituisce l'encoder convoluzionale con un Vision Transformer (ViT) per modellare le dipendenze globali tramite meccanismi di auto-attenzione (self-attention) su patch di immagini.
• Swin UNETR: Integra blocchi di Swin Transformer nel framework U-Net, utilizzando l'apprendimento della rappresentazione gerarchica e la self-attention a finestra spostata (shifted window self-attention) per catturare efficientemente sia i dettagli locali che il contesto globale.

Tutti i modelli sono stati addestrati con un preprocessing coerente (normalizzazione, padding/cropping) e strategie di data augmentation. Le prestazioni sono state valutate utilizzando metriche quantitative tra cui i punteggi Dice, Precision, Recall, punteggi F1 e la convergenza della perdita di addestramento (training loss).

Contributi Chiave
L'articolo fornisce tre contributi primari:

1. Valutazione Comparativa: Un confronto sistematico di cinque varianti di U-Net (convoluzionali, residue, basate sull'attenzione e basate sul transformer) attraverso compiti di imaging cerebrale MRI volumetrico e di imaging retinico 2D.
2. Analisi dei Meccanismi: Un'analisi di come specifici meccanismi architettonici — l'apprendimento residuo, i gate di attenzione e la self-attention dei transformer — impattino le prestazioni di segmentazione in diversi scenari di imaging medico.
3. Guida Pratica: Approfondimenti pratici nella selezione del modello, distinguendo tra architetture più adatte alla segmentazione volumetrica complessa rispetto a quelle efficaci per l'estrazione di strutture fini.

Risultati Sperimentali

• BraTS 2023 (Tumore Cerebrale): Swin UNETR ha ottenuto le migliori prestazioni complessive con un punteggio Dice medio di 0,8965 (WT: 0,8966, ET: 0,9088, TC: 0,8841). UNETR è seguita con un Dice medio di 0,8540. La Residual U-Net (0,7341), la 3D U-Net (0,6890) e la Attention U-Net (0,6522) hanno ottenuto prestazioni inferiori. I risultati indicano che la modellazione del contesto globale tramite transformer è altamente benefica per la segmentazione volumetrica dei tumori.
• DRIVE (Vasi Retinici): Swin UNETR ha ottenuto nuovamente il punteggio Dice più alto (0,8078), seguita da vicino da Residual U-Net (0,8069) e UNETR (0,8015). L'Attention U-Net ha ottenuto 0,7855, mentre la 3D U-Net ha raggiunto il Dice più basso (0,5967) nonostante avesse la perdita di addestramento più bassa (0,1030). Questa discrepanza evidenzia che una minore perdita di addestramento non garantisce una migliore qualità di segmentazione, specialmente quando l'architettura è disallineata alla struttura dei dati (ad esempio, modelli 3D su strutture 2D fini).

Significatività e Rivendicazioni
L'articolo sostiene che Swin UNETR sia l'architettura più efficace in termini assoluti, riuscendo a bilanciare la preservazione del dettaglio spaziale locale con la modellazione del contesto a lungo raggio attraverso il suo encoder gerarchico e la self-attention a finestra spostata. Lo studio suggerisce che:

• Le varianti basate su Transformer (specificamente Swin UNETR) sono superiori per compiti di segmentazione complessi che richiedono una comprensione del contesto globale, come la delineazione del tumore cerebrale.
• L'apprendimento residuo rimane altamente efficace per compiti che richiedono la preservazione di dettagli strutturali fini, come la segmentazione dei vasi retinici, dove la Residual U-Net si è dimostrata competitiva con i modelli transformer.
• La selezione del modello deve dipendere dal compito: approcci basati su transformer sono raccomandati per compiti volumetrici e complessi, mentre approcci residui o ibridi possono essere sufficienti per l'estrazione di strutture fini.

Gli autori concludono che il lavoro futuro dovrebbe esplorare architetture ibride che combinino la modellazione contestuale basata su transformer con moduli convoluzionali o residui per migliorare ulteriormente la robustezza e la generalizzazione attraverso diversi dataset clinici.