Source-linked AI summary
Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis
Yucheng Tang, Dong Yang, Wenqi Li, Holger Roth, Bennett Landman, Daguang Xu, Vishwesh Nath, Ali Hatamizadeh
TL;DR
Medical image analysis has faced domain and volumetric-context challenges that limit direct benefit from general computer-vision advances. The paper introduces Swin UNETR with tailored self-supervised tasks, pre-trains it on CT data, and fine-tunes it for segmentation, achieving first-place public leaderboard results on MSD and BTCV.
Problem
Medical image analysis has a domain gap from natural images and lacks cross-plane context when transformers are applied to volumetric CT or MRI data.
Method
The paper pre-trains a 3D Swin UNETR with tailored proxy tasks including image inpainting, 3D rotation prediction, and contrastive learning.
Results
The pre-trained model achieves state-of-the-art performance and top public leaderboard rankings on both MSD and BTCV segmentation benchmarks.
Takeaways & Limitations
The framework supports transfer to medical image analysis tasks using large collections of unlabeled CT images and tailored anatomical proxy tasks.
Takeaways & Limitations
The discussion states that the framework can be extended with more proxy tasks and augmentation transformations, and that the pre-trained encoder can benefit other tasks such as classification and detection.
Abstract
from arXiv · showhide
Vision Transformers (ViT)s have shown great performance in self-supervised learning of global and local representations that can be transferred to downstream applications. Inspired by these results, we introduce a novel self-supervised learning framework with tailored proxy tasks for medical image analysis. Specifically, we propose: (i) a new 3D transformer-based model, dubbed Swin UNEt TRansformers (Swin UNETR), with a hierarchical encoder for self-supervised pre-training; (ii) tailored proxy tasks for learning the underlying pattern of human anatomy. We demonstrate successful pre-training of the proposed model on 5,050 publicly available computed tomography (CT) images from various body organs. The effectiveness of our approach is validated by fine-tuning the pre-trained models on the Beyond the Cranial Vault (BTCV) Segmentation Challenge with 13 abdominal organs and segmentation tasks from the Medical Segmentation Decathlon (MSD) dataset. Our model is currently the state-of-the-art (i.e. ranked 1st) on the public test leaderboards of both MSD and BTCV datasets. Code: https://monai.io/research/swin-unetr
NVIDIA
The paper reports that its model ranked first on the public test leaderboards of both MSD and BTCV datasets.
- The model ranked first on the public test leaderboards of both MSD and BTCV datasets.
1. Introduction
Vision Transformers learn global and local representations that can transfer to downstream applications, but the paper situates medical imaging as a distinct setting for these advances.
- Vision Transformers learn global and local information across layers for transferable visual representations.
- The paper identifies medical imaging as a distinct application context for these transformer advances.
NVIDIA
The paper proposes a 3D Swin UNETR framework that combines hierarchical transformer encoding with tailored self-supervised proxy tasks for medical image analysis. It pre-trains on CT data and evaluates transfer through segmentation benchmarks.
- The pre-training framework combines randomly cropped and augmented CT sub-volumes with masked inpainting, contrastive learning, and rotation prediction.
- The framework uses image inpainting, 3D rotation prediction, and contrastive learning to learn anatomical patterns from 3D medical images.
- Swin UNETR uses a hierarchical transformer encoder and CNN-based decoder connected through multi-resolution features.
- The proposed encoder and proxy tasks were pre-trained on 5,050 publicly available CT images from various applications.
- Fine-tuning on MSD and BTCV public benchmarks achieved state-of-the-art performance on both test leaderboards.
2. Related Works
Related work covers transformer-based medical image segmentation and pre-training, while positioning this paper's combination of 3D transformer encoding and tailored self-supervision as distinct.
- Medical Segmentation with Transformers: Medical segmentation methods use transformer blocks in architectures designed for dense, multi-scale inference.
- Medical Segmentation with Transformers: Single-patch-size vision transformers are difficult to apply to high-resolution, high-dimensional volumetric images.
- Pre-training in Medical Image Analysis: Prior medical-image pre-training studies used labeled data, while self-supervised learning offers a way to use unlabeled data.
- Pre-training in Medical Image Analysis: The paper combines multiple 3D-imaging pre-text tasks with a transformer-based encoder for representation learning.
3. Swin UNETR
Swin UNETR uses a 3D Swin Transformer encoder with hierarchical, multi-resolution representations and a CNN decoder connected through skip connections for downstream segmentation.
- Swin UNETR combines a Swin Transformer encoder using 3D patches with a CNN-based decoder connected by skip connections at different resolutions.
- The encoder begins by projecting 3D patches into a C-dimensional embedding space before local self-attention is computed within windows.
- The encoder partitions 3D tokens into non-overlapping windows and alternates regular and shifted-window self-attention for efficient token interaction.A 3D cyclic shift supports efficient batch computation of shifted windows.
- The encoder has four transformer stages, each with two blocks, while patch merging halves spatial resolution between stages and produces hierarchical representations.The stages operate at progressively reduced resolutions, supporting multi-scale feature extraction.
- Each stage and the bottleneck are reshaped into multi-resolution features, processed by residual convolutional blocks, upsampled, and concatenated with preceding-stage features.
4. Pre-training
Swin UNETR is pre-trained with multiple self-supervised proxy tasks designed to encode anatomical information from 3D medical sub-volumes. The framework combines masked volume inpainting, rotation prediction, and contrastive coding in a multi-objective loss.
- The pre-training framework attaches three projection heads to the encoder and optimizes masked volume inpainting, 3D rotation prediction, and contrastive coding.The projection heads are removed when the full model is fine-tuned for downstream tasks.
- Masked Volume Inpainting: Masked volume inpainting reconstructs randomly masked regions using a transpose-convolution reconstruction head and an L1 loss between the input and reconstruction.The task extends masked inpainting from 2D images to volumetric 3D medical images.
- 3D Rotation Prediction: 3D rotation prediction classifies whether a sub-volume was rotated by 0°, 90°, 180°, or 270° around the z-axis.An MLP classification head predicts the rotation category with a cross-entropy loss.
- Contrastive Coding: Rotation and cutout transformations also provide augmented views for contrastive learning.
- Contrastive Coding: Contrastive coding maps augmented sub-volumes to latent representations and maximizes similarity between views from the same sub-volume while minimizing similarity between views from different sub-volumes.The loss uses cosine similarity and is intended to strengthen intra-class compactness and inter-class separability.
5. Experiments
Experiments evaluate Swin UNETR on BTCV and MSD benchmarks, showing strong leaderboard performance and consistent benefits from self-supervised pre-training, including in data-limited settings.
- Experimental Setup: 5,050 CT subjects from five public datasets were used for unlabeled pre-training, covering chest, abdomen, and head/neck volumes.Existing annotations and labels were not used during pre-training.
- Experimental Setup: The BTCV benchmark contains 30 contrast-enhanced CT subjects with annotations for 13 abdominal organs.The dataset includes 80–225 slices per scan, with 512×512 pixels and 1–6 mm slice thickness.
- BTCV Results: On BTCV, Swin UNETR achieved state-of-the-art Dice of 0.908, outperforming the second-, third-, and fourth-ranked baselines by 1.6%, 2.0%, and 2.4%.Improvements were especially pronounced for splenic and portal veins, pancreas, and adrenal glands.
- MSD Results: Across all ten MSD tasks, Swin UNETR achieved the best average Dice of 78.68% and ranked first on the public leaderboard.It achieved state-of-the-art performance on BrainTumour, Lung, Pancreas, and Colon tasks, with comparable results on the remaining tasks.
- Efficacy of Pre-training: Pre-training improved MSD CT performance over training from scratch, including Task10 Colon from 34.83% to 43.38% Dice.Task03 Liver improved from 75.27% to 77.77%, while Task08 Hepatic Vessel improved from 64.63% to 68.52%.
- Data Efficiency: Using 10% of BTCV labels produced approximately 10% improvement with pre-training, while the full-data setting showed 1.3% higher average Dice.The scratch-model Dice of 83.13 achieved with all labels was matched using pre-training with 60% of the data.
- Ablations: Pre-training remained robust across 100, 3,000, and 5,000 unlabeled scans, with larger pre-training datasets providing additional benefit.The ablation study found that combining all proxy tasks achieved the best Dice of 84.72% on BTCV.
6. Discussion and Limitations
The framework achieves strong MSD and BTCV leaderboard results and can exploit unlabeled medical images without annotation effort. However, pre-training was limited to CT, and direct transfer to MRI was not sufficiently demonstrated.
- Discussion: Fine-tuning pre-trained Swin UNETR achieves higher accuracy, faster convergence, and reduced annotation effort than training from randomly initialized weights.The framework is presented as scalable to additional proxy tasks and augmentation transformations.
- Discussion: The pre-trained encoder may support transfer to other medical imaging tasks, including classification and detection.This consequence is stated as a potential benefit of the framework.
- Limitations: Pre-training has only been performed on CT, and the experiments do not demonstrate sufficient direct transferability to MRI.The paper identifies this modality gap as a limitation of the current framework.
7. Conclusions
The work presents a self-supervised 3D medical-image framework combining Swin UNETR with large-scale CT pre-training and reports state-of-the-art segmentation performance on BTCV and MSD.
- Swin UNETR combines transformer-encoded spatial representations with convolution-based decoders for 3D medical-image segmentation.
- 5,050 CT volumes from multiple public datasets cover diverse anatomical regions, including head, neck, chest, abdomen, and pelvis.
- Self-supervised pre-training followed by fine-tuning achieves state-of-the-art performance on the BTCV multi-organ and MSD segmentation challenges.
A. Pre-training Datasets
The pre-training cohort combines five public datasets containing 5,050 CT scans spanning diverse body regions, while fine-tuning uses BTCV and MSD segmentation benchmarks.
- The pre-training cohort contains 5,050 CT scans from head, neck, chest, abdomen, and pelvis regions of interest.
- The chest CT cohort includes 888 LUNA16 scans, 761 TCIA Covid19 scans, and 475 LiDC scans.
- BTCV provides 30 CT scans with 13 annotated anatomies for single-task multi-organ segmentation.
- MSD contains 10 segmentation tasks involving multiple organs, sources, and imaging modalities.
- BTCV preprocessing uses isotropic voxel spacing of 1.5 × 1.5 × 2.0 mm and random 96 × 96 × 96 voxel sampling.
B.2. MSD Dataset
The MSD experiments span heterogeneous CT and MRI tasks with task-specific preprocessing and augmentation. Qualitative and quantitative comparisons report stronger overall segmentation results for Swin UNETR, including the best average MRI performance.
- MSD Dataset: MSD comprises 10 tasks, including 6 CT and 4 MRI datasets covering multiple organs and modalities.
- Preprocessing: MRI preprocessing varies by task, including four-channel BrainTumour inputs, isotropic Heart resampling, and task-specific normalization and cropping.
- Qualitative Comparisons: Swin UNETR scores 0.94 versus DiNTS at 0.85 in the displayed BrainTumour and Heart qualitative comparison.
- Qualitative Comparisons: Swin UNETR reduces under-segmentation and over-segmentation outliers observed for DiNTS, while achieving better segmentation on several MSD tasks.
- Quantitative Comparisons: 2% average Dice improvement over the second-best performance is reported for BrainTumour MRI segmentation.
- Quantitative Comparisons: 82.14% Dice and 94.66% NSD are the best average results reported across four MRI datasets.
- Model Complexity: Swin UNETR uses 61.98M parameters, fewer than TransUNet’s 96.07M and SETR’s 86.03M in the cited BTCV comparison.
- Pre-training: Self-supervised pre-training uses masked volume inpainting, 3D rotation prediction, and contrastive coding on 5,050 scans for approximately 6 GPU days.