Diff-UNet：一种用于体积分割的扩散嵌入网络

摘要

近年来，降噪扩散模型在生成语义上有价值的逐像素表示方面取得了显著成功，用于图像生成建模。 在本研究中，我们提出了一种新颖的端到端框架，称为 Diff-UNet，用于医学体积分割。 我们的方法将扩散模型集成到标准 U 形架构中，以有效地从输入体积中提取语义信息，从而为医学体积分割生成出色的像素级表示。 为了增强扩散模型预测结果的鲁棒性，我们还在推理过程中引入了基于步长不确定性的融合 (SUF) 模块，以结合每个步骤的扩散模型输出。 我们在三个数据集上评估了我们的方法，包括 MRI 中的多模态脑肿瘤、肝肿瘤和多器官 CT 体积，并证明了 Diff-UNet 显著优于其他最先进的方法。 我们的实验结果也表明了所提模型的通用性和有效性。 提出的框架有可能通过实现更精确的解剖结构分割，促进对医疗状况的准确诊断和治疗。 Diff-UNet 的代码可在 https://github.com/ge-xing/Diff-UNet 获得。

1 引言

医学体积分割是医学图像分析中的一个关键任务 [12, 9, 20]，涉及在像素级识别高维医学图像数据集中的病灶区域。 更准确的分割结果可以为医生提供有价值的信息，帮助他们诊断疾病。 传统的 3D 医学分割算法通常采用编码器-解码器结构 [19, 15, 3]，并结合跳跃连接以使解码器能够重用编码器提取的特征。 目前许多针对模型结构设计的 3D 医学图像分割算法都取得了令人鼓舞的分割结果。 例如，SegResNet [16] 使用变分自动编码器 [10] 添加重建分支，从而提高模型的特征提取能力。 然而，由于其结构基于卷积神经网络，它可能无法有效地提取全局特征。

最近，Transformer 结构由于其全局自注意力机制 [23, 21] 在建模全局特征方面越来越受欢迎。 TransBTS [24] 利用 3D-CNN 提取局部空间特征，然后应用 Transformer 来对高级特征中的全局依赖关系进行建模。 UNETR [7] 利用 ViT [4] 作为编码器直接对全局特征进行建模，并使用基于 CNN 的解码器和跳跃连接输出分割结果。 然而，上述方法由于 Transformer 结构的计算复杂度，在提取多尺度特征方面的能力有限。 SwinUNETR [6] 利用 Swin-Transformer [13] 作为编码器提取多尺度特征，并采用基于 CNN 的解码器生成输出，从而实现了最先进的医学图像分割结果。

去噪扩散模型 [8, 22, 17] 在各种生成任务中取得了显著的成功，包括医学图像分割。 例如，MedSegDiff [26] 通过对去噪 UNet 进行分割并通过傅立叶变换交互结构间信息，实现了 2D 医学图像分割。 Wolleb 等人 [25] 采用扩散模型来解决二维医学图像分割问题，并在测试期间通过求和的方式融合每个扩散步骤的输出结果，从而提高了分割结果的鲁棒性。 然而，这些方法仅限于二维分割，并且扩散模型无法直接生成多标签分割。

与传统的分割方法相比，扩散模型在输入中引入噪声并迭代地预测分割标签图，这两者都可以提高扩散模型预测的鲁棒性。 为了利用扩散模型的潜力，我们提出了一种通用的基于扩散的端到端三维医学图像分割算法，称为 Diff-UNet，用于解决高维医学图像分割问题。 然而，传统的扩散模型只能解决二元分割问题。 为了分割多个类别，我们设计了一种标签嵌入操作，将分割标签图转换为独热标签。 这使 Diff-UNet 能够同时分割多个目标。 为了从输入体积中提取语义信息，我们设计了一个降噪模块，该模块包含一个降噪 UNet 和一个独立的特征编码器来学习降噪过程。 该模块从噪声标签图输出清晰的分割标签图。 最后，我们设计了一个基于步骤不确定性的融合 (SUF) 模块，该模块在测试阶段融合来自降噪模块的多个预测，以获得更稳健的分割结果。 在 BraTS2020 多模式脑肿瘤分割数据集 [14, 2]、BTCV 多器官分割数据集 [11] 和 MSD 肝脏和肝脏肿瘤分割数据集 [1] 上进行的大量实验表明，我们的方法明显优于最先进的方法。

2 方法

