MICCAI’24｜随机步骤对齐的Diffusion模型，解决无监督域适应图像分割问题

前言

像生成对抗网络（GANs）这样的图像生成器由于其复杂的内部工作机制，可以被视为“黑箱”，只能访问最终生成的图像，使得它们无法在执行领域对齐时使用生成过程中产生的中间分布的知识。为了解决这个问题，作者提出了一个新的UDA框架，利用扩散模型来捕获和转移一定量的域间知识，从而减轻领域偏移问题。设计了一个耦合的结构保持扩散模型，通过多步合成中间图像，使中间图像分布可访问。

论文信息

题目：Diffusion-based Domain Adaptation for Medical Image Segmentation using Stochastic Step Alignment

通过基于扩散的随机步骤对齐增强医学图像分割中的无监督域适应

作者：Wen Jiang and Albert C. S. Chung

论文链接：https://papers.miccai.org/miccai-2024/paper/0195_paper.pdf

论文创新点

本文提出了一种新的无监督领域适应（UDA）框架，其创新点主要包括：

1. 提出了一种耦合的结构保持扩散模型：作者设计的模型能够通过两个双向逐步图像投影序列生成互补图像，并保留语义信息。这种模型能够推导出从原始图像到最终生成图像的所有中间图像，提供了更有效的域知识。
2. 开发了随机步骤域对齐策略：这一策略通过多级生成对抗学习，针对整个生成过程中的数据进行领域差异的减少。特别是，通过利用采样的中间图像，这些图像包含了丰富的源域和目标域之间的转移知识，从而提高了跨模态分割模型的适应能力。
3. 实现了特征和预测空间的多级对抗适应：作者采用了多级生成对抗学习方法来对齐扩散模型整个生成过程中生成的数据。这种方法不仅对齐了特征分布，还在预测空间上实施了一致性约束，进一步提升了模型的泛化能力。
4. 在腹部多器官分割上验证了方法的有效性：通过在腹部多器官分割任务上的实验，作者证明了所提出方法的有效性，并在多个评估指标上达到了最先进的性能，展示了模型在处理领域偏移问题时的优势。

摘要

本研究的目的是利用从源域到目标类似域的中间图像分布来改进无监督领域适应（UDA），在图像生成过程中进行这一操作。然而，像生成对抗网络（GANs）这样的图像生成器由于其复杂的内部工作机制，可以被视为“黑箱”，我们只能访问最终生成的图像。这一限制使得它们无法在执行领域对齐时使用生成过程中产生的中间分布的知识。为了解决这个问题，作者提出了一个新的UDA框架，利用扩散模型来捕获和转移一定量的域间知识，从而减轻领域偏移问题。设计了一个耦合的结构保持扩散模型，通过多步合成中间图像，使中间图像分布可访问。进一步开发了随机步骤对齐策略来对齐特征分布，从而提高了适应能力。通过在腹部多器官分割上的实验，证明了所提出方法的有效性。

关键词

无监督领域适应 · 扩散模型 · 跨模态分割

方法

在UDA场景中，我们提供了来自两个不同域的数据：带有相应标签y ∈ Y的源数据xsrc ∈ Xs，以及未标记的目标数据xtgt ∈ Xt。我们的目标是学习一个模型，该模型可以在目标数据上表现良好。我们提出的方法的整体框架如图1所示。它使用两个未配对的图像作为输入，旨在将未标记的目标域的分布适应到标记的源域。该框架由三个关键组件组成。首先，引入耦合的结构保持扩散模型作为图像生成器，将图像从源域合成到目标域，反之亦然。其次，我们使用一个分割网络来提取两个原始图像、两个相关的步随机生成图像和源域的最终生成图像的特征，从而总共获得五个图像的特征。然后我们使用两个原始图像和最终生成图像的特征来预测分割结果。最后，我们应用生成对抗学习在特征空间和预测空间上执行领域适应。特别是，步随机生成图像的特征被用来对齐整个生成过程中的分布，以提高跨模态分割模型的适应能力。

结构保持扩散模型用于图像合成

与基于GAN的图像生成器不同，基于扩散的生成器逐步合成图像，从而推导出从原始图像到最终生成图像的所有中间图像。为了解释结构保持扩散模型的构建，我们以源域为例。在训练期间，用于源域的生成扩散模型是用目标数据训练的。在采样期间，提供源图像作为参考图像，扩散模型将其逐步投影到目标域。因此，我们获得了一系列的生成图像，这些中间生成的图像包含了源域和目标域之间的大量分布知识。因此，对于基于生成方法的UDA任务，扩散模型更适合作为生成器。

