【技术追踪】基于扩散模型的医学图像合成与测量指导（TPAMI-2024）

本文从数据分布的角度对以前的指导及其对进一步应用的贡献进行了分析。为了合成有助于下游应用的样本，本文在每个采样步骤中引入不确定性指导，并设计了一个不确定性引导扩散模型。在四个医学数据集上进行实验，在生成样本集上训练10个经典网络，为本文方法的实际贡献提供了全面的评价。此外，还为扩散模型中的一般梯度指导提供了理论保证，这将有助于进一步研究面向特定生成任务其他形式的测量指导。

风巽·剑染春水

1346人浏览 · 2024-07-23 10:04:39

风巽·剑染春水 · 2024-07-23 10:04:39 发布

不确定性引导条件：从生成类别医学图像到生成没那么确定的类别医学图像，增加合成图像的信息量~

论文：Measurement Guidance in Diffusion Models: Insight from Medical Image Synthesis
代码：https://github.com/yangqy1110/MGDM

0、摘要

在医疗领域，样本的获取通常受到多种因素的限制，包括成本、费时费力的标注、隐私问题和辐射危害，因此，合成感兴趣的图像是数据扩充的重要工具。
近年来，扩散模型在各种合成任务中取得了最先进的成果，嵌入能量函数（embedding energy functions）已被证明可以有效地指导预先训练好的模型合成目标样本。

然而，目前方法的开发和验证仍然仅限于改进指标，如 Fréchet Inception Distance score（FID）和 Inception Score（IS），并没有对下游任务进行深入研究，如疾病分级和诊断。
此外，现有的分类器指导作为能量函数（energy function）的一种特例，仅对合成数据集的分布产生单一影响，这可能导致生成的样本对下游模型的优化帮助有限。

本文从数据分布的角度对以前的指导及其对进一步应用的贡献进行了分析。
为了合成有助于下游应用的样本，本文在每个采样步骤中引入不确定性指导，并设计了一个不确定性引导扩散模型。

在四个医学数据集上进行实验，在生成样本集上训练10个经典网络，为本文方法的实际贡献提供了全面的评价。此外，还为扩散模型中的一般梯度指导提供了理论保证，这将有助于进一步研究面向特定生成任务其他形式的测量指导。

1、引言

1.1、扩散模型与医学图像合成

（1）追求高性能的扩散模型用于生成任务并非易事，因为从训练目标到采样过程可能会遇到多种挑战；
（2）从实用性的角度来看，临床医生对视觉感知以及下游应用（如疾病分级和诊断）的性能提升有更高的要求；
（3）扩散模型的发展不仅应关注快速和大规模的生产，还应关注高质量的生成，若合成的医学样本不能对准确诊断或治疗计划做出实际贡献，其价值是有限的；

1.2、扩散模型应被引导生成什么样的样本

（1）从自动诊断的角度来看，高质量的样本是那些具有可靠的标签和高信息量的样本；
（2）之前的扩散模型尝试通过在训练过程中嵌入类别信息来拟合特定类别样本的条件分布，然而，对于医学图像合成任务，这种策略在完全保证可靠性方面是有限的，因为合成的样本有时与标签不匹配；
（3）为保证合成样本的标签可靠性，可在扩散模型的采样过程中采用分类器指导以生成类别一致的图像样本；
（4）分类器提供的概率梯度有时可能会在渐进采样过程中将合成结果带到不同的极端，使得扩散模型倾向于生成更健康或更严重的样本；
（5）分布内样本对下游模型优化的影响有限，并可能在某种程度上浪费计算资源；
（6）在实际诊断中，即使是同一种疾病，不同患者的严重程度也不同，而那些中间病例是样本集中最有信息量的资源；
（7）通过另一种指导合成各种等级的样本具有重要的临床意义，将分类器指导扩展到一般测量指导有助于实现可控生成；

本文中，高质量医学图像合成的关键是在扩散模型的采样过程中使用适当的指导以生成对下游任务更有帮助的样本。在采样步骤中引入了两种不确定性指导（熵和边缘），并设计了一种旨在解决生成过多分布内样本挑战的不确定性指导扩散模型。之前的分类器指导仅对改变增强数据集的分布有单一影响，本文的不确定性指导提供了相反的效果，使这种分布更具可控性。

