[Paper Review] GCtx-UNet: Efficient Network for Medical Image Segmentation

Khaled Alrfou, Tian Zhao|arXiv (Cornell University)|Jun 9, 2024

Brain Tumor Detection and Classification6 citations

TL;DR

GCtx-UNet is a lightweight UNet-like architecture that blends GC-ViT global/local attention with CNN-based downsampling/upsampling to achieve competitive segmentation accuracy with lower model complexity across multiple medical imaging datasets.

ABSTRACT

Medical image segmentation is crucial for disease diagnosis and monitoring. Though effective, the current segmentation networks such as UNet struggle with capturing long-range features. More accurate models such as TransUNet, Swin-UNet, and CS-UNet have higher computation complexity. To address this problem, we propose GCtx-UNet, a lightweight segmentation architecture that can capture global and local image features with accuracy better or comparable to the state-of-the-art approaches. GCtx-UNet uses vision transformer that leverages global context self-attention modules joined with local self-attention to model long and short range spatial dependencies. GCtx-UNet is evaluated on the Synapse multi-organ abdominal CT dataset, the ACDC cardiac MRI dataset, and several polyp segmentation datasets. In terms of Dice Similarity Coefficient (DSC) and Hausdorff Distance (HD) metrics, GCtx-UNet outperformed CNN-based and Transformer-based approaches, with notable gains in the segmentation of complex and small anatomical structures. Moreover, GCtx-UNet is much more efficient than the state-of-the-art approaches with smaller model size, lower computation workload, and faster training and inference speed, making it a practical choice for clinical applications.

Motivation & Objective

• Address the need for accurate medical image segmentation without high computational cost.
• Integrate global context modeling with local attention in a UNet-like architecture.
• Leverage MedNet pre-training to improve in-domain performance over ImageNet pre-training.
• Evaluate on multiple datasets (Synapse, ACDC, Polyp) to demonstrate generalization and efficiency.

Proposed method

• Use GC-ViT blocks that combine local self-attention with global context queries to model long- and short-range dependencies.
• Incorporate a downsampling (Fused-MBConv) module to inject inductive bias and inter-channel modeling.
• Employ a GC-ViT based encoder–bottleneck–decoder with skip connections in a U-shaped architecture.
• Pre-train GC-ViT on MedNet (medical images) and compare with ImageNet pre-training.
• Use a patchify layer to create overlapping patches, followed by embedding projections in the encoder.

Experimental results

Research questions

• RQ1 Can GC-ViT-based blocks achieve competitive segmentation performance with fewer parameters than state-of-the-art CNN/Transformer-based models?
• RQ2 Does pre-training on domain-specific medical data (MedNet) improve segmentation accuracy over natural-image pre-training (ImageNet)?
• RQ3 How does GCtx-UNet perform across diverse medical imaging tasks (CT, MRI, polyp images) in terms of DSC and HD?
• RQ4 What is the impact of the upsampling/downsampling design and hyperparameters on segmentation performance?
• RQ5 Is the proposed architecture efficient in terms of model size, FLOPs, training time, and inference speed compared to peers?

Key findings

• GCtx-UNet achieves competitive to state-of-the-art results while having the smallest model size (12.34M parameters) and lowest FLOPs among surveyed methods.
• MedNet-pretrained GCtx-UNet generally outperforms ImageNet-pretrained variants on several datasets.
• On Synapse, GCtx-UNet2 (MedNet-pretrained) attains DSC 82.39% and HD 15.94 mm, among the best reported results with lower computation than many peers.
• On ACDC, GCtx-UNet2 achieves DSCs of 91.23 (RV), 89.88 (Myocardium), and 87.25 (LV), outperforming several Transformer-based and hybrid models.
• In polyp datasets, GCtx-UNet2 shows strong generalization to unseen datasets (CVC-ColonDB, ETIS-LaribDB, CVC-300) and often the top or near-top DSC scores.
• Ablation studies identify optimal loss combination (dice 0.3, cross-entropy 0.7) and learning rate (0.0001), and show transposed-convolution upsampling with SE blocks yields best performance.