在这项工作中，我们基于去噪扩散概率模型（DDPM）[9]构建了扩散模型。DDPM是一类潜变量模型，从分布x0 ∼ q(x0)中采样数据点x0开始。扩散的前向过程可以定义为一个马尔可夫链，我们在T步中逐渐向样本x0添加少量高斯噪声：

这里 , 其中 t=1 是方差计划。当时，xT 是一个各向同性的高斯分布。

对于逆过程，由于不能轻易估计，因此学习了一个深度网络pθ来近似条件概率。因此，给定，逆过程被制定为具有学习均值和固定方差的马尔可夫链：

设 和 , 然后

扩散模型的简化目标可以写成：

尽管DDPM能够从源域合成与目标域外观相匹配的图像，但观察到采样图像的结构被扭曲，导致其语义内容发生变化。这对于UDA任务是不可取的，因为跨模态分割的性能可能会受到不利影响。此外，直接使用采样图像作为分割网络的输入将导致由于噪声的存在而不稳定的训练。

为了在保持原始域内容的同时将图像从一个域投影到另一个域，受[5]的启发，我们引入了迭代潜在细化过程来指导图像的结构。具体来说，我们采用了线性低通过滤操作ϕN(.)和一系列N倍的下采样和上采样来捕获图像的结构信息。通过将DDPM的前向过程中的图像序列表示为（x0, x1, …, xT），将DDPM的逆过程中的图像序列表示为（x ′ T , x ′ T −1, …, x ′ 0），^ 是基于xt和x′t细化的样本。我们使采样图像^ 在每s步的间隔中参考图像xt的结构。最终的逆过程可以定义为：

然后，为了使分割稳定，我们遵循Tweedie公式[13]来获得干净的图像：

随机步骤领域对齐

借助扩散模型，我们可以双向投影图像，这允许我们在特征和预测空间中传输图像的外观。然而，由于UDA中存在显著的领域差距，模型的适应能力仍然需要改进。为了解决这个问题，我们采用了多级生成对抗学习方法来对齐扩散模型整个生成过程中生成的数据。特别是，为了充分利用采样的中间图像，它包含了源域和目标域之间的丰富转移知识，我们提出了随机步骤领域对齐策略，用于多步扩散模型。

对于扩散模型的生成过程的开始，从两个原始图像和开始，我们使用分割网络分别提取它们的特征 和 。随后，我们引入了一个鉴别器来对齐和的特征分布。这种对齐的目标是将未标记的目标数据的分布带到标记的源数据的分布附近。为了实现这种对齐，我们在分割网络的训练中最小化对抗损失。这个损失迫使主干网络提供与源数据分布接近的特征以欺骗鉴别器，而鉴别器则尝试使用交叉熵来分类不同域的数据。损失可以表示为：

进一步地，我们观察到尽管源数据和目标数据的外观存在显著差异，但预测的结构，即，在腹部多器官分割中是一致的。因此，我们执行相同的生成对抗策略，使用鉴别器在预测空间上实施一致性约束。

对于扩散模型采样的中间图像，即去噪图像序列 和 ,，我们提出了随机步骤领域对齐策略，以建立多级对抗适应。我们随机选择一个t从均匀分布中为分割网络训练的每次迭代。然后我们将选定的步随机图像T放入分割网络中，以获得相应的特征。以适应源域为例，我们希望分割网络为来自目标图像的生成中间图像产生接近源域的特征分布，以欺骗鉴别器。

对抗和判别损失可以表示为：

相应地，另一个鉴别器旨在尽可能区分和。这些对齐可以使网络产生的特征分布尽可能接近。因此，尽管我们只能访问源数据的注释，网络仍然可以在目标数据上表现良好。

对于最终生成的去噪图像，我们让它与原始图像共享相同的注释，以计算分割损失，它由交叉熵和广义Dice损失组成。因此，分割网络被迫在预测空间中保持语义一致性。最终的分割损失由以下公式计算：

最后，我们将腹部多器官分割和对抗学习整合到UDA任务的统一框架中。整体目标函数定义为所有先前定义的损失函数的加权求和：

3 实验

公众号后台回复“数据集”获取100+深度学习各方向资源整理