MiM: Mask in Mask Self-SupervisedPre-Training for 3D Medical Image Analysis

发布于:2025-05-09 ⋅ 阅读:(14) ⋅ 点赞:(0)

Abstract

Vision Transformer在3D医学图像分析的自监督学习(Self-Supervised Learning,SSL)中展现了卓越的性能。掩码自编码器(Masked Auto-Encoder,MAE)用于特征预训练,可以进一步释放ViT在各种医学视觉任务中的潜力。然而,由于3D医学图像具有更大的空间尺寸和更高的维度,缺乏层次化设计的MAE可能会阻碍下游任务的性能。在本文中,我们提出了一种新颖的3D医学图像掩码内掩码(Mask in Mask,MiM)预训练框架,旨在通过从不同尺度的层次化视觉标记中学习区分性表示,来提升MAE的性能。我们引入了多个粒度级别的掩码输入,这些输入来自体数据(volume),并且同时在精细和粗糙级别上进行重建。此外,我们应用了一种跨层次对齐机制,对相邻层次的体数据进行对齐,以强制执行层次化的解剖学相似性。此外,我们采用了一种混合骨干网络,以在预训练期间高效地增强层次化表示学习。MiM在大规模可用的3D体图像上进行了预训练,例如包含各种身体部位的计算机断层扫描(Computed Tomography,CT)图像。在十二个公共数据集上进行的广泛实验表明,MiM在器官/肿瘤分割和疾病分类方面优于其他自监督学习方法。我们进一步将MiM扩展到包含超过1万个体数据的大规模预训练数据集上,结果表明大规模预训练可以进一步提升下游任务的性能。代码和检查点将在被接受后提供。

索引词——CT、自监督学习、分割、分类、3D医学图像。

I. INTRODUCTION

深度学习的出现推动了医学图像分析领域的前所未有的进步,尤其是在监督学习范式中。然而,这一范式的基本局限性在于其对大量标注训练数据的依赖,这对于3D医学图像来说尤其具有挑战性,因为标注需要大量的专业知识、时间和资源 [1]–[6]。自监督学习(SSL)作为一种变革性解决方案应运而生,它提供了一种强大的机制,可以从无标注数据中学习到鲁棒且可迁移的表示,从而显著提升下游任务的性能 [7]–[16]。这种范式转变不仅推动了医学图像分析领域的重大进步,还大幅减轻了标注的负担 [2], [4]。

在自监督学习(SSL)领域中,掩码自编码器(Masked Autoencoders,MAE)[17] 通过其创新的通过视觉变换器(Vision Transformer)架构重建失真视图的方法,在自然图像分析中取得了显著的成功 [17], [18]。然而,将这些方法扩展到3D医学成像领域带来了独特的挑战,因为体积数据具有高维度和跨越多个尺度的复杂解剖结构等固有复杂性。最近的研究进展在解决这些挑战方面取得了重大突破。MAE3D [19] 通过重建裁剪子体积中的非重叠块,开创性地适应了3D医学成像,为理解复杂解剖结构奠定了基础。GL-MAE [20] 通过引入局部-全局对比学习进一步推动了这一方向,以增强多粒度解剖结构建模,而SwinUNETR [9] 则通过混合变换器 [21] 纳入了多尺度特征学习。尽管这些方法展示了有希望的结果,但它们对单尺度裁剪体积的依赖本质上限制了它们捕捉全面解剖关系的能力。最近的努力探索了替代策略来解决这些限制:SwinMM [22] 利用多视图一致性来增强体积表示,而Alice [7] 则利用预训练模型 [23] 来捕捉丰富的体积内关系。尽管取得了这些进展,但三个关键挑战仍然存在:①体积裁剪所施加的空间上下文限制阻碍了对完整解剖结构的理解;②缺乏对不同解剖尺度间层次关系的显式建模;③以及处理全分辨率3D体积的巨大计算需求。这些相互关联的挑战突显了需要一种更复杂的方法,能够高效地捕捉细粒度解剖细节和全局上下文信息,同时保持实际的计算需求。

