PaperlayerView exact revision on arXiv

Source-linked AI summary

Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, Manning Wang

arXiv:2105.05537v1eess.IVcs.CV

TL;DR

Medical image segmentation still lacks effective global and long-range semantic modeling with CNNs. Swin-Unet uses a pure Swin Transformer U-shaped encoder-decoder with skip connections and outperforms convolutional and hybrid methods on multi-organ and cardiac tasks.

  • Problem

    CNN-based segmentation methods struggle to model explicit global and long-range semantic interactions, limiting their ability to meet medical applications’ accuracy requirements.

  • Method

    Swin-Unet tokenizes image patches and processes them with a symmetric Swin Transformer U-shaped encoder-decoder using patch expansion and skip connections.

  • Results

    Swin-Unet outperformed full-convolution and transformer-convolution methods across multi-organ and cardiac segmentation tasks, achieving 79.13% DSC and 21.55% HD on Synapse.

  • Takeaways & Limitations

    The results support pure Transformer U-shaped architectures as effective approaches for medical image segmentation across the evaluated multi-organ and cardiac tasks.

  • Takeaways & Limitations

    The method uses ImageNet-pretrained weights in a potentially suboptimal way and is evaluated on 2D images, leaving end-to-end pretraining and 3D segmentation for future work.

Abstract

from arXiv · show

In the past few years, convolutional neural networks (CNNs) have achieved milestones in medical image analysis. Especially, the deep neural networks based on U-shaped architecture and skip-connections have been widely applied in a variety of medical image tasks. However, although CNN has achieved excellent performance, it cannot learn global and long-range semantic information interaction well due to the locality of the convolution operation. In this paper, we propose Swin-Unet, which is an Unet-like pure Transformer for medical image segmentation. The tokenized image patches are fed into the Transformer-based U-shaped Encoder-Decoder architecture with skip-connections for local-global semantic feature learning. Specifically, we use hierarchical Swin Transformer with shifted windows as the encoder to extract context features. And a symmetric Swin Transformer-based decoder with patch expanding layer is designed to perform the up-sampling operation to restore the spatial resolution of the feature maps. Under the direct down-sampling and up-sampling of the inputs and outputs by 4x, experiments on multi-organ and cardiac segmentation tasks demonstrate that the pure Transformer-based U-shaped Encoder-Decoder network outperforms those methods with full-convolution or the combination of transformer and convolution. The codes and trained models will be publicly available at https://github.com/HuCaoFighting/Swin-Unet.

1 Introduction

Medical image segmentation supports computer-aided diagnosis and image-guided surgery, but CNN-based methods struggle to model global and long-range semantic interactions. Swin-Unet addresses this limitation with a pure Transformer U-shaped architecture using Swin Transformer blocks, patch tokens, patch expanding, and skip connections.

  • Accurate and robust medical image segmentation is important for computer-aided diagnosis and image-guided clinical surgery.
  • Existing methods mainly use fully convolutional U-shaped networks with symmetric encoder-decoder structures and skip connections.U-Net extracts deep features through convolution and down-sampling, then restores input resolution through decoder up-sampling for pixel-level prediction.
  • CNN-based approaches remain limited because convolutional locality makes explicit global and long-range semantic interaction difficult.Despite strong segmentation performance, these methods do not fully satisfy medical applications’ strict accuracy requirements.
  • Swin-Unet is proposed as a pure Transformer-based U-shaped architecture comprising an encoder, bottleneck, decoder, and skip connections.The architecture is built from Swin Transformer blocks and is presented as a first pure Transformer-based U-shaped design for 2D medical image segmentation.
  • Swin-Unet tokenizes non-overlapping image patches, applies Transformer encoding, and uses patch expanding plus skip connections to restore spatial resolution for segmentation.Decoder features are fused with multi-scale encoder features before segmentation prediction; experiments cover multi-organ and cardiac segmentation datasets.

2 Related work

