PaperlayerView exact revision on arXiv

Source-linked AI summary

D-Former: A U-shaped Dilated Transformer for 3D Medical Image Segmentation

Yixuan Wu, Kuanlun Liao, Jintai Chen, Jinhong Wang, Danny Z. Chen, Honghao Gao, Jian Wu

arXiv:2201.00462v2cs.CVcs.AI

TL;DR

3D medical-image Transformers offer long-range modeling but face high computational costs, while slice-wise processing can lose volumetric context. D-Former addresses this with alternating local and dilated global attention, dynamic position encoding, and a U-shaped architecture, achieving strong segmentation performance with lower computational cost and training from scratch.

  • Problem

    Transformer segmentation models can lose 3D contextual information through slice-wise processing and incur high costs from global self-attention on volumetric images.

  • Method

    D-Former is a 3D U-shaped hierarchical encoder-decoder that alternates local attention with dilated global attention and uses dynamic position encoding.

  • Results

    D-Former achieves state-of-the-art segmentation performance on Synapse and ACDC while reducing model parameters and FLOPs.

  • Takeaways & Limitations

    The design supports direct 3D segmentation with long-range dependency modeling at reduced computational cost.

  • Takeaways & Limitations

    Transformer pre-training is commonly important, but D-Former is trained from scratch rather than initialized with pre-trained weights.

Abstract

from arXiv · show

Computer-aided medical image segmentation has been applied widely in diagnosis and treatment to obtain clinically useful information of shapes and volumes of target organs and tissues. In the past several years, convolutional neural network (CNN) based methods (e.g., U-Net) have dominated this area, but still suffered from inadequate long-range information capturing. Hence, recent work presented computer vision Transformer variants for medical image segmentation tasks and obtained promising performances. Such Transformers model long-range dependency by computing pair-wise patch relations. However, they incur prohibitive computational costs, especially on 3D medical images (e.g., CT and MRI). In this paper, we propose a new method called Dilated Transformer, which conducts self-attention for pair-wise patch relations captured alternately in local and global scopes. Inspired by dilated convolution kernels, we conduct the global self-attention in a dilated manner, enlarging receptive fields without increasing the patches involved and thus reducing computational costs. Based on this design of Dilated Transformer, we construct a U-shaped encoder-decoder hierarchical architecture called D-Former for 3D medical image segmentation. Experiments on the Synapse and ACDC datasets show that our D-Former model, trained from scratch, outperforms various competitive CNN-based or Transformer-based segmentation models at a low computational cost without time-consuming per-training process.

1 Introduction

D-Former addresses the loss of 3D context and high computational cost that limit Transformer-based segmentation of volumetric medical images. It combines local and dilated global attention with dynamic position encoding in a 3D U-shaped architecture, achieving strong performance with reduced complexity.

  • 3D medical-image Transformers may lose volumetric context by processing slices independently and incur high costs during global self-attention.
  • D-Former directly processes 3D images and predicts volumetric segmentation masks using a hierarchical U-shaped encoder-decoder with skip connections.
  • Alternating local and global scope modules capture fine and long-range information, while dilated global attention expands interaction scope without increasing the patches involved.
  • Dynamic position encoding learns position information from the input, helping alleviate permutation-equivalence in Transformers.
  • D-Former outperforms state-of-the-art segmentation methods across CT and MRI domains with smaller model sizes and fewer FLOPs.

2 Related Work

Prior segmentation research progressed from CNN-based architectures to vision Transformers, but the reviewed approaches motivate D-Former's dilated self-attention design for efficient long-range interaction.

  • Dilated convolution enlarges receptive fields by inserting gaps between kernel elements, inspiring D-Former's patch-skipping self-attention.
  • Transformers were adapted from sequence modeling to vision because self-attention can model relationships among spatial patches.
  • Medical segmentation Transformers commonly supplement convolutional feature extraction with Transformer-based encoding.
  • Figure 1 presents D-Former as a multi-stage architecture composed of dynamic position encoding, local scope modules, and global scope modules.

3.1 The Overall Architecture

D-Former uses a hierarchical U-shaped encoder-decoder that alternates local and global attention modules across four blocks in both pathways, with skip connections between matching levels.

  • The encoder contains four D-Former blocks separated by three down-sampling layers, while the decoder mirrors it with four blocks and three up-sampling layers.
  • The first, second, and fourth blocks contain one LSM and one GSM, whereas the third contains three of each arranged alternately.
  • Skip connections concatenate encoder and decoder feature maps at corresponding levels to transfer information across the hierarchy.
  • The architecture includes patch embedding at the input and patch expanding with a segmentation head at the output.

3.2 Patch Embedding & Patch Expending

Patch embedding converts a 3D medical image into a projected patch sequence for Transformer processing, while patch expanding restores spatial resolution for pixel-wise mask prediction.

  • The embedded feature map represents the 3D image as a sequence whose length equals the number of patches.
  • The final patch expanding layer restores the feature map to the original input size and feeds a segmentation head for pixel-wise masks.

3.3 D-Former Blocks

