Diffusion Transformer U-Net for MedicalImage Segmentation

用于医学图像分割的扩散变压器U-Net

摘要：

扩散模型在各种发电任务中显示出其强大的功能。在将扩散模型应用于医学图像分割时，存在一些需要克服的障碍:扩散过程调节所需的语义特征与噪声嵌入没有很好地对齐;这些扩散模型中使用的U-Net骨干网对上下文信息不敏感，而上下文信息在反向扩散过程中对于精确的像素级分割至关重要。

为了克服这些限制，我们提出了一个交叉注意模块来增强源图像的条件反射，以及一个基于转换器的U-Net，该U-Net具有多尺寸窗口，用于提取各种尺度的上下文信息。在Kvasir-Seg、CVC Clinic DB、ISIC 2017、ISIC 2018和Refuge 5个不同成像模式的基准数据集上进行了评估，我们的扩散变压器U-Net具有很强的泛化能力，并且在这些数据集上优于所有最先进的模型。

1 介绍

卷积神经网络(CNN)和视觉变形(ViT)等深度学习(DL)方法已经应用于医学图像分割[7,8,17]，并取得了良好的效果。然而，这些深度学习方法在其网络架构上有一些固有的限制。例如，cnn能够提取局部特征，但不能直接提取全局特征，而vit采用固定窗口，这限制了它们提取精确像素级分割所必需的精细上下文细节的能力。

近年来，去噪扩散概率模型(Denoising Diffusion Probabilistic Model, DDPM)[9]在各种条件和无条件生成任务中表现出了很好的性能，也被应用于医学图像分割[23,24]。尽管取得了成功，但仍存在一些不足:(1)从源图像中提取的语义嵌入与扩散中的噪声嵌入不能很好地对齐

工艺，导致调理不良，性能欠佳;(2)在反向扩散(去噪)过程中，这些基于ddpm的方法中的UNet主干对各种尺度的上下文信息不敏感，在cnn和vit中也观察到这一点。

基于这些局限性，我们提出了一种扩散变压器U-Net，并做出了以下贡献:

提出了一种具有前向和后向过程的条件扩散模型来训练分割网络。在反向去噪过程中，通过一个新的交叉注意模块将噪声图像的特征嵌入与条件源图像的特征嵌入对齐。然后，通过分割网络将其去噪为源图像的分割掩码。

一个基于变压器的U-Net多尺寸窗口，命名为MT U-Net，它旨在提取像素级和全局上下文特征，以实现良好的分割性能。

扩散模型训练的MT U-Net在各种成像模式上具有很强的泛化能力，在结肠镜图像中的息肉分割[1,10]、皮肤镜图像中的皮肤病变分割[4,5]、视网膜眼底图像中的光学杯分割[14]等5个基准数据集上优于目前所有的最先进技术。

图1所示。用交叉注意扩散模型训练MT U-Net。

2 方法

2.1扩散模型

扩散有正向和反向两个过程(图1)。在正演过程中，通过T个时间步长逐渐加入高斯噪声，将地真值M0变换为带噪声的地真值MT。在相反的过程中，首先，源图像I和噪声映射m_t +1通过一个编码器E(两个残差初始化块[18])得到嵌入fI∈Rh×w×c1和fM∈Rh×w×c2(下标I和M表示图像和映射)，其中h、w和c1 (c2)分别是嵌入的高度、宽度和通道。然后，在特征空间中通过交叉注意(Cross-Attention, CA)模块对两个嵌入进行对齐。将对齐后的特征映射作为带噪声的输入输入到MT U-Net中进行恢复M^t

这个反向过程从t = t−1迭代到t = 0(即，当t = t−1时，初始的m_t +1, m_t, i = s, i = s Mt, a + n和m_0最终恢复，预计与基本真理M0相同)。

