ImgX-DiffSeg：基于 DDPMs 的 3D 医学图像分割
ImgX-DiffSeg：基于 DDPMs 的 3D 医学图像分割
目录
- Abstract
- Current Limitations
- ImgX-DiffSeg Architecture
-
Overview
-
DDPM with Variance Schedule Resampling
-
Diffusion Model for Segmentation
-
实验
-
总结
-
参考
-
本文首发于 GiantPandaCV，未经允许不得转载！！
前言
本文继续深入探讨医学图像处理中的diffusion技术。在之前的讨论中，我们分别探讨了diffusion技术在自监督和有监督分割任务中的应用。链接：
https://mp.weixin.qq.com/s/7g3_4hHfKCAp2WQibkbzgA
https://mp.weixin.qq.com/s/TR6TaTGAdzQZNPUgSRARrQ
本文并非单纯的扩散模型应用，而是对训练与推理策略进行了优化，以3D医学图像分割任务为目标，参考链接在文末。
目前存在的问题
- 目前带有 diffusion model 的架构训练和推理耗时。
- 在一些分割任务中，并不确定 diffusion model 预测噪声推断分割图和直接预测分割图哪个效果更好。
- 模型过度依赖先前时间步中的信息。
ImgX-DiffSeg 架构
概述
DDPM 是一种生成式模型，主要用于图像降噪和分割。其工作原理是模拟干净图像的概率分布，随后在图像中逐步添加不同级别的噪声，生成噪声版本。相反，模型通过去除添加的噪声来实现图像降噪。在图像分割任务中，模型会生成分割掩码，根据输入图像的特征，将图像划分为多个区域。更详细的信息，请参考前置文章（强烈推荐）。
对于ImgX-DiffSeg的整体架构流程图如图所示。该架构直接预测分割掩码而非采样噪声，并采用Dice Loss作为优化目标。这表明，与传统方法不同，ImgX-DiffSeg可以直接预测分割图，无需先生成噪声并依赖其进行推断。随后，从上一个时间步中恢复的预测掩码被用于生成经过噪声污染的掩码。这有助于减少模型对前一时间步信息的过度依赖，从而避免可能出现的问题。最后，该架构将训练所需的扩散步骤减少至五个，与推理过程所需的步骤一致。
DDPM with Variance Schedule Resampling
训练过程的公式和 DDPM 基本是保持一致的，下面的公式分别表示反向过程中预测噪声和原图：
\begin{aligned} & L_{\text {simple }, \epsilon_l}(\theta)=\mathbb{E}_{t, \mathbf{x}_0, \epsilon_t}\left\|\epsilon_t-\epsilon_{t, \theta}\left(\mathbf{x}_t\left(\mathbf{x}_0, \epsilon_t\right), t\right)\right\|_2^2,\left(\text { Predict } \epsilon_t\right) \\ & L_{\text {simple }, \mathbf{x}_0}(\theta)=\mathbb{E}_{t, \mathbf{x}_0, \epsilon_t}\left\|\mathbf{x}_0-\mathbf{x}_{0, \theta}\left(\mathbf{x}_t\left(\mathbf{x}_0, \epsilon_t\right), t\right)\right\|_2^2 .\left(\text { Predict } \mathbf{x}_0\right) \end{aligned}
推理过程中，DDPM 中的生成过程从正常噪声开始，由变量 xT 表示。该初始噪声是从平均值为 0 且方差为 1 的正态分布中采样的。在生成过程的每个步骤中，使用预测的平均值 µ 对变量 xtk-1 进行采样。下标 k-1 表示上一个时间步。这意味着每步 x 的值取决于上一步中 x 的值以及分布的预测平均值。
\begin{aligned} p_\theta\left(\mathrm{x}_{t-1} \mid \mathrm{x}_t\right) & =\mathcal{N}\left(\mathrm{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\mathrm{x}_t, t\right), \sigma_t^2 \mathrm{I}\right), \\ \mu_\theta\left(\mathrm{x}_t, t\right) & \left.=\frac{1}{\sqrt{\alpha}}\left(\mathrm{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_{t, \theta}\left(\mathrm{x}_t, t\right)\right), \text { (Predict } \epsilon_t\right) \\ \mu_\theta\left(\mathrm{x}_t, t\right) & =\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1-\bar{\alpha}_t} \mathrm{x}_{0, \theta}\left(\mathrm{x}_t, t\right)+\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t} \sqrt{\alpha_t} \mathrm{x}_t .\left(\text { Predict } \mathrm{x}_0\right) \end{aligned}
上面这些流程均和 DDPM 相似，我们就不展开说明了。重点关注 Variance Schedule Resampling 的过程，也就是如何实现训练的扩散过程减少到五个步骤的。可以理解为对方差值子序列进行采样的过程。给出了方差表 {βt} Tt=1，子序列 {βk} Kk=1 可以用 {tk} Kk=1 进行采样。简而言之，在训练或推理过程中，给出方差值序列，并对这些值的子序列进行采样。子序列中的值是根据先前的值和重新计算的值计算的。目标是通过在训练或推理期间调整方差值来优化模型的性能。如果是在图像去噪任务中进行方差的重采样，一定会影响结果，但在分割任务中经过验证是有效的。
Diffusion Model for Segmentation
上一部分专门处理 DDPM 的方差重采样问题，而未涉及图像分割。在分割任务的优化方面，ImgX-DiffSeg 通过时间步的分析，采用预测噪声与采样噪声的 L2 损失进行训练。此外，ImgX-DiffSeg 进一步计算了预测掩码与金标准之间的特定分割损失，如Dice Loss或CE Loss。
损失函数 L，其基于 \mathrm{scg} 和初始样本 \mathbf{x}_0，参数为 \theta，等于对时间步 t、初始样本 \mathbf{x}_0、噪声 \epsilon_l 和条件变量 I 的联合分布取期望，得到的该段落损失函数 \mathcal{L}_{\mathrm{seg}} 的输出。该段落损失函数 \mathcal{L}_{\mathrm{seg}}，其输入为原始样本 \mathrm{x}_0 和通过参数 \theta 生成的样本 \mathrm{x}_{0, \theta}(\mathrm{x}_t, t, I)。
在训练阶段，现有方法通过插值噪声生成器和金标准对噪声掩模进行采样，这可能导致数据信息丢失。为了解决这一问题，模型利用上一个时间步的预测来替代金标准。回收的噪声掩码通过以下方程进行计算。
\begin{aligned} \mathbf{x}_t & =\sqrt{\bar{\alpha}_t} \mathbf{x}_{0,\theta}\left(\mathbf{x}_{t+1}, t+1, I\right)+\sqrt{1-\bar{\alpha}_t} \epsilon_t \\ \mathbf{x}_{t+1} & =\sqrt{\bar{\alpha}_{t+1}} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_{t+1}} \epsilon_{t+1} \end{aligned}
其中，x0θ 是利用金标准计算的上一个时间步的预测分割掩码，xt 和 xt+1 是两个独立的采样噪声生成器。梯度停止应用于 xt+1，以防止通过回收的噪声掩码进行反向传播。αt 是超参数。第一个方程基于先前的预测和当前噪声计算 xt，而第二个方程基于金标准和下一个噪声计算 xt+1。
实验
本研究采用MRI和CT图像数据集作为实验基础，值得注意的是，本研究中的ImgX-DiffSeg在三维空间中的性能显著优于其二维数据集版本，具体结果可见表1。
该图用于比较非扩散分割模型与扩散概率模型之间的差异关系。其中，t代表时间步序列中的一个阶段，模型通过持续反向传播至初始时间步的效果最为显著。
下表列出了四种消融实验，包括预测噪声推断分割图与直接预测分割图的对比、损失函数的对比、是否在上一个时间步中回收预测的掩码的对比，以及训练过程的时间步数量的对比。实验结果表明，5 个 steps 的效果优于 1000 个 steps，这表明 Variance Schedule Resampling 具有显著的效果。
总结
ImgX-DiffSeg 是首个专为 3D 体积多分类分割设计的 DDPM 模型。相较于现有扩散基线方法，该模型在性能上实现了显著提升。然而，其在性能上也未能超越传统非扩散分割方法，仍需进一步优化以提升适用性。