Source-linked AI summary
Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation
Lei Zhou, Huidong Liu, Joseph Bae, Junjun He, Dimitris Samaras, Prateek Prasanna
TL;DR
Medical image analysis needs contextual representations, but suitable large-scale medical pre-training data is difficult to acquire. The paper pre-trains a ViT with MAE on each target task’s own training set, then fine-tunes task-specific heads, and reports improved classification and segmentation across chest X-ray, abdominal CT, and MRI brain tumor tasks.
Problem
Suitable ImageNet-scale medical image data for pre-training is unavailable, while medical image analysis depends on contextual relationships among anatomical structures.
Method
The method applies MAE to the target dataset’s training set, then fine-tunes a ViT with a linear classifier or UNETR-based segmentation decoder.
Results
MAE self pre-training significantly improves classification and segmentation over random initialization and surpasses ImageNet pre-training across the evaluated datasets.
Takeaways & Limitations
MAE self pre-training improves medical image analysis across classification, 2D imaging, and 3D CT and MRI segmentation, including small datasets.
Abstract
from arXiv · showhide
Masked Autoencoder (MAE) has recently been shown to be effective in pre-training Vision Transformers (ViT) for natural image analysis. By reconstructing full images from partially masked inputs, a ViT encoder aggregates contextual information to infer masked image regions. We believe that this context aggregation ability is particularly essential to the medical image domain where each anatomical structure is functionally and mechanically connected to other structures and regions. Because there is no ImageNet-scale medical image dataset for pre-training, we investigate a self pre-training paradigm with MAE for medical image analysis tasks. Our method pre-trains a ViT on the training set of the target data instead of another dataset. Thus, self pre-training can benefit more scenarios where pre-training data is hard to acquire. Our experimental results show that MAE self pre-training markedly improves diverse medical image tasks including chest X-ray disease classification, abdominal CT multi-organ segmentation, and MRI brain tumor segmentation. Code is available at https://github.com/cvlab-stonybrook/SelfMedMAE
1. INTRODUCTION
The paper proposes MAE self pre-training for medical image analysis, motivated by anatomical interdependencies and the need for contextual representations when suitable medical pre-training data is scarce. A ViT is pre-trained by reconstructing masked image regions, then adapted to classification or segmentation, improving performance across evaluated tasks.
- Medical image analysis must account for functional and mechanical relationships among anatomical structures and surrounding regions.Examples include cardiac and mediastinal changes informing chest X-ray pathology classification and tumor-associated tissue changes aiding brain tumor segmentation.
- Masked image modeling trains a ViT to infer hidden patches by aggregating contextual information from visible image regions.MAE uses a visible-token ViT encoder and a lightweight decoder that reconstructs masked patches with mean square error.
- The paper proposes MAE self pre-training by applying MAE to the downstream task’s own training set.This approach targets settings where suitable pre-training data is difficult to acquire and unifies pre-training and fine-tuning domains.
- The pre-trained ViT is equipped with task-specific heads and fine-tuned for chest X-ray classification, abdominal CT segmentation, and MRI brain tumor segmentation.Classification uses a linear head, while segmentation follows the UNETR decoder design.
- MAE self pre-training significantly improves medical image classification and segmentation over random initialization and surpasses ImageNet pre-training across the evaluated datasets.The reported experiments cover diverse medical image tasks rather than a single modality or task.
2. METHODOLOGY
The methodology uses a ViT-based MAE pipeline that reconstructs images from visible patches, then transfers the encoder to task-specific classification or segmentation networks. The downstream experiments include UNETR-based abdominal multi-organ segmentation and comparative evaluation against alternative initialization strategies.
- Vision Transformer: The ViT converts 3D medical volumes into flattened patch sequences, adds positional embeddings, and processes them with Transformer blocks.The method uses sine-cosine positional embeddings during MAE pre-training and initializes downstream learnable positional embeddings from them.
- Masked Autoencoder: MAE randomly separates patches into visible and masked groups, encodes only visible patches, and reconstructs the full input using decoder mask tokens.Encoding partial observations encourages the representation to capture global information needed to restore masked content.
- Downstream Architectures: After pre-training, a linear classifier supports multi-label chest X-ray prediction, while a UNETR decoder uses multi-resolution encoder features for 3D segmentation.The segmentation decoder reshapes encoder representations and repeatedly upsamples and concatenates them with shallower features.
3. EXPERIMENTS AND RESULTS
Experiments evaluate MAE self pre-training across chest X-ray classification, abdominal CT segmentation, and brain tumor MRI segmentation, including qualitative reconstruction and ablation analyses. Results show improvements over random initialization and ImageNet pre-training, including on small medical datasets.
- Datasets and implementation: The study evaluates chest X-ray disease classification, BTCV abdominal multi-organ segmentation, and MSD brain tumor segmentation.ChestX-ray14 uses multi-class AUC; BTCV and MSD report average DSC and 95% HD.
- MAE reconstruction: MAE reconstructs masked image regions from random context, while visible patches remain blurrier because reconstruction loss applies only to masked patches.The authors emphasize downstream task performance rather than high-quality reconstructions.
- Qualitative segmentation: Qualitative segmentation shows fewer false positives, more complete stomach masks, and improved necrotic-core capture after MAE pre-training.The figure reports these changes for BTCV and MSD BraTS examples.
- Lung disease classification: 0.8%: MAE self-pre-training outperforms ImageNet-pre-trained ViT on lung disease classification.The paper also reports that MAE self-pre-trained ViT outperforms the compared CNN-based methods.
- Abdominal multi-organ segmentation: 83.5% average DSC: MAE self-pre-training improves the UNETR baseline from 78.8% on BTCV multi-organ segmentation.The method also surpasses ImageNet pre-training, including on the small BTCV dataset of 30 subjects.
4. CONCLUSION
The study concludes that MAE pre-training improves classification and segmentation across diverse medical imaging tasks, including small datasets and 3D CT and MRI data. It identifies prognosis and outcome prediction as future evaluation targets.
- Conclusion: MAE pre-training improves classification and segmentation performance across diverse medical image analysis tasks.The conclusion covers the study’s evaluated task range without claiming state-of-the-art performance.
- Conclusion: MAE self pre-training outperforms existing methods on small datasets, including ImageNet-transfer learning.The conclusion highlights this result as important for medical imaging tasks.
- Conclusion: The study demonstrates MAE effectiveness on 3D medical images including CTs and MRIs.Future work will test MAE pre-training for prognosis and outcome prediction.
5. COMPLIANCE WITH ETHICAL STANDARDS
The retrospective study used open access images and required no additional approval.
- Compliance with ethical standards: The research was conducted retrospectively using open access images, so additional approval was not required.