【Diffusion分割】Cold SegDiffusion：医学图像分割的扩散模型

Cold SegDiffusion: A novel diffusion model for medical image segmentation

摘要：

随着深度学习的发展，扩散模型在医学图像分割任务中表现出了卓越的性能。然而，传统的分割扩散模型通常采用随机高斯噪声生成分割掩膜，导致分割掩膜不唯一，无法保证分割结果的可重复性。为解决这一问题，本文介绍了一种基于扩散模型的用于普通医学图像分割的新方法 Cold SegDiffusion。在该方法中，医学图像分割被概念化为一个去噪问题。覆盖医学图像的分割掩码作为分割编码器的输入，解决了因噪声随机性而产生非唯一掩码的难题。此外，对比度增强模块旨在将特征转换到频域，以解决医学图像中对比度低和边界消失的问题。此外，建议的条件交叉注意模块利用条件编码器和交叉注意权重来增强分割编码器输出的重要特征，从而提高网络聚焦目标区域的能力。我们在三个不同模式的医学图像分割数据集上对所提出的方法进行了验证。实验结果表明，Cold SegDiffusion 优于主流的分割方法。

1. Introduction：

最近，基于扩散模型的大型生成网络 Stable Diffusion [13] 和 DALL-E2 [14]，在生成高质量图像方面表现出了卓越的生成能力。此外，一些研究也证明了扩散模型在医学图像分割中的功效。例如，Wolleb 等人[15] 采用去噪扩散概率模型（DDPM）进行脑肿瘤分割，证明扩散模型可应用于医学图像分割。该方法在训练和采样阶段将医学图像作为先验知识，引导扩散模型从噪声分布中生成精确的分割掩膜。同样，Wu 等人[16]提出了一种用于普通医学图像分割的 MedSegDiff，该方法采用特征频率解析器来减轻噪声添加过程中引入的高频噪声成分的影响。虽然现有的基于扩散模型的分割方法已经取得了可喜的成果，但仍存在一些局限性和改进空间。例如，在利用 DDPM 进行图像分割时，高斯噪声的随机性会给遮罩的生成带来不确定性，从而影响分割结果的可重复性。此外，提高医学图像的分割准确性仍然是当前分割算法面临的重大挑战。如图 1 所示，医学图像通常存在噪声、边界模糊和对比度低等问题。这可能导致通过分割方法对 ROI 分割不足或过度[17]。在分割模型中整合注意力机制是缓解这些问题的有效策略之一。例如，Huang 等人[18] 设计了动态定位注意力，以增强网络学习本地信息的能力。Fan 等人[19]提出了一种 CSAP-UNet 方法，用于带有边缘增强的医学图像分割，该方法集成了注意力融合模块，以融合 CNN 和 Transformer 分支的输出特性。Ates 等人[20]提出了一种双交叉注意模块，通过捕捉全局通道和空间依赖关系来增强特征表示。因此，将注意力机制整合到分割模型中可以加强其聚焦 ROI 和提取图像全局信息的能力，提高分割结果的准确性和鲁棒性[21]。

为应对上述挑战，我们提出了基于扩散模型的 Cold SegDiffusion，用于自动分割医学图像。扩散模型可以理解为利用朗格文动力学围绕图像密度函数的随机行走[22]。这种行走开始时处于高温状态，噪声很大，然后逐渐退火，进入噪声最小的低温状态。因此，不需要高斯噪声或随机性的扩散模型被归类为冷扩散模型[23]。所提出的冷扩散模型包括扩散和去噪两个过程。在扩散过程中，医学图像被视为噪声，并逐渐叠加到分割掩膜上。在去噪过程中，将覆盖噪声（即医学影像）的分割掩膜输入 Cold SegDiffusion 以消除噪声。这一系列步骤将医学影像还原为分割掩膜，减轻了高斯噪声的随机性引起的分割掩膜的不确定性。解决生成分割掩码时的不确定性对于确保分割结果的可重复性至关重要，可进一步提高分割模型的一致性、可靠性和可追溯性[24]。同时，分割编码器包括对比度增强模块（CEM），利用可训练的频域滤波器来增强频域空间内特征的边缘和细节信息。该模块中集成的注意力机制可强化重要的空间和通道特征，同时抑制背景信息。此外，条件编码器和分割编码器的输出特性都是条件交叉注意模块（CCAM）的输入。在该模块中，条件编码器的输出用于查询分割编码器，以获得交叉注意力权重。这些注意力权重可以引导网络关注医学图像中的 ROI。最后，分割解码器对从 CCAM 接收到的输出特性进行解码，以获得与医学图像相对应的分割掩码。

