[论文解读] VIS-MAE: An Efficient Self-supervised Learning Approach on Medical Image Segmentation and Classification

Zelong Liu, Andrew Tieu+10 more|arXiv (Cornell University)|Feb 1, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

VIS-MAE 是一种基于自监督视觉变换器的奠基模型，在涵盖 CT、MR、PET、X 射线和超声等多种模态的 250 万张未标注医学影像上进行预训练。该模型在医学图像分割和分类任务中实现了最先进性能，具有极高的标签效率，即使在仅使用 50%–80% 的标注数据进行微调时，仍优于各类基准模型。

ABSTRACT

Artificial Intelligence (AI) has the potential to revolutionize diagnosis and segmentation in medical imaging. However, development and clinical implementation face multiple challenges including limited data availability, lack of generalizability, and the necessity to incorporate multi-modal data effectively. A foundation model, which is a large-scale pre-trained AI model, offers a versatile base that can be adapted to a variety of specific tasks and contexts. Here, we present VIsualization and Segmentation Masked AutoEncoder (VIS-MAE), novel model weights specifically designed for medical imaging. Specifically, VIS-MAE is trained on a dataset of 2.5 million unlabeled images from various modalities (CT, MR, PET,X-rays, and ultrasound), using self-supervised learning techniques. It is then adapted to classification and segmentation tasks using explicit labels. VIS-MAE has high label efficiency, outperforming several benchmark models in both in-domain and out-of-domain applications. In addition, VIS-MAE has improved label efficiency as it can achieve similar performance to other models with a reduced amount of labeled training data (50% or 80%) compared to other pre-trained weights. VIS-MAE represents a significant advancement in medical imaging AI, offering a generalizable and robust solution for improving segmentation and classification tasks while reducing the data annotation workload. The source code of this work is available at https://github.com/lzl199704/VIS-MAE.

研究动机与目标

• 为解决标注医学影像数据有限的问题，开发一种可泛化的奠基模型。
• 提升模型在多样化医学影像模态（CT、MR、PET、X 射线、超声）中的泛化能力。
• 提升下游分类与分割任务中的标签效率。
• 在预训练阶段无需人工标注，实现有效的多模态表征学习。
• 为临床医学影像应用提供一种稳健且可迁移的人工智能基础模型。

提出的方法

• VIS-MAE 采用掩码自编码（MAE）框架，将输入医学图像的随机图像块进行掩码处理，模型通过可见图像块重建被掩码的区域。
• 模型在包含 250 万张未标注医学影像的大型多模态数据集上进行预训练。
• 采用视觉变换器主干网络，以捕捉医学图像中的长距离依赖关系与空间层次结构。
• 在微调阶段，仅使用部分标注数据，将模型适配至分类与分割任务。
• 架构中引入模态特定的归一化与数据增强策略，以提升在不同影像类型下的鲁棒性。
• 自监督预训练使模型在微调前即可学习到可泛化的特征。

实验结果

研究问题

• RQ1 在多样化、未标注的医学影像上预训练的自监督奠基模型，是否能提升下游分割与分类任务的性能？
• RQ2 当使用有限标注数据进行微调时，VIS-MAE 的性能与现有预训练模型相比如何？
• RQ3 VIS-MAE 在不同影像模态与临床任务中的泛化能力达到何种程度？
• RQ4 在医学影像上使用掩码自编码是否能带来更鲁棒且可迁移的表征？
• RQ5 VIS-MAE 是否能在不牺牲准确性的前提下，减轻医学影像人工智能中的标注负担？

主要发现

• VIS-MAE 在多个医学图像分割与分类基准测试中达到最先进性能，无论在域内还是域外设置下，均优于现有模型。
• 即使仅使用 50% 的标注数据进行微调，模型仍保持高性能，展现出卓越的标签效率。
• VIS-MAE 在包括 CT、MR、PET、X 射线和超声在内的多种影像模态中表现出强大的泛化能力。
• 与使用有限数据进行监督预训练相比，在 250 万张未标注图像上进行自监督预训练，显著提升了下游任务的准确率。
• 模型在分类与实例分割任务中均取得具有竞争力的结果，表明其具有广泛的应用潜力。
• 源代码已公开，支持可复现性，并推动医学人工智能领域的进一步研究。