PaperlayerView exact revision on arXiv

Source-linked AI summary

UNETR: Transformers for 3D Medical Image Segmentation

Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Landman, Holger Roth, Daguang Xu

arXiv:2103.10504v3eess.IVcs.CVcs.LG

TL;DR

FCNNs struggle to learn long-range spatial dependencies because convolutional receptive fields are local. UNETR reformulates 3D segmentation as sequence-to-sequence prediction with a transformer encoder and CNN decoder, achieving state-of-the-art BTCV results and outperforming competing MSD approaches.

  • Problem

    FCNNs have limited ability to learn long-range spatial dependencies because convolutional layers use localized receptive fields.

  • Method

    UNETR reformulates volumetric segmentation as sequence-to-sequence prediction using a transformer encoder over embedded 3D patches and a CNN decoder connected through skip connections.

  • Results

    UNETR achieves new state-of-the-art performance on BTCV and outperforms competing approaches for brain tumor and spleen segmentation on MSD.

  • Takeaways & Limitations

    UNETR captures global and local dependencies for volumetric medical image segmentation across CT and MRI tasks.

Abstract

from arXiv · show

Fully Convolutional Neural Networks (FCNNs) with contracting and expanding paths have shown prominence for the majority of medical image segmentation applications since the past decade. In FCNNs, the encoder plays an integral role by learning both global and local features and contextual representations which can be utilized for semantic output prediction by the decoder. Despite their success, the locality of convolutional layers in FCNNs, limits the capability of learning long-range spatial dependencies. Inspired by the recent success of transformers for Natural Language Processing (NLP) in long-range sequence learning, we reformulate the task of volumetric (3D) medical image segmentation as a sequence-to-sequence prediction problem. We introduce a novel architecture, dubbed as UNEt TRansformers (UNETR), that utilizes a transformer as the encoder to learn sequence representations of the input volume and effectively capture the global multi-scale information, while also following the successful "U-shaped" network design for the encoder and decoder. The transformer encoder is directly connected to a decoder via skip connections at different resolutions to compute the final semantic segmentation output. We have validated the performance of our method on the Multi Atlas Labeling Beyond The Cranial Vault (BTCV) dataset for multi-organ segmentation and the Medical Segmentation Decathlon (MSD) dataset for brain tumor and spleen segmentation tasks. Our benchmarks demonstrate new state-of-the-art performance on the BTCV leaderboard. Code: https://monai.io/research/unetr

1. Introduction

Image segmentation is an integral first step in quantitative medical image analysis, and U-shaped encoder-decoder networks have achieved state-of-the-art results across medical semantic segmentation tasks.

  • Image segmentation is often the first step in quantitative medical image analysis.
  • FCNNs, particularly U-shaped encoder-decoder architectures, have achieved state-of-the-art results in medical semantic segmentation.
  • The introduction frames medical image segmentation as central to analyzing anatomical structures quantitatively.

NVIDIA

UNETR applies a transformer encoder directly to embedded 3D volumes and connects it to a CNN decoder through multi-resolution skip connections. It targets long-range dependencies in volumetric segmentation and achieves strong results across BTCV and MSD benchmarks.

  • NVIDIA: The proposed model is motivated by the limited long-range dependency learning of convolutional networks with localized receptive fields.This limitation can produce sub-optimal segmentation for structures with variable shapes and scales.
  • NVIDIA: UNETR reformulates 3D segmentation as a 1D sequence-to-sequence prediction problem using a transformer encoder.The encoder learns contextual information from embedded input patches.
  • NVIDIA: Unlike approaches that place transformers in a decoder or bottleneck, UNETR connects transformer representations directly to a CNN-based decoder through skip connections.The CNN decoder is used because transformers do not properly capture localized information.
  • NVIDIA: UNETR is validated on volumetric CT and MRI segmentation tasks using the BTCV and MSD datasets.The evaluated tasks include BTCV multi-organ segmentation and MSD brain tumor and spleen segmentation.
  • NVIDIA: UNETR achieves new state-of-the-art performance on both Standard and Free BTCV leaderboard competitions and outperforms competing approaches on MSD tasks.
  • NVIDIA: The architecture directly uses embedded 3D volumes to capture long-range dependencies and combines representations across resolutions for segmentation prediction.

2. Related Work

Prior work spans CNN-based volumetric segmentation, multi-scale frameworks, vision transformers, and transformer-assisted 2D and 3D segmentation. UNETR differs by directly connecting transformer-encoded representations to its decoder.

  • CNN-based Segmentation Networks: CNN-based methods have achieved state-of-the-art results in diverse 2D and 3D medical image segmentation tasks.Volumetric approaches include tri-planar architectures and methods that directly process full 3D volumes.
  • CNN-based Segmentation Networks: Multi-scale and assembled frameworks extract varying-resolution information to address spatial-context and low-resolution challenges in 3D segmentation.
  • Vision Transformers: Vision transformer research introduced pure and hierarchical transformer models for computer vision, including architectures with changing feature resolutions and spatial embeddings.
  • Vision Transformers: Transformer-based semantic segmentation methods have used CNN decoders, bottleneck transformers, parallel CNN-transformer streams, and axial attention, primarily for 2D tasks.
  • Vision Transformers: Existing 3D approaches commonly combine a CNN backbone or encoder with a transformer in the bottleneck or encoded-representation processing stage.

3. Methodology

