Large-Scale 3D Medical Image Pre-training with Geometric Context Priors

大规模三维医学图像预训练:利用几何上下文先验

摘要 ------标注数据的稀缺性给医学图像分析带来了重大挑战,尤其对于高维的三维医学图像而言,获取标注需要放射科医生投入大量精力。大规模预训练作为一种有望实现标签高效利用的解决方案,凭借对大规模数据、大模型和先进预训练技术的整合,展现出巨大潜力。然而,其在医学图像领域的应用仍处于初步探索阶段。主要挑战在于如何有效利用大规模无标注数据,在缺乏人工标注的情况下学习高层次语义信息。我们观察到,三维医学图像具有稳定的几何上下文特性,即不同器官之间存在一致的几何关系,这为学习一致性的表征提供了可行路径。受此启发,我们提出了一种简洁而高效的体素对比(Volume Contrast, VoCo)框架,通过利用几何上下文先验实现自监督学习。给定一个输入体数据,我们从不同区域提取"基准裁剪块"(base crops),构建用于对比学习的正负样本对;随后,通过将一个随机裁剪块与这些基准裁剪块进行相似性对比,预测其在整体中的上下文位置。通过这种方式,VoCo 将固有的几何上下文隐式编码到模型表征中,从而在无需标注的情况下促进高层次语义的学习。

为全面评估所提方法的有效性,我们开展了以下工作:(1)构建了迄今规模最大的医学图像预训练数据集 PreCT-160K,包含 16 万例覆盖多种解剖结构的计算机断层扫描(CT)三维体数据;(2)研究了模型缩放规律,并提出了针对不同医学任务适配不同模型规模的实用指南;(3)建立了一个涵盖 48 项医学任务的综合基准,包括分割、分类、配准以及视觉-语言任务。大量实验结果充分验证了 VoCo 的优越性,其在未见过的模态和数据集上展现出卓越的迁移能力。尤其值得注意的是,VoCo 在标注样本稀缺的数据集上显著提升了性能,并大幅加快了微调过程的收敛速度。

1 引言

近年来,人工智能驱动的医学图像分析取得了显著进展 [2]--[7],但其发展严重受限于专家标注的高昂成本,尤其对于包含丰富体素信息的大规模三维(3D)医学图像而言更是如此 [8]--[11]。为应对这一困境,自监督学习(Self-Supervised Learning, SSL)[12]--[16] 在基础模型预训练中的应用展现出巨大潜力:它能够在无需人工标注的情况下学习有效的特征表示,为缓解3D医学图像分析中的标注瓶颈提供了有前景的解决方案 [8], [9], [17]--[19]。

近期研究 [12], [20]--[24] 表明,视觉基础模型的成功依赖于三大关键要素:大规模数据、大模型以及先进的预训练技术。然而,这些要素在3D医学图像预训练中的迁移效果尚未得到充分探索。如图1所示,当前研究在以下三个方面仍存在明显局限:

(1)数据:现有方法 [8], [9], [17], [18], [25]--[28] 受限于数据规模(最多仅使用1万例体数据)。例如,UniMiss [9], [27] 创新性地通过融合二维胸部X光片来增强胸部CT的预训练效果,但其在其他解剖区域的可扩展性仍有待探索。

(2)模型:以往工作 [8], [9], [17], [18], [25]--[28] 所采用的模型规模仍然较小,参数量仅在千万级别。医学图像预训练中模型容量的缩放规律尚未被系统研究。

(3)预训练技术:SuPreM [26] 聚焦于有监督预训练,并为此标注了一个腹部分割数据集 [29]。尽管其性能优于此前方法,但仍受限于标注数据的规模,未能有效整合来自多样化解剖区域的大规模无标注数据。

在自监督学习领域,大多数现有方法 [8], [9], [10], [17], [19], [25], [30], [31] 主要依赖低层次信息重建来学习对数据增强不变的表征,通常通过对图像施加增强操作后再重建原始信息。然而,这种策略在预训练阶段缺乏高层次语义信息,限制了其在各类下游任务中的表现。

当前的核心挑战在于:如何在大规模无标注数据上引入高层次语义进行预训练。我们指出,3D医学图像中蕴含的几何上下文先验(geometric context priors)可被有效利用。如图2所示,我们观察到,在3D医学图像中,不同器官(即语义区域)之间呈现出相对稳定的几何关系,并具有相似的解剖特征。这种器官间一致的几何上下文为我们在无标注条件下学习一致的语义表征提供了可行路径。

为此,本文提出了一种简洁而高效的体素对比(Volume Contrast, VoCo)框架 ,旨在利用几何上下文先验进行对比学习。VoCo引入了一项新颖的预训练任务------上下文位置预测(contextual position prediction),旨在将不同器官间的几何关系编码到模型表征中。具体而言,VoCo首先从输入体数据的不同区域提取一组互不重叠的"基准裁剪块"(base crops)。随后,将一个随机裁剪块(random crop)与这些基准裁剪块配对,构建用于对比学习的正负样本对:若随机裁剪块与某个基准裁剪块存在空间重叠,则视为正样本;否则为负样本。接着,通过计算随机裁剪块与各基准裁剪块之间的相似度,预测其在整体中的上下文位置。直观上,相似度越高意味着重叠区域越大,因此可通过相似度判断该随机裁剪块所属的解剖区域。具体实现中,我们将随机裁剪块与各基准裁剪块之间的重叠比例作为位置标签,用于监督位置预测任务。通过学习预测上下文位置,VoCo在无需人工标注的情况下,将固有的几何上下文隐式编码进模型表征中。

