[论文解读] Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models

Hyungjin Chung, Dohoon Ryu+3 more|arXiv (Cornell University)|Nov 19, 2022

Advanced Neuroimaging Techniques and Applications被引用 24

一句话总结

提出 DiffusionMBIR，一种将二维扩散先验与基于模型的先验相结合的方法，在单个 GPU 上高效解决三维医学图像重建任务。

ABSTRACT

Diffusion models have emerged as the new state-of-the-art generative model with high quality samples, with intriguing properties such as mode coverage and high flexibility. They have also been shown to be effective inverse problem solvers, acting as the prior of the distribution, while the information of the forward model can be granted at the sampling stage. Nonetheless, as the generative process remains in the same high dimensional (i.e. identical to data dimension) space, the models have not been extended to 3D inverse problems due to the extremely high memory and computational cost. In this paper, we combine the ideas from the conventional model-based iterative reconstruction with the modern diffusion models, which leads to a highly effective method for solving 3D medical image reconstruction tasks such as sparse-view tomography, limited angle tomography, compressed sensing MRI from pre-trained 2D diffusion models. In essence, we propose to augment the 2D diffusion prior with a model-based prior in the remaining direction at test time, such that one can achieve coherent reconstructions across all dimensions. Our method can be run in a single commodity GPU, and establishes the new state-of-the-art, showing that the proposed method can perform reconstructions of high fidelity and accuracy even in the most extreme cases (e.g. 2-view 3D tomography). We further reveal that the generalization capacity of the proposed method is surprisingly high, and can be used to reconstruct volumes that are entirely different from the training dataset.

研究动机与目标

• 激励在不训练三维先验的情况下，使用扩散模型解决三维反问题。
• 利用二维预训练扩散模型重建连贯的三维体积。
• 将 MBIR 风格的正则化通过 ADMM 与扩散去噪结合，以在切片之间实现数据一致性。

提出的方法

• 逐切片使用二维扩散模型进行去噪，通过三维数据一致性保持 z 轴的一致性。
• 通过对三维体积的 ADMM 更新，使用 z 方向的 TV 先验来增强扩散先验。
• 将重构表述为一个带数据保真项和 z 方向正则化的三维 MBIR 风格优化。
• 用共轭梯度步和软阈值化来求解 ADMM 子问题，实现 z 方向的稀疏性。
• 使用一种快速版本，变量共享，在 SDE 步之间复用 ADMM 状态以降低计算。

实验结果

研究问题

• RQ1 能否在不训练三维扩散模型的情况下，使用预训练的二维扩散模型来解决三维反问题？
• RQ2 将二维扩散先验与三维 MBIR 风格的 z 方向正则化相结合，是否在所有轴上产生连贯的三维重建？
• RQ3 提出的 DiffusionMBIR 是否足够内存高效，能够在单个普通 GPU 上完成如 SV-CT、LA-CT 和 CS-MRI 之类的任务？
• RQ4 相较于现有的基于扩散的方法和完全监督的方法，DiffusionMBIR 在分布内和分布外的三维医学成像数据上的表现如何。

主要发现

• DiffusionMBIR 在稀疏视角 CT、有限角度 CT 和压缩感知 MRI 上实现了最先进的重建。
• 该方法在轴位、矢状位和冠状位上实现连贯的三维重建。
• DiffusionMBIR 保持高 PSNR 和 SSIM，并显示出对分布外的强泛化能力。
• 一种极简迭代（最少 ADMM/CG 迭代，M=1，K=1）的快速变体在较低计算量下提供高保真结果。