【技术追踪】通过潜在扩散和先验知识增强时空疾病进展模型（MICCAI-2024）

  向扩散模型中引入先验知识，实现疾病进展预测，扩散模型开始细节作业了~

---

论文：Enhancing Spatiotemporal Disease Progression Models via Latent Diffusion and Prior Knowledge
代码：https://github.com/LemuelPuglisi/BrLP

---

0、摘要

  本文介绍了基于潜在扩散的新型时空疾病进展模型—— Brain Latent Progression（BrLP）。BrLP 旨在预测个体在三维脑 MRI 上的疾病进展。（由当前图像获得疾病进展后的图像）。

  为完成这一任务而开发的现有深度生成模型主要以数据驱动为主，面临学习疾病进展的挑战。BrLP 通过结合疾病模型的先验知识来增强预测的准确性，解决了这些挑战。

  为了实现这一点，本文建议整合一个辅助模型来推断不同脑区的体积变化。此外，本文引入了潜在平均稳定化（Latent Average Stabilization，LAS），一种新的技术来提高预测进展的时空一致性。

  BrLP 在一个大规模数据集上进行训练和评估，该数据集包含来自 2805 名受试者的 11730 张 T1 加权脑部 MRI 图像，这些图像采集自三个公开的纵向阿尔茨海默病（AD）研究。

  实验中，在横断面和纵向设置中，将 BrLP 生成的 MRI 扫描与受试者实际可用的随访 MRI 进行比较。BrLP 在体积准确性和图像相似性方面均显著优于现有方法，其中在与阿尔茨海默病（AD）相关的脑区体积准确性提高了 22%，与真实扫描的图像相似性提高了 43%。

  BrLP 能够在受试者水平生成条件化的三维扫描，并且通过整合先验知识来提高准确性，这在疾病进展建模方面代表了一个重大进步，并为精准医疗开辟了新的途径。

---

1、引言

1.1、研究意义与当前挑战

  （1）相比于生物标志物的表征，利用丰富的、高维的影像表示疾病进展，便于可视化和精确定位复杂结构变化的模式，从而提供对疾病动态更详细的了解；
  （2）挑战 1： 通过基于受试者特定的元数据进行条件化，以提高个体化程度；
  （3）挑战 2： 在有条件的情况下使用纵向扫描；
  （4）挑战 3： 增强时空一致性，以实现空间和时间维度上的平滑进展；
  （5）挑战 4： 管理由于使用高分辨率三维医学图像而导致的高内存需求；

1.2、本文贡献

  （1）解决挑战 1： 提出将 LDM 和 ControlNet 结合起来，根据可用的受试者数据生成个体化的脑部 MRI；
  （2）解决挑战 2： 提出通过使用一个辅助模型来整合疾病进展的先验知识，该模型旨在推断不同脑区的体积变化，从而在有条件的情况下使用纵向数据；
  （3）解决挑战 3： 提出了 LAS，这是一种用于改善预测进展的时空一致性的技术；
  （4）解决挑战 4： 使用脑部 MRI 的潜在表示，以减少处理三维扫描的内存需求；

  本文通过训练 BrLP 来学习不同认知状态个体脑部的渐进性结构变化，从而评估其性能：认知正常（Cognitively Normal, CN）、轻度认知障碍（Mild Cognitive Impairment, MCI）以及阿尔茨海默病（Alzheimer’s Disease）。

  为此，本文使用了一个大规模数据集，包含来自 2805 名受试者的 11730 张 T1 加权脑部 MRI 图像，这些图像来自三个公开的阿尔茨海默病纵向研究。据本文作者所知，本文是首次提出一种用于脑部 MRI 的三维条件生成模型，该模型将疾病进展的先验知识整合到图像生成过程中。

挑战与贡献对应的写法很棒的感觉~

---

2、方法

2.1、Background - Diffusion Models

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

2.2、Proposed Pipeline - Brain Latent Progression (BrLP)

  BrLP 架构由四个关键组件构成：一个 LDM（潜在扩散模型）、一个 ControlNet（控制网络）、一个辅助模型以及一个 LAS（时空一致性增强）模块。这四个组件共同解决了引言部分提到的挑战，其架构总体如 图1 所示。具体来说：

  （1）LDM 旨在生成符合特定协变量的随机三维脑部 MRI 图像；
  （2）ControlNet 则专注于将这些 MRI 图像细化到特定受试者的解剖结构；
  （3）辅助模型利用疾病进展的先验知识，以提高对特定脑区体积变化预测的精确性；
  （4）LAS 模块在推理过程中被用于增强时空一致性；

  
Figure 1 | BrLP训练和推理过程概述：

