Does Your ViT Still Need U-Net for Segmentation?

Technical Summary: Does Your ViT Still Need U-Net for Segmentation?

Problem Statement
Medical image segmentation has been dominated for over a decade by U-Net-style encoder–decoder architectures. While these designs successfully address the loss of spatial details through skip connections and symmetric decoding, they rely on convolutional operations that are constrained by limited receptive fields, making global context modeling difficult. Although Vision Transformers (ViTs) offer self-attention mechanisms to capture long-range dependencies, early ViT-based segmentation methods (e.g., TransUNet, SwinUNet) retained U-Net decoders because pre-trained ViT representations were deemed insufficient for accurate dense prediction.

Recent advances in large-scale pretraining (e.g., DINOv2, SigLIP) have significantly redefined the representation capabilities of ViTs, enabling them to learn powerful visual features from massive unlabeled datasets. This evolution prompts a critical question: Is the U-Net paradigm still necessary for medical image segmentation when modern, strongly pre-trained ViT backbones are available? If not, how should an encoder-only framework be designed to effectively predict segmentation masks without a decoder?

Methodology: EoSeg
The authors propose EoSeg (Encoder-only Segmentation), a framework that eliminates the U-Net decoder in favor of a query-based design built upon modern pre-trained ViT backbones. The methodology is developed through a systematic exploration of key architectural choices:

1. Backbone Selection: The study evaluates various pre-trained ViT backbones (including ViT-L/14, SigLIP, MTP, DINOv2, and DINOv3) in an encoder-only setting. DINOv2 is selected as the optimal backbone, achieving the highest baseline performance (81.86% mDice on Synapse) without any decoder.
2. Mask Generation:
   • Deconvolution vs. Query-based: Initially, a deconvolutional module was tested to upsample low-resolution ViT features. However, the authors found that a query-based mask prediction approach was superior. Instead of pixel-wise decoding, learnable query embeddings act as global semantic probes.
   • Mechanism: Query vectors are projected into a mask embedding space and interact with the upsampled global feature map via a channel-wise dot product to generate binary masks directly.
3. Multi-level Query Modeling: Inspired by skip connections in U-Nets, the framework leverages representations from the last three Transformer blocks rather than just the final block. Learnable queries are injected before these blocks and refined through self-attention with image tokens, allowing the model to aggregate semantic information from multiple stages.
4. Learnable Block Fusion: To combine predictions from the three intermediate blocks, a learnable fusion module is introduced. A fusion gate uses the query representation to predict adaptive weights for each block's output, allowing the model to dynamically determine the contribution of each level to the final prediction.
5. Query Configuration: The framework adopts a query-based classification paradigm where each query predicts a mask and a class label. Experiments determined that assigning two queries per category yields optimal performance, balancing the need to capture intra-class variations without over-parameterization.
6. Supervision Strategy: Contrary to common practices in medical segmentation, deep supervision (applying loss to intermediate outputs) was found not to improve performance in this specific encoder-only setup and was excluded from the final design.

Key Contributions

• Architectural Re-evaluation: The paper provides an extensive evaluation demonstrating that a U-Net-style decoder is no longer a prerequisite for high-quality medical image segmentation when leveraging modern, large-scale pre-trained ViT backbones.
• EoSeg Framework: The authors present EoSeg, a query-based encoder-only framework that integrates multi-level query modeling and learnable block fusion to enable direct mask prediction from encoder features.
• Comprehensive Validation: The proposed design is validated across seven benchmark datasets spanning five distinct medical imaging modalities: CT (Synapse), MRI (ACDC), Histopathology (GlaS, MoNuSeg), Endoscopy (Kvasir-Seg), and Dermoscopy (ISIC-2016, ISIC-2017).

Experimental Results
EoSeg consistently outperforms state-of-the-art CNN-based, Transformer-based, and hybrid encoder–decoder methods across all tested modalities:

• Synapse (CT): Achieved 85.50% mean Dice (mDice), surpassing the previous best (EViT-UNet at 80.87%) by 4.63 points. It also achieved top scores for specific organs like the Aorta (90.80%) and Spleen (94.13%).
• ACDC (MRI): Achieved 91.73% mDice, exceeding the previous best by 0.24 points.
• GlaS (Histopathology): Achieved 93.27% mDice and 87.79% mIoU, outperforming the previous best by 0.83 and 1.29 points, respectively.
• MoNuSeg (Histopathology): Achieved 80.51% mDice and 67.45% mIoU.
• Kvasir-Seg & ISIC: Demonstrated superior performance in polyp and lesion segmentation, achieving high mIoU and mDice scores with smooth boundaries and coherent structures.

Qualitative analysis confirms that EoSeg produces smoother boundaries and fewer artifacts compared to competing methods, effectively preserving fine structural details without decoder-based feature restoration.

Significance and Claims
The paper claims that the dominance of the U-Net paradigm in medical image segmentation is no longer justified by the limitations of early ViT backbones. Instead, the strong representations learned through large-scale pretraining allow for encoder-only designs that can directly predict accurate segmentation masks.

The authors emphasize that the success of EoSeg is not merely due to removing the decoder, but rather the adoption of an appropriate query-based prediction paradigm combined with multi-level modeling and adaptive fusion. These findings suggest that future medical image segmentation frameworks built on large-scale pre-trained vision models can move away from complex decoder architectures, relying instead on the intrinsic representational power of modern ViTs.