LSegDiff: A Latent Diffusion Model for Medical ImageSegmentation

LSegDiff:用于医学图像分割的潜在扩散模型

摘要

扩散模型最初是为图像生成而设计的，它也可以有效地应用于各种任务，包括语义分割。然而，现有的基于扩散的语义分割方法大多在高维像素空间中运行，在训练和推理过程中需要大量的计算和内存资源。本文首次尝试利用潜在扩散模型进行语义分割。具体而言，我们提出了一种快速有效的潜在扩散模型，并对其在医学图像分割任务中的应用进行了评价。首先，我们训练了一个变分自编码器(vae)网络，将二值图像掩码转换为紧凑的潜在向量。扩散过程可以在这个低维潜在空间中执行，从而大大加速。随后，我们使用V AE解码器从扩散过程产生的潜在输出向量重建精确的预测图。最后，我们通过使用形态学操作的简单后处理步骤来细化最终的分割结果。我们报告了两个公共数据集的结果，包括结肠息肉图像和皮肤癌图像。实验表明，与传统扩散模型相比，我们的方法在具有更好的训练和推理速度以及更有效的内存消耗的同时，获得了具有竞争力的准确性。

1 介绍

医学图像分割是一个专注于精确描绘医学图像中特定结构的专业领域。

特别是，我们的工作集中在两个关键领域，包括息肉和皮肤病变的分割。两特定域分割是指在医学图像中识别和描绘异常区域，如息肉和皮肤病变的过程，特别是在内窥镜和皮肤镜等医学成像背景下。息肉是一种异常生长的组织，可以出现在身体的各个部位，如结肠、胃或呼吸道。这个过程在医疗诊断和治疗计划中至关重要，因为准确地识别和测量息肉可以帮助医生发现疾病的早期迹象，如结肠直肠癌，并对患者护理做出明智的决定。此外，皮肤病变是指皮肤表面的异常和病理区域，包括痣、溃疡、皮疹和皮肤癌等多种情况。在本研究中，医学图像分割算法通常利用先进的图像处理技术、机器学习和深度学习方法来分析医学图像，准确定位内窥镜图像中息肉的边界和皮肤镜图像中病变皮肤的区域。

解决这一挑战的传统机器学习方法主要依赖于手工制作的特征来提取图像属性，如颜色、形状和纹理。然而，由于病变类型的多样性，这些方法在准确表示和检测息肉时经常遇到困难。近年来，深度神经网络，特别是UNet[19]架构，已经成为息肉分割的主要方法。这些网络可以捕获复杂和抽象的特征，从而产生令人印象深刻的性能结果。然而，深度神经网络的采用也引入了复杂性方面的权衡，因为模型可能非常大(包含数百万个参数)，并可能导致显著的推理延迟。

去噪扩散模型在各种生成任务中取得了显著的成就，包括将医学图像细分为有意义的组件。MedSegDiff[23]利用去噪扩散模型对二维医学图像进行分割。这是通过利用名为

去噪-UNet，它将去噪原理与UNet神经网络的框架相结合。此外，利用傅里叶变换结合了结构间的洞察。这种技术组合使MedSegDiff能够有效地隔离和突出医学图像中感兴趣的区域，从而获得准确和详细的2D分割结果。另一项值得注意的研究[22]采用扩散模型来解决分割二维医学图像的挑战。在他们的方法中，采用扩散模型逐步细化分割结果。具体来说，每个扩散步骤的输出结果在测试阶段通过求和进行汇总。这种聚合策略增强了分割结果的弹性，提高了最终分割区域的一致性和可靠性。

虽然去噪扩散模型可以获得准确的分割结果，但它们在推理过程中使用的迭代细化过程可能导致与计算复杂性、内存使用以及准确性和速度之间的权衡相关的挑战。

稳定扩散[18]技术可以通过加速迭代细化过程的收敛来解决扩散图像模型去噪时推理缓慢的局限性。

这可以在不牺牲分割或去噪结果质量的情况下实现更快、更有效的推理。隐空间中去噪扩散模型的一般概念包括在应用迭代去噪过程之前将带噪声的输入图像转换为低维隐空间表示。

该方法结合了降维、降噪和稳定收敛的优点，实现了准确高效的去噪结果。

提出了一种利用潜在扩散模型进行语义分割的新方法。具体而言，我们提出了一种快速有效的潜在扩散模型LSegDiff，并评估了其在医学图像分割任务中的性能。

LSegDiff基于三个阶段:将图像从像素空间转换到低维潜在空间，在潜在空间上训练条件扩散模型，将扩散模型的输出从潜在空间转换回原始像素空间。具体而言，我们的主要贡献如下:

-据我们所知，我们提出的方法LSegDiff是第一次尝试使用潜在扩散模型进行语义分割。

我们训练了一个V AE模型，将像素级二值图像掩模映射到潜在空间向量。这个预训练的V AE允许我们在低维空间中执行扩散过程。

-在几个基准医学图像分割数据集上的实验表明，与传统的像素空间扩散模型相比，我们的方法在速度和计算资源方面具有优越的性能，同时保持相同的精度水平。

本文的其余部分组织如下。在第二节中，我们回顾了相关工作。第3节描述了我们提出的模型。

实验结果见第4节。最后，我们对本文进行了总结，并在第5节中提出了未来的工作。

2 相关工作

2.1 图像语义分割

图像语义分割是计算机视觉领域的一项重要任务，具有广泛的实际应用。分割模型应用于医学图像诊断、自动驾驶汽车、场景理解等多个任务。已经提出了许多不同的方法来解决这个任务。这些方法可以根据不同的网络架构分为四类:基于cnn的架构、基于变压器的架构、基于mlp的架构和其他。针对该任务提出的最著名的UNet模型[19]显著提高了性能。另一种利用变压器架构[24]或使用卷积核变体[4]的方法也可以在语义分割数据集上获得SoTA结果。一些基于扩散的架构，如SegDiff[1]、MedSegDiff[23]和DiffUNet[25]，与传统方法相比，也取得了显著的效果。

对于特定的任务，如医学语义分割，一些基于扩散的架构，如MedSegDiff[23]或DiffUNet[25]表现非常好，并在医疗数据集上获得了显着的结果。

2.2 条件扩散模型

与gan[8]和基于流的模型[10]等其他深度图像生成模型类似，条件扩散模型提出了一种控制去噪过程输出的方法。条件模块取决于所添加条件的类型。对于text-to-image，文本条件可以通过Imagen[16]中的T5或DallE[17]中的Transformers等大型预训练语言模型嵌入。在image-to-image模型中，图像条件由Stable Diffusion中的CLIP等视觉语言模型嵌入[15]。

2.3 稳定扩散

长时间的训练和推理是扩散模型最大的缺点。LDM[18]提出了一种在潜在空间而不是像素空间执行扩散过程的方法，因为它的训练成本更低，推理速度更快。这种类型的模型被称为潜在扩散模型。通过利用变分自动编码器(Variational AutoEncoder, V AEs)[12]将图像数据压缩到被称为“潜空间”的隐藏空间中，扩散过程在这个小维空间中运行，从而加快了训练过程。

2.4 去噪扩散隐式模型

 去噪扩散概率模型(Denoising Diffusion Probabilistic Model, DDPM)[9]提出了一种用扩散模型进行训练和推断的方法，但该方法在训练和推理速度上存在缺点，因为模拟一个反向扩散过程的马尔可夫链需要花费很多时间步，可能高达一步或几千步。此后不久，论文《去噪扩散impit Model (Denoising Diffusion impit Model, DDIM)》[20]提出了一种新的采样方法，通过使去噪过程非马尔可夫化，从而减少去噪时的时间步数，使去噪过程可以跳过步骤。使用ddim，可以按照DDPM方法训练扩散模型，使用任意数量的前向步骤，但在生成过程中仅使用几个步骤进行降噪。

                                                        图1:我们提议的LSegDiff的整体架构。

3  方法

3.1 总体架构

典型的医疗分割数据集由一对输入图像和掩码图像组成，表示为 𝐷 = (𝑦𝑖, 𝑥𝑖 )𝑁 𝑖 。RGB 模式的图像 𝑦𝑖 和二进制模式的掩码图像 𝑥𝑖。这些样本来自未知条件分布𝑝 (𝑥|𝑦)，这意味着一对一的映射，即每个输入图像只对应一个掩码图像。我们的目标是通过一个参数模型来逼近𝑝 (𝑥|𝑦)，该模型经历了一个随机迭代改进过程，将输入图像𝑦映射到掩码𝑥上。潜在分割扩散模型 LSegDiff 的整体架构如图 1 所示。1.我们的架构由三个主要模块组成：

掩码压缩模块：利用变异自动编码器架构，该模块包括两个部分：编码器和解码器。编码器将地面实况掩码𝑥𝑖 压缩为潜在表示𝑧𝑖，而解码器则将潜在向量𝑧𝑖 变回掩码图像𝑥𝑖。在训练整体架构之前，先用大量地面实况掩码对该模块进行预训练。在训练过程中，我们使用编码器𝜉 来预先计算潜在𝑧，而解码器𝐷 在与去噪 UNet 模块结合训练时则完全冻结。