2.2.1、LDM：学习脑 MRI 分布

  在文献 [12] 的基础上，本文训练了一个潜在扩散模型（LDM），其目标是生成反映特定协变量 $c=⟨s,v⟩$ 的三维脑部 MRI 图像。其中，$s$ 包括受试者特定的元数据（如年龄、性别和认知状态），而 $v$ 则涵盖了与阿尔茨海默病（AD）进展相关的指标，例如与 AD 进展相关的脑区（海马体、大脑皮层、杏仁核、大脑白质和侧脑室）的体积。

  LDM 的构建是一个两阶段的过程。首先，训练一个自编码器 $(\mathcal{E},\mathcal{D})$（图1中的A部分），该自编码器旨在为数据集中的每个脑部 MRI $x$ 生成一个潜在表示 $z=\mathcal{E}(x)$。随后训练一个条件 UNet（图1中的B部分），表示为 ${ϵ}_{\theta }$，网络参数为 $\theta$，旨在估计从 $z_t$ 恢复到 $z_{t-1}$ 所需的噪声 ${ϵ}_{\theta }(z_t,t,c)$，通过最小化损失函数 ${\mathcal{L}}_{ϵ}$ 来训练 ${ϵ}_{\theta }$。协变量 $c$ 作为条件通过交叉注意力机制集成到网络中。

  生成过程首先从采样随机高斯噪声 $z_T\sim \mathcal{N}(0,I)$ 开始，然后迭代地逆向执行每个扩散步骤 $z_t\to z_{t-1}$，其中 $t=T,\dots ,1$。对最终步骤 $t=1$ 的输出 $z_0$ 进行解码，可以得到一个符合指定协变量 $c$ 的合成脑部 MRI 图像 $\hat{x}=\mathcal{D}(z_0)$。

2.2.2、ControlNet：基于受试者脑部 MRI 的条件

  LDM 仅通过协变量 $c$ 对生成的脑部 MRI 提供有限程度的控制，并且不允许根据单个解剖结构对模型进行调节。这个模块的目的是扩展 LDM 的能力，使其能够包含这种额外的控制功能。为了实现这一点，本文使用了 ControlNet（图1中的C部分），这是一个设计为与 LDM 协同工作的神经网络。

  将 ControlNet 和 LDM 概念化为一个统一的网络 ${ϵ}_{\theta ,\varphi }$，其中 $\theta$ 表示 LDM 的固定网络参数，$\varphi$ 表示 ControlNet 的可训练网络参数。与 LDM 一样，${ϵ}_{\theta ,\varphi }$ 仍用于预测反向扩散步骤 $z_t\to z_{t-1}$ 中的噪声 ${ϵ}_{\theta ,\varphi }(z_t,t,c,z)$，现在将 $z=\mathcal{E}(x)$ 也作为条件，以涵盖生成过程中目标脑 $x$ 的结构。

  为了训练 ControlNet，本文使用来自同一患者在不同年龄 $A<B$ 时拍摄的脑部 MRI 对的潜在表示 $z^{(A)}$ 和 $z^{(B)}$。与 $z^{(B)}$ 相关的协变量 $c^{(B)}$ 已知并用作目标协变量。每次训练迭代包括：

  （1）采样 $t\sim U[1,T]$
  （2）执行 $t$ 步前向扩散 $z^{(B)}\to z_t^{(B)}$
  （3）预测噪声 ${ϵ}_{\theta ,\varphi }(z_t^{(B)},t,c^{(B)},z^{(A)})$ 以恢复 $z_t^{(B)}\to z_{t-1}^{(B)}$
  （4） 最小化损失 ${\mathcal{L}}_{ϵ}$

2.2.3、提出的辅助模型 - 利用疾病先验知识

  AD 相关区域随时间而缩小或扩大，且速率不同。基于深度学习的时空模型试图以黑盒方式直接从脑部 MRI 中学习这些进展速率，这可能非常具有挑战性。为了帮助这一过程，本文提出将有关体积变化的先验知识直接整合到流程中。

  为此，本文利用辅助模型 $f_{\psi }$（图1中的D部分），能够预测与 AD 相关的区域体积随时间的变化，并通过进展相关的协变量 $v$ 将这些信息提供给 LDM。辅助模型的选择针对两种情况，使 BrLP 在横断面和纵向数据中都具有灵活性。

  对于在年龄 $A$ 时有单次扫描的受试者，本文采用回归模型来估计年龄 $B$ 时的体积变化 ${\hat{v}}^{(B)}=f_{\psi }(c^{(A)})$。对于可以访问在年龄 $A_1,...,A_n$ 时的 $n$ 次过去访问的受试者，使用疾病进程映射 (DCM) 来预测 ${\hat{v}}^{(B)}=f_{\psi }(c^{(A_1)},...,c^{(A_n)})$，该模型专门用于疾病进展。DCM 旨在根据受试者可用的体积变化历史提供更准确的轨迹。虽然使用 DCM 作为潜在的解决方案，但任何合适的疾病进展模型都可用于 BrLP。

