【Diffusion分割】体素医学图像分割的通用半监督框架

       

         在三维医学图像中进行体积标注是一项耗时且需要专业知识的任务。因此，人们对使用半监督学习（SSL）技术利用有限的标注数据训练模型越来越感兴趣。然而，所面临的挑战和实际应用已超出了 SSL 的范围，扩展到了无监督领域适应（UDA）和半监督领域泛化（SemiDG）等环境。这项工作旨在开发一个通用的 SSL 框架，以处理所有这三种情况。在现有的 SSL 框架中，我们发现了实现这一目标的两个主要障碍：1）捕捉分布不变特征的弱点；2）未标记数据被标记数据淹没的趋势，导致训练过程中对标记数据的过度拟合。为了解决这些问题，我们提出了一个聚合与解耦框架。聚合部分由扩散编码器组成，该编码器通过从多个分布/领域的聚合信息中提取分布不变特征来构建通用知识集。解耦部分由三个解码器组成，将训练过程与标注数据和非标注数据解耦，从而避免对标注数据、特定领域和类别的过度拟合。我们在 SSL、Class-imbalanced SSL、UDA 和 SemiDG 四个基准数据集上评估了我们提出的框架。结果表明，在所有四种情况下，与最先进的方法相比，我们的框架都有明显的改进，这表明我们的框架有潜力应对更具挑战性的 SSL 场景。

介绍：

        标注体素医学图像需要专业知识，而且是一个耗时的过程。因此，使用半监督学习（SSL）来训练标注数据有限的模型是非常可取的。各种 SSL 技术已被提出，特别是在半监督体积医学图像分割（SSVMIS）领域，以充分利用标记和非标记数据。然而，目前的 SSVMIS 方法 都假设标记数据和非标记数据来自同一领域，这意味着它们具有相同的分布。在实践中，医学图像通常是使用不同的扫描仪从不同的临床中心采集的，这就造成了显著的域偏移。这些偏移是由患者群体、扫描仪和扫描采集设置的差异造成的。因此，这些 SSVMIS 方法在实际应用场景中存在局限性，经常出现过拟合问题，导致结果不理想。

        为解决这一局限性，研究人员越来越多地关注无监督域适应（UDA）技术。这些技术同时利用标记数据（源域）和非标记数据（目标域）但这些数据来自不同的领域。此外，半监督领域泛化（SemiDG）这种更严格的方案也引起了人们的极大兴趣。半监督领域泛化（SemiDG）在训练过程中利用了来自多个领域的标记和非标记数据，并在一个未见过的领域中进行评估。目前，针对这三种情况的方法都是单独优化的，还没有一种方法能在统一的框架内解决所有三种情况。不过，考虑到所有训练阶段都涉及标注和未标注数据，探索一种通用的。

        基于 SSL 的框架可以处理所有设置，无需针对特定任务进行复杂的设计。因此，本文旨在开发一个通用框架，以应对现实世界场景中的现有挑战，包括：

场景 1：SSL（图 1(a)）：用于训练和测试的样本数据来自同一领域，代表标准 SSL 设置。

场景 2：UDA（图 1(b)）：采样数据来自两个域，目标域的标签无法访问，代表 UDA 设置。

