Source-linked AI summary

Denoising Diffusion Models for Plug-and-Play Image Restoration

Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, Luc Van Gool

arXiv:2305.08995v1cs.CVeess.IV

TL;DR

Plug-and-play image restoration has mainly used discriminative Gaussian denoisers, leaving diffusion models as generative priors insufficiently explored and prior diffusion-restoration methods costly or ineffective. DiffPIR integrates plug-and-play data-prior splitting into diffusion sampling and evaluates it across three restoration tasks. It reports state-of-the-art faithfulness and perceptual quality on FFHQ and ImageNet with no more than 100 NFEs.

  • Problem

    Plug-and-play restoration largely relies on discriminative Gaussian denoisers, while diffusion models' potential as generative denoising priors remains insufficiently explored.

  • Method

    DiffPIR uses HQS-based diffusion sampling with an off-the-shelf diffusion model as the plug-and-play denoising prior and solves the data subproblem in the clean image manifold.

  • Results

    DiffPIR achieves superior restoration quality across super-resolution, deblurring, and inpainting on FFHQ and ImageNet with no more than 100 NFEs.

  • Takeaways & Limitations

    DiffPIR combines diffusion models' generative denoising ability with plug-and-play restoration's flexibility across degradation models and inverse-problem tasks.

  • Takeaways & Limitations

    Prior diffusion-restoration methods can suffer from low sampling speed, and the discussion restricts the derivation to DDPM without loss of generality.

Abstract

from arXiv · show

