MedSegDiff-V2: Diffusion based Medical Image Segmentation with Transformer

MedSegDiff-V2:基于变压器的扩散医学图像分割

摘要

扩散概率模型(Diffusion Probabilistic Model, DPM)最近在计算机视觉领域获得了广泛的应用，这要归功于它的图像生成应用，如Imagen、Latent Diffusion Models和Stable Diffusion，这些应用已经展示了令人印象深刻的能力，并在社区内引发了许多讨论。最近的研究进一步揭示了DPM在医学图像分析领域的应用，正如医学图像分割模型在各种任务中表现出的令人称赞的性能所强调的那样。尽管这些模型最初是由UNet体系结构支持的，但是通过集成视觉转换机制来增强它们的性能存在着潜在的途径。然而，我们发现简单地结合这两个模型会导致性能低于标准。

为了有效地整合这两种前沿技术进行医学图像分割，我们提出了一种新的基于变压器的扩散框架，称为MedSegDiffV2。在20个不同图像模态的医学图像分割任务中验证了该方法的有效性。通过综合评估，我们的方法优于先前的最先进的(SOTA)方法。代码发布在https://github.com/KidsWithTokens/MedSegDiff

1 介绍

医学图像分割是将医学图像分割成不同的感兴趣区域。这是许多医疗应用的关键一步，如诊断和图像引导手术。近年来，人们对自动分割方法越来越感兴趣，因为它们有可能提高结果的一致性和准确性。随着深度学习技术的进步，一些研究已经成功地应用了基于神经网络的模型，包括经典卷积神经网络(cnn) (Ji et al 2021;Wu et al . 2022b)和最近流行的视觉变压器(ViTs)(Chen et al . 2021;Wang et al . 2021b)，用于医学图像分割任务。

最近，扩散概率模型(Diffusion Probabilistic Model, DPM)(Ho, Jain, and Abbeel 2020)作为一种强大的生成模型得到了普及，能够生成高质量和多样化的图像(Ramesh et al 2022;Saharia

等人2022;Rombach et al . 2022)。受其成功的启发，许多研究将DPM应用于医学图像分割领域(Wu et al . 2022c;Wolleb等2021;Kim, Oh，和Y, 2022;郭等人2022;Rahman et al . 2023)。

他们中的许多人通过使用DPM在几个基准测试中报告了新的SOTA。该模型的卓越性能源于其固有的随机抽样过程(Wu et al . 2022c;Rahman et al . 2023)。DPM能够通过多次运行生成不同的分割预测。这些样本之间的多样性直接捕获了医学图像中与目标相关的不确定性，其中器官或病变通常具有模糊的边界。然而，值得注意的是，所有这些方法都依赖于经典的UNet主干。与日益流行的视觉变压器相比，经典的UNet模型在分割质量上有所妥协，这可能导致在集合中生成发散但不正确的掩模，最终引入永久阻碍性能的噪声。

下一步自然是将基于变压器的UNet(如TransUNet(Chen et al 2021))与DPM结合起来。然而，我们发现以一种直接的方式实现它会导致性能欠佳。一个问题是变压器抽象的条件特征与扩散主干的特征不兼容。变压器能够从原始图像中学习深度语义特征，而扩散主干从损坏和噪声掩模中抽象特征，使特征融合更具挑战性。此外，变压器的动态和全局特性使其比cnn (Naseer等)更敏感

2021)。因此，以前基于扩散的方法(Wu et al . 2022c)中使用的自适应状态策略将导致变压器整定的较大方差。这导致了更多的集合和收敛困难。

为了克服上述挑战，我们设计了一种新的基于变压器的医学图像分割扩散框架，称为MedSegDiff-V2。主要思想是在扩散过程中对原始图像的主干采用两种不同的调节技术。

一种是锚定条件，它将条件分割特征集成到扩散模型编码器中，以减小扩散方差。我们设计了一种新的不确定空间注意(U-SA)集成机制，该机制用更多的不确定空间注意放宽了条件分割特征

图1:MedSegDiff-V2的插图，它从(a)管道的概述开始，并继续使用单个模型的放大图，包括(b) SS-Former，和(c) NBP-Filter。

