[Paper Review] UNETR: Transformers for 3D Medical Image Segmentation

Ali Hatamizadeh, Tang, Yucheng+6 more|arXiv (Cornell University)|Mar 18, 2021

Radiomics and Machine Learning in Medical Imaging52 references210 citations

TL;DR

UNETR uses a transformer encoder to process 3D medical volumes as a sequence of patches, connecting to a CNN-based decoder with skip connections for accurate 3D segmentation, achieving state-of-the-art results on BTCV and MSD datasets.

ABSTRACT

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

Motivation & Objective

• Motivate leveraging transformers to capture long-range 3D context in medical image segmentation.
• Propose a UNETR architecture that connects a transformer encoder directly to a CNN decoder via skip connections.
• Demonstrate effectiveness on BTCV multi-organ segmentation and MSD brain tumor and spleen segmentation datasets.

Proposed method

• Represent a 3D volume as non-overlapping patches and project to a K-dimensional embedding.
• Process the patch sequence with a ViT-B16 style transformer encoder (L=12, K=768, patch size 16^3).
• Add 1D positional embeddings and omit a class token since the task is semantic segmentation.
• Extract intermediate transformer representations (z3, z6, z9, z12), reshape to spatial tensors, and fuse with a CNN-based decoder via skip connections.
• Use 3x3x3 convolutions to project transformer features to the decoder at multiple resolutions; apply deconvolutions to upsample; final 1x1x1 conv with softmax for voxel-wise predictions.
• Train with a combined soft Dice and cross-entropy loss; employ patch-based sliding-window inference with 0.5 overlap.

Experimental results

Research questions

• RQ1 Can a transformer encoder trained on 3D patches capture long-range dependencies in volumetric medical images for segmentation?
• RQ2 Does connecting transformer-derived features to a CNN-based decoder via multi-resolution skip connections improve segmentation accuracy over CNN- or transformer-only baselines?
• RQ3 What is the impact of decoder design, patch resolution, and model size on segmentation performance in 3D medical images?

Key findings

• UNETR achieves state-of-the-art performance on BTCV in both Standard and Free Competitions.
• In MSD brain tumor and spleen segmentation, UNETR outperforms competing approaches, especially in small structures like gallbladder and adrenal glands.
• On BTCV, average Dice scores show clear improvement over baselines, with notable gains for small organs.
• On MSD, UNETR yields higher Dice scores across brain tumor subregions and spleen segmentation than strongest baselines.
• The model has about 92.58M parameters and 41.19G FLOPs, with competitive inference time (~12.08s avg) compared to other transformer-based methods.