MedSegDiff-V2: Diffusion-Based Medical Image Segmentation with Transformer 论文总结

标题：MedSegDiff-V2: Diffusion-Based（基于扩散模型）Medical Image Segmentation（医学图像分割）with Transformer

论文（AAAI）：https://ojs.aaai.org/index.php/AAAI/article/view/28418/28816

源码：KidsWithTokens/MedSegDiff: Medical Image Segmentation with Diffusion Model (github.com)

一、摘要

研究背景：扩散概率模型（DPM）最近在计算机视觉领域获得了流行，这得益于其图像生成应用，如Imagen、潜扩散模型和稳定扩散，这些应用展示了令人印象深刻的能力，并在社区内引发了许多讨论。最近的研究进一步揭示了 DPM 在医学图像分析领域的效用，医学图像分割模型在各种任务中表现出的良好性能强调了这一点。

研究问题：尽管这些模型最初是由UNet架构支撑的，但存在通过集成视觉transformer机制来提高其性能的潜在途径。然而，本文发现简单地组合这两个模型会导致性能不佳。

主要工作：为了有效地将这两种前沿技术集成到医学图像分割中，本文提出一种新的基于transformer的扩散框架，称为MedSegDiffV2。

实验效果：在具有不同图像模态的20个医学图像分割任务上验证了其有效性。通过综合评估，所提出方法证明了比之前最先进的(SOTA)方法的优越性。

二、引言

三、相关工作

四、方法

