Diffusion【1】：SDSeg——基于Stable Diffusion的单步扩散分割！

文章目录

前言

和大家分享一下我们发表在 MICCAI 2024 （CCF-B，医学图像处理方面的顶会，最后得分 54） 上的论文：Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process。

项目主页：
Stable-Diffusion-Seg

欢迎大家在 arxiv 上阅读：
Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process

代码已经开源！！！期待您的 Star！！！
Stable Diffusion Segmentation (SDSeg)

Abstract

扩散模型已经在各种生成任务中证明了它们的有效性。然而，当应用于医学图像分割时，这些模型遇到了几个挑战，包括显著的资源和时间需求。它们还需要一个多步骤的逆向过程和多个样本来产生可靠的预测。为了应对这些挑战，我们介绍了第一个潜在扩散分割模型，名为 SDSeg，建立在稳定扩散（SD）之上。SDSeg 采用了一种简单的潜在估计策略，以促进单步逆向过程，并利用潜在融合连接去除了对多个样本的需求。广泛的实验表明，SDSeg 在五个具有不同成像方式的基准数据集上超越了现有的最先进方法。值得注意的是，SDSeg 能够通过单一的逆向步骤和样本生成稳定的预测，体现了模型以其名称所暗示的稳定性。

Introduction

图像分割是医学图像分析中的一项关键任务。为了减轻医学专业人员的工作负担，已经开发了许多自动化的医学图像分割算法。各种神经网络架构的有效性，如卷积神经网络（CNN）和视觉变换器（ViT），强调了深度学习是医学图像分割的一种有前途的方法。

研究人员对扩散概率模型（DPM）的近期兴趣导致了基于 DPM 的分割方法中对图像级扩散模型的关注。图像级扩散模型通过前向过程向图像引入噪声，并通过学习逐步解码噪声添加的逆过程来生成新图像。基于 DPM 的分割方法利用图像条件生成分割预测。然而，这些方法面临限制：(1) 在像素空间生成分割图是不必要的，可能导致优化效率低下和高计算成本，因为与普通图像相比，二元语义图具有稀疏的语义信息；(2) 扩散模型通常需要多个逆步骤来实现详细和多样的生成，以前的扩散分割模型需要多个样本来平均以获得稳定的预测。

为了克服这些挑战，我们提出了一个简单而高效的分割框架，称为 SDSeg，具有以下贡献：

Methods

SDSeg 的框架如图 1 所示。对于医学图像，我们引入了一个可训练的视觉编码器 τ θ \tau_\theta τθ​，将图像 C ∈ R H × W × 3 C\in \mathbb{R}^{H\times W\times 3} C∈RH×W×3 编码为其潜在表示 z c = τ θ ( C ) z_c=\tau_\theta(C) zc​=τθ​(C)。对于分割图，我们利用一个自编码器进行感知压缩。如图 1 所示，给定像素空间中的分割图 X ∈ R H × W × 3 X\in \mathbb{R}^{H\times W\times 3} X∈RH×W×3，编码器 E \mathcal{E} E 将 X X X 编码为一个潜在表示 z = E ( X ) z=\mathcal{E}(X) z=E(X)，解码器 D \mathcal{D} D 则从潜在表示中恢复分割图，得到重构 X ~ = D ( z ) = D ( E ( X ) ) \widetilde{X}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(X)) X =D(z)=D(E(X))，其中 z ∈ R h × w × c z\in \mathbb{R}^{h\times w\times c} z∈Rh×w×c。

实际上，我们注意到 SD 提供的自编码器对二值分割图表现良好，如图 2 所示。因此，在训练阶段我们保持自编码器冻结，这使得 SDSeg 成为一个端到端的方法。SDSeg 的扩散过程是在潜在空间中进行的。

Latent Estimation

在训练阶段，第一时间步的分割图潜在表示 z 0 z_0 z0​会添加 t t t个时间步的高斯噪声，得到 z t z_t zt​。扩散的前向过程可以表示为：

z t = α ˉ t z 0 + 1 − α ˉ t n z_t = \sqrt{\bar{\alpha}_t}z_0+\sqrt{1-\bar{\alpha}_t}{n} zt​=αˉt​ ​z0​+1−αˉt​ ​n

