Ambiguous Medical Image Segmentation using Diffusion Models
事实证明,在临床任务中,专家小组的集体见解总是优于个人的最佳诊断。对于医学影像分割任务,现有的基于人工智能的替代方法研究更多地侧重于开发能够模仿最佳个体的模型,而不是利用专家组的力量。在本文中,介绍了一种基于单一扩散模型的方法,该方法通过学习群体洞察力的分布产生多种可信的输出。提出的模型利用扩散固有的随机抽样过程,只需最少的额外学习就能生成分割掩膜的分布。在三种不同的医学图像模式(CT、超声波和核磁共振成像)上证明,模型能够生成多种可能的变体,同时捕捉它们的出现频率。综合结果表明,提出的方法在准确性方面优于现有的最先进的模糊分割网络,同时保留了自然发生的变异。还提出了一种新的指标来评估分割预测的多样性和准确性,这与临床实践中的集体洞察兴趣相一致。
Introduction
近来,用于不同医学图像分割任务的专门确定性模型有了显著改进 。确定性模型即使存在不确定性,也会选择最有可能的假设,这可能会导致次优分割。为了克服这一问题,一些模型在分割任务中加入了像素方面的不确定性,但它们产生的输出结果并不一致。条件变异自动编码器(c-VAE)是一种条件生成模型,可与确定性分割网络融合,通过从以输入图像为条件的潜在空间采样,产生无限量的预测结果。概率 U 网及其变体在推理过程中使用了这种技术。在这里,潜空间是从先验网络中采样的,而先验网络经过训练后与 c-VAE 相似。这种对先验网络的依赖以及只在分割网络的最高分辨率注入随机性的做法,会导致分割预测的多样性和模糊性降低。为了克服这个问题,引入了一个单一的固有概率模型,不需要任何额外的先验网络,它代表了多位专家的集体智慧,可以在诊断管道中利用多种可信的假设。
扩散概率模型是一类生成模型,由使用变分推理训练的马尔可夫链组成。该模型通过对潜在空间的扩散过程建模,学习数据集的潜在结构。通过学习反向扩散过程,可以训练神经网络对使用高斯噪声模糊的噪声图像进行去噪处理 。最近,人们发现扩散模型在图像生成和内绘等各种任务中取得了广泛成功。也有人提出了一些利用扩散模型进行语义分割的方法。在这里,使用同一预训练模型的扩散模型的每个采样步骤中的随机元素为从单个输入图像生成多个分割掩码铺平了道路。然而,尽管扩散模型具有很大的潜力,但在使用扩散模型进行模糊医学影像分割方面仍缺乏探索。在本文中,提出了集体智能医学扩散模型(CIMD),以解决医学成像中的模糊分割任务。首先,介绍了一种新颖的基于扩散的概率框架,它可以从单个输入图像生成多个逼真的分割掩膜。这是因为可以利用扩散模型的随机抽样过程来抽取多个可信的注释。随机抽样过程还消除了推理阶段对单独 "先验 "分布的需求,这对于基于 c-VAE 的分割模型对模糊分割的潜在分布进行抽样至关重要。模型的分层结构还可以控制每个时间步骤的多样性,从而使分割掩码更加真实和异质。最后,为了评估模棱两可的医学图像分割模型,最常用的指标之一是广义能量距离(GED),它将地面真实分布与预测分布相匹配。在现实世界的模糊医学影像分割场景中,地面真实分布的特征只有一组样本。在实践中,GED 指标被证明会奖励样本多样性,而不考虑生成样本的保真度或与地面实况的匹配度,这在临床应用中可能是有害的。在医疗实践中,单个评估会被人工合并成一个诊断,并根据灵敏度进行评估。当进行实时小组评估时,参与者之间会达成共识。最后,在临床环境中还要考虑放射科医生之间的最小一致和最大一致。受当前集体洞察医学实践的启发,创造了一个新的指标,即 CI 分数(集体洞察),它考虑了总灵敏度、普遍共识和放射科医生之间的差异。
Proposed Method
Collectivly Intelligent Medical Diffusion
假设 b 是尺寸为 C × H × W 的给定图像,xb 是相应的分割掩码。在经典扩散模型中,输入图像 xb 需要用于训练,这将导致在推理过程中从噪声中采样时产生任意的分割掩码 x0。与此相反,要为给定的图像 b 生成一个分割掩码 xb,0,需要在输入图像中加入一个额外的通道。这将诱导解剖信息,将其连接为 xb 之前的图像,从而定义 X := b ⊕ xb。在噪声处理过程 q 中,噪声只添加到地面真实分割 xb 中。由于采样过程是随机的,扩散模型会为图像 b 生成不同的分割掩码 xb,0。
现在,使用模糊建模网络(AMN)对地面实况的模糊性进行建模。AMN 对输入图像的基本真相掩码的分布进行建模。在给定图像 b 和基本真相 xb 的情况下,将 AMN 的权重参数化为 ν,从而将分割掩码的模糊性嵌入潜在空间。该概率分布(用 Q 表示)被建模为高斯分布,其均值 μ(b, xb; ν) ∈ RN,方差 σ(b, xb; ν) ∈ RN×N,其中 N 表示低维潜在空间。潜空间的特征是:
同样,使用模糊控制网络(ACN)对预测掩码的模糊性进行建模。ACN 对来自输入图像的扩散模型调节的噪声输出进行建模。对于每个时间步长 t,假设 ˆ xb = fθ( ̃ xb, t),通过将 ACN 的权重参数化来估计扩散模型的模糊性,ω 为均值μ(b, ˆ xb; ω)∈ RN,方差σ(b, ˆ xb; ω)∈ RN×N 的概率分布 P,如下所示:
AMN 和 ACN 两个网络都使用对轴对齐高斯分布建模,其协方差矩阵为对角线。两种网络的结构细节可参见补充资料。通过库尔贝-莱伯勒发散(Kullback-Leibler divergence)对两种分布之间的差异进行惩罚:
最后,通过修改,将所有损失合并为一个加权和,正则化系数 β 为 :
在采样过程中,对于 Xt := b ⊕ xb,t,公式 修改为:
其中,z = N (0,I) 和 t = T,...,1。拟议方法的图形模型如图 2 所示。