PaperlayerView exact revision on arXiv

Source-linked AI summary

DS-TransUNet:Dual Swin Transformer U-Net for Medical Image Segmentation

Ailiang Lin, Bingzhi Chen, Jiayu Xu, Zheng Zhang, Guangming Lu

arXiv:2106.06716v1cs.CV

TL;DR

Medical image segmentation needs models that capture long-range context while preserving pixel-level structure, limitations associated with CNNs and patch-based Transformers. DS-TransUNet combines dual-scale Swin encoders, TIF self-attention fusion, and Swin Transformer decoding, and experiments across four tasks show it significantly outperforms state-of-the-art methods.

  • Problem

    CNNs have limited receptive fields for long-range dependencies, while Transformer patch division can ignore pixel-level intrinsic structural features.

  • Method

    DS-TransUNet uses dual-scale Swin Transformer encoder subnetworks, TIF self-attention fusion, and Swin Transformer blocks in the decoder.

  • Results

    Across four typical medical image segmentation tasks, DS-TransUNet significantly outperforms state-of-the-art methods, especially for polyp segmentation.

  • Takeaways & Limitations

    The framework combines hierarchical Swin Transformers with a U-shaped architecture to improve medical image segmentation quality through multiscale and long-range feature modeling.

Abstract

from arXiv · show

Automatic medical image segmentation has made great progress benefit from the development of deep learning. However, most existing methods are based on convolutional neural networks (CNNs), which fail to build long-range dependencies and global context connections due to the limitation of receptive field in convolution operation. Inspired by the success of Transformer in modeling the long-range contextual information, some researchers have expended considerable efforts in designing the robust variants of Transformer-based U-Net. Moreover, the patch division used in vision transformers usually ignores the pixel-level intrinsic structural features inside each patch. To alleviate these problems, we propose a novel deep medical image segmentation framework called Dual Swin Transformer U-Net (DS-TransUNet), which might be the first attempt to concurrently incorporate the advantages of hierarchical Swin Transformer into both encoder and decoder of the standard U-shaped architecture to enhance the semantic segmentation quality of varying medical images. Unlike many prior Transformer-based solutions, the proposed DS-TransUNet first adopts dual-scale encoder subnetworks based on Swin Transformer to extract the coarse and fine-grained feature representations of different semantic scales. As the core component for our DS-TransUNet, a well-designed Transformer Interactive Fusion (TIF) module is proposed to effectively establish global dependencies between features of different scales through the self-attention mechanism. Furthermore, we also introduce the Swin Transformer block into decoder to further explore the long-range contextual information during the up-sampling process. Extensive experiments across four typical tasks for medical image segmentation demonstrate the effectiveness of DS-TransUNet, and show that our approach significantly outperforms the state-of-the-art methods.

I. INTRODUCTION

DS-TransUNet addresses CNNs’ limited long-range context and patch-based Transformers’ loss of pixel-level structure by combining dual-scale Swin Transformer encoding, TIF fusion, and Swin decoding for medical image segmentation.

  • CNNs struggle to build long-range dependencies because convolutional kernels focus on sub-regions and lose global context.
  • DS-TransUNet uses dual-scale Swin Transformer encoder subnetworks to extract coarse and fine-grained representations at different semantic scales.Images are sliced into non-overlapping patches at large and small scales.
  • Swin Transformer blocks in both encoder and decoder model long-range dependencies and multi-scale context during down-sampling and up-sampling.
  • The TIF module fuses coarse and fine-grained features through self-attention to establish global dependencies and semantic consistency.
  • Across four medical image segmentation tasks, DS-TransUNet consistently outperforms previous state-of-the-art methods, especially for polyp segmentation.

II. RELATED WORK

Related work progresses from CNN-based U-Net variants to Transformer architectures and multi-scale vision Transformers. DS-TransUNet builds on these directions with a dual-branch Swin encoder, TIF fusion, and a U-shaped Transformer design.

  • CNN-based methods: CNN-based U-Net variants use encoder-decoder pathways, skip connections, attention, dense connections, or residual structures for medical image segmentation.
  • Vision Transformers: Vision Transformers introduced self-attention-based image modeling, while Swin Transformer added shifted-window attention and hierarchical representations for dense prediction.
  • Multi-scale feature representations: Multi-scale Transformer research includes dual-branch, cross-attention, and pyramid-based approaches for extracting or connecting features at different scales.
  • Proposed method: DS-TransUNet extends these directions with a dual-branch Swin encoder and a Transformer Interactive Fusion module for multiscale feature fusion.
  • Proposed method: The method section presents a U-shaped architecture whose encoder and decoder are based on Swin Transformer blocks.