Plug-and-play Image Restoration (IR) has been widely recognized as a flexible and interpretable method for solving various inverse problems by utilizing any off-the-shelf denoiser as the implicit image prior. However, most existing methods focus on discriminative Gaussian denoisers. Although diffusion models have shown impressive performance for high-quality image synthesis, their potential to serve as a generative denoiser prior to the plug-and-play IR methods remains to be further explored. While several other attempts have been made to adopt diffusion models for image restoration, they either fail to achieve satisfactory results or typically require an unacceptable number of Neural Function Evaluations (NFEs) during inference. This paper proposes DiffPIR, which integrates the traditional plug-and-play method into the diffusion sampling framework. Compared to plug-and-play IR methods that rely on discriminative Gaussian denoisers, DiffPIR is expected to inherit the generative ability of diffusion models. Experimental results on three representative IR tasks, including super-resolution, image deblurring, and inpainting, demonstrate that DiffPIR achieves state-of-the-art performance on both the FFHQ and ImageNet datasets in terms of reconstruction faithfulness and perceptual quality with no more than 100 NFEs. The source code is available at {\url{https://github.com/yuanzhi-zhu/DiffPIR}}

1. Introduction

Plug-and-play image restoration separates data fidelity from an implicit denoising prior, but existing approaches largely rely on discriminative Gaussian denoisers. DiffPIR incorporates diffusion models as generative denoiser priors within this framework and evaluates them across several restoration tasks.

  • Existing plug-and-play restoration: Existing plug-and-play methods commonly integrate Gaussian denoisers through iterative variable-splitting procedures such as ADMM and HQS.These methods use denoisers to address the prior subproblem.
  • Existing plug-and-play restoration: Plug-and-play IR separates the data term from the prior term in an optimization problem constrained by the degradation model.The data term enforces consistency with measurements, while the prior term enforces the desired data distribution.
  • Motivation: Generative-model denoisers are expected to better handle ill-posed inverse problems, but earlier GAN-, flow-, and VAE-based priors were not designed for denoising.Their generative capabilities can be hindered when used as plug-and-play priors.
  • DiffPIR: DiffPIR decouples data and prior terms within diffusion sampling, solving the data term independently and using an off-the-shelf diffusion model as the denoising prior.This design allows different degradation operators H and a wide range of degradation models.
  • Evaluation: DiffPIR is evaluated on super-resolution, deblurring, and inpainting using FFHQ and ImageNet, with reported efficient restoration and superior visual quality.The paper highlights restoration comparisons and visual examples in Figure 1.

2. Background

Diffusion models progressively add Gaussian noise in a forward process and reverse it using a learned score or noise predictor. DDPM and DDIM provide discrete reverse-sampling formulations, while conditional generation incorporates measurement or condition information through posterior terms.

  • Score-based diffusion models: The forward diffusion process gradually perturbs images with Gaussian noise, while the reverse process removes noise to generate samples.The forward process can be represented by an Itô SDE with drift and diffusion coefficients.
  • Score-based diffusion models: The reverse SDE depends on the score ∇x log p_t(x), which is modeled by a learned time-dependent score function sθ(x_t,t).Evaluating this function at intermediate timesteps enables generation from an initially Gaussian-noise state.
  • Denoising diffusion probabilistic models: DDPM uses a noise schedule β_t to form noisy states, with x_t = √(ᾱ_t)x_0 + √(1−ᾱ_t)ε at arbitrary timestep t.Here α_t = 1−β_t and ᾱ_t is the product of the preceding α_s values.
  • Denoising diffusion probabilistic models: A DDPM reverse step uses a learned noise predictor εθ(x_t,t), which estimates the total noise between x_t and x_0.Score prediction and noise prediction are approximately related and can represent diffusion models interchangeably in subsequent discussion.
  • Efficient diffusion sampling: DDIM generalizes reverse diffusion from Markovian to non-Markovian sampling, with ση_t controlling the stochasticity of each step.The first reverse-update term uses the predicted clean image at timestep t.
  • Conditional generation: For conditional generation, the reverse process combines the image score with the conditional likelihood term log p_t(y|x), enabling posterior sampling from an unconditional diffusion model.Conditional information can be incorporated after an unconditional reverse step.

3. Proposed Method

DiffPIR decouples data fidelity and image priors within diffusion sampling, combining diffusion denoising with iterative data-consistent correction for inverse restoration.

  • DiffPIR framework: DiffPIR uses HQS to split the data and prior terms into iterative subproblems within the diffusion sampling framework.The auxiliary variable z enables separate prior and data updates.
  • DiffPIR sampling: Each reverse-diffusion step predicts a clean image, applies the data proximal correction, and then derives x_t−1 by adding controlled noise.The injected-noise variance is controlled by ζ, with deterministic sampling when ζ = 0.
  • Diffusion prior: Diffusion models serve as generative denoiser priors, combining generation in early steps with denoising in later steps.This generative ability is intended to address challenging ill-posed problems such as large-mask inpainting.
  • Data consistency: The data subproblem enforces measurement consistency through a proximal update involving the degradation operator H.For deblurring, inpainting, and super-resolution, a fast solution is available; otherwise, a first-order proximal approximation can be used.
  • Evaluation: The method is evaluated against diffusion-based and plug-and-play baselines on noisy Gaussian-deblurring, motion-deblurring, and 4× super-resolution settings.The referenced comparison includes DiffPIR, DPS, DDRM, and DPIR qualitative results.
  • Related methods: DiffPIR can handle arbitrary degradation operators H, unlike DDRM, which supports only linear operators and may depend on feasible fast SVD.DPS is described as suffering from slow sampling and reduced reconstruction faithfulness with few sampling steps.

4. Experiments

Experiments evaluate DiffPIR across super-resolution, deblurring, and inpainting on FFHQ and ImageNet, measuring fidelity, perceptual similarity, and visual quality. Results show strong perceptual performance with limited NFEs, diverse reconstructions, and sensitivity to sampling and guidance hyperparameters.

  • Experimental Setup: Experiments use FFHQ 256×256 and ImageNet 256×256, evaluating 100 hold-out validation images per dataset with pre-trained diffusion models.The same linear noise schedule and unchanged diffusion-model settings are used across methods.
  • Experimental Setup: Evaluations cover 4× super-resolution, Gaussian and motion deblurring, and inpainting under specified noisy and noiseless degradation settings.Inpainting is evaluated only without noise, while deblurring and super-resolution include both noisy and noiseless settings.
  • Quantitative Results: DiffPIR achieves superior FID and LPIPS on both datasets for noisy restoration, while remaining competitive in PSNR except for SR LPIPS.The SR LPIPS exception is attributed to accumulated sampling errors from inaccurate approximated bicubic kernels.
  • Quantitative Results: For noiseless FFHQ restoration, DiffPIR with 100 NFEs significantly outperforms comparison methods in FID and LPIPS, while 20 NFEs remain competitive.DPIR has higher PSNR on noiseless tasks, but its generated images often have lower perceptual quality.
  • Qualitative Results: DiffPIR reconstructs intricate motion-deblurring details and requires far fewer NFEs than DPS for faithful reconstruction.Unlike DDRM and DPIR, which tend to generate blurry images, DiffPIR preserves detailed content in the reported comparison.
  • Qualitative Results: DiffPIR generates diverse, semantically aligned inpainting reconstructions under strong degradation, including masks covering 75% of the image.The sampling method is described as producing diverse reconstructions similarly to DDPM.
  • Ablation Study: In noisy 4× SR, LPIPS is lowest for 100–500 sampling steps, and DiffPIR produces detailed images with fewer than 100 NFEs.The default number of NFEs is set to 100.
  • Ablation Study: Starting reverse diffusion at tstart = 400 reduces NFEs without loss of quality in noisy Gaussian deblurring.The study also compares tstart values of 200 and 400.

5. Conclusions

DiffPIR is introduced as a diffusion-based plug-and-play image-restoration sampler using off-the-shelf diffusion models as denoising priors and solving the data subproblem on the clean-image manifold.

  • DiffPIR combines plug-and-play image restoration with diffusion-model sampling.It uses an HQS-based diffusion sampling approach and off-the-shelf diffusion models as plug-and-play denoising priors.
  • The method solves the data subproblem in the clean image manifold.
  • Experiments emphasize DiffPIR’s flexibility, efficiency, and generalizability against competitive methods.

A. Other HQS-based Sampling Methods

This section formulates conditional reverse diffusion as an HQS-based optimization process, connecting its prior subproblem to diffusion-model score updates.

  • Each conditional reverse diffusion step is formulated as a MAP estimation problem at noise level βt.
  • HQS separates the data and prior terms, allowing the prior-related update to be approximated through a first-order proximal step.
  • The score function sθ is linked to the prior gradient through ∇xP(x) = −∇x log p(x) = −sθ(x).
  • The resulting approximation recovers the reverse-process form used by DDPM.

A.2. DPS as a Special Case

The section derives DPS as a special case of the broader HQS-based sampling framework by inserting measurement guidance into the reverse-diffusion update.

  • The data-consistency update uses a gradient of the measurement residual as an approximation to the data proximal subproblem.
  • DPS guidance approximates the noisy measurement likelihood gradient using the denoised estimate x̂0.
  • The derivation reproduces the DPS sampling strategy and distinguishes the original DPSy0 algorithm from the variant DPSyt.
  • A noise-level measurement yt is defined by adding diffusion noise to y and treating it as the measurement corresponding to xt.
  • DPSyt alternates reverse diffusion updates with data-proximal-gradient steps before returning x0.

B.1. Hyperparameters Values

The section specifies task-dependent degradation models and data-subproblem solvers for super-resolution, inpainting, and deblurring, including iterative and closed-form options.

  • Hyperparameters: The experiments list hyperparameter settings across FFHQ and ImageNet tasks, including box and random inpainting.
  • Image Inpainting: Inpainting assumes noiseless masked observations, y = M ⊙ x, and recovers missing pixels from known pixels.M is a user-defined boolean mask, and the closed-form update is x0 = (M ⊙ y + ρt z0)/(M + ρt).
  • Image Deblurring: Deblurring uses the linear degradation model y = x ⊗ k + n with Gaussian noise.When the convolution is circular, the data subproblem admits an analytical Fourier-domain solution.
  • Single Image Super-Resolution: Bicubic single-image super-resolution uses y = x↓^bicubic_sf + n, where sf is the downscaling factor.The data subproblem can be solved by iterative back-projection or an approximated-kernel closed-form method.
  • Single Image Super-Resolution: The closed-form super-resolution solution should outperform iterative solutions in quantitative metrics because it has fewer hyperparameters.

C. Additional Ablation Study

The reverse diffusion process is examined through intermediate results, revealing an initial phase where the analytical solution offers no help. This motivates skipping that phase using an experimentally determined end timestep, tstart.

  • Intermediate results illustrate the reverse diffusion process and reveal an initial phase where the analytical solution offers no help.
  • The method skips this initial phase, with tstart selected as its end timestep based on experiments.This choice is described in connection with Section 4.4.

D. Additional Visual Results

Additional visual examples show DiffPIR and related DPS variants across super-resolution, deblurring, and inpainting settings. The examples include diverse reconstructions, multiple blur kernels, and different masks.

  • Additional examples on FFHQ and ImageNet demonstrate DPSyt and DPSy0 on deblurring and super-resolution tasks.
  • DiffPIR produces diverse 8× and 16× super-resolution reconstructions with 100 NFEs under multiple noise levels.
  • DiffPIR is shown handling both motion and Gaussian blur kernels with 100 NFEs.
  • DiffPIR is shown handling inpainting with different masks using 100 NFEs.
Loading 2305.08995v1…