图2展示了我们的CA模块的架构，该模块用于校准fM和fI，以改善扩散模型的调节。首先，将fM和fI分割成小块，并通过小块编码(Patch Encoding, PE)层将其平面化。然后，使用位置编码层(PoE)获取补丁的位置信息，并将其添加到原始补丁嵌入中以保持其位置信息。两个位置内置补丁嵌入的维度使用线性投影(LP)层对齐，并通过层归一化(LN)进行归一化，将两个LN之后的输出表示为f p M∈Rd和f p I∈Rd(补丁的d-dim特征向量)。第三，利用自关注进行特征融合;

图2所示。交叉关注(CA)模块的体系结构。

其中f p M是查询(Q)， f p M和f p I的连接是键(K)和值(V)。？表示转置。第四，根据[20]，我们通过层归一化(LN)和二层多层感知器(MLP)对LSA的输出进行编码，以提取更多的上下文信息。辅助连接(残差)用于增强信息传播。最后，我们应用一个重塑(RS)层来重塑并组装成与fM相同大小的补丁。

2.2 多规格变压器U-Net (MT U-Net)

图3(a)展示了我们的MT U-Net的架构，包括编码和解码部分。编码部分由Patch Partitioning层、Linear Embedding层、PoE和四个编码器块组成。Patch Partitioning层将输入分割成不重叠的Patch, Patch大小为2 × 2。

图3所示。提出的MT U-Net的体系结构，以及MT模块。为清晰起见，图中没有给出时间步长嵌入。

使用线性嵌入层将这些补丁与时间嵌入一起平面化为D × 1维线性嵌入。然后在通过四个编码器块之前，将从PoE获得的位置信息添加到线性嵌入中。除了最后一个编码器块只包含MT模块外，每个编码器块由一个多尺寸变压器(MT)模块和一个补丁合并层组成。MT模块提取多尺度上下文特征(稍后详细说明)，Patch merge层对特征图进行下采样。受U-Net[15]的启发，我们采用了一种s - k - p连接，利用来自编码器的多尺度上下文信息来克服降采样过程中空间信息的丢失。

与编码器块类似，每个解码器块由一个MT模块和一个补丁扩展层组成，除了第一个解码器块只包含MT模块。补丁扩展层对特征映射进行上采样和重塑操作。最后，我们使用线性投影层来获得像素级预测。

所提出的多尺寸变压器(MT)模块(图3(b))与传统变压器不同[6]。MT模块由多尺寸窗口和可移动窗口两部分组成。多尺度窗口部分提取多尺度上下文信息，移动窗口部分丰富了提取的信息。多尺寸窗口部分有K个平行分支，每个分支由一个层归一化(LN)、多头自注意(SA)、辅助连接(残量)和一个多层感知器(MLP)组成，该感知器有两层，后面跟着GELU激活函数。多头自注意使用不同的窗口大小来提取多尺度的上下文特征。这些单独分支的输出被合并，并被发送到移动窗口部分。移动窗口部分的结构类似于在多尺寸窗口中使用单个分支，但在自关注(SW-SA)中使用移动窗口。

2.3 训练与推理

在训练过程中，将源图像及其分割地面真值图作为扩散模型的输入。扩散模型使用噪声预测损失(LNo ise)[12]和交叉熵损失(LCE)进行训练。

在推理过程中，从高斯分布中采样的噪声图像与测试图像一起作为反向过程的输入。

3 实验结果

3.1 数据集和评估指标

为了评估该方法的有效性和泛化能力，我们测试了不同的医学图像分割任务，包括:(1)结肠镜图像的息肉分割(Kvasir-SEG (KSEG) [10]， CVC- clinic DB (CVC)[1])，(2)皮肤镜图像的皮肤病变分割(ISIC 2017 (IS17 ') [5]， ISIC 2018 (IS18 ')[4,19])，以及(3)视网膜眼底图像的光学杯分割(REFUGE (REF)[14])。Dice Coefficient (DC)和Intersection over Union (IoU)被用作评价指标。

