Medical Diffusion - Denoising Diffusion Probabilistic Models for 3D Medical Image Generation

Medical Diffusion - Denoising Diffusion Probabilistic Models for 3D Medical Image Generation

论文链接：https://arxiv.org/abs/2211.03364

摘要

计算机视觉的最新进展已经在图像生成方面显示出有希望的结果。扩散概率模型尤其从文本输入中生成了真实的图像，如DALL-E 2、Imagen和Stable Diffusion所示。然而，它们在医学中的应用尚未得到系统评估，因为医学中的图像数据通常包括三维体积。合成图像可能在保护隐私的人工智能中起着至关重要的作用，也可以用于增强小数据集。这里我们展示了扩散概率模型可以合成高质量的医学成像数据，我们展示了磁共振图像（MRI）和计算机断层扫描（CT）图像。我们通过与两位医学专家的读者研究，对合成图像的质量进行了三个分类：真实图像外观、解剖正确性和切片之间的一致性，从而对其性能进行了定量测量。此外，我们证明了合成图像可以用于自我监督的预训练，并在数据稀缺的情况下提高乳房分割模型的性能（dice score 0.91 vs. 0.95 without vs. with synthetic data）。

介绍

深度学习（DL）在医学成像中的作用正在稳步增加¹。DL可以解决的一个典型问题涉及图像的分类，即将图像中包含的高维数据浓缩为单个类。相反的作用，即从低维非图像输入生成医学图像的研究较少，但具有巨大的潜力：合成图像可以用于在站点之间共享受保护的数据，或用于教育目的，甚至可以用于预测射线照相中的医学疾病进展^2，3。这些研究是在二维（2D）图像上进行的，但不是三维（3D）体积⁴。然而，现代医学中最重要的诊断成像方式，磁共振成像（MRI）或计算机断层扫描（CT）可以产生3D数据。因此，对2D数据的集中是一个严重的限制，它基本上忽略了可用于提高这些图像的手动或自动评估的性能和弹性的有用数据。因此，需要生成合成3D数据的方法。

先前的工作已经采用了生成对抗网络^5，6（GAN），但这种技术存在严重的局限性：首先，这些模型的训练很困难，模式崩溃是一个常见问题⁷，这意味着神经网络无法生成不同的样本。第二，即使没有发生模式崩溃，这些模型生成的图像的多样性也是有限的⁸。第三，GAN和类似模型只关注图像域，从文本生成图像或反之亦然并不简单。另一方面，扩散模型在非医学领域取得了巨大成功，因为它能够生成多种多样的图像，并将图像和非图像数据^9，10链接起来。尽管扩散模型的性能大大提高，但尚未系统地用于医学中的3D图像生成。

在这项工作中，我们检查了在医学中是否存在用于生成3D数据的扩散模型的潜力。我们提出了一种适用于潜在空间的扩散模型的新架构，并在四个公开可用的数据集上对其进行训练，这些数据集包括来自广泛解剖范围的数据：大脑MRI、胸部MRI、乳房MRI和膝关节MRI。我们在用户研究中调查这些图像对医学专家来说是否合理，并根据结构相似性指数定量检查其多样性。最后，为了弥合与医学应用的差距，我们研究了对生成的合成图像进行预训练是否有助于在有限的数据设置中改进分割模型。

相关工作

由于潜在扩散模型成功地在非医学领域生成高质量图像，并将图像数据与文本数据相关联，因此最近受到了关注¹¹。人们对将这一成功转化为医学领域仍感兴趣，因为此类模型可用于多种任务，包括数据匿名化、教育和培训、，数据增强和发现新的形态关联¹²。