条件图像嵌入器。与其他条件扩散模型一样，我们使用一个简单的卷积网络来提取输入图像中的特征图向量信息，然后将这些向量与去噪 UNet 模块中早期层的输出连接起来。这确保了生成过程的每一步都受到输入图像的限制。此外，该模块与去噪 UNet 模块进行端到端训练。

去噪 UNet 模块。如前所述，我们的方法旨在通过参数模型逼近条件分布𝑝（𝑦|𝑥）。从𝑥预先计算出潜在向量𝑧0后，我们使用 UNet 架构来学习反向扩散过程，其分布为𝑝 (𝑦|𝑧), 𝑧 = 𝜉(𝑥)。在反向扩散过程的每个时间步𝑡，UNet 会预测输入 ˜𝑧𝑡的噪声𝜖，然后输出一个噪声较小的向量 ˜𝑧𝑡-1 = ˜𝑧𝑡 - 𝜖。最后，我们计算 𝑧0 和 ˜𝑧0 之间的损耗，𝑥 和 ˜𝑥 = 𝐷( ˜𝑧0)。

3.2  将掩码图像压缩到潜空间

我们的掩码压缩模型基于之前的研究成果[18]，由感知损失组合训练的自动编码器组成。这确保了重构局限于图像流形，强化了局部真实感，避免了单纯依赖像素空间损失（如 L2 或 L1 目标）带来的模糊。更确切地说，给定二进制掩码图像 𝑥∈ 𝑅 𝐻×𝐻×𝑅，编码器 𝜉 将 x 编码为潜在表示 𝑧 = 𝜉(𝑥)、解码器 𝐷 从潜在表示重建图像，得到 𝑥 = 𝐷(𝑧) = 𝐷(𝜉(𝑥))，其中 𝑧 ∈ 𝑅 类型×𝑤×𝑐。重要的是，编码器对图像的下采样系数𝑓 = 𝐻/ h= W/𝑤，我们研究了不同的下采样系数𝑓 = 2𝑚，其中𝑚∈ 𝑁。我们决定选择变异自动编码器（VAEModel）[12] 进行压缩。根据文献[6]，自编码器模型是以对抗的方式训练的，因此判别器𝐷𝜓是最优化的，以区分原始图像和重建图像𝐷(𝜉(𝑥))。具体来说，自动编码器模型的目标函数为：

3.3 前向扩散过程

假设我们有一个来自分布 𝐷(𝑧)的样本 𝑧0。使用前向马尔可夫扩散过程 q，在𝑇迭代过程中逐渐向潜在向量𝑧0 中添加高斯噪声，不同的水平𝑡∈{0，1，...，𝑇} 如下：

其中，𝑎1:𝑇 是标量超参数，约束条件为 0 < 𝛼𝑡 < 1，决定了每次迭代时添加的噪声方差。请注意，𝑧𝑡-1 被 √ 𝛼𝑡 削弱，以确保随机变量的方差在𝑡 →∞ 时保持受限。例如，如果𝑧𝑡-1 的方差为 1，那么𝑧𝑡 的方差也为 1。

重要的是，将中间步骤边际化后，给定𝑧0 的𝑧𝑡 的分布可以描述如下：

这样就可以得出给定（𝑧0, 𝑧𝑡 ）的𝑧𝑡-1 的后验分布，这对反向链的参数化和反向链的对数概率的变分下限都很有用。

3.4  优化去噪模型

我们利用条件图像 𝑦 形式的额外信息来帮助逆转扩散过程。因此，利用条件图像 y 和噪声目标潜掩码 𝑧𝑡 优化神经去噪模型 f 非常重要。下式中的噪声目标潜像掩码𝑧𝑡 与前向扩散过程不同步骤中噪声潜像掩码的边际分布相一致：

去噪模型 𝑓𝜃 (𝑦, 𝑧𝑡, 𝛾)通过输入源图像 𝑦、噪声目标图像 𝑧𝑡 和噪声方差的充分统计量 𝛾来增强，然后训练它来预测噪声向量 𝜖。通过对标量 𝛾 的调节，去噪模型可以了解噪声的程度。训练 𝑓𝜃 的目标函数为：

其中，𝜖 ∼ N(0,𝐼)，𝛾 ∼ 𝑝 (𝛾)。𝛾的分布对模型和生成输出结果的质量有显著影响。

另一种方法是将𝑓𝜃的输出回归到𝜎，而不是将𝑓𝜃的输出回归到𝑦0。根据𝛾和˜𝑧的值，可以确定𝜎和𝑧0 的值。然而，改变回归目标会影响到 𝜎 和 𝑧0 的规模。

