【Diffusion分割】Diff-SFCT：空间频率交叉Transformer扩散模型

Diff-SFCT: A Diffusion Model with Spatial-Frequency Cross Transformer for Medical Image Segmentation

        现有的大多数语义分割方法主要采用带有判别模型的监督学习。虽然这些方法简单明了，但却忽略了底层数据分布的建模。在本文中，提出了一种基于扩散模型的新型医学图像分割框架，称为 Diff-SFCT。将语义分割表述为分割掩膜的生成问题，用潜在先验学习过程取代传统的像素判别学习，从而产生更准确的分割结果。Diff-SFCT 采用了卷积神经网络（CNN）和变换器相结合的骨干网络，利用了 CNN 的局部感知和变换器的全局信息建模能力。在 Diff-SFCT 中，设计了一个语义编码器，能有效地从真实图像中提取细粒度语义特征。同时，提出了一种新颖的空间-频率交叉变换器（SFCT）框架，它能有效地对漫反射噪声掩模的全局特征和真实语义特征进行建模和交互，缩小两者之间的域差距，增强模型的表征能力。此外，为了保留扩散模型中的空间和频率信息，设计了一个空间-频率注意模块（SFAM），作为卷积块的一部分。该模块提高了模型的空间和频率感知能力，同时产生的计算开销可以忽略不计。实验结果表明，DiffSFCT 大大优于其他分割方法，在各种医学影像分割数据集上表现出卓越的性能。

METHOD

Overall Architecture

        图 2 展示了 Diff-SFCT 在时间步长 t 上的整个过程。Diff-SFCT 由四个部分组成：语义编码器、扩散编码器、扩散解码器和空间-频率交叉变换器（SFCT）。语义编码器学习真实图像的语义特征，然后将其添加到相应尺度的扩散编码器特征中，从而获得多尺度融合特征。随后，SFCT 从扩散编码器和语义编码器中提取并交互全局特征信息，整合噪声掩码特征和真实语义特征。SFCT 采用傅立叶变换和交叉变换器来计算频域交叉注意，有效弥合了噪声掩码和真实语义之间的域差距。此外，它还利用变换器的长距离依赖建模能力，抽象出更强的特征表示。语义编码器和扩散编码器是 ResUNet 编码器的改进版本，而扩散解码器则是 ResUNet 解码器的改进版本。此外，在每个卷积块中，还引入了一个名为 "空间-频率注意模块"（SFAM）的新型注意模块，用于长期捕捉卷积结构之间的空间和频率信息，其计算开销可忽略不计。

        首先，通过添加高斯噪声，在扩散过程中破坏分割掩码 x0，为 T 个扩散步骤生成噪声掩码。对于多类分割，我们采用one-hot编码（one-hot encoding）将大小为 1 × W × H 的单通道分割掩码转换为包含 K 个类别的多通道分割图：x0 ∈ K × W × H，其中每个通道包含二进制标签（0，1），用于两类分割。在训练过程中，噪声掩码 xt∈ K × W × H 和真实图像 c∈ C × W × H（其中 C 代表真实图像中的模态数）作为 Diff-SFCT 模型的参数化输入，在反向扩散过程中预测下一步的噪声掩码 xt-1。在推理过程中，以真实图像为条件信息，反向扩散过程反复进行 T 步，从高斯噪声 xT 中生成最终的分割结果 x0。 

Spatial-Frequency Cross Transformer

     

   

        SFCT 是核心模块，每个 SFCT 块都由两部分组成：语义空间变换块（SSTB）和频率交叉变换块（FCTB），如图 3 所示。具体来说，采用两个连续的变换器块，为局部特征学习彻底引入全局指导。利用 SSTB 从语义编码器中提取全局特征，并通过 FCTB 整合扩散模型的骨干特征，从而促进信息感知以及语义信息与扩散噪声之间的融合。通过将全局语义信息与骨干网络的频率信息相结合，进一步增强了模型生成分割掩码的能力。此外，还重新设计了 FCTB 中的前馈网络频率 MLP，使 SFCT 在学习全局频率信息时更好地与原始图像保持一致。第 l 层 SFCT Block 的实现公式如下：

