Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors (2405.18782v2)
Abstract: Diffusion models (DMs) have recently shown outstanding capabilities in modeling complex image distributions, making them expressive image priors for solving Bayesian inverse problems. However, most existing DM-based methods rely on approximations in the generative process to be generic to different inverse problems, leading to inaccurate sample distributions that deviate from the target posterior defined within the Bayesian framework. To harness the generative power of DMs while avoiding such approximations, we propose a Markov chain Monte Carlo algorithm that performs posterior sampling for general inverse problems by reducing it to sampling the posterior of a Gaussian denoising problem. Crucially, we leverage a general DM formulation as a unified interface that allows for rigorously solving the denoising problem with a range of state-of-the-art DMs. We demonstrate the effectiveness of the proposed method on six inverse problems (three linear and three nonlinear), including a real-world black hole imaging problem. Experimental results indicate that our proposed method offers more accurate reconstructions and posterior estimation compared to existing DM-based imaging inverse methods.
- Plug-and-play methods for magnetic resonance imaging: Using denoisers for image recovery. IEEE Signal Processing Magazine, 37(1):105–116, 2020.
- Brian. D. O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12:313–326, 1982.
- Towards a theory of non-log-concave sampling:first-order stationarity guarantees for langevin monte carlo. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2896–2923. PMLR, 02–05 Jul 2022.
- Conditional image generation with score-based diffusion models. ArXiv, abs/2111.13606, 2021.
- Generative plug and play: Posterior sampling for inverse problems, 2023.
- Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends in Machine Learning, 3:1–122, 01 2011.
- Tweedie moment projected diffusions for inverse problems, 2023.
- Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
- Phase retrieval from coded diffraction patterns. Applied and Computational Harmonic Analysis, 39, 10 2013.
- Monte carlo guided diffusion for bayesian linear inverse problems. ArXiv, abs/2308.07983, 2023.
- Interferometric imaging directly with closure phases and closure amplitudes. The Astrophysical Journal, 857(1):23, apr 2018.
- Plug-and-play admm for image restoration: Fixed point convergence and applications. IEEE Transactions on Computational Imaging, PP, 05 2016.
- Improved analysis for a proximal algorithm for sampling. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2984–3014. PMLR, 02–05 Jul 2022.
- Ilvr: Conditioning method for denoising diffusion probabilistic models. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 14347–14356, 2021.
- Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11472–11481, June 2022.
- Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representations, 2023.
- Score-based diffusion models for accelerated mri. Medical Image Analysis, page 102479, 2022.
- Plug-and-play split gibbs sampler: embedding deep generative priors in bayesian inference, 2023.
- It has potential: Gradient-driven denoisers for convergent solutions to inverse problems. In Neural Information Processing Systems, 2021.
- The Event Horizon Telescope Collaboration. First m87 event horizon telescope results. iv. imaging the central supermassive black hole. The Astrophysical Journal Letters, 875(1):L4, apr 2019.
- The Event Horizon Telescope Collaboration. First m87 event horizon telescope results. v. physical origin of the asymmetric ring. The Astrophysical Journal Letters, 875(1):L5, apr 2019.
- Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024.
- Bradley Efron. Tweedie’s formula and selection bias. Journal of the American Statistical Association, 106:1602–1614, 12 2011.
- Improved dimension dependence of a proximal algorithm for sampling. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 1473–1521. PMLR, 12–15 Jul 2023.
- What’s in a prior? learned proximal networks for inverse problems. In The Twelfth International Conference on Learning Representations, 2024.
- Regularization by denoising: Bayesian model and langevin-within-split gibbs sampling, 2024.
- Score-based diffusion models as principled priors for inverse imaging. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10520–10531, October 2023.
- James R. Fienup. Phase retrieval algorithms: a comparison. Applied optics, 21 15:2758–69, 1982.
- A dual algorithm for the solution of nonlinear variational problems via finite element approximation. Computers & Mathematics With Applications, 2:17–40, 1976.
- Nonlinear image recovery with half-quadratic regularization. IEEE transactions on image processing : a publication of the IEEE Signal Processing Society, 4 7:932–46, 1995.
- Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840–6851. Curran Associates, Inc., 2020.
- Robust compressed sensing mri with deep generative priors. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 14938–14954. Curran Associates, Inc., 2021.
- Plug-and-play methods for integrating physical and learned models in computational imaging: Theory, algorithms, and applications. IEEE Signal Processing Magazine, 40:85–97, 01 2023.
- Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
- A style-based generator architecture for generative adversarial networks. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 4396–4405, 2018.
- Denoising diffusion restoration models. In Advances in Neural Information Processing Systems, 2022.
- SNIPS: Solving noisy inverse problems stochastically. Advances in Neural Information Processing Systems, 34:21757–21769, 2021.
- Bayesian imaging using plug & play priors: When langevin meets tweedie. SIAM Journal on Imaging Sciences, 15(2):701–737, 2022.
- Structured logconcave sampling with a restricted gaussian oracle. In Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 2993–3050. PMLR, 15–19 Aug 2021.
- Decoupled data consistency with diffusion purification for image restoration, 2024.
- Dolce: A model-based probabilistic diffusion framework for limited-angle ct reconstruction. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10498–10508, October 2023.
- Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
- Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
- A variational perspective on solving inverse problems with diffusion models. In The Twelfth International Conference on Learning Representations, 2024.
- Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. 2017 IEEE International Conference on Computer Vision (ICCV), pages 1799–1808, 2017.
- prdeep: Robust phase retrieval with a flexible deep network. In International Conference on Machine Learning, 2018.
- Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8162–8171. PMLR, 18–24 Jul 2021.
- The split gibbs sampler revisited: Improvements to its algorithmic structure and augmented target distribution. SIAM Journal on Imaging Sciences, 16(4):2040–2071, 2023.
- Probabilistic mass-mapping with neural score estimation. Astronomy & Astrophysics, 672, 12 2022.
- Solving linear inverse problems provably via posterior sampling with latent diffusion models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- Plug-and-play methods provably converge with properly trained denoisers. In International Conference on Machine Learning, 2019.
- Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, PP:1–14, 09 2022.
- Solving inverse problems with latent diffusion models via hard data consistency. In The Twelfth International Conference on Learning Representations, 2024.
- Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
- Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023.
- Loss-guided diffusion models for plug-and-play controllable generation. In Proceedings of the 40th International Conference on Machine Learning, ICML’23. JMLR.org, 2023.
- Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations, 2022.
- Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
- He Sun and Katherine L. Bouman. Deep probabilistic imaging: Uncertainty quantification and multi-modal solution characterization for computational imaging. In AAAI Conference on Artificial Intelligence (AAAI), 2021.
- An online plug-and-play algorithm for regularized image reconstruction. IEEE Transactions on Computational Imaging, PP:1–1, 01 2019.
- Provable probabilistic imaging using score-based generative priors. ArXiv, abs/2310.10835, 2023.
- Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9(86):2579–2605, 2008.
- Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Neural Information Processing Systems, 2019.
- Plug-and-play priors for model based reconstruction. pages 945–948, 12 2013.
- Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23:1661–1674, 2011.
- Split-and-augmented gibbs sampler—application to large-scale inference problems. IEEE Transactions on Signal Processing, 67(6):1648–1661, 2019.
- High-dimensional gaussian sampling: A review and a unifying approach based on a stochastic proximal point algorithm. SIAM Review, 64(1):3–56, 2022.
- Zero-shot image restoration using denoising diffusion null-space model. The Eleventh International Conference on Learning Representations, 2023.
- Practical and asymptotically exact conditional sampling in diffusion models. arXiv preprint arXiv:2306.17775, 2023.
- Online regularization by denoising with applications to phase retrieval. pages 3887–3895, 10 2019.
- Simba: Scalable inversion in optical tomography using deep denoising priors. IEEE Journal of Selected Topics in Signal Processing, 14(6):1163–1175, 2020.
- Diffir: Efficient diffusion model for image restoration. ICCV, 2023.
- Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction, 2024.
- On a class of gibbs sampling over networks. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 5754–5780. PMLR, 12–15 Jul 2023.
- Plug-and-play image restoration with deep denoiser prior. IEEE Transactions on Pattern Analysis and Machine Intelligence, PP:1–1, 06 2021.
- Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE Transactions on Image Processing, 26(7):3142–3155, 2017.
- Learning deep cnn denoiser prior for image restoration. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 2808–2817, 2017.
- Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023.
- Denoising diffusion models for plug-and-play image restoration. In IEEE Conference on Computer Vision and Pattern Recognition Workshops (NTIRE), 2023.