A. Swin Transformer block

The Swin Transformer block combines local window attention with shifted windows to model cross-window interactions efficiently. DS-TransUNet uses these blocks within its U-shaped architecture.

  • Standard Transformer blocks combine multi-head self-attention, layer normalization, MLPs, and residual connections.
  • DS-TransUNet adopts Swin Transformer blocks as part of its encoder-decoder segmentation architecture.
  • Swin Transformer limits self-attention to non-overlapping local windows to reduce the quadratic cost of global token relationships.
  • Shifted-window self-attention addresses the lack of interaction between neighboring windows without additional computation.
  • The shifted partition uses cyclic shifting, while both attention variants incorporate relative position bias in similarity computation.

B. Encoder

The encoder extracts hierarchical features from non-overlapping image patches using four Swin Transformer stages, while the decoder combines upsampling, skip connections, and Swin blocks.

  • Encoder: The encoder converts image patches into tokens and processes them through four hierarchical Swin Transformer stages.
  • Encoder: Patch merging concatenates each group of 2 × 2 neighboring patches, reducing token count by 4 while doubling feature dimension.
  • Decoder: The decoder has three stages that upsample features by 2, concatenate same-stage encoder skip features, and apply Swin Transformer blocks.
  • Decoder: Additional low-level features are extracted at H×W and H 2 × H 2 resolutions and used through skip connections for final mask prediction.

D. Multi-Scale Feature Representations

DS-TransUNet uses dual-scale Swin Transformer branches to capture complementary coarse- and fine-grained features, then fuses them through Transformer Interactive Fusion.

  • Motivation: Multi-scale Swin Transformer feature extraction addresses the loss of pixel-level structure caused by patch division.
  • Dual-scale features: Large patches capture coarse-grained features, whereas small patches obtain fine-grained features, making the two scales complementary.
  • Dual-scale features: The two encoder branches produce feature maps at different resolutions across multiple stages.
  • Transformer Interactive Fusion: TIF uses multi-head self-attention to enable interaction between features from the two scales instead of simple concatenation followed by convolution.
  • Transformer Interactive Fusion: TIF transforms one branch into a global token, concatenates it with pixel-level tokens from the other branch, and computes global self-attention.
  • Transformer Interactive Fusion: The resulting connections let fine-grained features obtain coarse-grained information from the large-scale branch.

IV. EXPERIMENTS

The study evaluates DS-TransUNet across four common medical image segmentation tasks using publicly available datasets and comparisons with state-of-the-art methods.

  • Experiments assess the learning and generalization ability of DS-TransUNet on four common medical image segmentation tasks.
  • The evaluation uses several publicly available datasets and compares DS-TransUNet with other state-of-the-art methods.

A. Datasets

The study evaluates medical image segmentation across polyp, skin-lesion, gland, and nuclei datasets, using standard segmentation metrics including mDice, mIoU, precision, and recall.

  • Polyp Segmentation: Polyp experiments use Kvasir, CVC-ColonDB, EndoScene, ETIS, and CVC-ClinicDB under three published split configurations.The configurations include Kvasir-only, ClinicDB-only, and training on Kvasir plus ClinicDB with five test datasets.
  • ISIC 2018 Dataset: The ISIC 2018 dataset contains 2,596 images resized to 256 × 256, with 2,076 for training and 520 for testing.The split is repeated five times and averaged.
  • GLAS Dataset: GLAS contains 165 histology images, split into 85 training and 80 testing images, and resized to 128 × 128.
  • 2018 Data Science Bowl: The 2018 Data Science Bowl dataset contains 670 images, divided into 80% training, 10% validation, and 10% testing after resizing to 256 × 256.
  • Evaluation Metrics: Comparisons use mean Dice coefficient, mean Intersection over Union, precision, and recall, based on true-positive, true-negative, false-positive, and false-negative values.

C. Implementation Details

