[WACV2023] Masked Image Modeling Advances 3D Medical Image Analysis

Masked Image Modeling Advances 3D Medical Image Analysis

摘要

1. 最近，掩码图像建模(MIM)由于其从大量未标记数据中学习的能力而获得了相当大的关注，并已被证明在涉及自然图像的各种视觉任务中是有效的。同时，由于大量的未标记图像以及质量标记的费用和难度，自监督学习在三维医学图像建模中的潜力是巨大的。然而，MIM在医学图像上的适用性仍然不确定。
2. 在本文中，我们证明了掩码图像建模方法除了自然图像外，还可以推进3D医学图像分析。我们研究了掩码图像建模策略如何从3D医学图像分割的角度来利用性能，作为一个具有代表性的下游任务:i)与naive对比学习相比，掩码 图像建模方法加速监督训练的收敛速度更快(1.40×)，最终产生更高的骰子分数;Ii)以较高的掩蔽率和相对较小的补丁大小预测原始体素值是医学图像建模的非平凡的自监督借口任务;iii)用于重建的轻量级解码器或投影头设计对三维医学图像上的掩模图像建模具有鲁棒性，从而加快训练速度并降低成本;iv)最后，我们还研究了MIM方法在应用不同图像分辨率和标签数据比的不同实际场景下的有效性。

1 引言

近年来，由于深度学习和硬件计算能力的技术进步，对对3D医学图像数据进行分析任务的深度神经网络的需求急剧增长。3D医疗体积图像在医疗保健领域显示出巨大潜力，有助于提高诊断患者病情的速度和准确性。例如，通过MRI/CT扫描正确、快速地发现和测量肿瘤病变对于疾病预防、早期发现和治疗方案优化至关重要，也将刺激更成功的临床应用的发展，最终改善患者的生活。然而，专家注释的高成本经常阻碍使用深度学习方法利用临床结果的进步。放射科医生对大规模3D医学图像的标注是有限的、昂贵的和耗时的。3D医学成像的另一个障碍是数据量，这是由3D图像维数和分辨率的增加所驱动的，导致了显著的处理复杂性。因此，在随机初始化的三维医学图像上训练深度学习模型需要繁重的计算和数据需求。

​ 作为一种可行的替代方法，自监督学习从数据本身获得监督信号，最近已被证明能够成功地解决对数据的需求，并能够学习输入的可泛化密集表示。在当前的方法中，掩码信号建模就是这样一个学习任务:掩码输入信号的子集，并尝试预测掩码信号。这种范式在NLP中非常成功，因为基于掩码语言建模任务的自监督学习算法在很大程度上彻底改变了该学科，证明了像BERT和GPT这样的大型模型可以在未标记的文本数据上学习，并适应于各种各样的应用。更重要的是，随着Vision transformer (ViT)的引入，使用协同运算应用掩码token并不直观的架构差距不再是一个障碍。遵循这一理念，基于掩码图像建模(MIM)的最新方法已经在可扩展视觉模型的开发中证明了其有效性。尽管取得了这些成就，基于掩码图像建模的算法在医学成像建模中受到的关注很少，其适用性还没有得到彻底的研究。自然，我们想知道是否掩码图像建模也将推动3D医学成像分析。在这项工作中，我们的目标是从以下几个方面来解决这个问题:

•对比学习已在一些研究中被证明能够学习利用3D分割和分类等下游任务的医学图像的一般表示A。在医学图像上比较掩码图像建模和对比学习方法(参见图1a进行说明)是值得的。

•自然图像是原始的、低水平的信号，具有显著的空间冗余度;恢复一些缺失的patch可以通过直接复制周围的patch来完成，而不需要对对象和场景有很高的了解。大多数背景组织与实体瘤的某些CT/MRI扫描相当，这使得模型更难了解病变区域的有用特征。因此，我们评估了几种掩蔽策略(掩蔽patch大小和掩蔽比)，以确定最有效的方法来促进对低级数据的整体理解，同时避免过度关注纹理和材料等特征。

•在实践中，医学图像分析用于各种不同的上下文中，具有不同数量的注释数据，可访问的未标记数据，甚至图像分辨率。因此，对我们来说，广泛分析这些元素如何影响相关的以及下游任务的性能也是至关重要的。

