Transferring SAM-pretrained 2D ViTs for semi-supervised 3D medical image segmentation

Abstract

The scarcity of labeled data has made semi-supervised learning essential for medical image segmentation. Recently, Vision Transformers (ViT), pre-trained on large-scale 2D natural images have shown remarkable performance in 2D image segmentation tasks. It is natural to transfer the learned knowledge in ViT to the data-limited semi-supervised medical image segmentation task. However, directly applying ViT faces several challenges, including adapting 2D-pretrained models to 3D medical data and addressing performance degradation in ViT when trained on small datasets. To tackle these challenges, this paper proposes a method for medical image segmentation. The method leverages the strengths of both ViT and Convolutional Neural Networks (CNN) via a hybrid architecture. The CNN-based encoder and decoder play a projector and tokenizer role for the ViT, while the architecture of ViT is fully retained to preserve the knowledge in the pre-trained model derived from SAM (Segment Anything Model) as much as possible. In addition, pseudo-labeling serves as the core guidance for our method to learn from unlabeled data. Experimental results show that our approach outperforms the state-of-the-art on the Pancreas-CT dataset by a large margin and enables rapid transfer learning from 2D-pretrained models to 3D medical tasks with few labeled data, making it especially valuable for rare disease diagnosis.