-基于扩散模型的 Cold SegDiffusion 被提出用于医学图像分割。该方法解决了传统扩散模型生成的分割掩膜不唯一的难题，确保了分割结果的可重复性；

-所设计的对比度增强模块（CEM）利用可学习的频域滤波器增强编码特征的边缘和细节信息，同时利用空间和通道注意力增强任务相关特征，抑制无关特征；

-建议使用条件交叉注意模块（CCAM）来融合分割和条件编码器的输出特性，利用交叉注意权重，从而引导网络关注医学图像中的目标区域。

2. Related work：

2.1. Diffusion model：

最近的研究表明，扩散模型在图像生成、超分辨率重建和图像增强等一系列任务中表现出良好的性能。例如，Rombach 等人[13]提出了一种潜在扩散模型，该模型将图像形成过程分解为连续的去噪自动编码器应用，取得了卓越的合成效果。Zhou 等人[25]提出了一种多实例生成控制器，利用坐标和文本描述在单幅图像中生成具有不同控制的实例。受 DDPM 和去噪分数匹配的启发，Saharia 等人[26] 开发了 SR3 方法，通过迭代细化实现超分辨率。Zhou 等人[27] 提出了一种高效的金字塔扩散模型，用于恢复低照度图像中被噪声覆盖的细节。该模型利用全局校正器来减轻扩散模型可能导致的全局劣化。在这些成功经验的基础上，研究人员开始探索将 DDPMs 应用于图像分割任务。尽管做出了许多努力，但利用扩散模型进行图像分割的研究仍然相对有限。Baranchuk 等人[28] 提出了一种基于 DDPM 的两阶段图像分割算法。然而，这种非端到端分割方法容易造成误差累积，而且无监督的预训练过程无法捕捉特定任务的语义信息[29]。Amit 等人[30]提出的另一种方法是 SegDiff，它采用条件编码器引导网络从高斯噪声中恢复原始图像的分割掩膜。同时，Wu 等人[31]提出了用于医学图像分割的 MedSegDiff-V2 方法。该方法利用基于变换器的条件 U-Net 框架来提取扩散噪声和条件语义特征。然而，上述方法在测试过程中输入图像的高斯噪声会带来不确定性，导致分割掩码不唯一。因此，这些分割掩码的整合过程可能会导致过度分割或分割不足，从而降低基于扩散模型的分割方法的可靠性和可解释性。

2.2. Attention mechanism：

在各种计算机视觉任务中，注意力机制对提高深度学习算法的性能起着至关重要的作用 [32]。常用的注意机制包括空间注意、通道注意和自我注意。大量研究表明，整合注意力机制可以提高深度学习算法的分割性能。例如，Yang 等人[33] 通过整合多尺度前景-背景信息，提出了一种半监督视频分割方法。该方法采用实例级关注机制来捕捉全局信息，有助于缓解小感受野造成的局部模糊性，并增强不同物体尺度的鲁棒性。Valvano 等人[34] 在多尺度生成对抗网络中提出了一种注意力门控机制，通过对抗信号动态调整注意力机制。此外，Yang 等人[35] 提出了一种将物体与可扩展变换器关联起来的新方法，该方法集成了可扩展的长期短期变换器。这些变换器结合了基于层识别的注意力和可扩展的监督，证明了视频对象分割在线架构的可扩展性。Zhao 等人[36]为多病灶医学图像分割设计了一种先验注意力网络，通过将分割过程分解为两个阶段来增强其可解释性。此外，Rasti 等人[37] 开发了一种自适应双注意力模块，能够从空间和通道特征中提取上下文细节。该模块能自适应地考虑空间和信道的相互依存关系，利用两个可学习的参数突出重要的空间和信道特征。虽然将注意力机制整合到医学图像分割方法中已经取得了令人满意的成果，但人们对探索将其整合到扩散模型中以进一步增强医学图像分割效果的兴趣却日益高涨。因此，后续的研究方向涉及将注意力机制整合到扩散模型中。这将引导网络强调投资回报率，减轻医学图像分析中与分割不足和过度分割相关的挑战。

