Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Solving General Noisy Inverse Problem via Posterior Sampling: A Policy Gradient Viewpoint

Haoyue Tang, Tian Xie, Aosong Feng, Hanyu Wang, Chenyang Zhang, Yang Bai

TL;DR

This work proposes Diffusion Policy Gradient (DPG), a tractable computation method by viewing the intermediate noisy images as policies and the target image as the states selected by the policy.

Abstract

Solving image inverse problems (e.g., super-resolution and inpainting) requires generating a high fidelity image that matches the given input (the low-resolution image or the masked image). By using the input image as guidance, we can leverage a pretrained diffusion generative model to solve a wide range of image inverse tasks without task specific model fine-tuning. To precisely estimate the guidance score function of the input image, we propose Diffusion Policy Gradient (DPG), a tractable computation method by viewing the intermediate noisy images as policies and the target image as the states selected by the policy. Experiments show that our method is robust to both Gaussian and Poisson noise degradation on multiple linear and non-linear inverse tasks, resulting into a higher image restoration quality on FFHQ, ImageNet and LSUN datasets.

Solving General Noisy Inverse Problem via Posterior Sampling: A Policy Gradient Viewpoint

TL;DR

This work proposes Diffusion Policy Gradient (DPG), a tractable computation method by viewing the intermediate noisy images as policies and the target image as the states selected by the policy.

Abstract

Solving image inverse problems (e.g., super-resolution and inpainting) requires generating a high fidelity image that matches the given input (the low-resolution image or the masked image). By using the input image as guidance, we can leverage a pretrained diffusion generative model to solve a wide range of image inverse tasks without task specific model fine-tuning. To precisely estimate the guidance score function of the input image, we propose Diffusion Policy Gradient (DPG), a tractable computation method by viewing the intermediate noisy images as policies and the target image as the states selected by the policy. Experiments show that our method is robust to both Gaussian and Poisson noise degradation on multiple linear and non-linear inverse tasks, resulting into a higher image restoration quality on FFHQ, ImageNet and LSUN datasets.
Paper Structure (26 sections, 3 theorems, 28 equations, 17 figures, 3 tables, 1 algorithm)

This paper contains 26 sections, 3 theorems, 28 equations, 17 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction and Problem Formulation
  2. Methodology
  3. Preliminaries: Diffusion Models
  4. Solving Inverse Problems as Posterior Sampling
  5. Computing $\nabla_{\bf x}\log p_i({\bf x}_i|{\bf y})$ as Policy Gradient
  6. Implementation Details
  7. Tractable Monte Carlo sampling
  8. Selection of $p_0({\bf y}|{\bf x}_0)$
  9. Reward Shaping
  10. Score Function Re-scaling
  11. Algorithm Description
  12. Connection with Diffusion Posterior Sampling (DPS)
  13. Experiments
  14. Quantitative Results on Noisy Linear Inverse Problems
  15. Quantitative Experiments on Non-Linear Image Inverse Problems
  16. ...and 11 more sections

Key Result

Theorem 1

Suppose $p_0({\bf x}_0)$ is the probability measure of $N_{\text{train}}, N_{\text{train}}<\infty$ high quality training images. Then for all $i\in[N]$, we can compute the score function ${\tilde{\bf s}}_i({\bf x}_i, {\bf y})$ from equation eq:pgdef as follows:

Figures (17)

  • Figure 1: Examples on solving noisy image inverse problems on ImageNet validation set using our proposed method without task specific model finetuning or training.
  • Figure 2: Evolution of the reconstruction loss $\Vert{\bf y}-\mathcal{A}({\bm \mu}_i({\bf x}_i))\Vert_2$ of the DPS and DPG method.
  • Figure 3: Image generation procedure and the reconstruction loss $\ell:=\Vert {\bf y}-\mathcal{A}({\bm \mu}_i({\bf x}_i))\Vert_2$ by using DPG and DPS methods in super-resolution.
  • Figure 4: Results on solving linear noisy inverse problems (inpainting, super-resolution and Gaussian deblurring) on ImageNet Dataset. The input image is distorted by random Gaussian noise $\sigma_{\bf y}=0.05$.
  • Figure 5: Image Restoration Results for Motion Deblurring on ImageNet256$\times$256.
  • ...and 12 more figures

Theorems & Definitions (3)

  • Theorem 1: Leibniz Rule
  • Lemma 1
  • Corollary 1