ACCELERATING DIFFUSION MODELS VIA PRE-SEGMENTATION DIFFUSIONSAMPLING FOR MEDICAL IMAGE SEGMENTATION

基于预分割扩散采样的加速扩散模型医学图像分割

我觉得这个创新点很好，就是没有代码555555

摘要

基于去噪扩散概率模型(DDPM)，医学图像分割可以被描述为一个条件图像生成任务，该任务允许计算分割的逐像素不确定性映射，并允许隐式集成分割以提高分割性能。然而，DDPM需要许多迭代去噪步骤来从高斯噪声中生成分割，导致推理效率极低。为了缓解这一问题，我们提出了一种有原则的加速策略，称为预分割扩散采样DDPM (PD-DDPM)，该策略专门用于医学图像分割。其关键思想是基于单独训练的分割网络获得预分割结果，并根据前向扩散规则构造噪声预测(非高斯分布)。然后，我们可以从噪声预测开始，使用更少的反向步骤来生成分割结果。

实验表明，PD-DDPM在显著减少反向步数的情况下，比代表性基线方法的分割效果更好。此外，PD-DDPM与现有的高级分割模型是正交的，可以结合起来进一步提高分割性能。

1 介绍

去噪扩散概率模型(Denoising Diffusion Probabilistic Models, DDPM)[1,2]是深度生成模型的一个类别，由于在无条件和条件生成任务中都有很好的效果，近年来成为计算机视觉领域最热门的话题之一[2 - 4]。基于去噪扩散概率模型(DDPM)，医学图像分割可以被描述为一个条件图像生成任务，该任务允许计算逐像素的不确定性映射，该分割方法允许隐式的分割集合来提高分割性能[5]。特别是在后续诊断或治疗依赖于分割的医疗应用中，仅提供最可能假设的算法可能导致误诊和次优治疗。如果提供了多个一致的假设，它们可以用来建议进一步的诊断测试来解决歧义，或者获得额外信息的专家可以为后续步骤选择适当的假设[5,6]。

消噪扩散概率模型(DDPM)由两条马尔可夫链组成。在正向扩散过程中，干净图像逐渐受到高斯噪声的干扰，直至近似于高斯分布[1]。

在逆扩散过程中，从采样的高斯噪声中，利用训练好的去噪深度神经网络对数据进行迭代去噪，得到干净的图像。

因此，通过对采样的高斯噪声进行迭代去噪，可以实现对DDPM的合成。对于医学图像分割，我们可以在地面真值分割的基础上训练DDPM，并在训练和采样过程的每一步中使用图像作为先验[5]。

然而，普通DDPM的一个主要问题是推理效率低下。因为从DDPM中获得干净的分割通常需要数百甚至数千个去噪步骤，每个步骤都涉及到去噪神经网络的前向预测。

为了加速ddpm，提出了几种方法。

Song等[7]试图通过使用非马尔可夫逆过程来减少扩散步骤的数量。与DDPM将噪声尺度定义为常数不同的是，San-Roman等[8]提出了一种自适应噪声调度方法，用于估计在推理时给定当前输入的噪声参数，所需的步骤更少。在[9]中，作者提出将整个采样过程提炼成一个更快的采样器，只需要一半的步骤。一些方法[10,11]使用预训练的自编码器将扩散过程转移到潜在空间。然而，上述方法不能在不牺牲质量的情况下实现显著的加速

发电的。其他一些方法[12,13]通过截断正向和反向扩散过程来提高采样效率，同时提高了性能。

但该方法需要结合GAN[14]或VAE[15]等难以训练的模型。打破了传统DDPM在扩散和采样过程中图像维数保持不变的特点。综上所述，以上所有方法都没有实现专门针对分割任务的加速采样。