医学图像的固有层次性,尤其是那些具有广阔空间维度的图像,需要一个复杂的多级分析框架来进行全面的临床解读 [9], [24]。如图1所示,我们引入了MiM,这是一个新颖的层次化框架,从根本上推进了基于MAE的3D医学图像表示学习。我们的框架通过三个协同的组件系统地解决了现有方法的关键局限性。首先,为了克服裁剪体积的有限上下文,我们提出了一个多级体积生成策略,该策略同时在多个尺度上处理3D体积的较大视图,使我们的模型能够捕捉到细粒度的解剖细节及其更广泛的上下文关系。其次,为了明确地建模层次化表示,我们设计了一个复杂的多级重建机制,该机制在不同的解剖尺度上运行。这种机制在不同的粒度上保留了关键的解剖细节,同时通过一种先进的跨层次对齐策略强制执行一致性,确保局部结构及其全局上下文之间的一致性解读。第三,为了应对处理3D医学图像较大视图的计算挑战,我们纳入了一个高效的混合骨干设计,该设计受到MCMAE [25]的启发。

这种架构显著降低了计算开销,同时保留了基于Transformer模型的优势,使得分析高分辨率3D医学图像变得切实可行。这些创新协同作用,使MiM能够有效地对完整的解剖层次结构进行建模,同时保持计算效率。

本工作的主要贡献有三个方面:
1. 我们提出了MiM,这是一个计算效率高的自监督学习(SSL)框架,通过层次化设计推进了掩码自编码器(MAE)在3D医学图像预训练中的应用。我们的方法有效地管理了3D医学数据的复杂性,同时能够同时捕捉多个尺度上的解剖特征,这对于准确的医学图像分析至关重要。
2. 我们引入了一种全面的方法,通过两个协同的代理任务——多级重建和跨级对齐——来编码多级视觉信息。这种设计通过我们新颖的跨级对齐机制,在各个尺度上保持解剖学一致性的同时,实现了强大的局部和全局表示学习。
3. 通过在十二个多样化的数据集上进行广泛的实验验证,使用从1k到10k体积不等的预训练集,我们展示了最先进的性能,并建立了预训练数据集规模与模型有效性之间的明确相关性。我们的全面评估表明,在各种医学成像任务中都有显著改进,我们的多级方法始终优于单级替代方案。

II. RELATED WORKS

医学图像分析领域中自监督学习(SSL)的最新进展经历了不同的范式演变,其中像分割这样的密集预测任务对于3D医学图像来说至关重要,但同时也极具挑战性。在本节中,我们首先回顾医学成像中的SSL方法,重点关注从对比学习方法到生成式方法的演变及其在这些密集预测任务中的有效性。这种演变引导我们讨论掩码图像建模(Masked Image Modeling),这是一种在3D医学图像分析中展现出潜力的强大生成式SSL技术。然而,当前的掩码图像建模方法在捕捉3D医学数据的复杂层次结构方面存在困难,这促使我们审视层次化的SSL设计。通过这一综述,我们展示了现有方法在有效建模3D医学图像的多粒度特征方面存在的不足,从而为我们的层次化掩码内掩码(MIM)框架的提出奠定了基础。

医学图像分析中的自监督学习(SSL)医学成像中的SSL方法主要经历了两种范式:基于对比的方法和基于生成的方法 [26]–[28]。基于对比的方法:这些方法专注于在特征空间中对正样本对进行对齐,同时将负样本对分开 [29]。像SimCLR [29]、MoCov2 [30]和MoCov3 [31]这样的开创性工作已经在医学分类任务中取得了成功 [32], [33],并且通过特定领域的特征金字塔网络(FPN)创新进一步增强了它们在医学成像中的有效性。例如,[34]引入了基于子体积的对比学习,而[23], [35]通过在局部和全局尺度上对比特征,推动了对多粒度的理解。然而,这些方法专注于实例级对齐,限制了它们在像分割这样的密集预测任务中的有效性 [32]。基于生成的方法:这些方法通过明确建模空间结构并保持局部-全局一致性来解决上述限制 [36]。该领域从使用3D U-Net进行基本重建任务 [37], [38]发展到结合现代架构(如带有修复功能的3D Swin Transformers)的复杂方法 [9], [34]。这一演变最终促成了掩码图像建模(Masked Image Modeling)的发展,它在各种3D医学成像应用中展现出了强大的性能 [17], [19], [20], [39], [40]。

医学中SSL主要是两种方法:基于对比和生辰发的方法,对比、生成、掩码图像建模,这样的演变过程

医学图像分析中的掩码图像建模