然而，在CT和MRI的情况下，医学图像数据由于其三维性质而比2D图像数据更具挑战性。在目前的时间点上，我们所知道的只有一个同时进行的研讨会演示，它使用了潜在空间中的扩散模型来基于大型脑扫描数据库¹³生成3D MRI数据。其他几个小组已经致力于利用生成性对抗网络来生成3D数据¹⁴。相比之下，我们的方法可以被视为潜在扩散模型的扩展，其中我们将扩散概率模型附加到VQ-GAN¹⁵的潜在空间以生成高分辨率3D图像。与将扩散模型直接应用于3D数据^16，17相比，该方法具有几个优点：（1）我们可以减少训练模型所需的计算资源，因为它应用于具有较小维度的压缩潜在空间。（2） 虽然扩散模型擅长生成高质量的图像，但它们的采样速度——尤其是与GAN相比——相当低。通过在低维潜在空间中采样图像，我们可以减少生成新样本所需的时间。（3） 潜在空间封装了关于图像内容的更抽象的信息，覆盖了图像的整个区域从而使其可用于诸如预测未来图像外观而不管图像旋转或平移³的高级应用。

结果

Medical Diffusion Models can be Robustly Trained

我们在四个不同解剖领域的公开可用数据集上训练扩散模型：阿尔茨海默病神经成像倡议（ADNI）的脑部MRI检查、癌症成像档案（LIDC）的胸部CT检查、杜克大学（Duke）的乳房MRI检查和斯坦福大学（MRNET）的膝盖MRI检查。为了展示我们的方法的能力，即使是小数据集，这四个模型仅在1250（膝盖MRI）、998（脑部MRI）、1844（胸部MRI）和1010（胸部CT）图像上进行了训练。

尽管数据集相对较小，但每个模型都会收敛并生成真实的合成图像，而无需微调任何超参数（图1）。特别是，我们没有在任何训练课程中观察到模式崩溃。此外，模型架构能够适应多种分辨率，包括分辨率为64x64x64体素的脑MRI、分辨率为128x128x128的胸部CT以及分辨率为256x256x32体素的各向异性乳腺和膝关节MRI。在四个数据集中的每一个中，都可以生成真实的三维数据（图1）。

Medical Diffusion Models can Generate High Quality Medical 3D Data

我们根据三个不同的类别评估了人类专家的合成图像：1）整体图像外观的质量，2）切片之间的一致性，3）解剖正确性。两名分别具有9年（读者A）和5年（读者B）经验的放射科医生被要求对四个数据集中的每一个数据集的50幅图像进行Likert评分，见表1。

经验更丰富的放射科医生5对200项检查中的189项进行了评分，认为总体上真实，只有较小的不真实区域或更好（ADNI为50/50，LIDC为40/50，DUKE为50/60，MRNET为49/50）。200项检查中有191项被评定为在大多数切片中显示出切片之间的一致性（ADNI为50/50，LIDC为41/50，DUKE为50/60，MRNET为50/40），185/200项仅显示出轻微或无解剖不一致性（ADNI为50/50，LIDC为40/50，DUKE为50/40，MRNET为45/50）。具有5年经验的放射科医生对类似评分进行了分配（图2）。这些数据一起表明，我们的架构可以生成合成图像，对该领域的专家来说，这些图像看起来很逼真。

The Dimension of the Latent Space is Important for High Quality Image Generation

为了分析潜在维度对图像生成质量的影响，我们使用两种不同的压缩因子训练VQ-GAN自动编码器。我们发现，当将每个空间维度压缩8倍（即，256x256x32大小的图像的潜在维度为32x32x4）时，相关的解剖特征会丢失（图3）。当使用较小的压缩因子4训练VQ-GAN自动编码器时（即，尺寸为256x256x32的图像具有64x64x8的潜在尺寸），解剖特征被更精确地重建。对于所有四个数据集，我们发现最大压缩因子4（即，与图像的原始维度相比，潜在空间中的每个维度都小了四倍）包含了正确的解剖细节，如放射学专家在每个数据集的20个样本图像的测试集中所评估的。

Medical Diffusion Models outperform GANs in Terms of Image Diversity

