VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis 阅读

论文:https://arxiv.org/abs/2402.17300

code:Luffy03/VoCo：[CVPR 2024] VoCo：一种简单而有效的 3D 医学图像分析体积对比学习框架 (github.com)

摘 要

1. 介绍

近年来，在三维医学图像分析领域中应用的深度学习方法已取得了显著成果[52,21,…

自监督学习（Self-Supervised Learning,…

现有研究主要集中在基于信息重建的方法上

这种方法的基本假设是物体各部分之间存在某种几何关系…

最近提出的一些改进方法也沿此思路展开…

此外mask reconstruction等替代方案也在被探索……

尽管这些工作取得了一定成效……然而现有研究表明缺乏足够的高级语义知识会显著影响 downstream任务性能……为此我们提出了一种新的预训练策略……

_图1所示。(a)在三维医学图像中描绘了器官间的几何关联关系。(b)基于这些先验知识进行预训练阶段后，在图1(b)的基础上提出了体积对比度(VoCo)框架。具体而言，在本研究中提出了一种简洁且高效的体积对比度(VoCo)框架。该框架首先从不同的位置裁剪一组不重叠的体积块，并强制要求其间的特征差异显著。随后将这些体积块表示为学习得到的高维空间中的基底集合，并在该空间中将它们作为不同位置的类别归属依据。为此我们需要随机裁剪子卷并利用它们与不同碱基的相似性来进行分类预测任务。在此基础上构建了一个三维医学图像SSL环境下的上下文位置预测任务模型。

基于对上下文位置的学习, 我们通过隐式建模将高级语义作为先验知识融入网络结构, 这一过程显著提升了 downstream 任务的表现. 在六个不同领域的广泛实验中发现, 所提出的 VoCo 体系在 3D 医学图像领域表现出超越现有最先进 SSL 方法的优势.

图2展示了典型的对比学习框架。具体而言，在示例级对比学习中（标记为(a)），我们通过施加强烈的输入数据增强或模型扰动来生成实例的不同表征，并通过正则化机制保证这些表征的一致性。（b）在 prototype-level对比学习中，则采用两种不同的途径：一是进行在线聚类或者二是进行随机初始化，并随后通过在线更新过程生成prototype作为基准。接着利用prototype对输入图像进行对比。（c）本研究提出的VoCo方法正是继承了 prototype-level对比学习的思想精髓。具体而言，在实现上并非依赖于耗时的在线聚类与更新程序，在此我们创新性地基于丰富的3D医学图像上下文位置信息，并基于基础作物作为基准构建prototype完成这一任务。

2. 相关的工作

在本节中

constrasive学习。如图2(c)所示，为了解决上述存在的问题，VoCo不是随机初始化和更新原型，而是利用3D医学图像中有价值的上下文位置先验来生成基础作物作为原型，这也不需要在大数据集上进行耗时的聚类。用于医学图像分析的SSL。由于在标签高效学习方面的巨大潜力[29,56,59,57,58,37]，SSL在医学图像分析领域也受到了极大的关注[68,32,31,50,19]。现有方法主要基于比较SSL[69]。具体而言，Zhou等[67]将Mixup[64]与MoCo[29]结合，了解InfoNCE中阳性和阴性样本的多样性[43]。Azizi等人使用多实例学习来比较来自每个患者的多个图像视图。还有许多方法[25,68,69]通过从原始图像中恢复低级信息来监督模型。在三维医学图像分析中，重建原始信息是学习表征的常用借口任务[49,50,69]。现有的方法主要是基于增强图像的信息重构。这些方法首先进行强数据增强，如旋转[50,73,51]、多视图作物[68,69,32]、掩码[13,71,54]，然后通过重建原始3D信息来监督模型。尽管已经证明了有希望的结果，但这些方法中的大多数仍然在很大程度上忽略了将高级语义集成到模型表示中的重要性，这严重阻碍了下游任务的性能。Position-related SSL。在自然图像领域，与位置相关的SSL方法也在许多先前的工作中进行了探索[8,9,40,46,42,17,62,66]。Noroozi等[42]提出预测一组洗牌的顺序。Zhai等[62]和Caron等[8]提出训练ViT[18]来预测每个输入patch的位置。然而，由于自然图像中不同物体的几何关系并不十分一致，因此仅给出视觉外观仍然难以有效地学习一致的位置表示(如[66]所述)。此外，以往的作品[62,8,66]主要训练线性层直接输出位置，采用黑盒方式。在本文中，我们将上下文位置预测的托辞任务引入到三维医学图像领域，其中不同器官之间的几何关系是相对一致的，这指导我们在预训练中学习一致的语义表示。与以往的方法不同，本文引入了一种完全不同的位置预测范式。具体来说，我们不是直接使用线性层来输出位置，而是基于体积对比度来预测上下文位置，这更加直观和有效。

3. 方法

在当前章节里主要介绍了小节3.1提出的VOCo总体架构。随后将在第3.2小节详细阐述上下文位置预测的具体过程。最后将在第4.2小节深入讨论VOCo框架中的体积对比正则化机制。

3.1. 总体框架

本研究开发的VoCo整体架构如图3所示。该架构主要由上下文位置推断分支和正则化增强模块组成。其中，
推断分支旨在推断不同裁剪块之间的相对位置关系。
对于输入数据块，在每个采样点上将其分割成多个非重叠的小块。
随后，在每个分割点上随机提取一个小样本块，并通过主流网络模型（如CNN[30]或Transformer[18]）对其进行特征提取。
我们的目标是推断随机抽取的小样本块与其所在的基本块之间的相对位置关系。
与以往研究相比，在这里我们摒弃了线性分类器的传统做法。
为此目的，
本研究设计了损失函数Lpred用于评估推断效果。
此外，
为了增强特征区分度，
在后续阶段我们引入了另一个损失项Lreg。
这些技术方案的相关细节将在下文部分进行详细阐述。

见图3所示。VoCo的整体架构（如图3所示）主要由两个关键组件构成：一个是预测模块（Predictive Block），另一个是正则化模块（Regularization Block）。其中，预测模块的主要功能是通过分析各子块间的相对位置关系来推导其对应的类别信息；而正则化模块则是通过增强各基底间的特征区分度来优化分类性能，在此过程中实现了对各类别样本的精准分类目标。

3.2. Contextual Position Prediction 情境位置预测

给定输入的空间体素（spatial voxels），我们首先将之划分为n个互不重叠的基础体素集合（base blocks），确保能够覆盖整个输入的空间范围。随后，在每个基础块内提取出表征该区域典型特性的特征向量z（我们称之为基向量），这些基向量能够在不同位置上反映原始数据的基本属性信息。在此基础上，在先前的研究工作[10,11,29]基础上引入了一种线性变换模型（linear mapper），用于将基向量z映射至潜在的空间属性表示q中进行进一步分析与编码存储。随后，在整个过程中又进行了另一轮的数据增强操作：即对原始数据集进行随机采样操作以获得多样化的子样本集（random sub-volumes sampling），并将这些子样本嵌入到高维表征空间p中进行深度学习建模与分类训练任务的具体实现策略研究与算法设计工作流程研究

Detailed volume-based contrast is employed to enhance the accuracy of contextual position prediction, which is particularly useful in scenarios requiring precise spatial understanding.

基于主干网络与投影机提取特征后，在借鉴以往的SSA研究[10,11,29]的基础上进行三维自适应均值池化处理。随后将该操作作用于样本数据集上，并将其尺寸缩减至单维空间即p∈R 1×C、q∈R 1×C（其中C代表通道数量）。我们经验性地选择C=2048作为通道数量，并未采用固定参数而是根据实验结果进行了优化（如图2所示）。在此基础上通过余弦相似度计算得到logits值l = cos(p,q)

其中qi代表每种基础作物的投影特征... Li则表示气中p的相关性指数其取值范围介于0至1之间... 在计算Eq. 1的过程中为了避免出现特征崩溃现象我们需要暂停q的学习过程这一措施有助于提高模型稳定性... 当li较高时则意味着p与qi之间存在更大的重叠区域概率这为我们后续构建关联模型提供了理论基础即通过li值高低我们可以明确区分不同位置对相似度的影响进而实现更加直观的位置编码机制为了避免混淆我们没有训练一个不可知的黑色盒子线性层而是采用了基于体积对比的方法直接预测上下文位置这一方法具有更高的可解释性和实践意义。

Position labels generation过程如图4所示。在生成位置标签的过程中，在n = 4×4的基础作物中进行计算时（如图4所示），会得到n个类作业（即对应于每个基础作物的位置标签）。随后将随机裁剪区域与基础作物之间的重叠面积进行计算，并将重叠区域的比例分配为位置标签y值（其范围从0到1）。通过预测逻辑l与位置标签y之间的距离进行监督学习是该模型的关键步骤之一。其中关于基础作物数量n的具体设置将在第4.4节中详细讨论。

图4展示了完成位置标签的过程。该过程通过计算重叠区域占总面积的比例来确定位置标签y的值。例如，在图中随机裁剪的体积分别以权重比例分配给各个类别：具体来说，在类别的体积被赋予为 0.25 的情况下，则被赋予为 0.25 的概率；类别的体积被赋予为 0.1 的情况下，则被赋予为 0.1 的概率；类别的体积被赋予为 0.5 和 0.15 的情况下，则分别被赋予为这些数值的概率。

该研究引入了一个新的损失函数用于上下文位置预测的任务。该预测损失函数Lpred被设计为基础于熵的概念以确保其有效性。具体而言，在模型训练过程中我们首先通过评估预测逻辑l与真实位置标签y之间的相似性度量d从而建立了一个衡量标准以优化模型性能的指标

值得注意的是，VoCo预测了一个卷的上下文位置（与其所有上下文高度相似），因此无需一一对应：例如，在图4所示的位置中，高值li同时对应于yi>0的位置（i=5、6、9、10）。接下来我们计算了li与yi之间的距离（如公式2所示）。

3.3. Volume Contrast for Regularization 正则化的体积对比

我们的研究目标在于探究多样化的课堂作业（基数）及其与体积对比的关系。直观上而言，在体态分析中存在明显语义差异的不同器官通常会对应于不同的体态变化范围。基于此观察，在设计过程中我们需要最大限度地扩大不同基底之间的高维特征差异。为此，在设计正则化损失函数时我们将重点放在放大不同基z之间的特征差异上。首先，在模型训练阶段我们采用投影空间中的基向量q作为参考点并结合余弦相似度评估机制来衡量各基向量间的相似性关系；具体来说通过计算不同qi与qj之间的余弦相似度sij来辅助模型优化过程（如图所示）。

我们旨在通过归一化sij使其达到零值；通过在不同基之间的比较来强化这种差异。为此，我们定义损失函数Lreg为：

|。|表示绝对值。对于损失Lreg，我们的目标是将q优化为线性无关的基:

基于正则化项 L_{\text{reg}}, 我们旨在估计一组相互独立的基础向量, 以表征多维度特征的不同方向 [6]. 为此, 我们将估计一组更加能区分不同类别的课堂作业, 以监督最终的位置预测. 总体而言, 总体损失 L 是由 L_{\text{reg}} 和 L_{\text{pred}} 的组合而成.

其中λ被用来平衡这两个损失项之间的相对重要性，并且在经验实验中将其设定为1.0这一数值是因为我们假设这两种损失具有同等的重要程度。关于λ相关机制的研究内容则包含在附录部分中

4. Experiments 实验

本节首先阐述了在预训练以及下游任务中所采用的数据集。接着对VoCo的实现细节进行了概述。最后展示了我们的方法与当前最先进的SSL技术在3D医学图像处理方面的详细实验结果。更详细的补充信息可参见附录部分。

4.1. 数据集

在公平比较的基础上（为了与之前的工作 [50,54,68,69,13,71] 进行对比），我们还在相同的三个公共基准数据集中进行了预训练实验（即 BTCV [35], TCIA covid [14], 和 LUNA [47] 数据集），总共涵盖了约 1.6k CT 扫描用于预训练工作）。值得注意的是，在 BTCV 下游任务中（即 BTCV 下游实验 [35]），我们仅使用 BTCV [35] 和 TCIA Covid-19 [14] 进行预训练工作；而对于其他下游任务，则采用了所有三个基准数据集进行预训练工作。详细内容请参阅补充资料。

为了验证 VoCo 的有效性（即 VoCo 的有效性评估），我们在 BTCV [35], LiTs [4], MSD �脾脏 [1], MMWHS [72], BraTS 21 [48], 和 CC-CCII [65] 六个公共基准数据集中进行了下游实验（包括分割和分类任务）。其中前五个基准数据集用于分割任务（即分割任务），最后一个基准数据集用于 COVID-19 分类任务（即 COVID-19 分类）。需要注意的是：只有 BTCV [35] 被用作预训练阶段的数据；其余所有基准数据集在预训练过程中均不可见；此外，在评估跨模态泛化能力时（即将 CT 数据集上的模型转移至 MRI 数据集合 BraTS 21 上），我们采用了与现有文献一致的方法设置（即一致设置方法）；同时我们还在二维医学图像基准数据库上进行了性能评估（即性能评估）。详细内容请参阅补充资料。

4.2. Implementation details 实现细节

基于先前的研究工作（引用文献编号：[50,54]），本研究采用SwinUNETR架构进行预训练及下游应用）。本研究采用了AdamW优化器（引用文献编号：[38]）配合余弦衰减学习率策略（引用文献编号：[26]）来完成所有的实验工作）。在预训练阶段设定10万次迭代（引用文献编号：10^5），并采用滑动窗口推理策略（引用文献编号：[71]），以实现对现有研究的一致性评估（引用文献编号：[50,54,13,71]）。为了排除基础模型及其后处理组件的影响（引用文献编号：[34,36]），本研究特别关注模型核心性能的评估。详细内容可参考补充材料中的相关内容

该研究采用了自监督学习方法对医学图像进行处理，并与现有相关技术进行了系统对比分析。具体而言，在与Masked Autoencoder(MAE)[28,13]以及Masked Cross-attention(MoCo v3)[29,12]等主流自监督架构进行对比的基础上，我们还引入了其他先进方法如SimCLR[10]和SimMIM[61]作为参照指标。此外，在采用统一实验设置的基础上，我们进一步验证了Jiasaw[9]和PositionLabel[66]的有效性指标。通过对现有文献中大多数先进医用自监督 SSL 方法的系统性对比实验表明，在相同条件下取得更好的性能表现。

4.3. 下游任务实验

在BTCV数据集上的实验结果显示优于现有方法。本研究通过系统性实验验证了该方法在BTCV上的优势。具体而言，在比较现有技术时使用UNetR[27]作为基础架构，并引入统一架构Swin-UNETr[26]作为基准模型，在此前提下与其他研究工作进行了对比分析。从表1的数据结果可以看出，在通用SSL方法中表现最突出的是MoCo v3[29,12]这一模型，在所有测试指标中均未超过79.54%的Dice分数。这一性能表现低于预期，并且该模型严重依赖于大规模的数据以获取足够的负样本支持，在处理具有巨大计算负担的3D医学图像时显得不够实用。此外，在不同图像间的负相关关系限制了其在医学图像分类任务中的应用效果[29,12]。基于此，在当前研究框架下MAE[28,13]、SimCLR[10]以及来自SimMIM[61]的结果也未能显著超越我们的VoCo模型的表现（来自文献[13]的研究结果）。值得注意的是，在医用SSL任务中引入位置编码等额外信息（如Jigsaw[9]和PositionLabel[66])并未显著提升分类性能。综合以上分析可知，在考虑医学图像特性的情况下选择合适的SSL技术至关重要

