Diffusion Models for Implicit Image Segmentation Ensembles

隐式图像分割集成的扩散模型

摘要

扩散模型在图像生成建模方面表现出令人印象深刻的性能。本文提出了一种新的基于扩散模型的语义分割方法。通过改进训练和采样方案，我们证明了扩散模型可以对医学图像进行病灶分割。为了生成特定于图像的分割，我们在地面真值分割上训练模型，并在训练和采样过程的每一步中使用图像作为先验。利用给定的随机采样过程，我们可以生成一个分割掩码的分布（确定每个像素或每个区域在图像或空间中是否属于特定的类别或对象。）。此属性允许我们计算分割的像素不确定性映射，并允许隐式的分割集合，从而提高分割性能。我们在BRATS2020数据集上评估了我们的方法用于脑肿瘤分割。与最先进的分割模型相比，我们的方法产生了良好的分割结果，此外，还提供了详细的不确定性图。

1. 介绍

语义分割是医学图像分析中一个重要且探索得很好的领域(Rizwan I Haque and Neubert, 2020)。利用机器学习对医学图像中的病变进行自动分割已经显示出良好的性能(Isensee et al .， 2021)，并且已经准备好用于临床应用以支持诊断(Sharrock et al .， 2021)。在医学应用中，测量给定预测的不确定性是非常有趣的，特别是当用于放射治疗等进一步治疗时。

在这项工作中，我们专注于BRATS2020脑肿瘤分割挑战(Menze et al, 2014;Bakas et al, 2017,2018)。该数据集为每个患者提供了四种不同的MR序列(即t1加权、t2加权、FLAIR和对比度增强的t1加权)，以及基于像素的地面真值分割。在附录A中可以找到一个示例图像。

我们提出了一种基于去噪扩散概率模型(DDPM)的新分割方法(Ho et al .， 2020)，该方法可以提供生成的分割掩模的不确定性映射。BRATS2020数据集图像的工作流程概述如图1所示。我们在分割蒙版上训练一个DDPM，并将原始脑磁共振图像作为图像添加，从而诱导出解剖信息。作为抽样

ddpm在每个采样步骤中都有一个随机元素，我们可以为相同的输入图像和相同的预训练模型生成许多不同的分割掩码。这种分割集合允许我们计算逐像素的方差图，从而可视化生成的分割的不确定性。此外，在平均映射中集成分割，提高了分割性能。

我们将自己与最先进的分割算法进行比较，并直观地将我们的方差图与常见的不确定性图进行比较。代码将在https://github.com/JuliaWolleb/Diffusion-based-Segmentation上公开。

相关工作

在医学图像分割中，常用的方法是应用U-Net (Ronneberger et al .， 2015)或SegNet (Badrinarayanan et al .， 2017)来预测每个输入图像的分割掩码。这种方法已成功应用于许多不同的任务(Habijan等人，2019;Kumar等人，2019;肖等人，2020)。最先进的技术是由nnU-Nets提供的(Isensee et al, 2021)，其中为每个特定数据集自动选择最佳架构和超参数。

不确定性量化在深度学习研究中很受关注(Abdar等人，2021)，通常使用贝叶斯神经网络(Kendall等人，2015b;Mitros and Mac Namee, 2019;Gal and Ghahramani, 2016)。我们可以区分认知不确定性(epistemic uncertainty)和任意不确定性(aleatoric uncertainty)，前者指模型参数的不确定性，后者指数据的不确定性。如(Kendall and Gal, 2017)所述，分割模型的认知不确定性可以用蒙特卡罗Dropout近似，而任意不确定性可以用MAP推理建模。这些方法也应用于各种医疗任务(Wang et al .， 2019;Nair等，2020;DeVries

和Taylor, 2018)，包括脑肿瘤分割(Sagar, 2020;Jungo and Reyes, 2019;Mehta et al, 2020)。