分类器引导扩散模型合成的四个典型的分布样本，其病变区域很容易识别，并且数据分布位于类别的中心：

在这里插入图片描述

1.3、本文贡献

（1）对之前的分类器指导及其对增强数据分布的影响进行了全面分析。为了实现可控生成，将其扩展到任何测量，并为一般梯度指导提供了理论保证。此外，利用了两种不确定性测量（熵和边缘），并在扩散模型的每个采样步骤中提出了不确定性指导；

（2）提出了一种新的高质量医学图像合成扩散模型，即不确定性指导扩散模型（uncertainty-guided diffusion models，UGDM）。为了在较少的扩散步骤中学习更丰富的视觉概念，在训练过程中结合了感知优先权重（perception prioritized weight，P2W）和可学习方差。此外，为了生成既具有可靠标签又具有高信息量的样本，设计了一个两阶段采样过程，其中每个阶段都由特定的测量梯度指导；

（3）为了评估合成医学图像样本对进一步任务的实际贡献，除了常用的指标外，在增强数据集上训练了十个经典网络，提供了更直接的准确性比较。实验结果表明，所提出的扩散模型在四个公共医学图像数据集上可以实现更高的平均胜率；

2、相关工作

2.1、扩散模型

（1）与GANs相比，扩散模型有更多完善的数学解释，无对抗性的训练策略，以及实现稳定和可控的生成的能力；
（2）扩散概率模型的原理：利用生成的马尔可夫链将一个简单的已知分布转换为一个目标分布；
（3）DDPM：高质量的图像生成，主要局限性是其沿马尔可夫链的长时间渐进采样；
（4）DDIM：采样加速，提供了一种通过噪声隐变量的确定性抽样的条件生成策略；
（5）Score-based：从随机微分方程（SDEs）的角度创新性地研究了扩散模型；
（6）扩散模型在图像任务中的应用：重建、配准、分割、去噪、2/3D 生成、异常检测；

2.2、医学图像合成

（1）传统的图像增强方法，如缩放、旋转、仿射和变形变换，很难提供医学图像数据集的内在多样性；
（2）当深度学习训练需要大规模数据时，合成数据可以有效地补充传统的数据增强；
（3）理想的合成数据在形态和纹理上应该与真实数据相似，但同时在视觉外观和数据特征上具有更大的多样性；

2.3、不确定性

（1）基于学习的模型中的不确定性表明了该模型对其预测的确定程度，这可以作为一个评价指标，以说明鉴别模型对其对给定输入的正确输出的无知或模糊性；
（2）引入不确定性理论对医学图像进行分类和分割，可以取得明显更好的性能；

3、基础工作（Preliminaries）

扩散模型基本原理：原文略，可参考【Diffusion综述】医学图像分析中的扩散模型（一）中2.2节；

4、方法

之前的研究创新地指导预先训练的扩散模型来生成具有强类别条件的样本，这启发了本文的假设：分类器指导只是一个特殊的情况，本节中，将分类器指导扩展到一般测量 ${\mathcal{M}}$ ，以便实现可控生成；

4.1、模型（Model）

用于准确诊断的高质量样本应该具有可靠的标签和高信息量，也就是说，在类别信息可靠的情况下，我们需要引入一些不确定性。如图2所示，本文提出了一种新的基于扩散模型的医学图像合成方法，并设计了一个两阶段的采样过程，每个阶段都由特定的梯度引导。

不确定性引导扩散模型（UGDM）的 pipeline 作为一个数据增强的工具：

在这里插入图片描述

在获得在具有类别信息的原始样本集上预先训练的扩散模型和分类器后，使用两阶段采样过程对下游应用进行大规模数据增强：
（1）第一阶段，即生成阶段，利用每个采样步骤中的分类器引导，从随机高斯噪声中生成具有可靠标签的初始样本（分布内样本）；
（2）第二阶段，即扰动阶段，通过不确定性引导采样，将出分布结果转换到潜在空间，然后通过再次采样合成出分布结果；

从数据集分布的角度来看，第一阶段的分类器指导是指接近原始样本集的类别中心的扩充样本，而第二阶段的不确定性指导是指适当地将合成的样本推到类别边界；

4.2、训练目标（Training Objectives）

