【论文笔记】Enhancing Spatiotemporal Disease Progression Models via Latent Diffusion and Prior Knowledge

Abstract

本文提出了Brain Latent Progression (BrLP)，这是一种基于潜变量扩散的全新时空疾病进展模型(spatiotemporal disease progression model)，旨在预测个体层面的3D脑部MRI疾病演变。
现有用于该任务的深度生成模型主要依赖数据驱动，在学习疾病进展方面存在挑战。BrLP通过结合疾病模型的先验知识来提升预测准确性，从而克服这些问题。
为实现这一目标，本文提出集成一个辅助模型(auxiliary model)，以推断不同脑区的体积变化。此外，我们引入Latent Average Stabilization (LAS)，用于提高预测进展的时空一致性。
本文在实验中，分别在横断面（单次扫描）和纵向（随访扫描）两种场景下，将BrLP生成的MRI扫描与受试者的真实随访MRI进行比较。结果表明，BrLP在AD相关脑区的体积预测准确性上提高了22%，在与真实扫描的图像相似度上提高了43%，相较于现有方法取得了显著提升。
BrLP能够在个体层面生成条件3D扫描，并首次将先验知识融入模型以提升预测精度，这一突破代表了疾病进展建模的重大进展，为精准医学开辟了新的研究方向。

1 Introduction

神经退行性疾病已成为全球健康挑战，影响数百万人，并导致广泛的疾病负担和高死亡率。随着全球人口老龄化加剧，这一问题给医疗系统和社会带来了更大压力。因此需要开发新的方法，以更深入地理解疾病的发展过程，从而为更有针对性和个性化的治疗策略奠定基础。

最初的研究主要基于标量生物标志物进行疾病进展建模。尽管这些标志物的表达较为粗略，但此类方法仍然有助于加深对疾病的理解。
随着研究的自然演进，空间-时间模型逐步发展，通过丰富的高维影像生物标志物直接作用于医学扫描数据来表征疾病进展。与标量生物标志物不同，这些方法可以可视化并精确定位结构变化的复杂模式，从而提供更详细的疾病动态分析。

近年来，研究人员利用深度生成技术，如VAE、GAN和扩散模型，以个体层面推断疾病进展。其中，DaniNet是一种最先进的模型，它结合了对抗学习和生物学约束，以提供个性化的脑部MRI预测。为降低内存需求，DaniNet通过生成2D切片并利用超分辨率模块将其组装成3D体数据。另一个方法是CounterSynth，它是一种基于GAN的反事实合成方法，可在脑部MRI中模拟衰老和疾病进展等不同状态。此外，SADM是一种基于扩散模型的方法，采用自回归采样机制，通过一系列先前的MRI扫描生成纵向影像。
DaniNet
Degenerative adversarial neuroimage nets for brain scan simulations: Application in ageing and dementia
CounterSynth
Equitable modelling of brain imaging by counterfactual augmentation with morphologically constrained 3d deep generative models
SADM
Sadm: Sequence-aware diffusion model for longitudinal medical image generation

这些方法面临的主要挑战包括：
1）通过结合个体特定的元数据来提高个性化程度；
2）在可用时充分利用纵向扫描数据；
3）增强时空一致性，以确保在空间和时间维度上实现平滑的疾病进展；
4）应对高分辨率 3D 医学影像带来的高内存需求。
具体而言，DaniNet和CounterSynth在有纵向数据的情况下仍无法直接利用，而 SADM 无法结合个体特定的元数据，同时对内存的需求较高。此外，CounterSynth和SADM均未提供保证时空一致性的解决方案。

针对这些挑战，本文提出BrLP，这是一种新的时空模型，并作出以下关键贡献：
i）结合潜在扩散模型（LDM）和 ControlNet 生成基于个体数据的个性化脑部 MRI，解决挑战 1；
ii）通过引入辅助模型推断不同脑区的体积变化，以结合疾病进展的先验知识，使得在有纵向数据时能够充分利用，解决挑战 2；
iii）提出 LAS 技术，以提高预测进展的时空一致性，解决挑战 3；
iv）利用脑部 MRI 的潜在表示来降低 3D 扫描处理的内存需求，解决挑战 4。

3 Methods - Brain Latent Progression (BrLP)

图1：BrLP训练和推理过程概述。

BrLP架构由四个关键组件组成：LDM、ControlNet、辅助模型和 LAS 模块。

