[Paper Review] MIS-FM: 3D Medical Image Segmentation using Foundation Models Pretrained on a Large-Scale Unannotated Dataset

Guotai Wang, Jianghao Wu+4 more|arXiv (Cornell University)|Jun 29, 2023

Radiomics and Machine Learning in Medical Imaging23 citations

TL;DR

MIS-FM introduces Volume Fusion, a self-supervised pretraining strategy for 3D medical image segmentation that uses unannotated data to pretrain a full segmentation model, plus a Parallel Convolution and Transformer Network (PCT-Net) for effective transfer to downstream tasks.

ABSTRACT

Pretraining with large-scale 3D volumes has a potential for improving the segmentation performance on a target medical image dataset where the training images and annotations are limited. Due to the high cost of acquiring pixel-level segmentation annotations on the large-scale pretraining dataset, pretraining with unannotated images is highly desirable. In this work, we propose a novel self-supervised learning strategy named Volume Fusion (VF) for pretraining 3D segmentation models. It fuses several random patches from a foreground sub-volume to a background sub-volume based on a predefined set of discrete fusion coefficients, and forces the model to predict the fusion coefficient of each voxel, which is formulated as a self-supervised segmentation task without manual annotations. Additionally, we propose a novel network architecture based on parallel convolution and transformer blocks that is suitable to be transferred to different downstream segmentation tasks with various scales of organs and lesions. The proposed model was pretrained with 110k unannotated 3D CT volumes, and experiments with different downstream segmentation targets including head and neck organs, thoracic/abdominal organs showed that our pretrained model largely outperformed training from scratch and several state-of-the-art self-supervised training methods and segmentation models. The code and pretrained model are available at https://github.com/openmedlab/MIS-FM.

Motivation & Objective

• Motivate pretraining segmentation models on large-scale unannotated 3D medical images to improve performance when annotated data are limited.
• Propose Volume Fusion as a self-supervised pseudo-segmentation pretext task that yields paired inputs and segmentation labels without manual annotation.
• Introduce PCT-Net, a network combining parallel convolution and transformer blocks for efficient transfer to multi-scale segmentation tasks.
• Demonstrate effectiveness across head-and-neck, thoracic, and abdominal organ segmentation with large-scale pretraining.
• Release pretrained models on 110k CT scans to enable broader clinical transferability.

Proposed method

• Volume Fusion (VF): fuse a foreground sub-volume from one scan with a background sub-volume from another using a discrete fusion coefficient map, turning voxel-wise fusion categories into a segmentation pretext task.
• Fusion coefficients alpha_i come from a discrete set {0, 1/K, ..., 1}, creating C = K+1 classes for voxels; the model predicts the voxel class Y corresponding to alpha_i.
• Pretraining objective L_sup combines Dice loss and cross-entropy loss to train the full segmentation model end-to-end on the pseudo-segmentation task.
• PCT-Net: a three-level, pyramid architecture that fuses local CNN features with long-range Transformer-based context using a dual-branch PCT block (local Conv branch and global self-attention branch).
• Embedding module uses 2D convs at high resolution and 3D convs at lower resolutions to handle anisotropic 3D CT data; predictions are produced at multiple scales with deep supervision.
• Training setup uses 110k unannotated CT volumes (PData-110k) for pretraining, followed by fine-tuning on downstream segmentation datasets using standard Dice and ASSD metrics.

Figure 1 : Overview of our proposed self-supervised pretraining method based on volume fusion.

Experimental results

Research questions

• RQ1 Does Volume Fusion pretraining improve segmentation performance on downstream 3D CT tasks compared to training from scratch and existing SSL methods?
• RQ2 How does the fusion parameter K affect transfer performance and context learning in VF pretraining?
• RQ3 What is the impact of pretraining data scale (1k, 10k, 110k volumes) on downstream segmentation performance?
• RQ4 Can PCT-Net efficiently leverage VF-pretrained features across organs with different scales (head/neck, thoracic, abdominal)?

Key findings

• VF pretraining consistently outperforms training from scratch and other SSL methods on SegTHOR for thoracic organs, achieving Dice of 88.30% and ASSD of 1.78 mm.
• On MICCAI 2015 Head-Neck data, PCT-Net with VF achieves higher average Dice (82.74) and lower ASSD (0.77 mm) than several baselines including nnU-Net, TransUNet, and UNETR++.
• VF with K=4 yielded the best trade-off, attaining the lowest ASSD (1.78 mm) and competitive Dice across tested datasets; large K values did not always improve performance.
• Increasing pretraining data size (1k → 10k → 110k) generally improves downstream Dice for both 3D U-Net and PCT-Net, with larger gains observed using the PCT-Net architecture.
• Compared to Patch Swapping, Model Genesis, and MIM, VF achieves the highest average Dice (88.30%) and the lowest ASSD (1.78 mm) on SegTHOR.
• VF+PCT-Net demonstrates improved downstream results over either VF alone or baseline architectures, illustrating beneficial transfer and context learning.

Figure 2 : 2D illustration of fused images with different $K$ values. The colorful images are fusion coefficient maps, and they are converted into label maps for the pretext segmentation task.