其中，Xl s 代表第 l 层的空间语义信息，Xl f 代表第 l 层的漫反射噪声掩码信息。

        语义空间变换器块。在这项工作中，采用标准变换器编码器块作为全局语义特征提取器，用于学习语义编码器的全局语义特征。传统的变换器编码器由 N 个相同的块组成，每个块由多头自注意（MHSA）和多层感知器（MLP）构成。此外，在每个 MHSA 模块和每个 MLP 模块之前都应用了层规范（LN），每个模块之后都应用了残差连接。 

        频率交叉变换块。在扩散模型的骨干网络中，为了更好地提取全局特征，使用多头交叉注意（MHCA）和快速傅立叶变换（FFT）建立长程依赖关系，以提取全局特征。这涉及在傅立叶空间而非欧几里得空间中融合真实语义信息和噪声掩码信息并使之相互作用。FCTB 的结构如图 4 所示，由两部分组成：左侧是频域和全局语义信息之间的交叉注意模块，右侧是过滤高频和低频相关信息的频率前馈网络。在频率交叉注意模块中，我们采用二维快速傅里叶变换将特征从空间域转换到频率域，以提取全局信息。随后，通过交叉关注，来自扩散编码器的主要信息可作为查询信息，而来自 SSTB 的全局语义信息则可作为键和值。这样就能将语义信息融合到扩散模型的噪声中。注意力机制的实现公式如下：

其中，F (-) 表示二维 FFT，F-1 (-) 表示二维 IFFT，Xf 表示扩散模型骨干的噪声掩码特征，Xs 表示语义编码器的真实语义特征，Q = F (Xf )，K = F (Xs)，V = Xs。

        频率前馈网络。在变换器框架内，采用了基于多层感知器（MLP）的前馈网络（FFN）来提取更丰富的语义信息。因此，开发一种在频域运行的 FFN 将有助于提高 SFCT 从噪声中恢复清晰图像特征的能力。由于并非所有低频和高频信息都有利于生成可信的清晰图像，因此引入了频域前馈网络（Frequency FFN），它能自适应地决定保留哪些频率信息。在 MLP 的基础上，加入了傅立叶变换，将多头交叉注意力学习到的特征转换到频域。引入了一个可学习的量化矩阵，以决定保留哪些频率信息，随后进行反傅里叶变换，将特征从频域转换回空间域。

Convolution Block and Spatial-Frequency Attention Module 

        在 Diff-SFCT 的骨干网络中，卷积块应用于语义编码器、扩散编码器和扩散解码器。本文采用的卷积块由两个 ResNet 块、一个注意力块和一个向下或向上块组成，如图 5(a) 所示。ResNet 块是 ResUNet 块的改进版。在每个 ResNet 块中，从两个卷积之间的扩散模型中引入了一个时间嵌入，旨在学习当前扩散步骤的时间信息，如图 5(b) 所示。

        在扩散模型的去噪过程中，空间和频率信息的利用至关重要，因此提出了一个名为 "空间-频率注意模块 "的新型注意力模块。 可应用于所有卷积块，以提高特征提取性能。SFAM 的结构如图 5（c）所示。SFAM 由并行串联的空间注意模块和频率注意模块组成。对于空间注意模块，从卷积块注意模块（CBAM）的设计中汲取了灵感，CBAM 是前馈卷积神经网络的一种简单而有效的注意模块，能够提取信道和空间维度上有意义的特征。CBAM 的整体工作流程可概括如下：

        其中，σ 表示 sigmoid 函数，f 7×7 表示滤波器大小为 7×7 的卷积运算。Pavg 表示平均汇集，Pmax 表示最大汇集。频率注意模块将模型的注意力引向有意义的频率范围，而空间注意模块则将模型的注意力引向有意义的区域。整个过程可以概括为 ：

其中，⊙ 表示元素相乘（又称哈达玛乘积），W 是与 F 维度相同的可学习矩阵，可以表示频域中的任何滤波器。此外，SFAM 还加入了残差块，以实现快速、稳定的训练，并将低频特征转移到末端。SFAM 的整个过程可概括如下：

Loss Function 

        标准扩散模型是利用负对数似然的变分上界和 L2 损失最小化来训练的。在本任务中，对扩散模型进行端到端训练，并采用 DDPM 的简单去噪方法。因此，可以通过最小化噪声 ε 预测的均方误差（MSE）损失来优化模型权重θ：

        除了用 MSE 损失优化扩散模型的目标之外，还利用 Dice 损失和二元交叉熵（BCE）损失作为语义分割的损失函数，这两种损失函数都广泛应用于医学图像分割任务中。Dice Loss 根据预测的分割结果与地面实况注释之间的 Dice 系数重叠度来优化网络，从而有效缓解前景与背景之间的不平衡。BCE 损失通常用于二元分类问题中的损失计算。Dice Loss 和 BCE Loss 的结合可以确保网络在专注于学习前景像素的同时，不会完全忽略与背景像素之间的关系，从而有效提高分割精度。因此，Diff-SFCT 的总损失计算公式如下：