掩码图像建模(Masked Image Modeling)作为一种强大的医学图像分析技术,建立在生成式自监督学习(SSL)的成功基础之上。最初由[17]提出,这种方法证明了高比例掩码通过原始像素恢复创建了一个有效的自监督任务。该领域通过掩码策略[41]和重建目标[18]的创新得到了发展。尽管早期的应用在2D医学任务中显示出潜力,包括疾病分类[44]和分割[43],但将其扩展到3D医学成像带来了新的挑战。MAE3D[19]率先将掩码自编码技术应用于体积数据,而GL-MAE[20]通过引入全局-局部一致性对比学习,增强了对解剖结构的理解。SwinUNETR[9]通过认识到多尺度表示学习的重要性并纳入混合变换器[21],进一步推动了该领域的发展。然而,这些方法[9], [19], [20], [22]尽管有创新,但它们在单一尺度上处理3D医学图像,具有有限的感受野,或者仅依赖于骨干级别的多尺度处理。如图1所示,当处理3D医学图像中显著更大的空间维度和变化的解剖结构时,这种方法变得有问题。

层次化自监督学习(Hierarchical SSL)层次化结构在自监督学习(SSL)中的重要性已经在自然图像和医学图像领域得到了认可。在自然图像处理中,像[18]和[39]这样的方法已经将掩码图像建模与层次化骨干网络(如Swin Transformer [45])相结合。[46]探索了多尺度特征重建,而[25]通过利用未掩码的上下文信息,推进了掩码块预测。在医学成像领域,Adam [47]引入了基于ResNet [48]的多粒度对比学习。然而,这些方法要么受到2D图像有限维度的限制[25], [47],要么独立处理层次化级别[47],未能捕捉到对3D医学图像至关重要的跨层次语义关系。

与这些方法不同,我们的MiM(掩码内掩码)框架通过以下四项关键创新推进了3D医学图像分析的最新技术:
1. 我们通过设计一种强大的多级体积生成策略,极大地超越了MAE3D [19],该策略能够同时在细粒度和粗粒度尺度上重建特征,解决了现有单尺度方法感受野有限的局限性 [9], [19], [20], [22]。
2. 受到[20], [49]的启发,我们为3D医学体积开发了一种复杂的跨层次对齐机制该机制在层次化级别之间强制执行解剖学一致性,从而在解剖特征学习方面取得了显著改进。
3. 我们从2D自然图像[25]中改编的创新3D混合骨干架构,在预训练期间以卓越的效率捕捉多尺度特征,同时减少了计算需求。
4. MiM通过成功在超过10,000个体积的超大规模数据集上进行预训练,展现了卓越的可扩展性,超越了现有3D SSL医学成像方法的规模 [8], [9], [19], [20], [32], [38]。

这些创新的有效性通过在十二个公共数据集上进行的广泛实验得到验证,涵盖了器官/肿瘤分割和疾病分类任务,确立了MiM在该领域的一项变革性进步。

III. METHODOLOGY

本节概述了我们提出的MiM方法。 首先,第III-A节介绍了MiM方法的总体框架。 其次,在第III-B节中介绍了多级重建的过程。 然后,在第III-C节中描述了我们提出的MiM方法中通过对比学习进行的跨层次对齐。 最后,在第III-D节中介绍了预训练阶段MiM方法的主干。

A. Overall framework

本文提出的MiM方法的整体框架如图2所示,包括多级重建模块和跨级对齐模块。为了使用MiM对模型进行预训练,我们首先从输入的3D医学图像中生成多级体积。然后,将输入体积裁剪成不重叠的块,这些块分为未掩码块和掩码块。未掩码块通过典型的骨干网络(CNN [48]和Transformer [50])被转换到高维特征空间掩码块用于生成下一级的掩码体积,以进行多级学习。目标是从不同级别的掩码体积中恢复掩码块。在本文中,与单一级别的掩码体积[41], [46]不同,我们提出通过多级掩码体积来简化这一目标。我们开发了一个损失函数L_R来监督最终预测。此外,我们进一步使用一个损失函数L_C来对齐跨级体积之间的共享语义,旨在学习全局语义内容以及局部细节。更多细节将在第III-B节和第III-C节中介绍。

B. Multi-level reconstruction

因此,来自不同级别的所有未掩码标记都被输入到骨干网络中以提取高维特征z。根据[17], [19]的方法,使用一个轻量级的解码器将z与可学习的掩码标记一起投影到潜在特征q。为了重建,解码器后面跟着一个简单的预测头用来重建掩码标记y。