3. Methodology：

3.1. Overall framework：

Cold SegDiffusion 的学习目标可表述如下：

其中，x0 表示分割掩码，D 表示扩散过程，R 表示去噪过程。与标准 DDPM 不同，医学图像的分布并不遵循高斯分布。因此，建议的方法无法通过最小化高斯噪声与覆盖在掩膜上的医学图像之间的误差来恢复分割掩膜[38]。为解决这一难题，提出的方法采用了以下损失函数：

其中，f 表示神经网络，m 表示医学图像的数量，xi,t 表示第 i 个分割掩膜经过 t 步扩散后得到的图像。因此，该损失函数迫使网络去除前一时间步的覆盖噪声，以重建 t-1 时刻的分割掩膜。然而，这种方法往往会带来较高的计算复杂度和误差累积。因此，我们简化了扩散过程，以减轻该方法的计算负荷。简化后的扩散过程如下：

其中，x0 表示分割掩膜，z 表示分割掩膜对应的医学图像，αt = ∏t s=0（1 - βs）是分割掩膜的权重。同样，去噪过程的相应损失函数简化如下：

这一损失函数驱动去噪网络直接将医学图像还原为初始分割掩码。然而，在去噪过程中，随着医学图像权重参数的增加，噪声掩膜会逐渐与医学图像更加相似。因此，将噪声掩膜还原为初始分割掩膜变得越来越具有挑战性。为了提高网络在训练过程中处理复杂样本的能力，建议网络将重点放在与时刻 T 非常相似的噪声掩膜上。改进后的损失函数表述如下：

其中，γ 表示时间步长 t 的指数因子。γ 值越大，表示去噪网络越重视时刻 T 附近的噪声掩膜。同时，在采样过程中，医学图像通过一系列采样步骤逐渐转化为分割掩膜[23]。这种方法有利于细化输出并减少误差的积累，算法 1 演示了迭代过程。

Cold SegDiffusion 遵循扩散模型的标准实施方法，采用修改后的 ResUNet [39] 作为去噪网络的骨干。去噪网络由三个部分组成：分割编码器、分割解码器和条件编码器。首先，噪声掩码 xt 和医学图像 xT 分别输入到分割编码器和条件编码器。分割编码器通过对比度增强模块将特征转换到频域空间，并利用可学习的频域滤波器增强细节和边缘信息。随后，来自分割编码器和条件编码器的编码特征将作为条件交叉注意模块的输入。分割编码器的编码特征会被条件编码器的语义信息查询，从而获得交叉注意权重，自适应地增强网络集中于 ROI 的能力。最后，来自条件交叉注意模块的特征被分割解码器解码，从而得到精确的分割掩码。所提出的方法解决了传统分割扩散模型固有的生成分割掩码的非唯一性问题。它确保了分割结果的可重复性，大大提高了扩散模型的一致性和可靠性。

3.2. Contrast enhancement module：

医学图像分割任务面临的主要挑战是纹理模糊和边界消失。为了解决这些问题，一种典型的技术是利用频域处理来增强医学图像的纹理细节和边界信息[40]。基于这种技术，我们提出了对比度增强模块（CEM），用于在分割网络的编码阶段放大纹理和边界特征的对比度。如图 3 所示，该模块主要包括两个部分：利用快速傅立叶变换（FFT）增强频域信息，以及通过注意力机制在抑制背景信息的同时改善特征信息。

CEM 利用二维 FFT 层将编码特征从空间域转换到频率域，从而获得振幅和相位频谱。

