【Diffusion分割】医学图像分割中去噪扩散技术

Noise Reduction of Diffusion Processes in Latent Space for Medical Imagery Classification

扩散模型（DPM）在图像生成领域展现出卓越的能力，在处理图像生成任务方面表现尤为突出。自其诞生以来，扩散模型已成功应用于多种判别性任务中，并以其强大的噪声去除与图像恢复管道技术闻名于世。近年来，在医学成像领域取得了显著的应用突破。然而，在实际应用过程中仍面临诸多挑战：因内存占用较高且采样过程耗时较长等现实问题制约着其推广使用。针对这些问题，在本研究中我们提出了一种创新的条件生成建模框架 （LDSeg）。该框架基于潜空间设计了一套独特的条件扩散机制，在此过程中实现了对多标签对象的精准n维图像分割，并显著提升了传统扩散模型在处理域转移问题时的表现能力

INTRODUCTION

通过引入高斯噪声来破坏分割掩码, 将会使得底层分布呈现出不自然的扭曲, 其中由于掩码仅由有限的模式组成（取决于存在的语义类别）, 这种扭曲尤其明显. 因此, 由于各模式之间缺乏渐变过渡, 去噪器的学习过程变得异常困难. 为了获得最终稳定的分割掩码, 需要执行额外的阈值处理. 然而, 由于高频噪声的存在, 最终得到的结果可能会表现出类似孔洞等特征结构[11]. 此外,Wu 等人提出的在去噪器隐层中加入频率解析器块的方法[9], 虽然有助于抑制高频噪声的影响, 但仍无法确保采样后结果干净可靠. Bogensperger 等人则将离散化的分割掩码转换为符号距离函数（SDFDDPM）[11], 其中每个像素点代表其所在位置与最近对象边界之间的符号欧氏距离. 这种方法的一个显著缺陷在于其对多类图像的距离图呈现模糊特性. 最近,Zaman 等人提出了将分割掩码重新参数化为图结构的方法[13], 该方法旨在保证图上连续表面距离的变化具有自然性. 然而, 该模型同样面临多类目标表示问题: 不同物体表面之间的界限变得难以区分. 因此, 需要一种既能同时处理多类目标又能在状态转移上保持平稳的技术方案

DPM方法的主要困难在于优化扩散采样的效率。研究者们开发了多种技术来降低自然图像生成所需的时间 [14-17]。在医学领域中，由于每个参与者通常拥有大量的三维扫描数据集 [20] ，这不仅增加了存储需求还提升了扩散模型用于生成高质量分割结果所需的计算时间 为此 一些学者开始探索如何通过学习潜在空间中的结构加速扩散采样过程 [18-19] 。例如 PNVR等人提出了一种基于潜在空间的学习方法用于去噪器训练 并在此基础上构建了一个高效的Unet架构用于分割任务 同时结合注意力机制以更好地利用扩散模型的特性 [20] 。为了整合潜在空间扩散的优势 在本研究中我们提出了一种新的条件扩散生成框架（LDSeg）专注于医学图像分割任务 该框架通过学习目标物体形状的单变量高斯潜表征以及源图像嵌入特征实现了精确分割目标 并在此过程中充分考虑了模型参数与计算资源之间的平衡关系

据所知，这项研究首次将基于物体形状的单变量高斯潜在空间与去噪器相结合，并通过优化设置显著提高了采样效率

连续潜在空间使得常规扩散技术得以直接应用于前向与反向过程，在多类对象分割中有效解决了标签分割遮罩上非自然型噪声的注入问题。

METHOD

我们提出的LDSeg框架主要由两部分构成：第一部分为掩膜自编码器（Mask Autoencoder），它被用来学习目标物体形状的低维潜在表示；第二部分为基于源图像生成的嵌入信息作为条件，在每个时间步上推断出对应的噪声分布情况。其中图像嵌入是通过图着色自编码器进行学习的过程。其工作流程已在图1中进行了详细展示。

Mask Autoencoder

