前言
和大家分享一下我们发表在 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，具有以下贡献：
- SDSeg 建立在稳定扩散（SD）上，SD 是一种潜在扩散模型（LDM），在感知上等价的低分辨率潜在空间中举行扩散过程，使扩散过程对计算更为友爱
- 引入了一个简单的潜在估计损失，使 SDSeg 能够在单步逆过程中生身分割效果，并提出了一种连接潜在融合技术，以消除对多个样本的需求
- 将条件视觉编码器设置为可训练，以学习图像特征举行分割并顺应多个医学成像范畴
- 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
由于扩散模型是生成模型，它们生成的样本可能会显现出一定的变革性。然而，在医学分割模型的背景下，多样性并不被视为一个优点特征，因为医疗专业人员需要人工智能的帮助以保持一致性和可靠性。给定一个训练好的模型和固定的测试数据，我们评估基于扩散的分割模型在以下两个任务中的稳定性：
- Dataset-level Stability：举行重复推断测试数据，使用 LPIPS（Zhang等人，2018）指标来衡量不同推断之间的变异性
- Instance-level Stability：通过在固定条件下举行重复推断，检验模型在不同初始噪声下的一致性，使用 PSNR、SSIM 和 MS-SSIM 作为评估指标
表 5 展示了 SDSeg 在这些测试中的显著稳定性，突显了其在不同初始噪声条件下在分割任务中的可靠性。
Ablation Study
我们的消融研究评估了 SDSeg 内每个组件的贡献，详见表 6。基准模型依靠稳定扩散与交叉留意力生成基于图像的分割图。引入潜在融合串联显著增强了性能，有助于有用学习空间信息和特征。此外，可训练编码器通过从分割目标中提取相关语义特征显著提升了性能。
虽然潜在估计损失函数在一定程度上提升了性能，但其重要优势在于显著加快了逆向过程，使得 SDSeg 能够放弃传统的采样器，采用单步逆向过程，如图 3 所示。
总结
本文提出了 SDSeg，一个利用稳定扩散举行医学图像分割的新奇高效框架。我们引入了一种潜在估计计谋，实现了单步潜在猜测，从而消除了多步逆向过程的需求。模型采用串联潜在融合，将学习到的图像潜在有用地引导分割任务。此外，可训练的视觉编码器增强了模型学习图像特征并顺应多样化图像模态的能力。SDSeg 在五个分割数据集上实现了最先进的性能，显著减少了训练资源需求，加快了推理过程，并保持了杰出的稳定性。
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