从而为扩散过程提供了更大的灵活性，以进一步校准预测。另一种是将条件嵌入与扩散嵌入相结合的语义条件。为了有效地弥合这两种嵌入之间的差距，我们提出了一种新的变压器机制，称为频谱空间变压器(SS-Former)，用于嵌入集成。SS-Former是一种频域交叉注意链，采用时间步长自适应神经带通滤波器(NBP-Filter)对每次的噪声和语义特征进行对齐。

简而言之，本文的贡献如下:•我们首次将transformer集成到基于扩散的通用医学图像分割模型中。

•我们提出了一个带有U-SA的锚定条件来减轻扩散方差。

•我们提出了基于SS-Former的语义条件来建模分割噪声和语义特征的相互作用。

•我们在包括5种图像模式的20种器官分割任务上实现了SOTA性能。

相关工作

基于变压器的医疗分割

先前的研究强调了基于变换的模型在医学图像分割中实现SOTA结果的潜力。一个值得注意的例子是TransUNet(Chen等)

2021)，将变压器与UNet结合作为瓶颈特征编码器。从那时起，几部作品提出将尖端变压器技术纳入

医学图像分割模型的主干，包括swan - unet (Cao et al . 2022)、swan - unet (Tang et al . 2022)

2022)和DS-TransUNet(Lin et al . 2022)。由于最近基于UNet的扩散分割模型在医学图像分割中实现了新的SOTA，因此探索将公认的变压器架构集成到这一强大的新主干中的方法是值得的。

医学分割的扩散模型

扩散模型最近在各种分割任务中显示出巨大的潜力，包括医学图像(armatto III et al . 2011;Caron等2021;曹等

2022;Chen, Ma, and Zheng 2019)。事实上，这些模型利用随机抽样过程来生成隐式的分割集合，从而提高了分割性能(Zhai et al . 2022)。然而，如果没有对多样性的有效控制，集合往往难以收敛，因此，提高每次采样迭代的样本精度至关重要。

方法

MedSegDiff-V2的扩散过程

我们基于(Ho, Jain, and Abbeel 2020)中提到的扩散模型设计了我们的模型。扩散模型是由两个阶段组成的生成模型:正向扩散阶段和反向扩散阶段。在前向过程中，通过一系列步骤T逐步将高斯噪声添加到分割标签x0中。在逆向过程中，训练神经网络来恢复原始状态数据通过反向加噪处理。这可以用数学方法表示如下:

式中θ表示逆向过程的参数。

从高斯噪声分布出发，pθ(xT) = N (xT;0, In×n)，其中I为原始图像，反向过程将潜在变量分布pθ(xT)转换为数据分布pθ(x0)。为了保持与正演过程的对称性，反向过程逐步恢复噪声图像，最终得到最终清晰的分割。

根据DPM的标准实现，我们利用编码器-解码器网络进行学习。为了实现分割，我们在原始图像的先验信息上设置阶跃估计函数ε。这种条件可以表示为:

这里，T ransF表示基于变压器的注意机制。EIt表示条件特征嵌入，在我们的例子中，它对应于原始图像的嵌入。Ext表示当前步骤的分割映射的特征嵌入。这两个组件通过变压器合并在一起，并通过UNet解码器D进行重构。步骤索引t与组合嵌入和解码器特征集成，每个步骤索引使用共享学习查找表嵌入，遵循(Ho, Jain, and Abbeel 2020)中描述的方法。

总体架构

MedSegDiff-V2的总体流程如图1所示。为了介绍这个过程，我们考虑扩散过程的单个步骤t。首先将噪声掩码xt输入到一个称为扩散模型的UNet中。扩散模型是通过另一种标准的UNet，即条件模型，从原始图像中提取分割特征。扩散模型采用了锚定条件和语义条件两种不同的条件作用方式。根据输入流，首先对扩散模型的编码器施加锚定条件。它将锚点分割特征(即条件模型的解码分割特征)集成到扩散模型的编码特征中。这允许扩散模型通过一个粗略但静态的参考来初始化，这有助于减少扩散方差。然后对扩散模型的嵌入施加语义条件，将条件模型的语义分割嵌入集成到扩散模型的嵌入中。

这种条件集成由SS-Former实现，它弥补了噪声和语义嵌入之间的差距，并利用变压器的全局性和动态性的优势抽象出更强的表示。