在过去的一年中，ddpm由于其在图像生成方面的惊人表现而获得了很多关注(Dhariwal和Nichol, 2021)。图像由高斯噪声采样生成。该采样方案遵循随机过程，因此从相同的噪声图像中采样不会产生相同的输出图像。通过去噪扩散隐式模型(DDIM)引入了一种不同的采样方案(Song等人，2020)，其中采样是确定性的，可以跳过多个步骤来完成。

此外，图像之间可以实现有意义的插值。(Nichol and Dhariwal, 2021)和(Dhariwal and Nichol, 2021)进一步改进了DDPM，其中损失目标的改变、架构的改进和采样过程中的分类器引导提高了输出图像质量。

虽然一些新工作将扩散模型应用于图像到图像的翻译(Sasaki等人，2021)、风格迁移(Choi等人，2021)或绘画任务(撒哈拉等人，2021)等任务，但到目前为止，关于语义分割的工作很少。最近，(Baranchuk et al .， 2021)提出了一种使用扩散模型进行语义分割的方法。训练DDPM来重建需要分割的图像。然后，对模型的特征进行MLP分类，得到对原始图像的分割掩码。与此方法相反，我们直接训练DDPM来生成分割掩码。

同时，独立于我们，(Amit et al, 2021)开发了一种类似于我们的图像分割方法。但是，它们使用单独的编码器进行图像和分割。对于医学图像分析来说，训练一个更大的模型可能会很困难，因为可能会有大的输入图像，比如3D数据。我们的方法只使用一个编码器对图像信息和分割掩码进行编码。

2. 方法

目标是训练一个DDPM来生成分段掩码。我们遵循(Nichol和Dhariwal, 2021)中提出的想法和实施。扩散模型的核心思想是，对于许多时间步长T，将噪声添加到图像x中。这将产生一系列噪声图像x0, x1，…， xT，其中噪声电平从0(无噪声)稳步增加到T(最大噪声)。该模型遵循U-Net的架构，并从xt预测任意步t∈{1，…， t}。在训练过程中，我们知道xt−1的基本真值，并且用MSE损失训练模型。在采样过程中，我们从噪声xT ~ N (0, I)开始，采样T步，直到得到假图像x0。

以下公式的完整推导可参见(Ho et al, 2020;Nichol and Dhariwal, 2021)。扩散模型的主要组成部分是前向去噪过程q和反向去噪过程p。根据(Ho et al .， 2020)，给定图像x在步长t处的前向去噪过程q为

其中I表示单位矩阵，β1，…， βT为正向过程方差。其思想是，在每一步中，少量的高斯噪声被添加到图像中。做对于t步，我们可以这样写

利用重参数化技巧，我们可以直接把xt写成x0的函数：

逆过程pθ由模型参数θ学习得到，由式给出

如(Ho et al .， 2020)所示，我们可以用xt预测xt−1

其中σt表示模型可以学习的方差方案，如(Nichol and Dhariwal, 2021)中提出的。从方程5中我们可以看到，抽样有一个随机分量z，这导致了一个随机抽样过程。注意∊θ是我们训练的U-Net，输入xt =√αtx0 +√1 - αt?在根据式5的采样过程中，将从xt中减去的噪声格式∊θ(xt, t)必须由模型学习。这个U-Net是根据(Nichol和Dhariwal, 2021)中给出的损失目标进行训练的。

我们现在修改这个想法，使用扩散模型进行语义分割。图2给出了脑肿瘤分割任务的可视化工作流程。

设b为给定尺寸为(c, h, w)的脑MR图像，其中c表示通道数，(h, w)表示图像高度和图像宽度。输入图像b的肿瘤的ground truth分割记为xb，其维数为(1,h, w)，训练DDPM生成分段掩码。在经典的DDPM方法中，xb将是我们训练所需的唯一输入，当我们在推理期间从噪声中采样时，这将导致任意分割掩码x0。与此相反，我们提出的方法的目标不是生成任何分割掩码，而是对给定图像b想要一个有意义的分割掩码xb,0。为了实现这一目标，我们在输入中添加了额外的通道:我们通过将b添加为xb之前的图像来诱导b中存在的解剖信息。我们通过连接b和xb来实现这一点，并定义X:= b⊕xb。因此，X的维数为(c + 1, h, w)。

