MedSegDiff-V2: Diffusion-Based Medical Image Segmentation with Transformer

摘要： 

       最近的研究揭示了 DPM 在医学图像分析领域的实用性，医学图像分割模型在各种任务中表现出的出色性能就证明了这一点。尽管这些模型最初是以 UNet 架构为基础的，但通过整合视觉Transformer机制，仍有可能提高其性能。然而，我们发现，简单地将这两个模型结合在一起会导致性能不佳。为了将这两种尖端技术有效地整合到医学图像分割中，我们提出了一种新颖的基于变换器的扩散框架，称为 MedSegDiffV2。我们在 20 个不同图像模式的医学图像分割任务中验证了它的有效性。通过综合评估，我们的方法证明优于先前的最先进（SOTA）方法。

Introduction

        我们为医学图像分割设计了一种新颖的基于变换器的扩散框架，称为 MedSegDiff-V2。在扩散过程中，我们在原始图像的主干上采用了两种调节技术。一种是锚定条件（Anchor Condition），它将条件分割特征整合到扩散模型编码器中，以减少扩散方差。我们设计了一种新颖的 "确定空间注意力"（U -SA）,一种是整合机制，它放宽了具有更多不确定性的条件分割特征，从而为扩散提供了更大的灵活性，以进一步校准预测结果。另一种是将条件嵌入整合到扩散嵌入中的语义条件。为了有效弥合这两种嵌入之间的差距，我们提出了一种新颖的转换机制，称为频谱-空间转换器（SS-Former），用于嵌入整合。SSFormer 是一个频域交叉注意链，每次都有一个时间步长自适应神经带通滤波器（NBP-Filter）来对齐噪声和语义特征。

        -我们首次将变换器集成到基于扩散的模型中，用于普通医学图像分割。

        -我们提出了一种带有 U-SA 的锚条件，以减轻扩散方差。

        -我们提出了语义条件与 SS-Former 模型，以模拟分割噪声和语义特征的相互作用。

Related Work

Diffusion Model for Medical Segmentation

        最近，扩散模型在包括医学图像在内的各种分割任务中展现出巨大潜力（Armato III 等人，2011 年；Caron 等人，2021 年；Cao 等人，2022 年；Chen、Ma 和 Zheng，2019 年）。事实上，这些模型利用随机抽样过程生成隐含的分割集合，从而提高了分割性能（Zhai 等人，2022 年）。然而，如果不对多样性进行有效控制，集合往往难以收敛，导致多次耗时的采样迭代。此外，这些分歧样本不仅无法达到预期目标，还会引入噪声，影响分割质量。因此，提高每次采样迭代的采样精度至关重要。

Method

Diffusion Process of MedSegDiff-V2 

        我们的模型是根据（Ho、Jain 和 Abbeel，2020 年）中提到的扩散模型设计的。扩散模型是一种生成模型，由两个阶段组成：正向扩散阶段和反向扩散阶段。通过一系列步骤 T，高斯噪声被逐渐添加到分割标签 x0 中。在反向过程中，神经网络经过训练，通过逆转噪声添加过程来恢复原始数据。这可以用数学方法表示如下：

 

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

 

这里，T ransF 表示基于变换器的注意力机制。EI t 表示条件特征嵌入，在我们的例子中，它对应于原始图像的嵌入。Ex t 表示当前步骤的分割图特征嵌入。这两个部分通过变换器整合在一起，并通过 UNet 解码器 D 进行重构。步骤索引 t 与嵌入和解码器的组合特征整合在一起，每个步骤索引按照（Ho、Jain 和 Abbeel，2020 年）中描述的方法，使用共享学习的查找表进行嵌入。