MedSegDiff-V2使用DPM后的标准噪声预测损失Ln (Ho, Jain, and Abbeel 2020)和监督条件模型的锚点损失Lanc进行训练。lancs是软骰子损失lice和交叉熵的结合

损失特性。具体来说，总损失函数表示为:Lttotal = Ltn + (t≡0 (mod α))(lice + β lce)(3)其中t≡0 (mod α)通过超参数α控制对条件模型的监督次数，交叉熵损失由超参数β加权，分别设为5和10。

带有U-SA的锚点条件

没有卷积层的感应偏置，transformer块具有更强的表示能力，但在训练数据有限的情况下，对输入方差也更敏感(Naseer et al 2021)。直接将变压器块添加到扩散模型中会导致每次输出的方差很大。为了克服这种负面影响，我们采用MedSegDiff(Wu et al . 2022c)的结构，并在扩散模型中引入锚定条件操作。

锚条件从条件模型中提供了一个粗略的锚特征，并将其集成到扩散模型中。这为扩散模型提供了一个正确的预测范围，同时也允许它进一步完善结果。具体来说，我们将条件模型的解码分割特征集成到扩散模型的编码器特征中。我们提出了U-SA机制用于特征融合，以表示给定条件特征的不确定性。形式上，考虑我们将最后一个条件特征f−1c整合到第一个扩散特征f0d中。U-SA可以表示为:

其中*表示滑动窗口内核操作，·表示一般元素操作。在方程中，我们首先应用一个可学习的高斯核kG / f - 1c来平滑激活，因为f - 1c作为锚点，但可能不完全准确。然后，我们选择平滑映射和原始特征映射之间的最大值，以保留最相关的信息，从而得到平滑的锚点特征图。然后，我们将ffc集成到f0d中，得到增强的f0d特征。具体来说，我们首先应用1×1卷积k1×1conv将锚点特征通道减少到1，并在Sigmoid激活后将其与f0d相乘，然后将其添加到f0d的每个通道中，类似于空间注意的实现(Woo et al 2018)。

使用SS-Former的语义条件

扩散模型预测噪声掩模输入的冗余噪声，导致其嵌入与条件分割语义嵌入之间存在域间隙。当使用矩阵操作时，这种差异会影响性能，例如在绞合变压器中。为了解决这一挑战，我们提出了一种新的频谱空间变压器(SS-Former)。我们的核心思想是在频域中学习条件语义特征和扩散噪声特征的相互作用。我们使用一种称为神经带通滤波器(NBP-Filter)的滤波器将它们对齐到一个统一的范围

频率，即频谱。NBP-Filter学习通过特定的频谱，同时限制其他频谱。由于噪声水平(频率范围)对每一步都是特定的，因此我们以自适应的方式对扩散时间步长进行了学习。

SS-Former的鸟瞰图如图1 (b)所示，它由N个具有相同架构的块组成。

本文设N = 4。每个块由两个类似交叉注意力的模块组成。首先将扩散噪声嵌入编码为条件语义嵌入，下一个对称模块将最后一个语义嵌入编码为扩散噪声嵌入。这允许模型学习噪声和语义特征之间的相互作用，并实现更强的表示。形式上，考虑c0是条件模型的最深特征嵌入，e是扩散模型的最深特征嵌入。我们首先把c0和e转移到傅里叶空间，分别记为F (c0)和F (e)注意，特征映射都是按照标准视觉转换方法进行拼接和线性投影的。然后以e为查询，c0为键，计算傅里叶空间上的关联权映射，可以表示为M = (F (c0)Wq)(F (e)Wk)T，其中Wq和Wk是傅里叶空间中可学习的查询权和键权。

然后，我们应用nbp滤波器来对齐频率的表示。我们注意到M中的每个点现在代表一个特定的频率，并且由于我们需要控制一个连续的频率范围，因此直观地建立一个从特征映射位置到频率幅度的平滑投影。为了实现这一点，我们使用神经网络从坐标图中学习权重图。通过这样做，网络的归纳偏置将促进平滑投影的学习，因为相似的输入自然会产生相似的输出(Sitzmann et al 2020;Wu and Fu 2019)。这个想法被广泛应用于3D视觉任务，被称为神经辐射场(NeRF)(Mildenhall et al 2020)。但与原始NeRF不同的是，我们进一步用时间步长信息来约束它。具体来说，该网络以一个坐标图作为输入，并产生一个注意力图作为过滤器，两者都具有相同的大小m。我们使用一个简单的卷积块堆栈和中间层归一化来实现它。为了用时间步长信息约束网络，我们利用扩散模型的时间步长嵌入对归一化特征进行缩放和移位。