图 1: 提出的 Diff-UNet 的概述。 (A) 是 Diff-UNet 的训练阶段，通过降噪模块学习降噪函数。 (B) 是通过迭代方式生成分割结果的测试阶段。 (C) 是我们 SUF 模块的计算过程。

图 1 显示了训练阶段、测试阶段以及我们提出的 Diff-UNet 的两个子模块：降噪模块和基于步长不确定性的融合模块 (SUF)。 与直接将体积数据输入以预测相应分割标签图的传统医学图像分割方法不同，扩散模型学习降噪过程。 扩散模型以体积图像和带有噪声的分割标签图作为输入，并学习去除噪声以生成清晰的分割结果。

2.1 标签嵌入

通常使用独热编码将多分类标签转换为多个二分类标签。 独热向量 $v$ 是一个长度为 $c$ 的二进制向量，其中只有一个条目可以为 1，所有其他条目必须为 0。 例如，假设有 3 个分割目标，我们将分割标签 $(0,1,2)$ 通过独热编码转换为 $((0,0,1),(0,1,0),(1,0,0))$。 传统的扩散模型只能生成连续数据，无法预测多目标标签。 因此，我们首先将大小为 $D\times W\times H$ 的单通道标签图转换为多通道标签：$x_0\in {ℝ}^{N\times D\times W\times H}$，通过独热编码，其中 $N$ 是标签数量，$(D,W,H)$ 是体积医学图像的空间分辨率。 然后，我们为转换后的多通道标签添加连续的 $t$ 步噪声 $ϵ$，称为扩散前向过程。

$${𝐱}_t=\sqrt{{\overline{\alpha }}_t}{𝐱}_0+\sqrt{1-{\overline{\alpha }}_t}\mathit{ϵ}.$$ (1)

在获得带有 $t$ 步噪声的标签图 $x_t$ 之后，我们的目标是根据 $x_t$ 和原始体积数据通过降噪模块预测清晰的标签图 $x_0$。

2.2 降噪模块

如图 1 (A) 所示，降噪模块包含特征编码器 (FE) 和降噪-UNet (DU)，是 Diff-UNet 的主要部分。 降噪-UNet 也包含两个部分，一个编码器和一个解码器。 首先，给定体积数据 $I\in {ℝ}^{M\times D\times W\times H}$，其中 $M$ 是模态图像的数量，$I$ 和噪声 one-hot 标签 $x_t$ 按通道连接到 DU 的编码器以获得多尺度特征 ${\widehat{I}}_f:{[{ℝ}^{if\times \frac{D}{i}\times \frac{W}{i}\times \frac{H}{i}}]}_{i=1}^{16}$，其中 $f$ 是特征大小，$i$ 是尺度。 同时，为了更好地引入原始体积图像特征，我们通过特征编码器提取体积数据的多尺度特征 ${\tilde{I}}_f$，该编码器与 DU 的编码器具有相同的大小。 由于 ${\tilde{I}}_f$ 和 ${\widehat{I}}_f$ 包含相同数量和大小的特征，我们对对应尺度的特征进行求和以获得融合特征。 之后，我们将融合的多尺度特征输入到 DU 网络的解码器，以获得预测结果 $\widehat{x_0}\in {ℝ}^{N\times D\times W\times H}$：

$$\widehat{x_0}=\mathrm{DU}(\mathrm{cat}(I,x_t),t,{\tilde{I}}_f).$$ (2)

经典扩散使用 ${ℒ}_2$ 去噪损失进行训练。 在此任务中，我们将医学图像分割任务建模为离散数据生成问题，并直接预测 $x_0$ 而不是噪声 $ϵ$。Diff-UNet 通过结合 Dice 损失、BCE 损失和 MSE 损失进行训练，因此我们 Diff-UNet 的总损失 ${ℒ}_{total}$ 为：

$${ℒ}_{total}={ℒ}_{dice}(\widehat{x_0},x_0)+{ℒ}_{bce}(\widehat{x_0},x_0)+{ℒ}_{mse}(\widehat{x_0},x_0).$$ (3)

2.3 基于步骤不确定性的融合

扩散模型在测试阶段通过 Denoising Diffusion Implicit Models (DDIM) 方法迭代 $t$ 次。 在传统的生成任务中，最后一个预测被作为最终的生成结果，而 Diff-UNet 的每次迭代都会生成一个分割图。 随着预测时间步长的增加，预测结果越准确，预测不确定性越低。 因此，为了提高 Diff-UNet 模型的分割鲁棒性，我们根据预测步骤的数量和不确定性来融合输出。

