PaperlayerView exact revision on arXiv

Source-linked AI summary

nnFormer: Interleaved Transformer for Volumetric Segmentation

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

arXiv:2109.03201v6cs.CV

TL;DR

Medical segmentation methods often treat transformers as auxiliary modules on convolutional representations, leaving their long-range modeling capacity incompletely exploited. nnFormer makes transformer processing central through an interleaved 3D architecture with local and global volume attention and skip attention. Across three datasets it outperforms prior transformer-based models and is better than nnUNet on HD95 while remaining comparable on DSC; averaging nnFormer and nnUNet predictions further improves performance.

  • Problem

    Most transformer-based medical segmentation approaches use transformers as assisted modules on convolutional backbones, potentially limiting their ability to model long-term dependencies with spatial and hierarchical representations.

  • Method

    nnFormer is a 3D volumetric segmentation transformer using interleaved convolution and self-attention, local and global volume-based attention, and skip attention in U-Net-like connections.

  • Results

    nnFormer significantly outperforms previous transformer-based models on three datasets and is better than nnUNet on HD95 while producing comparable DSC results.

  • Takeaways & Limitations

    Averaging nnFormer and nnUNet predictions produces large performance gains, indicating that the two models are highly complementary for ensembling.

Abstract

from arXiv · show

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.

I. INTRODUCTION

nnFormer addresses the limited use of transformers in medical segmentation by making them a central part of a 3D hybrid architecture. It combines interleaved convolution and self-attention, volume-based attention, and skip attention, with favorable comparisons against prior models and nnUNet.

  • Research gap: Most prior medical segmentation methods use transformers as auxiliary modules on convolutional backbones, limiting their exploitation of long-term dependencies.Convolutional representations retain precise spatial information and hierarchical concepts, while shallow transformer components may not fully entangle these with long-range dependencies.
  • Core architecture: nnFormer interleaves convolution and self-attention so convolution encodes precise spatial information while transformer blocks model long-term dependencies across hierarchical scales.A convolutional embedding layer precedes interleaved transformer and convolutional down-sampling blocks.
  • Attention mechanisms: LV-MSA and GV-MSA construct feature pyramids and provide local and global 3D receptive fields with lower computational complexity than naive MSA.Skip attention also replaces concatenation or summation in U-Net-like skip connections and further improves segmentation results.
  • Results: nnFormer significantly improves over previous transformer-based medical segmentation models on three established datasets.The contribution statement reports improvements across the three datasets but gives no numerical values here.
  • Results: nnFormer achieves lower HD95 and slightly better DSC than nnUNet, while averaging their predictions greatly boosts overall performance.The authors attribute the ensemble gains to the models’ high complementarity.

II. RELATED WORK

Related work divides transformer-based medical image segmentation methods by whether their stems are primarily convolutional or transformer-based. These approaches range from convolution-assisted transformers to convolution-free and hierarchical transformer architectures.

  • Organization: The review categorizes hybrid medical segmentation architectures according to whether most of the stem is convolutional or transformer-based.
  • Convolution-based stem: Convolution-based methods use convolutional feature extractors before applying transformer modules to encode broader dependencies.TransUNet applies patch embedding to bottleneck feature maps, while related work adds attention and expansion mechanisms for diversified representations.
  • Transformer-based stem: Transformer-based methods include gated axial attention, convolution-free 3D transformers, U-shaped SwinUNet, and multi-scale DS-TransUNet.These models progressively extend transformer backbones for volumetric, hierarchical, or multi-scale segmentation.
  • nnFormer architecture: nnFormer retains a U-shaped encoder–bottleneck–decoder structure with embedding, transformer, sampling, and patch-expanding components.Its architecture can vary slightly with input scan size.

B. Encoder