为了缓解这一问题，我们提出了一种有原则的加速策略，称为预分割扩散采样DDPM (PD-DDPM)，该策略专门用于医学图像分割。其关键思想是基于单独训练的分割网络获得预分割结果，并根据前向扩散规则构造噪声预测(非高斯分布)。然后，我们可以从噪声预测开始，使用更少的反向步骤来生成分割结果。PD-DDPM不仅在不打破任何假设的情况下提高了香草DDPM的效率，而且作为一个额外的好处，提高了香草DDPM的分割性能。实验表明，PD-DDPM在显著减少反向步数的情况下，比代表性基线方法的分割效果更好。此外，PD-DDPM与现有的高级分割模型是正交的，可以结合起来进一步提高分割性能。由于预分割和扩散操作都可以在一步中轻松实现，因此预训练的分割网络只会给PD-DDPM带来很小的计算开销。

2 方法

2.1. ddpm的背景

在DDPM[2]中，正向扩散过程是一条一阶马尔可夫链，通过在时刻t逐渐加入高斯噪声来扰动数据分布q(x0)，方差βt∈(0,1)，直到数据分布收敛到标准高斯分布。形式的前进过程可以概括如下:

式中x1:T为变量集合x1, x2，…， xT。T=1000∽4000是大多数工作的典型选择。以小扩散速率为极限(即βt保持足够小)，反向分布q(xt−1|xt)也服从高斯分布。因此，逆向过程可以用神经网络参数化高斯分布pθ来近似，从p(xT) = N (xT;0,我):

训练神经网络来模拟方程1中定义的扩散过程的反向过程。为了从相反的过程中生成图像，我们首先通过从p(xT)(选择为各向同性高斯分布)中采样潜点(与训练数据点x0大小相同)从底层数据分布中采样xT，然后依次从pθ(xT−1| xT)中提取样本xT−1,t = t, t−1，…， 1，直到得到新的数据x0。DDPM的生成过程非常缓慢，因为它需要从转移分布pθ(xt−1|xt)中迭代采样，这涉及到对神经网络输出的大量评估。

2.2. 预分割扩散采样去噪扩散概率模型

从香草DDPM中提出的一个自然问题可能是:我们可以将反向过程削减到T ' (< T)步骤吗?如果我们能提前得到噪声样本xT '，就可以实现。

为了生成特定于图像的分割，我们在ground truth上训练DDPM，并在训练和抽样期间使用图像作为先验和抽样。为了生成特定于图像的分割结果，我们在地面真值分割上训练DDPM，并在训练和采样过程中使用医学图像作为条件信息。因此，我们可以利用医学图像条件信息训练一个独立的分割网络。

在采样过程中，我们首先可以通过图1所示的预分割网络fψ得到预分割结果。然后，根据公式1，将预分割结果扩散到T '步，得到xT '的近似样本xT '。在这种情况下，反向过程可以重新定义为:

图1:PD-DDPM与香草DDPM的比较。该方法的训练和抽样程序。在每一步t中，通过将医学图像I连接到带噪分割掩码xI,t来诱导条件信息。

因此，我们可以使用去噪神经网络将非高斯分布xT '去噪到一个干净的分割x0，比根据公式4的香草采样过程的步骤更少。由于预分割结果不是真实的ground truth分割，因此在xT '和xT '之间会有一些误差。然而，我们通过实验实证发现，PD-DDPM提高了香草DDPM的分割精度。我们将预分割扩散采样DDPM命名为PD-DDPM。需要强调的是，PD-DDPM并没有打破传统DDPM去噪过程的高斯假设。这样大大减少了所需的去噪步骤，从而有效缓解了DDPM推理成本高的痛点。

3. 实验和结果

3.1. 数据集

我们在MICCAI 2017中白质高强度分割挑战提供的WMH数据集上评估了我们的方法[16]。它包括60例来自三个不同机构/扫描仪的脑MRI图像(3D t1加权图像和2D多层FLAIR图像)，并手工注释白质高强度(二值掩膜)。并在FLAIR图像上定义了手动参考标准。

因此，通过对3D T1加权图像进行重采样以与FLAIR匹配，生成二维多层T1图像。在本文中，所有的案例被随机分成五组。然后我们随机地将这五倍分配到训练集(3倍)、验证集(1倍)和测试集(1倍)中。