我们使用两个MLP层将当前时间步嵌入投影到代表均值和方差的两个值，这两个值分别用于缩放和移位。我们总共堆叠了R = 6个这样的块和一个Sigmoid函数来产生最终的过滤器。最后，筛选器按元素与管道中的关联映射M相乘。NBP-Filter以端到端的方式与整个管道进行训练。

然后使用快速反傅立叶变换(IFFT)将过滤后的关联映射M '转移回欧几里德空间，并应用于值为f = f−1(M ')(c0wv)的条件特征，其中wv是可学习的值权重。

我们还使用MLP进一步细化注意力结果，得到最终的特征≈c0。下面的注意模块与第一个注意模块是对称的，但使用组合特征~ c0作为查询，噪声嵌入e作为键和值，以便将分割特征转换到噪声域。变换后的特征c1将作为条件嵌入下一个块。

实验

数据集

我们在五种不同的医学图像分割数据集上进行了实验。使用两个数据集来验证一般分割性能，它们是带有16个解剖结构的公共AMOS2022(Ji et al 2022)数据集和带有12个解剖结构注释的公共BTCV(Fang and Yan 2020)数据集，用于腹部多器官分割。其他四个公共数据集REFUGE-2 (Fang et al .

2022)、BraTs-2021数据集(Baid等2021)、ISIC 2018数据集(Milton 2019)和TNMIX数据集(Pedraza等)

使用2015)验证模型在眼底图像的光学杯分割、MRI图像的脑肿瘤分割和超声图像的甲状腺结节分割等多模态图像上的性能。关于数据集的更多细节见附录。

实现细节

所有实验均使用PyTorch平台进行，并在4个NVIDIA A100 gpu上进行训练/测试。所有图像都统一调整为256×256像素的分辨率。

使用批量大小为32的AdamW(Loshchilov and Hutter 2017)优化器以端到端方式训练网络。初始学习率设置为1 ×10−4。

我们采用了100个扩散步骤进行推理。我们为集成运行了10次模型，这比MedSegDiff中的25次要少得多(Wu et al 2022c)。然后我们使用STAPLE算法(Warfield, Zou, and Wells 2004)融合不同的样本。我们通过Dice分数，IoU和HD95指标来评估分割性能。

主要结果

比较SOTA在腹部多器官分割中的应用

为了验证一般医学图像分割的性能，我们在多器官分割数据AMOS和BTCV上比较MedSegDiff-V2和SOTA分割方法。Dice得分的定量结果分别如2和3所示。在表中，我们与社区中广泛使用和认可的分割方法进行了比较，包括基于cnn的方法nnUNet(Isensee et al . 2021)，基于变压器的方法TransUNet(Chen et al . 2021)， UNetr(Hatamizadeh et al . 2022)， swun -UNetr(Jiang et al . 2022)和基于扩散的方法EnsDiff (Wolleb et al . 2021)， SegDiff(Amit et al . 2021)， MedSegDiff (Wu et al . 2022c)。我们还比较了扩散和变压器的简单组合模型。我们将MedSegDiff中的UNet模型替换为TransUNet，并在表中表示为“MedSegDiff + TransUNet”。

如2和3所示，先进的网络架构和复杂的设计对于实现良好的性能至关重要。考虑到结构，基于变压器的模型(如swan - unet)优于精心设计的基于cnn的模型(nnUNet)。基于扩散的MedSegDiff模型再次优于基于变压器的模型

图2:在BTCV上与SOTA分割模型的视觉对比