如图1所示,现有工作 [8], [9], [17], [18], [19], [26] 仍受限于数据规模,距离构建强大的医学视觉基础模型仍有显著差距。为此,我们从公开来源整理构建了大规模数据集 PreCT-160K ,这是目前医学图像预训练领域规模最大、覆盖最全面的数据集。如图1(a)所示,PreCT-160K 包含超过16万例CT体数据(总计逾4200万张切片),覆盖人体完整的3D解剖结构。该数据集还包含大量标注数据,使我们能够结合自监督与半监督学习,实现全监督预训练(omni-supervised pre-training)。本文进一步提出了一个全监督预训练框架,以充分释放标注与无标注医学图像的联合潜力。

我们还深入探索了模型容量的缩放规律,并提出了针对不同医学任务适配不同模型规模的实用指南。具体而言,我们构建了一个大规模医学图像预训练评估基准。与以往研究 [8], [9], [17], [18], [19], [26] 在评估数据和任务上的局限性不同,我们的基准涵盖48个下游数据集,覆盖分割、分类、配准及视觉-语言等多种任务。在48个下游任务上的大量实验结果表明,所提出的VoCo方法显著优于现有方法,并在多项任务上达到新的当前最优(state-of-the-art)性能。

本研究的初步版本曾发表于CVPR 2024 [1],命名为VoCo-v1。本文在此基础上进行了重大且实质性的改进,但仍沿用VoCo这一名称。本文的新贡献包括但不限于以下几点:

  • 相较于仅关注体内对比学习 (intra-volume contrastive learning)的VoCo-v1 [1],我们进一步引入了体间对比学习(inter-volume contrastive learning),通过基于动量的教师-学生模块,实现不同体数据之间一致表征的学习。
  • 我们探索了自监督与半监督学习的结合,用于全监督预训练,有效利用标注与无标注数据。
  • 我们发布了当前规模最大的医学图像预训练数据集 PreCT-160K,将数据规模从1万例 [1] 大幅扩展至16万例。该数据集有望推动医学图像预训练领域的未来研究。
  • 我们构建了当前最大的医学图像预训练评估基准,涵盖48个下游数据集的多样化任务。我们开源的各类医学任务实现也将惠及该领域的后续研究者。
  • 我们深入研究了模型缩放规律,并发布了参数量从3100万到12亿(1.2B)的系列预训练模型,同时提出了针对不同医学任务选择合适模型规模的实用指南。
  • 我们提供了详尽且富有洞见的分析,深入剖析了VoCo的核心组件。这些实验进一步凸显了大规模预训练的重要性,为医学图像预训练领域的未来研究提供了宝贵启示。

2 相关工作

2.1 大规模视觉预训练

视觉预训练为利用大规模视觉数据开辟了广阔前景,在大型视觉基础模型的发展中扮演着关键角色 [12], [20], [22], [23], [24], [32]。其主要挑战在于设计有效的预训练方法。尽管有监督预训练是一种直观的途径,但它面临人工标注稀缺的问题,且构建大规模标注数据集需要巨大的工程投入。Deng 等人 [33] 构建了著名的 ImageNet 数据集,而基于 ImageNet 的预训练已被证明能显著提升下游任务性能。SAM [22] 引入了包含超过10亿个分割掩码的 SA-1B 数据集用于有监督预训练,从而构建了一个强大的分割基础模型。然而,高昂的标注成本以及对大规模无标注数据的忽视,仍然制约了有监督预训练的进一步发展。

为此,自监督学习(Self-Supervised Learning, SSL)被提出,旨在无需人工标注的情况下学习鲁棒特征 [12], [13], [14], [20], [21],近年来受到广泛关注。

典型的 SSL 方法 。SSL 已在各类视觉任务中展现出良好效果 [12], [13], [14], [20], [21], [34]。DINO [12], [20] 提出融合先进的 SSL 方法,在无标注条件下学习鲁棒特征,已成为当前研究中预训练骨干网络的主流选择。当前最先进的 SSL 方法大致可分为两类:生成式方法 [21], [35]--[37] 和 对比学习方法 [13], [14], [34], [38]--[43]。

(1)生成式方法主要基于从增强后的图像中重建原始信息。例如,MAE [21] 提出对输入图像的随机图像块进行掩码,并重建缺失的像素。

(2)对比学习方法则通过构建样本的正负对,学习具有一致性的表征。

向医学图像分析的迁移 。尽管上述方法在自然图像领域取得了显著成果,但将其预训练模型直接迁移到医学图像时仍面临挑战,主要源于领域差异(domain gap)[8], [9], [17], [19], [27], [28]。DINO [12], [20] 预训练了一系列强大的二维 Vision Transformer [44],在 X 光片和病理图像等二维医学图像上展现出良好的迁移能力 [45]--[47]。然而,在需要提取体素信息的复杂三维医学任务中,强大的预训练三维模型仍处于初步探索阶段 [8], [9], [26], [48]。

