Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

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 combine 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.

[1]  A. Dimakis,et al.  Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models , 2023, ArXiv.

[2]  R. Fergus,et al.  Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC , 2023, ICML.

[3]  Mengdi Wang,et al.  Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-Dimensional Data , 2023, ICML.

[4]  S. Shakkottai,et al.  A Theoretical Justification for Image Inpainting using Denoising Diffusion Probabilistic Models , 2023, ArXiv.

[5]  Yinhuai Wang,et al.  Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model , 2022, ICLR.

[6]  Michael T. McCann,et al.  Solving 3D Inverse Problems Using Pre-Trained 2D Diffusion Models , 2022, 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[7]  Michael T. McCann,et al.  Diffusion Posterior Sampling for General Noisy Inverse Problems , 2022, ICLR.

[8]  Ben Poole,et al.  DreamFusion: Text-to-3D using 2D Diffusion , 2022, ICLR.

[9]  J. C. Ye,et al.  Improving Diffusion Models for Inverse Problems using Manifold Constraints , 2022, Neural Information Processing Systems.

[10]  Tero Karras,et al.  Elucidating the Design Space of Diffusion-Based Generative Models , 2022, NeurIPS.

[11]  M. Uecker,et al.  Bayesian MRI reconstruction with joint uncertainty estimation using diffusion models , 2022, Magnetic resonance in medicine.

[12]  Michael Elad,et al.  Denoising Diffusion Restoration Models , 2022, NeurIPS.

[13]  Jeffrey A. Fessler,et al.  Sparse-View Cone Beam CT Reconstruction Using Data-Consistent Supervised and Adversarial Learning From Scarce Training Data , 2022, IEEE Transactions on Computational Imaging.

[14]  Jong-Chul Ye,et al.  Come-Closer-Diffuse-Faster: Accelerating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction , 2021, 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[15]  S. Ermon,et al.  Solving Inverse Problems in Medical Imaging with Score-Based Generative Models , 2021, ICLR.

[16]  Jong-Chul Ye,et al.  Score-based diffusion models for accelerated MRI , 2021, Medical Image Anal..

[17]  Alexandros G. Dimakis,et al.  Robust Compressed Sensing MRI with Deep Generative Priors , 2021, NeurIPS.

[18]  Il-Chul Moon,et al.  Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation , 2021, ICML.

[19]  Prafulla Dhariwal,et al.  Diffusion Models Beat GANs on Image Synthesis , 2021, NeurIPS.

[20]  Abhishek Kumar,et al.  Score-Based Generative Modeling through Stochastic Differential Equations , 2020, ICLR.

[21]  Jiaming Song,et al.  Denoising Diffusion Implicit Models , 2020, ICLR.

[22]  Matteo Ronchetti,et al.  TorchRadon: Fast Differentiable Routines for Computed Tomography , 2020, ArXiv.

[23]  Pieter Abbeel,et al.  Denoising Diffusion Probabilistic Models , 2020, NeurIPS.

[24]  Aaron Defazio,et al.  End-to-End Variational Networks for Accelerated MRI Reconstruction , 2020, MICCAI.

[25]  Yang Song,et al.  Generative Modeling by Estimating Gradients of the Data Distribution , 2019, NeurIPS.

[26]  Pascal Vincent,et al.  fastMRI: An Open Dataset and Benchmarks for Accelerated MRI , 2018, ArXiv.

[27]  Michael Unser,et al.  Deep Convolutional Neural Network for Inverse Problems in Imaging , 2016, IEEE Transactions on Image Processing.

[28]  Lei Zhang,et al.  Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising , 2016, IEEE Transactions on Image Processing.

[29]  Guangyong Chen,et al.  An Efficient Statistical Method for Image Noise Level Estimation , 2015, 2015 IEEE International Conference on Computer Vision (ICCV).

[30]  Michael Elad,et al.  ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: Where SENSE meets GRAPPA , 2014, Magnetic resonance in medicine.

[31]  Stephen P. Boyd,et al.  Proximal Algorithms , 2013, Found. Trends Optim..

[32]  Z. Strakos,et al.  Krylov Subspace Methods: Principles and Analysis , 2012 .

[33]  B. Efron Tweedie’s Formula and Selection Bias , 2011, Journal of the American Statistical Association.

[34]  Pascal Vincent,et al.  A Connection Between Score Matching and Denoising Autoencoders , 2011, Neural Computation.

[35]  D. Donoho,et al.  Sparse MRI: The application of compressed sensing for rapid MR imaging , 2007, Magnetic resonance in medicine.

[36]  K. T. Block,et al.  Undersampled radial MRI with multiple coils. Iterative image reconstruction using a total variation constraint , 2007, Magnetic resonance in medicine.

[37]  Aapo Hyvärinen,et al.  Estimation of Non-Normalized Statistical Models by Score Matching , 2005, J. Mach. Learn. Res..

[38]  D. Keyes,et al.  Jacobian-free Newton-Krylov methods: a survey of approaches and applications , 2004 .

[39]  Jeffrey A. Fessler,et al.  Nonuniform fast Fourier transforms using min-max interpolation , 2003, IEEE Trans. Signal Process..

[40]  J. Kautz,et al.  Pseudoinverse-Guided Diffusion Models for Inverse Problems , 2023, ICLR.