The encoder combines convolutional embedding and down-sampling with 3D local-volume transformer blocks. This design preserves spatial detail, builds hierarchical representations, and efficiently models interactions within and across local volumes.

  • Embedding layer: The embedding layer uses successive small-kernel convolutions to encode pixel-level spatial information before transformer processing.The embedding block contains four convolutional layers with kernel size 3, with GELU and layer normalization after all but the last.
  • Local attention: LV-MSA computes self-attention within 3D local volumes, while shifted LV-MSA introduces interactions between neighboring volumes.Each local transformer block contains two successive transformer layers, with the second using shifted local volumes.
  • Architecture overview: The encoder’s attention mechanisms are organized alongside convolutional embedding and sampling components in nnFormer’s overall architecture.Figure 2 defines K as convolutional kernel size, DK as deconvolutional kernel size, S as stride, and Norm as layer normalization.
  • Down-sampling: Strided convolution replaces patch merging for down-sampling, producing hierarchical multi-scale representations; stride 2 is typical but may be reduced when slice counts are limited.Using an excessive stride along a limited slice dimension can be harmful.

C. Bottleneck

The bottleneck applies global 3D self-attention after down-sampling has reduced the spatial dimensions. This makes broader receptive fields computationally feasible before features are passed to the decoder.

  • Global attention: GV-MSA extends vision-transformer attention to 3D and provides a larger receptive field than local-volume attention.Its computational cost grows substantially when the full volume dimensions are much larger than local-volume dimensions.
  • Computational placement: Down-sampling reduces hwd in the bottleneck until it is similar in scale to the local-volume token count, enabling GV-MSA there.The encoder instead uses local attention to handle large inputs efficiently.
  • Bottleneck design: nnFormer uses three global transformer blocks, corresponding to six GV-MSA layers, in the bottleneck to provide receptive field to the decoder.

D. Decoder

The decoder upsamples low-resolution features and uses skip attention to integrate encoder representations before producing segmentation masks.

  • Strided deconvolution upsamples decoder feature maps, which are merged with encoder representations through skip attention.This combines semantic and fine-grained information in the decoder.
  • The final patch-expanding block uses deconvolution to produce the mask predictions.
  • Skip attention replaces concatenation or summation in conventional convolutional skip connections with an attention mechanism.The mechanism is inspired by attention used in machine translation.
  • Encoder features are linearly projected and split into key and value matrices for decoder-side attention.The projected representations are denoted K_l* and V_l*.
  • The upsampled decoder feature map supplies the query for LV/GV-MSA with the encoder-derived keys and values.The decoder applies the same volume-based self-attention operation used elsewhere in nnFormer.

IV. EXPERIMENTS

Experiments evaluate nnFormer across brain tumor, multi-organ, and cardiac segmentation using established datasets, train/validation/test splits, and HD95 and Dice-based metrics.

  • Experiments cover brain tumor segmentation, Synapse multi-organ segmentation, and ACDC automated cardiac diagnosis.Each experiment is repeated ten times, with average results and p-values reported.
  • Brain tumor segmentation: The brain tumor task uses 484 four-channel MRI images and evaluates whole tumor, enhancing tumor, and tumor core regions.The data are split into 80% training, 15% validation, and 5% testing, with HD95 and Dice score as metrics.
  • The reported transformer comparisons use tables that identify best and second-best results and calculate p-values against the strongest baseline.Brain tumor, Synapse, and ACDC use separate comparison tables with their stated metrics.
  • Multi-organ segmentation: Synapse contains 30 abdominal CT cases, with 18 used for training and 12 for testing across eight abdominal organs.Performance is evaluated using HD95 and DSC.
  • Automated cardiac diagnosis: ACDC includes 100 patients and segments the right ventricle, left-ventricle myocardium, and left ventricle.The split is 70 training, 10 validation, and 20 testing samples, with HD95 and Dice score evaluated.

A. Implementation details