​ 本文研究了如何利用基于掩码图像建模的自监督学习来改进三维医学图像分析。它通过在两个真实世界的基准数据集上进行广泛的实验来实现这一点:多器官分割和脑肿瘤分割。

​ 我们的实验结果表明，掩码图像建模有利于3D医学图像建模，可以显著加快训练收敛速度(例如，最多节省1.4倍的训练成本才能达到相同的骰子分数)，并最终提高下游性能(例如，通过简单的训练配方，在两个分割方面都有超过5%的改进)。

2 相关工作

掩模图像建模：掩码图像建模是一种通过恢复被掩码损坏的图像来学习表示的自监督学习方法。它的发展与NLP中的传销任务一致，但在很长一段时间内仍处于主流之外。DAE是该领域的开创性工作，将掩蔽作为一种噪声类型。上下文编码器通过绘制源图像的一个大矩形区域来预测缺失的像素。最近基于Transformer的技术是由NLP的成功所激励的。iGPT将像素值分组到不同的簇中，对未知像素进行分类。ViT研究通过预测图像的平均颜色来研究自监督学习的掩码补丁预测。BEiT最近使用dVAE网络将像素值标记并预测为离散数。最近，MAE坚持原始像素恢复的精神，首次证明遮蔽高比例的输入图像可以产生非平凡和有意义的自我监督任务。它采用了自动编码器和轻量级解码器的设计，进一步降低了训练成本。SimMIM更进了一步，用一个线性投影层代替了整个解码器，得到了类似的结果。data2vec和CAE等方法在潜在表示空间中从可见patch到隐藏patch进行预测，试图使MIM成为自监督学习的通用框架。尽管如此，上面描述的技术只被证明对自然图像建模有用。在这项工作中，我们的目标是研究MIM方法是否也可以推进3D医学图像分析。

医学图像中的迁移学习：从自然图像中迁移学习被广泛应用于医学图像分析，无论图像统计量、尺度和任务相关特征是否存在差异。Raghu等和表明，来自ImageNet的迁移学习可以加速医学图像的收敛，这在医学图像训练数据有限的情况下尤其有用。使用领域特定数据的迁移学习也可以帮助解决领域差异问题。例如，表示对来自同一域的标记数据进行预训练后性能有所提高。然而，这种策略对于需要标记数据的各种医疗场景通常是不切实际的，这些数据收集成本高且耗时长。

Self-Supervised学习：自监督学习的早期工作侧重于从未标记的数据中学习表示，以便低容量分类器可以使用这些嵌入实现高精度。多年来，对比学习作为最流行和最广泛的自我监督学习策略之一受到了广泛的关注。它对两个或多个视图之间的图像相似性和不相似性(或仅相似性)进行建模，数据增强对于对比和相关方法至关重要。根据之前的几篇文献，自我监督学习也被应用于医学领域。例如，研究了特定领域的借口任务，而其他工作则侧重于根据医疗数据定制对比学习。特别是，Taleb等人深入研究了一系列用于3D医学成像的自监督学习策略。MICLe表明，在ImageNet上预训练的模型也可以推进皮肤病学图像分类。Tang等进一步将inpainting与对比学习结合起来进行医学细分。尽管所有这些方法都在医学成像中显示出了希望，但基于掩模图像建模的方法在这一学科中尚未得到实质性的研究。

3 方法

掩码图像建模方法，一般来说，掩码部分输入图像或编码图像标记，并鼓励模型重新创建掩码区域。许多现有的MIM模型采用编码器-解码器设计，然后是投影头，如BEiT和MAE。编码器有助于建模潜在特征表示，而解码器有助于将潜在向量重新采样到原始图像。投影头随后将编码或解码嵌入与掩码区域的原始信号对齐。值得注意的是，解码器组件已被建议以轻量级的方式设计，以便最大限度地减少训练时间。根据我们的经验，轻量级解码器降低了计算复杂性，增加了编码器学习更通用表示的能力，解码器可以快速掌握、翻译和传递。因此，而编码器是更关键的(只有编码器将被继承进行微调)，像SimMIM这样的方法通过用单个投影层消除整个解码器来简化架构。在这项工作中，我们深入研究了不同MIM模型在3D医学成像数据上的有效性。以下组件提供了更多详细信息:

3.1 掩码策略

