医学图像分割与分类的基础模型 VISION-MAE

VISION-MAE: A Foundation Model for Medical Image Segmentation and Classification

2402.01034 (arxiv.org)

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 multimodal 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 a novel foundation model, VISION-MAE, specifically designed for medical imaging. Specifically, VISION-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. VISION-MAE has high label efficiency, outperforming several benchmark models in both in-domain and out-of-domain applications, and achieves high performance even with reduced availability of labeled data. This model 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.

人工智能（AI）具有彻底改变医学成像中诊断和分割技术的潜力。然而，其开发和临床实施面临多重挑战，包括数据可用性有限、缺乏通用性，以及需要有效整合多模态数据。基础模型是一种大规模的预训练AI模型，提供了一个多功能的基础，可以适应各种特定任务和上下文。在此，本文介绍了一种专为医学成像设计的新型基础模型VISION-MAE。具体来说，VISION-MAE使用自监督学习技术，在包含250万张来自不同模态（CT、MR、PET、X光和超声）的未标记图像的数据集上进行训练。然后，它使用明确的标签适应于分类和分割任务。VISION-MAE具有高标签效率，在域内和域外应用中均优于多个基准模型，即使在标记数据可用性降低的情况下也能实现高性能。该模型代表了医学成像AI领域的重大进展，为改善分割和分类任务提供了一种通用且稳健的解决方案，同时减轻了数据标注的工作负担。

Introduction

人工智能（AI）领域的最新进展催生了基础模型的发展，这些机器学习模型在大规模、多样化的数据集上进行训练，可以适应各种下游任务。虽然传统的深度学习模型是专为特定应用（如间质性肺病或COVID-19的分类）而训练的，并且在重新用于其他任务时表现不佳，但基础模型提供了更通用和可适应的能力。在初始训练后，基础模型能够针对多种不同任务进行微调。在医学成像领域，基础模型具有广泛的应用潜力，如提高诊断准确性、提高阅片效率以及预测疾病结果。将基础模型纳入临床工作流程有可能彻底改变医疗保健服务的领域。

Segment Anything Model（SAM）是一种在包含1100万张自然图像和超过10亿个图像掩码的数据集上训练的基础模型，展示了通过提示自动分割任何图像的能力。随后，随着MedSAM的发布，这种能力被应用于医学图像，MedSAM是在一个包含超过100万对医学图像-掩码对的大规模医学成像数据集上训练的，涵盖了广泛的模态和协议，从而为医学成像分析领域打开了众多可能性。自MedSAM发布以来，已有越来越多的医学成像基础模型被报道用于各种应用。例如，RETFound是一种在130万张视网膜图像上训练的基础模型，具有检测多种视网膜疾病的通用能力。其他医学基础模型能够执行运动结构的分割、体积分割、眼部疾病的诊断和预后，以及临床病理学图像的评估。尽管取得了这些进展，但在基础模型的开发中仍然存在一些挑战。此外，通过传统的监督学习方法训练基础模型受到需要大量标记数据的限制，这通常是一个耗时且成本高昂的过程。此外，高质量和开源的医学成像数据普遍缺乏，进一步阻碍了开发模型的通用性和准确性。

为了应对这些挑战，当前的医学成像分析方法采用了自监督学习（SSL）的方法，这是一种训练模型通过未标记数据的潜在特征来推断标签的方法。SimCLR和Masked Autoencoder等SSL方法已被证明能够取得与监督学习方法相当甚至更好的结果。最近，Swin MAE作为一种使用Swin Transformer作为骨干的掩码自编码器，展示了在较小数据集上无需使用预训练模型即可实现此类结果的可行性。

鉴于医学影像基础模型的潜力以及近期在自监督学习（SSL）方面的进展，本文提出了VISION-MAE基础模型，该模型采用基于Swin Transformer的掩码自编码器，可用于多种下游任务。本文的模型在包含250万张临床图像的数据集上进行训练，这些图像涵盖CT、MRI、PET/CT、放射学（X光）和超声（US）等多种模式。本文已将VISION-MAE应用于一系列复杂的分割和分类任务，并通过使用特定任务标签进行微调来优化其性能。