这四个组件的作用在图1中进行了总结，共同解决了引言中提出的挑战。
具体而言：
LDM旨在生成符合特定协变量的随机3D脑部MRI；
而ControlNet则用于将这些MRI适配到个体的特定解剖结构；
辅助模型利用疾病进展的先验知识，提高对特定脑区体积变化的预测精度；
在推理过程中使用 LAS 模块，以增强时空一致性。

LDM - Learning the Brain MRIs Distribution

基于前者的工作(Brain imaging generation with latent diffusion models)，本文训练了一个 LDM，用于生成符合特定协变量$c=⟨s,v⟩$的 3D 脑部 MRI。其中，$s$包括个体特定的元数据（年龄、性别和认知状态），而$v$则包含与AD进展相关的指标，如海马体、大脑皮层、杏仁核、大脑白质和侧脑室的体积。

LDM 的构建分为两个阶段。
1）训练一个自编码器$(\mathcal{E},\mathcal{D})$（图1中的A模块），用于为数据集中的每个脑部 MRI $x$ 生成潜在表示$z=\mathcal{E}(x)$。
2）训练一个条件UNet（图1中的B模块），带有网络参数$\theta$，表示为${ϵ}_{\theta }$，用于估计从$z_t$逆推至$z_{t-1}$所需的噪声${ϵ}_{\theta }(z_t,t,c)$。通过最小化损失函数$L_{ϵ}$(Eq.1)来训练${ϵ}_{\theta }$：
$${\mathcal{L}}_{ϵ}:={\mathbb{E}}_{t,x_t,ϵ\sim \mathcal{N}(0,I)}[\Vert ϵ-{ϵ}_{\theta }(x_t,t){\Vert }^2]\text{\hspace{1em}(1)}$$
协变量$c$作为条件信息通过交叉注意力机制(cross-attention)整合到网络中。生成过程从随机高斯噪声$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}=D(z_0)$。

ControlNet - Conditioning on Subject Brain MRI

本文采用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)$作为条件，以在生成过程中包含目标脑部 MRI $x$的结构信息。
为了训练ControlNet，使用来自同一患者、不同年龄（$A<B$）的两张脑部MRI组成的配对数据，分别提取其潜在表示$z^{(A)}$和$z^{(B)}$。其中，$z^{(B)}$关联的协变量$c^{(B)}$是已知的，并用作目标协变量。训练的每次迭代包括以下步骤：
i) 采样扩散步数$t\sim U[1,T]$；
ii) 对$z^{(B)}$执行$t$轮前向扩散，得到$z_t^{(B)}$；
iii) 预测噪声${ϵ}_{\theta ,\varphi }(z_t^{(B)},t,c^{(B)},z^{(A)})$以逆向还原$z_t^{(B)}\to z_{t-1}^{(B)}$；
iv) 最小化损失函数$L_{ϵ}$ (Eq.1)。

Proposed Auxiliary Model - Leveraging Disease Prior Knowledge

AD相关脑区会随时间以不同速率萎缩或膨胀。基于深度学习的时空建模方法通常试图通过黑箱方式直接从脑部MRI中学习这些进展速率，但这极具挑战性。为此，本文提出将体积变化的先验知识直接融入到BrLP的模型流程中。
具体而言，引入一个辅助模型$f_{\psi }$（图1中的D模块），用于预测AD相关脑区的体积如何随时间变化，并将该信息作为进展相关协变量$v$传递给LDM。这一辅助模型能够适应不同的数据情况，使BrLP既适用于横断面数据，也适用于纵向数据。

• 对于仅有单次扫描（年龄$A$）的个体，使用回归模型来估计其在未来年龄$B$时的体积变化，即$\hat{v}(B)=f_{\psi }(c^{(A)})$。
• 对于具有$n$次历史扫描（年龄$A_1,\dots ,A_n$）的个体，使用疾病进程映射（DCM） 预测$\hat{v}(B)=f_{\psi }(c^{(A_1)},\dots ,c^{(A_n)})$。DCM 专为疾病进展建模而设计，可根据个体过去的体积变化趋势提供更准确的轨迹预测。
  虽然在BrLP中采用DCM作为示例方案，但该框架支持替换为任何适合的疾病进展模型，从而提供更大的灵活性。

Inference Process