为了将我们的扩散模型与已建立的GAN进行比较，我们遵循Kwon等人的工作⁵，并选择具有梯度惩罚（WGAN-GP）的Wasserstein GAN作为基线。由于我们在使用更高分辨率的图像训练WGAN-GP时发现了差异行为，因此我们限制了对生成的64x64x64大小的脑MRI图像的比较。我们通过对同一数据集的1000个合成样本对的结果进行平均，根据多尺度结构相似性度量¹⁸（MS-SSIM）对两个模型进行了比较。因此，较高的MS-SSIM得分表明模型之间更为相似，而MS-SSIM评分越低则表示相反。我们发现，GAN模型无法生成不同的图像，如其0.9996的高MS-SSIM评分所示，导致合成图像通常相同。相比之下，扩散模型的MS-SSIM得分为0.8557，更接近原始数据的MS-SSI得分（0.8095）。总之，这些数据表明，扩散模型能够生成代表原始数据分布的更多样的样本，因此这些模型可能更适合后续项目，例如用于分类模型的训练。

Synthetic Data can be Used to Train Neural Networks

在机构a希望与机构B合作以提高神经网络性能而不共享任何原始数据的情况下，我们评估了合成数据的可用性。为此，我们使用在DUKE数据集上训练的扩散模型生成了2000张合成图像，并在合成数据的自我监督设置中预先训练了Swin UNETR¹⁹。然后，我们利用来自机构B的可用分割数据对预先训练的网络进行微调，以在MRI扫描中分割乳房区域。为了展示有限数据环境下的绩效提升，我们进行了多次培训，其中我们使用了来自机构B的越来越多的可用数据（5%、10%、20%、40%、80%和100%）。为了进行比较，当没有使用合成数据进行预训练时，我们训练相同的神经网络来执行相同的任务。我们发现，使用另一机构的合成数据进行预训练可以在很大程度上提高骰子得分方面的分割性能，尤其是在可用标记训练较少的情况下（0.91没有预训练，0.95有5%可用数据的预训练，见图4和图5）。

讨论

随着非医学领域生成模型质量的提高，医学数据的合成成为一个可实现的目标，在教育、数据匿名化、数据扩展和开发新的DL算法^3，12。特别是扩散模型已被证明可以与人类在图像生成方面的能力相媲美^9，10。

在这项工作中，我们首次对MRI和CT数据进行了潜在扩散模型的大规模评估。我们证明，这种模型可以生成真实的3D体积数据，该数据在连续3D结构的合成中是一致的，并且能够准确反映人体解剖结构。我们表明，即使在大约1000个样本的相对较小的数据集上进行训练，这种复杂数据的训练对于我们的模型也会稳健地收敛。

这与通常需要大量超参数调整和大型数据集才能成功训练的GAN形成对比。更重要的是，即使可以成功地训练GAN，我们发现我们的扩散模型能够更准确地涵盖医疗实践中遇到的图像的多样性。这对于在AI方法的开发中使用这种合成图像非常重要。我们还通过在合成数据上预训练用于人类乳房MRI检查的分割模型，证明了潜在扩散模型的潜在医学应用，并表明这种预训练可以帮助使分割模型更健壮。

我们的工作有局限性：首先，我们在大约1000次考试的相对较小的数据集上评估了我们的模型。这部分是出于设计-在有限的数据可用时展示潜在扩散模型的可能性-部分是由于有限的计算资源。可以预期，当在更大的数据集上训练时，相似的模型可以生成更逼真、分辨率更高的图像¹³。其次，生成的3D体积不具有完整的诊断分辨率。这是由于可用的公共数据在分辨率上是有限的，并且在所有情况下都不能反映图像分辨率的最新水平。我们证明了图像质量和图像分辨率之间存在权衡，并且潜在空间的压缩因子对于获得真实图像至关重要。如果大型数据集可用于训练此类扩散模型，例如通过使用联合方法20，我们的实验表明，如果不将潜在空间中的压缩选择得太高，则可以提高图像分辨率。

总之，我们已经证明，与GAN相比，潜在扩散模型是生成合成3D医学数据的一种优越方法，并且可以为开发合成MRI或CT数据的AI方法奠定基础。

Online Methods

数据集说明