情景 3：半数据采集（图 1(c)）：采样数据包含多个域，其中只有有限的几个域有标签，代表 SemiDG 设置。

        潜在的相似之处可归纳如下：(1）在训练阶段，既使用了标注数据，也使用了未标注数据；（2）在实际应用领域的场景中，无论是 SSL 中的分布偏移，还是 UDA 和 SemiDG 中的领域偏移，都可以视为采样偏差，即主要区别在于我们如何对图 1 中的数据进行采样。

        现在，我们不禁要问，现有的 SSVMIS 方法是否足够强大，可以处理这一一般任务。实验结果表明，现有的 SSL 方法在 UDA 和 SemiDG 设置下效果不佳。其中一个主要障碍在于这些模型存在严重的过拟合现象，这是由于训练过程中标注数据占主导地位造成的。具体来说，最先进的 SSVMIS 方法主要基于两个框架：(1) 教师-学生框架[1]，即首先用标注数据训练学生模型，然后由学生模型的 EMA 生成伪标签，再用标注数据重新训练学生模型，见图 2(a)；(2) 交叉伪监督（CPS）框架，即利用两个扰动模型之间的一致性，将其中一个网络生成的伪标签用于训练另一个网络，见图 2(b)。这两个主要框架中的预测模块都是通过标注数据和非标注数据进行训练的；不过，与非标注数据相比，标注数据在精确的地面真相监督下收敛得更快。因此，如图 3 所示，训练过程很容易被有监督的训练任务所淹没。另一个挑战在于，现有的 SSVMIS 方法无法解决分布变化的问题，更不用说域变化了，这就限制了捕捉不受分布变化影响的特征。 

        基于主流 SSVMIS 方法的相似性和主要弱点，我们认为，如果能解决过拟合问题并设计一种强大的方法来捕捉分布不变的特征，那么通用框架是可能的。为了解决上述问题，并设计出适用于现实世界应用场景的通用 SSVMIS 方法，本研究提出了一种新颖的聚合与解耦（A&D）框架。具体来说，A&D 由聚合阶段和解耦阶段组成。在聚合阶段，基于最近成功的扩散模型[19, 20]，我们提出了一个 Diff-VNet 将多域特征聚合到一个共享编码器中，以构建一个共同的知识集，从而提高捕捉分布不变特征的能力。为了解决过拟合问题，在解耦阶段，我们将解码过程解耦为：（1）有标注数据训练流，主要更新 Diff-VNet 解码器和难度感知 V-Net 解码器，以生成高质量的伪标签；（2）无标注数据训练流，主要更新另一个虚V-Net 解码器，并对伪标签进行监督。Diff-VNet 解码器的去噪过程提供了领域无偏的伪标签，而难度感知 V-Net 解码器则提供了类无偏的伪标签。我们还提出了一种重新参数化与平滑化相结合的策略，以进一步提高伪标签的质量。 

        我们工作的主要贡献可归纳如下：(1) 我们用一个通用框架统一了用于体积医学图像分割的 SSL、Class Imbalanced SSL、UDA 和 SemiDG；(2) 我们指出了当前 SSL 方法的过拟合问题，并提出通过高效的数据增强策略和解耦解码器来分别解决标记数据和非标记数据的过拟合问题；(3) 我们引入了扩散 V-Net 来学习不同领域的底层特征分布，从而将 SSL 方法推广到更现实的应用场景中；(4) 所提出的聚合与解耦框架在 SSL、类不平衡 SSL、UDA 和 SemiDG 任务的代表性数据集上达到了最先进水平。值得注意的是，我们的方法在 Synapse 数据集（Dice 值为 12.3）和 MR 到 CT 设置中的 MMWHS 数据集（Dice 值为 8.5）上取得了显著改进。为了验证所提方法的有效性，我们进行了广泛的消融研究。 

方法：

聚合与解耦框架概述：

        在本节中，我们将介绍我们的聚合与解耦 (A&D) 框架，如图 4 所示，该框架由聚合阶段和解耦阶段组成。训练流水线如算法 1 所示：

        聚合阶段的目的是根据所有数据都具有共同的基础高级知识（如纹理信息）这一理念，构建跨领域的通用知识集。

        通过聚合来自多个领域的信息并进行联合训练，编码器可以捕捉底层的领域不变特征。为此，我们引入了一种强大而高效的基于采样的体积数据扩增（SVDA）策略，以扩大分布多样性，并利用扩散模型捕捉多样化数据的不变特征。 

        在现有 SSL 方法的解码过程中，解码器需要同时使用标注数据和非标注数据进行训练，这导致了耦合和过拟合问题，进一步阻碍了向一般 SSL 的扩展。解耦阶段旨在通过解耦有标签和无标签数据训练流来解决这些问题。具体来说，对于标注数据流，(1) 扩散解码器主要用于引导扩散编码器通过扩散后向过程学习分布不变的表示，从而生成无领域偏差的伪标签；(2) 虚V-Net解码器采用所提出的难度感知重加权策略，主要用于避免模型过度拟合易类和多数类，从而生成无类偏差的伪标签。然后，对于未标记的数据流，通过提出的重参数化与平滑（RS）策略对领域和类别无偏的伪标签进行集合，以生成高质量的伪标签。最后，伪标签将用于监督仅用于预测的附加 V-Net 解码器的训练。

Aggregating Stage：

        假设整个数据集由 NL 个已标注样本 {(xl i, yi)}NL i=1 和 NU 个未标注样本 {xu i }NU i=1 组成，其中 xi∈ RD×H×W 是输入量，yi∈ RK×D×H×W 是具有 K 个类别的地面实况注释。聚合阶段的目标是用 SVDA 增强数据，并将已标注（xl, y）和未标注数据 xu 编码为高级分布不变特征，用于去噪已标注数据流 hl;ξ、难度感知的已标注数据流 hl;ψ和未标注数据流 hu。

        基于采样的体积数据扩增（SVDA） 与文献[69]中使用的级联所有扩增操作的耗时的传统数据扩增方法不同，SVDA 建立在扩增集的基础上，对 Naug 操作进行随机采样，将其应用于已标记和未标记的数据。增强集包括三维空间变换（随机裁剪、随机旋转、随机缩放）和体素变换（高斯模糊、亮度、对比度、伽玛）。Naug 根据经验设定为 3。

        利用扩散捕捉不变特征 我们沿用 Diff-UNet ，使用扩散模型进行感知，但将其修改为 V-Net 版本，并删除了额外的图像编码器。给定带标签的体积数据 xl ∈ RD×W ×H 及其标签 y ∈ RD×W ×H，我们首先将标签转换为one-hot格式 y0 ∈ RK×D×W ×H，然后连续添加 t 步噪声 ε，得到噪声标签 yt ∈ RK×D×W ×H，这就是扩散前向过程：

        然后，将噪声标签与图像 xl 连接起来，作为 Diff-VNet 的输入。具体来说，不同数据流的高级特征是不同的。对于去噪流，即D(xl; ξ) 作为解码器，扩散编码器以 concat([yt, xl]) 和时间步长 t 作为输入，生成嵌入时间步长的多尺度特征 hl;ξ i∈ Ri×F × D 2i × W 2i × H 2i，其中 i 为阶段，F 为基本特征大小。对于难度感知训练流和无标记数据流，即作为解码器的 D(xl; ψ) 和 D(xu; θ)，编码器只将 xl 和 xu 作为输入，分别获得多尺度特征 hl;ψ i , hu i。请注意，hl;ψ i , hu i 与 hl;ξ i 的形状相同。

Decoupling Stage ：

        解耦阶段包括四个步骤：监督去噪训练生成无领域偏差的伪标签、监督难度感知训练生成无类别偏差的伪标签、伪标签集合两个伪标签以及无监督训练得到最终预测结果。

        使用扩散解码器 D(xl; ξ) 进行有监督的去噪训练 以 hl;ξ i 为输入，D(xl; ξ) 对特征进行解码，以预测清晰标签 y0 作为域无偏差的伪标签。目标函数定义如下：

        使用 D(xl; ψ) 进行有监督的难度感知训练 为缓解 SSVMIS 中常见的类不平衡问题，我们设计了难度感知再加权策略（DRS），以强制模型关注最难的类（即学习速度较慢、性能较差的类）。通过扩散解码器 D(xl;ξ)生成的概率图 pl;ξ，可以从两个方面对难度进行建模：学习速度和性能。我们使用种群稳定指数[70]来衡量每个类在eth迭代后的学习速度： 

其中，λk 表示第 k 个类别在伦理迭代中 pl;ξ 的骰子得分，△ = λk,e - λk,e-1，duk,e 和 dlk,e 表示未学习和伦理迭代后学习的类别。那么，我们将经过伦理迭代后的第 k 个类别的难度定义为：

 

        如果学习速度较快的类别具有较小的 dk,e，那么损失函数中相应的权重也会变小，从而减慢学习速度。在这种情况下，骰子得分较低的类别在损失函数中的权重会较大，这就迫使模型更加关注这些类别。第 k 个类别的总体难度感知权重定义为： 

其中，α 在实验中根据经验设定为 1 / 5，以减少异常值。有监督的难度感知训练的目标函数定义如下 :

        使用重参数化与平滑（RS）策略进行伪标记 利用去噪扩散隐含模型（DDIM）方法迭代扩散模型（E(xl, xu; ξ)+D(xl; ξ)）t 次，生成域无偏差的 pu;ξ 概率图。通过 D(xl;ψ)的停止梯度前向传递，可以得到类无偏差的 pu;ψ 概率图。我们将 pu;ξ 和 pu;ψ 组合在一起，生成高质量的伪标签。然而，当我们将这两个图组合起来时，发现去噪概率图 pu;ψ 过于稀疏，即每个类别的置信度都非常高。这一特性有利于完全监督任务，但在这种情况下，它会抑制 pu;ψ，对噪声和误差不具有鲁棒性。因此，我们使用 Gumbel-Softmax 对 pu;ψ 进行重新参数化，以增加一些随机性，并使用高斯模糊核来消除这一操作带来的噪声。最终的伪标签为 :

        使用 D(xu; θ) 进行无监督训练 最后，我们可以使用伪标签 yξ,ψ 对 D(xu; θ) 进行无监督训练。无监督训练的目标函数定义为:

 

        为了更好地利用领域和类别无偏差特征，我们还采用了知识蒸馏策略：θ = wema × θ + (1 - wema) × (ξ + ψ)/2, wema = 0.99。A&G 框架的总体训练函数为:

 

原文地址：https://blog.csdn.net/qq_43583311/article/details/139616748

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！