Implementation uses fixed preprocessing and augmentation, deep supervision across decoder resolutions, and convolutional network configurations designed for efficient segmentation.

  • Training uses Python 3.6, PyTorch 1.8.1, Ubuntu 18.04, and a single NVIDIA 2080 GPU with 11GB memory.Optimization uses SGD with momentum 0.99, weight decay 3e-5, and summed cross-entropy and Dice losses.
  • The learning rate follows poly decay: lr = initial lr × (1 − epoch id / max epoch)^0.9.
  • Images are resampled to a common target spacing before ordered augmentation with rotation, scaling, noise, blur, intensity, resolution, gamma, and mirroring transformations.
  • Deep supervision: Deep supervision applies cross-entropy and Dice losses to decoder outputs at three resolutions after down-sampling ground-truth masks to match each prediction.
  • Deep supervision: Loss weights halve with each decrease in resolution and are finally normalized to 1.The implementation sets α2 = α1 / 2 as part of this weighting scheme.
  • Network configurations: Compared with nnUNet, nnFormer achieves better segmentation results using smaller-sized input patches.

B. Comparison with transformer-based methodologies

Across brain tumor, Synapse, and ACDC tasks, nnFormer outperforms transformer-based methods, with particularly strong HD95 and DSC results; it also complements nnUNet in ensembling.

  • Multi-organ segmentation (Synapse): On Synapse, nnFormer surpasses LeViT-UNet-384s and TransUNet by over 6 mm in average HD95 and 2 percentage points in average DSC.It achieves the highest DSC in six organs, including the pancreas and stomach.
  • Comparison with nnUNet: nnFormer and nnUNet are highly complementary because simply averaging their predictions can greatly improve overall performance.The comparison and ablation tables report these model-combination and module investigations.
  • Automated cardiac diagnosis (ACDC): On ACDC, nnFormer surpasses LeViT-UNet-384s in all classes and by nearly 1.7 percentage points in average DSC.
  • Statistical significance: Across all three datasets, p-values are smaller than 1e-2 for both HD95 and DSC against the best-performing baseline.The reported tests provide strong evidence against the null hypothesis of no nnFormer advantage.

C. Comparison with nnUNet and Discussion

nnFormer consistently improves boundary-sensitive HD95 performance over nnUNet across all three public datasets, while achieving higher DSC on two datasets. Their prediction diversity also makes simple averaging highly effective.

  • nnFormer achieves lower HD95 than nnUNet on all three public datasets and higher DSC on two of three datasets.
  • 11 of 16 class-specific HD95 results favor nnFormer, compared with 9 of 16 class-specific DSC results.
  • HD95 differences are statistically significant on all three datasets, whereas DSC is comparable on tumor and multi-organ segmentation and significant on ACDC.
  • Averaging nnFormer and nnUNet predictions produces large performance gains, including the best tumor-segmentation results across all classes and nearly 30% HD95 improvement on Synapse.

D. Ablation study

The ablation study on ACDC with DSC as the default metric shows incremental gains from convolutional embedding and down-sampling, global attention, and larger receptive fields.

  • The ablation experiments use ACDC and DSC as the default evaluation setting.
  • The baseline combining LV-MSA, patch merging, and embedding layers already achieves higher average DSC than LeViT-UNet-38.
  • The proposed convolutional embedding layer improves average DSC by approximately 0.4 percent over the baseline embedding.
  • Replacing patch merging with convolutional down-sampling further improves overall performance by 0.3 percent.
  • Replacing bottleneck LV-MSA with GV-MSA improves performance by 0.2 percent, indicating benefits from a larger receptive field.

E. Visualization of segmentation results

Visual comparisons show nnFormer produces fewer false positives than UNETR and retains advantages over nnUNet in boundary and shape delineation, while nnUNet remains highly competitive.

  • nnFormer greatly reduces false-positive predictions compared with UNETR across visualizations from three public datasets.
  • On ACDC, nnFormer avoids right-ventricle predictions outside the myocardium in an example where UNETR produces many such errors.
  • nnUNet remains much better than UNETR in nearly all displayed examples, demonstrating its competitive segmentation quality.
  • nnFormer better captures kidney and stomach shapes than nnUNet on Synapse and preserves more of a non-enhancing tumor than nnUNet in a brain-tumor example.
  • The conclusion attributes nnFormer’s representation to interleaved convolution and self-attention, with convolution encoding precise spatial information and building hierarchical concepts.
Loading 2109.03201v6…