Abstract
自我监督学习(SSL)提供了一个令人兴奋的机会,可以释放大量尚未开发的临床数据集的潜力,为各种下游应用程序提供标签数据的稀缺。虽然SSL已经彻底改变了自然语言处理和计算机视觉等领域,但它们在3D医学图像计算中的应用受到三个关键陷阱的限制:预训练数据集规模小,3D医学图像分析的架构不足,以及评估实践不足。我们通过 i)利用44k 3D脑MRI体积的大规模数据集来解决这些问题。ii)在最先进的nnU-Net框架内使用残差编码器U-Net架构。iii)一个强大的开发框架,包含5个开发和8个测试脑MRI分割数据集,允许性能驱动的设计决策来优化3D cnn的掩码自动编码器(MAEs)的简单概念。所得到的模型不仅超过了以前的SSL方法,而且比强大的nnU-Net基线平均高出约3个Dice点。 此外,我们的模型表现出卓越的稳定性,在7种方法中获得了最高的平均排名2,而第二好的方法的平均排名为3。 这里提供了我们的代码。
1 INTRODUCTION
近年来,自监督学习(SSL)的概念已经成为数据丰富领域的一个驱动因素,实现了大规模的预训练,从而促进了鲁棒性和可转移的通用表示的学习。这种范式转变有助于推动各个领域的发展,特别是在具有丰富标记数据的领域,如NLP或自然视觉。 在3D医学图像计算领域,这一趋势尚未流行起来。
目前,该领域主要集中在从头开始训练模型,主要使用nnU-Net框架,或者采用有监督预训练,但这种方式受到标注数据成本的限制。有监督预训练的使用表明该领域愿意采用预训练方法,这也引发了对当前领域中已建立的自监督学习(SSL)方法的质疑。我们认为,之前SSL方法未能广泛采用的原因可以归结为之前该领域SSL研究的三个主要缺陷:
P1 - Limited dataset size: 许多SSL方法都是在很少的未标记卷上开发的,通常在少于10,000个图像上进行训练,几乎接近监督数据集的规模。这些数据集往往来自公开可用的带注释的数据集,因为获取更大的数据集有一些障碍。虽然许多医院拥有数以百万计的3D医学图像,但由于患者隐私问题,这些图像被封锁在公众之外。虽然存在一些较大的开源数据集,例如美国国立卫生研究院(NIH)的青少年大脑认知衰退(ABCD)数据集(N=40k)或英国生物银行(UK- bb) (N=120k),但它们限制访问,等待内部审查委员会的批准,这对开放科学构成了障碍。最近,UK-BB停止允许将数据下载到本地硬件,这表明社区共享数据的意愿在公众中有所下降。
P2 - Outdated Backbones: 许多研究在非最先进的架构上开发SSL方法,例如,利用transformer。虽然transformer在2D自然成像领域很普遍(Dosovitskiy, 2020),但到目前为止,利用注意力的最新架构(Vaswani, 2017)尚未能够在3D医疗分割中达到最先进的性能。事实上,配置良好的3D U-Net启发cnn主导了3D医学图像分割,在很大程度上优于基于transformer的模型。这强调了SSL方法的必要性,这种方法可以与cnn无缝集成,以充分利用其在下游任务中的医学图像分析潜力。
P3 - Insufficient Evaluation: 现有的方法往往缺乏严格的评估,掩盖了方法的有效性(或缺乏有效性)。 这是通过:i)对太少的数据集进行评估,无法显示预训练的泛化。ii)叠加多个贡献,例如,新的体系结构与新的预训练,如果没有体系结构,则不允许得出预训练是否有效的结论。现有的方法常常缺乏严格的评估,从而掩盖了这些方法的有效性(或无效性)。这主要体现在以下几个方面:1. 在过少的数据集上进行评估 ,无法证明预训练方法的泛化能力。2. 将多种贡献叠加在一起,例如将新的架构与新的预训练方法结合,这使得人们无法判断预训练方法在没有新架构的情况下是否有效(例如Wang等人,2023)。3. 与配置不佳、过时的基线模型进行比较,而不是与配置良好的3D nnU-Net CNN基线模型进行比较(例如Cicek等人,2016年的模型而不是Isensee等人,2021年的模型)。4. 在预训练所使用的数据上评估他们的方法。我们想要强调的是,我们无意指责,而是想提高大家对评估重要性的认识。评估不足会导致无法明确哪些方法是最有效的。这一观察结果与Isensee等人(2024)最近的研究相似,该研究显示在从头开始训练的医学图像分割领域,这种情况也很常见。
在本文中,我们在探索**3D cnn的掩模自动编码器(MAE)**范式时,小心地避免了所有这些陷阱,Tian等人(2023)最近引入了适应性; Woo等人(2023),并强调在适当配置的情况下,MAEs在3D医学图像分割中超过了当前最先进的SSL方法。我们的贡献可以总结如下:
(1)我们通过利用44k 3D MRI体积的数据集集来开发我们的自监督预训练,从而避免了P1,在规模上超过了大多数3D医学图像分割SSL方法.
(2)我们通过利用来自Isensee等人(2024)的最先进的残差编码器U-Net CNN架构作为主干来规避P2。 此外,我们将此主干用于所有基线SSL方法-允许我们量化预训练cnn的预训练方案的有效性。
(3)我们通过使用5个开发和7个测试数据集,跨越不同的下游目标集,来规避P3。这包括头颈部器官和病理,具有在预训练期间未见的新图像模式的数据集,以及在不同中心获得的相同病理的数据集。
通过基于开发数据集上的性能表现,逐步优化适用于卷积神经网络(CNNs)的掩码自编码器(MAEs),我们提出了新颖的稀疏MAE启发式预训练范式,释放了医学领域中大量未标注数据的潜力。我们证明了选择最佳微调策略至关重要。此外,我们在一系列关键的3D医学场景中评估了我们的方法,包括低数据量情况、加速微调计划,以及在未见中心和模态上的泛化能力。
2 DEVELOPMENT FRAMEWORK
本文的目标是开发一种稳健的自监督学习(SSL)预训练方法。由于在3D医学领域相关的先前工作有限,许多设计选择需要做出。我们通过在五个下游开发数据集上依次验证每种方法论贡献,然后在八个未接触过的测试数据集上测试最终配置,来解决这一问题。为了缩小搜索范围并将SSL预训练的影响与其他基本设计选择分开,我们选择根据该领域的最佳实践保持一些参数固定。
(i)所使用的架构始终是最新且最先进的残差编码器U-Net架构(ResEnc U-Net)。
(ii)输入图像块的尺寸为160×160×160。
(iii)所有图像都被重新采样为目标间距[1×1×1] mm³(Roy等人,2023)。
(iv)所有图像都经过z分数归一化,使其均值为零,方差为单位方差。
(v)作为优化器,我们使用带有递减多项式学习率的随机梯度下降(SGD),遵循nnU-Net的设置。
(vi)在预训练期间,我们始终采用随机采样,无论不同磁共振(MR)模态的出现频率如何。
Pre-training Dataset 为了开发我们的预训练方法,我们使用了一个专有的脑部MRI数据集,该数据集来自超过44个中心,包含超过9000名患者的约4.4万张3D MRI扫描图像。由于数据来源的多样性,该数据集包含了来自10多种不同磁共振扫描仪的图像,涵盖了多种磁共振成像模态,并且包含了多样化的人群。关于数据分布的更多细节,请参见图2。