其中 n n n 是随机高斯噪声， α ˉ t \bar{\alpha}_t αˉt​ 是控制前向过程的超参数。每个训练步骤中，去噪 U-Net 的目标是估计随机高斯噪声 n n n 的分布，表示为 n ~ = f ( z t ; z c ) \tilde{n} = f(z_t;z_c) n~=f(zt​;zc​)，其中 f ( ⋅ ) f(\cdot) f(⋅) 表示去噪 U-Net。噪声预测损失可以表示为 L n o i s e = L ( n ~ , n ) \mathcal{L}_{noise}=\mathcal{L}(\tilde{n}, n) Lnoise​=L(n~,n)。

在旨在生成多样且语义丰富图像的任务中，逆向过程中逐步应用噪声估计可以逐步改进结果。然而，我们认为分割图过于简单，不需要通过复杂的逆向过程来实现显著的改进。相反，经过有效训练的去噪U-Net能够恢复包含所有必要结构和空间特征的潜在特征，适用于分割图。因此，在获得估计噪声 n ~ \tilde{n} n~ 后，我们可以通过简单的变换直接推导出相应的潜在估计：

z ~ 0 = 1 α ˉ t ( z t − 1 − α ˉ t n ~ ) \tilde{z}_0=\frac{1}{\sqrt{\bar{\alpha}_t}}(z_t-\sqrt{1-\bar{\alpha}_t}{\tilde{n}}) z~0​=αˉt​ ​1​(zt​−1−αˉt​ ​n~)

这种技术使得可以轻松添加一个监督分支，优化目标是最小化预测的 z ~ 0 \tilde{z}_0 z~0​ 与真实 z 0 z_0 z0​ 之间的差异，潜在损失函数定义为 L l a t e n t = L ( z ~ 0 , z 0 ) \mathcal{L}_{latent}=\mathcal{L}(\tilde{z}_0, z_0) Llatent​=L(z~0​,z0​)。因此，最终的损失函数可以表示为：

L = L n o i s e + λ L l a t e n t \mathcal{L}=\mathcal{L}_{noise}+\lambda\mathcal{L}_{latent} L=Lnoise​+λLlatent​

其中 λ \lambda λ 表示潜在损失函数的权重。在实践中， λ \lambda λ设为1， L n o i s e \mathcal{L}_{noise} Lnoise​ 和 L l a t e n t \mathcal{L}_{latent} Llatent​ 均为平均绝对误差。

值得注意的是，仅使用 L n o i s e \mathcal{L}_{noise} Lnoise​ 和多次DDIM抽样可以生成出色的分割结果。引入 L l a t e n t \mathcal{L}_{latent} Llatent​ 的最大贡献在于消除生成分割图的冗余逆向过程，从而显著提高推理阶段的速度。

Concatenate Latent Fusion

稳定扩散引入了交叉注意力机制以支持多模态训练和生成。然而，对于图像到图像的分割模型，优先考虑从图像中提取语义特征和结构信息至关重要，而多模态能力可能并不会带来额外的优势。此外，将交叉注意力嵌入到去噪 U-Net 的多个块中会增加额外的计算成本。因此，寻找一种更有效的潜在融合方法对于 SDSeg 变得至关重要。

此外，我们在图 2 中的观察显示，分割图的潜在表示与其对应的潜在表示具有显著的空间相关性，可能包含有助于分割任务的必要结构和特征信息。因此，受传统语义分割方法的启发，我们采用串联方法将分割图的潜在表示与图像切片的潜在表示合并。串联在诸如 U-Net 和 DeepLabV3+ 等模型中广泛应用，并被验证为一种有效的整合图像语义特征的策略。

Trainable Vision Encoder

在语义分割中，一个有效的视觉编码器能够从图像中提取必要的结构和语义特征，从而提升分割结果。作为一种基于图像条件的生成模型，SDSeg 采用可训练的视觉编码器来捕获图像间丰富的语义特征。

视觉编码器 τ θ \tau_\theta τθ​ 与编码器 E \mathcal{E} E 具有相同的架构，并使用其预训练权重进行初始化。虽然我们发现，简单使用在自然图像上预训练的冻结图像编码器可以带来显著的结果，但我们选择将视觉编码器设置为可训练的，使得 SDSeg 能够适应不同的医学图像数据集模态，增强其多样性和有效性。

Experiment

Datasets and Evaluation Metrics

为全面评估SDSeg的有效性和泛化能力，我们在两个 2D 分割任务的 RGB 数据集和两个 3D 分割任务的 CT 数据集上进行实验，如表 1 所示：

