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个基准数据集上优于目前所有的最先进技术。