我们计算不确定性的方式类似于 Monte Carlo Dropout (MC Dropout) [5]，它激活网络的 dropout 层，然后执行 $S$ 次前向传递以估计不确定性图。 另一方面，Diff-UNet 在测试阶段初始化一个随机噪声 $x_t$（图 1（B）），因此它可以在不激活 dropout 层的情况下将随机性引入网络。 与 Monte Carlo dropout 一样，Diffusion 的测试过程包括 $t$ 步，每一步预测 $S$ 个输出，这些输出用于计算不确定性。 公式如下：

$$u_i=-{\overline{p}}_{i}log({\overline{p}}_i),\text{where}{\overline{p}}_i=\frac{1}{S}\sum _{s=1}^S{p}_i^s.$$ (4)

结合预测步骤数量和不确定性的融合权重计算为 $w_i=e^{\sigma (\frac{i}{scale})\times (1-u_i)}$，其中 $\sigma$ 是 sigmoid 函数，$i$ 表示当前预测步骤，$u$ 是不确定性矩阵。 我们使用 $w$ 来加权每个步骤的预测结果，以获得最终的融合结果 $Y$，该结果被用作我们网络的输出。 最后，$Y$ 由以下给出：$Y={\sum }_{i=1}^t{w}_i\times {\overline{p}}_i$。

3 实验

实现细节。   我们的网络在 4 台 $\times$ NVIDIA A100 GPU 上使用 Pytorch 和 MONAI 实现。 在训练阶段，损失函数结合了 DICE 损失、BCE 损失和 MSE 损失。 我们采用 AdamW 优化器，权重衰减为 10-5。 预热设置为总 epochs 数的 1/10，学习率使用余弦退火调度进行更新。 每次迭代随机抽取 n 个 patch（patch 大小为 96$\times$96$\times$96）进行训练。 引入随机翻转、旋转、强度缩放和平移以进行数据增强。 在测试中，我们将 DDIM 抽样步骤数设置为 10，每个样本的大小为 96$\times$96$\times$96。 滑动窗口重叠率为 0.5，直到整个体积被预测。

3.1 数据集和评估指标

为了评估我们方法的三维分割性能，我们利用了三个公开可用的分割数据集，包括 BraTS2020 [14, 2]、MSD Liver [1] 数据集和腹部多器官分割数据集 BTCV [11]。 此外，采用 Dice 分数和 95% Hausdorff 距离 (HD95) 进行定量比较。

BraTS2020 数据集 包含 369 个对齐的四模态 MRI 数据（即 T1、T1ce、T2、FLAIR），以及专家分割掩码（即 GD 增强肿瘤、肿瘤周围水肿和肿瘤核心）。 每个模态都有 155×240×240 的体积，所有模态图像都已重新采样并配准。 分割任务的目标是分割整个肿瘤 (WT)、增强肿瘤 (ET) 和肿瘤核心 (TC) 区域。 训练集、验证集和测试集的拆分比例分别为 0.7、0.1 和 0.2。

MSD Liver 数据集 共有 131 例 3D 肝脏图像，每个 3D 肝脏图像有 1 个模态和 2 个分割目标 (肝脏和肝脏肿瘤)。 所有数据都重新采样到相同的空间 (2.0, 2.0, 2.0)。 MSD 肝脏数据集根据 0.7、0.1 和 0.2 的比例划分为训练集、验证集和测试集。

BTCV 数据集 包含 30 例 3D 腹部多器官图像，每个 3D 图像有 13 个器官分割目标。 所有数据都重新采样到相同的空间 (2.0, 1.5, 1.5)。 遵循 TransUNet，18 例用于训练，其余 12 例用于测试。

3.2 与 SOTA 方法的比较

表 1: BraTS2020 数据集上的定量比较。

Methods	WT			TC			ET				Average
	Dice$\uparrow$	HD95$\downarrow$		Dice$\uparrow$	HD95$\downarrow$		Dice$\uparrow$	HD95$\downarrow$			Dice$\uparrow$	HD95$\downarrow$
SwinUNETR [6]	91.68	2.856		82.60	4.314		74.85	4.503			83.04	3.891
UNETR [7]	90.15	4.305		81.26	5.740		73.23	4.643			81.55	4.896
TransBTS [24]	91.06	3.360		83.60	2.986		74.03	3.403			82.90	3.249
SegResNet [16]	91.54	3.2275		83.61	3.769		73.04	3.486			82.73	3.494
Attention-UNet [18]	84.49	15.174		78.17	16.380		71.62	9.095			78.09	13.549
ModelsGenesis [27]	91.98	2.799		84.31	2.836		73.84	4.333			83.38	3.096
Our Diff-UNet	92.23	2.588		86.94	3.596		76.87	3.984			85.35	3.389