实验表格展示了BTCV的研究成果（见参考文献[35]）。最高水平的结果已被成功获取（注：本研究采用无监督预训练方法作为基准模型）。*†*这一符号用于标识本研究中对所述方法进行了重新实现。多数研究来自文献[13,63,71]及其团队的研究成果。

表2展示了BTCV实验的结果[35]。最佳性能得以实现。“从零开始”表示未经过自监督预训练的有监督基准。“*”表示我们重新实现了该方法。这些结果主要源自文献[13、63、71]及其团队的研究工作。

本研究中基于MSD脾（文献编号1）与MMWHS（文献编号72）所开展的实验研究结果。其中符号†表示我们重新实现该方法并记录了相应的图像分割结果数值。

表4中 BraTS 21的研究结果[48]。WT对应于肿瘤整体，TC对应于肿瘤核心，ET对应于增强肿瘤。†表明我们重新实现了该方法。

划痕swwin - unetr[26]仅达到80.53%的Dice Score。通过VoCo预训练，我们以83.85%的Dice Score获得了3.32%的改进，这也明显优于现有的方法。在比较的方法中，GL-MAE[71]的Dice Score最高(82.01%)。我们的VoCo比它高出1.84%的Dice Score，这在这个数据集中是一个明显的改进。在未见过的数据集上有前景的性能。我们进一步在预训练的未见数据集上进行实验，即LiTs[4]、MSD脾脏[1]和MM-WHS[72]。LiTs的结果[4]如表2所示。我们根据[68,69,63]报道了比较方法的结果。由于scratch swwin - unetr[26]可以获得更高的Dice Score(93.42%)，我们进一步基于VoCo预训练了一个3D UNet[45]，目的是进行公平的比较。可以看出，通过VoCo预训练，swwin - unetr[26]提高了3.10%，Dice Score达到96.52%。以3D UNet[45]为骨干，VoCo也实现了96.03%的Dice Score，证明了VoCo在不同网络架构下的有效性。MSD脾脏[1]和MM-WHS[72]数据集的结果如表3所示。在以往的方法中，GVSL[32]在MSD脾脏[1]和MM-WHS[72]数据集上的Dice Score分别为95.47%和88.27%，而我们的VoCo分别以96.34%和90.54%优于以往的所有方法。MRI数据集泛化能力。为了验证在MRI数据集上的泛化能力，我们进一步评估了VoCo在BraTS 21上的性能[48]。如表4所示，VoCo达到78.53%的Dice Score，优于现有的最先进的方法，证明了VoCo的跨模型泛化能力。COVID-19分类评估。我们在表5中进一步评估了CCCCII[65]数据集上分类任务的性能。由于现有的SSL方法没有在这个数据集上进行实验，我们重现了相关的方法进行比较。可以看出，VoCo也可以达到90.83%的准确率，取得了优异的结果，证明了其在分类任务中的有效性。

