【Diffusion分割】HiDiff：医学图像分割混合扩散模型

HiDiff: Hybrid Diffusion Framework for Medical Image Segmentation

摘要：

        随着深度学习（DL）技术的快速发展，医学影像分割技术得到了长足的进步。现有的基于深度学习的分割模型通常是判别性的，即它们旨在学习从输入图像到分割掩膜的映射。然而，这些判别方法忽视了底层数据分布和内在类别特征，导致特征空间不稳定。在这项工作中，我们建议利用生成模型中的底层数据分布知识来补充判别式分割方法。为此，我们提出了一种用于医学影像分割的新型混合扩散框架，称为 HiDiff，它可以协同现有的判别分割模型和新的生成扩散模型的优势。HiDiff 包括两个关键部分：判别分割器和扩散细化器。首先，我们利用任何传统的训练有素的分割模型作为判别分割器，为扩散细化器提供分割掩码先验。其次，我们提出了一种新颖的二进制伯努利扩散模型（BBDM）作为扩散细化器，它可以通过对底层数据分布建模，有效、高效、交互式地细化分割掩码。第三，我们以交替协作的方式训练分割器和 BBDM，使其相互促进。在腹部器官、脑肿瘤、息肉和视网膜血管分割数据集（涵盖四种广泛使用的模式）上的大量实验结果表明，HiDiff 的性能优于现有的医疗分割算法，包括最先进的基于变换器和扩散的算法。此外，HiDiff 还擅长分割小物体，并能推广到新的数据集。 

I. INTRODUCTION

        近来，扩散概率模型（DPM）[18]、[22]、[23] 在各种图像生成任务[24]-[26]中取得了令人瞩目的成果，具有强大的（非）条件数据分布建模能力，这也为医学影像分割领域的探索提供了便利。一方面，人们可以直接应用扩散模型将图像分割作为生成任务[13]、[27]、[28]。另一方面，一些方法探索整合现有判别方法和生成方法的优势来解决分割任务。例如，PD-DDPM [19] 和 MedSegDiff-V2 [20] 尝试用现有的判别方法驯服 DPM 的力量，以提高分割性能。然而，大多数基于 DPM 的分割方法都依赖高斯噪声作为扩散核，忽略了分割任务固有的离散性。此外，DPM 的迭代去噪过程也使得扩散过程耗时较长。

        在这项工作中，我们探讨了如何有效、高效、互动地协同现有的判别式分割模型和生成式 DPM 的优势。为此，我们提出了一种用于医学图像分割的新型混合扩散框架，称为 HiDiff，如图 1 所示。拟议的 HiDiff 包括两个关键组件：判别分割器和扩散细化器。我们的想法是将现有分割器的判别能力与扩散细化器的生成能力相结合，以改进医学图像分割：判别分割器为扩散模型提供分割掩膜先验，然后扩散细化器有效、高效、交互式地细化分割掩膜。我们使用现有的医学图像分割模型作为判别分割器，并提出了一种新颖的二进制伯努利扩散模型（BBDM）作为扩散细化器。BBDM 的新颖之处在于三个方面：(i) 基于伯努利的扩散核增强了扩散模型对分割任务中离散目标的建模能力；(ii) 二值化扩散细化器以可忽略不计的计算成本显著提高了推理效率；(iii) 交叉变换器实现了扩散生成特征与判别特征之间的交互式交换。此外，我们还以交替协作的方式训练 HiDiff 的判别分割器和扩散细化器，以相互促进。

        -我们提出了一种用于医学图像分割的新型混合扩散框架（HiDiff），它可以协同现有判别分割模型和生成扩散模型的优势。

        -我们提出了一种新颖的二进制伯努利扩散模型（BBDM）作为扩散细化器，它可以通过对底层数据分布建模，有效、高效、交互式地细化分割掩码。

        -我们引入了一种交替协作训练策略来训练分辨式分割器和扩散式细化器，它们可以在训练过程中相互提高。

        这项工作的初步版本--伯努利扩散（Bernoulli Diffusion 或 BerDiff）--已发表在 2023 年医学影像计算和计算机辅助干预国际会议上[29]。在本文中，我们进一步扩展了 BerDiff [29]，并做了以下主要改进。首先，我们将 BerDiff [29] 扩展为用于医学影像分割的新型混合扩散框架，同时利用判别分割模型和生成扩散模型，并使用二值化扩散细化器来提高计算效率和可用性。其次，HiDiff 利用所提出的交替协作训练策略，促进判别式分割器和扩散细化器的同步训练，从而在训练过程中相互提高性能。第三，HiDiff 是一个多功能框架，可以与现有的判别式分割模型无缝集成，从而适用于各种医学影像分割任务。第四，我们扩展了 BerDiff，使其进一步兼容多目标分割任务。