在ViT之后，图像被划分为规则的不重叠的patch(例如，一个96× 96× 96的三维体将被划分为216个16× 16× 16的小块)，这些小块通常被认为是Vision Transformer的主要处理单元。在之前的文献中已经提出了多种随机掩蔽方法:1)InPainting引入了中心区域掩蔽策略;2) BEiT提出了一种复杂的分段屏蔽策略;3)最近的方法如MAE和SimMIM在研究不同的掩码斑块大小和掩码比时，采用了更直接的斑块级均匀随机掩码方法(分别见图1b和图1c)。许多随机屏蔽方案都是基于补丁的，因为在一个补丁一个补丁的基础上操作屏蔽更方便，其中一个补丁要么完全可见，要么被屏蔽。这些工作表明，采用高掩蔽比的均匀随机采样，可以有效地消除冗余，从而产生一种无法通过可见相邻补丁外推来轻松解决的自我监督任务。同时，均匀分布避免了潜在的中心偏差(即图像中心附近有更多的屏蔽补丁)。最后，稀疏输入允许开发一个高效的编码器，这将在下文中讨论。在这项工作中，为了简单和有效，我们还使用了随机贴片掩蔽方法。

3.2 编码器

编码器负责对被掩盖斑块的潜在特征表示进行建模，然后利用这些特征表示来预测被掩盖区域中的原始信号。习得的编码器应该能够适应广泛的视觉任务。我们在本文中考虑了各种架构，包括两个基本的视觉Transformer架构:vanilla ViT和SwinTransformer，以及一个注意力视觉网络VAN，它继承了注意力机制，以派生类似于SwinTransformer的层次表示，但使用纯卷积。所有模型都被重新实现为3D版本，以适应3D体积数据。我们将这些模型称为ViT3D、SwinTransformer3D和VAN3D。

3.3 解码器

对于遵循自动编码器设计来重建图像的方法，解码器采用整个已编码标记的集合，包括1)已编码的可见补丁和2)掩码标记。每个随机初始化的掩码token都是一个联合优化的可学习向量，以揭示掩码补丁。绝对位置嵌入或相对位置嵌入也应用于与主干体系结构相对应的掩码token。此外，所有被屏蔽的补丁对编码器来说都是不可见的，只有解码器才能看到所有的标记。这可以节省更多的计算和内存，同时不干扰训练。同时，解码器主干独立于编码器主干，编码器主干同样是可选的(见图1b)。默认情况下，我们遵循[23]并使用另一系列Transformer块进行解码。

3.4 重建目标

原始体素值预测： 对于3D医学图像，通过估计每个掩码标记的原始体素值来重建输入是简单而直观的。在体素空间中，恢复图像与原始图像之间的距离可以使用l1损失或l2损失的损失函数来计算。此外，损失只计算在被遮盖的补丁上，使得模型无法进行自我重建，这可能会主导学习过程，最终阻碍知识学习。值得注意的是，大多数vision Transformer拓扑将降低原始图像分辨率。对于三维医学图像，将96×体积分辨率降低到9×(即199*9≈768，使用vvi - base)，使用SwinTransformer或VAN，降低到3×。因此，对于vanilla ViT，我们应用单个线性投影层将潜在嵌入转换到原始体素空间;对于SwinTransformer和VAN，我们应用两层卷积转置将压缩嵌入上采样到原始分辨率。使用SimMIM和MAE分别重建TCIA-COVID19的三维肺部CT扫描，如图2和图3所示

其他的预测:许多早期的研究将屏蔽信号转换为集群或类，而不是原始像素值。例如，iGPT使用k-means将RGB值划分为512个簇，并鼓励模型预测每个像素属于哪个簇。BEiT[4]采用离散VAE (dVAE)将图像补丁转换为离散token。然后，预测目标基于令牌标识。另一方面，医学图像通常是稀疏的，体素值不是密集缩放的。用离散类目标替换原始信号可能会丢失细粒度纹理或材料信息。因此，为了简单和稳健性，我们在这项工作中专注于预测原始体素值。

4 3D分割上的实验

我们评估了两个单独的三维分割任务的蒙面图像建模方法，包括CT和MRI成像模式.