在降噪过程q中，我们只对地面真值分割xb加入噪声:

定义Xt:= b⊕xb,t。然后将式5修改为

得到维数为(1,h, w)的稍微去噪的xb,t−1。在推理过程中，我们遵循算法1中给出的过程，这是一个随机过程。因此，对相同的脑MR图像b进行两次采样，不会得到相同的分割掩码预测xb,0。利用这一特性，我们可以隐式地生成分割掩码集合，而无需训练新模型。这种集成可以用来提高分割性能。

3. 数据集和训练细节

我们在BRATS2020数据集上评估了我们的方法。如第1节所述，为每个患者提供四种不同MR序列的图像，这些图像被堆叠到4个通道。

我们将三维磁共振扫描结果轴向切片。由于肿瘤很少发生在大脑的上部或下部，我们排除了最低的80片和最高的26片。对于强度归一化，我们将像素强度的顶部和底部分别切割一个百分位数。我们将图像裁剪为(4,224,224)的大小。所提供的基础真值标签包含四类，分别是背景、gd增强肿瘤、肿瘤周围水肿以及坏死和非增强肿瘤核心。我们将三种不同的肿瘤类别合并为一类，因此将分割问题定义为逐像素的二值分类。我们的培训

集合包括来自332名患者的16,298张图像，测试集包括来自37名患者的非空地面真值分割的1,082张图像。没有应用数据增强。

我们的DDPM模型的超参数在(Nichol和Dhariwal, 2021)的附录中进行了描述。我们选择T = 1000步的线性噪声调度。该模型使用混合损失目标进行训练，Adam优化器的学习率为10−4，批大小为10。第一层的通道数选择为128，我们使用一个分辨率为16的注意头。我们在NVIDIA Quadro RTX 6000 GPU上训练了60,000次迭代的模型，大约需要一天的时间。对比方法的培训细节见附录B。

4. 结果与讨论

对于测试集的每个图像，我们采样5个不同的分割掩码。这通过对5个掩模进行平均并将其阈值设为0.5来隐式地定义一个集合。我们在表1的第二行报告了这个集合的结果。我们已经看到5个集合增加了我们的方法的性能。

在表1的最后一列中，我们计算了模型产生空分割掩码的情况。这导致骰子为零，HD95无法计算。如果我们忽略这些情况，我们报告HD95分数，平均骰子分数和Jaccard指数在表1的方括号中报告。

作为基线，我们报告了nnU-Net和SegNet的分割分数。默认情况下，nnU-Net是一个5重交叉验证的集合。我们还像(Kendall et al .， 2015b)中提出的那样，使用蒙特卡罗dropout实现贝叶斯隔离网。通过在推理过程中采样五次，我们可以再次合成生成的分割掩码。

表1的最后一行报告了这个集合的分数。

用我们的方法生成一个样本需要48秒，而用SegNet计算分割掩码需要13毫秒。为了加快采样过程，我们将在未来的工作中考虑使用DDIM方法进行采样。

为了使不确定性图可视化，我们从测试集中选择了三个示例图像b1、b2和b3。更多的例子在附录c中给出

图3:100次抽样运行生成的均值和方差图示例。

不确定性映射，我们为每个图像采样100个分割掩码，并计算逐像素方差。在图3中，我们给出了原始脑MR图像b的一个通道，地面真值分割，两个不同的采样分割掩码，以及均值和方差图。我们可以清楚地识别出模型不确定的区域。此外，通过将均值映射的阈值设定为0.5，我们可以生成集成的分割掩码。在表2中，我们报告了这个集合掩码的分割分数和，以及100个样本的平均分数。我们看到集成可以提高示例b1, b2和b3的性能。