由于这些数据直接来自临床检查,它包括空白或破碎的图像,低质量的图像和所谓的侦察扫描,用于确定患者在MR中的视野。
由于这些数据直接来源于临床检查,因此其中包含了空白或损坏的图像、质量较差的图像以及用于确定患者在磁共振成像中视野范围的所谓"定位扫描"图像。由于这些扫描图像不用于诊断,因此我们通过以下方式对这些图像进行筛选:丢弃(a)在任意轴向上视野范围小于50毫米的图像,(b)在任意方向上间距大于6.5毫米的图像,以及(c)文件大小小于200KB的图像,因为这些图像表明是空白图像。此外,由于磁共振血管造影(MR Angiography)、磁敏感加权成像(SWI)和质子密度(PD)加权图像的数量较少,我们将训练数据限制为仅包括T1、T2、T1 FLAIR和T2 FLAIR图像,最终形成了包含39,168张MRI图像的预训练数据集。
Development Datasets
在预训练后,我们对5个数据集进行微调,并计算每个数据集的平均骰子相似系数(DSC)来评估预训练的有效性。 多个数据集对于确保我们的设计选择不会过度拟合特定的MRI模式或病理目标至关重要。 具体来说,我们利用:
(1)MS FLAIR (D1):在t2加权FLAIR图像上一致描绘多发性硬化症(MS)病变。
(2)脑转移(D2):在斯坦福大学医院通过T1、增强梯度回波T1ce、增强自旋回波T1和T2 FLAIR序列成像的脑转移灶
(3)海马(D3):海马体数据集,医学分割十项全能(MSD)的任务4 (Antonelli等人,2022),包含T1加权MRI中海马前侧和后侧的描绘(Simpson等人,2019)。
(4)Atlas22 (D4): T1加权图像上的脑卒中后病变解剖图(ATLAS)。 我们使用Liew等人(2022)的Atlas R2.0数据集。
(5)cross - smoda (D5):通过增强T1加权MRI对金属内和金属外前庭神经鞘瘤肿瘤和耳蜗的描绘(Dorent等,2023)。
在所有这些数据集中,我们在方法开发过程开始之前留出了一个包含所有图像20%的保留测试集。 剩下的图像被进一步分成80/20的训练集和验证集用于开发过程。
Test Datasets
此外,八个保留测试集被用来评估我们学习表征在微调它们以分割其他目标结构时的有效性。