大多数当前最先进的 SSL 方法 [14], [15], [21], [34], [38] 在三维医学图像上的表现往往不尽如人意,主要原因在于忽视了三维医学图像的独特特性 [8], [9], [18], [28]。具体而言:

(1)自然图像中的对比学习方法 [15], [38] 通常在训练批次内构建正负样本对:将同一图像的不同增强视图视为正样本,其他图像视为负样本。然而,对于具有相似解剖结构的三维医学图像而言,这种构建负样本的方式难以奏效 [10], [11], [17], [28],因为不同样本之间可能存在高度语义相似性,导致负样本"不够负"。

(2)掩码图像建模方法(Masked Image Modeling)[13], [21] 通过掩码并重建缺失像素进行学习。然而,三维医学图像通常具有高维度、大尺寸以及大量背景区域的特点,这类方法容易使模型倾向于重建无关的背景区域 [16], [17], [31], [49]--[51],从而削弱对语义区域(如器官)的理解能力。

因此,面向三维医学图像的先进自监督学习技术的发展,必须细致考虑其独特的图像特性 ,并制定量身定制的策略

2.2 大规模医学图像预训练

医学图像预训练已被证明是缓解医学任务中标注稀缺问题的有效手段 [8], [9], [18], [26], [52], [53]。早期尝试 [25], [54], [55] 主要在二维X光图像上进行预训练 [56], [57],在胸部病理识别和气胸分割等任务上取得了性能提升。相比之下,三维医学图像(如CT和磁共振成像MRI)为临床诊断提供了更丰富的体素信息,近年来在医学图像分析中受到越来越多关注 [10], [11], [16], [26], [48], [58]。然而,三维医学图像固有的复杂性也给预训练带来了显著挑战。尽管近期工作 [8], [9], [17], [18], [19], [26], [28] 已验证了三维医学图像预训练的有效性,但在数据规模、模型容量和预训练方法等方面仍面临重大挑战。

2.2.1 大规模数据

与二维X光图像相比,获取CT等三维医学图像更为困难,原因包括成像速度较慢、辐射暴露更高以及成本更高等因素 [59], [60]。如图1(b)所示,大多数现有方法 [8], [17], [18], [19], [25], [26] 仅使用有限规模的三维数据进行预训练。FreeTumor [61] 首次在包含1.1万例CT的肿瘤分割任务中研究了数据缩放规律。Wang 等人 [51] 构建了一个包含10万例CT的预训练数据集,但该数据集并未公开用于研究。因此,要为预训练收集大规模三维数据,有必要整合来自不同地区和国家多家医院的多源数据 [29], [62]。然而,这种整合会带来图像特性多样、成像质量不一致等问题,为预训练引入新的挑战。

此外,以往方法主要针对特定身体部位收集数据进行预训练,例如:PCRL [8], [25] 和 UniMiss [9], [27] 聚焦于胸部区域,Alice [18] 和 SuPreM [26] 针对腹部区域,GVSL [28] 则专注于心脏区域。然而,鉴于不同解剖区域具有显著差异的特性,仅在一个区域上预训练的模型在迁移到其他区域时可能表现受限 [1], [16], [49]。

本文构建了大规模数据集 PreCT-160K,涵盖多样化解剖结构。然而,来自不同解剖区域的数据在成像参数(如图像尺寸、体素间距和强度范围)上存在显著差异,这对预训练中学习一致性表征提出了新的挑战。

2.2.2 大模型

早期的三维医学图像预训练工作 [25], [54], [55] 受限于模型容量,通常仅包含数千万参数。近期研究 [20], [21], [23], [44], [63] 已充分验证了缩放定律(scaling law)的惊人效果:在大规模数据上训练的大模型展现出卓越的智能能力。

本文收集了大规模三维医学图像数据集,其中包含来自多源的多样化图像特征。如此丰富的数据为训练大模型提供了前所未有的机会。

考虑到医学任务的多样性,有必要在全面的基准 上评估大模型的性能。以往方法 [8], [17], [18], [26], [27], [54], [55] 通常仅在少数下游任务(如分割或分类)上评估预训练模型。STU-Net [64] 是首个尝试评估大模型的工作,但仍局限于分割任务。本文则深入探索了不同医学任务中的缩放规律,为针对不同任务定制合适模型规模提供了实用指导。

2.2.3 先进的预训练技术

三维医学图像的自监督学习(SSL)。现有方法 [8], [10], [19], [30], [31], [65] 多基于信息重建来学习对数据增强不变的表征:首先对图像施加强增强,然后重建原始信息。"旋转-重建"(Rotate-and-reconstruct)[10], [17], [30], [66] 方法通过随机旋转三维体数据并学习恢复原始方向,促进旋转不变特征的学习。近期方法 [8], [18], [25], [27], [28], [67] 进一步探索了从不同视角恢复低层信息。例如,PCRL [8], [25] 同时裁剪全局和局部图像块,并进行多尺度重建;GVSL [28] 则通过仿射增强和匹配,挖掘多扫描图像间的几何相似性。掩码重建方法 [16], [19], [31], [49], [50] 源自 MAE [21],通过掩码图像并重建缺失像素来学习表征。