本文的评估过程首先从九个不同任务的分割开始，这些任务涵盖了各种成像模式、解剖结构和多种疾病。这些任务包括CT和MRI中的腹部器官分割、MRI中的心脏结构分割、超声中的乳房和甲状腺区域分割、前列腺MRI以及中风影响区域、胶质瘤和皮肤病变的分割。此外，本文还在五个不同的分类任务中评估了VISION-MAE。这些任务包括COVID-19的诊断、结节病的识别、14种不同肺部疾病的检测、骨关节炎严重程度的评估以及前交叉韧带（ACL）撕裂的评估。这种全面的方法确保了对VISION-MAE在广泛医学影像挑战中的能力进行彻底检查。

本文证明，在分割任务中，提出的模型能够与其他最先进的模型（如nnU-Net和TransUNet）相比达到相当或更好的性能。在分类任务中，与采用监督学习策略的方法（如使用RadImageNet和ImageNet预训练权重的方法）相比，VISION-MAE展示了相似或增强的有效性。总体而言，本文展示了VISION-MAE SSL基础模型能够在各种应用中提高性能和效率，提供了一种可应用于多种成像模式和解剖结构的通用工具。

图1提供了VISION-MAE开发和应用的全面概述。在开发阶段，本文编制了一个包含2,486,425张图像的大型数据集，这些图像跨越五种成像模式。该数据集包括1,199,904张MR图像、570,943张CT图像、657,313张PET图像、211,326张超声图像和438,521张X光图像，这些图像均来自2005年至2022年期间纽约市一家门诊放射科诊所。在使用这些放射学图像进行自监督学习后，本文评估了VISION-MAE在各种分割和分类任务中的性能和泛化能力。此次评估涉及9个公开可用的分割数据集和5个分类数据集。特别地，对于使用PET/MR成像进行结节病预测的任务，由于PET/MR数据集的稀缺性，本文使用了2017年至2022年间在西奈山医院收集的另一个队列的数据集。

Figure 1. 研究设计概述。首先，使用基于掩码的SSL策略收集了一个庞大的上游队列，用于开发自监督学习基础模型VISION-MAE。其次，利用VISION-MAE的预训练权重，针对两种医学成像应用（分类和分割）进行了优化。在分类任务中，使用VISION-MAE的预训练编码器权重对模型的骨干网络进行了微调。在分割任务中，VISION-MAE的编码器和解码器权重都在下游应用中进行了微调。这种方法利用了VISION-MAE学习到的特征，以提高在特定医学成像任务中的性能。

本文评估了VISION-MAE在医学图像分割任务中与两个领先模型（nnU-Net和TransUNet）以及在分类任务中与两个监督学习模型（RadImageNet和ImageNet）的性能和标签效率。所有这些模型都采用了不同的预训练权重进行分类，但在下游应用中使用了相同的架构框架和微调过程。nnU-Net和TransUNet的配置采用了它们原始出版物中推荐的默认参数。有关参数设置的详细信息，请参阅方法部分和补充材料。RadImageNet是一个全面的放射学数据集，包含135万张图像，跨越165个病理标签和14个解剖区域，涵盖CT、MR和超声三种成像模式。另一方面，ImageNet则通过在ImageNet-1k上进行监督学习的预训练来采用传统的迁移学习，ImageNet-1k包含约130万张跨越1000个类别的自然图像。TransUNet也利用了来自ImageNet-1k的预训练权重。此外，本文还将VISION-MAE与SimCLR进行了比较，SimCLR是VISION-MAE框架内的另一种流行的SSL策略。SimCLR采用基于对比学习的方法，并在分割任务中保持与VISION-MAE相同的架构和参数。本文的评估涵盖了9个分割数据集和5个分类数据集。对于每个包含少于4000张图像的数据集，本文进行了五折交叉验证。本文使用接收者操作特征曲线（ROC）下的面积（AUC）来衡量分类任务的模型性能，使用Dice相似系数（DSC）来衡量分割任务的模型性能。本文还进行了双侧t检验，以确定VISION-MAE与每个比较模型在每个任务上的差异是否具有统计学意义。