本研究采用了图6展示了BTCV[35]数据集的分割结果定性可视化。我们对比分析了VoCo、MoCo v3（引用编号12）、PositionLabel（引用编号66）、SwinUNETR（引用编号50及26）以及GVSL（引用编号32）等模型。

六项不同下游任务的整体对比分析展示了图5中的结果。具体而言，在六项不同下游任务上系统性对比分析了现有的先进方案。研究表明，在性能指标上超越现有先进方案的性能表现突出

4.4. 消融实验。

为了进一步验证消融效果，在VoCo框架下系统性地考察其损失函数及其配置参数的优化效果，并通过BTCV[35]和MM-WHS[72]两个数据集上的实验验证来评估消融性能

损失函数。我们首先研究两个损失函数Lpred和Lreg的重要性，如表6所示。

研究表明，在应用我们所提出的Lpred损失函数后，在BTCV和MM-WHS数据集上的实验结果表明性能提升显著。具体而言，在BTCV上性能提升了80.53%至82.96%，而在MM-WHS上则实现了86.11%至88.82%的改进。这些数据验证了我们提出的位置预测借口任务的有效性。进一步采用Lreg的正则化损失策略能够进一步优化模型性能。实验证明，在VoCo框架中学习提取关键识别基是实现目标的关键因素。