Single-level reconstruction.我们的重建目标是每个级别体积中掩码标记的像素值。
利用从骨干网络提取的特征以及带有预测头(即线性层)的解码器,按照先前的MAE方法[19],我们将预测结果和重建目标重塑为一维向量,即 \hat{y} \in \mathbb{R}^{1 \times C}y \in \mathbb{R}^{1 \times C},其中C 是维度的数量。具体来说,我们根据经验将C设置为768,如在[19],[20]中所做的那样。

多级重建的损失函数。为了学习多粒度的细节,我们对每个级别的掩码体积图像应用单级重建。因此,多级重建损失 LR​ 可以如下公式化:

C. Cross-level alignment

跨层次体积之间共享语义块的对齐可以从细粒度到粗粒度强制执行解剖学上的相似性。图4展示了跨层次对齐的过程。由于我们从更粗粒度的体积中生成更细粒度的体积(例如,从第1级体积的掩码块生成第2级体积),这些体积必须共享语义上下文,这可以被视为正样本对。相比之下,粗粒度体积中的非重叠块(例如,第1级体积中的其余块)被视为负样本块。为了扩大共享解剖结构块(即正样本对)的高维特征一致性以及非重叠块(即负样本对)之间的差异,我们对上下文向量 x和块p应用对比学习[29]。特别地,我们将上下文向量和块重塑为一维向量c \in \mathbb{R}^{1 \times D}\{p_i \in \mathbb{R}^{1 \times D}, i \in N\},其中D 是维度的数量,根据经验将D 设置为2048,如在[49]中所做的。

D. Backbone

虽然以前的混合变换器(如Swin Transformer [21])能够生成金字塔形的特征,但在掩码图像建模(Masked Image Modeling)预训练期间,它们会处理所有的标记(解释一下:在重建的时候会处理掩码和未掩码的,但是未掩码的是不需要重建的),导致计算效率低下[25]。我们通过将MCMAE的骨干网络[25]扩展到3D医学图像,通过增加一个深度维度来解决这个限制,使我们的模型仅在变换器层处理未掩码的标记。这种优化显著提高了计算效率和可扩展性,这一点在第四节D.1中得到了验证。如图5所示,我们的FPN(特征金字塔网络)架构采用了MCMAE的[25]层次化设计,在四个尺度(H/2×W/2到H/16×W/16)上处理特征,通道维度为(C到4C)。每个尺度都使用StrideConv[25]进行下采样,然后被分割成块,接着进行掩码操作以生成未掩码块。自底向上的路径通过横向连接和求和来整合特征,同时保持计算效率,生成全面的多尺度表示。

我们的MiM(Mask in Mask)框架通过多级重建 LR​ 和跨级对齐 LC​ 引入了一种针对3D医学图像表示学习的层次化设计方法。我们根据经验将掩码体积的层级数设置为 L=3,因为三级掩码体积可以在表示学习和计算效率之间提供良好的平衡。关于 L 的消融研究在第四节D.1中进行。具体来说,方程3中的多级重建损失可以进一步展开如下:

然后,由于相邻水平体之间施加的跨水平对齐损失,Eq. 6展开为:

因此,总损失函数L是这两种损失的组合,如式10所示。

其中超参数α用于平衡这两种损失的相对贡献。 根据实验结果,我们将α设为0.1。 超参数α的消融研究见第IV-D.2节。

IV. EXPERIMENTS

本节将从介绍数据集和评估指标开始。 然后,我们将详细阐述我们的MiM的实现细节。 最后,我们将展示MiM与现有方法的实验结果,并对我们提出的方法进行分析。

A. Datasets and Evaluation

预训练数据集。为了与之前的研究[8]、[9]、[19]、[20]、[22]、[32]、[38]、[51]、[52]进行公平的比较,我们也在两个公共数据集上进行了预训练实验,即BTCV[53]和TCIA Covid19[54],并将它们合并成一个新的数据集,命名为1k。此外,为了探索我们提出的方法与之前最先进方法[9]、[19]的扩展能力,我们收集了八个公开可访问的3D医学图像数据集,包含10,502次CT扫描,以建立我们的预训练数据集,我们将其命名为10k。需要注意的是,10k仅用于探索我们提出方法的扩展能力,而我们主要关注1k数据集,以便对我们提出的方法进行公平的比较和分析。表I提供了每个收集的数据集来源的摘要。

