BerDiff:用于医学图像分割的条件伯努利扩散模型

摘要：

医学图像分割是一项具有挑战性的任务，具有固有的模糊性和高度的不确定性，这主要是由于肿瘤边界不明确和多个似是而非的注释等因素。分割口罩的准确性和多样性对于在临床实践中为放射科医生提供有价值的参考至关重要。虽然现有的扩散模型在各种视觉生成任务中显示出强大的能力，但在分割中处理离散掩码仍然是一个挑战。

为了实现准确多样的医学图像分割掩模，提出了一种新的条件伯努利扩散医学图像分割模型(BerDiff)。我们首先提出使用伯努利噪声作为扩散核来代替高斯噪声来增强扩散模型对二值分割任务的能力，从而得到更精确的分割掩模。其次，通过利用扩散模型的随机特性，我们的BerDiff对初始伯努利噪声和中间潜在变量进行多次随机采样，以产生一系列不同的分割掩模，这可以突出显示感兴趣的突出区域，从而为放射科医生提供有价值的参考。

此外，我们的BerDiff可以有效地从反向扩散的整体轨迹中采样子序列，从而加快了分割过程。在两种不同模式的医学图像分割数据集上的广泛实验结果表明，我们的BerDiff优于其他最近发表的最先进的方法。

我们的研究结果表明，扩散模型可以作为医学图像分割的强大支柱。

1 介绍

医学图像分割在更好的诊断、手术计划和图像引导手术中起着至关重要的作用[8]。医学图像固有的模糊性和高度不确定性给准确分割带来了重大挑战[5]，原因包括脑磁共振成像(MRI)图像中肿瘤边界不清晰，肺结节计算机断层扫描(CT)图像中存在多种似是而非的注释。现有的医学图像分割工作通常提供一个单一的，确定性的，最有可能的假设口罩，这可能导致误诊或不理想的治疗。因此，为放射科医生提供准确多样的分割掩模作为有价值的参考[16]，在临床实践中至关重要。

最近，扩散模型[10]在各种视觉生成任务中显示出强大的能力[19,20]。然而，如何更好地与离散分割任务相结合需要进一步考虑。虽然许多研究[1,24]将扩散模型与分割任务结合起来并进行了一些修改，但这些方法没有充分考虑到分割任务的离散性，仍然使用高斯噪声作为其扩散核。为了实现准确多样的医学图像分割，提出了一种新的条件伯努利扩散模型(BerDiff)。我们首先提出使用伯努利噪声作为扩散核来代替高斯噪声来增强扩散模型的分割能力，从而得到更精确的分割掩模。此外，通过利用扩散模型的随机特性，我们的BerDiff对初始伯努利噪声和中间潜在变量进行多次随机采样，以产生一系列不同的分割掩模，这可以突出显示感兴趣的突出区域(ROI)，这可以作为放射科医生的有价值的参考。此外，基于去噪扩散隐式模型(Denoising diffusion Implicit Models, DDIM)[23]背后的原理，我们的BerDiff可以有效地从反向扩散的总体轨迹中采样子序列，从而加快分割过程。

本工作的贡献总结如下。

1)本文提出了一种新的基于伯努利噪声的条件扩散模型来代替高斯噪声用于离散二值分割任务，实现了准确多样的医学图像分割掩模。

2)我们的BerDiff可以有效地从反向扩散的整体轨迹中采样子序列，从而加快了分割过程。

3)在CT和MRI两种医学图像上的实验结果，特别是LIDC-IDRI和BRATS 2021数据集，表明我们的BerDiff优于其他最先进的方法。

2 方法

在本节中，我们首先描述了问题的定义，然后演示了我们的BerDiff的伯努利正向和多种反向过程，如图1所示。最后，我们提供了培训和抽样程序的概述。

图1所示。我们的BerDiff的伯努利正反过程的说明。

2.1 问题定义

设x∈RH×W×C为输入的医学图像，其空间分辨率为H×W，通道数为C。ground-truth mask表示为y0∈{0,1}H×W，其中0代表背景，1代表ROI。受去噪扩散概率模型(DDPM)和DDIM等基于扩散的模型的启发，我们提出了一种新的条件伯努利扩散模型，该模型可以表示为pθ(y0:T |x):= ʃ pθ(y0:T |x)dy1:T，其中y1，…， yT是与掩码y0大小相同的潜在变量。用于医学二值分割任务中，我们的BerDiff的多样化反向过程从初始伯努利噪声yT ~ B(yT;1 2·1)，通过受输入医学图像x约束的中间潜变量进行推进，生成分割掩模，其中1表示大小为H ×W的全一矩阵。

2.2 伯努利正演过程