基数设置对VoCo性能的影响是一个关键研究点。我们进一步分析了VoCo中碱基数n的不同取值方案。在对比消融研究中的不同设置时（如表7所示），我们发现由于Z轴方向上ROI大小不一的特性，在该方向上裁剪多个碱基并非实际可行方案。因为裁剪后必须调整体积大小以维持一致的尺度，在Z轴方向上进行多次裁剪会导致尺寸不协调的问题。此外，在计算资源有限的前提下增大n值会显著提升计算开销。根据表7的数据可以看出：当n设定为2×2×1时，在BTCV和MM-WHS上的Dice Score分别为81.56%和86.73%；而将n提升至3×3×1和4×4×1时模型性能明显增强。具体而言，在n=4×4×1的情况下BTCV和MM-WHS上的准确率分别达到了83.85%和90.54%。然而值得注意的是使用更高值如5×5×1并不能带来更好的效果表现（如表7所示）。为了全面评估模型性能我们还对Z轴方向上的基础作物生成能力进行了验证结果表明采用3×3×2和4×4×2的设置并未带来改进效果（如图8所示）。因此基于性能与效率平衡考虑我们将VoCo中的n参数稳定化为4×4×1这一最佳配置选择可见合理的n值设定对模型性能至关重要。

4.5. 可视化的结果

BTCV数据库中的分割效果展示[35]如图6所示。通过实验观察到, VoCo显著提升了分割的准确率与完整性.具体其他实验结果可在附录中查阅.

5. 结论及未来发展方向

在本文中,我们成功研发了一个简洁有效的SSL框架VoCo,用于3D医学图像分析领域。考虑到体内不同器官之间具有相对固定的上下文位置关系,我们在预训练阶段巧妙地引入了上下文位置先验,从而学习一致的语义表征。具体而言,我们从输入体的不同方位裁剪切片并将其作为基向量集合表示不同方向的特征信息;接着通过计算其与基向量集的相似度来预测任意裁剪切片的上下文位置信息;这种操作使得VoCo能够将上下文位置先验转化为模型内部表示,从而显著提升了对需依赖高级语义理解的任务性能支持能力。大量实验结果表明,该框架在六个典型的数据集上均展现出卓越的性能指标。展望未来研究,我们将重点探索以下几大扩展方向:(1)通过构建更具规模的数据集来进一步验证和提升VoCo框架的表现;(2)在更多数据集上开展系统性实验研究;(3)评估其在半监督学习等新领域的适用性;(4)深入研究其在发现体部组织间关系方面的潜力和局限性