下游数据集。我们在十二个公共数据集上进行实验,即BTCV[53]、MM-WHS[60]、Spleen[61]、Flare22[62]、Amos22[63]、MSD Task03[61]、MSD Task06[61]、MSD Task07[61]、MSD Task08[61]、MSD Task10[61]、BrasTS 21[61]和CC-CCII[64]。这些数据集包括分割和分类任务,前十个数据集用于器官分割,第十一个数据集用于病变分割,第十二个数据集用于肿瘤分割,最后一个数据集用于疾病分类。对于BTCV[53]数据集,我们严格遵循先前研究[9]、[11]、[20]中定义的数据分割,仅包括训练集和验证集。训练分割用于预训练和微调,而验证分割则从预训练中排除,仅用于评估。所有其他数据集在预训练期间都是未见过的。此外,为了评估跨模态泛化能力,我们将在CT扫描上预训练的模型转移到MRI数据集BrasTS 21[61]上。我们采用了与之前工作[9]、[19]、[20]一致的设置。

评价指标。 在[9],[20]之后,我们使用骰子相似系数(DSC)和归一化表面距离(NSD)来评估分割任务。 然后我们利用准确率(ACC)和曲线下面积(AUC)来评估疾病分类任务。

B. Implementation Details

在预训练过程中,我们遵循了之前研究[9]、[19]的设置,并在表II中提供了我们的MiM(多模态)预训练设置的详细信息。具体来说,第一级(level-1)体积是从整个CT体积中随机裁剪的。我们使用了[25]的backbone作为编码器,以高效地处理tokens。与之前的自监督学习(SSL)[20]、[49]一样,预测头和投影头都使用MLP层来对齐维度。在微调过程中,我们根据之前的研究[9]、[20],使用Swin-UNETR[21]进行分割任务的微调,使用Swin-ViT进行分类任务的微调。具体来说,对于分割任务,我们丢弃了解码器,仅在微调期间使用backbone。对于分类任务,我们严格遵循了一般计算机视觉[17]、[49]和3D医学成像[20]中先前使用的方法,仅使用最后一层的特征进行预训练。虽然结合多尺度特征可以提高分类性能[45]、[67]、[68],但我们选择不采用这种方法,以确保与其他方法进行公平比较。因此,我们仅使用最后一层的特征,结合全局平均池化(GAP)层和简单的MLP分类器来预测类别。我们用预训练过程中学习到的参数初始化网络的编码器,并对整个网络进行微调。对于这些数据集的推理,我们应用了重叠的滑动窗口推理,以便与之前的研究[9]进行公平比较。需要注意的是,为了评估我们提出方法的纯粹有效性,我们没有使用任何基础模型或后处理技术[7]、[69]。

比较方法。我们将我们的MiM方法与通用自监督学习(SSL)方法和医学自监督学习方法进行了比较。首先,我们比较了典型的SSL方法MoCov3[66]和MAE[17]、[19],因为它们代表了两种主流的SSL范式。我们还根据[19]、[20]报告了SimCLR[29]的结果。我们进一步评估了SimMIM[39]、HPM[22]、localMIM[19]和MCMAE[25]的性能,因为它们与我们先进的混合MAE相关。我们还与Adam[47]进行了比较,因为它与我们的多粒度层次设计相关。此外,我们还与大多数现有的最新医学SSL方法进行了比较。按照自然图像SSL[17]、[29]、[49]、[66]、[70]和3D医学图像[9]、[19]、[20]、[32]、[52]的常见做法,我们对所有方法进行了一轮预训练和微调,以获得结果。

C. Experiments on downstream tasks

1) Comparison on the BTCV dataset:我们首先在BTCV[53]数据集上进行了实验,结果展示在表III中。在比较的方法中,SimCLR[29]、MoCov3[66]、DINO[49]、localMIM[46]、HPM[46]、MAE3D[19]、GL-MAE[20]采用了UNETR[65]作为网络架构。大多数其他方法,包括我们的MiM方法,根据之前的研究[9],使用了Swin-UNETR[21]。表III包括了3D UNet[37]、UNETR[65]和Swin-UNETR[21]的backbone(主干网络)的详细信息。UNETR[65]使用了ViT[50],而Swin-UNETR[21]采用了Swin Transformer[45]。对于所有实验,我们为UNETR[65]使用了ViT-Base[50],为Swin-UNETR[21]使用了Swin-Base[45],以平衡性能和计算效率。这些预训练的编码器被用来初始化相应分割网络的编码器。