Related work progresses from traditional and CNN-based medical image segmentation toward Transformer-based and hybrid CNN–Transformer approaches. U-shaped architectures remain a central benchmark, while ViT’s pretraining requirement motivates further architectural development.

  • CNN-based methods: Medical image segmentation evolved from contour-based and traditional machine-learning methods to U-Net and numerous U-shaped variants, including Res-UNet, Dense-UNet, U-Net++, UNet3+, and 3D-Unet.U-Net’s simplicity and superior performance contributed to the continued emergence of Unet-like methods.
  • Transformer-based methods: Transformers achieved state-of-the-art performance across NLP tasks, motivating vision applications such as ViT, which offered an impressive speed-accuracy trade-off for image recognition.Transformer was initially proposed for machine translation before its broader adoption in NLP and vision.
  • Transformer-based methods: Compared with CNN-based methods, ViT requires pre-training on its own large dataset, presenting a drawback for vision applications.The supplied passage identifies this pretraining requirement as ViT’s drawback relative to CNN-based methods.
  • Benchmark architectures: U-shaped architectures with skip connections provide a benchmark comparison for developing Transformer-based methods in medical image analysis.The supplied passage explicitly frames such architectures as a benchmark for Transformer development.
  • Hybrid methods: Recent studies have introduced self-attention into CNNs and combined CNNs with Transformers to improve medical image segmentation beyond purely CNN-based methods.Examples include additive attention gates in U-shaped skip connections and approaches combining CNN and Transformer components.

3 Method · 3.1 Architecture overview · 3.2 Swin Transformer block

Swin-Unet is a U-shaped architecture composed of encoder, bottleneck, decoder, and skip connections, with all main stages built from Swin Transformer blocks. Its encoder tokenizes 4 × 4 image patches, while shifted-window attention and patch-based self-attention support hierarchical feature processing.

  • 3.1 Architecture overview: The encoder splits medical images into non-overlapping 4 × 4 patches and represents each patch with 4 × 4 × 3 = 48 features before linear embedding.Patch partitioning transforms the input image into sequence embeddings.
  • 3.1 Architecture overview: The basic computational unit of Swin-Unet is the Swin Transformer block.The architecture overview identifies the Swin Transformer block as the network’s basic unit.
  • 3.1 Architecture overview: Swin-Unet comprises an encoder, bottleneck, decoder, and skip connections, with encoder, bottleneck, and decoder constructed from Swin Transformer blocks.The architecture is presented as an encoder–decoder design with skip connections.
  • 3.2 Swin Transformer block: Unlike conventional multi-head self attention, the Swin Transformer block uses shifted windows and alternating W-MSA and SW-MSA modules across consecutive blocks.The window partitioning mechanism formulates continuous Swin Transformer blocks.
  • 3.2 Swin Transformer block: Each Swin Transformer block contains LayerNorm, multi-head self attention, residual connection, and a 2-layer MLP with GELU non-linearity.These components define the internal structure of each block.
  • 3.2 Swin Transformer block: The block’s outputs are denoted ˆz_l and z_l for the (S)W-MSA and MLP modules, respectively, with self-attention computed using query, key, and value matrices.The notation distinguishes the intermediate attention output from the block output.
  • 3.2 Swin Transformer block: In the attention formulation, Q, K, and V ∈R^M 2×d denote query, key, and value matrices, while M 2 is the number of patches in a window and d is the query or key dimension.The formulation also uses values from a bias matrix B̂ ∈ R^(2M−1)×(2M+1).

3.3 Encoder

The encoder applies consecutive Swin Transformer blocks for representation learning, then repeatedly uses patch merging to down-sample resolution and double feature dimensions.

  • Encoder: Two consecutive Swin Transformer blocks learn representations while preserving feature dimension and resolution.The encoder repeats this procedure three times.
  • Encoder: Patch merging increases feature dimension to 2× the original dimension after down-sampling.Concatenation first increases feature dimension by 4×, followed by linear projection.
  • Encoder: Patch merging divides input patches into four parts, concatenates them, and reduces feature resolution by 2×.A linear layer follows concatenation to unify the feature dimension.

Patch merging layer · 3.4 Bottleneck

The bottleneck uses two successive Swin Transformer blocks because excessive Transformer depth can prevent convergence. Its feature dimension and resolution remain unchanged while learning deep feature representations.

  • 3.4 Bottleneck: The bottleneck contains two successive Swin Transformer blocks to learn deep feature representations.The paper states that Transformer is too deep to be converged, motivating this limited depth.
  • 3.4 Bottleneck: The bottleneck keeps the feature dimension unchanged.No dimensional transformation is applied within the bottleneck.
  • 3.4 Bottleneck: The bottleneck keeps the feature resolution unchanged.Spatial resolution is preserved across the bottleneck blocks.
  • 3.4 Bottleneck: The bottleneck depth is restricted because the Transformer is described as too deep to converge.This convergence concern explains why only two successive blocks are used.
  • 3.4 Bottleneck: Deep feature representation learning is performed without changing feature dimension or resolution.The bottleneck combines representation learning with unchanged feature geometry.
  • Patch merging layer: The provided passage does not specify the operation or configuration of a separate patch merging layer.No additional patch-merging mechanism is described in the supplied evidence.