如图2所示，应该在具有类别信息的原始样本集上训练一个扩散模型，之前的扩散采用了经典的目标 ${L_{classic}}$ ，对于每一步 ${t}$ ，去噪分数匹配损失是两个高斯分布之间的距离，可以用噪声预测器 ${ϵ_θ}$ 来重写；
在扩散模型训练中嵌入类别信息可以有效地拟合与类别相关的条件分布，在训练过程中，类别信息 ${y}$ 和 ground truth 真实信息成对地输入网络，因此网络可以学习输入和输出之间的相关性；
当模型容量足够、数据分布合理时，可以在很大程度上保证合成样本的标签可靠性；

本文中，为进一步保证生成的稳定性和提高学习能力，扩散模型应该注重学习更多的视觉概念。在扩散模型训练过程中，图像样本在低噪声水平下仅被轻微损坏，恢复它不需要对图像上下文的先验知识。然而，当样本被严重损坏时，其内容是可识别的（应该不可识别吧，奇怪~），需要优先解决这些重要噪声水平的代理任务（pretext task）。

恢复被一定噪声水平破坏的数据，为扩散模型学习更丰富的视觉概念提供了适当的代理任务，在扩散模型训练中使用它可以显著提高扩散模型在生成任务中的性能；

基于 improved DDPM 中学习方差的方法，本文的模型可以用较少的扩散步骤产生样本。
为进一步改进扩散模型训练，本文对 ${L_{classic}}$ 进行了优化，使用扩散训练目标如下：

在这里插入图片描述

其中 ${W}$ 表示 P2W，被引入训练目标中，以学习更丰富的视觉概念，此处， ${W}$ 的计算方法为：

在这里插入图片描述

另一方面，为了实现引导抽样，应该在具有时间 ${t}$ 信息嵌入的原始样本集上预先训练一个分类器 ${p_ϕ(y | x_t)}$ ，其中 ${y}$ 表示样本标签，本文使用了一个分类器架构，它是一个简单的具有 attention pool 的 U-Net 模型中的降采样主干，其训练目标函数如下， ${||·||_c}$ 为交叉熵损失函数：

在这里插入图片描述

4.3、测量指导（Measurement Guidance）

分类器引导是最近提出的一种技术，通过使用预先训练的分类器模型 ${p_ϕ(y | x_t)}$ 的梯度，减少类条件扩散模型 ${ϵ_θ(x_t|y)}$ 的多样性，其中 ${y}$ 是类别标签，它是辅助扩散模型生成可靠样本的一种简单、计算高效的方法。
具体而言，对于给定的类别标签 ${y_i}$ ，分类器引导通过在每个采样步骤中替换 ${ϵ_θ(x_t|y_i)}$ 来进行分类器引导采样：
在这里插入图片描述
按照之前的假设，分类器引导只为引导采样提供了一种特殊的度量，将其扩展到利用任何度量 ${\mathcal{M}}$ 具有重要意义；
具体来说，正向的 SDE 是：

在这里插入图片描述
其对应的反向 SDE 为：

在这里插入图片描述
当 ${σ_t^2 = 0}$ 时，它对应于使用 DDIM 的反向 SDE：

在这里插入图片描述
其中：

从数学上讲，指导是指利用条件 ${Y}$ （即类标签或图像/文本嵌入）来对先验的数据分布 ${p_θ(x_t)}$ 进行调节，从而产生 ${p_θ(x_t|Y)}$ 。
对于本文的测量指导，其反向条件 SDE 为：

在这里插入图片描述
其中，利用贝叶斯公式，可以得到：

在这里插入图片描述

为了通过更高值的 ${\mathcal{M}(x_0)}$ 获得采样结果 $x_0$ ，使 ${p(\mathcal{M}|x_t)∝e^{γ·\mathcal{M}(x_t)}}$ ：

在这里插入图片描述
此外：

在这里插入图片描述
与之前的分类器引导类似，通过在每个采样步骤中替换 ${ϵ_θ(x_t|y)}$ 来执行测量引导采样：

在这里插入图片描述

因此，新的采样过程倾向于产生具有较大 ${\mathcal{M}(x_t)}$ 值的 $x_t$ ，这最终导致 $x_0$ 具有较大 ${\mathcal{M}(x_0)}$ ；

