Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Score-Based Variational Inference for Inverse Problems

Zhipeng Xue, Penghao Cai, Xiaojun Yuan, Xiqi Gao

TL;DR

By analyzing the probability density evolution of the conditional reverse diffusion process, this work proves that the posterior mean can be achieved by tracking the mean of each reverse diffusion step and establishes a framework termed reverse mean propagation (RMP) that targets the posterior mean directly.

Abstract

Existing diffusion-based methods for inverse problems sample from the posterior using score functions and accept the generated random samples as solutions. In applications that posterior mean is preferred, we have to generate multiple samples from the posterior which is time-consuming. In this work, by analyzing the probability density evolution of the conditional reverse diffusion process, we prove that the posterior mean can be achieved by tracking the mean of each reverse diffusion step. Based on that, we establish a framework termed reverse mean propagation (RMP) that targets the posterior mean directly. We show that RMP can be implemented by solving a variational inference problem, which can be further decomposed as minimizing a reverse KL divergence at each reverse step. We further develop an algorithm that optimizes the reverse KL divergence with natural gradient descent using score functions and propagates the mean at each reverse step. Experiments demonstrate the validity of the theory of our framework and show that our algorithm outperforms state-of-the-art algorithms on reconstruction performance with lower computational complexity in various inverse problems.

Score-Based Variational Inference for Inverse Problems

TL;DR

By analyzing the probability density evolution of the conditional reverse diffusion process, this work proves that the posterior mean can be achieved by tracking the mean of each reverse diffusion step and establishes a framework termed reverse mean propagation (RMP) that targets the posterior mean directly.

Abstract

Existing diffusion-based methods for inverse problems sample from the posterior using score functions and accept the generated random samples as solutions. In applications that posterior mean is preferred, we have to generate multiple samples from the posterior which is time-consuming. In this work, by analyzing the probability density evolution of the conditional reverse diffusion process, we prove that the posterior mean can be achieved by tracking the mean of each reverse diffusion step. Based on that, we establish a framework termed reverse mean propagation (RMP) that targets the posterior mean directly. We show that RMP can be implemented by solving a variational inference problem, which can be further decomposed as minimizing a reverse KL divergence at each reverse step. We further develop an algorithm that optimizes the reverse KL divergence with natural gradient descent using score functions and propagates the mean at each reverse step. Experiments demonstrate the validity of the theory of our framework and show that our algorithm outperforms state-of-the-art algorithms on reconstruction performance with lower computational complexity in various inverse problems.
Paper Structure (30 sections, 3 theorems, 57 equations, 7 figures, 6 tables, 2 algorithms)

This paper contains 30 sections, 3 theorems, 57 equations, 7 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Inverse Problems
  4. Variational Inference
  5. Score-Based Generative Models
  6. Diffusion Process and Posterior Estimation
  7. Score-Based Variational Inference
  8. RMP as Variational Inference
  9. Variational Inference by Natural Gradient Descent
  10. Score-Based Gradient Calculation
  11. Fixed Precision Update
  12. Experiments
  13. Gaussian Mixture Model
  14. Image Reconstruction
  15. Complexity vs Performance
  16. ...and 15 more sections

Key Result

Proposition 1

For diffusion models with forward process (diffusion_process), the reverse conditional $p_{k}(\bm{x}_k|\bm{x}_{k+1},\bm{y})$, $\forall k = 0\cdots T-1$, is Gaussian when $\Delta t \rightarrow 0$. For VE and VP diffusion, the mean and covariance of $p_{k}(\bm{x}_k|\bm{x}_{k+1},\bm{y})$ are tractable where $\bm{V}_{k,1} =(\sigma_{k}^2\bm{I} + \bm{C}_{\bm{x}_{0}}) (\sigma_{k+1}^2\bm{I}+\bm{C}_{\bm{x

Figures (7)

  • Figure 1: An illustration of Reverse Mean Propagation (RMP) for Gaussian mixture model. In the experiment, the data prior is $p(x_0) = \frac{1}{2}(\mathcal{N}(x_0;\mu_1,v_1^2)+ \mathcal{N}(x_0;\mu_2,v_2^2))$ and measurement $y = a x + v_0 \varepsilon$ where $\mu_1=-1$, $\mu_2=1$, $v_1=v_2=0.2$, $v_0=0.5$, $a=1$, $\varepsilon\sim \mathcal{N}(0,1)$ and $T=1000$. RMP is deterministic when $y$ and $\bm{x}_T$ are given and the final output converges to $\mathrm{E}_{p(x_0|y)}[x_0]$.
  • Figure 2: Illustration of the process of VP-based RMP for Gaussian mixture model. Top-left: Various measurement $y$ with fixed $x_T = 0$. Top-right: Random $x_T$ with $y=0.2$. Bottom-left: $(x_T, y)= (-1,-1.5)$, $(x_T, y)= (0,0.2)$ and $(x_T, y)= (1,1.5)$. Bottom-right: MMSE estimation $\mathrm{E}_{p(x_0|y = \bar{y})}[x_0]$ and VP-RMP outputs for different measurement $y$.
  • Figure 3: Part of the results on solving inverse problems with Gaussian noise ($\epsilon=0.05$).
  • Figure 4: Performance comparisons of algorithms with different NFE on SR task.
  • Figure 5: Running time of RMP and DPS for different image reconstruction tasks.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Proposition 1
  • Definition 1
  • Theorem 1
  • Proposition 2