HybridMIM: A Hybrid Masked Image Modeling Framework for 3D Medical Image Segmentation

TL;DR

HybridMIM introduces a two-level masked image modeling framework for 3D medical image segmentation that learns pixel-, region-, and sample-level semantics. By combining partial region reconstruction, patch-masking perception, and dropout-based contrastive learning, it delivers improved segmentation performance while accelerating pre-training and supporting both CNN and transformer backbones. Across BraTS2020, BTCV, MSD Liver, and MSD Spleen, HybridMIM consistently outperforms state-of-the-art supervised and SSL methods and reduces annotation requirements. The approach offers a practical, versatile pre-training paradigm for high-dimensional medical imaging tasks and is publicly available in the authors’ repository.

Abstract

Masked image modeling (MIM) with transformer backbones has recently been exploited as a powerful self-supervised pre-training technique. The existing MIM methods adopt the strategy to mask random patches of the image and reconstruct the missing pixels, which only considers semantic information at a lower level, and causes a long pre-training time.This paper presents HybridMIM, a novel hybrid self-supervised learning method based on masked image modeling for 3D medical image segmentation.Specifically, we design a two-level masking hierarchy to specify which and how patches in sub-volumes are masked, effectively providing the constraints of higher level semantic information. Then we learn the semantic information of medical images at three levels, including:1) partial region prediction to reconstruct key contents of the 3D image, which largely reduces the pre-training time burden (pixel-level); 2) patch-masking perception to learn the spatial relationship between the patches in each sub-volume (region-level).and 3) drop-out-based contrastive learning between samples within a mini-batch, which further improves the generalization ability of the framework (sample-level). The proposed framework is versatile to support both CNN and transformer as encoder backbones, and also enables to pre-train decoders for image segmentation. We conduct comprehensive experiments on four widely-used public medical image segmentation datasets, including BraTS2020, BTCV, MSD Liver, and MSD Spleen. The experimental results show the clear superiority of HybridMIM against competing supervised methods, masked pre-training approaches, and other self-supervised methods, in terms of quantitative metrics, timing performance and qualitative observations. The codes of HybridMIM are available at https://github.com/ge-xing/HybridMIM

Figures (6)

• Figure 1: Illustration of our key idea in the 2D form. We regularly divide the input image into two-levels of patches, masking sub-regions or patches randomly. The masking information is encoded into binary, providing the locations and number of masked patches. In addition, local regions are selected for reconstruction, which facilitates a faster pre-training efficiency for high-dimensional medical images.
• Figure 2: Overviwe of the HybridMIM pre-training framework. Input 3D medical images (demonstrated in 2D form) are randomly masked with a two-level masking strategy, then fed to the encoder twice to obtain two sets of feature representations. We use patch-masking perception, partial region reconstructions, and dropout-based contrastive learning as proxy tasks to learn contextual representations of input images.
• Figure 3: Qualitative visualizations of the proposed HybridMIM and baseline methods. "Ours" is the HybridMIM(Swin) method. The three rows of visual comparison results are from BraTS2020, Liver, and BTCV datasets. Our proposed method is better for segmenting tiny lesions (first row) and has higher segmentation integrity (second row, third row).
• Figure 4: Effect of different labeled data sizes on migration learning results. We selected 10%, 20%, 40%, 60%, 80%, and 100% of the training data in the BraTS2020 dataset to verify the transfer learning ability in different self-supervised learning methods.
• Figure 5: Comparison of pre-training time consumption for different SSL methods. "Not partial" denotes that the partial region prediction scheme is not used.