尽管这些方法取得了一定成果,但大多数仍忽视了将高层次语义信息融入模型表征,从而限制了下游任务性能的进一步提升。

预训练中的高层次语义。对于医学图像而言,高层次语义信息主要来源于人工标注,因其高度依赖专家知识。先前工作 [26], [64], [68], [69] 指出,有监督预训练效率更高,在更少训练时间和标注数据下即可达到更优性能 [26]。然而,标注数据的稀缺性仍是持续存在的挑战,限制了模型在不同医学任务、解剖结构以及大量未见数据集上的迁移能力。

本文旨在将大规模无标注数据纳入预训练过程,因此提出利用医学图像的固有特性作为高层次语义先验,用于自监督学习。

全监督学习(Omni-supervised Learning)。尽管自监督学习使我们能够利用大规模无标注数据进行预训练 [1], [9], [16],但它往往忽略了已有的标注数据。全监督学习 [70]--[73] 提出了综合利用多种监督信息的理念。具体而言,半监督学习 [74]--[77] 在联合利用标注与无标注数据方面展现出强大效能。

本文提出了一种简洁而高效的全监督预训练框架,将自监督与半监督学习相结合,充分释放标注与无标注医学图像的联合潜力。

3 方法

3.1 用于自监督的位置标签生成

关键步骤是生成用于自监督的位置标签。我们提出利用3D医学图像中固有的几何上下文先验。如图3所示,给定一个输入体数据 VV,我们首先随机裁剪出一个子体块 kk,目标是围绕 kk 构建用于对比学习的正负样本对。

具体而言,我们采用位置编码策略生成 nn 个互不重叠的基准裁剪块(base crops)qiqi​(i∈{1,...,n}i∈{1,...,n})。例如,图3中生成了 4×4=164×4=16 个基准裁剪块,每个 qiqi​ 代表输入体数据中的一个特定区域。

在人体解剖结构中,不同器官位于不同区域,这为我们构建正负样本对提供了天然依据。如图3所示,随机裁剪块 kk 与正样本基准块 qposqpos​ 存在空间重叠,而负样本基准块 qnegqneg​ 则无重叠(或重叠极少),更可能包含不同的器官(并非绝对)。例如,图3中 kk 与 qposqpos​ 均包含胃、胰腺、静脉、主动脉和下腔静脉,而 kk 与 qnegqneg​ 则呈现不同的器官信息。因此,我们可以利用位置编码来构建用于对比学习的正负对。

以往的对比学习方法 [14], [15], [38], [43] 主要采用 InfoNCE 损失 [78] 来最大化正样本对之间的互信息。本文则提出生成具有具体数值的位置标签,以监督正样本对之间的相关程度------即用标签反映 kk 与 qposqpos​ 的相似度。可以观察到,kk 与 qposqpos​ 的相关性与其重叠比例密切相关:重叠区域越大,二者越相似。

因此,如图3所示,我们将 kk 与每个基准裁剪块 qiqi​ 之间的重叠比例作为位置标签 yiyi​ 的值,用于衡量其相似度;而对于负样本 qnegqneg​,其位置标签 yiyi​ 被设为 0。通过这种方式,我们利用 kk 与 qq 之间的重叠比例来监督上下文位置预测结果。

3.2 用于上下文位置预测的体素对比(Volume Contrast)

VoCo 的整体框架如图4所示。我们提出了一项新颖的预训练任务------上下文位置预测 ,通过体素对比来预测随机裁剪块 kk 的上下文位置。该任务包含两个部分:

(1)体内对比 (intra-volume contrast):在同一个体数据内,对 kk、qposqpos​ 和 qnegqneg​ 进行对比;

(2)体间对比 (inter-volume contrast):在不同体数据 VAVA​ 与 VBVB​ 之间建立对比,通过典型的教师-学生模块(student-teacher module)[14], [15], [38], [43] 实现一致性正则化。

上下文位置预测。如图4(a)所示,给定一个输入体数据,我们首先提取一个随机裁剪块 kk 和一组基准裁剪块 qq,其对应的位置标签 yiyi​ 按照第3.1节和图3的方式生成。随后,我们将 kk 和 qq 输入模型以提取高维特征。

特征提取后,我们采用典型的基于动量的教师-学生模块 [15], [38] 分别对 kk 和 qq 进行投影。具体而言,教师投影器 ptpt​ 在训练过程中参数冻结,其参数 θtθt​ 通过学生投影器 psps​ 的参数 θsθs​ 进行指数移动平均(EMA)更新:

θt←ρθt+(1−ρ)θs,θt​←ρθt​+(1−ρ)θs​,

其中动量因子 ρρ 经验性地设为 0.9。该机制在对比学习中被证明有效 [15], [38],可实现稳定训练并避免特征坍缩(feature collapse)[14], [34], [43]。