其中，Me 表示编码特征，F (⋅) 表示二维 FFT，Mf 表示经过二维 FFT 处理后的频域特征。在频域中，振幅频谱通常传达全局信息，如图像的纹理和颜色。相反，相位频谱包含局部信息，如图像的轮廓和边缘[41]。该模块采用两个可学习的频域滤波器来权衡振幅和相位信息，以提高提取特征的对比度并捕捉高频细节。随后，加权特征经过反快速傅里叶变换（IFFT），实现从频域到空间域的重新转换，计算公式如下：

其中，F-1(⋅) 表示二维 IFFT，Mi 表示 IFFT 层的输出。振幅谱和相位谱分别用 Ma 和 Mp 表示。此外，wa 和 wp 分别表示振幅和相位频谱的可学习频域滤波器。可学习频域滤波器可以调整频率成分，从而有效捕捉上下文信息。

从二维 IFFT 层获得的输出特性输入到特征增强部分。在这一部分，信道注意模块（CAM）会压缩特征图中的空间信息，强调不同信道特征的语义信息[42]。它利用信道注意力权重加强与任务相关的信道信息。另一方面，空间注意模块（SAM）会压缩信道特征信息，强调特征图中的上下文空间信息[42]。该模块利用空间注意力权重，重点关注对分割结果有重大影响的关键区域。联合通道和空间注意力模块增强了任务相关的特征，同时抑制了任务无关的特征。随后，经过不同线性变换后的增强特征将分别用作通道查询 Mc、空间查询 Ms 和值 Mv。通道和空间查询被输入 CAM 和 SAM，分别获得通道和空间注意力权重。这些关注权重强调了数值中关键的信道和空间信息。注意力权重的计算方法如下：

其中，σ 表示 sigmoid 函数，f 7×7 表示核大小为 7×7 的卷积层。AvgPool 和 MaxPool 分别是全局平均池化层和全局最大池化层。同时，wc 和 ws 表示通道权重和空间注意力权重。这些权重相乘得出注意力分数，综合了空间注意力和通道注意力。将 Mv 与注意力分数加权后就得到了输出结果。计算过程如下：

其中 Mo 表示 CEM 的输出。该模块用于分割编码器的不同阶段，以增强不同比例的特征图。然后，这些增强的特征通过残差连接传送到分割解码器，帮助其生成准确的分割掩码。

3.3. Conditional cross-attention module：

利用注意力机制引导网络关注医学图像中的 ROI，可以提高其分割性能[43]。在 Cold SegDiffusion 中，引入了条件编码器和条件交叉注意模块（CCAM），以增强网络专注于目标区域的能力。具体来说，在训练过程中，分割编码器会收到覆盖不同程度医学图像的分割掩码。然而，当遮罩上的医学图像覆盖率较低时，分割编码器可能会面临捕捉重要目标信息的挑战。相比之下，条件编码器的输入是包含全面目标信息的初始医疗图像。因此，这一输入补充了分割编码器提取的特征。同时，CCAM 用于融合条件编码器和分割编码器的编码特征。在特征融合过程中，CCAM 对分割编码器的编码特征进行查询和增强，从而提高了建议方法聚焦于指定目标区域的能力。CCAM 模块的结构如图 4 所示。

来自分割编码器和条件编码器的编码特征被输入 CCAM。如公式 (14) 和 (15) 所示，这些特征会被复制成三份。

其中，Men 表示分割编码器的输出，Mce 表示条件编码器的输出。repeat(⋅, 3) 表示重复操作，即把变量复制三份。随后，来自分割编码器的编码特征通过三个不同的卷积层进行线性变换，核大小为 1 × 1，从而得出与分割编码器相对应的查询、键和值。同样，来自条件编码器的编码特征通过三个不同的卷积层进行线性变换，核大小为 1 × 1，从而生成与条件编码器相对应的查询、键和值。随后，将条件编码器生成的查询应用于分割编码器的密钥，并计算两者之间的相似度，即可得到关注度得分。注意力分数量化了条件编码器在不同位置上对分割编码器编码特征的关注程度。注意力分数按以下公式计算：