在上一代相关的扩散模型中，高斯噪声随着时间步长t的增加而逐渐增加。然而，对于分割任务，基底真值掩模由离散值表示。为了解决这个问题，我们的BerDiff使用噪声时间表逐渐增加更多的伯努利噪声β1，…， βT，如图1所示。我们的BerDiff的伯努利前向过程q(y1:T |y0)是一个马尔可夫链，可以表示为:

式中B为伯努利分布，其概率参数为(1 - βt)yt - 1 + βt/2。使用αt =1−βt和¯αt = Ⅱt τ=1ατ的符号，我们可以在任意时间步长t上有效地对yt进行如下采样:

为了确保第2.4节中描述的目标函数易于处理和计算，我们使用采样的伯努利噪声∊ ~ B(?;1−¯αt2·1)将Eq.(3)的yt重新参数化为y0⊕∊，其中⊕表示“exclusive”的逻辑运算。

或(XOR)”。另外，让⊙表示元素积，Norm(·)表示沿通道维度对输入数据进行规范化，然后返回第二个通道。具体的伯努利后验可以表示为:

2.3 多样化的逆向过程

不同的反向过程pθ(y0:T)也可以看作是从伯努利噪声yT ~ B(yT;1/2·1)，通过受输入医学图像x约束的中间潜变量，产生不同的分割掩码，如图1所示。我们的BerDiff的具体多样化逆向过程可以表示为:

具体来说，我们利用估计的伯努利噪声。∊ʌ(yt, t, x)通过校准函数FC参数化µʌ(yt, t, x)，如下:

其中|·|为绝对值运算。

2.4 详细步骤

在这里，我们概述了算法1和算法2中的训练和抽样过程。在训练阶段，给定图像和掩码数据对{x, y0}，我们从均匀分布{1，…T}中采样随机时间步长t，用来对伯努利噪声∊进行采样。

然后我们用∊从q(yt | y0)中抽取yt，得到伯努利后验q(yt−1 | yt, y0)。我们通过估计的伯努利噪声∊ʌ(yt, t, x)通过校准函数FC参数化pθ(yt−1 | yt, x)。

基于以往扩散模型[3]中负对数似然的变分上界，我们采用Kullback-Leibler (KL)散度和二进制交叉熵(BCE)损失对我们的BerDiff进行如下优化:

最后，总体目标函数表示为:

其中λBCE在实验中设为1。

在采样阶段，我们的BerDiff首先对初始潜在变量yT进行采样，然后迭代计算不同t下yT−1的概率参数。在算法2中，我们对潜在变量yT−1提出了DDPM和DDIM两种不同的采样策略。最后，我们的BerDiff能够产生不同的分割掩码。通过对这些掩模取均值，我们可以进一步得到一个显著性分割掩模，以突出ROI，为放射科医生提供有价值的参考。值得注意的是，我们的BerDiff提出了一种新的参数化技术，即校准函数，来估计yt的伯努利噪声，这与以前基于离散状态扩散的模型[3,11,22]不同。

3 实验

3.1 实验设置

数据集与预处理 

本实验使用的数据来自LIDC-IDRI[2,7]和BRATS 2021[4]数据集。LIDC-IDRI包含1,018个肺部CT扫描，由四位放射科医生注释的合理分割掩码。

我们采用了标准的肺部CT扫描预处理流水线和训练验证-测试分割，与之前的工作一样[5,14,21]。BRATS 2021由每个患者的四种不同序列(T1、T2、FlAIR、T1CE) MRI图像组成。

所有3D扫描切片为轴向切片，丢弃底部80片和顶部26片。需要注意的是，我们遵循之前的工作[23]，将原来的四种脑肿瘤类型作为一种类型，将多目标分割问题转化为二值分割。我们的训练集包括来自1126名患者扫描的55174张2D图像，测试集包括来自125名患者扫描的3991张2D图像。

最后，将来自LIDC-IDRI和BRAST 2021的图像分别调整为128 × 128和224 × 224。

实现细节 我们使用PyTorch库实现所有方法，并在NVIDIA V100 gpu上训练模型。所有网络都使用批大小为32的AdamW[17]优化器进行训练。最初的

BRATS 2021的学习率为1 × 10−4,LIDC-IDRI的学习率为5 × 10−5。

我们的BerDiff图1中的伯努利噪声估计U-net网络与之前基于扩散的模型[18]相同。对于所有扩散模型，我们采用T = 1000时间步长的线性噪声调度。并采用DDIM的子序列采样策略来加快分割过程。在LIDC-IDRI的小批量训练过程中，我们的BerDiff通过从每个图像的四个带注释的分割蒙版中随机抽取一个来学习不同的专业知识。三个指标用于性能评估，包括广义能量距离(GED)，匈牙利匹配交集优于联合(HM-IoU)和骰子系数。我们使用不同数量的分割样本(1、4、8和16)计算GED，使用16个样本计算HM-IoU。