从投影器输出的特征经3D自适应平均池化(3D adaptive average pooling)压缩为一维向量,即 k∈R1×Ck∈R1×C、q∈R1×Cq∈R1×C,其中 CC 为特征维度。随后,我们计算随机裁剪块 kk 与每个基准裁剪块 qiqi​ 之间的相似度 sisi​,采用余弦相似度:

si=CosSim(k,qi)∈[0,1].si​=CosSim(k,qi​)∈[0,1].

直观上,更高的 sisi​ 表明 kk 与 qiqi​ 共享更大重叠区域的概率更高。因此,我们可通过计算相似度 ss 来预测上下文位置,并利用生成的位置标签 yy 监督该预测。

预测损失函数 LpredLpred​ 基于熵设计。首先计算相似度 ss 与位置标签 yy 之间的距离 dd:

d=∣s−y∣,d=∣s−y∣,

其中 ∣⋅∣∣⋅∣ 表示绝对值(注意 s,y∈[0,1]s,y∈[0,1])。随后,损失函数定义为:

Lpred=−log⁡(1−d).Lpred​=−log(1−d).

备注:尽管我们将所有负样本 qnegqneg​ 的位置标签 yy 设为 0,但在实际中,随机裁剪块 kk 有时可能与某些 qnegqneg​ 相似。由于预训练阶段无真实标签,构建绝对理想的负样本对在对比学习中仍具挑战性 [15], [38], [79]。然而,负样本对之间的整体距离通常较大。因此,我们遵循以往方法 [15], [38], [78],采用式(4)中距离的平均熵形式。

体内对比。如图4(b)所示,我们在三元组(kk、qposqpos​、qnegqneg​)上进行体内对比:拉近 kk 与 qposqpos​,同时推远 kk 与 qnegqneg​、以及 qposqpos​ 与 qnegqneg​。对于 kk,我们使用位置标签 yy 监督对比过程;对于基准裁剪块 qq,我们设计了一个正则化损失 LregLreg​,以增强任意两个 qiqi​ 与 qjqj​ 之间的特征差异性:

Lreg=1n(n−1)∑i≠jsij,Lreg​=n(n−1)1​i=j∑​sij​,

其中 sij=CosSim(qi,qj)sij​=CosSim(qi​,qj​)。

体间对比 。如图4(c)所示,我们从体数据 VAVA​ 中提取随机裁剪块 kAkA​,从另一同批次的体数据 VBVB​ 中提取基准裁剪块 qBqB​,构建体间对比。值得注意的是,VAVA​ 与 VBVB​ 来自相同解剖区域(如同为腹部或胸部)。

我们采用一种简洁高效的一致性正则化方法 [15], [38], [76]。首先对 kAkA​ 和 qBqB​ 应用特征增强(图4中标注为"aug."),此处采用简单的 Dropout [80](同 [76]),得到增强后的 kA′kA′​ 和 qB′qB′​。随后,将增强前后的特征分别输入学生投影器 psps​ 和教师投影器 ptpt​。投影后计算余弦相似度:

ys=CosSim(kA,qB),yt=CosSim(kA′,qB′).ys​=CosSim(kA​,qB​),yt​=CosSim(kA′​,qB′​).

体间对比损失定义为:

Linter=−1n∑i=1nlog⁡(1−∣ys−yt∣).Linter​=−n1​i=1∑n​log(1−∣ys​−yt​∣).

SSL 总体损失函数。综上,自监督学习的总损失为:

LSSL=Lpred+Lreg+Linter,LSSL​=Lpred​+Lreg​+Linter​,

其中三项损失权重经验性地设为相等 [1],因我们认为它们同等重要。

3.3 面向全监督预训练(Omni-supervised Pre-training)

算法1:全监督预训练
输入 :标注分割数据 (XL,YL)(XL​,YL​),无标注数据 XUXU​
输出:预训练模型 MM

  1. 第一阶段

  2. 有监督训练:M←(XL,YL)M←(XL​,YL​);

  3. 自监督训练:M←XUM←XU​(使用式9的 LSSLLSSL​);

  4. 第二阶段

  5. 生成伪标签:YU←(M,XU)YU​←(M,XU​);

  6. 半监督训练:M←(XL,YL,XU,YU)M←(XL​,YL​,XU​,YU​);

  7. 自监督训练:M←XUM←XU​(再次使用 LSSLLSSL​)。

如图5所示,有监督学习与自监督学习各有优劣:

(a) 有监督学习可在标签指导下学习判别性强的决策边界,但受限于标注数据稀缺;

(b) 自监督学习可利用大规模无标注数据,但缺乏标注监督,难以学习清晰的类别边界。

为此,我们提出全监督预训练(见算法1),有效融合标注与无标注数据的优势。该框架结合了有监督与自监督学习的优点,充分释放两类数据的潜力。

构建标注分割数据集 (XL,YL)(XL​,YL​) 。PreCT-160K 包含大量标注分割数据,但许多数据集标签不完整 [2], [6], [29](例如,一个数据集仅有肝脏标签,另一个仅有胰腺标签)。因此,我们首先集成多个模型为 XLXL​ 生成完整标签 YLYL​,并从中整理出一个小型标注子集,命名为 VoComni。该子集包含2万例体数据,覆盖20种器官和肿瘤类别,将随 PreCT-160K 一同发布,以促进后续研究。所有验证集和测试集在 PreCT-160K 与 VoComni 中均不可见。