在分割标签上注入高斯噪声并非自然操作，因为标签/掩码图像仅包含少量特定模式（对应于对象类别数量）。当数据分布呈现为多变量高斯混合时，去噪器仍难以有效学习中间噪声分布特性。为此我们提出了一种解决方案：建议通过学习标签图像的单变量低维高斯表示来缓解这一固有问题。具体而言，在这种框架下我们希望实现两个目标：第一，在潜在空间中构建一个单变量高斯分布模型；第二，在潜在空间中构建一个对应的反函数模型用于重建输入掩码空间信息。为此我们建议采用一个简单的Res-Unet[21]自编码器进行处理：该编码器结构设计上未引入跳跃连接层以减少复杂度并提高泛化能力。通过该编码器可得到enc(.)≈h(.)；同时解码器可实现dec(.)≈h⁻¹(.)以便完成从潜在空间到输入空间的重建过程。假设M∼pdata(M)代表真实掩码/标签图像，则其潜在表示m0及其重构版本̄M可通过以下方式获得：

我们的目标旨在掌握联合概率密度模型 l_\theta(m_0,\overline{M} \mid M)；自动编码器所采用的损失函数是一种多分类交叉熵损失。

其中，
N 表示样本总数，
C 表示分类总数，
yi,j 表示实例i中类j的真实标签，
pi,j 则表示实例i中类j的预测概率。

Coder的最后一层是层归一化层，在此过程中确保潜在向量 m0 是单一维度的零均值高斯分布。从本质上讲，在这一架构中设计了这样的机制：通过Mask编码器能够提取出被遮蔽区域中物体形状所对应的低维潜在表示信息，并通过Mask解码器能够重构出与原图接近度较高的图像内容

Conditional Denoiser (CD)

标准的 DPM 去噪器有两个输入：一个是输入图像的噪声版本；另一个是相应的时间步值。为了实现分割任务；denoiser需要额外提供一些条件。这些条件可以来自两个方面：其一是源图像 [11, 9]；其二是通过指示目标对象的文本 [20] 来实现。我们建议使用图像嵌入作为 denoiser 的条件。具体而言；图像嵌入是一种能够提取源图像低维潜在表示的技术；它是通过一个编码器学习得到；该编码器架构与掩码编码器相似；唯一区别在于末尾没有归一化层。在前向过程中；给定 m0 和噪声方差表参数 α 和 β 的情况下；通过高斯块 G 可以生成时间步 t 的噪声版本 mt [4, 14]。（如图 2 所示）

去噪器基于标准的Unet架构设计，并具备时间编码机制以及自注意力机制。具体而言，在文献[4]中提出了类似的去噪器架构设计方案。在协同训练过程中，默认情况下图像编码器与去噪器共同进化以优化性能参数，并致力于学习不同时间步t=1,...,T时的状态转移概率分布dθ(mt-1|mt,I,t)，其中T表示最终的时间步长且mt∼N(0,I)（t=1,...,T）。在这里,T代表最终的时间步长,m_t服从标准正态分布。
我们通过最小化以下损失函数来训练条件去噪器：

\mathcal{L}(\theta)=\mathbb{E}_{t,\mathbf{x}_t,\epsilon}\left[\|\epsilon-\epsilon_\theta(\mathbf{x}_{t-1},\mathbf{x}_t,t;\theta)\|^2\right]

使用经过训练的掩码自动编码器的条件去噪器训练算法如算法 1 所示。

Reverse Process For Segmentation

由于该图像编码器独立运作于去噪器之外，在反向过程的初始阶段获取条件图像嵌入即可。在反向过程中我们的核心目标是生成基于图像嵌入的潜在表示 m0。与传统的DPM模型不同，在这一任务中我们从零均值单位方差的高斯分布 N(0, I) 开始，并将其作为时间步 T 的潜在掩码表示̃mT。随后我们对每个时间步 t 进行迭代去噪操作直至完成全部 T 步骤后停止并获得最终掩码表示 ̃m0。此时我们将此掩码表示作为训练好的掩码解码器输入并经过解码运算即可获得最终分割结果 S = dec(̃m0)。如图 2 所示 在下行部分展示了一个完整的反向过程示例 而算法 2 则给出了具体的实现流程图