Does Your ViT Still Need U-Net for Segmentation?

Technische Samenvatting: Heeft uw ViT nog een U-Net nodig voor segmentatie?

Probleemstelling
Medische beeldsegmentatie wordt al meer dan een decennium gedomineerd door U-Net-achtige encoder-decoderarchitecturen. Hoewel deze ontwerpen succesvol ruimtelijke details aanpakken via skip-verbindingen en symmetrische decodering, vertrouwen ze op convolutionele operaties die beperkt zijn door een beperkt receptief veld, waardoor het modelleren van globale context moeilijk is. Hoewel Vision Transformers (ViTs) self-attention-mechanismen bieden om langetermijnrelaties te vangen, behielden vroege ViT-gebaseerde segmentatiemethoden (bijv. TransUNet, SwinUNet) U-Net decoders omdat pre-trained ViT-representaties als onvoldoende werden beschouwd voor nauwkeurige dense voorspellingen.

Recente vooruitgang in grootschalige pretraining (bijv. DINOv2, SigLIP) heeft de representatiecapaciteiten van ViTs fundamenteel hergedefinieerd, waardoor ze krachtige visuele kenmerken kunnen leren van enorme ongelabelde datasets. Deze evolutie roept een kritische vraag op: Is het U-Net paradigma nog wel noodzakelijk voor medische beeldsegmentatie wanneer moderne, sterk pre-trained ViT-backbones beschikbaar zijn? Zo niet, hoe moet een encoder-only framework worden ontworend om effectief segmentatiemaskers te voorspellen zonder decoder?

Methodologie: EoSeg
De auteurs stellen EoSeg (Encoder-only Segmentation) voor, een framework dat de U-Net decoder elimineert ten gunste van een query-gebaseerd ontwerp gebouwd op moderne pre-trained ViT-backbones. De methodologie is ontwikkeld door middel van een systematische verkenning van cruciale architecturale keuzes:

1. Backbone Selectie: De studie evalueert verschillende pre-trained ViT-backbones (inclusief ViT-L/14, SigLig, MTP, DINOv2 en DINOv3) in een encoder-only setting. DINOv2 wordt gekozen als de optimale backbone, die de hoogste baseline prestatie bereikt (81,86% mDice op Synapse) zonder enige decoder.
2. Mask Generatie:
   • Deconvolutie vs. Query-gebaseerd: Aanvankelijk werd een deconvolutionale module getest om low-resolution ViT-features te upsamplen. De auteurs ontdekten echter dat een query-gebaseerde mask voorspelling superieur was. In plaats van pixel-wise decodering fungeren leerbare query-embeddings als globale semantische probes.
   • Mechanisme: Query-vectoren worden geprojecteerd in een mask-embedding ruimte en interageren met de upgesamplede globale feature map via een kanaal-gewijze dot product om direct binaire maskers te genereren.
3. Multi-level Query Modellering: Geïnspireerd door skip-verbindingen in U-Nets, maakt het framework gebruik van representaties van de laatste drie Transformer blocks in plaats van alleen het laatste block. Leerbare queries worden geïnjecteerd vóór deze blocks en worden verfijnd door self-attention met image tokens, waardoor het model semantische informatie uit meerdere stadia kan aggregeren.
4. Learnable Block Fusion: Om voorspellingen van de drie tussenliggende blocks te combineren, is een learnable fusion module geïntroduceerd. Een fusion gate gebruikt de query-representatie om adaptieve gewichten voor de output van elk block te voorspellen, waardoor het model dynamisch kan bepalen welke bijdrage elk niveau levert aan de uiteindelijke voorspelling.
5. Query Configuratie: Het framework hanteert een query-gebaseerd classificatieparadigma waarbij elke query een mask en een klasse-label voorspelt. Experimenten bepaalden dat het toewijzen van twee queries per categorie optimale prestaties oplevert, wat een balans vindt tussen het vangen van intra-klasse variaties en het voorkomen van over-parameterisatie.
6. Supervisie Strategie: In tegen tegenstelling tot gangbare praktijken in medische segmentatie, bleek deep supervision (het toepassen van loss op tussenliggende outputs) in deze specifieke encoder-only setup niet de prestaties te verbeteren en werd het uit het definitieve ontwerp uitgesloten.

Belangrijkste Bijdragen

• Architecturale Herwaardering: Het artikel biedt een uitgebreide evaluatie die aantoont dat een U-Net-achtige decoder niet langer een vereiste is voor hoogwaardige medische beeldsegmentatie wanneer moderne, grootschalige pre-trained ViT-backbones worden ingezet.
• EoSeg Framework: De auteurs presenteren EoSeg, een query-gebaseerd encoder-only framework dat multi-level query modellering en learnable block fusion integreert om directe mask voorspelling vanuit encoder features mogelijk te maken.
• Uitgebreide Validatie: Het voorgestelde ontwerp wordt gevalideerd over zeven benchmark-datasets die vijf verschillende medische beeldvormingsmodaliteiten beslaan: CT (Synapse), MRI (ACDC), Histopathologie (GlaS, MoNuSeg), Endoscopie (Kvasir-Seg) en Dermoscopie (ISIC-2016, ISIC-2017).

Experimentele Resultaten
EoSeg presteert consistent beter dan state-of-the-art CNN-gebaseerde, Transformer-gebaseerde en hybride encoder-decoder methoden over alle geteste modaliteiten:

• Synapse (CT): Behaalde 85,50% mean Dice (mDice), waarmee het de vorige beste (EViT-UNet met 80,87%) met 4,63 punten overtrof. Het behaalde ook top scores voor specifieke organen zoals de Aorta (90,80%) en Milt (94,13%).
• ACDC (MRI): Behaalde 91,73% mDice, waarmee het de vorige beste met 0,24 punten overtrof.
• GlaS (Histopathologie): Behaalde 93,27% mDice en 87,79% mIoU, waarmee het respectievelijk de vorige beste met 0,83 en 1,29 punten overtrof.
• MoNuSeg (Histopathologie): Behaalde 80,51% mDice en 67,45% mIoU.
• Kvasir-Seg & ISIC: Demonstreerden superieure prestaties bij polyp en laesie segmentatie, met hoge mIoU en mDice scores en vloeiende grenzen en coherente structuren.

Kwalitatieve analyse bevestigt dat EoSeg vloeiendere grenzen en minder artefacten produceert vergeleken met concurrerende methoden, waarbij fijne structurele details effectief worden behouden zonder decoder-gebaseerde feature restauratie.

Betekenis en Claims
Het artikel claimt dat de dominantie van het U-Net paradigma in medische beeldsegmentatie niet langer gerechtvaardigd is door de beperkingen van vroege ViT-backbones. In plaats daarvan maken de sterke representaties die geleerd zijn door grootschalige pretraining het mogelijk om encoder-only ontwerpen te gebruiken die direct nauwkeurige segmentatiemaskers kunnen voorspellen.

De auteurs benadrukken dat het succes van EoSeg niet enkel te danken is aan het verwijderen van de decoder, maar aan de adoptie van een passend query-gebaseerd voorspellingsparadigma gecombineerd met multi-level modellering en adaptieve fusie. Deze bevindingen suggereren dat toekomstige medische beeldsegmentatie frameworks, gebouwd op grootschalige pre-trained vision modellen, kunnen afstappen van complexe decoder-architecturen en kunnen vertrouwen op de intrinsieke representatiekracht van moderne ViTs.