半监督学习是一种可扩展的学习器 。为有效利用标注与无标注数据,我们采用半监督学习 [74]--[76],将标注数据中的知识迁移到大规模无标注数据上。值得注意的是,分割是监督训练中的关键技术 [26], [68], [69],因为许多医学任务需要像素级的精细理解以实现准确诊断。以往工作 [7], [68], [107] 仅使用数百例进行半监督分割,而复杂的半监督分割设计在我们的大规模数据上不可行。

本文采用最简半监督学习基线,并将数据规模扩展至16万例。我们发现,结合 VoCo 后,即使是最简单的半监督基线也能取得具有竞争力的结果。如算法1所示:

  • 第一阶段:从 PreCT-160K 中整理标注数据 (XL,YL)(XL,YL),进行有监督分割训练;
  • 第二阶段:为无标注数据 XUXU 生成伪标签 YUYU,在 (XL,YL)(XL,YL) 和 (XU,YU)(XU,YU) 上进行半监督训练;
  • 两个阶段均协同整合自监督学习(SSL)。

通过这种方式,我们融合了自监督与半监督学习的优势,迈向全监督预训练的新范式。

4 实验

4.1 数据集与实现细节

预训练数据集 。本文构建了目前规模最大的医学图像预训练数据集 PreCT-160K (见表1)。该数据集来源于多个机构,经过严格的预处理以确保训练时数据格式的一致性。具体而言,为应对不同解剖区域体数据在尺寸、体素间距和强度上的差异,我们设计了针对性的预处理流程。由于 PreCT-160K 中胸部区域数据占比较高,我们在预训练阶段通过采样平衡策略缓解类别偏差。

对于 VoComni 数据集 ,我们采用模型集成方法生成伪标签,并剔除预测置信度较低的样本。最终构建了一个包含 2万例伪标注体数据 的分割数据集,涵盖20种器官与肿瘤类别,用于本文提出的全监督预训练。

评估基准 。我们构建了一个大规模评估基准(见表2),包含 48个下游任务数据集 ,涵盖分割、分类、配准、视觉-语言等多种任务。如表3所示,该基准比以往工作 [8], [9], [17], [18], [19], [26], [28] 更加全面。部分数据集 [2], [84], [90], [95] 的结果在公开排行榜上进行评估;若测试集或排行榜不可用,则在相同的离线验证集划分上报告结果,以确保公平比较。

实验设置 。我们首先在 PreCT-160K 上进行预训练,随后在48个下游数据集(表2)上微调预训练模型进行评估。预训练骨干网络采用 SwinUNETR [129]nnUNet [48]。具体地,我们在 SwinUNETR 中分别使用 Swin-Base(B)、Swin-Large(L)和 Swin-Huge(H),其特征维度分别为48、96和192。

本项目依托 NVIDIA SuperPOD 硬件平台:预训练使用 8块 NVIDIA H800 GPU ,所有下游任务均可在单块 H800 或 RTX 3090 GPU 上完成。下游评估累计消耗超过 10,000 GPU 小时 。我们的代码已开源,全面支持 MONAI [130]nnUNet [48] 框架。


4.2.2 医学图像分类

CC-CCII [121]LUNA16 [83] 上的分类结果见表7。由于 LUNA16 上肺结节检测任务已接近性能上限,预训练带来的增益不明显。但在 CC-CCII 的新冠分类任务 中,VoCo 相比从零训练基线提升 2.76% ,相比 SuPreM [26] 提升 1.97%。值得注意的是,SuPreM [26] 仅在腹部数据集上进行有监督分割预训练,可能限制了其在胸部分类任务上的迁移能力。


4.2.3 医学图像配准

IXI [126]OASIS [127] 上的配准结果见表8。我们以 TransMorph [134] 为基线模型。本文聚焦于评估预训练的有效性,未提出新的配准算法 ,因此配准分析重点在于骨干网络对比(从零训练 vs 预训练)。我们发现,以往预训练方法 [17], [26] 在配准任务上表现不佳。而在脑部 MRI 配准数据集 OASIS 上,VoCo 带来了 2.6% 的 DSC 提升------这在配准任务中已是显著进步。


4.2.4 视觉-语言分析

如表3所示,本文首次系统评估 医学图像预训练在视觉-语言(Vision-Language, VL)任务中的有效性。具体包括:

  • CTRG-Chest [142] 上的医学报告生成任务;
  • CT-RATE [105] 上的词汇分类报告-体数据检索任务。

结果见表9和表10。本文聚焦医学图像预训练,因此通过替换视觉编码器验证其效果,语言模型部分则沿用 M2KT [142] 和 CT-CLIP [105] 的原始设置。

VoCo 在各项指标上均优于现有医学预训练方法 [8], [17], [26], [55]:

  • 报告生成(表9):VoCo(Swin-H)取得最高 BLEU-4 分数 37.91%
  • 词汇分类(表10):AUC 达 73.69%
  • 报告-体数据检索:Recall@1 达 24.12%

