Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images (2201.01266v1)

Abstract: 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

Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images

The paper presents the Swin UNEt TRansformers (Swin UNETR), a novel approach for semantic segmentation of brain tumors using MRI images. Utilizing advanced transformer architectures, this model targets the limitations observed in traditional Fully Convolutional Neural Networks (FCNNs), particularly in capturing long-range dependencies in 3D medical image segmentation.

Model Architecture

The Swin UNETR model builds upon the well-established U-Net architecture by integrating a Swin Transformer as the encoder. In contrast to traditional FCNNs that struggle with modeling long-range information due to limited kernel sizes, the Swin Transformer uses shifted windows to efficiently compute self-attention. This enables the model to capture global contextual information across multiple resolutions, thereby addressing common segmentation challenges, especially for tumors of varying sizes.

The model employs a hierarchical design:

• Encoder: It consists of a sequence-to-sequence prediction framework that processes multi-modal input data as a 1D sequence of embeddings. The Swin transformer encoder extracts features at five distinct resolutions, enhancing multi-scale feature capture.
• Decoder: Corresponding features are upsampled in a CNN-based decoder using skip connections. This facilitates detailed reconstruction of segmentation outputs, ensuring effective translation of encoded features into accurate labels for tumor sub-regions.

Results and Analysis

The Swin UNETR was evaluated on the BraTS 2021 challenge dataset, a comprehensive collection for brain tumor segmentation. In this competitive setting, Swin UNETR demonstrated superior performance, outperforming established models such as nnU-Net and SegResNet, as well as transformer-based approaches like TransBTS.

• Numerical Performance: Swin UNETR's Dice scores for Enhancing Tumor (ET), Whole Tumor (WT), and Tumor Core (TC) were reported as 0.858, 0.926, and 0.885 on the validation dataset, highlighting its efficacy. It shows an average increase of approximately 0.5% over competitive approaches across these metrics.
• Inferential Robustness: In addition to delivering strong quantitative results, the Swin UNETR offers qualitative improvements, with well-delineated segmentation outputs evident from visual analyses of the MRI images.

Implications and Future Directions

The integration of Swin Transformers represents a significant advancement in applying self-attention mechanisms within medical image analysis. By effectively modeling long-range information through a hierarchical architecture, Swin UNETR sets a precedent for leveraging transformer models in 3D image segmentation tasks.

The paper not only offers a novel methodological contribution but also opens avenues for future research in medical imaging. Possible extensions include adapting such architectures to other imaging modalities or pathologies and exploring synergistic effects with pre-existing CNN-based frameworks. Further investigation could focus on optimizing computational costs and enhancing model efficiency without compromising performance.

Conclusion

Swin UNETR exemplifies effective utilization of transformer encoders in semantic segmentation of complex medical images. Its success in the BraTS challenge underscores the potential of Swin Transformers to significantly enhance brain tumor segmentation tasks, offering substantial clinical utility. As the adoption of transformer-based models expands in medical image analysis, Swin UNETR provides a foundational architecture for future advancements in this domain.