当前位置：首页 > news >正文

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

news 来源：原创 2025/5/7 10:22:10

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

论文：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$ 生成一个潜在表示 $\mathcal E(x)$ 。随后训练一个条件 UNet（图1中的B部分），表示为 $ϵ_θ$ ，网络参数为 $θ$ ，旨在估计从 $z_t$ 恢复到 $z_{t-1}$ 所需的噪声 $ϵ_θ(z_t,t,c)$ ，通过最小化损失函数 $\mathcal L_ϵ$ 来训练 $ϵ_θ$ 。协变量 $c$ 作为条件通过交叉注意力机制集成到网络中。

生成过程首先从采样随机高斯噪声 $z_T∼\mathcal N(0,I)$ 开始，然后迭代地逆向执行每个扩散步骤 $z_t→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 概念化为一个统一的网络 $ϵ_{θ,ϕ}$ ，其中 $θ$ 表示 LDM 的固定网络参数， $ϕ$ 表示 ControlNet 的可训练网络参数。与 LDM 一样， $ϵ_{θ,ϕ}$ 仍用于预测反向扩散步骤 $z_t→z_{t−1}$ 中的噪声 $ϵ_{θ,ϕ}(z_t,t,c,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)}→z^{(B)}_t$
（3）预测噪声 $ϵ_{θ,ϕ}(z^{(B)}_t,t,c^{(B)},z^{(A)})$ 以恢复 $z^{(B)}_t→z^{(B)}_{t-1}$
（4）最小化损失 $\mathcal L_ϵ$

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

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

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

对于在年龄 $A$ 时有单次扫描的受试者，本文采用回归模型来估计年龄 $B$ 时的体积变化 $\hat v^{(B)} = f_ψ(c^{(A)})$ 。对于可以访问在年龄 $A_1, ..., A_n$ 时的 $n$ 次过去访问的受试者，使用疾病进程映射 (DCM) 来预测 $\hat v^{(B)} = f_ψ(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_ψ(c^{(A)})$ ；
（2）将此信息与受试者特定的元数据 $s^{(B)}$ 连接，形成目标协变量 $c^{(B)} =⟨s^{(B)},\hat v^{(B)}⟩$ ；
（3）计算潜在表示 $z^{(A)} = \mathcal E(x^{(A)})$ ；
（4）对随机高斯噪声 $z_T∼\mathcal N(0,I)$ 进行采样；
（5）通过预测噪声 $ϵ_{θ,ϕ}(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∼\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 计算得出，并以总脑体积的百分比表示，以考虑个体差异。