大部分的器官。然而，网络体系结构本身并不是性能的决定因素。例如，精心设计的基于cnn的模型nnUNet在表中显著优于基于变压器的模型TransUNet和UNetr。这对于基于扩散的模型也是成立的。我们可以看到，直接采用扩散模型进行医学图像分割，即EnsDiff和SegDiff的分割效果不如UNetr和swan -UNetr。变压器和扩散模型的简单组合，即MedSegDiff + TransUNet，获得的性能甚至比标准MedSegDiff更差。通过在扩散+变压器模型中引入锚定条件和SS-Former, MedSegDiff-V2克服了这些挑战，表现出了优异的性能。我们还在2中进行了定性比较。可以看出，MedSegDiff-V2预测的分割图具有更精确的细节，即使在低对比度或模糊区域也是如此。

与SOTA在多模态图像上的比较

我们还比较了MedSegDiff-V2和SOTA分割方法在三种不同图像模式下的具体任务。结果见1。表中提出了用于光学杯分割的ResUnet(Y u et al . 2019)和BEAL(Wang et al . 2019)， TransBTS(Wang et al . 2019)

采用MTSeg(Gong等人2021)和UltraUNet(Chu、Zheng、Zhou 2021)进行甲状腺结节分割，采用FA T- net (Wu等人2022a)和BA T(Wang等人2021a)进行皮肤病变分割。

从表中可以看出，MedSegDiff-V2在5个不同的任务中都超过了所有其他方法，突出了其在各种医学分割任务和图像模态上的卓越泛化能力。相比之下基于unet的MedSegDiff, MedSegDiff- v2在optical - cup上的得分提高了2.0%，在脑瘤上的得分提高了1.9%，在甲状腺结节上的得分提高了3.9%，强调了其基于变压器的主干的有效性。此外，当与MedSegDiff + TransUNet进行比较时，MedSegDiff- v2的性能优于MedSegDiff + TransUNet，这清楚地证明了所提出的U-SA和SS-Former在提高性能方面的功效。

消融研究

我们进行了全面的烧蚀研究，以验证所提出模块的有效性。结果如4所示，其中ac . cond。和Sem.Cond。分别表示锚定条件和语义条件。

如表所示，c。显著改进了香草扩散模型，提出的U-SA在所有数据集上都优于之前的空间注意。在Sem.Cond。，单独使用SS-Former只提供了微小的改进，但将其与NBP-Filter结合使用会产生显著的改进，证明了所提出的SS-Former设计的有效性。

分析与讨论

内隐集合效应

许多先前的研究证实了这一点(Wolleb et al . 2021;Wu等人2022c;阿米特等人

2021)，多次采样运行的隐式集成在基于扩散的方法中起着至关重要的作用。在扩散模型上下文中，隐式集成是指将来自单个扩散模型的多个采样的预测组合在一起，而不是融合来自不同模型的预测。

在本研究中，我们评估了各种基于扩散的医学分割模型的集成性能，如3所示。评估基于AMOS数据集计算的平均Dice Score。每个配置运行20次，并使用平均Dice Score作为性能

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

表2:通过Dice Score评估AMOS数据集，MedSegDiff-V2与SOTA分割方法的比较。最好的结果用粗体表示。

指标。在图中，我们将“MedSegDiff+TransUNet”设置表示为“MSD-Trans”。我们的发现表明了一个共同的趋势，即模型性能在最初的50个集合中迅速提高，然后稳定下来。一般来说，最好的表演是在大约50次合奏之后实现的。

当将MedSegDiff-V2变体与其他扩散方法进行比较时，我们发现它需要更少的集成来收敛。它开始时的性能明显更好，超过MedSegDiff 5%，并始终保持超过2%的领先优势。这突出了MedSegDiff-V2的效率，因为它即使使用更少的集成迭代也能获得令人满意的结果。此外，这表明一个更好的起点和更稳定的预测可以导致更高的性能上限。这与我们的假设一致，即低质量的样本会通过引入噪声不断降低模型的性能。这再次证明了引入U-SA进行散度控制和利用SS-Former获得更好起点的重要性。

不确定性分析

在5中，我们比较了refuge - cup数据集上的样本多样性。我们比较了以前基于dpm的方法，骨干网与单个提议的模块，以及最终的MedSegDiff-V2。我们使用广义En-来评估样本间的方差

能量距离(GED)和置信区间(CI)。GED是一种常用的度量标准，通过比较分割的分布来衡量预测和真实分布之间的一致性(Kohl et al 2018)。能量值越低，一致性越好。

