Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

TL;DR

This paper introduces a novel diffusion sampling method that combines Krylov subspace techniques with diffusion models, significantly improving efficiency and reconstruction quality in large-scale inverse problems like MRI and CT imaging.

Contribution

The authors propose a new diffusion sampling strategy leveraging Krylov subspace methods, eliminating the need for manifold-constrained gradients and achieving state-of-the-art results.

Findings

Achieves over 80x faster inference than previous methods.

Provides state-of-the-art reconstruction quality in medical imaging.

Applicable across different diffusion parametrizations and settings.

Abstract

Krylov subspace, which is generated by multiplying a given vector by the matrix of a linear transformation and its successive powers, has been extensively studied in classical optimization literature to design algorithms that converge quickly for large linear inverse problems. For example, the conjugate gradient method (CG), one of the most popular Krylov subspace methods, is based on the idea of minimizing the residual error in the Krylov subspace. However, with the recent advancement of high-performance diffusion solvers for inverse problems, it is not clear how classical wisdom can be synergistically combined with modern diffusion models. In this study, we propose a novel and efficient diffusion sampling strategy that synergistically combines the diffusion sampling and Krylov subspace methods. Specifically, we prove that if the tangent space at a denoised sample by Tweedie's formula forms a Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update to remain in the tangent space. This negates the need to compute the manifold-constrained gradient (MCG), leading to a more efficient diffusion sampling method. Our method is applicable regardless of the parametrization and setting (i.e., VE, VP). Notably, we achieve state-of-the-art reconstruction quality on challenging real-world medical inverse imaging problems, including multi-coil MRI reconstruction and 3D CT reconstruction. Moreover, our proposed method achieves more than 80 times faster inference time than the previous state-of-the-art method. Code is available at https://github.com/HJ-harry/DDS