2.2.4、推理过程

  将 $x^{(A)}$ 定义为年龄为 $A$ 的受试者的输入脑部 MRI，已知受试者特定元数据 $s^{(A)}$ 和从 $x^{(A)}$ 测量的进展相关体积 $v^{(A)}$。如图1中的E部分总结，为了推断年龄 $B>A$ 时的脑部 MRI $x^{(B)}$，执行 6 个步骤：

  （1）使用辅助模型预测进展相关体积 ${\hat{v}}^{(B)}=f_{\psi }(c^{(A)})$；
  （2）将此信息与受试者特定的元数据 $s^{(B)}$ 连接，形成目标协变量 $c^{(B)}=⟨s^{(B)},{\hat{v}}^{(B)}⟩$；
  （3）计算潜在表示 $z^{(A)}=\mathcal{E}(x^{(A)})$；
  （4）对随机高斯噪声 $z_T\sim \mathcal{N}(0,I)$ 进行采样；
  （5）通过预测噪声 ${ϵ}_{\theta ,\varphi }(z_t,t,c^{(B)},z^{(A)})$ 来运行反向扩散过程，以反转每个扩散步骤 $t=T,\dots ,1$；
  （6）使用解码器 $\mathcal{D}$ 在图像域中重建预测的脑 MRI $x^{(B)}=\mathcal{D}(z_0)$；

  这个推理过程被总结为一个紧凑的符号 ${\hat{z}}^{(B)}=\mathcal{I}(z_T,x^{(A)},c^{(A)})$ 和 ${\hat{x}}^{(B)}=\mathcal{D}({\hat{z}}^{(B)})$。

2.2.5、通过提出的潜在平均稳定化（LAS）增强推理

  初始值 $x_T\sim \mathcal{N}(0,I)$ 的变化可能导致推理过程产生的结果出现轻微差异。当对连续时间步进行预测时，这些差异尤为明显，表现为不规则的模式或进程的非平滑过渡。因此，本文引入了 LAS（图1中的F部分），这是一种通过平均推理过程的不同结果来提高时空一致性的技术。

  特别地，LAS 基于这样的假设：预测值 ${\hat{z}}^{(B)}=\mathcal{I}(z_T,x^{(A)},c^{(A)})$ 偏离理论平均值 ${\mu }^{(B)}=\mathbb{E}[{\hat{z}}^{(B)}]$。为了估计期望值 ${\mu }^{(B)}$，本文建议重复推理过程 $m$ 次，并对结果进行平均：

  与之前类似，本文对预测的扫描进行解码 ${\hat{x}}^{(B)}=\mathcal{D}({\mu }^{(B)})$。整个推理过程（其中 $m=4$）在消费级 GPU 上每 MRI 需要 4.8 秒。

---

3、实验与结果

3.1、数据

  收集了一个大型数据集，包含来自 2,805 名受试者的 11,730 张 T1 加权脑部 MR I扫描图像，这些数据来自多个公开的纵向研究：ADNI 1/2/3/GO (1,990名受试者）、OASIS-3 (573名受试者）和 AIBL (242名受试者）。每位受试者至少有两张 MRI 图像，且每次扫描均在不同的访问中获取。

  所有数据集中均可以获得年龄、性别和认知状态。平均年龄为 74±7 岁，53% 的受试者为男性。根据最终访视结果，43.8% 的受试者被归类为正常对照组，25.7% 表现出或发展为轻度认知障碍，30.5% 表现出或发展为阿尔茨海默病。

  将数据随机分为训练集（80%）、验证集（5%）和测试集（15%），且各组之间没有重叠的受试者。验证集用于训练过程中的提前停止。每张脑部 MRI 图像均经过预处理，包括：N4 偏场校正、颅骨剥离、到 MNI 空间的仿射配准、强度归一化以及重采样至 1.5mm³。用于评估进展的相关协变量体积是使用 SynthSeg 2.0 计算得出，并以总脑体积的百分比表示，以考虑个体差异。

3.2、评价指标

  均方误差（MSE）和结构相似性指数（SSIM）用于评估扫描之间的图像相似性；
  与 AD 相关的区域（海马体、杏仁核、侧脑室、脑脊液和丘脑）的体积指标评估模型在追踪疾病进展方面的准确性；

3.3、消融实验与对比实验

  
Table 1 | 消融研究结果以及与基线方法的比较：预测体积中的 MAE（± SD）表示为总脑体积的百分比；

3.4、示例展示

  
Figure 2 | 将 70 岁 MCI 受试者 15 年的真实进展与 BrLP 和基线方法的预测进行比较：每种方法显示预测的 MRI（左）及其与受试者真实脑 MRI 的偏差（右）；

---

  疾病进展还是太细节的变化了，怎样提升扩散模型的细节处理呢（我又被拒稿了，难过/(ㄒoㄒ)/~~）