[论文解读] Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation

Lei Zhou, Huidong Liu+4 more|arXiv (Cornell University)|Mar 10, 2022

Radiomics and Machine Learning in Medical Imaging被引用 8

一句话总结

本文提出一种自监督掩码自编码器（MAE）预训练方法，仅使用目标数据集的训练数据进行医学图像分析，无需依赖ImageNet等外部数据集。通过掩码12.5%至75%的3D图像块并利用视觉Transformer进行重建，该方法提升了上下文特征学习能力，在胸部X光片分类、腹部CT分割和MRI脑肿瘤分割任务中均达到最先进性能，显著优于随机初始化和ImageNet预训练方法。

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

研究动机与目标

• 为解决医学影像中缺乏大规模、领域特定预训练数据的问题，提出一种仅使用目标数据集训练数据的自监督预训练范式。
• 探究掩码自编码是否能增强医学图像中的上下文特征学习，其中解剖结构在功能和机械上相互关联。
• 在低数据环境下，评估MAE预训练在多样化3D医学影像任务（包括分类与分割）中的有效性。
• 比较MAE自预训练与ImageNet预训练及随机初始化在医学图像基准下游性能上的表现。
• 确定医学图像任务中最佳超参数（如掩码比率和预训练周期数）

提出的方法

• 该方法采用视觉Transformer（ViT）主干网络，并使用正弦-余弦位置编码，以避免在预训练期间与MAE重建产生干扰。
• 采用掩码自编码器架构，其中12.5%至75%的3D图像块被随机掩码，模型仅利用可见块和可学习掩码标记进行重建。
• 编码器仅处理可见块并生成块级表示，而轻量级解码器则基于这些表示和掩码标记重建完整图像。
• 模型通过在掩码块上使用均方误差损失进行训练，促使网络利用周围块的上下文信息推断缺失区域。
• 预训练完成后，将编码器权重用于初始化下游模型：分类任务使用线性分类器，分割任务使用类似UNETR的解码器。
• 微调阶段采用逐层权重衰减和余弦学习率调度，并为每个数据集配置特定任务的超参数。

Fig. 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 enc

实验结果

研究问题

• RQ1 仅在目标数据集上进行自监督掩码自编码预训练，是否能提升医学图像分类与分割任务的性能？
• RQ2 在小样本医学图像数据集上，MAE预训练是否优于ImageNet预训练？
• RQ3 在医学图像分割任务中，特别是针对3D CT和MRI数据，MAE预训练的最佳掩码比率是多少？
• RQ4 自预训练如何影响模型捕捉解剖连接区域（如肿瘤微环境）中上下文依赖关系的能力？
• RQ5 在医学MAE预训练中，位置编码的选择（正弦-余弦 vs. 可学习）是否会影响重建质量与下游性能？

主要发现

• 在BTCV腹部多器官分割数据集上，MAE自预训练将平均Dice相似系数（DSC）从基线的78.8%提升至83.5%，显著优于ImageNet预训练。
• 在BraTS MRI脑肿瘤分割任务中，MAE预训练使DSC从77.4%提升至78.91%，HD95从7.78mm降至7.22mm。
• 在CXR14胸部X光片分类任务中，MAE自预训练性能较ImageNet预训练提升0.8%，超越了采用自监督预训练的最先进CNN模型。
• 最佳分割性能在12.5%掩码比率下实现，而更高比率（如75%）效果较差，表明需针对任务进行优化。
• 较长的预训练周期可提升性能至一定限度，但在100–150个周期后出现过拟合，尤其在小样本数据集上。
• 与可学习位置编码相比，预训练期间使用正弦-余弦位置编码能提升重建质量与下游性能。

Fig. 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.