设$x^{(A)}$为个体在年龄$A$时的输入脑部MRI，已知该个体的特定元数据$s^{(A)}$以及从$x^{(A)}$中测量得到的进展相关体积$v^{(A)}$。如图1中的E模块所示，为了推断该个体在年龄$B>A$时的脑部 MRI $x^{(B)}$，执行以下六个步骤：
i) 使用辅助模型预测进展相关体积：${\hat{v}}^{(B)}=f_{\psi }(c^{(A)})$；
ii) 将该信息与个体特定元数据$s^{(B)}$连接，形成目标协变量$c^{(B)}=⟨s^{(B)},{\hat{v}}^{(B)}⟩$；
iii) 计算输入MRI的潜在表示：$z^{(A)}=\mathcal{E}(x^{(A)})$；
iv) 采样随机高斯噪声：$z_T\sim \mathcal{N}(0,I)$；
v) 运行反向扩散过程，通过预测噪声 ${ϵ}_{\theta ,\varphi }(z_t,t,c^{(B)},z^{(A)})$ 逐步逆转扩散步骤$z_t\to z_{t-1}$（对于$t=T,\dots ,1$）；
vi) 使用解码器$\mathcal{D}$在图像域中重建预测的脑部 MRI：${\hat{x}}^{(B)}=\mathcal{D}(z_0)$。

该推理过程可被简化为紧凑表示：$\hat{z}(B)=\mathcal{I}(z_T,x^{(A)},c^{(A)})$，以及最终的MRI预测结果 ${\hat{x}}^{(B)}=\mathcal{D}({\hat{z}}^{(B)})$。

Enhance Inference via Proposed Latent Average Stabilization (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$次，并对结果取均值：
$${\mu }^{(B)}={\mathbb{E}}_{z_T\sim \mathcal{N}(0,I)}[\mathcal{I}(z_T,x^{(A)},c^{(A)})]\approx \frac{1}{m}\sum ^m\mathcal{I}(z_T,x^{(A)},c^{(A)})\text{\hspace{1em}(2)}$$

与之前的方法类似，我们最终通过解码得到预测的脑部扫描图像：${\hat{x}}^{(B)}=D({\mu }^{(B)})$。完整的推理过程（设$m=4$）在消费级GPU上每张MRI需要约4.8秒。

4 Experiments and Results

Data

本文收集了一个大型数据集，包括11,730例T1加权脑部MRI扫描，来自2,805名受试者，这些数据涵盖了多个公开的纵向研究：ADNI 1/2/3/GO（1,990名受试者）、OASIS-3（573 名受试者） 和AIBL（242 名受试者）。每名受试者至少有两次MRI扫描，每次扫描均在不同的访视期间进行。首次与随访扫描之间的平均时间间隔为 4.3 年（标准差 = 3.1），最长跨度达 16 年。所有数据集均提供了受试者的年龄、性别和认知状态信息。

受试者的平均年龄为74 ± 7岁，其中53%为男性。根据最终访视结果，43.8%的受试者被归类为认知正常（CN），25.7%表现出或发展为轻度认知障碍（MCI），30.5%表现出或发展为阿尔茨海默病（AD）。数据集被随机划分为训练集（80%）、验证集（5%）和测试集（15%），确保不同子集中无重复受试者。验证集用于训练过程中的早停（early stopping）。

所有脑部 MRI 经过以下预处理步骤：N4偏场校正、颅骨剥离、仿射配准至 MNI 空间、强度归一化，以及重采样至1.5 mm³。用于表征疾病进展的体积特征及后续评估指标由 SynthSeg 2.0计算，并以总脑容量的百分比形式表达，以考虑个体差异。

Evaluation Metrics

本文使用基于图像和体积的指标来评估 BrLP，以比较预测的脑部 MRI 扫描与受试者的实际随访扫描之间的差异。具体而言，采用均方误差（MSE）和结构相似性指数（SSIM）来衡量扫描图像之间的相似性。
此外，在阿尔茨海默病相关区域（如海马、杏仁核、侧脑室、脑脊液（CSF）和丘脑）中，使用体积指标评估模型在追踪疾病进展方面的准确性。具体来说，报告了实际随访扫描与生成的脑部 MRI 之间的体积均绝对误差（MAE）。值得注意的是，CSF 和丘脑未被纳入进展相关的协变量，这使得能够分析预测中未受条件约束的脑区。