Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation

TL;DR

The paper introduces MAE-based self pre-training of Vision Transformers on the target medical datasets to learn rich contextual representations without relying on large external pretraining data. By pre-training ViT encoders on three medical tasks and transferring them to task-specific heads (linear for classification and UNETR-style decoders for segmentation), the approach achieves consistent improvements over random and ImageNet pretraining, including 3D CT/MRI analyses. Ablation results highlight the importance of masked-patch reconstruction and reveal task-dependent optimal mask ratios. Overall, MAE self pre-training offers a practical, data-efficient pathway to boost medical image classification and segmentation with ViTs, applicable to both 2D and 3D modalities and small datasets.

Abstract

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

Figures (3)

• Figure 1: Segmentation Pipeline with MAE Self Pre-training. Left: A ViT encoder is first pre-trained with MAE. A random subset of patches is input to the encoder and a transformer decoder reconstruct the full image. Right: The pre-trained ViT weights are transferred to initialize the segmentation encoder. Then the whole segmentation network, e.g., UNETR, is finetuned to segment. A linear classifier can also be appended to an MAE self pre-trained ViT for classification tasks.
• Figure 2: MAE reconstruction. First row: Original image. Second row: Masked image where masked regions are colored with gray/black. Third row: Reconstructed images from unmasked patches. Every two columns show the results of one dataset, i.e., CXR, BTCV and BraTS from left to right.
• Figure 3: Qualitative Results of Segmentation. Results on BTCV are shown in the first two rows. In the first row, note the absence of the false positive segmentation (orange arrow) after MAE pre-training. In the second row, the stomach segmentation (red asterisk) is incomplete when created by the UNETR approach compared to an MAE pre-trained UNETR. Results on MSD BraTS are shown in the last two rows. In the third row, only subtle improvements are noticeable in the segmentation such as peripheral yellow necrotic core segmentations (white arrows) being captured after MAE pre-training. In the fourth rows, necrotic core segments are nearly absent without MAE pre-training.