尽管 SuPreM [26] 在腹部分割任务中表现优异,但在胸部 VL 任务中提升有限。VoCo 的结果凸显了强大视觉编码器在 VL 任务中的关键作用------它能为语言模型提供更精确的视觉语义信息。


4.2.5 讨论

整体性能提升。如图6所示,在相同骨干网络下,VoCo 显著优于从零训练基线。在众多已有方法 [8], [9], [17], [27], [48], [54], [55], [66], [68], [69] 中,SuPreM [26] 表现最佳,但 VoCo 仍全面超越:

  • 在24个器官分割数据集上平均提升 2.93%
  • 在14个肿瘤分割数据集上提升 3.72%
  • 在15个胸部分析任务上提升 2.57%
  • 在28个预训练未见数据集 上提升 2.18%
  • 在13个跨模态数据集 上提升 3.52%
  • 在18个标注稀缺分割数据集 (<50例)上提升 2.72%

在全部48个下游任务上的持续改进,有力验证了 VoCo 的有效性。

对未见数据集的迁移能力 。如表2所示,评估基准包含28个预训练阶段未出现的数据集 。图6显示,VoCo 在这些数据集上平均比 Swin-B 基线提升 3.53%,展现出强大的泛化能力。

跨模态迁移能力 。我们在 CT 数据上预训练,随后迁移到另一种3D医学成像模态------MRI 。评估基准包含13个 MRI 数据集(涵盖分割、配准等任务)。图6显示,VoCo 在这些数据集上平均提升 3.52% ,证明其具备良好的跨模态迁移能力

标注高效性 (Label-efficient)。由于标注成本高昂,许多3D医学数据集仅有少量标注样本。表2显示,有18个分割数据集微调时标注样本少于50例。图6表明,VoCo 在此类标注稀缺场景下表现尤为突出,是理想的标注高效解决方案。

预训练骨干网络对比 。我们同时使用 nnUNet [48] 和 SwinUNETR [129] 进行预训练。尽管 nnUNet 是强大的分割基线,但其架构不可扩展(仅3100万参数)[64]。因此,我们重点研究 SwinUNETR 的缩放规律。分析表明,预训练的 SwinUNETR 在34个分割数据集上平均 DSC 提升 +3.34% ,而预训练的 nnUNet 仅提升 +1.98%(见表5、6)。nnUNet 提升有限可能源于其独特的"数据指纹"预处理策略与本文流程存在差异。

高效微调 。以往研究 [12], [20], [26] 已证明,强预训练模型可显著加速收敛 ,在更少训练轮次下达到更优性能。如图7所示,VoCo 在 BTCV [81]、SegTHOR [113] 和 TotalSegmentator [88] 上大幅加快收敛速度,该现象在全部48个任务中普遍存在。这对计算资源密集型任务(如 TotalSegmentator [88])尤为重要,VoCo 预训练模型有望显著降低医学图像分析的计算成本,推动高效学习范式的发展。

失败案例分析 。尽管 VoCo 在48个数据集中均取得至少1%的提升,但在少数任务中增益有限:5个数据集提升小于1.5%。例如 AutoPET [122] (正电子发射断层扫描,PET)数据集因成像原理与 CT 差异巨大,存在显著域偏移(domain gap),限制了预训练效果的迁移。这类挑战性任务提示未来需进一步探索多模态联合预训练策略。

4.3 医学图像分析中的缩放规律(Scaling Law)

更大的模型总是更好吗?

在医学任务中,答案似乎是否定的。从图8可以看出,对于某些特定任务,较小的模型反而能取得更优性能 。本文深入探究了影响模型容量缩放规律的关键因素,包括:微调样本数量、数据多样性以及任务难度

如图8所示:

  • (a) TotalSegmentator [88] 是一个极具挑战性的数据集,包含1,200例样本和104个分割类别。在此任务中,最大的模型 VoCo-H 取得了最佳结果。
  • (b) BTCV [81] 仅提供24例用于微调,有限的数据容易导致大模型过拟合,从而损害验证性能。
  • (c) CC-CCII [121] 虽有4,200例训练样本,但仅为简单的二分类任务(准确率已超90%),表明使用过大的模型并无必要。
  • (d) OASIS [127] 是一个脑部MRI配准数据集,仅有约400例,且结构多样性有限。在此场景下,最小的模型 VoCo-B 表现最佳。
  • (e) CT-RATE [105] 包含5万例样本,用于18类词汇分类任务。得益于大规模训练数据,更大的模型展现出更强的性能。

为不同医学任务定制模型规模

基于上述实验洞察,我们经验性地提出以下简单而合理的模型选择指南

  1. 微调阶段拥有大量标注数据的任务,更可能从大模型中受益;
  2. 覆盖多样化解剖区域的任务,适合使用大模型以捕捉复杂空间关系;
  3. 类别数量多、任务难度高的场景(如多器官分割),更适合大模型。

尽管这些准则已在我们构建的全面基准上得到验证,但由于医学领域的高度多样性,它们未必适用于所有任务。因此,我们开源了多种规模的预训练模型(VoCo-B/L/H),以帮助研究者根据具体需求选择最合适的模型。


4.5 定性可视化结果

