Self Pre-training with Adaptive Mask Autoencoders for Variable-Contrast 3D Medical Imaging

TL;DR

This paper introduces a 3D adaptive masked autoencoder architecture for pre-training Vision Transformers on variable-contrast MRI data, improving infarct segmentation accuracy in medical imaging.

Contribution

It proposes a novel adaptive Masked Autoencoder model that handles varying numbers of input contrasts in 3D medical images, enhancing pre-training effectiveness.

Findings

Pretraining with AMAE improves infarct segmentation by 2.8%-3.7%.

The method effectively handles variable input contrasts in MRI data.

Self pre-training enhances model performance in medical image analysis.

Abstract

The Masked Autoencoder (MAE) has recently demonstrated effectiveness in pre-training Vision Transformers (ViT) for analyzing natural images. By reconstructing complete images from partially masked inputs, the ViT encoder gathers contextual information to predict the missing regions. This capability to aggregate context is especially important in medical imaging, where anatomical structures are functionally and mechanically linked to surrounding regions. However, current methods do not consider variations in the number of input images, which is typically the case in real-world Magnetic Resonance (MR) studies. To address this limitation, we propose a 3D Adaptive Masked Autoencoders (AMAE) architecture that accommodates a variable number of 3D input contrasts per subject. A magnetic resonance imaging (MRI) dataset of 45,364 subjects was used for pretraining and a subset of 1648 training, 193 validation and 215 test subjects were used for finetuning. The performance demonstrates that self pre-training of this adaptive masked autoencoders can enhance the infarct segmentation performance by 2.8%-3.7% for ViT-based segmentation models.