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，因此探索将公认的变压器架构集成到这一强大的新主干中的方法是值