Method

上游数据集

本文用于开发VISION-MAE的数据集是从2005年至2022年间纽约市一家私人放射科诊所收集的，共包含2,486,425张图像，涵盖五种成像模式：磁共振成像（MR，1,199,904张）、计算机断层扫描（CT，570,943张）、正电子发射断层扫描/计算机断层扫描（PET/CT，65,731张）、X射线（438,521张）和超声（211,325张）。这一数据集是在本文之前的RadImageNet项目基础上构建的，后者包含了135万张MR、CT和超声图像。在此扩展数据集中，本文纳入了更多成像模式和解剖区域。MR、CT和PET/CT图像经过精心挑选，由阅片放射科医生选择了显示主要病理的图像。对于X射线和超声，本文纳入了每个研究中的所有图像。数据集的年龄分布特征为均值45岁，标准差5岁。

下游数据集

对于下游任务，本文在9个公开可用的分割数据集上评估了VISION-MAE的预训练权重，这些数据集包括：包含13个腹部器官标注的BTCV-Abdomen数据集（2178个CT切片）、包含5个病变标签的stroke数据集（4656个非对比增强CT切片）、包含左心室、右心室和心外膜标注的ACDC心脏MR数据集（2978个切片）、包含12个腹部器官标注的AMOS腹部MR数据集（2476个切片）、包含三个标签的prostate数据集（3554个MR切片）、包含胶质瘤异常掩码的brain segmentation数据集（1373个MR切片）、breast segmentation数据集（647张超声图像）、thyroid segmentation数据集（17641张超声图像）以及ISIC 2016皮肤病变分割数据集（1279张皮肤镜图像）。

此外，本文还在5个分类任务上评估了VISION-MAE：包含9050张胸部CT图像的COVID-19分类数据集、包含1231张PET/MR心脏图像以分类结节病的内部结节病数据集、包含1021张MR膝关节图像的前交叉韧带（ACL）数据集、包含4个标签的1650张膝关节X射线图像的膝关节骨关节炎数据集，以及包含14个肺部类别的112,120张X射线图像的NIH Chest X-ray数据集。对于图像数少于4000张的数据集，本文实施了五折交叉验证，在此过程中，数据以9:1的比例分为训练集和验证集。对于更大的数据集，本文分配了72%的图像用于训练，8%用于验证，20%用于测试。这种方法确保了模型评估的稳健性和在各种医学成像任务中的泛化能力。

VISION-MAE的开发

在开发VISION-MAE时，本文构建了一个基于Swin Transformer驱动的掩码自编码器架构的基础模型，该模型包括编码器和解码器。在预处理阶段，数据集中的每张图像都被调整大小并分割成4x4的块，然后随机选择75%的块进行掩码处理。编码器被设计用于处理掩码块和非掩码块，并输出高级特征表示。该编码器采用了修改后的微型Swin Transformer（Swin-T），包含四个层，每层都配备了块合并块和两个Swin Transformer块。每个Swin Transformer块的结构包括LayerNorm（LN）层、多头自注意力模块以及随后的两个MLP（多层感知器）层和一个残差连接。解码器的作用是基于这些高级图像特征进行扩展，并重建块以形成完整的图像。由于编码器和解码器的权重都旨在在下游应用中使用Swin-UNet架构进行微调，因此VISION-MAE中的解码器也采用了Swin Transformer块。这使得通过最终的线性投影进行图像重建成为可能。解码器架构与Swin Transformer框架的一致性至关重要，确保了将VISION-MAE的预训练权重应用于各种下游医学成像任务时的无缝过渡和兼容性。

适应下游任务