为了证明医学扩散模型的性能和稳健性，我们在四个不同的公开数据集上对其进行训练：MRNet²¹数据集包含来自1199名患者的1250次膝关节MRI检查，每个检查都包含轴向、矢状和冠状面扫描。为了演示的目的，我们只在脂肪饱和的矢状面T2扫描上训练我们的模型。阿尔茨海默病神经成像倡议（ADNI）²²数据集包含2733名患者的脑部MRI扫描。ADNI于2003年启动，由首席研究员Michael W。

ADNI的主要目标是测试序列磁共振成像（MRI）、正电子发射断层扫描（PET）、其他生物标志物以及临床和神经心理学评估是否可以结合起来测量轻度认知障碍（MCI）和早期阿尔茨海默病（AD）的进展。我们在998个标记为认知正常（CN）的3D MP RAGE序列上训练我们的模型。此外，我们在从922名乳腺癌患者获取的乳腺癌MRI数据集²³上评估了我们的模型，其中我们使用了每个患者的非脂肪饱和T1加权序列。为了证明我们的模型的通用性，我们还训练医学扩散模型来合成CT图像。为此，我们使用了来自肺图像数据库联盟（LIDC）和图像数据库资源倡议（IDRI）24的1010（n=1010名患者）低剂量肺CT。我们还使用了200（n=200名患者）T1加权乳腺MRI的内部数据集和相应的乳房区域地面真实掩模，以评估在自我监督的预训练方法中合成乳腺图像的使用。

数据预处理

通过将高分辨率图像平面缩放到256x256像素并对图像应用基于直方图的强度归一化²⁵，对来自MRNet数据集的膝关节MRI扫描进行预处理。该程序由数据集提供者²¹执行。此外，我们将每个图像中心裁剪为256x256x32（高度、宽度、深度）的统一形状。对来自ADNI数据集的脑MRI序列进行预处理，以去除MRI图像的非脑区域。此过程由数据集提供程序完成。为了在扩散模型和GAN之间进行比较，我们遵循了Kwon等人的方法⁵，在将大脑MRI输入神经网络之前，将其大小调整为64x64x64。来自乳腺癌数据集的图像通过首先将所有图像重采样到公共体素间距（0.66mm、0.66mm和3mm），然后使用勾勒出乳房轮廓的相应分割掩模裁剪出感兴趣的区域来进行预处理。然后，这些图像被分成两半，这样左乳房和右乳房就在不同的图像上。最后，将图像大小调整为256x256x32的均匀形状。因此，首先对肺CT进行重新采样，使其在所有方向上的公共体素间距为1mm。随后，将像素值转换为Hounsfield单位，并将图像中心裁剪为320x320x320的形状，然后再调整为128x128x128。将所有数据集的图像最小值最大值归一化为-1和1之间的范围。此外，我们在训练期间通过垂直翻转图像来增强所有数据集，概率为50%。

Architecture

医学扩散架构基于两步方法，我们首先将图像编码到低维潜在空间，然后根据数据的潜在表示训练扩散概率模型。在下文中，我们首先提供矢量量化自动编码器的背景信息，特别是VQ-GAN¹⁵和去噪扩散概率模型²⁶：

VQ-GAN

为了将图像编码为有意义的潜在表示，矢量量化自动编码器被证明是一种可行的选择，因为它们缓解了可变自动编码器^27、28中输出模糊的问题。它们通过将自动编码器瓶颈中的潜在特征矢量映射到从学习的码本中获取的量化表示来操作。Esser等人提出的VQ-GAN架构¹⁵可以看作是一类矢量量化的自动编码器，其图像重建质量通过在其输出端施加鉴别器损失而进一步提高。更准确地说，图像被输入编码器以构建潜在代码$z_e\in {\mathbb{R}}^{(H/s)\times (W/s)\times (k)}$。其中H表示高度、W表示宽度、C表示通道数、k表示潜在特征图数和s表示压缩因子。在矢量量化步骤中，通过将每个潜在特征矢量替换为其最接近的对应码本矢量来量化潜在特征矢量$e_n$包含在学习的码本Z中。然后通过将量化的特征向量馈入解码器G来重建图像。学习目标被定义为重建损失$L_{rec}$、码本损失$L_{codebook}$和承诺损失$L_{commit}$的最小化。正如原作者所定义的，我们使用感知损失作为重建损失，并使用直通估计器来克服不可微分量化步长。承诺损失被定义为未量化的潜在特征向量与相应的码本向量之间的均方误差。注意，仅针对连续潜在特征向量计算梯度，以强制更接近量化码本向量。可学习码本向量通过在映射到它的所有潜在向量上保持指数移动平均来优化。此外，在输出端使用基于块的鉴别器以获得更好的重建质量。为了扩展此架构以支持3D输入，我们遵循Ge等人²⁹的方法，并用3D卷积替换2D卷积。此外，我们将原始VQ-GAN模型中的鉴别器替换为切片鉴别器和3D鉴别器，前者将图像体积的随机切片作为输入，后者将整个重建体积作为输入。我们还遵循他们的方法添加特征匹配损失以稳定GAN训练。