在每个采样步骤中：

在这里插入图片描述

理论上，为保证 Eq. 14 中的反向 SDE 具有唯一的平滑解，对应 Eq. 19 中离散近似所描述的域转移过程具有一定的平滑性， ${\mathcal{M}(·)}$ 应满足对状态 $x_t$ 和时间 $t$ 的可微性。
根据实分析中的 Rademacher 定理，利用 Lipschitz 条件的符号来描述光滑唯一解存在的充分条件。遵循这个符号系统，并希望为扩散模型的受控生成提供更多的理论应用实例。
采用文献[22]附录中给出的 ${\mathcal{M}(·)}$ 满足正则性条件的充分条件描述为：

在这里插入图片描述
考虑到原始采样过程：

在这里插入图片描述
重要的是：

在这里插入图片描述
可以将 Eq. 21 和 Eq. 22 代入 Eq. 19：

在这里插入图片描述
综上所述，这相当于在每个采样步骤对 ${\mathcal{M}}$ 进行梯度上升；
随后，可以使用与常规 DDIM 相同的采样程序，这为扩散模型的一般梯度制导提供了理论保证，并为今后研究针对生成任务的其他形式的测量制导铺平了道路。

在这里插入图片描述

4.4、采样和不确定性指导（Sampling and Uncertainty Guidance）

算法1说明了本文用于指导的整体两阶段采样过程，结合图2可以更好地理解；
在第一阶段，选择较大的分类器引导强度（本文中 s = 10）来合成具有高类别置信度的样本；
在第二阶段，使用 DDIM 采样和不确定引导采样的反向过程，进一步合成高信息量的样本，由于这个扰动阶段，合成样本在数据集分布上更接近类边界；顺便提一下，第二阶段也可以利用人工标记的样本集作为数据增强工具；

第二阶段包括 DDIM 反演和不确定性引导采样，有研究表明，使用 DDIM 采样重构输入图像的次要细节，有利于近乎精确的重构和真实图像处理，受此启发，本文通过反向运行 DDIMs 的确定性生成过程来推断出一个随机潜在编码：

在这里插入图片描述
因此，本文的目标是通过另一个指导 ${\mathcal{M}(x_t)}$ 合成不同级别的医学图像样本，用于下游任务，如疾病分级和诊断。
如图2所示，在第一阶段的分类器引导 ${∇_{x_t}log p_θ(x_t)}$ 引导从随机高斯噪声中采样生成类一致的图像。在第二阶段，另一个引导 ${\mathcal{M}(x_t)}$ 将这些可靠的样本引导到一个轻微的类不确定性样本。

因此，在第二次采样过程中，应该使用符合式(20) 假设的不确定度测量来提供相反的指导，基于预先训练的分类器 ${p_ϕ(y | x_t)}$ ，引入了两种不同的不确定度测量，即熵 ${\mathcal{M}_e}$ 和边缘 ${\mathcal{M}_m}$ ：

在这里插入图片描述
其中 ${p_ϕ(\hat y_1 | x_t)}$ 和 ${p_ϕ(\hat y_2 | x_t)}$ 分别表示 ${x_t}$ 最有可能属于的两个类别概率；
由于这两个测量值都是基于预训练的分类器，因此它们明显符合 Eq.(20) 的假设，那么我们可以通过上述指导获得高质量的样本。由于这两个指导（分类器和不确定性）存在一定的互斥性，同时设置两个超参数不方便。本文设计了两阶段的采样过程，并在不同的阶段使用了具体的指导。

对于引入不确定性指导是否会改变样本属性的问题，在参数设置和数据分布合理的情况下，本文认为答案是不会的。在第二阶段，基于预训练分类器的不确定性梯度只是对采样过程的干扰。
当不确定度尺度 ${γ=0}$ （无干扰）时，扩散模型本身将合成的样本属于 ${y}$ ；
当不确定度尺度 ${γ > 0}$ 时，随着 ${\mathcal{M}(x_t)}$ 的增加，梯度变小，并逐渐趋于0，因此，条件 ${y}$ 在采样的后期阶段中绝对占主导地位；如果一个函数的定义域和值域是有限的，则可以证明其倒数是有界的。因此，在本文的扩散模型中提出的不确定性制导是可控的；

