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

TL;DR

This paper introduces a novel method combining 2D diffusion models with model-based priors to effectively solve 3D inverse problems like medical image reconstruction, achieving high fidelity with low computational cost.

Contribution

The paper presents a new approach that extends 2D diffusion models to 3D inverse problems by integrating model-based priors, enabling high-quality reconstructions on a single GPU.

Findings

Achieves state-of-the-art reconstruction quality in 3D medical imaging tasks.

Operates efficiently on a single GPU, reducing computational requirements.

Demonstrates strong generalization to volumes outside training data.

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.