对于所有在(D6-D11)上进行微调的测试数据集,我们使用80/20分割进行微调和测试,因为我们只对每种方法进行一次微调,而不进行任何干预。 D13使用所有数据,没有进行培训。
3 REVISITING 3D MAES
Masked autoencoders (MAEs)在自然成像领域和医学图像分割领域是一种成熟的预训练范式。 在本节中,我们研究了这种范例,并使用Isensee等人(2024)的ResEnc U-Net架构对其进行了3D医学图像分割优化。
Default parameters掩码自编码器(MAEs)通过将输入图像掩码到一定程度,并训练网络重建被遮挡的区域,从而最小化重建图像与原始图像之间的偏差。在我们的实验中,我们在归一化的 z 分数体素空间中使用 L2 损失函数训练 MAE,并且仅在被掩码的区域计算重建损失。此外,我们没有移除跳跃连接,这与 Woo 等人(2023 年)、Tian 等人(2023 年)和 He 等人(2022 年)的普遍共识一致。默认的超参数(如表 1 中灰色标注的模型所用)为:学习率 1e-2,权重衰减 3e-5,批量大小 6,使用带有 Nesterov 动量的 SGD 优化器(动量值为 0.99),掩码比例为 75%,采用 PolyLR 学习率调度训练 250k 步(这相当于在 nnU-Net 框架中的 1000 个 epoch),并进行了轻微的空间增强,包括仿射缩放、旋转和镜像。
Sparsification
在掩码输入图像时,卷积神经网络(CNNs)无法像变换器(transformers)那样忽略掩码区域。为了解决这一问题,Tian等人(2023年)提出对CNN架构进行调整,使其更适合稀疏输入:
-
稀疏卷积和归一化:通过卷积的感受野,掩码区域会从其边界开始逐渐被侵蚀。通过在每次卷积后重新应用掩码区域,可以解决这一问题。此外,掩码可能会由于引入的零值而导致归一化层统计信息出现偏差。为了解决这一问题,归一化被限制为仅考虑非掩码值。
-
掩码标记(Mask Token):与其将带有零值掩码区域的轻量级解码器特征图输入到网络中,不如用一个可学习的掩码标记来填充这些区域,从而简化解码器的重建任务。
-
密集化卷积:在用掩码标记填充掩码区域,并在将特征图传递给解码器之前,除了最高分辨率外,对每个分辨率的特征图应用一个3×3×3的卷积,以准备用于解码的表示形式。
这些改变的结果在表1a中进行了可视化。这些改进是逐步引入的,也就是说,"掩码标记(MaskToken)"的消融实验仅与稀疏卷积和归一化一起应用。可以观察到,完整的改进措施在我们的开发数据集上平均提高了0.3个DSC点的性能。随后,所有改进措施都被保留,并且后续的评估都是在应用了这些改变的情况下进行的。

