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

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.

14 pages, 10 figures

References in corpus (4)

• Diffusion Models Beat GANs on Image Synthesis
• Score-Based Generative Modeling through Stochastic Differential Equations
• Improving Diffusion Models for Inverse Problems using Manifold Constraints
• Image Prediction for Limited-angle Tomography via Deep Learning with Convolutional Neural Network

Cited by in corpus (8)

• CoreDiff: Contextual Error-Modulated Generalized Diffusion Model for Low-Dose CT Denoising and Generalization
• BSDM: Background Suppression Diffusion Model for Hyperspectral Anomaly Detection
• Physics-informed Score-based Diffusion Model for Limited-angle Reconstruction of Cardiac Computed Tomography
• Score-Based Generative Models for PET Image Reconstruction
• 3DGR-CT: Sparse-View CT Reconstruction with a 3D Gaussian Representation
• Geometry-Aware Attenuation Learning for Sparse-View CBCT Reconstruction
• Accurate Patient Alignment without Unnecessary Imaging Dose via Synthesizing Patient-specific 3D CT Images from 2D kV Images
• MRI Reconstruction with Regularized 3D Diffusion Model (R3DM)