Source-linked AI summary
Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images
Ali Hatamizadeh, Vishwesh Nath, Yucheng Tang, Dong Yang, Holger Roth, Daguang Xu
TL;DR
Brain tumor segmentation needs models that can capture long-range information beyond the limited receptive fields of convolutional networks. Swin UNETR addresses this with a hierarchical Swin transformer encoder and multiresolution CNN decoder, achieving strong BraTS 2021 validation performance and competitive testing results.
Problem
CNN-based segmentation has limited ability to model long-range information, which can hinder segmentation of tumors with variable sizes.
Method
Swin UNETR projects multimodal 3D MRI data into a token sequence, encodes it with a hierarchical shifted-window Swin transformer, and connects five resolutions to an FCNN decoder through skip connections.
Results
Swin UNETR outperformed competing approaches by 0.7% for ET, 0.6% for WT, 0.4% for TC, and 0.5% averaged across classes and folds.
Takeaways & Limitations
The model ranked among the top-performing methods across more than 2,000 BraTS 2021 validation submissions and was the first transformer-based model to place competitively in BraTS challenges.
Abstract
from arXiv · showhide
Semantic segmentation of brain tumors is a fundamental medical image analysis task involving multiple MRI imaging modalities that can assist clinicians in diagnosing the patient and successively studying the progression of the malignant entity. In recent years, Fully Convolutional Neural Networks (FCNNs) approaches have become the de facto standard for 3D medical image segmentation. The popular "U-shaped" network architecture has achieved state-of-the-art performance benchmarks on different 2D and 3D semantic segmentation tasks and across various imaging modalities. However, due to the limited kernel size of convolution layers in FCNNs, their performance of modeling long-range information is sub-optimal, and this can lead to deficiencies in the segmentation of tumors with variable sizes. On the other hand, transformer models have demonstrated excellent capabilities in capturing such long-range information in multiple domains, including natural language processing and computer vision. Inspired by the success of vision transformers and their variants, we propose a novel segmentation model termed Swin UNEt TRansformers (Swin UNETR). Specifically, the task of 3D brain tumor semantic segmentation is reformulated as a sequence to sequence prediction problem wherein multi-modal input data is projected into a 1D sequence of embedding and used as an input to a hierarchical Swin transformer as the encoder. The swin transformer encoder extracts features at five different resolutions by utilizing shifted windows for computing self-attention and is connected to an FCNN-based decoder at each resolution via skip connections. We have participated in BraTS 2021 segmentation challenge, and our proposed model ranks among the top-performing approaches in the validation phase. Code: https://monai.io/research/swin-unetr
1 Introduction
Brain tumor segmentation supports diagnosis, surgical planning, progression analysis, and life-expectancy prediction. CNN-based methods perform strongly but have limited ability to model long-range dependencies, motivating transformer-based architectures such as Swin UNETR.
- Brain tumor segmentation supports volumetric progression analysis, pre-surgical planning, and prediction of life expectancy.
- CNN-based segmentation methods achieve strong benchmark performance but struggle to learn long-range dependencies needed for tumors with varied shapes and sizes.
- Vision Transformers model long-range information through pairwise self-attention between token embeddings, producing local and global contextual representations.
- Swin UNETR combines a hierarchical Swin transformer encoder with a CNN decoder connected through skip connections at different resolutions for multimodal 3D brain tumor segmentation.
- The model was validated in BraTS 2021, ranking among the top-performing validation approaches and showing competitive testing performance.
2 Related work
Prior work explored CNN ensembles and hybrid or standalone transformer components for medical image segmentation. Swin UNETR instead uses a hierarchical Swin transformer encoder with multiresolution decoder connections to model multiscale context and long-range dependencies.
- Earlier BraTS approaches included ensembles of U-Net-shaped CNN architectures, including 3D U-Net, 3D FCN, DeepMedic, and SegResNet.
- Related transformer methods inserted ViT components into bottlenecks, CNN encoder-decoder pipelines, or deformable transformer layers between CNN stages.
- Swin UNETR uses a U-shaped architecture with a Swin transformer encoder and CNN decoder connected by skip connections at multiple resolutions.
- Swin transformers provide multiscale contextual representations and long-range dependency modeling through hierarchical shifted-window self-attention.
3 Swin UNETR
Swin UNETR combines a hierarchical shifted-window transformer encoder with a convolutional decoder to process multimodal 3D MRI for brain-tumor segmentation.
- Encoder: Multimodal MRI is partitioned into 3D tokens and projected into an embedding space before transformer processing.The input has four MRI channels, with a 2 × 2 × 2 patch size and embedding dimension C = 48.
- Encoder: Self-attention operates within non-overlapping 3D windows, whose regions are shifted in subsequent layers to support efficient token interaction.A 3D cyclic shift is used to compute attention efficiently.
- Encoder: The encoder contains four stages with two transformer blocks per stage, while patch merging halves feature resolution between stages.The encoder therefore has L = 8 transformer layers and maintains a hierarchical multi-resolution representation.
- Decoder: Encoder representations are reshaped and passed to a convolutional decoder through skip connections at each resolution.The decoder uses residual blocks with 3 × 3 × 3 convolutions, deconvolutional upsampling, concatenation, and a final 1 × 1 × 1 convolution.
- Training and data: Training uses soft Dice loss computed voxel-wise across classes and applies normalization, random cropping, flipping, and intensity augmentation.Models use 128 × 128 × 128 training patches and are trained for 800 epochs with warmup and cosine annealing.
- Training and data: The BraTS 2021 setup uses four aligned MRI modalities, 1251 training cases, 219 validation cases, five-fold cross-validation, and no additional data.The final result is obtained using an ensemble of 10 Swin UNETR models.
4 Results and Discussion
Swin UNETR outperformed competing approaches in five-fold cross-validation and ranked among the top-performing methods on BraTS 2021 validation, while testing performance remained similar for ET and WT but decreased for TC.
- Cross-validation: The cross-validation benchmarks were reported as mean Dice scores for Enhancing Tumor, Whole Tumor, and Tumor Core.ET, WT, and TC denote the three semantic classes evaluated.
- Cross-validation: Swin UNETR outperformed all competing approaches across five folds and all semantic classes, with average gains of 0.7% for ET, 0.6% for WT, and 0.4% for TC.The average improvement across all classes and folds was 0.5%.
- Model comparison: nnU-Net and SegResNet remained competitive, while TransBTS performed sub-optimally compared with the other models.The authors attribute TransBTS’s lower performance to using ViT only in the bottleneck without decoder connections at different resolutions.
- Validation: On the BraTS 2021 validation set, Swin UNETR ranked among the top-performing methods across more than 2000 submissions.The validation results used mean Dice score and Hausdorff distance benchmarks.
- Validation: Validation segmentation outputs were well-delineated for all three tumor sub-regions, consistent with the quantitative benchmarks.The three sub-regions correspond to ET, WT, and TC.
- Testing: On the BraTS 2021 testing set, ET and WT performance was similar to validation, whereas TC performance decreased by 0.9%.Testing benchmarks reported mean Dice scores and Hausdorff distances.
5 Conclusion
The paper introduces Swin UNETR for semantic segmentation of brain tumors from multi-modal MRI images and validates it in the BraTS 2021 challenge. The model uses a U-shaped design with a Swin-transformer encoder and CNN decoder connected through multi-resolution skip connections.
- Swin UNETR is a novel architecture for semantic segmentation of brain tumors using multi-modal MRI images.
- The model combines a U-shaped network, Swin-transformer encoder, and CNN-based decoder connected by skip connections at different resolutions.
- Swin UNETR was validated in the BraTS 2021 challenge and ranked among the top-performing approaches in validation.It also demonstrated competitive performance in testing.
- The authors propose Swin UNETR as a potential foundation for transformer-based models with hierarchical encoders for brain tumor segmentation.