表 2: MSD 肝脏数据集上的定量比较。

Methods	Liver			Tumor			Average
	Dice$\uparrow$	HD95$\downarrow$		Dice$\uparrow$	HD95$\downarrow$		Dice$\uparrow$	HD95$\downarrow$
SwinUNETR [6]	95.47	0.392		49.94	20.906		72.70	10.645
UNETR [7]	93.75	1.080		38.43	24.87		66.09	12.979
TransBTS [24]	95.11	0.403		44.99	17.463		70.05	8.933
SegResNet [16]	95.30	0.418		46.39	19.424		70.85	9.921
Attention-UNet [18]	95.32	0.499		48.43	20.273		71.88	10.386
ModelsGenesis [27]	95.04	0.934		50.04	31.146		72.54	15.823
Our Diff-UNet	95.72	0.222		51.65	17.280		73.69	8.751

表 3: BTCV 多器官分割数据集上的定量比较。

Framework				Average		Aorta	Gallbladder	Kidney(L)	Kidney(R)	Liver	Pancreas	Spleen	Stomach
Encoder	Decoder			Dice$\uparrow$	HD95$\downarrow$
VNet				68.81	-	75.34	51.87	77.10	80.75	87.84	40.05	80.56	56.98
DARR				69.77	-	74.74	53.77	72.31	73.24	94.08	54.18	89.90	45.96
R50	U-Net			74.68	36.87	84.18	62.84	79.19	71.29	93.35	48.23	84.41	73.92
R50	AttUNet			75.57	36.97	55.92	63.91	79.20	71.71	93.56	49.37	87.19	74.95
ViT	None			61.50	39.61	44.38	39.59	67.46	62.94	89.21	43.14	75.45	69.78
ViT	CUP			67.86	36.11	70.19	45.10	74.70	67.40	91.32	42.00	81.75	70.44
R50-Vit	CUP			71.29	32.87	73.73	55.13	75.80	72.20	91.51	45.99	81.99	73.95
TransUNet				77.48	31.69	87.23	63.13	81.87	77.02	94.08	55.86	85.02	75.62
Our Diff-UNet				83.75	8.115	89.30	76.23	85.20	84.73	95.90	74.25	89.75	74.65

对于 BraTS2020 和 MSD 肝脏数据集，我们将我们的 Diff-UNet 与最先进的分割方法进行比较，包括 SwinUNETR、UNETR、TransBTS、SegResNet、Attention-UNet 和 ModelsGenesis。 为了公平比较，所有方法都使用公开可用的实现。 对于 BTCV 数据集，我们遵循 TransUNet 的相同实验设置，利用相同的训练和测试数据集，并与最先进的方法进行比较。

BraTS2020。 表 1 报告了所有方法在 BraTS2020 数据集上三个区域（WT、TC、ET）的 Dice 和 HD95 分数以及平均分数。 显然，我们提出的 Diff-UNet 方法在所有三个区域和它们的平均 Dice 分数方面明显优于比较的最先进方法。 三个区域的平均 Dice 为 85.35%，比第二名 ModelsGenesis 提高了 1.97%。 虽然我们在三个区域的平均 HD95 分数达到 3.3898，排名第三，但它略小于前两个 HD95 结果（即 3.2499 和 3.0961）。

默沙东肝脏。 表 2 报告了我们提出的 Diff-UNet 和最先进方法在 MSD 肝脏数据集上的 Dice 和 HD95 性能。 与其他比较方法相比，我们的 Diff-UNet 方法在肝脏区域、肿瘤区域及其平均值方面具有更高的 Dice 分数和更低的 HD95 分数。 具体来说，我们的 Diff-UNet 在肝脏分割方面实现了 95.72% 的 Dice 分数和 0.222 的 HD95 分数，在肝脏肿瘤分割方面实现了 51.65% 的 Dice 分数和 17.280 的 HD95 分数。 两个区域的平均 Dice 和 HD95 分数分别为 73.69% 和 8.751。

