Self-Supervised Modality-Agnostic Pre-Training of Swin Transformers

TL;DR

This paper introduces SwinFUSE, a modality-agnostic pre-training method for Swin Transformers that learns from multiple medical imaging modalities, improving generalization and out-of-distribution performance in 3D segmentation tasks.

Contribution

It proposes a novel multi-modal pre-training framework with a domain-invariance module, enabling Swin Transformers to effectively learn from diverse medical imaging modalities.

Findings

Achieves 1-2% performance trade-off on in-distribution data.

Surpasses single-modality models by up to 27% on out-of-distribution data.

Demonstrates strong generalizability across different medical imaging tasks.

Abstract

Unsupervised pre-training has emerged as a transformative paradigm, displaying remarkable advancements in various domains. However, the susceptibility to domain shift, where pre-training data distribution differs from fine-tuning, poses a significant obstacle. To address this, we augment the Swin Transformer to learn from different medical imaging modalities, enhancing downstream performance. Our model, dubbed SwinFUSE (Swin Multi-Modal Fusion for UnSupervised Enhancement), offers three key advantages: (i) it learns from both Computed Tomography (CT) and Magnetic Resonance Images (MRI) during pre-training, resulting in complementary feature representations; (ii) a domain-invariance module (DIM) that effectively highlights salient input regions, enhancing adaptability; (iii) exhibits remarkable generalizability, surpassing the confines of tasks it was initially pre-trained on. Our experiments on two publicly available 3D segmentation datasets show a modest 1-2% performance trade-off compared to single-modality models, yet significant out-performance of up to 27% on out-of-distribution modality. This substantial improvement underscores our proposed approach's practical relevance and real-world applicability. Code is available at: https://github.com/devalab/SwinFUSE