PaperlayerView exact revision on arXiv

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

arXiv:2111.14791v2cs.CVcs.AIcs.LG

TL;DR

Medical image analysis faces a gap between natural-image methods and volumetric context, while large expert-annotated 3D datasets are difficult to create. This paper introduces self-supervised pre-training with Swin UNETR and tailored proxy tasks, achieving state-of-the-art results on BTCV and MSD segmentation leaderboards.

  • Problem

    Medical image analysis faces a natural-to-medical image domain gap, limited cross-plane context in 2D models, and difficulty creating expert-annotated 3D datasets at scale.

    Source: NVIDIA · p. 2

  • Method

    The framework pre-trains a 3D Swin UNETR using tailored self-supervised proxy tasks, including inpainting, 3D rotation prediction, and contrastive learning.

    Source: NVIDIA · p. 2

  • Results

    Swin UNETR achieves state-of-the-art results on both BTCV and MSD; across MSD’s ten tasks, it reports 78.68% average Dice.

    Source: §5.4.1 · p. 6, §5.4.2 · p. 7, §7 · p. 8

  • Takeaways & Limitations

    For BTCV, pre-training reaches the scratch-trained model’s 83.13 Dice using 60% of the data, indicating at least 40% less annotation effort.

    Source: §5.5.2 · p. 8

    The demonstrated evaluation is segmentation on BTCV and MSD; transfer to classification and detection is presented as a potential application.

    Source: §6 · p. 8

Abstract

from arXiv · show

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 · 1. Introduction

The introduction presents ViTs [20] as driving a trend in computer vision [14] [62] and medical image analysis [7] [24], and notes transformers’ ability to learn pre-text tasks and global and local information across layers.

1. Introduction

  • Vision Transformers (ViT)s [20] have started a revolutionary trend in computer vision [14] [62] and medical image analysis [7] [24]; transformers show exceptional pre-text-task learning and effectively learn global and local information across layers.Source: §1 · p. 1

NVIDIA

  • The work was completed during an internship at NVIDIA.Source: §1 · p. 1

NVIDIA

The paper introduces Swin UNETR and anatomy-focused self-supervised proxy tasks for pre-training on CT images, then fine-tunes the model for 3D segmentation. The pre-trained model achieves state-of-the-art results on the public test leaderboards of MSD and BTCV.

NVIDIA

  • Medical imaging presents a domain gap from natural images, while 2D transformers lack cross-plane context for volumetric data and expert-annotated 3D datasets are costly to create.Source: NVIDIA · p. 2
  • The Swin Transformer encoder directly uses 3D input patches and extracts representations at multiple resolutions for pre-training.The proposed Swin UNETR has a hierarchical encoder.Source: NVIDIA · p. 2
  • Proxy tasks exploit anatomical consistency in CT images, using patch queries from head, neck, lung, abdomen, and pelvis to learn features across anatomical contexts, organs, tissues, and shapes.Source: NVIDIA · p. 2
  • Swin UNETR is pre-trained on 5,050 public CT images using contrastive learning, masked volume inpainting, and 3D rotation prediction.Contrastive learning distinguishes ROIs; inpainting learns masked-region texture, structure, and correspondence to surrounding context; rotation prediction learns structural content and generates sub-volumes for contrastive learning.Source: Figure 1 · p. 1, NVIDIA · p. 2
  • The model is fine-tuned for segmentation using hierarchical encoder features at four resolutions connected to a CNN decoder and segmentation head.The task is reformulated as a 1D sequence-to-sequence prediction task.Source: NVIDIA · p. 2
  • Fine-tuned Swin UNETR achieves state-of-the-art on the public test leaderboards of both MSD and BTCV.The evaluation uses 3D medical image segmentation as the downstream task.Source: NVIDIA · p. 2

2. Related Works

Earlier medical segmentation transformers often used transformers only as bottlenecks or add-on modules, while single-scale ViTs are impractical for high-resolution 3D volumes. This work instead uses hierarchical volumetric context and jointly trains tailored self-supervised tasks with a transformer encoder.

Medical Segmentation with Transformers

  • Earlier medical segmentation transformers used bottleneck or add-on modules with limited spatial-context exploitation; this work directly embeds high-dimensional volumes as 3D patches with positional embeddings.The approach is contrasted with prior work using transformers as secondary encoders [7] [55].Source: §2 · p. 2
  • Swin Transformer [36] [37] uses shifted windows for hierarchical encoding to avoid quadratic attention costs, and prior 2D systems include Swin UNet [5] and DS-TransUNet [34].The paper extends this direction to broader medical segmentation by incorporating hierarchically volumetric context.Source: §2 · p. 3