损失函数。如果对𝑝 (𝛾) 进行调整以考虑损失函数的规模，我们预计这两种变体都会有相当好的效果。

3.5  反向扩散过程

在 DDPM [9]中，推理是通过反向马尔可夫过程进行的，其运行方向与正向扩散过程相反。该过程从高斯噪声𝑧𝑇开始，条件图像𝑦以相反的顺序进行：

推理过程使用各向同性的高斯条件分布𝑝𝜃 (𝑧𝑡-1|𝑧𝑡, 𝑦)，这些分布是从 DDPM 中学习的。通过设置正向过程步骤的噪声方差尽可能小（例如，𝛼1:𝑇 ≈ 1），最佳反向过程 𝑝 (𝑧𝑡-1|𝑧𝑡, 𝑦)将近似高斯。因此，在推理过程中选择高斯条件可以合理地近似于真实的反向过程。值得注意的是，1 - 𝛾𝑇 应该足够大，以确保 𝑧𝑇 遵循近似先验分布，即 𝑝 (𝑧𝑇) = N(𝑧𝑇; 0, 𝑰)。这使得采样过程可以从纯高斯噪声开始。我们注意到，𝑓𝜃 是经过训练的去噪模式，可以在任何噪声图像𝑧𝑡（包括𝑧𝑡）的情况下估计𝜖。因此，在给定𝑧𝑡 的情况下，我们可以通过重排近似地计算𝑧0：

将估计值 ˆ𝑧0 代入𝑧𝑡-1|𝑧0，𝑧𝑡 的后验分布，即可将𝑝𝜃 (𝑧𝑡-1|𝑧𝑡, 𝑦) 的均值参数化为：

并将𝑝𝜃 (𝑧𝑡-1|𝑧𝑡，𝑦) 的方差设为 (1-𝛼𝑡 )，默认值为正向过程的方差。根据这一参数化，DDPM 模型下反向推理过程的每次迭代形式如下：

我们将从变分下界的角度来论证概率模型训练目标的选择。如果把前向扩散过程看成是推理过程的一个固定近似后验，我们就可以得出以下边际对数似然的变分下界：

3.6  后期处理

在进行实验时，我们发现所提出的 LSegDiff 的分割预测图经常包含一些小洞。为了解决这个问题，我们对模型的输出进行了后处理，利用直接的闭合形态学操作来填补这些小洞。该技术先将图像扩张，然后侵蚀扩张后的图像，并通过使用核滑过在整个图像中反复使用白色像素值填充空白区域。

4  实验和讨论

在本文中，我们实验的主要重点是在表 1 所示的六个数据集上，将所提出的方法与同样利用扩散模型及其变体的类似方法进行比较：

ISIC 2018 数据集[5]：由国际皮肤成像协作组（ISIC）发布，是一个大规模的皮肤镜图像数据集。该任务 1 数据集是对病变分割的挑战，包含 2594 张图像。

Kvasir-SEG [11] ：包含从内窥镜检查过程中获取的图像，包括由医疗专业人员标注和验证的图像。该数据集包含三个解剖位置（食道、胃和结肠）、三种疾病类别（食道炎、息肉、溃疡性结肠炎）和两个与息肉切除相关的附加类别（染色和切除息肉）的图像。总体而言，数据集包含 8,000 张内窥镜图像，每个类别有 1,000 个示例。

CVC-ClinicDB[2]：是一个开放存取的数据集，包括 31 个内窥镜序列中分辨率为 384×288 的 612 幅图像。该数据集用于医学图像分割，特别是内窥镜视频中的息肉检测。

CVC-300 ：是 CVC-ClinicDB 数据集的一部分。CVC-300 包含 60 幅专门用于模型评估的图像。

CVC-ColonDB ：包含来自 15 个不同内窥镜手术的 379 帧图像，每个手术至少包含一张含有息肉的图像。

ETIS-LaribPolypDB [21] 包括 196 幅分辨率为 1225 × 966 的静态图像，包含从 34 种不同内窥镜手术中获取的 44 种类型的息肉。

我们的实验主要是将所提出的方法与在这些数据集上使用扩散模型的其他类似方法进行比较。目的是在息肉分割的背景下，评估所提出的方法与现有技术相比的性能和有效性。

4.2 实验设置

