MedSegDiff-v2

本文的贡献包括：

  • 第一个将Transformer集成到基于扩散的通用医学图像分割模型中。
  • 提出了具有U-SA(Uncertain Spatial Attention)的Anchor Condition以减小扩散方差。
  • 提出了具有SS-Former的Semantic Condition，以建模分割噪声和语义特征的相互作用。
  • 在包括5种图像模态的20个器官分割任务上实现了SOTA性能。

总体架构image-20240104151754725

条件模型（绿色UNet)：从原始图像中提取分割特征

扩散模型（蓝色UNet)：有两个输入分别是锚点条件（蓝色箭头）和噪声分割信息（黑色箭头）

锚定条件与U-SA

锚点条件：将条件模型的解码分割特征整合到扩散模型的编码器特征中（蓝色箭头）

U-SA机制用于从条件模型中提取一个粗糙的锚点特征，并将其整合到扩散模型中，为扩散模型提供了一个正确的预测范围，同时也让它进一步完善了预测结果，具体做法如下：

image-20240104161300625

注1：是来自原始图像的条件分割特征，是噪声分割图像的扩散特征

注2：为高斯卷积核目的是做平滑处理，然后取最大值是为了保留最相关的信息并且消除一些噪声和不必要的细节

语义条件与SS-Former

SS-Former：使得模型可以学习条件语义特征与噪声信息的交互。

具体设计如下：

image-20240104164023507

注1：绿色为语义信息，蓝色为噪声信息

注2：和是可学习的权重参数，它们用于线性映射，以调整语义信息和噪声信息之间的关系

NBP-Filter:将语义信息和噪声信息交互得到的结果M调整到统一的频率范围(扩散模型生成的过程中具有一定的随机性，这种随机性会对数据进行一些扰动，NBP-Filter降低扩散生成过程中引入的随机性的影响)

NBP-Filter从坐标图学习权重图，坐标图的每个点代表了不同的频率分量，并且使用两个MLP层将时间信息投影到两个值，这两个值分别用于缩放和移位，然后使用得到的缩放因子和移位因子去调整坐标图的表示。

实验部分

表1：MedSegDiff-V2在不同图像模态上与SOTA分割方法的比较。灰色背景表示这些方法是为特定任务提出的。

image-20240105132907016

注：REFUGE2（眼底图像）、BraTs （脑肿瘤分割)、TNMIX(甲状腺结节分割)、ISIC（皮肤病变分割）

表2：使用Dice分数评估MedSegDiff-V2与SOTA分割方法在AMOS数据集上的比较。最佳结果以粗体表示。

image-20240105133848749

image-20240105134743328

表3：MedSegDiff-V2在BTCV数据集上与其他先进分割方法的比较结果的表格。Dice Score用于评估模型性能，最佳结果以粗体标注。

image-20240105135108153

消融实验

image-20240105135511006

提出的U-SA在所有数据集上均优于先前的空间注意力。

仅使用SS-Former提供了较小的改进，但与NBP-Filter结合使用则带来了显著的改进，证明了提出的SS-Former设计的有效性。

采样次数对模型效果的影响

扩散模型具有一定的随机性，所以需要进行多次采样，然后将采样后的结果集成在一起。

作者评估了各种基于扩散的医学分割模型的采样次数对模型效果的影响image-20240105140410979

最佳性能在大约50个集成后实现

MedSegDiff-V2其他扩散方法进行比较时，观察到它需要更少的采样次数来收敛。

MedSegDiff-V2更优越的起始点和更稳定的预测可以导致更高的性能上限

作者在REFUGE2-Cup数据集上讨论了样本多样性对模型效果的影响

image-20240105143354543

CI(置信区间):置信区间越大，一致性越好。

GED(广义能量距离)：GED越低，表示一致性越好。

U-SA在CI方面较低，而在GED方面较高，表明样本的多样性较大，表明其生成的样本大部分落在目标的不确定性区域内，效果不好。

而单独使用SS-Former而没有U-SA时，该模型在CI最高、GED最低的情况下达到了最好的一致性，未能充分利用扩散模型的多样性集成能力。

将U-SA和SS-Former组合成MedSegDiff-V2，性能得到了显著提高，表明SS-Former有助于减轻U-SA中生成的噪声，而U-SA为模型提供了更多的多样性，从而相互改进。