[Paper Review] STA-Unet: Rethink the semantic redundant for Medical Imaging Segmentation

Vamsi Krishna Vasa, Wenhui Zhu+4 more|arXiv (Cornell University)|Oct 13, 2024

Radiomics and Machine Learning in Medical ImagingMedicine3 citations

TL;DR

This paper proposes STA-Unet, a U-Net architecture enhanced with Super Token Attention (STA) to reduce semantic redundancy in shallow layers of vision transformers, improving multi-organ segmentation. By replacing standard self-attention with sparse, superpixel-inspired token processing, STA-Unet achieves state-of-the-art performance on four medical imaging datasets, with Dice scores up to 2.86% higher than U-Net and 3.22% higher than TransUNet.

ABSTRACT

In recent years, significant progress has been made in the medical image analysis domain using convolutional neural networks (CNNs). In particular, deep neural networks based on a U-shaped architecture (UNet) with skip connections have been adopted for several medical imaging tasks, including organ segmentation. Despite their great success, CNNs are not good at learning global or semantic features. Especially ones that require human-like reasoning to understand the context. Many UNet architectures attempted to adjust with the introduction of Transformer-based self-attention mechanisms, and notable gains in performance have been noted. However, the transformers are inherently flawed with redundancy to learn at shallow layers, which often leads to an increase in the computation of attention from the nearby pixels offering limited information. The recently introduced Super Token Attention (STA) mechanism adapts the concept of superpixels from pixel space to token space, using super tokens as compact visual representations. This approach tackles the redundancy by learning efficient global representations in vision transformers, especially for the shallow layers. In this work, we introduce the STA module in the UNet architecture (STA-UNet), to limit redundancy without losing rich information. Experimental results on four publicly available datasets demonstrate the superiority of STA-UNet over existing state-of-the-art architectures in terms of Dice score and IOU for organ segmentation tasks. The code is available at \url{https://github.com/Retinal-Research/STA-UNet}.

Motivation & Objective

• Address the inherent feature redundancy in shallow layers of transformer-based U-Net architectures, which hampers efficient learning and increases computation.
• Investigate the under-explored issue of semantic redundancy in vision transformers applied to medical image segmentation.
• Enhance U-Net performance by integrating Super Token Attention (STA) to maintain rich semantic information while minimizing redundant attention computation.
• Demonstrate generalization and robustness across diverse medical imaging tasks, including multi-organ and gland/nucleus segmentation.
• Provide a computationally efficient alternative to over-parameterized transformer-based models without sacrificing segmentation accuracy.

Proposed method

• Integrate the Super Token Attention (STA) mechanism into the U-Net encoder-decoder architecture to replace standard multi-head self-attention layers.
• Transform pixel-level features into super tokens via superpixel-based clustering in the token space, reducing the number of tokens and focusing on salient visual regions.
• Apply sparse association learning and token space mapping to improve attention efficiency and reduce computation in shallow layers.
• Maintain skip connections between encoder and decoder, but enhance feature fusion using the STA module to preserve high-resolution contextual information.
• Use a learnable tokenization strategy that maps local image patches to compact, semantically meaningful super tokens, reducing redundant attention between nearby, similar pixels.
• Train the model end-to-end using cross-entropy and Dice loss to optimize pixel-wise segmentation performance on medical imaging benchmarks.

Experimental results

Research questions

• RQ1 To what extent does semantic redundancy in shallow layers of transformer-based U-Net models impair segmentation performance and efficiency?
• RQ2 Can the Super Token Attention (STA) mechanism effectively reduce redundant attention computation in early layers while preserving essential semantic features?
• RQ3 How does STA-Unet compare to state-of-the-art U-Net variants in terms of segmentation accuracy across diverse medical imaging datasets?
• RQ4 What is the impact of key hyperparameters—such as token size and number of attention heads—on the model’s performance and efficiency?
• RQ5 Can the proposed architecture generalize to challenging segmentation tasks, such as small or low-contrast organs (e.g., pancreas, kidneys) and fine-grained structures (e.g., nuclei, glands)?

Key findings

• STA-Unet achieves a 92.25% mean Dice score on the ACDC dataset, outperforming U-Net by 2.86 percentage points and TransUNet by 2.83 percentage points.
• On the MoNuSeg dataset, STA-Unet achieves a Dice score of 81.06%, a 6.03% improvement over U-Net and a 3.22% gain over TransUNet.
• For the Glas dataset, STA-Unet achieves 91.03% Dice score, surpassing U-Net by 6.53 percentage points and TransUNet by 2.97 percentage points.
• STA-Unet significantly improves segmentation of small and challenging organs, such as the pancreas and kidneys, outperforming SwinUNet and other models that fail to segment these structures.
• Visual analysis confirms that STA-Unet produces more accurate and complete foreground predictions than SwinUNet and LeViT-UNet, especially in complex regions like glands and nuclei.
• Ablation studies show that both token size and number of attention heads significantly influence performance, with optimal settings balancing accuracy and computational cost.