Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

nnFormer: Interleaved Transformer for Volumetric Segmentation

Hong-Yu Zhou, Jiansen Guo, Yinghao Zhang, Lequan Yu, Liansheng Wang, Yizhou Yu

TL;DR

nnFormer introduces a 3D transformer for volumetric medical image segmentation that interleaves convolution and self-attention, leveraging local-volume and global-volume self-attention to build multi-scale representations. The model uses LV-MSA and GV-MSA to capture local and global dependencies and employs skip attention to fuse encoder-decoder features, achieving superior HD95 and competitive DSC across three public datasets. It outperforms state-of-the-art transformer-based methods and shows complementary strengths when ensembled with nnUNet, underscoring the practical benefits of combining convolutional inductive bias with self-attention in medical image segmentation.

Abstract

Transformer, the model of choice for natural language processing, has drawn scant attention from the medical imaging community. Given the ability to exploit long-term dependencies, transformers are promising to help atypical convolutional neural networks to overcome their inherent shortcomings of spatial inductive bias. However, most of recently proposed transformer-based segmentation approaches simply treated transformers as assisted modules to help encode global context into convolutional representations. To address this issue, we introduce nnFormer, a 3D transformer for volumetric medical image segmentation. nnFormer not only exploits the combination of interleaved convolution and self-attention operations, but also introduces local and global volume-based self-attention mechanism to learn volume representations. Moreover, nnFormer proposes to use skip attention to replace the traditional concatenation/summation operations in skip connections in U-Net like architecture. Experiments show that nnFormer significantly outperforms previous transformer-based counterparts by large margins on three public datasets. Compared to nnUNet, nnFormer produces significantly lower HD95 and comparable DSC results. Furthermore, we show that nnFormer and nnUNet are highly complementary to each other in model ensembling.

nnFormer: Interleaved Transformer for Volumetric Segmentation

TL;DR

nnFormer introduces a 3D transformer for volumetric medical image segmentation that interleaves convolution and self-attention, leveraging local-volume and global-volume self-attention to build multi-scale representations. The model uses LV-MSA and GV-MSA to capture local and global dependencies and employs skip attention to fuse encoder-decoder features, achieving superior HD95 and competitive DSC across three public datasets. It outperforms state-of-the-art transformer-based methods and shows complementary strengths when ensembled with nnUNet, underscoring the practical benefits of combining convolutional inductive bias with self-attention in medical image segmentation.

Abstract

Transformer, the model of choice for natural language processing, has drawn scant attention from the medical imaging community. Given the ability to exploit long-term dependencies, transformers are promising to help atypical convolutional neural networks to overcome their inherent shortcomings of spatial inductive bias. However, most of recently proposed transformer-based segmentation approaches simply treated transformers as assisted modules to help encode global context into convolutional representations. To address this issue, we introduce nnFormer, a 3D transformer for volumetric medical image segmentation. nnFormer not only exploits the combination of interleaved convolution and self-attention operations, but also introduces local and global volume-based self-attention mechanism to learn volume representations. Moreover, nnFormer proposes to use skip attention to replace the traditional concatenation/summation operations in skip connections in U-Net like architecture. Experiments show that nnFormer significantly outperforms previous transformer-based counterparts by large margins on three public datasets. Compared to nnUNet, nnFormer produces significantly lower HD95 and comparable DSC results. Furthermore, we show that nnFormer and nnUNet are highly complementary to each other in model ensembling.

Paper Structure

This paper contains 14 sections, 8 equations, 4 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Overview
  5. Encoder
  6. Bottleneck
  7. Decoder
  8. Experiments
  9. Implementation details
  10. Comparison with transformer-based methodologies
  11. Comparison with nnUNet and Discussion
  12. Ablation study
  13. Visualization of segmentation results
  14. Conclusion

Figures (4)

  • Figure 1: The interleaved stem used in the encoder of nnFormer.
  • Figure 2: Architecture of nnFormer. In (a), we show the overall architecture of nnFormer. In (b), we present more details of the embedding layers on three publicly available datasets. In (c), (d), (e), we display how to implement the down-sampling, up-sampling and expanding layers, respectively. In practice, the architecture may slightly vary depending on the input scan size. In (b)-(e), K denotes the convolutional kernel size, DK stands for the deconvolutional kernel size and S represents the stride. Norm refers to the layer normalization strategy.
  • Figure 3: Three types of attention mechanism in nnFormer. Norm denotes the layer normalization method. MLP is the abbreviation for multi-layer perceptron, which is a two-layer neural network in practice.
  • Figure 4: Visualization of segmentation results on three well-established datasets. We mainly compare nnFormer against nnUNet and UNETR. In addition to segmentation results, we also provide ground truth masks for better comparison.