在将VISION-MAE模型适应于各种下游任务的过程中，本文根据其编码器和解码器预训练权重的选择和应用，针对每个任务的具体要求进行了修改。对于涉及放射学图像的分割任务，本文对VISION-MAE的编码器和解码器预训练权重都进行了微调。而对于处理彩色医学图像的分割任务，本文仅利用了VISION-MAE的编码器部分，并在下游模型的解码器中初始化随机权重。在处理分类任务时，本文采用了VISION-MAE编码器中的预训练权重，并补充了一个随机初始化的分类头层。

下游任务架构配置

下游任务的架构配置与上游VISION-MAE模型略有不同。具体而言，下游模型中的编码器遵循了具有四层的Swin-T架构，分别包含2、2、6和2个Swin Transformer块。值得注意的是，在下游模型的第三层中，两个块使用了VISION-MAE编码器中的预训练权重，而其余四个块则进行了随机初始化。此外，分割模型的解码器部分加入了跳跃连接，这是一种旨在提高生成预测掩码准确性和质量的策略性增强。这种在模型架构上的定制方法确保了下游应用能够针对其特定任务进行优化配置，有效地利用了VISION-MAE的基础优势，同时适应了每个应用的独特需求。

模型训练

针对上游VISION-MAE模型的训练方案进行了精心设计，为每种成像模式（如MR、CT和PET/CT、超声和X射线）开发了不同的预训练权重。这些预训练权重均使用640的批量大小在800个周期内进行了训练。本文采用了AdamW优化器，并在此过程中使用了8个NVIDIA DGX A100 GPU。初始学习率设置为0.0001，并随着训练周期的进行，遵循半周期余弦函数递减。此外，训练过程中还包含了一个10个周期的初步预热阶段。

在开发下游分割任务模型时，本文将预热期标准化为40个周期，总训练周期为150个。基本学习率和批量大小是可调参数，在测试集上进行最终评估之前，先在验证集上进行微调以优化性能。对于开发下游分类任务模型，训练包括在总共50个训练周期内的10个预热周期，基本学习率为0.001。在这里，批量大小是唯一需要微调的变量参数。所有下游应用均一致使用了AdamW优化器。分割模型的训练监控采用了双指标方法，结合了Dice系数和交叉熵损失。相比之下，分类模型的训练监控则仅通过交叉熵损失进行。

基准模型

为了严格评估VISION-MAE的性能，本文在分割和分类任务中与其他基准模型进行了比较分析。对于分割任务，本文采用了三种不同的策略来开发和评估模型：nnU-Net、TransUNet和SimCLR。在nnU-Net模型的情况下，本文将其配置为与VISION-MAE的数据分布相匹配，并进行了相同数量的训练周期。nnU-Net以其独特的预处理方法来管理训练过程而脱颖而出。另一方面，TransUNet模型则使用默认的随机梯度下降（SGD）优化器进行训练，学习率设置为0.01，共训练150个周期。SimCLR预训练权重是使用与VISION-MAE相同的上游数据，通过SimCLR自监督学习策略，采用Swin Transformer架构开发的。然后，针对下游任务，SimCLR预训练模型在训练参数上与VISION-MAE保持一致，以确保公平和一致的比较。

对于分类任务，本文将VISION-MAE与使用了监督学习Swin Transformer预训练权重的模型进行了基准比较，特别是那些在RadImageNet和ImageNet数据集上训练的模型，以及自监督学习的SimCLR。这些预训练模型使用与VISION-MAE在下游任务中相同的参数进行了微调。

统计方法

对于分割任务，本文采用Dice分数作为主要评价指标。为了建立Dice分数结果的可靠性，本文使用包含1000次重采样的自助法计算了95%置信区间（CI）。在分割任务的统计分析中，本文使用了配对t检验来确定最具竞争力的模型，确保性能差异不仅明显，而且具有统计学意义。分类任务则通过接收者操作特征曲线下面积（AUROC）进行评价。在二分类领域，本文采用了DeLong检验来确定不同模型在各个标签上性能差异的95%置信区间和统计学意义。对于涉及多类和多标签分类的更复杂场景，本文再次使用了包含1000次重采样的自助法。该方法用于计算95%置信区间并确定不同模型之间的p值。