Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

Hyungjin Chung, Suhyeon Lee, Jong Chul Ye

TL;DR

This work introduces Decomposed Diffusion Sampling (DDS), a diffusion-model-based inverse problem solver that integrates Krylov subspace optimization with Tweedie-denoised tangent-space updates. By showing that the tangent space at a denoised sample can be represented as a Krylov subspace, DDS uses multi-step Conjugate Gradient updates to stay within the tangent space, eliminating the need for computationally expensive manifold-constrained gradients. The approach applies to both variance-preserving and variance-exploding diffusion setups and yields state-of-the-art reconstructions in multi-coil MRI and 3D CT while delivering dramatic speedups (up to 80x faster) over prior diffusion-based solvers. Extensive experiments demonstrate improved PSNR/SSIM and robustness across forward models, including non-Cartesian MRI, with open-source code provided.

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

Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

TL;DR

This work introduces Decomposed Diffusion Sampling (DDS), a diffusion-model-based inverse problem solver that integrates Krylov subspace optimization with Tweedie-denoised tangent-space updates. By showing that the tangent space at a denoised sample can be represented as a Krylov subspace, DDS uses multi-step Conjugate Gradient updates to stay within the tangent space, eliminating the need for computationally expensive manifold-constrained gradients. The approach applies to both variance-preserving and variance-exploding diffusion setups and yields state-of-the-art reconstructions in multi-coil MRI and 3D CT while delivering dramatic speedups (up to 80x faster) over prior diffusion-based solvers. Extensive experiments demonstrate improved PSNR/SSIM and robustness across forward models, including non-Cartesian MRI, with open-source code provided.

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
Paper Structure (40 sections, 6 theorems, 75 equations, 9 figures, 8 tables, 6 algorithms)

This paper contains 40 sections, 6 theorems, 75 equations, 9 figures, 8 tables, 6 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Krylov subspace methods
  4. Diffusion models
  5. DDIM
  6. Decomposed Diffusion Sampling
  7. Conditional diffusion for inverse problems
  8. Key observation
  9. Experiments
  10. Accelerated MRI
  11. 3D CT reconstruction
  12. Conclusion
  13. Ethics Statement
  14. Reproducibility Statement
  15. Preliminaries
  16. ...and 25 more sections

Key Result

proposition 1

Suppose the clean data manifold ${\mathcal{M}}$ is represented as an affine subspace and assumes the uniform distribution on ${\mathcal{M}}$. Then, for some $\zeta_t>0$, where ${\mathcal{P}}_{{\mathcal{M}}}$ denotes the orthogonal projection to ${\mathcal{M}}$.

Figures (9)

  • Figure 1: Parallel imaging MR reconstruction evaluation PSNR vs. NFE (log scale). Reconstruction from 1D uniform random $\times 4$ acceleration zbontar2018fastmri.
  • Figure 2: Representative reconstruction results. (a) Multi-coil MRI reconstruction, (b) 3D sparse-view CT. Numbers in parenthesis: NFE. Yellow numbers in bottom left corner: PSNR/SSIM.
  • Figure 3: Illustration of the imaging forward model used in this work. (a) 3D sparse-view CT: the forward matrix ${\boldsymbol A}$ transforms the 3D voxel space into 2D projections. (b) Multi-coil CS-MRI: the forward matrix ${\boldsymbol A}$ first applies Fourier transform to turn the image into $k$-space. Subsequently, sensitivity maps are applied as element-wise product to achieve multi-coil measurements. Finally, the multi-coil measurements are sub-sampled with the masks.
  • Figure 4: DDS reconstruction of CS-MRI on radial sampling trajectory. Col 1: sampling trajectory, 2: Zero filled reconstruction + density compensation, 3: DDS (99 NFE), 4: ground truth.
  • Figure 5: Evolution of the reconstruction error through time. $\pm 1.0 \sigma$ plot. (a) VE parameterized with ${\boldsymbol s}_\theta$, (b) VP parameterized with ${\boldsymbol \epsilon}_\theta$, (c) Visualization of $\hat{{\boldsymbol x}}_t$.
  • ...and 4 more figures

Theorems & Definitions (10)

  • proposition 1: Manifold Constrained Gradient
  • lemma 1: Tweedie's formula
  • proof
  • lemma 2: Total noise
  • proof
  • proposition 1: Manifold Constrained Gradient
  • proof
  • proposition 2: VE Decomposition
  • proof
  • proposition 3: VE-DDIM Decomposition