Training uses multi-scale inputs, weighted IoU and binary cross-entropy losses, deep supervision, and fixed optimization and training schedules across two model variants.

  • Training Strategy: All experiments use multi-scale training instead of data augmentation, with weighted IoU and binary cross-entropy losses.
  • Deep Supervision: Deep supervision additionally supervises encoder stage 4 output S2 and decoder stage 1 output S3.
  • Optimization: The total loss weights are α=0.6, β=0.2, and γ=0.2, while SGD uses momentum 0.9, weight decay 1e-4, and learning rate 0.01.
  • Implementation: Models train for 100 epochs with early stopping and cosine annealing on an NVIDIA RTX 3090 GPU using PyTorch.
  • Model Variants: DS-TransUNet-B uses Swin-Base as its primary encoder branch, whereas DS-TransUNet-L uses Swin-Large; both use Swin-Tiny as the complementary branch.

V. RESULT

DS-TransUNet is evaluated against state-of-the-art methods across four medical segmentation tasks and reports strong quantitative and qualitative performance across multiple datasets.

  • Experimental Scope: The experiments compare DS-TransUNet with state-of-the-art methods across polyp, skin-lesion, gland, and nuclei segmentation tasks.
  • Polyp Segmentation: On Kvasir, DS-TransUNet-L achieves mDice 0.913, mIoU 0.859, recall 0.936, and precision 0.916, improving over HarDNet-MSE by 0.9%, 1.1%, 1.3%, and 0.9%.
  • Polyp Segmentation: On ClinicDB, DS-TransUNet-L improves over FANet by 0.67% F1, 0.02% mIoU, and 1.6% recall, while achieving comparable precision to DoubleU-Net.
  • Polyp Segmentation: Across five polyp datasets, DS-TransUNet-L achieves the top performance on all five and improves average mDice by about 1.9% over TransFuse.
  • ISIC 2018: On ISIC 2018, DS-TransUNet-L reaches F1 0.9132, mIoU 0.8523, and recall 0.9217, improving over the state-of-the-art method by 1.70%, 3.11%, and 4.37%.
  • GLAS: On GLAS, DS-TransUNet-L improves over KiU-Net by 3.94% mDice and 5.67% mIoU despite only 85 training samples.
  • 2018 Data Science Bowl: On the 2018 Data Science Bowl, DS-TransUNet-L achieves F1 0.9219, mIoU 0.8612, and recall 0.9378.

B. Ablation study

Ablations evaluate Swin Transformer components, decoder design, multi-scale branches, and TIF, showing that the complete architecture performs best among the tested variants.

  • Ablation Setup: The ablation study compares four DS-TransUNet variants on polyp segmentation using mean Dice and mean IoU, with TransFuse as the baseline.
  • Effect of Swin Transformer: Swin Transformer provides the encoder representation through window-based and shifted-window self-attention, producing results close to TransFuse.
  • Decoder Design: Swin U-Net improves average mean Dice by 0.4% over TransFuse, while adding Swin Transformer blocks after upsampling further improves decoder performance.
  • Decoder Design: The Swin Decoder improves average mDice by 1.5% over TransFuse across five challenging polyp datasets.
  • Multi-Scale Fusion: Adding a second encoder branch without TIF fails to improve performance, indicating that convolution-based multi-scale fusion is insufficient in this comparison.The passage attributes this to ineffective multi-scale fusion and harder convergence.
  • Transformer Interactive Fusion: Adding TIF yields the best variant, improving mDice over the single-branch Swin Decoder by 0.6%, 0.2%, 0.7%, and 1.3% on Kvasir, ClinicDB, EndoScene, and ETIS.It also improves ColonDB mIoU by 0.5%.

VI. CONCLUSION

DS-TransUNet is a dual-scale Swin Transformer U-shaped framework that enhances medical image segmentation by integrating Swin Transformer blocks throughout the encoder-decoder and fusing multi-scale features with TIF. Experiments across four segmentation tasks report significant improvements over state-of-the-art methods, especially for polyp segmentation.

  • DS-TransUNet is a U-shaped encoder-decoder framework based on hierarchical Swin Transformer for medical image segmentation.
  • Its encoder uses dual-branch Swin Transformers to extract multi-scale feature representations.
  • Swin Transformer blocks are added to the decoder to model long-range contextual information during up-sampling.
  • The Transformer Interactive Fusion module builds long-range dependencies between different-scale features through self-attention and fuses encoder features.
  • Experiments across four medical image segmentation tasks show significant improvement over state-of-the-art methods, especially for polyp segmentation.
Loading 2106.06716v1…