上下文位置预测。如图10所示,我们展示了上下文位置预测的可视化结果。预训练过程中,位置预测损失稳定收敛。VoCo 生成的位置预测结果与真实位置高度一致,验证了所提预训练任务的有效性。

图10说明

(1) 左侧展示了上下文位置预测结果。我们对预测 logits 设定阈值,输出最可能的位置。预测结果与随机裁剪块的原始位置高度吻合。左下角为一个失败案例------两个区域具有相似解剖结构,导致模型混淆。

(2) 右侧显示,位置预测损失在预训练过程中稳定下降。我们进一步对比了从零训练模型和 SuPreM [26] 预训练模型的位置预测能力。值得注意的是,SuPreM 通过有监督分割预训练 也隐式提升了上下文位置预测能力,这间接表明:分割性能与我们提出的上下文位置预测任务之间存在正相关性

定性分割结果 。图11展示了不同解剖区域的分割可视化结果,涵盖腹部、胸部、盆腔等多个部位。结果表明,我们的方法能够广泛适用于各类下游医学图像分析任务,展现出良好的泛化性和鲁棒性。

5 局限性与未来方向

尽管我们的预训练方法在多种医学任务中展现出令人鼓舞的结果,但仍存在若干局限性,值得在未来进一步探索:

  • 用于提升数据质量的数据引擎 :当数据规模从1万例扩展到16万例时,性能提升趋于边际效应。尽管我们已对预训练数据集进行了整理与预处理,PreCT-160K 中仍不可避免地包含大量低质量样本。数据质量在预训练中至关重要 ,是充分发挥大规模数据潜力的关键 [20], [22], [24], [143]。未来,我们将聚焦于构建数据引擎(data engines),以系统性提升数据集质量。

  • 增强数据多样性以涵盖独特的成像特性 :如第4.2.5节所述,VoCo 在少数具有独特成像特征的下游任务上提升有限。鉴于医学影像模态和成像协议的高度多样性,我们将在未来进一步扩充预训练数据集的多样性,涵盖更多解剖区域、病理类型和成像设备。

  • 多模态预训练 :本研究仅使用CT数据进行3D医学图像预训练。未来,我们将构建大规模MRI预训练数据集,并与CT数据结合,推动多模态3D医学图像预训练的发展,以提升模型在跨模态任务中的泛化能力。

  • 可扩展的网络架构 :本研究并未聚焦于设计新型3D医学图像分析网络架构。当前采用的骨干网络 [48], [129] 在大规模预训练场景下可能缺乏良好的可扩展性。未来,我们将深入探索面向大规模预训练的先进网络架构,以更好地匹配数据与模型的缩放需求。

  • 进阶的全监督学习策略 :如第4.4节所示,全监督预训练(omni-supervised pre-training)已优于纯自监督学习(SSL)。未来,我们将进一步优化全监督学习策略,更高效地协同利用标注与无标注数据,释放两者互补潜力。

6 结论

本文提出了一种简洁而高效的体素对比 (Volume Contrast, VoCo)框架,用于大规模3D医学图像预训练。受不同器官之间稳定几何关系 的启发,我们利用几何上下文先验 来学习一致的语义表征,从而提升自监督学习效果。VoCo 还可无缝集成到半监督学习框架中,实现全监督预训练

为推动大规模3D医学图像预训练研究,我们构建了当前规模最大的医学图像预训练数据集 PreCT-160K ,包含16万例CT体数据(约4200万张切片),覆盖多样化解剖结构。我们进一步深入研究了模型容量的缩放规律 ,并提出了针对不同医学任务定制模型规模的实用指南

为全面评估预训练效果,我们建立了涵盖48个下游数据集的综合评估基准,覆盖分割、分类、配准、视觉-语言等多种任务。大量实验表明,VoCo 在各项任务上均显著优于现有方法,验证了其有效性与泛化能力。

相关推荐
大模型服务器厂商2 小时前
适配的 GPU 服务器能让 AI 模型充分发挥算力优势
人工智能
AscendKing2 小时前
LandPPT - AI驱动的PPT生成平台
人工智能·好好学电脑·hhxdn.com
FreeCode2 小时前
LangChain1.0智能体开发:流输出组件
人工智能·langchain·agent
故作春风2 小时前
手把手实现一个前端 AI 编程助手:从 MCP 思想到 VS Code 插件实战
前端·人工智能
人工智能训练3 小时前
在ubuntu系统中如何将docker安装在指定目录
linux·运维·服务器·人工智能·ubuntu·docker·ai编程
掘金一周3 小时前
没开玩笑,全框架支持的 dialog 组件,支持响应式| 掘金一周 11.6
前端·人工智能
电鱼智能的电小鱼3 小时前
基于电鱼 ARM 边缘网关的智慧工地数据可靠传输方案——断点续传 + 4G/5G冗余通信,保障数据完整上传
arm开发·人工智能·嵌入式硬件·深度学习·5g·机器学习
Juchecar3 小时前
翻译:Agentic AI:面向企业应用的智能
人工智能
武子康3 小时前
AI研究-121 DeepSeek-OCR 研究路线:无限上下文、跨模态抽取、未来创意点、项目创意点
人工智能·深度学习·机器学习·ai·ocr·deepseek·deepseek-ocr