Does Your ViT Still Need U-Net for Segmentation?

Sintesi Tecnica: Il tuo ViT ha ancora bisogno di U-Net per la segmentazione?

Problematica
La segmentazione di immagini mediche è stata dominata per oltre un decennio da architetture encoder-decoder in stile U-Net. Sebbene questi design affrontino con successo la perdita di dettagli spaziali attraverso connessioni skip e una decodifica simmetrica, essi si affidano a operazioni convoluzionali che sono vincolate da campi recettivi limitati, rendendo difficile la modellazione del contesto globale. Sebbene i Vision Transformer (ViT) offrano meccanismi di auto-attenzione per catturare dipendenze a lungo raggio, i primi metodi di segmentazione basati su ViT (ad es., TransUNet, SwinUNet) hanno mantenuto decoder di tipo U-Net perché le rappresentazioni pre-addestrate dei ViT erano ritenute insufficienti per una densa predizione accurata.

Recenti progressi nel pre-addestramento su larga scala (ad es., DINOv2, SigLIP) hanno ridefinito significativamente le capacità di rappresentazione dei ViT, consentendo loro di apprendere potenti caratteristiche visive da massicci dataset non etichettati. Questa evoluzione pone una domanda critica: il paradigma U-Net è ancora necessario per la segmentazione di immagini mediche quando sono disponibili moderni backbone ViT fortemente pre-addestrati? Se non lo è, come dovrebbe essere progettato un framework encoder-only per predire efficacemente le maschere di segmentazione senza un decoder?

Metodologia: EoSeg
Gli autori propongono EoSeg (Encoder-only Segmentation), un framework che elimina il decoder U-Net a favore di un design basato su query costruito su moderni backbone ViT pre-addestrati. La metodologia è sviluppata attraverso un'esplorazione sistematica delle scelte architettoniche chiave:

1. Selezione del Backbone: Lo studio valuta vari backbone ViT pre-addestrati (inclusi ViT-L/14, SigLIP, MTP, DINOv2 e DINOv3) in un setting encoder-only. DINOv2 è selezionato come il backbone ottimale, raggiungendo la prestazione di baseline più elevata (81,86% mDice su Synapse) senza alcun decoder.
2. Generazione della Maschera:
   • Deconvoluzione vs. Basata su Query: Inizialmente è stato testato un modulo deconvoluzionale per l'upsampling delle caratteristiche a bassa risoluzione del ViT. Tuttavia, gli autori hanno scoperto che un approccio di predizione della maschera basato su query era superiore. Invece di una decodifica pixel-wise, gli embedding di query apprendibili agiscono come sonde semantiche globali.
   • Meccanismo: I vettori di query vengono proiettati in uno spazio di embedding della maschera e interagiscono con la mappa delle caratteristiche globali upsampled tramite un prodotto scalare channel-wise per generare direttamente maschere binarie.
3. Modellazione Multi-livello delle Query: Ispirandosi alle connessioni skip delle U-Net, il framework sfrutta le rappresentazioni degli ultimi tre blocchi Transformer anziché solo del blocco finale. Le query apprendibili vengono iniettate prima di questi blocchi e perfezionate attraverso l'auto-attenzione con i token dell'immagine, permettendo al modello di aggregare informazioni semantiche da molteplici stadi.
4. Fusione di Blocchi Apprendibile: Per combinare le predizioni dei tre blocchi intermedi, viene introdotto un modulo di fusione apprendibile. Una porta di fusione (fusion gate) utilizza la rappresentazione della query per predire pesi adattivi per l'output di ciascun blocco, consentendo al modello di determinare dinamicamente il contributo di ogni livello alla predizione finale.
5. Configurazione delle Query: Il framework adotta un paradigma di classificazione basato su query in cui ogni query predice una maschera e un'etichetta di classe. Gli esperimenti hanno determinato che assegnare due query per categoria produce prestazioni ottimali, bilanciando la necessità di catturare le variazioni intra-classe senza sovra-parametrizzazione.
6. Strategia di Supervisione: Contrariamente alle pratiche comuni nella segmentazione medica, è stato scoperto che la supervisione profonda (applicare la loss agli output intermedi) non migliora le prestazioni in questo specifico setup encoder-only ed è stata esclusa dal design finale.

Contributi Chiave

• Rivalutazione Architetturale: Il paper fornisce una valutazione estesa dimostrando che un decoder in stile U-Net non è più un prerequisito per una segmentazione di alta qualità delle immagini mediche quando si sfruttano moderni backbone ViT pre-addestrati su larga scala.
• Framework EoSeg: Gli autori presentano EoSeg, un framework encoder-only basato su query che integra la modellazione multi-livello delle query e la fusione di blocchi apprendibile per consentire la predizione diretta della maschera dalle caratteristiche dell'encoder.
• Validazione Completa: Il design proposto è validato attraverso sette dataset benchmark che coprono cinque distinte modalità di imaging medico: CT (Synapse), MRI (ACDC), Istopatologia (GlaS, MoNuSeg), Endoscopia (Kvasir-Seg) e Dermoscopia (ISIC-2016, ISIC-2017).

Risultati Sperimentali
EoSeg supera costantemente i metodi basati su CNN, Transformer e ibridi encoder-decoder in tutte le modalità testate:

• Synapse (CT): Ha raggiunto l'85,50% di mean Dice (mDice), superando il precedente migliore (EViT-UNet con 80,87%) di 4,63 punti. Ha inoltre ottenuto punteggi top per organi specifici come l'Aorta (90,80%) e la Milza (94,13%).
• ACDC (MRI): Ha raggiunto il 91,73% di mDice, superando il precedente migliore di 0,24 punti.
• GlaS (Istopatologia): Ha raggiunto il 93,27% di mDice e 87,79% di mIoU, superando il precedente migliore rispettivamente di 0,83 e 1,29 punti.
• MoNuSeg (Istopatologia): Ha raggiunto l'80,51% di mDice e 67,45% di mIoU.
• Kvasir-Seg & ISIC: Ha dimostrato prestazioni superiori nella segmentazione di polipi e lesioni, raggiungendo alti punteggi di mIoU e mDice con confini fluidi e strutture coerenti.

L'analisi qualitativa conferma che EoSeg produce confini più fluidi e meno artefatti rispetto ai metodi concorrenti, preservando efficacemente i dettagli strutturali fini senza il restauro delle caratteristiche basato su decoder.

Significato e Rivendicazioni
Il paper sostiene che la dominanza del paradigma U-Net nella segmentazione di immagini mediche non è più giustificata dai limiti dei primi backbone ViT. Al contrario, le forti rappresentazioni apprese attraverso il pre-addestramento su larga scala consentono design encoder-only che possono predire direttamente maschere di segmentazione accurate.

Gli autori sottolineano che il successo di EoSeg non è dovuto meramente alla rimozione del decoder, quanto piuttosto all'adozione di un appropriato paradigma di predizione basato su query combinato con la modellazione multi-livello e la fusione adattiva. Questi risultati suggeriscono che i futuri framework di segmentazione di immagini mediche costruiti su modelli di visione pre-addestrati su larga scala possono allontanarsi da architetture decoder complesse, facendo affidamento invece sul potere rappresentativo intrinseco dei moderni ViT.