我们使用 PyTorch 在计算机上实现了我们的模型，计算机的规格如下：AMD Ryzen 3970X 3.7GHz CPU、128GB 内存、NVIDIA GeForce RTX 3090 GPU，操作系统为 Ubuntu 20.04 64 位。训练过程分为三个阶段。首先，使用分割的遮罩图像数据训练遮罩压缩模型。该数据集由 52,000 张图像组成，其中包括 2,000 张真实图像和 50,000 张由 GAN [8] 模型生成的假图像。大小为 1x384x384 的二进制掩码𝑥被压缩成尺寸为 8 × 24 × 24 × 24 的潜向量。实验使用的批量大小为 8，亚当进行了优化，学习率为 5 × 10-5。其次，我们使用条件图像嵌入模型将 3 × 384 × 384 图像分别转换为尺寸为 32 × 24 × 24、64 × 12 × 12 和 128 × 6 × 6 的向量。然后将这些向量与 UNet 网络编码器模块中每个卷积层之后的潜向量进行串联，用于反向扩散过程。正向扩散过程使用从第一阶段获得的潜向量。计算潜向量与反向扩散过程后获得的向量之间的均方误差 (MSE) 损失。最后，我们使用从第一阶段预训练的 VAE 掩膜压缩模型中获得的解码器对输出向量𝑧0 进行解码。在获得预测分割图 ˆ𝑥 后，我们会计算预测图 ˆ𝑥 和地面实况 𝑥 之间的骰子损失和二元交叉熵（BCE）损失。一般损失值计算为三项损失之和：MSE 损失、Dice 损失和 BCE 损失之和。然后在整个网络中执行反向传播，但冻结的 VAE 解码器组件除外。

4.3 结果与讨论

我们从准确性、速度和计算资源三个方面来评估我们的方法。用于基准准确性的指标是IoU和Dice分数。在速度方面，我们评估了模型的每个历元的训练时间、收敛时间和推理时间。最后，在计算资源方面，我们评估了运行训练过程时模型可训练参数和GPU内存消耗。

从表2可以看出，LSegDiff的准确率略好于DiffUNet[25]的结果，但与其他架构相比仍然较低。可以解释为LSegDiff计算的是隐空间中的扩散过程，隐空间的维数较小，相对于图像空间中的输入，会导致信息泄漏，从而导致度量降低。

从图2可以看出，LSegDiff的结果趋向于平滑分割掩码的边界。除此之外，在预测图内的一些微小黑洞导致的结果比原来的DiffUNet[25]差。为了改善这一弱点，我们可以寻找另一种方法将图像转换为隐空间，将输入从图像空间转换为隐空间时的信息泄漏最小化。

图3说明了如何推断预测的分割掩码。

由于扩散过程是在潜在空间中执行的，因此我们使用预训练的V AE解码器将生成的潜在向量转换为像素空间并将结果可视化。

从表3可以看出，我们提出的模型LSegDiff在使用ISIC[5]数据集和Polyps数据集运行时，每个epoch的训练时间和总训练时间都更快，我们模型的推理速度也优于原始模型DiffUNet。

4.4  消融研究

冻结 VAE 解码器的效果。与 DiffUNet 和其他架构相比，LSegDiff 的局限性在于训练过程使用潜在向量来计算损失值，因此只有在训练过程中才会出现损失。

使用均方误差函数，而语义分割任务包含其他损失。我们可以通过将潜在向量转换到图像空间并计算Dice损失来结合MSE损失来改进损失计算。此外，我们实验中使用的息肉和皮肤癌数据集包含一个类，因此s形函数可以用于BCE损失。我们在ISIC数据集上运行这个实验，并比较不同设置下的LSegDiff。

表3和表4显示，在训练中结合三个损失项的LSegDiff略微提高了IoU和Dice分数，但是

由于V AE译码模块参数较多，训练时间和推理时间也增加。

潜在大小对模型性能的影响。在本实验中，我们考察了6 × 6、12 × 12、24 × 24等不同潜值对掩模图像压缩的影响。我们根据Dice和IoU分数对每个配置进行基准测试。从表5中我们观察到，8 × 24 × 24的潜在向量产生的效果最好。

5 结论

本文提出了一种利用潜在扩散模型来处理医学语义分割任务的开创性方法。我们的模型首先将分割掩模映射到潜在向量中，然后在这个潜在空间中执行扩散过程来加速去噪过程。最后，我们使用预训练的V AE解码器从扩散过程的输出中重建预测的分割结果。总的来说，我们的模型在不牺牲性能的情况下显著提高了训练和测试速度以及内存消耗。

在未来的工作中，我们希望探索快速傅里叶变换(FFT)[13]，将从条件图像嵌入器中提取的信息融合到频率空间中。我们还致力于在分割掩码中集成各种感知压缩模型技术，如VQ-V AE[14]和VQ-GAN[7]。