Masking strategy 掩码区域是通过在CNN的瓶颈层(形状为[5x5x5])中随机采样,并把这些区域上采样到输入分辨率来确定的,以确保掩码在CNN架构的瓶颈层对齐。这导致在输入中形成了[32x32x32]体素的不重叠区域被掩码。作为采样策略,我们遵循随机掩码,因为之前的研究表明,对于图像或视频,结构化掩码并没有带来好处(He等人,2022年;Feichtenhofer等人,2022年)。在开发阶段,我们探索了5种静态掩码比例,范围在30%到90%之间,并评估了一种动态掩码比例,该比例在60%到90%之间随机掩码。
结果在表1b中展示,突出显示60%、75%的掩码比例以及60%到90%的动态掩码比例表现同样出色。由于性能高度相似,我们选择使用动态掩码比例而不是静态掩码比例进行下一步,因为我们预计这种掩码更难学习,并且考虑到即将进行的扩展实验。我们将这种模型称为Spark3D-Base(以下简称S3D-B)。

Scaling已知MAEs(掩码自编码器)从扩展中受益。我们通过将批量大小增加8倍至48,学习率增加至3e−2,以及迭代次数增加4倍至100万次来评估复合扩展的效果。我们将这个模型称为S3D-Large(S3D-L),以表示更高的计算资源,但请注意,架构保持不变,以便于参数的轻松适应。结果在表1c中展示。可以观察到,这种计算量增加32倍的结果仅导致性能略有提高,DSC点数增加了0.1。
Fine-tuning strategy给定一个预训练模型,一个关键问题出现了:哪些权重需要转移,以及如何安排微调?我们研究了各种不同的计划。关于权重转移,我们研究了(i)同时转移编码器和解码器的权重,或(ii)仅转移编码器的权重,而解码器则随机初始化。关于微调计划,我们研究了是否使用12.5k步的学习率预热,将学习率提高到最大学习率(LR)。当仅转移编码器时,我们还研究了仅对解码器进行额外预热,以使随机初始化的解码器适应预训练的编码器。在某些配置中,这导致两个学习率预热,每个12.5k步。此外,我们还研究了是否在整个微调过程中保持编码器冻结,或者也对编码器权重进行微调。最后,我们研究了是否将峰值学习率降低到1e-3、1e-4,或者保持默认的1e-2。
结果在表2中呈现,并允许我们做出三个重要的观察:(i)预热阶段是必不可少的:不应用预热步骤会显著降低性能。对编码器和解码器都进行预热可以将准确性提高0.6到1个DSC点。(ii)学习率调整很重要:在微调期间将峰值学习率降低到1e-3,比默认的1e-2始终能获得更好的结果,当同时对编码器和解码器进行微调且使用较低的学习率时,表现最佳。(iii)冻结编码器权重是有害的:在微调过程中编码器不应保持固定。允许编码器进行微调比仅微调解码器可以提高模型的性能。