3.5 Decoder

The decoder mirrors the encoder with Swin Transformer blocks and uses patch expanding layers to up-sample deep features while reducing their channel dimension.

  • 3.5 Decoder: The symmetric decoder is built from Swin Transformer blocks corresponding to the encoder.
  • 3.5 Decoder: Patch expanding layers replace encoder patch merging to reshape adjacent feature-map dimensions into higher-resolution maps.
  • 3.5 Decoder: 2× up-sampling increases spatial resolution while reducing the feature dimension to half its original size.

Patch expanding layer

The patch expanding layer upsamples feature maps by first increasing channel dimensionality with a linear layer, then rearranging patches to enlarge spatial resolution while reducing feature dimensionality.

  • Patch expanding layer: A linear layer first doubles the input feature dimension before up-sampling.The transformation increases the feature dimension from 8C to 16C.
  • Patch expanding layer: A rearrange operation then doubles the input resolution and reduces the feature dimension to one quarter of the input dimension.The example maps W_32 × H_32 × 16C to W_16 × H_16 × 4C.

3.6 Skip connection

Swin-Unet uses U-Net-like skip connections to fuse encoder multi-scale features with up-sampled decoder features, reducing spatial-information loss from down-sampling while preserving feature dimensionality.

  • 3.6 Skip connection: Skip connections concatenate shallow encoder features with deep up-sampled features to fuse multi-scale information and reduce spatial-information loss caused by down-sampling.A subsequent linear layer keeps the concatenated feature dimension equal to that of the up-sampled features.

4 Experiments

Swin-Unet is evaluated on Synapse multi-organ CT and ACDC cardiac MRI segmentation, achieving strong performance across modalities. Ablation studies examine up-sampling, skip connections, input size, and model scale, while noting limitations of ImageNet initialization and 2D inputs.

  • Synapse multi-organ segmentation: On Synapse, Swin-Unet achieves 79.13% DSC and 21.55% HD, outperforming the compared state-of-the-art methods.The dataset uses 18 training and 12 testing cases and evaluates eight abdominal organs with average DSC and HD.
  • ACDC cardiac segmentation: On ACDC, Swin-Unet achieves 90.00% average DSC on cardiac MRI segmentation, indicating generalization across imaging modalities.ACDC contains 70 training, 10 validation, and 20 testing samples with LV, RV, and MYO labels.
  • Ablation studies: The patch expanding layer outperforms bilinear interpolation and transposed convolution for decoder up-sampling on Synapse.It performs up-sampling while increasing feature dimensions, corresponding to the encoder’s patch merging layer.
  • Ablation studies: Segmentation performance increases as the number of skip connections rises from zero to three across the 1/4, 1/8, and 1/16 resolution scales.The three-scale configuration is selected to improve model robustness.
  • Ablation studies: Increasing input resolution from 224 × 224 to 384 × 384 slightly improves accuracy but increases computational load, whereas scaling the model adds cost with little performance gain.The Tiny-based model is adopted as an accuracy-speed trade-off.

5 Conclusion

Swin-Unet is a pure Transformer-based U-shaped encoder-decoder for medical image segmentation, using Swin Transformer blocks for feature representation and long-range semantic interaction. Experiments on multi-organ and cardiac segmentation tasks show excellent performance and generalization ability.

  • 5 Conclusion: Swin-Unet introduces a pure Transformer-based U-shaped encoder-decoder for medical image segmentation.The model is presented as a novel architecture.
  • 5 Conclusion: Swin Transformer blocks provide feature representation and long-range semantic information interaction within the proposed architecture.The Swin Transformer block is used as the basic unit.
  • 5 Conclusion: Experiments on multi-organ and cardiac segmentation tasks demonstrate excellent performance and generalization ability.The conclusion reports these results across both task types.
Loading 2105.05537v1…