3.2 实现细节

通过交叉验证，MT模块中的分支数设置为3，窗口大小分别为4、8、16。扩散变压器U-Net使用SGD优化器进行40000次迭代训练，动量为0.6，批量大小为16，学习率设置为0.0005。在扩散中，我们使用T = 1000步的线性噪声调度器。为了与最近基于扩散的分割模型[23,24]进行公平的比较，在推理过程中，25个预测的平均集合被认为是最终预测。所有的实验都是使用NVIDIA Tesla V-100 GPU和32gb RAM进行的。

3.3 性能对比

首先，我们将我们的方法与几个众所周知的U-Net和/或transformer相关的分割模型进行定量比较，包括U-Net[15]、unnet++[26]、Attention U-Net[13]、Swin U-Net[2]、Trans U-Net[3]和SegFormer[25]。利用源代码对这些模型进行训练，并在实验数据集上进行评估。为了公平比较，所有模型对每个数据集使用相同的实验方案。定量结果如表1所示。

表1。与U-Net和/或变压器相关的最先进方法的比较。KSEG、CVC、IS18采用80:10:10’(训练:验证:测试数据分割)实验协议;分别在REF和IS17 '上使用默认分割。

图4所示。与SOTA方法在KSEG[10]、CVC[1]、IS18 '[4,19]、IS17 '[5]和REF[14]数据集上的定性比较。蓝色等高线代表真实情况，绿色等高线代表预测结果。

我们的扩散变压器U-Net在不同成像模式的五个数据集上优于所有其他U-Net或Transformer相关模型，验证了其有效性和泛化能力。

其次，我们将我们的扩散变压器U-Net与其他U-Net或变压器相关模型进行定性比较。从图4中随机抽样的测试图像中，我们观察到其他模型产生过分割(例如，Trans U- net, SegFormer)或欠分割结果(例如，U- net, SegFormer)

表2。与SOTA结果比较。' - ':未报告结果。“*”:图像数量。

Net, unet ++， Attention U-Net, Swin U-Net)，我们的分割掩码最接近地面真相，证明了我们方法的有效性。

最后，我们将我们的扩散变压器U-Net与五个数据集上所有最新的最佳模型进行了比较，如表2所示。除MedSegDiff和MedSegDiff- v2外，引用方法的结果均直接从其论文中复制。这两种方法被重新训练，并在REF数据集上进行评估。注意，由于一些方法在IS18数据集上使用不同的实验协议。为了公平比较，我们使用两种不同的协议训练/交叉验证/测试我们的方法，并将我们的方法与使用相同协议的其他方法进行比较。如表2所示，我们的方法在这五个数据集上的表现始终优于目前所有的最佳模型，再次验证了我们的方法的有效性和优越性。

3.4 消融研究

我们进行了一系列烧蚀研究，以评估扩散变压器U-Net中每个模块的贡献，如表3所示:

-我们用简单的连接操作取代了扩散模型中的交叉注意(CA)，并应用这种简化的扩散模型来训练U-Net。即使这种简化的扩散模型(第2行)也可以提高第1行U-Net的性能，显示扩散的有效性。

-使用我们的带有CA模块的扩散模型(第3行)，与基本的串联操作(第2行)相比，性能得到了进一步提高，这验证了CA模型在扩散模型去噪过程中对特征嵌入进行对齐的贡献。

-使用CA模块的扩散模型，我们将没有多尺寸窗口的基本变压器单元[6]添加到U-Net中(第4行)。与第3行相比，这也提高了分割性能，这表明变压器可以帮助U-Net进行分割。

-基于第4行模型，我们在变压器中添加了多尺寸窗口(即第5行扩散变压器U-Net)。与烧蚀研究中的其他配置相比，这提供了最高的性能。

表3。消融对KSEG[10]、CVC[1]、IS18′[4,19]、IS17′[5]、REF[14]的影响。