4 RESULTS AND DISCUSSION
这段话描述了作者将他们最终的模型S3D-B和S3D-L与其他几种模型进行比较的过程。这些模型包括VoCo (Wu et al., 2024),VolumeFusion (VF) 由Wang et al. (2023)提出,以及MG Models Genesis (Zhou et al., 2021)。这些基线模型使用相同的框架、相同的数据、相同的主干网络和尽可能相同的超参数进行预训练,并针对A100 40GB GPU进行了优化,以充分利用其性能,目标是进行250k步训练。作者在附录A中提供了基线方法和配置的详细信息。此外,他们还与两个从头开始训练的基线模型进行了比较。第一个是'No (Dyn.)',它代表一个非预训练的(即从头开始)默认nnU-Net 3D全分辨率架构,该架构针对每个下游数据集单独进行规划和预处理,可能会导致不同的架构、数据处理和间隔。第二个从头开始训练的基线是'No (Fixed)',它使用与预训练相同的计划和预处理进行nnU-Net训练。在测试数据集上,这些方法的DSC(Dice Similarity Coefficient)和NSD(Normalized Surface Dice)平均值在表3中提供。此外,通过自助法评估了这些方法的排名稳定性,并在图3中提供。
4.1 OBSERVATIONS
SSL预训练在所有测试数据集上,SSL预训练方法显示了改进的下游分割性能。 将我们的S3D-B方法与最相似的从头开始的基线"No (Fixed)"进行比较,我们观察到11个测试数据集中有10个DSC得分更高,平均增加+2 DSC点和+1.6 NSD点。 这种改进并不局限于我们的方法; MG和VF还实现了比基线更高的性能,这表明在将SSL方法应用于足够的数据和最先进的体系结构时可以使用SSL方法。
MAEs dominate在测试数据集池中,使用掩码图像建模范式(MG,S3D-B和S3D-L)的SSL方案始终排名高于对比学习VoCo或基于伪分割的VolumeFusion预训练方法。鉴于"Models Genesis"的年代久远------发表于2019年------它能够超越更近期的VoCo或VF,这令人感到惊讶。我们将此归因于两个因素的结合:1. Models Genesis最初发布和训练在一个过时的3D-UNet(C¸ic¸ek et al., 2016)上,并且是在功能强大的nnU-Net框架(Isensee et al., 2021)之外。这突出了避免陷阱2的重要性:在最先进的主干网络上进行训练。2. VoCo和VF是与它们优化的架构一起引入的。通过将它们转移到CNN设置中,为原始架构预训练组合选择的超参数可能对新的CNN主干不是最优的。
S3D-B在DSC(Dice Similarity Coefficient,Dice相似性系数)和Normalized Surface Distance(NSD,归一化表面距离)方面排名第一,而S3D-L和MG并列第二。尽管S3D-L和MG在平均DSC上非常接近,但S3D-L在所有数据集上始终获得较低的排名。此外,根据bootstrapped aggregated rank(自助法聚合排名),S3D-L的排名明显高于MG。这些结果表明,与目前建立的CNN预训练方法相比,我们的预训练方法总体上是有效的。
Impact of dynamic configuration比较"无(动态)"(No (Dyn.))和"无(固定)"(No (Fixed))两种配置,这两种配置都是从头开始训练的,结果显示为每个数据集选择适当的配置可以显著影响性能。例如,在数据集D2和D11上,动态配置分别比固定配置高出+7和+5个DSC点,而在D6上,固定配置比动态配置高出+18个DSC点。在大多数固定配置相对于动态nnU-Net表现不佳的数据集中,预训练有助于恢复性能。然而,在某些情况下,例如D5,动态默认的nnU-Net仍然表现出色。
4.2 ABLATION EXPERIMENTS
Low-Data Regime研究者们探讨了在低数据量环境下,预训练模型相较于从零开始训练的基线模型的性能。在医学领域,许多应用只能获取到非常少量的标记图像。为了衡量在这种低数据量环境下预训练的益处,研究者们人为地将可用于训练的数据总量减少到10、20、30或40张标记图像。
表4中展示的结果表明,作者们预训练的S3D-B模型在这种设置下相比于从零开始训练的模型能带来更好的下游性能。仅使用40张训练图像,微调后的模型几乎可以匹配在完整数据集上训练的从零开始模型的性能。未来的研究可以探索是否优化训练持续时间或学习率计划可以防止预训练网络在如此有限的训练图像上进行微调时过拟合。

Generalization Performance 为了评估所提出的预训练方法的泛化能力,研究者们测试了两种场景。首先,他们在一个未见过的模态上微调了他们的方法,使用了TOF血管造影动脉瘤数据集(D12)。如表7所示,没有预训练的情况下,固定配置的性能下降了20个Dice点。这种性能下降归因于下游任务的中位间距([0.50, 0.43, 0.43] mm)与预训练实验中使用的固定目标间距[1, 1, 1] mm之间存在显著差异,后者的分辨率更高。这种较低的分辨率可能增加了分割小动脉瘤的难度。尽管如此,预训练减轻了部分性能下降,并且与从头开始训练相同配置相比,预训练被证明是非常有益的。有趣的是,ModelGenesis取得了最佳结果,这可能是由于其在预训练期间使用了强度增强,这增加了对亮度变化的鲁棒性,例如在泛化到不同的MRI序列时。其次,他们在D2数据集上仅使用T1对比增强(T1ce)序列进行微调,并将这些模型直接应用于D13,而无需额外的微调。虽然动态配置在D2的分布内验证案例中表现最佳,但D13的结果表明,MAE预训练改善了跨不同中心的泛化能力,S3D模型产生了最佳性能。
Pre-training Time先前的研究已经证明了延长训练计划对下游任务学习表征质量的积极影响(He et al., 2022; Feichtenhofer et al., 2022)。 为了探索这一因素在3D医学领域的相关性,我们进行了一个类似的实验,评估从62.5k到1M步的训练持续时间。 我们的结果表明,在250k步之后,较长的训练计划的好处开始降低,如表5所示。 这可以解释为什么我们的比例模型没有实现进一步的性能改进。