Diffusion Models

扩散模型是一类生成模型，通过潜在变量上的马尔可夫链定义$x_1,...,x_T$²⁶。主要思想是，从图像$x_0$开始，我们通过在若干时间步长T内增加方差的高斯噪声来连续扰动图像。然后训练以时间步长t和时间步长本身的图像的噪声版本$p\left(x_{t-1}\mid x_t\right)$为条件的神经网络，以学习用于干扰图像的噪声分布。当变得足够大时，我们可以通过先验分布$\mathcal{N}(0,I)$进行近似$p(x_t)$，从该分布中采样，然后反向遍历马尔可夫链，这样我们就可以从学习的分布中采样新的图像$p_{\theta }\left(x_0\right):=\int p_{\theta }\left(x_{0:T}\right)d{x}_{1:T}$。用于建模噪声的神经网络通常被选择为U-Net30。为了支持3D数据，我们将U-Net中的2D卷积替换为3D卷积。此外，我们遵循Ho等人的方法³¹，仅在高分辨率图像平面上使用卷积（即核大小为$3\times 3\times 1$），然后在该高分辨率平面上使用空间关注块（因此将深度维度视为批量大小的扩展），以提高计算效率。空间关注块之后是深度关注块，其中高分辨率图像平面轴被视为批次轴。

Putting It All Together

在第一步中，我们在整个数据集上训练VQ-GAN模型，以学习数据的有意义的低维潜在表示。由于馈入扩散模型的输入应归一化为-1和1之间的范围，我们必须保证图像的潜在表示也在该范围内26。假设VQ-GAN模型中的矢量量化步骤强制学习的码本矢量接近量化前的潜在特征矢量，我们通过学习码本中的最大值来近似未量化特征表示的最大值。类似地，我们在学习的码本中用最小值近似未量化特征表示的最小值。因此，通过对未量化的特征向量执行简单的最小-最大归一化，我们获得了值接近范围-1和1的潜在表示。然后可以使用这些值来训练3D扩散模型。然后，我们可以从从标准高斯采样的噪声开始，通过反向扩散过程来生成新的图像。然后使用VQ-GAN的学习码本对该过程的输出进行量化，并随后馈送到解码器中以生成相应的图像。所有型号均在配备24GB GPU RAM的NVIDIA Quadro RTX6000上进行训练，每种型号大约需要7天。有关每个模型的训练设置的更多详细信息，请参见补充表1。

数据可用性

我们对公众可访问的数据进行了扩散模型实验，以允许其他小组复制和测试我们的结果。只有用于测试合成数据的医学适用性的乳房分割模型依赖于私人数据。该数据可根据作者的要求与书面合作和数据保护协议提供。LIDC-IDRI和乳腺MRI（DUKE）数据集可在癌症成像档案馆（TCIA）³²获得。ADNI数据集可在图像和数据档案馆（IDA）³³免费获得。MRNet数据集可直接从数据集提供商处获得²¹。
协议提供。LIDC-IDRI和乳腺MRI（DUKE）数据集可在癌症成像档案馆（TCIA）³²获得。ADNI数据集可在图像和数据档案馆（IDA）³³免费获得。MRNet数据集可直接从数据集提供商处获得²¹。