BTCV。 遵循 TransUNet 的相同实验设置，我们在表 3 中报告了八个腹部器官的 Dice 分数以及平均 Dice 和 HD 分数，以比较我们的 Diff-UNet 和最先进的分割方法。 从表 3 中，我们可以发现我们的 Diff-UNet 在八个器官上具有最佳的平均 Dice 和 HD95 分数。 具体来说，我们的 Diff-UNet 在六个器官的 Dice 得分上排名第一，平均 Dice 和 HD95 得分分别为 83.75% 和 8.115。 这表明我们的 Diff-UNet 在 BTCV 数据集上比最先进的方法能够实现更准确的多器官分割性能。   虽然我们的 Diff-UNet 在脾脏上的 Dice 排名第二，在胃上的 Dice 排名第三，但它们的 Dice 分数（89.75% 和 74.65%）略低于最佳分数，脾脏的最佳分数为 89.90%，胃的最佳分数为 75.62%。

图 2： 我们在 BraTS2020 和 MSD 肝脏数据集上对我们网络和最先进方法生成的分割结果进行了视觉比较。 显然，我们的方法具有更准确的分割性能，并且与地面实况（标记为“GT”）一致。

视觉比较。 图 2 视觉比较了我们的 Diff-UNet 和 SOTA 方法在 BraTS2020 和 MSD 肝脏数据集上生成的分割结果。 我们没有在 BTCV 数据集上运行比较方法，而是直接使用了 TransUNet 中的结果表，因此我们没有显示 BTCV 数据集上的分割结果。 我们发现我们的 Diff-UNet 实现了更准确的分割结果，特别是在微小目标上，而比较方法往往会遗漏一些目标区域，或者在其分割结果中包含其他非目标区域（参见第二行中的 SwinUNETR）。

3.3 消融研究

Module		WT	TC	ET	Average
basic		91.62	85.02	75.10	83.91
basic+FE		91.52	85.85	75.59	84.32
basic+FE+SF		92.02	86.58	75.67	84.76
basic+FE+$\mathrm{SUF}$ (Ours)		92.23	86.94	76.87	85.35

表 4： BraTS2020 数据集上不同模块的消融研究。 FE 表示单独的特征编码器。 SF 表示简单融合。 SUF 表示基于步长不确定性的融合模块。

$S$		WT	TC	ET	Average
3		92.19	86.18	76.82	85.06
4 (Ours)		92.23	86.94	76.87	85.35
5		92.17	86.96	76.84	85.32
6		92.22	86.92	76.84	85.33

表 5： 对 BraTS2020 数据集的每个 DDIM 步中计算不确定性的预测数量 (S) 的消融研究。

主要模块的有效性。 我们在 BraTS2020 数据集上进行消融实验，以评估 Diff-UNet 中涉及的不同主要模块（即 FE 和 SUF）的作用，并在表 5 中显示其定量结果。 从表 5 中的定量结果可以看出，“basic+FE” 在三个区域（即 WT、TC 和 ET）的平均 Dice 值都大于“basic”，这表明将我们的 FE 作为图像编码器可以将更多图像信息引入扩散模型，从而提高分割精度。 同时，“basic+FE+SF” 比“basic+FE” 的 Dice 分数更高，这表明融合 DDIM 每个步骤预测的分割结果可以进一步提高扩散模型的分割精度。 此外，我们的方法比“basic+FE+SF” 的 Dice 分数更高，这表明为不同 DDIM 步骤的预测分配不同的权重可以进一步提高我们方法的分割性能。

设置 S。 此外，我们进行了一个消融研究实验来讨论如何设置 $S$ 的值（见式 4），它是每个 DDIM 步中计算不确定性的预测数量。 在这里，我们考虑 $S$ 的不同值，它们分别是 3、4、5 和 6，并在表 5 中显示了相应的结果。 显然，当 $S=4$ 时，我们的方法在三个区域上的平均 Dice 分数最高，它在 WT 上的 Dice 分数最高为 92.23%，在 TC 上的 Dice 分数第二高为 76.87%，在 ET 上的 Dice 分数最高为 85.35%。 因此，我们在我们的方法中经验性地设置了 $S=4$。

4 结论

在本文中，我们提出了第一种基于扩散模型的 3D 医学图像分割方法，名为 Diff-UNet，它将医学图像分割建模为离散数据生成任务。 所提出的算法引入了一种通用的端到端 3D 医学图像分割方法，利用扩散模型的优势来提高分割鲁棒性。 在不同基准数据集上的实验结果表明，我们的 Diff-UNet 比最先进的方法更优。 总体而言，我们的工作对医学图像分割领域做出了重大贡献，证明了扩散模型在 3D 医学图像分割任务中的有效性。 该方法有可能促进更精确和准确的医疗状况诊断和治疗，最终导致改善患者预后。