论文标题

使用预训练的2D扩散模型解决3D反问题

Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models

论文作者

Chung, Hyungjin, Ryu, Dohoon, McCann, Michael T., Klasky, Marc L., Ye, Jong Chul

论文摘要

扩散模型已成为具有高质量样本的新最先进的生成模型,具有吸引人的特性,例如模式覆盖范围和高灵活性。它们还被证明是有效的反问题解决者,是分布的先验,而可以在采样阶段授予正向模型的信息。但是,由于生成过程保持在相同的高维度(即与数据维度相同)空间中,因此由于非常高的内存和计算成本,模型尚未扩展到3D反问题。在本文中,我们将基于传统模型的迭代重建的想法与现代扩散模型相结合,这导致了一种高效的方法来解决3D医疗图像重建任务,例如稀疏视图,有限的角度层析成像,有限的角度层析成像,来自预先训练的2D扩散模型的压缩感测MRI。从本质上讲,我们建议在测试时间剩余方向上以基于模型的先验增强2D扩散,以便可以在所有维度上实现连贯的重建。我们的方法可以在单个商品GPU中运行,并建立新的最先进的方法,表明所提出的方法即使在最极端的情况下也可以执行高保真和准确性的重建(例如,2视图3D层析成像)。我们进一步揭示了所提出的方法的概括能力令人惊讶地很高,可用于重建与培训数据集完全不同的体积。

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.

扫码加入交流群

加入微信交流群

微信交流群二维码

扫码加入学术交流群,获取更多资源