备注:如表III所示,我们观察到通用自监督学习方法(General SSL methods)的表现不如医学自监督学习方法(Medical SSL methods)。具体来说,SimCLR[29]和MoCov3[66]分别只达到了73.85%和79.54%。这是因为这些方法依赖于大批量大小和负样本来避免平凡常数,这对于3D医学图像来说是不切实际的。此外,SimCLR[29]和MoCov3[66]中不同图像之间使用的负关系不适合3D医学图像。DINO[49]也只取得了有限的改进。我们提出的MiM方法,以显著的优势超越了基于MAE的方法,如MAE3D[19]、GL-MAE[20]、localMIM[46]、HPM[46]和MCMAE[25]。我们得出结论,通用自监督学习方法对3D医学图像不太适合,设计自监督学习方法时必须考虑3D医学图像的特性。

从头开始训练的Swin-UNETR[21]只达到了80.53%的DSC(Dice相似系数)。通过在1k个未标记数据集上预训练MiM,我们获得了3.93%的提升,达到了84.46%的DSC,这明显优于现有的方法。在比较的方法中,Swin-UNETR[9]和MAE3D[19]分别达到了最佳的82.58%和第二好的82.40%的DSC。我们的MiM分别以1.88%和2.06%的DSC超越了这两种方法,这在这个数据集上是一个明显的改进。

值得注意的是,扩展定律[71]也适用于3D医学图像预训练。通过在更大规模的未标记10k数据集上进行预训练,我们观察到Swin-UNETR[9]和MAE3D[19]分别达到了83.20%和83.52%的DSC分数。我们的MiM在10k数据集上达到了85.41%的DSC分数,这一致显著地超越了这两种方法。这些结果表明,扩展在3D医学图像预训练中起着重要作用,我们的MiM方法对于在更大规模的数据集上进行预训练是有效的。

定性结果:MiM被发现可以提高分割结果的完整性,如图6所示。使用MiM进行分割的结果比现有方法更好。

2) Comparison on the Unseen datasets:

我们进一步在预训练中未见过的测试集上进行了实验,即MM-WHS[60]、Spleen[61]、Amos22[63]和Flare22[62]。这四个数据集的结果在表IV中展示。可以观察到,我们的MiM(可能指某种模型或方法)在所有现有方法中持续表现优异,并且优势明显,这证明了其对未见数据集有良好的泛化能力。具体来说,MiM平均比现有方法至少高出1.89%的DSC(可能指某种性能指标,如Dice相似系数)。通过使用更大规模的未标记数据集10k进行预训练,Swin-UNETR[9]和MAE3D[19]分别提高了0.92%和0.94%的DSC,达到了86.32%和86.64%。我们的MiM也得到了提升,达到了88.34%的DSC,并持续优于这两种方法。MiM在微调时也显示出标签效率[9]。具体来说,使用50%标签的MiM与从头开始训练且使用100%标签的Swin-UNETR相比,性能相当,并且有明显的优势。

3) Comparison on the MSD datasets:

为了评估在器官分割任务上的泛化能力,我们在MSD数据集[61]上进行了五个基于CT的任务实验,即任务03肝脏、任务06肺部、任务07胰腺、任务08肝血管和任务10结肠。由于现有的方法没有使用相同的预训练数据集进行实验,我们重新实现了这些方法以便进行公平比较。从表V中可以观察到,MiM在所有任务中都取得了最好的平均DSC(70.07%)和NSD(78.75%)。由于从头开始训练的Swin-UNETR[21]在平均DSC(64.84%对比62.98%)和NSD(73.12%对比69.08%)方面的表现优于UNETR[65],我们进一步基于MiM对UNETR[65]进行了预训练以进行公平比较。可以观察到,通过MiM预训练,Swin-UNETR[21]在DSC和NSD方面分别获得了平均5.23%和5.63%的提升。使用UNETR[65]作为网络时,我们在DSC和NSD方面分别观察到了平均10.09%和9.29%的提升。此外,通过使用更大规模的未标记数据集10k进行预训练,MiM在DSC和NSD方面分别进一步提高到了70.76%和79.67%。

4) Comparison on CC-CCII dataset:

为了评估我们的MiM模型在分类任务上的泛化能力,我们在CC-CCII[64]数据集上对其进行了微调,并将其性能与最先进的通用和医学领域自监督学习(SSL)方法进行了比较。由于现有的SSL方法没有在这个数据集上进行实验,我们复现了相关方法并报告了结果。如表VI所示,可以观察到MiM在准确率(ACC)和曲线下面积(AUC)方面取得了最佳性能,分别达到了93.63%和99.39%,超越了所有其他方法。这些发现表明,MiM学习到的表征可以很好地转移到分类问题上,并可有效用于医学图像分类任务。通过使用更大规模的预训练数据集10k,我们的MiM在准确率和AUC方面分别进一步提高到了94.12%和99.52%,这显示了我们提出的方法在跨任务转移时的可扩展性。

5) Comparison on the BraTS 21 dataset:

为了评估我们的MiM模型在MRI数据集上的泛化能力,我们在BraTS 21[61] MRI肿瘤分割数据集上对其进行了微调,并将其性能与最先进的通用和医学自监督学习(SSL)方法进行了比较。WT、TC和ET分别代表整个肿瘤、肿瘤核心和增强肿瘤。从表VII中可以观察到,SSL方法都能提高模型在BraTS 21[61]数据集上分割肿瘤的性能。这是因为CT和MRI通常用于相同的任务,但目的不同,因此共享相似的解剖结构。因此,从无标签的CT数据集中学习的SSL方法的知识可以转移到MRI数据集[9],[72]。我们的MiM模型在所有其他方法中表现最佳,至少提高了1.34%,达到了79.28%的DSC(Dice相似系数)。通过使用更大规模的未标记数据集10k进行预训练,我们的MiM进一步提高到了79.92%,这显示了我们提出的方法在跨模态转移时的可扩展性。

D. Analysis ofour proposed method

所有模型在1k上进行预训练,然后用BTCV[53]和MM-WHS[60]对模型进行评估。

1) Ablation study:损失函数。 我们对BTCV[53]和MM-WHS[60]验证数据集进行了全面的消融研究,以评估我们分层设计的有效性,重点是多层次重建和跨水平对准组件。

生成后续层级体积的块类型。 在块化处理过程中(图3),每个层级-l体积(xl)被划分为掩码和非掩码块。 我们在框架中评估了两种块类型,用于生成后续层级-l + 1体积(xl+1)。如表IX所示,使用来自xl的掩码块来生成xl+1,其性能始终优于使用来自xl的非掩码块。这种优越性源于掩码块迫使模型恢复缺失信息,从而促进了跨层级重叠区域的有效重建和多尺度语义表示的学习。相比之下,非掩码块在同一迭代过程中直接向模型暴露原始体积信息,有效地创建了一个信息捷径。这个捷径通过允许模型简单地复制非掩码特征,而不是学习推断和重建它们,减少了学习挑战,最终限制了模型捕捉丰富语义细节和发展强大泛化能力的能力。

关于LC的负样本对。在公式6中,我们采用infoNCE损失[29]作为默认选择。这种损失函数最大化了跨层级图像之间的相似性,并将负样本推开。另一种损失函数是BYOL余弦损失[73]。主要区别在于是否使用负样本。如表X所示,负样本有助于学习更好的表示[29],因此我们将它们包含在我们的默认选择中。

MiM的效率。本研究采用了一种混合卷积-变换器(convolution-transformer)骨干网络,该网络将卷积块与变换器层相结合。通过引入卷积块,该架构增强了归纳偏置学习,并使得多尺度特征的重用成为可能,有效地支持了混合表示学习[25]。这种设计增强了变换器处理3D医学图像的效率和效果。因此,这个骨干网络是我们所提出框架的基础组成部分。表XI展示了在预训练期间不同方法的计算成本(即浮点运算次数flops和时间)的比较分析。我们的评估强调了基于MAE(掩码自编码器)的方法与UNETR结合时的计算效率,因为它们仅利用了来自掩码3D医学图像的未掩码标记。相比之下,其他与混合骨干网络(如Swin-UNETR)简单合作的基于MAE的方法则需要处理所有标记。我们提出的方法将混合骨干架构[25]扩展到3D医学图像的预训练阶段,在保持计算效率的同时实现了显著更优的性能。