Pre-training in Medical Image Analysis

  • Medical pre-training has used labeled-data transfer and self-supervised proxy tasks, whereas this framework jointly trains tailored 3D tasks with a transformer encoder.Prior proxies include jigsaw puzzles [40], inpainting [23] [43], Rubik’s cube solving, rotation prediction [21] [50], and contrastive coding [13] [41].Source: §2 · p. 3

3. Swin UNETR

Swin UNETR pairs a hierarchical 3D Swin Transformer encoder with a CNN decoder through skip connections at different resolutions, forming a U-shaped network for downstream segmentation.

3. Swin UNETR

  • Swin UNETR connects a 3D-patch Swin Transformer encoder to a CNN-based decoder through skip connections at different resolutions.The encoder’s hierarchical representations provide multi-scale features for downstream applications such as segmentation.Source: §3 · p. 3, §3.1 · p. 4

3.1. Swin Transformer Encoder

  • The encoder computes local self-attention within non-overlapping 3D windows and shifts the windows in the subsequent layer.It uses 2×2×2 patches, a 48-dimensional embedding, and four stages with two transformer blocks each; patch merging reduces resolution between stages.Source: §3.1 · pp. 3–4

3.2. Decoder

  • The decoder reshapes representations from each encoder stage and the bottleneck, processes them with residual convolutional blocks, then upsamples and concatenates features across stages.For segmentation, it combines the encoder output with processed input-volume features and uses a final 1×1×1 convolution with softmax to compute probabilities.Source: §3.2 · p. 4

4. Pre-training

Swin UNETR is pre-trained with three self-supervised proxy tasks—masked volume inpainting, 3D image rotation, and contrastive coding—combined in a multi-objective loss. Randomly cropped sub-volumes receive two augmented views, and the projection heads used during pre-training are removed for downstream fine-tuning.

4. Pre-training

  • The pre-training framework combines masked volume inpainting, 3D image rotation, and contrastive coding to learn ROI-aware representations of the human body.Three projection heads are attached during pre-training; random sub-volume crops receive two views with stochastic rotation and cutout augmentations, and the heads are removed for downstream fine-tuning.Source: §4 · p. 4

4.1. Masked Volume Inpainting

  • The 3D masked volume inpainting task extends prior 2D work by reconstructing the input volume from a cutout-masked sub-volume using an L1 loss.A transposed convolution layer attached to the encoder serves as the reconstruction head.Source: §4.1 · p. 4

4.2. Image Rotation

  • The image-rotation task classifies each sub-volume’s rotation as 0°, 90°, 180°, or 270° about the z-axis using a cross-entropy loss.Source: §4.2 · p. 4

4.3. Contrastive Coding

  • Contrastive coding uses cosine similarity to bring augmented views from the same sub-volume together while separating views from different sub-volumes.A linear layer maps each view to a latent representation; the loss uses a normalized temperature scale and strengthens intra-class compactness and inter-class separability.Source: §4.3 · p. 5

4.4. Loss Function

  • The encoder is trained by minimizing the combined losses for masked volume inpainting, 3D image rotation, and contrastive coding, with grid search selecting λ1 = λ2 = λ3 = 1.Source: §4.4 · p. 5

5. Experiments

Experiments evaluate self-supervised Swin UNETR pre-training on public CT data and fine-tuning on BTCV and MSD. The model achieves the top reported benchmark results, while pre-training improves segmentation and reduces BTCV annotation requirements.

5.1. Datasets; 5.2. Implementation Details; 5.3. Evaluation Metrics

  • Experiments use five public CT datasets with 5,050 subjects, without pre-training labels, and evaluate Dice and HD95.BTCV includes 30 subjects with 13 annotated organs; MSD covers ten segmentation tasks spanning different organs and modalities. Pre-training uses 450K iterations with 96^3 patches; BTCV and MSD models use five-fold cross-validation and ensemble the best fold models.Source: §5.1 · p. 5, §5.2 · p. 5, §5.3 · p. 5

5.4.1 BTCV Multi-organ Segmentation Challenge

  • Swin UNETR achieves a BTCV mean Dice of 0.908, outperforming the second-, third-, and fourth-ranked baselines by 1.6%, 2.0%, and 2.4%, respectively.Improvements over the prior state-of-the-art include 3.6% for splenic and portal veins, 1.6% for pancreas, and 3.8% for adrenal glands.Source: Table 1 · p. 6, §5.4.1 · p. 6

5.4.2 Segmentation Results on MSD

  • Swin UNETR ranks first on the MSD leaderboard, with the best average Dice of 78.68% across all ten tasks.It achieves state-of-the-art performance on BrainTumour, Lung, Pancreas, and Colon tasks, and comparable results on the other six. Pre-trained weights are used for CT tasks; MRI tasks are trained from scratch because of the CT–MRI domain gap.Source: Table 2 · p. 6, §5.4.2 · p. 7