A. Diffusion Process of MedSegDiff-V2 ( MedSegDiff-v2的扩散过程）

基于(Ho, Jain 和 Abbeel 2020) 中提到的扩散模型设计了本文的模型。扩散模型是由两个阶段组成的生成式模型：正向扩散阶段和反向扩散阶段。在正向过程中，通过一系列步骤 t 逐步将高斯噪声添加到分割标签 $x_0$ 中。在反向过程中，通过训练神经网络实现逆转噪声添加过程恢复数据。可以用数学表达式表示如下：

其中，θ 表示逆向过程中的参数。从高斯噪声分布出发， $p_\theta (x_T) = N(x_T;0, I_{n \times n})$ ，其中 $I$ 是原始图像，反向过程将潜变量分布 $p_\theta (x_T)$ 转换为数据分布 $p_\theta (x_0)$ 。为了保持与前向过程的对称性，反向过程逐步恢复噪声图像，最终获得最终的清晰分割。

Q：什么是这里指的潜变量分布？

A：Latent variable distribution（潜变量分布）指的是在统计建模中用来描述未直接观测到的、无法直接测量的变量的分布。

在本文和扩散模型中，Latent variable distribution（潜变量分布）常用来描述高斯噪声和图像加性后的数据分布，这是一个符合高斯分布的概率分布。

在DPM（扩散模型）的标准实现之后，本文利用编解码器网络进行学习。为了实现分割，本文以原始图像的先验信息为条件，采用步长估计函数ϵ。这个条件可以表示为（这个公式是对整个模型的概述，包括模型的输入，TransF模块和解码器的输入）：

这里，TransF 表示基于 transformer 的注意力机制。 $E_t^I$ 表示条件特征嵌入，在本文中，它对应于原始图像的嵌入。 $E^x_t$ 表示当前步骤的分割图的特征嵌入。这两个分量通过 transformer 合并在一起，并通过 UNet 解码器 D 进行重建。步骤索引 t 与组合的嵌入和解码器特征集成在一起，并且遵循(Ho，Jain，and Abbeel 2020)中描述的方法，使用共享的学习查找表来嵌入每个步骤索引。

B. Overall Architecture ( 整体架构 )

MedSegDiff-V2的整体流程如图所示。

整体流程概述：噪声掩码 $x_t$ 首先输入到UNet，称为 扩散模型（DIffusion Model）。扩散模型通过另一标准的UNet（称为 条件模型 Condition Model）从原始图像中提取的分割特征作为条件。将条件模型的解码分割特征作为锚点分割特征（先验特征）集成到扩散模型的编码特征中（辅助UNet分支解码器的先验特征和扩散模型逆恢复特征融合）。这使得扩散模型可以用一个粗略但静态的参考来初始化，有助于减小扩散方差。然后，将语义条件施加到扩散模型的嵌入中，将条件模型的语义分割嵌入集成到扩散模型的嵌入中（将辅助UNet分支的最底层全局语义特征与扩散模型逆恢复特征融合）。这种条件集成由SS-Former实现，弥合了噪声和语义嵌入之间的差距，并利用transformer的全局和动态特性的优势抽象出更强的表示。

输出结果：依旧是扩散模型中 $x_t \rightarrow x_{t-1}$ ，只不过这个 $x_{t-1}$ 受到了辅助 UNet 分支的指导，集合了先验特征和语义特征，使得扩散模型的反向过程更有方向性，会朝着指定的分割区域恢复。

MedSegDiff-V2主要由两步操作组成：

1. 辅助UNet分支解码器最顶层的先验特征和扩散模型逆恢复特征融合 —— U-SA
2. 将辅助UNet分支的最底层全局语义特征与扩散模型逆恢复特征融合 —— SS-Former

这两步操作的作用位置不同，U-SA作用于辅助UNet分支解码器顶层，以分割特征作为先验特征对逆恢复特征进行约束。SS-Former作用于辅助UNet分支的最底层，为逆恢复特征融合全局语义。

Q：这种融合方法不会破坏扩散模型的反向过程吗？毕竟扩散模型的反向过程是依赖于正向过程中的数据分布的，向反向过程中添加正向过程中不存在的数据，不会破坏反向过程的推理吗？

损失函数：MedSegDiff-V2在DPM(Ho、Jain和Abbeel 2020)之后使用标准噪声预测损失 $\mathcal{L}_n$ 和监督条件模型的锚点损失 $\mathcal{L}_{anc}$ 进行训练。 $\mathcal{L}_{anc}$ 是soft dice loss $\mathcal{L}_{dice}$ 和交叉熵损失 $\mathcal{L}_{ce}$ 的组合。具体地，总损失函数表示为：

其中 t ≡ 0 (mod α) 通过超参数 α 控制对条件模型的监督次数，交叉熵损失用超参数 β 加权，分别设为 5 和 10 。

C. Anchor Condition with U-SA ( 锚点条件与U-SA )

问题：没有卷积层的归纳偏差，transformer块具有更强的表示能力，但在训练数据有限的情况下，对输入方差也更敏感(Naseer et al. 2021)。直接将 transformer 块添加到扩散模型中会导致每次输出的方差很大。

方法：为了克服这种负面影响，本文调整了 MedSegDiff 的结构，并将锚点条件操作引入扩散模型。锚点条件从条件模型（辅助UNet分支）中提取出一种粗糙的锚点特征，并将其融入到扩散模型中。这为扩散模型提供了正确的预测范围，同时也允许它进一步完善结果。

过程：具体地，将条件模型解码后的分割特征融入扩散模型的编码器特征中。提出 U-SA 机制（高斯卷积 + 自注意力）进行特征融合，以表示给定条件特征的不确定性本质。形式上，将最后一个条件特征 $f^{-1}_c$ 集成到第一个扩散特征 $f^0 _d$ 中。U-SA 可以表示为：

其中，∗ 表示滑动窗口内核操作，· 表示常规的元素级操作。在方程中，首先在 $f^{-1}_c$ 上应用可学习的高斯核 Kg 来平滑激活，因为 $f^{-1}_c$ 用作锚，但可能不完全准确。因此，在平滑后的特征图和原始特征图之间选取最大值来保留最相关的信息，从而得到平滑的锚点特征图 fanc。然后将 fanc 融合到 $f^0_d$ 中，得到增强特征 $f^{'0}_d$ 。（本质上是一个对Key值进行高斯平滑处理的自注意力）

细节：具体地说，首先应用1×1卷积 $k_{1\times 1conv}$ 将锚点特征通道减少到1，并在Sigmoid激活后将其与 $f^0_d$ 相乘，然后将其添加到 $f^0_d$ 的每个通道，类似于空间注意的实现。

D. Semantic Condition with SS-Former ( 语义条件和SS-Former )

问题：扩散模型从噪声掩码输入中预测冗余噪声，导致其嵌入与条件分割语义嵌入之间的域间差距（频域和空间域之间的差距）。当使用矩阵操作时，这种差异会影响性能，例如在一个 transformer 中。

思想：为应对这一挑战，本文提出一种新的频谱空间Transformer (SS-Former)。本文的核心思想 是在频域上学习条件语义特征和扩散噪声特征的相互作用 。

方法：本文使用一个滤波器，称为神经带通滤波器 (NBP-Filter) 来将它们对齐到统一的频率范围，即频谱。NBP-Filter学习通过特定的频谱，同时约束其他频谱。以自适应扩散时间步长的方式学习该频谱，因为噪声水平 (频率范围) 对每个步骤都是特定的。这允许根据频率亲和度混合特征，也可以根据扩散步骤对齐特征。（通过 NBP-Filter 对每一个步长都学习一个特定的频谱，并将输入特征对齐到这一统一的频率范围）

结构：如图所示，由共享相同架构的 N 个块组成。本文中设置 N = 4。每个模块由两个类似交叉注意力的模块组成。

NBP-Filter设计来源：然后，应用 NBP-Filter 来对齐频率的表示。 $M$ 中的每个点现在表示一个特定的频率，由于我们需要控制一个连续的频率范围，直观的是建立一个从特征映射位置到频率幅度的平滑投影。为了实现这一点，本文使用神经网络从坐标映射中学习权重映射。通过这样做，网络的归纳偏差将促进平滑投影的学习，因为相似的输入将自然产生相似的输出(Sitzmann et al. 2020;吴和傅2019)。这种想法被广泛应用于3D视觉任务中，被称为神经辐射场(NeRF)(Mildenhall et al. 2020)。

过程：第一步将扩散噪声嵌入条件语义嵌入中，接下来的对称模块将上一个语义嵌入到扩散噪声嵌入中。这 允许模型学习噪声和语义特征之间的相互作用（选取最深层特征的作用），并实现更强的表示。形式上，考虑 $c^0$ 是条件模型中 最深的特征嵌入 ，e 是扩散模型中 最深的特征嵌入 。

首先将 $c^0$ 和 e 转移到傅里叶空间（FFT操作），分别记为 $F(c^0)$ 和 $F(e)$ 。请注意，特征图都是按照标准的vision transformer方法进行拼接和线性投影的。然后，以e为查询， $c^0$ 为键，计算傅里叶空间上的亲和度权重矩阵，该映射可以表示为 $M = (F(c^0)W^q)(F(e)W^k)^T$ ，其中 $W^q$ 和 $W^k$ 是傅里叶空间中可学习的查询权重和键权重。（傅里叶空间变换 + ViT ）

然后，应用 NBP-Filter 来对齐频率的表示。但与原始的NeRF不同，本文进一步用时间步长信息对其进行条件化。具体来说，该网络以一个坐标图（coordinate map）作为输入，并产生一个注意力图作为过滤器，两者都具有相同的大小 m。简单堆叠卷积块和中间层归一化来实现它。为了用时间步信息来调节网络，使用扩散模型的时间步嵌入来缩放和移动归一化特征。使用两个 MLP 层将当前时间步嵌入投影到表示均值和方差的两个值，分别用于缩放和移动。堆叠总共 6 个这样的块和一个 Sigmoid 函数来产生最终的滤波器。最后，将滤波器与亲和性矩阵逐元素相乘。
然后，将经过过滤的亲和力图 M' 使用逆快速傅里叶变换 ( IFFT ) 转换回欧几里得空间，并应用 value 值中的条件特征： $f = F^{-1}(M')(c^0w^v)$ ，其中 $w^v$ 是可学习的 value 权重 (卷积) 。还使用一个 MLP 进一步细化注意力结果，获得最终特征 $\tilde{c}^0$ 。接下来的注意力模块与第一个模块对称，但使用组合特征 $\tilde{c}^0$ 作为查询，噪声嵌入 e 作为键和值，以便将分割特征转换到噪声域。转换后的特征 c1 将作为下一个块的条件嵌入。

Q：坐标图（coordinate map）怎么获得的？由亲和度矩阵转化而来的吗？

五、实验

A. Dataset（数据集）

在五个不同的数据集上进行了实验。使用两个数据集来验证 通用分割性能 ，分别是包含16种解剖结构的AMOS2022公开数据集和包含12种被标注用于腹部多器官分割的解剖结构的BTCV公开数据集。另外使用四个公共数据集(包括:眼底图像的光学杯分割、MRI图像的脑肿瘤分割、超声图像的甲状腺结节分割)、BraTs-2021数据集、ISIC数据集和TNMIX数据集验证其 在多模态图像上的分割性能。（五种数据集，验证通用分割和多模态分割任务）

B. Implementation Details（实验细节）

实验设备：所有实验均使用PyTorch平台进行，并在4个NVIDIA A100图形处理器上进行培训/测试。

预处理：所有图像都被均匀地调整到256×256像素的分辨率。

优化器和批量大小：使用AdamW优化器以端到端的方式训练网络，批次大小为32。初始学习速率设置为1×10^−4。

扩散步数和执行次数：采用100个扩散步骤进行推理。为了集成，运行模型10次，这比MedSegDiff中的25次要少得多。然后使用STAPLE算法对不同的样本进行融合;

评价指标：通过Dice score，IoU，HD95指标对分割性能进行评估。

C. Main Results ( 主要结果 )

验证 MedSegDiff-V2 与SOTA方法在多器官分割数据集AMOS和BTCV上的分割性能。

对比方法：基于cnn的方法nnUNet，基于transformer的方法TransUNet、UNetr、Swin-UNetr以及基于扩散的方法EnsDiff、SegDiff、MedSegDiff。

结果分析：网络体系结构本身并不是性能的决定因素。例如，设计良好的基于 cnn 的模型 nnUNet 在表中大大优于基于 transformer 的模型 TransUNet 和 UNetr。对于基于扩散的模型也是如此。可以看到，直接采用扩散模型进行医学图像分割，即 EnsDiff 和 SegDiff ，其性能比 UNetr 和 Swin-UNetr 差。transformer和扩散模型的简单组合，即MedSegDiff + TransUNet，获得的性能甚至比标准MedSegDiff更差。

改进方法：通过引入锚定条件和SS-Former，MedSegDiff-v2克服了这些挑战，并表现出优越的性能。本文还在进行了定性比较。可以观察到，MedSegDiff-v2预测的分割图具有更精确的细节，即使在低对比度或模糊的区域。