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

Technical Summary: From Convolution to Transformer

Problem Statement
Medical image segmentation is critical for computer-aided diagnosis, treatment planning, and disease monitoring. While the U-Net architecture has become the standard for biomedical segmentation due to its encoder-decoder structure and skip connections, conventional convolution-based variants face limitations. Specifically, standard U-Nets rely on local convolutional operations, which restrict their ability to capture long-range dependencies and global contextual information. This limitation is particularly detrimental in complex tasks such as brain tumor segmentation, where regions exhibit irregular shapes, heterogeneous appearances, and ambiguous boundaries. Furthermore, the performance of various U-Net modifications (residual, attention-based, or transformer-based) can vary significantly depending on the imaging modality (volumetric MRI vs. 2D fundus) and the target structure (tumor subregions vs. thin retinal vessels).

Methodology
This study conducts a comparative evaluation of five representative U-Net-based architectures across two distinct benchmark datasets:

1. BraTS 2023: A volumetric, multimodal MRI dataset for brain tumor segmentation (involving whole tumor, enhancing tumor, and tumor core subregions).
2. DRIVE: A 2D fundus image dataset for retinal vessel segmentation, characterized by thin, branching, and low-contrast structures.

The evaluated models represent a spectrum of design philosophies:

• 3D U-Net: Extends the original 2D architecture to volumetric data using 3D convolutions to capture inter-slice anatomical continuity.
• Residual U-Net: Incorporates residual blocks to improve gradient propagation and feature learning stability, aiding in the extraction of complex and heterogeneous structures.
• Attention U-Net: Introduces attention gates to the skip connections to suppress irrelevant background information and emphasize task-relevant regions.
• UNETR: Replaces the convolutional encoder with a Vision Transformer (ViT) to model global dependencies via self-attention mechanisms on image patches.
• Swin UNETR: Integrates Swin Transformer blocks into the U-Net framework, utilizing hierarchical representation learning and shifted window self-attention to efficiently capture both local details and global context.

All models were trained with consistent preprocessing (normalization, padding/cropping) and data augmentation strategies. Performance was assessed using quantitative metrics including Dice scores, Precision, Recall, F1 scores, and training loss convergence.

Key Contributions
The paper provides three primary contributions:

1. Comparative Evaluation: A systematic comparison of five U-Net variants (convolutional, residual, attention-based, and transformer-based) across both volumetric brain MRI and 2D retinal imaging tasks.
2. Mechanism Analysis: An analysis of how specific architectural mechanisms—residual learning, attention gates, and transformer self-attention—impact segmentation performance in different medical imaging scenarios.
3. Practical Guidance: Practical insights into model selection, distinguishing between architectures best suited for complex volumetric segmentation versus those effective for fine-structure extraction.

Experimental Results

• BraTS 2023 (Brain Tumor): Swin UNETR achieved the best overall performance with an average Dice score of 0.8965 (WT: 0.8966, ET: 0.9088, TC: 0.8841). UNETR followed with an average Dice of 0.8540. Residual U-Net (0.7341), 3D U-Net (0.6890), and Attention U-Net (0.6522) performed lower. The results indicate that transformer-based global context modeling is highly beneficial for volumetric tumor segmentation.
• DRIVE (Retinal Vessels): Swin UNETR again achieved the highest Dice score (0.8078), closely followed by Residual U-Net (0.8069) and UNETR (0.8015). Attention U-Net scored 0.7855, while 3D U-Net achieved the lowest Dice score (0.5967) despite having the lowest training loss (0.1030). This discrepancy highlights that lower training loss does not guarantee better segmentation quality, particularly when the architecture is mismatched to the data structure (e.g., 3D models on 2D fine structures).

Significance and Claims
The paper claims that Swin UNETR is the most effective architecture overall, successfully balancing local spatial detail preservation with long-range contextual modeling through its hierarchical encoder and shifted window attention. The study suggests that:

• Transformer-based variants (specifically Swin UNETR) are superior for complex segmentation tasks requiring global contextual understanding, such as brain tumor delineation.
• Residual learning remains highly effective for tasks requiring the preservation of fine structural details, such as retinal vessel segmentation, where Residual U-Net performed competitively with transformer models.
• Model selection should be task-dependent: transformer-based approaches are recommended for volumetric and complex tasks, while residual or hybrid approaches may suffice for fine-structure extraction.

The authors conclude that future work should explore hybrid architectures combining transformer-based contextual modeling with convolutional or residual modules to further enhance robustness and generalization across diverse clinical datasets.