我们的评估包括三个主要方面：首先，使用 Dice 系数（DC）和交并比（IoU）指标跨数据集评估分割结果的准确性。其次，通过比较与其他基于扩散的分割方法的计算资源使用和推理速度来评估我们模型的效率。第三，使用 LPIPS、PSNR、SSIM 和M S-SSIM 评估我们生成的分割结果的稳定性，与其他扩散分割模型进行对比。此外，我们还进行了消融研究来验证我们提出的模块的有效性。

Implementation Details

Experimental Settings

SDSeg 在单个V100 GPU上训练，具有 16GB RAM。模型使用 AdamW 优化器进行训练，基础学习率为 1 × 1 0 − 5 1\times 10^{-5} 1×10−5，总共训练 100,000 个步骤。默认情况下，批大小设置为 4。我们使用 KL 正则化自编码器和下采样率 r = H h = W w = 8 r=\frac{H}{h}=\frac{W}{w}=8 r=hH​=wW​=8的LDM模型。SDSeg 将 RGB 图像作为像素空间输入，其中 H = W = 256 H=W=256 H=W=256，相应的潜在表示形状为 h = w = 32 h=w=32 h=w=32， c = 4 c=4 c=4。所有模型部分均使用稳定扩散提供的预训练权重进行初始化。用于串联输入的去噪 U-Net 的额外模型参数被初始化为零。

Inference Stage

在推理阶段，我们将随机生成的高斯噪声与医学图像的潜在表示进行串联。然后，去噪U-Net预测估计的噪声，使得 SDSeg 可以推导出潜在估计 z ~ 0 \tilde{z}_0 z~0​。接着，解码器 D \mathcal{D} D 将潜在估计转换为像素空间，得到最终的预测结果。如表 4 所示，SDSeg 无需外部采样器，只需进行单步逆向操作进行一次采样，即可实现稳定的预测。

Main Results

Comparison with State-of-the-Arts

我们在 REF、BTCV 和 STS 数据集上将我们的模型与几种语义分割方法进行了比较，如表 2 所示。此外，我们还在 CVC、KSEG 和 REF 数据集上将我们的模型与基于最新扩散技术的分割模型进行了比较，如表 3 所示。SDSeg 在这五个具有不同成像模态的数据集上表现优于所有其他模型，验证了其有效性和泛化能力。

Comparison of computing resource and time efficiency

表 4 展示了在BTCV数据集上对MedSegDiffs、Diff-U-Net 和 SDSeg 进行的效率评估结果。为了公平比较，这些模型都是在同一台服务器上使用它们的源代码进行训练的。结果突显了SDSeg在效率上的优势，其训练过程所需资源和时间显著减少。显著的是，SDSeg的推理过程比MedSegDiffs快约100倍，生成单个分割图的速度大约快28倍。

表 4 还比较了这些模型的逆向过程。潜在估计方案使得 SDSeg 能够在单步操作中生成分割图，并且串联潜在融合模块使得 SDSeg 仅需进行一次采样而不损害模型性能。此外，潜在估计使得 SDSeg 不再依赖任何外部采样器进行采样。

Stability Evaluation

由于扩散模型是生成模型，它们生成的样本可能会展现出一定的变化性。然而，在医学分割模型的背景下，多样性并不被视为一个优点特征，因为医疗专业人员需要人工智能的帮助以保持一致性和可靠性。给定一个训练好的模型和固定的测试数据，我们评估基于扩散的分割模型在以下两个任务中的稳定性：

表 5 展示了 SDSeg 在这些测试中的显著稳定性，突显了其在不同初始噪声条件下在分割任务中的可靠性。

Ablation Study

我们的消融研究评估了 SDSeg 内每个组件的贡献，详见表 6。基准模型依赖稳定扩散与交叉注意力生成基于图像的分割图。引入潜在融合串联显著增强了性能，有助于有效学习空间信息和特征。此外，可训练编码器通过从分割目标中提取相关语义特征显著提升了性能。

虽然潜在估计损失函数在一定程度上提升了性能，但其主要优势在于显著加速了逆向过程，使得 SDSeg 能够放弃传统的采样器，采用单步逆向过程，如图 3 所示。

总结

本文提出了 SDSeg，一个利用稳定扩散进行医学图像分割的新颖高效框架。我们引入了一种潜在估计策略，实现了单步潜在预测，从而消除了多步逆向过程的需求。模型采用串联潜在融合，将学习到的图像潜在有效地引导分割任务。此外，可训练的视觉编码器增强了模型学习图像特征并适应多样化图像模态的能力。SDSeg 在五个分割数据集上实现了最先进的性能，显著减少了训练资源需求，加速了推理过程，并保持了卓越的稳定性。

总结