5.5.4 Efficacy of Self-Supervised Objectives

  • On BTCV, all proxy tasks together achieve the best Dice of 84.72%.Inpainting gives the highest improvement among single-task objectives; pairing inpainting with contrastive learning yields Dice of 84.45% and HD of 24.37.Source: Table 4 · p. 7, §5.5.4 · p. 8

5.5.1 Efficacy of Pre-training

  • Pre-training improves Dice on all five MSD CT tasks, with Colon rising from 34.83% to 43.38%.The other scratch-to-pre-trained Dice pairs are Liver, 75.27% to 77.77%; Hepatic Vessel, 64.63% to 68.52%; Pancreas, 67.12% to 67.82%; and Spleen, 96.05% to 97.32%.Source: §5.5.1 · p. 7, Figure 6 · p. 8

5.5.2 Reduce Manual Labeling Efforts

  • With 60% of BTCV labels, pre-trained Swin UNETR matches the 83.13 Dice achieved by scratch training on the full dataset.With 10% labeled data, pre-training improves performance by approximately 10%; using all labeled data yields 1.3% higher average Dice, and the approach reduces annotation effort by at least 40%.Source: Figure 7 · p. 8, §5.5.2 · p. 8

5.5.3 Size of Pre-training Dataset

  • BTCV fine-tuning benefits from larger pre-training datasets, with results reported for models pre-trained on 100, 3,000, and 5,000 scans.The study reports robustness to the total number of pre-training CT scans and increasing benefit as unlabeled data increases.Source: Figure 8 · p. 8, §5.5.3 · p. 8

6. Discussion and Limitations

Swin UNETR’s leaderboard results support the proposed self-supervised framework, and pretraining improves fine-tuning outcomes over random initialization. Its broader transferability is promising but remains unestablished for other imaging modalities.

6. Discussion and Limitations

  • Swin UNETR achieved state-of-the-art test-leaderboard results on MSD and BTCV, and outperformed DiNTS [27] and C2FNAS [59] on MSD pancreas segmentation.The cited AutoML algorithms were designed to search for optimal architectures on the same segmentation task.Source: §6 · p. 8
  • Fine-tuning pretrained Swin UNETR achieved higher accuracy, faster convergence, and reduced annotation effort compared with training from randomly initialized weights.Source: §6 · p. 8
  • The framework can be extended with additional proxy tasks and augmentation transformations, while its pretrained encoder may support transfer to classification and detection tasks.Source: §6 · p. 8
  • Pretraining has used only CT images, and direct transfer to MRI or other modalities remains insufficiently demonstrated because of domain gaps and modality-specific input channels.The authors identify transfer to other modalities as a direction for future work.Source: §6 · p. 8

7. Conclusions · Appendix

The paper presents a self-supervised pre-training framework for 3D medical images and reports state-of-the-art performance on BTCV and MSD segmentation challenges. The appendix outlines supplementary dataset, implementation, comparison, complexity, and pre-training details.

7. Conclusions

  • Self-supervised pre-training enables Swin UNETR to achieve state-of-the-art performance on the BTCV multi-organ segmentation and MSD challenges.Source: §7 · p. 8
  • Swin UNETR combines a Swin Transformer encoder’s spatial representations with convolution-based decoders.The architecture draws on merging feature maps at different scales.Source: §7 · p. 8
  • The work presents transformer-based 3D medical image pre-training for fine-tuning segmentation tasks.Source: §7 · p. 8
  • Pre-training used 5,050 CT volumes assembled from multiple public datasets with diverse anatomical regions of interest.Source: §7 · p. 8

Appendix

  • The appendix describes the public datasets used for pre-training and fine-tuning preprocessing and implementation details for BTCV and MSD.Source: Appendix · p. 12
  • The appendix presents qualitative and quantitative MRI segmentation comparisons from MSD, including benchmarks of top-ranking methods on its test leaderboard.Source: Appendix · p. 12
  • The appendix also provides model complexity analysis and pseudocode for Swin UNETR self-supervised pre-training.Source: Appendix · p. 12

A. Pre-training Datasets · B. Preprocessing Pipelines · B.1. BTCV Dataset

Swin UNETR was pre-trained on 5,050 unlabeled CT scans spanning multiple body regions and datasets, with foreground patch sampling and intensity normalization. The merged sections also describe the BTCV and MSD benchmark scope and BTCV preprocessing.

