Improving Diffusion Models for Inverse Problems using Manifold Constraints

Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce suboptimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin. With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available https://github.com/HJ-harry/MCG_diffusion

[1]  David J. Fleet,et al.  Video Diffusion Models , 2022, NeurIPS.

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

[3]  L. Gool,et al.  RePaint: Inpainting using Denoising Diffusion Probabilistic Models , 2022, 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[4]  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).

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

[6]  Victor Lempitsky,et al.  Resolution-robust Large Mask Inpainting with Fourier Convolutions , 2021, 2022 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV).

[7]  Youngjune Gwon,et al.  ILVR: Conditioning Method for Denoising Diffusion Probabilistic Models , 2021, 2021 IEEE/CVF International Conference on Computer Vision (ICCV).

[8]  Kwanyoung Kim,et al.  Noise2Score: Tweedie's Approach to Self-Supervised Image Denoising without Clean Images , 2021, NeurIPS.

[9]  Michael Elad,et al.  SNIPS: Solving Noisy Inverse Problems Stochastically , 2021, NeurIPS.

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

[11]  Baining Guo,et al.  Aggregated Contextual Transformations for High-Resolution Image Inpainting , 2021, IEEE Transactions on Visualization and Computer Graphics.

[12]  Jing Liao,et al.  High-Fidelity Pluralistic Image Completion with Transformers , 2021, 2021 IEEE/CVF International Conference on Computer Vision (ICCV).

[13]  Dong Liu,et al.  Generating Diverse Structure for Image Inpainting With Hierarchical VQ-VAE , 2021, 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[14]  Valentin De Bortoli,et al.  Bayesian imaging using Plug & Play priors: when Langevin meets Tweedie , 2021, SIAM J. Imaging Sci..

[15]  Alexandros G. Dimakis,et al.  Intermediate Layer Optimization for Inverse Problems using Deep Generative Models , 2021, ICML.

[16]  Aggelos K. Katsaggelos,et al.  2-Step Sparse-View CT Reconstruction with a Domain-Specific Perceptual Network , 2020, ArXiv.

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

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

[19]  Alain Durmus,et al.  Maximum Likelihood Estimation of Regularization Parameters in High-Dimensional Inverse Problems: An Empirical Bayesian Approach. Part II: Theoretical Analysis , 2020, SIAM J. Imaging Sci..

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

[21]  Valentin De Bortoli,et al.  Maximum Likelihood Estimation of Regularization Parameters in High-Dimensional Inverse Problems: An Empirical Bayesian Approach Part I: Methodology and Experiments , 2019, SIAM J. Imaging Sci..

[22]  I. I. Rushakov,et al.  Computed Tomography , 2019, Compendium of Biomedical Instrumentation.

[23]  Ullrich Köthe,et al.  Guided Image Generation with Conditional Invertible Neural Networks , 2019, ArXiv.

[24]  Arno Solin,et al.  Applied Stochastic Differential Equations , 2019 .

[25]  Raja Giryes,et al.  Image-Adaptive GAN based Reconstruction , 2019, AAAI.

[26]  Ali Razavi,et al.  Generating Diverse High-Fidelity Images with VQ-VAE-2 , 2019, NeurIPS.

[27]  Timo Aila,et al.  A Style-Based Generator Architecture for Generative Adversarial Networks , 2018, 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[28]  W. Clem Karl,et al.  Deep Learning-Based Sinogram Completion for Low-Dose CT , 2018, 2018 IEEE 13th Image, Video, and Multidimensional Signal Processing Workshop (IVMSP).

[29]  Peyman Milanfar,et al.  Local Kernels That Approximate Bayesian Regularization and Proximal Operators , 2018, IEEE Transactions on Image Processing.

[30]  Alexei A. Efros,et al.  The Unreasonable Effectiveness of Deep Features as a Perceptual Metric , 2018, 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.

[31]  Oriol Vinyals,et al.  Neural Discrete Representation Learning , 2017, NIPS.

[32]  Sepp Hochreiter,et al.  GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium , 2017, NIPS.

[33]  Wangmeng Zuo,et al.  Learning Deep CNN Denoiser Prior for Image Restoration , 2017, 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[34]  Alexandros G. Dimakis,et al.  Compressed Sensing using Generative Models , 2017, ICML.

[35]  Alexei A. Efros,et al.  Image-to-Image Translation with Conditional Adversarial Networks , 2016, 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[36]  Raymond Y. K. Lau,et al.  Least Squares Generative Adversarial Networks , 2016, 2017 IEEE International Conference on Computer Vision (ICCV).

[37]  Jong Chul Ye,et al.  A deep convolutional neural network using directional wavelets for low‐dose X‐ray CT reconstruction , 2016, Medical physics.

[38]  Yinda Zhang,et al.  LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop , 2015, ArXiv.

[39]  Andrew Zisserman,et al.  Very Deep Convolutional Networks for Large-Scale Image Recognition , 2014, ICLR.

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

[41]  Stephen P. Boyd,et al.  Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers , 2011, Found. Trends Mach. Learn..

[42]  Fei-Fei Li,et al.  ImageNet: A large-scale hierarchical image database , 2009, 2009 IEEE Conference on Computer Vision and Pattern Recognition.

[43]  P. Massart,et al.  Adaptive estimation of a quadratic functional by model selection , 2000 .

[44]  B. Anderson Reverse-time diffusion equation models , 1982 .

[45]  C. Stein Estimation of the Mean of a Multivariate Normal Distribution , 1981 .

[46]  G. Herman,et al.  Algebraic reconstruction techniques (ART) for three-dimensional electron microscopy and x-ray photography. , 1970, Journal of theoretical biology.

[47]  Eero P. Simoncelli,et al.  Stochastic Solutions for Linear Inverse Problems using the Prior Implicit in a Denoiser , 2021, NeurIPS.

[48]  Michael I. Jordan,et al.  AUTO-ENCODING VARIATIONAL BAYES , 2020 .

[49]  Martin Turner,et al.  CCPi-Regularisation toolkit for computed tomographic image reconstruction with proximal splitting algorithms , 2019, SoftwareX.

[50]  Charles A. Bouman,et al.  5-29-2013 Plug-and-Play Priors for Model Based Reconstruction , 2013 .

[51]  Marc Teboulle,et al.  A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems , 2009, SIAM J. Imaging Sci..

[52]  H. Robbins An Empirical Bayes Approach to Statistics , 1956 .