Overall Architecture 

        MedSegDiff-V2 的整体流程如图 1 所示。 为了介绍该流程，请考虑扩散流程的单步 t。首先将噪声掩膜 xt 输入一个称为扩散模型的 UNet。扩散模型由从原始图像中提取的分割特征通过另一个标准 UNet（称为条件模型）进行调节。扩散模型采用两种不同的调节方式：锚点条件和语义条件。按照输入流程，锚点条件首先施加到扩散模型的编码器上。它将锚分段特征（即条件模型的解码分段特征）整合到扩散模型的编码特征中。这样，扩散模型就可以通过一个粗略但静态的参考来初始化，从而有助于减少扩散差异。然后将语义条件强加给扩散模型的嵌入，从而将条件模型的语义分割嵌入整合到扩散模型的嵌入中。这种有条件的整合是通过 SS 变压器实现的，它在噪声嵌入和语义嵌入之间架起了一座桥梁，并利用变压器的全局和动态特性优势抽象出了一种更强的表示方法。

 

        MedSegDiff-V2 采用标准噪声预测损失 Ln（DPM，Ho、Jain 和 Abbeel，2020 年）和监督条件模型的锚损失 Lanc 进行训练。Lanc 是软骰子损失 Ldice 和交叉熵损失 Lce 的组合。具体来说，总损失函数表示为 :

其中，t ≡ 0 (mod α) 通过超参数 α 控制对条件模型的监督次数，交叉熵损失通过超参数 β 加权，分别设为 5 和 10。

Anchor Condition with U -SA

        没有卷积层的感应偏差，变压器块具有更强的表示能力，但在训练数据有限的情况下，对输入方差也更加敏感（Naseer 等，2021 年）。在扩散模型中直接添加变换器块会导致每次输出的方差过大。为了克服这一负面影响，我们对 MedSegDiff（Wu 等，2022e）的结构进行了调整，并在扩散模型中引入了锚定条件操作。

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

其中，∗ 表示滑动窗口核操作，- 表示一般元素操作。在等式中，我们首先在 f -1 c 上应用可学习的高斯核 kG 来平滑激活，因为 f -1 c 可作为锚点，但可能并不完全准确。然后，我们选择平滑图与原始特征图之间的最大值，以保留最相关的信息，从而得到平滑锚特征 fanc。然后，我们将 fanc 整合到 f0 d 中，得到增强特征 f ′0 d。具体来说，我们首先应用 1 × 1 卷积 k1×1conv 将锚特征通道还原为 1，并与 Sigmoid 激活后的 f 0 d 相乘，然后将其添加到 f 0 d 的每个通道中，类似于空间注意力的实现（Woo 等人，2018 年）。 

Semantic Condition with SS-Former

         扩散模型预测了来自噪声掩码输入的冗余噪声，导致其嵌入与条件分割语义嵌入之间存在域差距。在使用矩阵操作（stranded transformer）时，这种差距会影响性能。我们提出了一种新颖的频谱空间变换器（SS-Former）。我们的主要想法是学习条件语义特征和扩散噪声特征在频域中的相互作用。我们使用一种名为神经带通滤波器（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 等，2020 年；Wu 和 Fu，2019 年）。这一想法被广泛应用于三维视觉任务中，并被称为神经辐射场（NeRF）（Mildenhall 等，2020 年）。但与最初的 NeRF 不同的是，我们进一步用时间步长信息对其进行调节。具体来说，该网络将坐标图作为输入，并生成一个注意力图作为滤波器，两者的大小相同，都是 M。为了用时间步信息对网络进行调节，我们使用扩散模型的时间步嵌入对归一化特征进行缩放和移动。我们使用两个 MLP 层将当前的时间步嵌入投影到代表均值和方差的两个值上，这两个值分别用于缩放和移动。我们总共堆叠了 R = 6 个这样的块和一个 Sigmoid 函数，以生成最终的滤波器。最后，滤波器与管道中的亲和图 M 按元素顺序相乘。NBP-Filter 是通过整个流水线以端到端的方式进行训练的。

        滤波后的亲和力图 M′通过反快速傅立叶变换（IFFT）转回到欧几里得空间，并应用于值条件特征：f = F -1(M ′)(c0wv)，其中 W v 是可学习的值权重。我们还使用 MLP 进一步完善注意力结果，得到最终特征 ̃ c0。下面的注意力模块与第一个模块对称，但使用综合特征̃ c0 作为查询，并使用噪声嵌入 e 作为键和值，以便将分割特征转换到噪声域。转换后的特征 c1 将作为下一个区块的条件嵌入。