II. METHODOLOGY 

        为了有效、高效、交互地协同现有判别式分割器和所提出的 BBDM 的优势，我们提出了一种用于医学图像分割的新型混合扩散框架，如图 2 所示。在本节中，我们首先在第二章 A 节和第二章 B 节中分别阐述了现有的判别式分割器和所提出的扩散细化器。最后，我们将在第 II-C 节中详细介绍采用交替协作训练策略的混合扩散框架。

A. Existing Discriminative Segmentors

        让 x∈R H×W 表示空间分辨率为 H×W 的输入医疗图像，y0∈{0, 1}H×W ×C 表示地面实况掩膜，其中 C 表示 H×W ×C 类别的数量。给定图像-掩码对（x，y0），现有的判别式分割方法通常采用深度神经网络 f (-) 来预测分割掩码的概率 f (x) ∈ RH×W ×C ，每个元素的范围为（0，1）。它们大多通过最小化交叉熵损失或骰子损失进行端到端训练。

        尽管这些判别实现通常表现令人满意，但它们仍然无法捕捉到潜在的数据分布和固有的类别特征，导致特征空间不稳定，给处理模糊边界和微妙对象带来了挑战。为了解决这些局限性，我们提出了一种混合扩散框架，以协同现有的判别式分割模型和所提出的 BBDM 的优势。在预训练 HiDiff 中使用的判别分割器时，我们使用交叉熵损失 LCE 和 Dice 损失 LDice 的组合作为判别目标函数：

B. Binary Bernoulli Diffusion Refiner 

        在此，我们提出了一种新颖的二进制伯努利扩散模型（BBDM）作为扩散细化器，通过对底层数据分布建模，有效、高效、交互式地细化任何判别分割器生成的先验掩码。BBDM 的新颖之处在于三个方面：(i) 基于伯努利扩散核，增强扩散模型在建模分割任务中的离散能力；(ii) 二值化扩散精炼器，以可忽略不计的计算成本显著提高效率；(iii) 交叉转换器，以交互方式交换扩散生成特征和判别特征，从而增强效果。下面，我们将详细介绍这三项新技术。 

1) Bernoulli-based Diffusion Model for Effective Refinement:

        我们采用 U-Net [30] 的变体 g(-) 作为扩散细化器，迭代细化由任何判别分割器 f (-) 生成的先验掩码 f (x)。关于先验掩码的作用，我们不仅将其用于扩散正向过程中的噪声添加，还将其作为扩散反向过程的起始采样点。BBDM 的整个扩散过程可以表示为：

其中 y1, ., yT 是与掩码 y0 大小相同的潜变量，q(-) 和 pθ(-) 分别代表正向扩散过程和反向扩散过程。

        现有的大多数基于扩散的分割方法都依赖高斯噪声作为扩散核，忽略了分割任务固有的离散性。因此，我们提出了一种基于伯努利的方法来解决这一局限性，详情如下。

前向扩散过程。在前向扩散过程中，我们的扩散细化器使用余弦噪声表 β1, ., βT 。伯努利前向过程 q(y1:T |y0, f (x)) 由以下公式给出： 

其中，B((1-βt)yt-1 + βtf (x))表示伯努利分布。其概率质量函数定义为：

 

使用符号 αt = 1 - βt 和 ，我们就能以封闭形式对任意时间步长 t 的 yt 进行采样：

其中，伯努利分布的均值参数可视为先验掩码 f (x) 与地面实况掩码 y0 之间的插值，随着时间步长的增加而增加。我们可以进一步使用伯努利采样噪声 ε ∼ B((1 - ̄ αt) |f (x) - y0|)，以先验掩码为条件，将公式 (6) 中的 yt 重新参数化为 y0 ⊕ ε，其中 | - | 表示绝对值运算，⊕ 表示（XOR）"逻辑运算。

        伯努利后验可表示为：

        扩散反向过程。扩散反向过程 pθ(y0:T |f (x))也可以看作是一个马尔可夫链，它从先验掩码 yT 开始--先验掩码 yT 取自伯努利分布，参数化为预先训练的分段模型 f (x) ∈ RH×W ×C ， yT ∼ B(f (x))--通过受先验掩码 f (x) 约束的中间潜变量来学习底层数据分布：

具体来说，我们的扩散精炼器 g(-) 通过校准函数 FC 对第 t 个时间步长下的伯努利噪声 ˆε(yt, t, f (x))进行估计，从而对 ˆμ(yt, t, f (x)) 进行重新参数化，如下所示：