3.2. 实现细节

在本文中，所有的网络使用Pytorch训练，使用NVIDIA TESLA V-100 (Pascal) gpu，内存为32gb。

图2:分割图和不确定性图的可视化。

表1:不同方法的分割得分。

我们用Adam优化器优化了所有配置，学习率为1e-4，权重衰减为1e-5。批量大小设置为12。对于WMH，图像和注释标签被随机裁剪为128 × 192块。我们选择T = 1000步的余弦噪声调度。选择第一层去噪网络的通道数为128个，使用一个分辨率为16的注意头。在PD-DDPM中，我们选择AttUnet[18]作为预分割模型。

3.3. 分割性能比较

DDPM的推理过程是一个随机过程。因此，我们可以隐式集成分割掩码来提高分割性能。对于测试集的每个图像，我们采样5个不同的分割掩码。然后对5个不同的分割掩码进行平均组合，阈值设为0.5，得到二值分割。表1给出了Dice得分、Jaccard指数、95% Hausdorff距离(HD95)和F1。

我们进行定量实验，将我们的方法与一系列有代表性的方法进行比较。在这里,

U-Net[17]、AttUnet[18]和U-Net++[19]是分割领域最具代表性的深度学习模型，但它们都不能估计分割的不确定性。Bayesian U-Net[20]和Probabilistic U-Net[21]是估计分割不确定性的代表性方法。我们还将PD-DDPM与其他加速ddpm进行了比较，包括TDPM[12]和ESDDPM[13]。需要强调的是，对比方法中集合的大小也被设置为5。

表1显示PD-DDPM(当T ' =300时)在所有四个指标上都取得了最佳结果。PD-DDPM优于普通DDPM。为了可视化分割和不确定性映射，我们从图2中的测试集中选择了三个图像I1、I2和I3。

3.4. 确定最优T '

本文分析了超参数T′对PD-DDPM的影响。通过在{50,100,200,300,400,500,600,700,800,1000}中改变T '，我们训练PD-DDPM用于WMH分割。如图3所示，当T ' = 300时，PD-DDPM的Dice得分最佳。我们还分析了T '对不确定性估计(即预测softmax输出的变化)的影响。图3显示当T < 500时，不确定性随T '增加。然后T′进一步增大，不确定度趋于饱和。

3.5. 集合大小的影响

集成的最优大小是一个需要优化的任务特定参数[22]。图4显示了(1)具有多个掩码的集成优于只有一个掩码的集成。

(2)随着系综尺寸的增大，性能趋于饱和。在我们的方法中，我们将集合大小设置为5。图4还显示了不同集成尺寸下分割性能的标准差。随着集合大小的增加，分割性能的变化在骰子度量上减小。它证明了合奏该模型不仅提高了分割性能，而且保证了分割结果的鲁棒性.

3.6. 预分割精度的影响

由于预分割结果不是真实的ground truth分割，因此在xT '和xT '之间会有一些误差。本文分析了预分割精度对PD-DDPM的影响。如表2所示，预分割Dice得分越高，PD-DDPM的性能越高(当T ' =300时)。因此，PD-DDPM可以与现有的高级分割网络相结合，进一步提高性能并获得不确定性估计。

图3:测试集上的骰子和不确定性相对于T '。横轴表示T '的大小。

图4:骰子在测试集上相对于集合大小的平均值和标准差(当T ' =300时)。

表2:不同预分割Dice分数的PD-DDPM的分割分数。

4. 结论

为了加速DDPM在医学图像分割中的应用，本文提出了一种专门用于医学图像分割的预分割扩散采样DDPM (PD-DDPM)。实证发现，当T ' = 300 (T=1000)时，PD-DDPM在本文的参数设置下效果最佳。实验表明，即使反向采样步骤数量明显减少，PD-DDPM也优于香草DDPM。与现有的几种加速方法相比，该方法也得到了最好的结果。此外，PD-DDPM与现有的高级分割模型正交，可以将它们组合起来进一步提高模型性能并获得不确定性估计。