2) Hyper-parameter analysis:LR和LC的权重。在公式10中,MiM的损失函数由两部分组成。因此,在图7中,我们增加了α的值。通过观察,α的最佳值为1e-1。进一步增加跨层对齐的权重并没有带来任何额外的好处。这可能是因为跨层对齐函数的幅度远大于重建损失的幅度,导致忽略了重建过程。不同层级LR的权重。为了确定多级重建的最佳策略,我们通过改变不同层级重建损失的权重,在预训练期间评估了四种不同的学习过程,如表XII所示。使用BTCV[53]和MM-WHS[60]数据集,我们发现我们的基线模型没有层次重建(即没有粗到细和细到粗,所有权重都设置为0)的性能最低,DSC得分为81.4%和89.42%。从粗到细的过程,最初关注较粗的层级( L^1_R L^2_R),然后再优先考虑较细的层级L^3_R,将性能提高到83.99%和90.70%的DSC。值得注意的是,从细到粗的过程,通过从细层级重建开始,颠倒了这个顺序,取得了更好的结果,DSC得分为84.26%和90.85%。同时进行的过程,在训练过程中保持所有层级的权重相等,成为最有效的方法,达到了最高的DSC得分84.66%和91.04%。这些结果表明,层次化学习策略显著增强了表示学习,同时进行的多级重建证明是最有益的。

多级数L的设置。我们采用三级架构的做法得到了实证证据和实际考虑的支持。从实证角度来看,如图9所示,我们的多级策略通过精心选择的分辨率保留了解剖结构:第一级(384×384×192)捕捉广泛的上下文信息,第二级(96×96×96)提取中级特征,第三级(16×16×16)保留细节信息。通过在图8中展示的广泛实验,我们研究了改变级数L的数量的影响,并发现L=3产生了最佳性能。增加第四级将导致分辨率极低(1×1×1),虽然可能简化重建过程,但由于信息损失过大,导致性能次优,这与类似架构中的发现一致[17],[39]。

3) Reconstruction results:我们在图10中提供了BTCV[53]上的MiM重建结果,其中第一和第二行分别代表体和被遮挡的三维医学图像,最后一行的重建结果表现出优异的性能。

4) Incorporating Parameter-Efficient Fine-Tuning Methods:参数高效的微调方法,如LoRA[74]和Ladder微调[75],在计算资源有限的现实场景中提供了宝贵的洞见[74],[75]。具体来说,我们使用我们提出的方法在3D医学图像上评估了LoRA[74]在分类和分割任务上的性能,通过官方代码库1实现了它。我们的实现在注意力和卷积层中引入了低秩矩阵来近似权重更新,同时在微调期间冻结原始参数。如表XIII和表XIV所示,LoRA在显著减少更新参数数量的同时,实现了与全微调相当的性能,这与之前的发现[74]和[75]一致。我们观察到,增加秩r可以提高性能,接近全微调的性能,尽管这会以增加训练时间和内存使用为代价。这些结果表明,我们的方法可以在微调期间有效地利用LoRA[74],在计算资源和性能之间实现最佳平衡。

V. CONCLUSION

在本文中,我们介绍了掩码中的掩码(Mask in Mask,简称MiM)预训练框架,它显著推进了3D医学图像分析的发展。通过结合层次化设计,即多级重建和跨级对齐,MiM有效地将结构和细节的多粒度视觉线索编码到表示中。为了促进对现有方法的公平和全面比较,我们收集了十个公共数据集,并策划了两个规模的预训练数据集,即1k和10k。结果表明,MiM框架的层次化设计对于实现3D医学图像的优越性能至关重要。我们进一步探索将预训练数据集扩展到10k。结果表明,通过扩展预训练数据集,MiM的性能可以进一步提高。这一发现强调了大规模预训练在构建3D医学图像基础模型中的重要性。

基于这项工作,可以探索几个有前景的方向。(1) 构建大规模的预训练数据集,例如,超过10万个体积。(2) 探索更多的下游任务,例如,3D医学图像配准。(4) 探索使用可学习的层级嵌入代替硬编码层级数量的潜力。(3) 探索与其他模态合作进行多模态预训练,例如,语言。实际上,(1) 是3D医学图像基础模型的基石,(2) 可以进一步理解评估。(3) 可以通过不同模态的信息来补充3D医学基础模型。(4) 可以进一步改进方法的设计。


网站公告

今日签到

点亮在社区的每一天
去签到