UNETR reformulates 3D segmentation around a transformer encoder that processes volumetric patches as a sequence, then combines multi-resolution representations with a convolutional decoder. The resulting U-shaped architecture produces voxel-wise semantic predictions using skip connections and a combined soft Dice/cross-entropy loss.

  • 3.1. Architecture: UNETR divides a 3D input volume into flattened, uniform, non-overlapping patches to form a 1D sequence.The sequence length is N=(H×W×D)/P^3, where P is the patch resolution.
  • 3.1. Architecture: A linear layer projects patches into a constant K-dimensional embedding space, which receives a learnable positional embedding before transformer processing.The class token is omitted because the backbone targets semantic segmentation rather than classification.
  • 3.1. Architecture: The encoder uses stacked transformer blocks containing multi-head self-attention and multilayer perceptron sublayers with normalization and GELU activations.Each self-attention sublayer uses parallel heads to learn query-key-value relationships across the input sequence.
  • 3.1. Architecture: Representations from transformer layers 3, 6, 9, and 12 are reshaped and projected into feature maps, then merged with the decoder through skip connections.Consecutive 3×3×3 convolutional and normalization layers project the representations into the input space.
  • 3.1. Architecture: The decoder progressively upsamples and concatenates transformer features until the original input resolution, where a 1×1×1 convolution with softmax produces voxel-wise predictions.At the bottleneck, a deconvolution first doubles resolution before subsequent concatenation and upsampling stages.
  • 3.2. Loss Function: The training objective combines soft Dice loss with cross-entropy loss computed voxel-wise across classes.The probability output and one-hot encoded ground truth are defined for each class at each voxel.

4. Experiments

Experiments evaluate UNETR on BTCV abdominal CT and MSD brain-tumor MRI and spleen CT tasks using Dice and 95% Hausdorff Distance. UNETR achieves leading quantitative performance and shows improved qualitative delineation across organs and tumor regions.

  • Datasets: BTCV contains 30 abdominal CT subjects with annotations for 13 organs, while MSD provides 484 multimodal MRI cases for brain tumors and 41 CT volumes for spleen segmentation.Brain-tumor segmentation uses three classes with four-channel input; spleen segmentation is binary with one-channel input.
  • Evaluation Metrics: Dice score and 95% Hausdorff Distance (HD) are used to evaluate segmentation accuracy.The 95% HD uses the 95th percentile of surface distances, reducing the influence of a very small subset of outliers.
  • Quantitative Evaluations: UNETR achieves an overall average Dice score of 0.899 in the BTCV Free Competition, exceeding the second-, third-, and fourth-ranked methods by 1.238%, 1.696%, and 5.269%.These comparisons are reported from the BTCV leaderboard.
  • Quantitative Evaluations: UNETR reaches 85.3% average Dice across BTCV organs and exceeds the second-best baselines by 1.043%, 0.830%, and 2.125% for spleen, liver, and stomach.For smaller organs, the reported margins are 6.382% for gallbladder and 6.772% for adrenal glands.
  • Quantitative Evaluations: On MSD, UNETR exceeds the closest brain-tumor baseline by 1.5% averaged across semantic classes and the best spleen competitor by at least 1.0% Dice.The brain-tumor advantage is particularly noted for the tumor-core subregion.
  • Qualitative Results: Qualitative comparisons show improved boundaries and spatial-context handling for abdominal organs, alongside finer-grained tumor-detail capture.Examples include separating liver from stomach, detecting kidney and adrenal glands against surrounding tissue, and delineating kidney from spleen.

5. Discussion

UNETR outperforms CNN- and transformer-based models by capturing global and local dependencies, with particularly strong results on small and difficult anatomical structures. Its decoder design is also evaluated against alternative upsampling architectures.

  • UNETR demonstrates superior segmentation performance over CNN- and transformer-based models by capturing both global and local dependencies.Qualitative examples include accurate long-range segmentation of the pancreas tail.
  • UNETR achieves new state-of-the-art benchmarks on the BTCV leaderboard.
  • Table 3 compares Naive UpSampling, Progressive UpSampling, and Multi-scale Aggregation decoder architectures.The supplied caption identifies the three alternatives evaluated for decoder performance.
  • UNETR significantly improves segmentation of adrenal glands and uniquely detects their branches among compared models.
  • For low-contrast gallbladder and portal-vein tissues, UNETR produces clear connected boundaries.

6. Ablation

Ablations examine decoder choice, patch resolution, and computational complexity. Smaller patches improve Dice performance but increase memory demands, while UNETR maintains moderate complexity and fast inference relative to competing models.

  • Decoder Choice: Alternative NUP, PUP, and MLA decoders yield sub-optimal performance, although MLA marginally outperforms NUP and PUP.
  • Patch Resolution: Decreasing patch resolution from 32 to 16 improves average Dice by 1.1% for spleen and 0.8% for brain segmentation.
  • Patch Resolution: Lower patch resolution increases sequence length and memory consumption because memory scales inversely with the cube of resolution.
  • Patch Resolution: Table 4 evaluates how patch resolution affects segmentation performance.
  • Model and Computational Complexity: UNETR uses 92.58M parameters and 41.19G FLOPs, providing moderate model complexity.
  • Model and Computational Complexity: UNETR outperforms listed CNN- and transformer-based models while having comparable complexity and the second-lowest averaged inference time after nnUNet.

7. Conclusion

UNETR reformulates volumetric medical image segmentation as sequence-to-sequence prediction with a transformer encoder for long-range and multi-scale context. It achieves state-of-the-art or superior performance across BTCV and MSD tasks in CT and MRI.

  • UNETR reformulates volumetric medical image segmentation as a 1D sequence-to-sequence prediction problem.
  • Its transformer encoder is designed to learn long-range dependencies and global contextual representations at multiple scales.
  • UNETR achieves new state-of-the-art performance on BTCV and outperforms competing approaches for brain tumor and spleen segmentation on MSD.
  • The authors report validation across volumetric CT and MRI segmentation tasks.
Loading 2103.10504v3…