FC 的目的是通过两个步骤将潜在变量 yt 校准为噪声较小的潜在变量 yt-1：(i) 通过计算 yt 与估计噪声 ˆε 之间的绝对偏差来估计掩码 y0；(ii) 利用公式 (7) 计算伯努利后验 q(yt-1 | yt, y0, f (x))，从而估计 yt-1 的分布。 

扩散目标函数。基于之前的 DPM [14]中的负对数似然的变分上界，给定一对图像-掩码（x, y0）和第 t 个潜变量 yt，我们采用 Kullback-Leibler (KL) 发散和焦点损失 [33] 来优化我们的扩散精炼器，如下所示：

其中，ε 和 ˆ ε 分别代表相应的地面真实噪声和估计伯努利噪声。γ用于平衡良好分类像素和错误分类像素的相对损失。当 γ = 0 时，LFocal 相当于 BerDiff  中使用的二元交叉熵损失 LBCE。最后，扩散目标函数定义为 ：

 2) 高效细化的二值化模块：

        为了减轻迭代扩散过程的计算负担，我们仿照文献[34]，提出将扩散精炼器二值化，使其轻量化，并定制了随时间变化的二值化（TB）和激活（TA）模块，从而以可忽略不计的资源实现高效精炼。值得注意的是，上标 b 和 r 分别表示二值和实值。

        二值化计算。在二值化计算中，我们首先将实值输入张量 Ur 和权值 Wr 分别二值化为 Ub 和 Wb，具体实现与 [35] 相同。然后，实值输入张量 Ur 和权重 Wr 之间计算量很大的浮点矩阵乘法，可以用二进制 Ub 和 Wb 之间轻量级的比特 XNOR 和 popcount 运算[36]来代替，定义如下：

         随时间变化的二值化模块。为了有效适应以时间步长为条件的 DPM 的迭代特性，受自适应实例归一化 [37] 的启发，我们设计了 TB 模块来对输入张量进行二值化，并设计了 TA 模块来动态激活输入张量。

        TB 模块通过与信道时间相关的二进制阈值 αi 来实现：

 其中，ub i 和 ur i 分别是第 i 个信道上同一输入张量元素的二值和实值表示。此外，αi 由轻量级全连接层生成，该层将时间步长 t 作为输入。

        类似地，TA 模块的实现方式为：

 其中，γi 和 ζi 是第 i 个信道上的可学习移位参数，由时间步长 t 线性转换而来；βi 是可学习缩放系数。

3) Cross Transformer for Interactive Enhancement:

        为了交互式地交换扩散生成特征和判别特征以增强效果，我们提出了一种新颖的交叉Transformer，称为 X-Former。X-Former 由两个交叉Transformer块（CTB）组成。第一个区块将从 U 型扩散精炼器瓶颈处提取的具有生成知识的特征编码到判别分割器的中间位置，即 fd, fp → f ′ p，而第二个区块的编码方向相反，即 f ′ p, fd → f ′ d，将具有判别知识的特征注入扩散精炼器。这种双向注入将实现更强的表征，从而生成更好的精炼掩码。

        如图 3 所示，我们的 X-Former 可以使用定制的 TA 和 TB 模块进行二值化，从而转换为对应的二值化 X-Former 或 BX-Former，后者由两个二值化交叉变压器块（BCTB）组成。这种二值化可用于减轻变压器模块的计算负担。

 C. Hybrid Diffusion Framework

        为了协同现有的判别式分割模型和建议的 BBDM 的优势，我们提出了一个混合扩散框架，将现有分割器的判别能力和扩散细化器的生成能力整合在一起，以改进医学影像分割，如图 2 所示。在判别分割过程中，判别分割器先提供分割掩码，而在扩散细化过程中，BBDM 则充当扩散细化器，有效、高效、交互式地细化分割掩码。下面，我们将详细介绍 HiDiff 的训练和推理过程，具体算法见 Algs.1和2。 

        在训练过程中，我们以交替协作的方式优化扩散细化器和判别分割器。具体来说，在优化扩散细化器时，我们可以冷冻判别式分割器并使用公式 (13) 中的扩散目标函数。在优化判别式分割器时，我们冻结了扩散细化器，并协同使用判别和扩散目标函数：

 

 

        在推理过程中，首先利用判别分割器生成先验掩码 f (x)，然后我们的 HiDiff 从先验掩码中采样初始潜变量 yT，接着进行迭代改进，以获得更好的掩码。需要注意的是，我们的 HiDiff 还能使用其他 DPM 加速策略，如去噪扩散隐含模型（DDIM）[23]，如 Alg. 2 所示。我们按照 [23] 的方法为我们的 HiDiff 推导出 DDIM 的采样策略，并将超参数 σt 设为 1- ̄ αt-1/ 1- ̄ αt，以减少反向过程的随机性。