Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing and Accelerating Diffusion-Based Inverse Problem Solving through Measurements Optimization

Tianyu Chen, Zhendong Wang, Mingyuan Zhou

TL;DR

This work tackles the challenge of solving inverse problems with diffusion-based priors under realistic measurement constraints. It introduces Measurements Optimization (MO), a plug-and-play module that alternates stochastic gradient Langevin dynamics with querying a pretrained diffusion prior, enabling multiple informative gradient steps per diffusion NFE and projection back onto the data manifold. By integrating MO into DPS and Red-diff, the authors achieve state-of-the-art or near-state-of-the-art results with as few as 50–100 NFEs across a suite of linear and nonlinear tasks, including phase retrieval, on FFHQ and ImageNet. The approach significantly reduces memory and wall-time while maintaining high reconstruction quality, broadening the practicality of diffusion-based inverse problem solving.

Abstract

Diffusion models have recently demonstrated notable success in solving inverse problems. However, current diffusion model-based solutions typically require a large number of function evaluations (NFEs) to generate high-quality images conditioned on measurements, as they incorporate only limited information at each step. To accelerate the diffusion-based inverse problem-solving process, we introduce \textbf{M}easurements \textbf{O}ptimization (MO), a more efficient plug-and-play module for integrating measurement information at each step of the inverse problem-solving process. This method is comprehensively evaluated across eight diverse linear and nonlinear tasks on the FFHQ and ImageNet datasets. By using MO, we establish state-of-the-art (SOTA) performance across multiple tasks, with key advantages: (1) it operates with no more than 100 NFEs, with phase retrieval on ImageNet being the sole exception; (2) it achieves SOTA or near-SOTA results even at low NFE counts; and (3) it can be seamlessly integrated into existing diffusion model-based solutions for inverse problems, such as DPS \cite{chung2022diffusion} and Red-diff \cite{mardani2023variational}. For example, DPS-MO attains a peak signal-to-noise ratio (PSNR) of 28.71 dB on the FFHQ 256 dataset for high dynamic range imaging, setting a new SOTA benchmark with only 100 NFEs, whereas current methods require between 1000 and 4000 NFEs for comparable performance.

Enhancing and Accelerating Diffusion-Based Inverse Problem Solving through Measurements Optimization

TL;DR

This work tackles the challenge of solving inverse problems with diffusion-based priors under realistic measurement constraints. It introduces Measurements Optimization (MO), a plug-and-play module that alternates stochastic gradient Langevin dynamics with querying a pretrained diffusion prior, enabling multiple informative gradient steps per diffusion NFE and projection back onto the data manifold. By integrating MO into DPS and Red-diff, the authors achieve state-of-the-art or near-state-of-the-art results with as few as 50–100 NFEs across a suite of linear and nonlinear tasks, including phase retrieval, on FFHQ and ImageNet. The approach significantly reduces memory and wall-time while maintaining high reconstruction quality, broadening the practicality of diffusion-based inverse problem solving.

Abstract

Diffusion models have recently demonstrated notable success in solving inverse problems. However, current diffusion model-based solutions typically require a large number of function evaluations (NFEs) to generate high-quality images conditioned on measurements, as they incorporate only limited information at each step. To accelerate the diffusion-based inverse problem-solving process, we introduce \textbf{M}easurements \textbf{O}ptimization (MO), a more efficient plug-and-play module for integrating measurement information at each step of the inverse problem-solving process. This method is comprehensively evaluated across eight diverse linear and nonlinear tasks on the FFHQ and ImageNet datasets. By using MO, we establish state-of-the-art (SOTA) performance across multiple tasks, with key advantages: (1) it operates with no more than 100 NFEs, with phase retrieval on ImageNet being the sole exception; (2) it achieves SOTA or near-SOTA results even at low NFE counts; and (3) it can be seamlessly integrated into existing diffusion model-based solutions for inverse problems, such as DPS \cite{chung2022diffusion} and Red-diff \cite{mardani2023variational}. For example, DPS-MO attains a peak signal-to-noise ratio (PSNR) of 28.71 dB on the FFHQ 256 dataset for high dynamic range imaging, setting a new SOTA benchmark with only 100 NFEs, whereas current methods require between 1000 and 4000 NFEs for comparable performance.

Paper Structure

This paper contains 33 sections, 1 theorem, 14 equations, 21 figures, 13 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Diffusion Models
  4. Inverse Problem under the Bayesian Framework
  5. Measurements Optimization
  6. Stochastic Gradient Langevin Dynamics (SGLD) as an Optimization Method.
  7. Querying the Diffusion Prior.
  8. Measurement Optimization: Integrating the Two Components.
  9. MO as a Plug-and-Play Module.
  10. Experiments
  11. Experiment Setup
  12. Datasets, Checkpoints, and Metrics.
  13. Inverse Problems.
  14. Baselines.
  15. Experiment results
  16. ...and 18 more sections

Key Result

Theorem 1

Assume that the training set of the pretrained mean-predicted diffusion model $D_\theta$ consists of a finite number of samples $\{\bm{z}_1, \bm{z}_2, \dots, \bm{z}_n\}$. Given sufficient data and model capacity, and $s(t)=1,\forall t$, then for any query $\bm{x}$ with noise level $\sigma$, the pret

Figures (21)

  • Figure 1: DPS-MO examples of restoration. We integrate the MO module into DPS framework and present the recovered images with measurements and ground truth to demonstrate the performance of our method. All results are achieved with 100 NFEs.
  • Figure 2: Workflow overview illustrating how our MO module can be integrated into sampling-based methods (Full algorithm is in Algorithm \ref{['alg:dps_mo']}). Solving the MO module with measurement $\bm{y}$ provides substantial information, reducing the NFE requirements for inference.
  • Figure 3: Comparison of different sampling schedules on the FFHQ phase retrieval task, with all other hyperparameters fixed.
  • Figure 4: Inpainting task with 170 $\times$ 170 box. Four independent runs are able to genrate different faces and provide diversity.
  • Figure 5: Performance comparison with respect to NFEs, showing that DPS-MO achieves high performance at an early stage, requiring fewer NFEs than DPS and Red-diff.
  • ...and 16 more figures

Theorems & Definitions (2)

  • Theorem 1: Adapted from Section B3 in EDM karras2022elucidating
  • Proof 1: Proof of Theorem \ref{['thm:edm']}