数据集：BTCV包括30名参与者，他们进行了腹部CT扫描，其中13个器官由翻译人员注释在范德比尔特大学医学中心的临床放射科医生的监督下。前24卷用于训练，6卷验证。BraTS包含387个多模态多位点MRI数据(FLAIR, T1w, T1gd, T2w)的训练集，带有胶质瘤分割的ground truth标签坏死/活动肿瘤和水肿用于脑瘤分割。此外，我们利用公共数据集TCIA-COVID19[20]，包括covid - 19感染患者的未增强胸部ct。共从661名患者中收集了771卷。在消融研究中，我们采用这个额外的无标记数据集进行自监督学习。有关数据集的更多信息和数据预处理细节，由于篇幅所限，可参考补充材料。

监督基线：

UNETR是一种U型编码器-解码器架构，用于医疗细分，采用ViT作为编码器骨干和遵循U-Net设计的卷积上采样解码器。它是医学影像分割领域中以Vision为骨干的SOTA模型之一。UNet - base表示vit - base作为编码器骨干。我们在消融研究中采用UNETRb作为默认的监督基线。对于其他产生分层特征的主干(SwinTransformer和VAN)，我们默认采用UPerNet作为解码头进行下游细分。利用骰子分数来评估分割的准确性

实现设置：所有的模型都在PyTorch中实现。我们使用MONAI进行数据转换和加载。在我们的消融实验中，我们使用ViTBase作为默认的编码器主干。对于器官分割的监督基线，我们使用了4个批量，AdamW优化器，以及0.0003的学习率，权重衰减为0.05(因为基于vit的架构非常庞大，很容易过拟合)，基于300个周期的线性预热和余弦退火调度器。训练在单个NVIDIA A10G GPU上进行，总共3000个epoch。对于脑肿瘤分割，在4个NVIDIA A10G gpu上进行1000 epoch的训练，batch size设置为8。我们使用100个周期的线性预热，优化器设置与器官分割兼容。我们的补充材料提供了更多的信息

4.1 不同方法的比较

我们首先评估1)掩码图像建模方法与对比学习方法的比较，2)不同的掩码图像建模方法如何使用MAE和SimMIM[56]以及传统的对比学习方法SimCLR进行比较。我们评估了一系列具有不同网络大小的编码器骨干，包括纯Vision Transformer，SwinTransformer和视觉注意网络(VAN)。对于MAE，我们使用了一个带有512-d的8层Transformer块作为解码器;对于SimMIM，我们使用单个线性层作为投影头。我们使用两层卷积转置作为投影头在Swin3D和VAN3D中使用预训练和UPerNet进行分割。在这次调查中，所有其他超参数都是相同的。此外，由于整个3D图像体积通常很难直接加载到GPU(内存爆炸)，我们采用了滑动窗口训练策略，其中原始图像被分为几个(96×96×96)小3D窗口。对于所有vit，默认情况下使用16的补丁大小。

从表1可以看出，蒙面图像建模方法总体上优于对比学习方法，MAE和SimMIM[56]的平均骰子分数都在0.752 ~ 0.758左右，而SimCLR的平均骰子分数在0.723左右，即比最佳方法低4.5%。表2中对brat的分割结果遵循类似的模式。掩模图像建模方法的平均骰子得分略大于0.80;而SimCLR得到的dice值为0.7739，比表1的最佳方法低4.37%。另一个值得注意的是，尽管两种MIM技术相似，SimMIM的性能比MAE略好，如表1和表2所示。对这种现象的一种解释是，即使编码器没有获得可推广的表示，高效的解码器(即使是轻量级解码器)也可能能够重建原始图像，从而周期性地缓解编码器学习更有效表示的动机。自监督学习的最终目标始终是学习数据的有效和可推广的表示，而不仅仅是自收敛。相比之下，SimMIM采用了更轻的设计，完全省略了解码器，这推动编码器执行更复杂的重建和学习任务

此外，掩码图像建模方法极大地提高了训练速度，降低了训练成本，如图4所示。基于simmim的架构在1.3k训练步可以获得1.76×更好的dice分数。此外，基于mim的方法可以达到0.7的骰子分数，训练时间比监督基线所需的训练时间少1.4倍。

4.2. 掩码策略