A. Pre-training Datasets

  • The pre-training cohort comprised 5,050 CT scans from head, neck, chest, abdomen, and pelvis datasets.It included 888 LUNA16, 761 TCIA Covid19, and 475 LiDC chest scans; 1,287 HNSCC head-and-neck scans; and 1,599 TCIA Colon abdomen-and-pelvis scans.Source: A. Pre-training Datasets · p. 12
  • Pre-training used only unlabeled images, with 5% of each dataset reserved for validation, yielding 4,761 training and 249 validation volumes.Source: A. Pre-training Datasets · p. 12
  • CT intensities were clipped from -1000 to 1000 and normalized to 0–1; 96 × 96 × 96 foreground sub-volumes were sampled, excluding full-air patches.Source: A. Pre-training Datasets · p. 12

B. Preprocessing Pipelines

  • Fine-tuning was evaluated on BTCV, with 30 CT scans and 13 annotated anatomies, and the MSD challenge [48], comprising 10 tasks across organs, sources, and modalities.Source: B. Preprocessing Pipelines · p. 12

B.1. BTCV Dataset

  • BTCV scans were interpolated to isotropic 1.5 × 1.5 × 2.0 mm voxel spacing.Source: §B.1 · p. 12
  • BTCV was formulated as 13-class segmentation covering large organs, vascular tissues, and small anatomies including the gallbladder, pancreas, and adrenal glands.The classes included liver, spleen, kidneys, stomach, esophagus, aorta, IVC, splenic and portal veins.Source: §B.1 · p. 12
  • BTCV preprocessing used a soft-tissue window, normalized intensities to 0–1, and randomly sampled 96 × 96 × 96 voxels; training augmentation used random flips, rotations, and intensity shifts.The respective augmentation probabilities were 0.1, 0.1, and 0.5.Source: §B.1 · p. 12

B.2. MSD Dataset … E. Pre-Training Algorithm Details

Swin UNETR achieves the best average results across four MSD MRI datasets, while its self-supervised pre-training uses inpainting, rotation prediction, and contrastive coding. The merged sections also report qualitative comparisons, task-specific MSD preprocessing, model complexity, and pre-training time.

B.2. MSD Dataset

  • MSD contains six CT and four MRI datasets, and the tasks use task-specific preprocessing and augmentation; BrainTumour combines four MRI modalities and constructs TC, WT, and ET labels.For BrainTumour, labels 2 and 3 form tumor core (TC), labels 1, 2, and 3 form whole tumor (WT), and label 2 forms enhancing tumor (ET).Source: §B.2 · p. 12

C.2. MSD Quantitative Comparisons

  • Dice of 82.14% and NSD of 94.66% are Swin UNETR’s best average results across four MSD MRI datasets; brain-tumor Dice improves by 2% over the second-best result.On Heart, Hippocampus, and Prostate, results are comparable to Models Genesis and nnUNet.Source: Table S.2 · p. 15, §C.2 · p. 15

C.1. MSD Qualitative Comparisons

  • Across MSD visualizations, Swin UNETR generally segments better than DiNTS and resolves observed under- and over-segmentation outliers.DiNTS under-segments Heart and Hippocampus, and tumors in BrainTumour, Liver, Pancreas, HepaticVessel, and Colon; on Lung it over-segments surrounding tissue. Dice is comparable for Heart and Spleen, though Swin UNETR better segments tissue corners.Source: Figure S.1 · p. 14, §C.1 · p. 15

D. Model Complexity and Pre-training Time

  • Swin UNETR has 61.98M parameters, fewer than listed transformer baselines but more than nnUNet and ASPP, with FLOPs and inference time comparable to 3D approaches.The listed parameter counts are TransUNet 96.07M, SETR 86.03M, UNETR 92.58M, nnUNet 19.07M, and ASPP 47.92M.Source: D. Model Complexity and Pre-training… · p. 15
  • Pre-training uses 5,050 scans from five public datasets and takes approximately 6 GPU days, with training capped at 45K iterations.Source: D. Model Complexity and Pre-training… · p. 16

E. Pre-Training Algorithm Details

  • Self-supervised pre-training combines masked volume inpainting, rotation prediction, and contrastive coding to learn representations from unlabeled 3D data.The proxy tasks target image context, geometry, and similarity, respectively.Source: E. Pre-Training Algorithm Details · p. 16
  • The algorithm generates two rotated views per input, applies cutout masking 30% of sub-volume voxels, and reconstructs the inputs from Swin Transformer output tokens.Random rotations are 0°, 90°, 180°, or 270°; local CT transformations are arranged across neighboring tokens to construct semantic targets in medical spatial context.Source: E. Pre-Training Algorithm Details · p. 16, D. Model Complexity and Pre-training… · p. 16

Related papers

Loading 2111.14791v2…