随着 ${γ}$ 值的增加，合成样本的类别置信度降低。然而，这种下降并不能动摇 DDIM 采样下的类别属性。图3给出了不同不确定度尺度 ${γ}$ （从0到10）下合成的胸片肺炎样本，及其使用 VGG19 的分类概率，可以清楚地看到，在这种情况下，样本的类别属性没有被改变。

采用 VGG19 对不同不确定度尺度 ${γ}$ （0~10）下的合成胸片肺炎样本进行类别概率比较：

在这里插入图片描述

4.5、评估（Evaluation）

本文作者认为，FID 与 IS 容易惩罚非 GAN 模型，且他们与没有显示出与评估下游任务性能的指标一致的相关性。
为了全面评价这些合成样本对其下游任务的确切贡献，本文使用了10个著名的网络，在分级或诊断准确性方面进行了比较。本文的扩散模型的有效性独立于这些网络，这10个验证网络具有巨大的影响力和广阔的市场前景。与最近提出的网络相比，它们所提供的精度具有更高的可靠性和更广泛的接受度。

5、实验与结果

5.1、数据集与设置

4个医学数据集与实施设置：
在这里插入图片描述
对于基线，使用原始样本集进行下游网络训练，以提供样本贡献的基本评估；
对于消融，将原始样本集与无制导扩散模型和分类器制导扩散模型合成的样本混合，生成新的样本集用于下游网络训练；
如果在原始样本和我们的合成样本上训练的网络在精度方面可以超过基线和消融，我们将相信所提出的不确定性指导有利于高质量的样本合成；

5.2、Messidor 数据集结果

Messidor数据集是一个具有连续分级的具有代表性的疾病分级任务；

两种不确定性引导的扩散模型使病变区域更小，对人眼不易区分；

生成样本可视化：

在这里插入图片描述

通过熵引导扩散模型和边缘引导扩散模型，其精度最多可分别提高5.00%和4.00%；

不同网络分级精度：

在这里插入图片描述

5.3、ISIC 数据集结果

ISIC 是一个典型的良恶性类别不平衡数据集，因此，只合成恶性样本来缓解这种不平衡；

所提出的不确定性指导略微减少了第一个样本中的病变区域，对于无指导的合成结果，这个恶性样本是错误的；

生成样本可视化：

在这里插入图片描述

在大多数情况下，本文的方法可以提高诊断的准确性，只有一小部分略有下降。由于基线非常高，因此精度的总体提高并不显著，这些合成样本也没有太大的贡献空间；

不同网络分级精度：

在这里插入图片描述

5.3、PneumaniaMnist 数据集结果

不确定性指导显著缩小了感染样本中白肺的范围;

生成样本可视化：

在这里插入图片描述

熵引导在诊断方面表现出明显更好的表现，而边缘引导显示出相对较差的结果，可能是因为熵在不确定性评估方面更加全面；
在SqueezeNet和VGG13结果中，仅分类器引导扩散模型合成的样本贡献很小，甚至降低了诊断精度，这种情况证明了过度置信度样本的存在；

不同网络分级精度：

在这里插入图片描述

5.4、BreastMnist 数据集结果

不确定性指导显著缩小了样本中病变的范围；

生成样本可视化：

在这里插入图片描述

熵引导和边缘引导扩散模型在下游诊断方面均明显优于其他方法；

不同网络分级精度：

在这里插入图片描述

5.5、平均胜率

熵引导扩散模型和边缘引导扩散模型击败了之前的扩散模型，平均胜率分别为7.25/10和6/10，熵比边缘表现得更好、更稳定；

UGDM的平均胜率：

在这里插入图片描述

5.6、FID 和 IS 比较

四种生成模型的FID和IS之间没有显著差异，FID和IS与评估下游应用性能的指标显示出有限的一致性相关性，因此，它们在面向诊断的医学图像合成任务中的评价价值有限；

四个数据集上 FID 和 IS 的比较：

在这里插入图片描述

5.7、不确定性指导是否真的有助于提高样本的信息量

