[Paper Review] nnFormer: Interleaved Transformer for Volumetric Segmentation

Hong-Yu Zhou, Jiansen Guo+4 more|arXiv (Cornell University)|Sep 7, 2021

Radiomics and Machine Learning in Medical Imaging40 references261 citations

TL;DR

nnFormer is a 3D transformer for volumetric medical image segmentation that interleaves convolution and self-attention, uses local and global volume-based self-attention, and introduces skip attention to outperform nnUNet on multiple datasets.

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.

Motivation & Objective

• Motivate the use of transformers for volumetric medical image segmentation to capture long-range dependencies beyond local convolutional biases.
• Develop a hybrid stem that interleaves convolutional embedding with transformer blocks to preserve spatial precision and hierarchical representations.
• Propose local volume-based (LV-MSA) and global volume-based (GV-MSA) self-attention to balance efficiency and receptive field.
• Introduce skip attention to improve skip connections in a U-Net–like architecture.
• Demonstrate that nnFormer achieves superior segmentation performance over transformer-based baselines and complements nnUNet in model ensembles.

Proposed method

• Embed 3D medical images with a lightweight convolutional embedding layer to preserve pixel-level spatial information.
• Use interleaved encoder blocks combining LV-MSA (local 3D self-attention) and convolutional down-sampling to build multi-scale features.
• Employ GV-MSA in the bottleneck to provide large receptive fields with reduced compute.
• Decode with a symmetric structure using up-sampling and skip attention to fuse encoder and decoder features.
• Replace traditional skip connections with skip attention to improve information flow between encoder and decoder.
• Train with deep supervision at multiple decoder stages and a joint loss over multiple resolutions.

Experimental results

Research questions

• RQ1 Can a 3D transformer with interleaved convolutional and self-attention blocks improve volumetric segmentation over purely CNN-based or standard transformer-based models?
• RQ2 Do local and global volume-based self-attention (LV-MSA and GV-MSA) provide complementary benefits for learning multi-scale volumetric representations?
• RQ3 Does skip attention enhance information transfer in U-Net–like architectures for 3D segmentation?
• RQ4 How does nnFormer compare to nnUNet and other transformer-based models across brain tumor, multi-organ, and cardiac segmentation tasks?

Key findings

• nnFormer significantly improves over prior transformer-based methods on three public datasets.
• Compared with nnUNet, nnFormer achieves lower HD95 and comparable or better DSC.
• Combining nnFormer with nnUNet yields complementary gains when averaged ensemble.
• Ablation studies show the effectiveness of interleaved conv+transformer stems, LV-MSA/GV-MSA, and skip attention in improving segmentation performance.