关键观察结果包括:
性能变化:随着预训练迭代次数的增加,模型的平均性能(Avg. D1-D5)并没有显著提升。例如,从62.5k迭代到1M迭代,平均性能仅从70.67提高到71.10,提升非常有限。
训练时间:预训练迭代次数的增加显著增加了训练时间。从62.5k迭代到1M迭代,训练时间从28小时增加到448小时。
性能稳定:在250k迭代时,模型已经达到了相对稳定的性能水平,进一步增加迭代次数并没有带来明显的性能提升。
此外,表格下方的注释提到了一些技术细节:
在位置预测任务中,输入图像块的大小为384x384x96,这些图像块被重新缩放和调整大小以适应4x4的64x64x64图像块网格。
由于选择的图像块大小160x160x160与64的立方体长度不兼容,因此调整了VoCo的图像块大小为192x192x64,以适应3x3的64x64x64图像块网格。
4x4网格导致内存限制超出,因此需要减少图像块的数量。此外,为了充分利用A100节点的40GB VRAM,将目标图像块大小从4增加到5,并将批量大小从6(其他实验中的默认值)增加到12。
Fine-tuning length从预训练权重开始初始化有可能减少网络适应新任务所需的计算资源。为了评估这一点,我们在开发数据集上测试了不同的预训练持续时间。在保持解码器和整个网络预热阶段各12.5k次迭代的同时,我们尝试了不同的后续训练长度。如表8所示,仅增加12.5k次额外迭代(总计37.5k次)就已经优于从头开始训练。然而,要达到最佳性能仍然需要完成整个微调计划。

Multi-Channel Input 在许多医学检查中,通常会进行多次扫描,因为临床医生通常需要具有不同特征的图像来进行准确的决策制定。因此,一些数据集,如D2、D8和D10,包含多种输入模态。虽然预训练可能涉及所有模态,但由于并非所有患者都有每种模态的扫描,我们一次只向网络提供一种模态。这就引出了如何处理具有多个注册图像的数据集的问题。为了解决这个问题,我们在D2开发数据集上进行了5折交叉验证。我们评估了每种输入模态的复制以及输入stem权重的随机初始化。此外,我们还测试了在解码器预热阶段冻结stem权重。如表6所示,最稳定且始终有效的方法是在解码器的预热期间复制预训练的stem并保持其冻结。

5 CONCLUSION
这项工作首次展示了在三维医学图像分割中正确配置的掩码自编码器(MAEs)的潜力。通过克服先前研究中的关键问题,例如有限的数据集规模、过时的架构和不充分的评估,我们展示了与以前的半监督学习(SSL)方法相比一致的性能提升。值得注意的是,我们首次实现了在动态的、数据集自适应的nnU-Net基线之上的一致增益,并在大量多样的开发和测试数据集上得到了验证。虽然我们的发现是有希望的,但仍有几个方向留待未来探索。特别是,增加训练时间和批量大小并没有带来性能提升,但问题仍然存在,即扩大预训练数据集的规模或模型参数是否能够释放新的潜力。此外,ModelGenesis SSL任务所采用的强度变化暗示了提高跨未见过的MRI模态泛化能力的有趣可能性,这需要对MAEs进行进一步探索。最后,以数据为中心的方法来策划SSL最相关数据代表了未来研究的一个令人兴奋的前沿。原始临床数据集通常包含不打算用于诊断目的的图像,例如用于扫描仪校准的图像,这可能会削弱预训练的效果。虽然我们应用了基本的过滤来排除低质量数据,但更复杂的过滤技术可以显著提高预训练过程的质量。
这项工作继承了先前研究的精神,如nnU-Net(Isensee等人,2021年)或其他研究,展示了稳健的开发策略、明智的模型配置和严格的验证可以带来真实且可持续的性能提升,这与当前追求使用和修改最新网络架构的热潮形成了对比。通过我们的动态开源框架,我们希望促进社区向验证驱动的开发转变,从而推动真正的科学进步。