Source-linked AI summary
Masked Image Modeling Advances 3D Medical Image Analysis
Zekai Chen, Devansh Agarwal, Kshitij Aggarwal, Wiem Safta, Samit Hirawat, Venkat Sethuraman, Mariann Micsinai Balan, Kevin Brown
TL;DR
Costly expert annotations and the processing demands of 3D medical images make learning from unlabeled data valuable, but MIM's medical-imaging applicability has been uncertain. This paper evaluates MIM for 3D segmentation and finds it outperforms contrastive learning while reaching Dice 0.7 with 1.4× less training time than the supervised baseline.
Problem
MIM's applicability to 3D medical images is uncertain, while expert annotation is costly and time-consuming and volumetric data impose substantial processing demands.
Source: §1 · pp. 1–2Method
The paper evaluates masked image modeling for 3D medical image segmentation against contrastive learning and across different masking strategies.
Source: §1 · p. 2Results
MIM outperformed SimCLR on both reported segmentation tasks and reached a Dice score of 0.7 with 1.4× less training time than the supervised baseline.
Source: §4.1 · p. 7Takeaways & Limitations
The findings support MIM for 3D medical image modeling in the segmentation tasks and settings examined.
Source: §5 · p. 9The effectiveness of self-supervised learning when medical images are weakly labeled rather than entirely unannotated remains an open question.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Recently, masked image modeling (MIM) has gained considerable attention due to its capacity to learn from vast amounts of unlabeled data and has been demonstrated to be effective on a wide variety of vision tasks involving natural images. Meanwhile, the potential of self-supervised learning in modeling 3D medical images is anticipated to be immense due to the high quantities of unlabeled images, and the expense and difficulty of quality labels. However, MIM's applicability to medical images remains uncertain. In this paper, we demonstrate that masked image modeling approaches can also advance 3D medical images analysis in addition to natural images. We study how masked image modeling strategies leverage performance from the viewpoints of 3D medical image segmentation as a representative downstream task: i) when compared to naive contrastive learning, masked image modeling approaches accelerate the convergence of supervised training even faster (1.40$\times$) and ultimately produce a higher dice score; ii) predicting raw voxel values with a high masking ratio and a relatively smaller patch size is non-trivial self-supervised pretext-task for medical images modeling; iii) a lightweight decoder or projection head design for reconstruction is powerful for masked image modeling on 3D medical images which speeds up training and reduce cost; iv) finally, we also investigate the effectiveness of MIM methods under different practical scenarios where different image resolutions and labeled data ratios are applied.
1. Introduction
The paper examines whether masked image modeling can advance 3D medical image analysis, where costly annotations and high data volume make training demanding. Experiments compare MIM with contrastive learning and study masking strategies and practical data settings.
1. Introduction
- 3D medical image analysis faces costly, time-consuming expert annotation and substantial processing complexity from image volume, resolution, and dimensionality.These demands make training from random initialization burdensome in compute and data.Source: §1 · p. 1
- MIM's effectiveness in medical imaging remains insufficiently investigated despite its success in scalable vision models.The paper asks whether MIM can advance 3D medical image analysis.Source: §1 · pp. 1–2
- The study compares MIM with contrastive learning and examines masking choices and varied annotation, unlabeled-data, and image-resolution settings.For masking, it seeks strategies that promote holistic understanding while avoiding excessive focus on low-level texture and materials.Source: §1 · p. 2
- Across two real-world benchmarks, multi-organ and brain-tumor segmentation, masked image modeling benefited 3D medical image analysis by significantly speeding convergence and improving downstream performance.It achieved at most 1.4× training cost saving to reach the same Dice score and over 5% improvements on both segmentation tasks without any hyperparameter tuning.Source: §1 · p. 2
2. Related Work
Prior medical-imaging work includes transfer learning and several self-supervised approaches, while masked image modeling has been developed chiefly for natural images. This paper investigates MIM's applicability to 3D medical image analysis.
Transfer Learning in Medical Image Analysis
- Transfer learning from ImageNet can accelerate convergence on medical images, while domain-specific labeled pretraining can improve performance but may be impractical when labels are costly.Self-supervised learning offers an alternative that can use unlabeled medical data, which is massive and commonly more accessible.Source: §2 · p. 2
Masked Image Modeling
- MIM methods evolved from masked-pixel prediction toward raw-pixel restoration and lightweight decoders or projection heads, but their reported use was in natural-image modeling.MAE uses a lightweight decoder, while SimMIM replaces the decoder with a single linear projection layer and reports comparable results.Source: §2 · pp. 2–3
Self-Supervised Learning
- Medical self-supervised work has explored domain-specific pretext tasks and contrastive learning, but MIM-based methods had not been substantially investigated in medical imaging.Examples include 3D medical imaging strategies, dermatology classification, and combining inpainting with contrastive learning for medical segmentation.Source: §2 · p. 4
3. Approach
The approach adapts masked image modeling to 3D volumes using random patch masking, 3D encoders, and reconstruction designs. It studies raw-voxel prediction and lightweight decoding or projection for the reconstruction task.
3. Approach
- MIM masks part of an image or its encoded tokens and trains the model to recreate the masked area; the paper examines multiple MIM designs for 3D medical data.The encoder models latent features, while the decoder resamples latent vectors and a projection head aligns embeddings with original signals in masked areas; lightweight decoding is used to reduce computation.Source: §3 · p. 4
3.1. Masking Strategies
- Uniform random patch masking avoids center bias and, at high masking ratios, removes redundancy that could make reconstruction solvable from neighboring patches.Patches are fully visible or masked; the study varies patch size and masking ratio.Source: §3.1 · p. 4
3.2. Encoders
- The study adapts ViT, SwinTransformer, and VAN architectures into 3D encoders for volumetric data.Encoders model latent representations of masked patches that are used to forecast the original signals.Source: §3.2 · p. 4
3.3. Decoders
- In autoencoder designs, the decoder receives encoded visible patches and learned mask tokens, while the encoder does not see masked patches.The paper follows a Transformer-block decoder design by default; decoder backbones are independent of encoder backbones.Source: §3.3 · pp. 4–5
3.4. Reconstruction Target
- The reconstruction target is raw voxel values, with L1 or L2 loss computed only on masked patches.A linear projection restores ViT embeddings to voxel space, while SwinTransformer and VAN use two transposed-convolution layers; discrete targets may lose fine-grained texture or material information.Source: §3.4 · p. 5
4. Experiments on 3D Segmentation
The experiments evaluate MIM on CT and MRI segmentation tasks, comparing it with contrastive learning and testing masking strategies, pretraining data, image resolutions, and labeled-data ratios. Across these settings, MIM improves segmentation performance and can reduce training time.
4.1. Comparison among Different Approaches
- MIM achieves average Dice scores around 0.752–0.758 versus 0.723 for SimCLR on multi-organ segmentation, and scores above 0.80 versus 0.7739 on brain-tumor segmentation.The main-result tables cover multi-organ segmentation and brain-tumor segmentation, respectively.Source: Table 1 · p. 6, Table 2 · p. 6, §4.1 · p. 7
- At the 1.3k training step, SimMIM architectures achieve a 1.76× better Dice score; MIM reaches Dice 0.7 with 1.4× less training time than the supervised baseline.Figure 4 compares validation average Dice across training steps for the supervised baseline and MIM methods with different masking ratios.Source: Figure 4 · p. 5, §4.1 · p. 7
4.2. Masking Strategy
- With MAE, a 75% masking ratio yields best Dice scores of 0.7183 with patch size 16 on multi-organ segmentation and 0.8041 with patch sizes 24 and 32 on brain-tumor segmentation.The ablations use BTCV-only pretraining for multi-organ segmentation and BraTS-only pretraining for brain-tumor segmentation.Source: §4.2 · p. 7, Table 3 · p. 7, Table 4 · p. 8
- For SimMIM, patch size 16 performs best at high masking, while patch size 32 performs best at the smallest tested masking ratio, 0.15.The authors suggest that larger patches can hide significant signals independently, making high masking more critical for small patches.Source: §4.2 · pp. 7–8, Table 3 · p. 7, Table 4 · p. 8
4.3. Data vs. Resolutions vs. Labeled Ratio
- More plentiful pretraining data consistently outperforms less data; at lower resolutions with half the labeled data, Dice 0.6919 is 5.6% higher than 0.6552.The data-volume comparisons use multi-organ segmentation as the downstream task.Source: §4.3 · p. 8, Table 5 · p. 8
- Higher pretraining resolutions generally perform better than lower ones; with the same data source and labeled ratio, Dice 0.7338 is 2.7% lower than at the greater resolution.Resolution is varied by downsampling the images in all dimensions.Source: §4.3 · p. 8, Table 5 · p. 8
- The experiments indicate that pretraining on large datasets followed by fine-tuning with small samples is feasible, and that MIM improves downstream performance across varied contexts.The paper notes that most medical images are unannotated because labeling is costly, while public data is abundant.Source: §4.3 · p. 8
5. Conclusion
The paper reports that masked image modeling outperforms traditional contrastive learning on the studied 3D segmentation tasks, while identifying weakly labeled data and harder downstream tasks as future questions.
5. Conclusion
- Masked image modeling outperforms traditional contrastive learning by accelerating convergence and improving downstream task performance on the studied segmentation tasks.The authors also investigate its effectiveness across varied practical situations.Source: §5 · p. 9
- Whether self-supervised learning helps with medical images that are weakly labeled rather than entirely unannotated remains an open question.The authors identify comparison with supervised learning using limited supervisory signals as future work.Source: §5 · p. 9
- The authors identify integration of self-supervised learning into more challenging downstream tasks as a future direction.Source: §5 · p. 9
6. Appendix
The appendix reports greater restoration power for the large model than the tiny model and includes reconstruction examples and experimental-setting tables for segmentation and pretraining.
6.1. Experimental Settings
- The large model has more restoration power than the tiny model, while the authors attribute information loss in small models to compression of the high-dimensional voxel space.Source: §6.1 · p. 11
- The appendix includes setting tables for BTCV and BraTS supervised baselines and for CT- and MRI-volume pretraining.Source: Table 6 · p. 11, Table 7 · p. 11, Table 8 · p. 11
- Figure 5 shows ground truth, masked input, and SimMIM reconstruction slices from a CT scan using ViT-Base, 16-voxel patches, and a 75% masking ratio.The slice index represents depth.Source: Figure 5 · p. 12
- Figure 6 shows the same image and reconstruction layout as Figure 5, but uses a ViT-Large encoder with the other settings unchanged.Source: Figure 6 · p. 12