此外，我们还研究了不同掩码补丁大小和掩码比例对自监督学习性能的影响。表3和表4总结了几种MIM技术在细化分割方面的性能。i)与原始MAE文献一致，我们得出的结论是，较高的掩蔽比是一个非平凡的自监督学习工作，它将不断驱动模型构建可有效转移到下游任务的可泛化表示。例如，当在多个补丁尺寸上使用0.75的掩蔽比时，在多器官分割和脑肿瘤分割任务中获得了最佳的骰子分数(例如，表3中16号补丁尺寸为0.7183，表4中24号和32号补丁尺寸为0.8041)。ii)当与SimMIM(类似于MAE)联合使用时，高掩蔽比结合小补丁尺寸同样会导致相对较好的性能。由表3和表4可以看出，当patch大小为16时，模型的性能最优，其骰子得分分别为0.7249和0.8077。iii)但是，随着patch大小的增加，SimMIM方法对这个掩蔽比的敏感性降低。例如，当补丁大小为32时，模型可以获得最高的骰子得分，掩蔽比为0.15，最小的可掩模比。一种假设是，医学图像通常是原始的、低水平的信号，具有很大程度的空间冗余;恢复一些缺失的补丁可以通过直接复制附近的补丁来执行，而不需要对物体和周围环境的全面了解。单个小的掩码补丁无法充分掩盖复杂和交叉的结构或位置，但大的掩码补丁可能能够单独隐藏更重要的信号。因此，小补丁尺寸的高掩蔽比比大补丁尺寸的高掩蔽比更关键。

4.3 数据、分辨率、标签比率

在本节中，我们分析结果以解决以下三个问题:i)增加预训练数据量是否会提高下游性能?ii)不同的预训练解决方案如何影响下游知识转移?此外，iii)当使用不同数量的标记数据时，掩码图像学习方法如何提高性能?表5中的所有预训练都基于MAE架构，该架构使用ViTBase/16作为主干，掩蔽率为75%，如表3和表4所示。不同的标记比率表明我们使用不同百分比的BTCV CT扫描(例如，50% = 12张图像，100% = 24张图像)用于下游微调，而6张图像的验证集是一致的。在大多数监督学习案例中，更多的训练数据会导致性能的提高。鉴于大多数医学图像从下到上的逻辑是相似的，我们想知道这是否适用于自监督学习，特别是在使用MIM进行3D医学分析时，通过训练前数据的大小可以获得多少好处。我们采用多器官分割作为示例下游任务，并创建了两个不同的训练场景:一个同时使用covid - 19和BTCV数据集，另一个仅使用BTCV。表5显示了在更丰富的预训练数据上训练的模型优于在更少预训练数据上训练的模型的持续趋势(例如，0.7534→0.7183:提高4.9%，0.7338→0.7018:提高4.6%)。这种优势在较低的图像分辨率下更加明显，因为当只使用一半标记的数据时，0.6919比0.6552高出5.6%。

在表5中，我们还探讨了不同的预训练图像分辨率如何影响下游任务性能。直观地说，更高的预训练分辨率应该会产生更好的分割结果[1]，因为图像包含更细粒度的信息。在这里，我们使用不同的下采样比来表示原始信号在每个体积的所有维度上被压缩的程度。具体来说，双线性插值函数与MONAI的spacingd变换结合使用。从表5可以看出，高分辨率(1.5x, 1.5x, 2.0x)的预训练模型通常比低分辨率(2.0x, 2.0x, 2.0x)的预训练模型表现更好。例如，0.7338骰子分数比使用相同数据源和标记比率但使用更高分辨率预训练的结果低2.7%。在实际情况下，大多数医学图像，如CT/MRI扫描，由于标注成本高，没有标注。然而，公共数据是免费和丰富的;上述结果再次说明，在大数据集上进行预训练，然后用小样本进行微调是可行的。研究还表明，掩码图像学习可以显著提高下游任务在各种环境下的性能。

5 结论

本文通过对两个样本分割任务进行大量实验，演示了自监督学习中的掩模图像建模方法如何利用3D医学图像建模。我们展示了蒙面图像建模如何通过加速收敛和显著提高下游任务性能来优于传统的对比学习。我们还展示了蒙面图像建模方法如何在各种情况下用于推进三维医学图像建模。然而，几乎所有的医学图像都是弱标记的(例如，只有很少几行用于描述的文本)，而不是完全没有注释，这是一个悬而未决的问题，我们希望在未来进一步研究。我们感兴趣的是比较自监督学习和有限监督信号的监督学习。