从表中可以看出，与之前的方法相比，我们提出的U-SA方法的CI较低，GED较高，表明样本多样性更大。然而，也观察到所提出的模型达到了更高或可比的性能，这表明其生成的样本大多落在目标的不确定性区域内。当单独使用SS-Former，不使用U-SA时，模型与最高CI和最低GED的一致性最好。虽然SS-Former获得了最大置信度的优良性能，但未能充分利用扩散模型的分集集成能力。通过将U-SA和SS-Former组合为MedSegDiff-V2，性能显著提高，置信度仍然很高。这表明SS-Former有助于减轻U-SA中产生的噪声，而U-SA为模型提供了更多的多样性，从而相互改进。

模型效率和复杂性

在5中，我们还提出了模型复杂性和Gflops与其他的比较

基于扩散的分割方法。报告的Gflops是在隐式集成中达到稳定之前单个256 × 256图像的处理速度。我们认为，过去10个集团的业绩差异小于0.1%是趋同的指标。该度量对于基于扩散的分割模型的实际应用具有重要意义，因为用户通常会迭代运行扩散模型以获得稳定的结果。

从表中可以看出，与传统的深度学习模型不同，由于存在隐式集成，基于扩散的模型中的参数数量与Gflops并不直接相关。例如，尽管MedSegDiff-V2集成了变压器块并占用了更多参数，但它在实现时需要更少的Gflops更少步骤的稳定性。相比之下，MedSegDiff- v2消耗的Gflops只有MedSegDiff的一半，但在各种分段任务中的表现优于MedSegDiff，如上所示。

这强调了MedSegDiff-V2的效率及其在实际应用中的潜力。

表3:MedSegDiff-V2与SOTA分割方法在Dice Score评估的BTCV数据集上的比较。

最好的结果用粗体表示。

表4:锚锚调节和SS-Former的消融研究。SA表示空间注意。

图3:基于dpm方法的集成效果比较。我们展示了随着采样次数的增加，它们在AMOS上的平均骰子分数的性能。

表5:模型参数、Gflops和生成样本不确定度的比较

结论

在本文中，我们通过在原始UNet主干(MedSegDiff-V2)中加入变压器机制来增强基于扩散的医学图像分割框架。我们提出了一种新的SS-Former架构来学习噪声和语义特征之间的相互作用。

对比实验表明，该模型在20种不同图像模态的医学图像分割任务上优于以往的SOTA方法。作为第一个基于变压器的医学图像分割扩散模型，我们相信MedSegDiff-V2将成为未来研究的基准。

通常，在训练、采样和评估生成模型时，会经历以下步骤：

1. 训练阶段：

   • 数据准备： 准备训练数据集，包括输入数据和相应的标签（如果有的话）。
   • 模型选择和设计： 选择适当的生成模型，并设计其架构。这可能涉及选择生成对抗网络（GAN）、变分自编码器（VAE）等。
   • 损失函数定义： 定义模型的损失函数，该函数衡量生成的输出与目标数据之间的差异。
   • 优化器选择： 选择用于优化模型参数的优化器，如随机梯度下降（SGD）或其他变体。
   • 训练过程： 使用训练数据集进行多轮训练，通过最小化损失函数来调整生成模型的参数。这可能需要一些时间，具体取决于数据集的大小和模型的复杂性。

2. 采样阶段：

   • 模型加载： 在训练完成后，加载已经训练好的生成模型。
   • 生成样本： 使用已训练的模型从潜在空间中生成新的样本。这通常涉及从模型的分布中采样，并将采样结果转换为可视化的输出，如图像或文本。
   • 样本保存： 将生成的样本保存到适当的位置，以备后续评估或应用。

3. 评估阶段：

   • 定量评估： 使用一些指标来定量评估生成模型的性能。这些指标可能包括生成样本的多样性、逼真度、覆盖率等。
   • 定性评估： 对生成的样本进行定性评估，通过直观观察来判断生成的样本是否合理、多样化，并符合预期。
   • 应用场景测试： 在实际应用场景中测试生成模型，看它是否能够满足具体任务的需求。
   • 可能的微调： 根据评估结果，可能需要微调模型的参数或架构，然后重新进行训练和评估。

这些步骤通常是循环的，可能需要多次迭代，直至生成模型达到满意的性能。在每一步中，都需要仔细监控模型的行为，以确保生成的样本质量和模型性能。