3.2 消融研究

我们首先进行烧蚀实验，以验证不同损失和估算目标的有效性，如表1所示。所有实验都在LIDC-IDRI上进行了21,000次训练迭代。我们首先进行了不同损耗的烧蚀研究，同时估计了前三排的伯努利噪声。我们发现KL散度和BCE损失相结合可以获得最佳性能。在此基础上，对下两行估计目标的选取进行了消融研究。我们观察到估计伯努利噪声比直接估计真值掩模更适合于我们的二值分割任务。这些发现与前人的研究结果一致[3,10]。关于采样策略和采样时间步长的额外消融研究请参见附录A。

在这里，我们用高斯噪声或伯努利噪声对我们的BerDiff进行烧蚀实验，结果如表2所示。对于离散的分割任务，我们发现使用伯努利噪声可以产生良好的效果

图2所示。LIDC-IDRI随机选取的两个肺结节的不同分割掩膜及相应的显著性掩膜。Xi0和xight分别是第i个生成的和ground-truth分割掩码。显著性掩码是多种分割掩码的均值。

训练迭代是有限的(例如21,000次迭代)，当训练迭代足够时(例如86,500次迭代)，甚至优于使用高斯噪声。

我们还在附录b中提供了基于伯努利和高斯的扩散模型在训练迭代上的更详细的性能比较。此外，我们在附录C中提供了一个玩具示例来证明伯努利扩散优于高斯扩散。

3.3 与其他先进方法的比较

在LIDC-IDRI上的结果  这里，我们在表3中给出了LIDC-IDRI的定量比较结果，发现我们的BerDiff在离散分割任务中表现良好。概率U-net (Prob.U-net)，层次问题。U-net (Hprob.U-net)和Joint prob.net。U-net (JPro.U-net)使用条件变分自编码器(cV AE)来完成分割任务。校准Adversar -

图3所示。BRATS 2021中随机选取的四张MRI图像的分割掩码。这里介绍的基于扩散的模型(SegDiff和我们的)的分割蒙版是显著性分割蒙版。

神经网络细化(CAR)采用生成对抗网络(GAN)来细化分割。PixelSeg是基于自回归模型，而SegDiff和MedSegDiff是基于扩散的模型。我们有以下两个观察结果:

1)对于离散分割任务，基于扩散的方法比基于VAE、GAN和自回归模型的传统方法具有显著的优势;

2)我们的BerDiff优于其他使用高斯噪声作为扩散核的基于扩散的模型。同时，我们给出了对比分割结果，如图2所示。与其他模型相比，我们的BerDiff可以有效地学习不同的专业知识，从而获得更多样化和准确的分割掩码。特别是对于可能产生歧义的小结节，例如左侧的肺结节，我们的BerDiff方法产生的分割掩模更符合ground-truth掩模。

BRATS 2021的结果  

在这里，我们分别在表4和图3中给出了BRATS 2021的定量和定性比较结果。

我们将我们的BerDiff与其他模型进行了比较分析，如nnU-net、基于变压器的模型(如TransU-net和Swin UNETR)以及基于扩散的方法(如SegDiff)。首先，我们发现与传统的U-net和基于变压器的方法相比，基于扩散的方法表现出优越的性能。此外，与我们的噪声估计网络具有相同架构的U-net所取得的高性能，突出了基于扩散模型的骨干设计的有效性。此外，我们提出的BerDiff优于其他使用高斯噪声作为扩散核的基于扩散的模型。最后，从图3中，我们发现我们的BerDiff片段在人眼难以识别的部分上更加准确，比如第3行的肿瘤。同时，我们还可以生成多种似是而非的分割掩码，从而得到显著性分割掩码。我们注意到，其中一些掩码可能是假阳性，如第一行所示，但由于显着性较低，它们可以被过滤掉。请参阅附录D，以获得由我们的BerDiff生成的不同分割掩码的更多示例。

4 结论

我们首先提出使用伯努利噪声作为扩散核来增强扩散模型对二值分割任务的能力，实现准确多样的医学图像分割结果。与其他基于扩散的模型相比，我们的BerDiff只关注二值分割任务，并且在迭代采样过程中花费大量时间;例如，我们的BerDiff对一幅医学图像的分割时间为0.4s，是传统U-net的10倍。在未来，我们将把我们的BerDiff扩展到多目标分割问题，并实施额外的策略来加快分割过程。