D-Former blocks alternate local and global self-attention modules to combine nearby and long-range patch interactions efficiently. Local attention partitions feature maps into adjacent units, while dilated global attention samples patches across units without increasing the number of attended patches.

  • 3.3 D-Former Blocks: Each D-Former block applies dynamic position encoding before alternating Local Scope Modules and Global Scope Modules for higher-level feature extraction.The blocks form the repeated processing units of the D-Former architecture.
  • 3.3.1 Local Scope Module and Global Scope Module: An LSM contains LayerNorm, LS-MSA, LayerNorm, and an MLP with residual connections, while a GSM replaces LS-MSA with GS-MSA.The two modules therefore share the surrounding normalization, perceptron, and residual structure.
  • 3.3.2 Local Scope Multi-head Self-Attention (LS-MSA): LS-MSA performs self-attention within non-overlapping local units, reducing complexity to linear in the total number of feature-map patches.Each unit contains ud × uh × uw patches, typically far fewer than the dhw patches in the whole feature map.
  • 3.3.2 Local Scope Multi-head Self-Attention (LS-MSA): In the local-scope design, self-attention operates on adjacent patches grouped within each unit.The figure illustrates a 6 × 6 × 6 feature map with 3 × 3 × 3 units.
  • 3.3.3 Global Scope Multi-head Self-Attention (GS-MSA): LS-MSA lacks cross-unit interaction, so GS-MSA forms units by selecting every g-th patch along each dimension in a dilated pattern.The construction continues until all patches are assigned to global units.
  • 3.3.3 Global Scope Multi-head Self-Attention (GS-MSA): The global-scope construction preserves the same number of units as the local modules by relating feature-map dimensions to unit sizes and sampling distances.The sampling distances gd, gh, and gw determine the nearest-patch spacing along depth, height, and width.

3.4 Down-sampling & Up-sampling

D-Former uses hierarchical down-sampling between encoder blocks and mirrored up-sampling in the decoder to fuse features at multiple scales.

  • Down-sampling: Each down-sampling layer concatenates 2 × 2 × 2 neighboring patches, reducing the patch count by 8 times before a fully connected layer reduces channels by 4 times.This channel adjustment allows the channel size to double across stages.
  • Up-sampling: The decoder reverses this process with four up-sampling layers that enlarge low-resolution feature maps and reduce channel dimensionality.The architecture uses this multi-scale processing to produce segmentation outputs.

3.5 The Dynamic Position Encoding Block

The dynamic position encoding block learns position information with 3D depth-wise convolution and adds it to each block’s input feature map.

  • Dynamic Position Encoding: Each D-Former block applies 3D depth-wise convolution once to input feature maps to learn position information.Depth-wise convolution uses one filter per input channel, reducing computational cost relative to common convolution.
  • Dynamic Position Encoding: The position-encoding operation resizes the input as needed, applies depth-wise convolution, resizes the result, and adds it residually to the original input.The output feature maps are embedded with the learned position information.
  • Dynamic Position Encoding: Because the position information is learned dynamically from the input, the block avoids the fixed-patch-number limitation of earlier approaches.The convolution’s translation invariance is also reported to improve stability and generalization performance.

4 Experiments

Experiments evaluate D-Former on Synapse and ACDC using average DSC, with comparisons spanning accuracy, qualitative segmentation, ablations, and computational complexity. D-Former achieves strong results across CT and MRI while being trained from scratch and using fewer FLOPs than several direct 3D segmentation models.

  • Experimental Setup: The experiments use average Dice Similarity Coefficient (DSC) to evaluate eight-organ Synapse CT segmentation and ACDC cardiac MRI segmentation.Synapse uses 18 training and 12 testing scans; ACDC uses 100 annotated cases split 7:1:2.
  • Experimental Setup: D-Former is trained from scratch with randomly initialized weights, despite pre-training commonly being important for Transformer-based models.The experiments use an NVIDIA GeForce RTX 3090 GPU, batch size 2 for training, and batch size 1 for inference.
  • Quantitative Results: On Synapse, D-Former improves average DSC by 1.43% over nnFormer and 6.87% over MISSFormer.Accuracy improves for almost every organ class except the pancreas and stomach.
  • Quantitative Results: On ACDC, D-Former achieves a best average DSC of 92.29%, improving by 0.51% over nnFormer and 4.39% over MISSFormer.It also reports 4.74% and 5.54% improvements over R50 U-Net and R50 Att-UNet, respectively, and the best average DSC for RV, Myo, and LV.
  • Ablation Studies: Ablations show that GSM exceeds LSM-only processing by 0.66% average DSC, GS-MSA exceeds shifted-window attention by 1.33%, and DPE exceeds APE, SPE, and RPE by 4.05%, 2.79%, and 2.48%.These comparisons assess global interaction and dynamic position encoding choices.

5 Conclusions

D-Former is a U-shaped 3D segmentation framework built on a dilated Transformer, combining efficient self-attention with long-range dependency modeling and dynamic position encoding. It reduces model complexity while achieving state-of-the-art segmentation performance on Synapse and ACDC.

  • D-Former combines a U-shaped encoder-decoder with a dilated Transformer for 3D medical image segmentation.
  • Its local scope module reduces the number of patches used in self-attention, while its global scope module captures long-range dependencies through dilated attention.
  • Dynamic position encoding enables the model to learn vital position information flexibly within input sequences.
  • The model reduces parameters and FLOPs while attaining state-of-the-art semantic segmentation performance on the Synapse and ACDC datasets.
Loading 2201.00462v2…