[论文解读] CoLa-Diff: Conditional Latent Diffusion Model for Multi-Modal MRI Synthesis

Lan Jiang, Mao Ye+3 more|arXiv (Cornell University)|Mar 24, 2023

Radiomics and Machine Learning in Medical Imaging被引用 4

一句话总结

CoLa-Diff 提出了一种新颖的条件潜在扩散模型，用于多模态 MRI 合成，该模型在潜在空间中运行，以减少内存使用并提高效率。通过整合基于脑区掩码的结构引导以及多条件平衡的自动权重自适应机制，该方法在合成缺失的 MRI 模态方面实现了最先进性能，具有更优的解剖学保真度和图像质量。

ABSTRACT

MRI synthesis promises to mitigate the challenge of missing MRI modality in clinical practice. Diffusion model has emerged as an effective technique for image synthesis by modelling complex and variable data distributions. However, most diffusion-based MRI synthesis models are using a single modality. As they operate in the original image domain, they are memory-intensive and less feasible for multi-modal synthesis. Moreover, they often fail to preserve the anatomical structure in MRI. Further, balancing the multiple conditions from multi-modal MRI inputs is crucial for multi-modal synthesis. Here, we propose the first diffusion-based multi-modality MRI synthesis model, namely Conditioned Latent Diffusion Model (CoLa-Diff). To reduce memory consumption, we design CoLa-Diff to operate in the latent space. We propose a novel network architecture, e.g., similar cooperative filtering, to solve the possible compression and noise in latent space. To better maintain the anatomical structure, brain region masks are introduced as the priors of density distributions to guide diffusion process. We further present auto-weight adaptation to employ multi-modal information effectively. Our experiments demonstrate that CoLa-Diff outperforms other state-of-the-art MRI synthesis methods, promising to serve as an effective tool for multi-modal MRI synthesis.

研究动机与目标

• 通过实现数据高效、多模态合成，解决临床实践中缺失 MRI 模态的挑战。
• 克服现有扩散模型在像素空间运行所导致的高内存消耗问题，并避免解剖结构的失真。
• 开发一种方法，通过条件学习有效融合多种 MRI 对比度（如 T1、T2、FLAIR），以合成目标缺失模态。
• 通过引入结构先验和自适应条件权重，提升模型的泛化能力与鲁棒性，以平衡多模态输入。
• 通过潜在空间操作与知识蒸馏，实现快速、高保真度的 MRI 合成，适用于临床部署。

提出的方法

• 利用变分自编码器（VAE）编码器和解码器在潜在空间中运行，压缩并重建 MRI 图像，从而降低内存消耗。
• 引入一种改进的潜在扩散网络，采用“相似协同滤波”机制，以减轻潜在空间去噪过程中的信息损失与噪声放大问题。
• 将脑区掩码作为结构先验，引导扩散过程，以保持生成结果的解剖一致性。
• 采用自动权重自适应机制，动态平衡多个输入模态（如 T1、T2、FLAIR）的影响，最大化相关特征的利用效率。
• 应用知识蒸馏技术以加速采样过程，实现在不损失合成质量的前提下提升推理速度。
• 采用基于多模态输入和结构掩码的条件 U-Net 架构，逐步在反向扩散过程中生成目标 MRI 模态。

实验结果

研究问题

• RQ1 基于扩散的模型是否能有效利用多种输入对比度，在多对一设置下合成缺失的 MRI 模态？
• RQ2 如何利用潜在空间降低多模态 MRI 合成中的内存使用量，并提升训练与推理效率？
• RQ3 结构先验（如脑区掩码）在多大程度上能提升生成 MRI 图像的解剖一致性？
• RQ4 在扩散过程中，如何有效平衡多个输入模态，以避免某一模态占据主导地位？
• RQ5 自动权重自适应能否提升模型利用相关多模态信息的能力，并增强在多样化输入组合下的泛化性能？

主要发现

• CoLa-Diff 在 BRATS 和 IXI 数据集上均达到最先进性能，PSNR 最高提升 6.01 dB，SSIM 最高提升 5.74%。
• 在 BRATS 数据集中，CoLa-Diff 从 T2+T1ce+FLAIR 输入合成 T1ce 模态，PSNR 达到 28.26±3.13 dB，SSIM 达到 93.65±3.02%，优于次佳模型 ProvoGAN，分别提升 0.81 dB 和 0.82%。
• 在 IXI 数据集中，CoLa-Diff 从 T1+T2 输入合成 PD 模态，PSNR 达到 32.24±2.95 dB，SSIM 达到 96.95±2.26%，优于最佳基线模型 Hi-Net，分别提升 0.45 dB 和 0.44 个百分点。
• 消融实验表明，自动权重自适应贡献最大，使 BRATS T1+T1ce+FLAIR→T2 任务的 PSNR 提升 1.9450 dB，SSIM 提升 4.0808%。
• 随着输入模态数量的增加，性能持续提升：在 BRATS 上，PSNR 从 T1 单一模态的 26.6355 dB 提升至 T1+T1ce+FLAIR 的 28.3126 dB；在 IXI 上，PSNR 从 32.1640 dB 提升至 32.8721 dB，证实了多模态信息的有效利用。
• 定性结果表明，该方法在保留复杂解剖结构（如脑沟和肿瘤边界）方面表现更优，误差图中的错误更少，优于基线方法。