对合成样本进行类置信度分析，如果不确定性指导可以降低类置信度，就可以确认该指导可以增加样本信息量；
大多数的平均分类概率变化是负的，也就是说，本文提出的不确定性指导可以有效降低分类器引导合成样本在统计总体意义上的分类置信度，并且可以将分布内样本的合成引导到更复杂的样本中。该方法可以缓解医学成像领域中常见的样本集不平衡问题，并将优化辅助医疗诊断模型的训练数据质量；

分类器（U-Net结构的编码器部分）在四个数据集上不确定性引导的概率分析：

在这里插入图片描述

原始数据分布（Original）的颜色比其他三个颜色要浅，表示样本数量明显增加；
对于没有指导的结果（None），一些合成样本位于相反的类，这意味着它们的类别信息是错误的。也就是说，无指导抽样过程已经产生了一个带有健康标签的受感染样本；
分类器引导（Classifier）纠正了这一错误，两类样本之间的距离变大，这就是为什么必须在第一阶段添加一个分类器引导来纠正这个类别错误；
对于不确定性引导结果（Classifier+Entropy），与之前的分类器引导相比，没有样本错误，两类样本之间的距离已经缩小，总体而言，该方法可以保证样本的可靠性，降低合成样本的类间可分性；

PneumoniaMnist 中的原始样本集和合成样本集的特征可视化：

在这里插入图片描述

5.8、引导强度是否影响下游诊断

图像合成的性能有时对 ${γ}$ 敏感，有时不敏感。如表2-5所示，不同的分类网络在不同的合成数据集上的诊断性能差异较大，在目前的研究中很难确定最佳的指导尺度；

用4个数据集对不同不确定性尺度 ${γ}$ 下的合成结果进行了可视化比较：

在这里插入图片描述

在采样过程中使用不同的分类器引导尺度时，平均诊断准确性没有显著差异，无论如何选择分类器引导尺度，都不影响在第二阶段的不确定性指导的平均有效性；
将分类器引导尺度设置为较大的值20时，本文的方法在所有分类网络上都表现出更好的提升性能，然而，不同的分类网络在不同的合成数据集上的诊断性能差异很大；

在 PneumoniaMnist 数据集上，用所提出的不确定性指导对分类器指导 (cg) 和无分类器指导 (cfg) 的指导尺度分析：

在这里插入图片描述

5.9、不确定性引导是否可以推广到无分类器引导的扩散模型

如表 9 所示，在 PneumoniaMnist 数据集上，熵和边缘引导对无分类器引导都有更好的促进作用；
${w}$ = 0.5 与 1 的平均诊断准确率没有显著差异。然而，当我们将 ${w}$ 设为 2 时，本文的方法的平均诊断准确率低于其他方法；

6、局限性与未来工作

6.1、局限性

（1）从自动化的角度来看，本文的扩散模型的指导尺度参数仍需手动设置，这使其难以成为广泛和主流的数据增强工具。此外，在通用性方面，不同的诊断网络（如VGG、ResNet和Stochastic Depth）在使用不同参数的合成样本上的适应性存在显著差异，尚未找到指导尺度设置的内在规律。根据对四个不同数据集的初步调查，不同合成样本的诊断结果有时对指导尺度敏感，有时则不敏感，且没有明确的趋势与这些参数对应；

（2）从评估的角度来看，当前的合成结果通常通过视觉表现的定性比较和使用评估指标（如 FID 和 IS ）的定量比较来评估。 然而，一方面，组织足够多的合格观察者来完成相对公平的测试是一项耗时的任务。另一方面，如表7所示，当前的指标对应用导向的合成任务的评估价值有限。此外，目前没有针对合成医学图像评估的特定 Inception Score；

（3）从理论的角度来看，本文中的分类器和不确定性指导都是测量指导的特殊情况。 作者认为，未来在扩散模型中会有各种测量指导，如相似性、不规则性和轮廓信息，以用于更多的下游应用。

（4）从工业的角度来看，本文中的2D合成技术无法支持3D医学体积合成，而3D医学体积合成在MRI、超声和CT等领域更为常用。 实现具有时间和空间信息的高质量 3D 合成是一项具有挑战性的任务。目前，作者的研究仅专注于开发特定成像模态和特定疾病样本的医学生成模型。随着人工智能生成内容（AIGC）的快速发展，一种更用户友好的合成技术可能在医院和学术机构中具有更多潜力。