Score-Based Variational Inference for Inverse Problems

## Abstract

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.

## 1 Introduction

Diffusion models [, , , , ] have shown impressive performance for image generation. For diffusion models such as diffusion denoising probability model (DDPM) [] and denoising score matching with Langevin dynamics (SMLD) [], the essential part is the learning of score functions of data distributions with large datasets. By approximating score functions with neural networks such as U-Net [, ], the prior of complex data distributions can be learned implicitly which encourages many applications. Inverse problems aim to recover an unknown state $\bm{x}_{0}$ from observation $\bm{y}$, which is fundamental to various research areas such as wireless communication, image processing and natural language processing. Recent works [, , , , , , , , , ] have shown that diffusion models can be used for solving inverse problems since the prior of data distribution is learned implicitly with score functions.

Based on Bayes’ rule, diffusion models are used for the generation of data from the posterior distribution with score functions of data and likelihood, and thus can be applied in solving reverse problems. The main difficulty of applying diffusion models to solving inverse problems is the calculation of likelihood score. SNIPS [] and DDRM [] are proposed to solve noisy linear inverse problems with diffusion process in the spectral domain. In these methods, the measurement and data to be estimated are transformed into the spectral domain via singular value decomposition (SVD), and the conditional score can be calculated with SVD explicitly. In [], the authors propose to approximate the likelihood score by a noise-perturbed pseudo-likelihood score which has a closed form under certain assumptions and can be efficiently calculated using SVD for noisy linear inverse problems. MCG [] circumvents the calculation of likelihood score by projections onto the measurement constrained manifold. In DPS [], the Laplace method is used for the approximation of the likelihood score for general inverse problems. Instead of directly approximating the likelihood score, a variational inference based method termed RED-Diff [] optimizes the reconstruction loss with score matching regularization.

The works mentioned above mainly focus on generating random samples from the posterior distribution rather than finding a statistical estimate such as the posterior mean, i.e. the MMSE estimation in Bayesian inference, which is preferred in many applications []. To get the posterior mean, we have to sample from the posterior multiple times and average out which is time-consuming. In this work, we show that the evolution of the conditional probability of the reverse diffusion process can be characterized by the reverse mean and covariance, and the posterior mean can be achieved by tracking the mean of each conditional reverse diffusion step. Base on this, we propose a variational inference based framework that minimizes the reverse KL divergence and propagate the reverse mean at each reverse step, referred to as Reverse Mean Propagation (RMP). By implementing RMP with stochastic natural gradient descent and approximating the gradient with score functions, we devise a practical algorithm for inverse problems. Instead of generating samples by running algorithms multiple times, RMP outputs an estimate which is closed to the posterior mean directly and thus can reduce complexity greatly. We demonstrate the validity of the theory of RMP by a toy example as illustrated in Fig 1 and show that by adopting the RMP algorithm, we achieve a large improvement over existing state-of-the-art algorithms on various image reconstruction tasks.

Figure: 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=ax+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}]$.
Refer to caption: x1.png

## 2 Background

### 2.1 Inverse Problems

An inverse problem is defined as the estimation of an unknown state or a latent $\bm{x}_{0}\in\mathbb{R}^{N\times 1}$ from measurement $\bm{y}\in\mathbb{R}^{M\times 1}$. Specifically, the measurement process can be described by a measurement operator $\mathcal{A}:\mathbb{R}^{N\times 1}\rightarrow\mathbb{R}^{M\times 1}$, and the final output is a noisy version of the measurement:

$$ $\displaystyle\bm{y}=\mathcal{A}(\bm{x}_{0})+\bm{w}_{0}$ (1) $$

where $\bm{w}_{0}\in\mathbb{R}^{M\times 1}$ is the measurement noise. Usually the measurement noise is assumed to be zero mean Gaussian with variance $\epsilon^{2}\bm{I}$. Other noise models may also apply. In linear inverse problems, the measurement function $\mathcal{A}$ is linear and can be represented by a linear transform $\bm{A}\in\mathbb{R}^{M\times N}$. We focus on inverse problems with known $\mathcal{A}$ in this paper.

### 2.2 Variational Inference

The inverse problem (1) can be formulated as a Bayesian estimation problem. The posterior distribution of $\bm{x}_{0}$ is given by $p(\bm{x}_{0}|\bm{y})=\frac{p(\bm{y}|\bm{x}_{0})p(\bm{x}_{0})}{p(\bm{y})}$,
where $p(\bm{x}_{0})$ is the prior distribution of $\bm{x}_{0}$ and $p(\bm{y}|\bm{x}_{0})$ is the conditional distribution of $\bm{y}$ given $\bm{x}_{0}$. The posterior mean, i.e., the MMSE estimator can be employed for the estimation of $\bm{x}_{0}$. However, the posterior $p(\bm{x}_{0}|\bm{y})$ is intractable in general since the prior $p(\bm{x}_{0})$ and the likelihood $p(\bm{y}|\bm{x}_{0})$ may be very complicated in real applications. An alternative way is to find an approximation of the posterior as in variational inference (VI). VI introduces a distribution $q_{\phi}(\bm{x}_{0}|\bm{y})$ and maximizes a lower bound of the log probability of marginal $p(\bm{y})$:

$$ $\displaystyle\begin{split}\log p(\bm{y})&=\log\int\frac{q_{\phi}(\bm{x}_{0}| \bm{y})}{q_{\phi}(\bm{x}_{0}|\bm{y})}p(\bm{y}|\bm{x}_{0})p(\bm{x}_{0})d\bm{x}_ {0}\\ &\geq\int q_{\phi}(\bm{x}_{0}|\bm{y})\log\frac{p(\bm{y},\bm{x}_{0})}{q_{\phi}( \bm{x}_{0}|\bm{y})}d\bm{x}_{0}=-\mathcal{F}_{\phi}(\bm{y})\\ &=-\text{KL}(q_{\phi}(\bm{x}_{0}|\bm{y})||p(\bm{x}_{0}|\bm{y}))+\log p(\bm{y}) \end{split}$ (2) $$

where the inequality is obtained by using the Jensen’s inequality. The lower bound is referred as the evidence lower bound (ELBO) or the negative of free energy $\mathcal{F}(\bm{y})$. It is worth noting that maximizing ELBO is equivalent to minimizing the KL divergence between $q_{\phi}(\bm{x}_{0}|\bm{y})$ and $p(\bm{x}_{0}|\bm{y})$ as shown in the last line of (2). Many methods have been proposed for the optimization of (2) such as mean field VI [], black box VI [], stochastic VI [] and normalizing flow VI []. However, these methods are difficult to be applied to real applications since prior $p(\bm{x}_{0})$ is complicated and is usually learned by neural networks. It is shown that learning the distribution of high dimensional data through score matching [] directly is inaccurate since the existance of low density data regions []. Also, perturbing data with Gaussian noise makes the data distribution more amenable to learn [] which is the core of score-based generative model.

### 2.3 Score-Based Generative Models

Score-based generative models or diffusion models generate samples of a data distribution from the reverse process of a diffusion process. The diffusion process is also called the forward process where Gaussian noise is added gradually to the original data distribution until the noisy data are approximately Gaussian-distributed. More specifically, the diffusion process is a Markov process with joint probability of its latent states $\{\bm{x}_{k}\}_{k=0}^{T}$ given by

$$ $\displaystyle p(\bm{x}_{T:0})=p(\bm{x}_{0})\prod_{k=0}^{T-1}p(\bm{x}_{k+1}|\bm {x}_{k}).$ (3) $$

Two classes of widely studied diffusion models, i.e., the variance preserving (VP) diffusion model [] and the variance exploding (VE) diffusion model [] are distinguished by Markov diffusion kernel $p(\bm{x}_{k+1}|\bm{x}_{k})$. For VE diffusion, $p(\bm{x}_{k+1}|\bm{x}_{k})=\mathcal{N}(\bm{x}_{k+1};\bm{x}_{k},(\sigma_{k+1}^{
2}-\sigma_{k}^{2})\bm{I})$, where $\sigma_{T}^{2}>\sigma_{T-1}^{2}>\cdots>\sigma_{1}^{2}>\sigma_{0}^{2}=0$, and for VP diffusion, $p(\bm{x}_{k+1}|\bm{x}_{k})=\mathcal{N}(\bm{x}_{k+1};\sqrt{1-\beta_{k+1}}\bm{x}
_{k},\beta_{k+1}\bm{I})$ where $\beta_{T}>\beta_{T-1}>\cdots>\beta_{1}>\beta_{0}=0$. The learning of score function $\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k})$ is essential to the sample generation for both diffusion models, and thus such models are called score-based generative models. In VP diffusion models, a variational reverse process is leaned to minimize the KL divergence between forward and reverse process. The score functions of perturbed data distributions are trained with a variational bound. Samples are generated from the learned reverse process with ancestral sampling method. In VE diffusion models, the score functions of perturbed
data distributions are learned with a score network using denosing score matching []. Samples from the desired distribution are generated with Langevin dynamics method.

## 3 Diffusion Process and Posterior Estimation

The measurement $\bm{y}$ in (1) and diffusion states $\{\bm{x}_{k}\}_{k=0}^{T}$ in (3) form a new Markov chain $\bm{y}\rightarrow\bm{x}_{0}\rightarrow\bm{x}_{1}\cdots\rightarrow\bm{x}_{T}$ and the reverse conditional $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ is given by $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\frac{p(\bm{x}_{k+1}|\bm{x}_{k})p(\bm{x}
_{k}|\bm{y})}{p(\bm{x}_{k+1}|\bm{y})}$, $\forall k=0,\cdots T-1$. In this part, we focus on the property of reverse conditional $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$. We relate the Markov chain $\{\bm{x}_{k}\}_{k=0}^{T}$ to continuous stochastic process $\{\bm{x}_{t}\}_{t=0}^{1}$ by letting $\bm{x}_{k}=\bm{x}_{t=k\Delta t}$, where $\Delta t=\frac{1}{T}$. Then, the discrete diffusion process (3) becomes a continuous process in the limit $\Delta t\rightarrow 0$. We have the following results in the limit of $\Delta t\rightarrow 0$.

###### Proposition 1

For diffusion models with forward process (3), 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 with mean given by

$$ $\displaystyle\begin{split}\bm{\mu}_{k}(\bm{x}_{k+1},\bm{y})&=\bm{V}_{k,1}\bm{x }_{k+1}+\bm{V}_{k,2}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\end{split}$ (4) $$

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}_{0}})^{-1}$ and $\bm{V}_{k,2}=(\sigma_{k+1}^{2}-\sigma_{k}^{2})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{
\bm{x}_{0}})^{-1}$ for VE diffusion, and $\bm{V}_{k,1}=\sqrt{\alpha_{k+1}}((1-\bar{\alpha}_{k})\bm{I}+\bar{\alpha}_{k}
\bm{C}_{\bm{x}_{0}})((1-\bar{\alpha}_{k+1})\bm{I}+\bar{\alpha}_{k+1}\bm{C}_{
\bm{x}_{0}})^{-1}$ and $\bm{V}_{k,2}=\sqrt{\bar{\alpha}_{k}}(1-\alpha_{k+1})((1-\bar{\alpha}_{k+1})\bm
{I}+\bar{\alpha}_{k+1}\bm{C}_{\bm{x}_{0}})^{-1}$ for VP diffusion. $\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$ and $\bm{C}_{\bm{x}_{0}}$ are the mean and covariance of $p(\bm{x}_{0}|\bm{y})$ respectively. $\bar{\alpha}_{k}=\prod_{i=0}^{k}\alpha_{i}$, $\alpha_{i}=1-\beta_{i}$. The covariance of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ for VE and VP diffusion models are given respectively by

$$ $\displaystyle\begin{split}\bm{C}_{k,\text{VE}}&=(\sigma_{k+1}^{2}-\sigma_{k}^{ 2})(\sigma_{k}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{ \bm{x}_{0}})^{-1}\\ \bm{C}_{k,\text{VP}}&=\frac{\beta_{k+1}}{1-\beta_{k+1}}((1-\bar{\alpha}_{k}) \bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0}})\left(\left(\frac{\beta_{k+1}}{1- \beta_{k+1}}+1-\bar{\alpha}_{k}\right)\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0 }}\right)^{-1}.\end{split}$ (5) $$

Proposition 1 generalizes the Gaussian property of $p_{t}(\bm{x}_{t}|\bm{x}_{t+\Delta t})$ to the conditional case $p_{t}(\bm{x}_{t}|\bm{x}_{t+\Delta t},\bm{y})$ when $\Delta t\rightarrow 0$. The essential is that the reverse diffusion process can also be expressed by a reverse SDE []. Based on Proposition 1, we obtain the following main result.

###### Definition 1

The reverse mean propagation chain of a diffusion process is defined as

$$ $\displaystyle\bm{\mu}_{T}\rightarrow\bm{\mu}_{T-1}(\bm{x}_{T}=\bm{\mu}_{T},\bm {y})\rightarrow\cdots\rightarrow\bm{\mu}_{1}(\bm{x}_{2}=\bm{\mu}_{2},\bm{y}) \rightarrow\bm{\mu}_{0}(\bm{x}_{1}=\bm{\mu}_{1},\bm{y})$ (6) $$

where $\bm{\mu}_{k}(\bm{x}_{k+1}=\bm{\mu}_{k+1},\bm{y})$ is the mean of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1}=\bm{\mu}_{k+1},\bm{y})$, $\forall k=0,\cdots,T-1$, and $\bm{\mu}_{T}$ and $\bm{\mu}_{0}$ are the initial point and the end point of the reverse chain respectively.

###### Theorem 1

For VE diffusion, when $\Delta t\rightarrow 0$, the end point of the reverse chain, i.e., $\bm{\mu}_{0}$ is given by

$$ $\displaystyle\bm{\mu}_{0}=(\sigma_{0}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{T }^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\bm{\mu}_{T}+(\sigma_{T}^{2}-\sigma_{0}^{ 2})(\sigma_{T}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\mathrm{E}_{p(\bm{x}_{0}|\bm {y})}[\bm{x}_{0}]$ (7) $$

and $\bm{\mu}_{0}\rightarrow\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$ as $\sigma_{T}\rightarrow\infty$. For VP diffusion, when $\Delta t\rightarrow 0$, $\bm{\mu}_{0}$ is given by

$$ $\displaystyle\begin{split}\bm{\mu}_{0}=\ &\sqrt{\bar{\alpha}_{T}}((1-\bar{ \alpha}_{0})\bm{I}+\bar{\alpha}_{0}\bm{C}_{\bm{x}_{0}})((1-\bar{\alpha}_{T}) \bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{x}_{0}})^{-1}\bm{\mu}_{T}\\ &+(1-\bar{\alpha}_{T})((1-\bar{\alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{ x}_{0}})^{-1}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\end{split}$ (8) $$

and $\bm{\mu}_{0}\rightarrow\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$ as $\bar{\alpha}_{T}\rightarrow 0$.

According to Theorem 1, the posterior mean can be obtained by tracking the mean at each reverse step. By calculating the reverse mean and following the reverse chain in (6), we get a posterior estimation framework termed Reverse Mean Propagation (RMP) for inverse problems as presented in Algorithm 1. We note that when the initial point of the reverse chain $\bm{\mu}_{T}$ and $\bm{y}$ are given, the reverse chain is deterministic and converges to the posterior mean $\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$.

Figure: Algorithm 1 Reverse Mean Propagation (RMP)

###### Proposition 1

###### Definition 1

###### Theorem 1

## 4 Score-Based Variational Inference

In this section, we propose a score-based variational inference method to implement the RMP framework. We show that tracking the mean of the reverse process of each step can be formulated as a sequential of variational inference problems which we prove to be equivalent to a variational inference problem for all latent variables. We solve the variational inference by stochastic natural gradient descent with approximations that simplify the calculation.

### 4.1 RMP as Variational Inference

In Section 3, we present the RMP framework based on the reverse diffusion process. However, in practice, the reverse mean in RMP cannot be calculated using (4) since $\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$ and $\bm{C}_{\bm{x}_{0}}$ are unknown. We now show that the RMP framework can be implemented using variational inference. Instead of applying variational inference on the conditional posterior of $\bm{x}_{0}$ as in (2), we focus on the joint conditional posterior of $\{\bm{x}_{k}\}_{k=0}^{T}$, i.e., $p(\bm{x}_{0:T}|\bm{y})$, which includes all the latent variables in the diffusion process. The variational reverse process with joint conditional is defined by

$$ $\displaystyle q(\bm{x}_{0:T}|\bm{y})=q(\bm{x}_{T}|\bm{y})\prod_{k=0}^{T-1}q_{k }(\bm{x}_{k}|\bm{x}_{k+1},\bm{y}),$ (9) $$

where $q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\mathcal{N}(\bm{x}_{k};\bm{\mu}_{k}(\bm{
x}_{k+1},\bm{y}),\bm{C}_{k}(\bm{x}_{k+1},\bm{y})),\forall k=0:T-1$. We set $q(\bm{x}_{T}|\bm{y})=\mathcal{N}(\bm{x}_{T};0,\bm{I})$. The variational reverse process is chosen as a Markov chain since the reverse of the diffusion process (3) is a Markov process. The KL divergence between variational joint posterior $q(\bm{x}_{0:T}|\bm{y})$ and and joint posterior $p(\bm{x}_{0:T}|\bm{y})$ is given by

$$ $\displaystyle KL(q||p)=\int q(\bm{x}_{0:T}|\bm{y})\log\frac{q(\bm{x}_{0:T}|\bm {y})}{p(\bm{x}_{0:T}|\bm{y})}d\bm{x}_{0:T}$ (10) $$

where the forward joint posterior $p(\bm{x}_{0:T}|\bm{y})=p(\bm{x}_{T}|\bm{y})\prod_{k=0}^{T-1}p(\bm{x}_{k}|\bm{x
}_{k+1:T},\bm{y})$. For VE diffusion $p(\bm{x}_{T}|\bm{y})=\mathcal{N}(\bm{x}_{T};0,\sigma_{T}^{2}\bm{I})$ and for VP diffusion $p(\bm{x}_{T}|\bm{y})=\mathcal{N}(\bm{x}_{T};0,\bm{I})$. The following proposition simplifies the minimization of $KL(q||p)$ with proof given in Appendix C.

###### Proposition 2

For a diffusion process with forward process (3), the KL divergence between variational $q(\bm{x}_{0:T}|\bm{y})$ and joint posterior $p(\bm{x}_{0:T}|\bm{y})$ defined in (10) equals

$$ $\displaystyle KL(q||p)=\sum_{k=T-1}^{0}\int q(\bm{x}_{k+1}|\bm{y})\int q_{k}( \bm{x}_{k}|\bm{x}_{k+1},\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y}) }{p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}d\bm{x}_{k}d\bm{x}_{k+1},$ (11) $$

and the minimization of $KL(q||p)$ is equivalent to the minimization of

$$ $\displaystyle KL(q_{k}||p_{k})=\int q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})\log \frac{q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm {y})}d\bm{x}_{k},\forall k=0,\cdots,T-1.$ (12) $$

According to Proposition 2, we can minimize the KL divergence between $q$ and $p$ by minimizing the KL divergence $KL(q_{k}||p_{k})$, i.e., solve the following VI problem at each reverse step $k$:

$$ $\displaystyle q_{k}^{\star}=\operatorname*{arg\,min}_{q_{k}}KL(q_{k}||p_{k}), \forall k=0,\cdots,T-1.$ (13) $$

By propagating the mean of $q_{k}$ and solving problem (13) at each reverse step $k$, $\forall k=0,\cdots,T-1$, we can approximate the RMP framework in Algorithm 1 based on variational inference, as detailed below.

###### Proposition 2

### 4.2 Variational Inference by Natural Gradient Descent

For $q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\mathcal{N}(\bm{x}_{k};\bm{\mu}_{k},
\Lambda_{k}^{-1}\bm{I})$, the KL divergence between $q_{k}$ and $p_{k}$ is given by

$$ $\displaystyle\begin{split}KL(q_{k}||p_{k})&=\int q_{k}(\bm{x}_{k}|\bm{x}_{k+1} ,\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm{ x}_{k+1},\bm{y})}d\bm{x}_{k}\\ &=-\frac{N}{2}\log(2\pi/\Lambda_{k})-\frac{N}{2}-\mathrm{E}_{q_{k}}[\log p_{k} (\bm{x}_{k}|\bm{x}_{k+1},\bm{y})].\end{split}$ (14) $$

A common practice to optimize $KL(q_{k}||p_{k})$ is to update variational parameters $\phi_{k}=\{\bm{\mu}_{k},\Lambda_{k}\}$ using mini-batch stochastic gradient descent which involves the calculation of $\nabla_{\phi_{k}}KL(q_{k}||p_{k})$. Since $q_{k}$ is Gaussian, the gradient of $\mathrm{E}_{q_{k}}[\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})]$ in (14) with respect to variational parameters $\phi_{k}=\{\bm{\mu}_{k},\Lambda_{k}\}$ have simple forms [] which are given by:

$$ $\displaystyle\begin{split}\nabla_{\bm{\mu}_{k}}\mathrm{E}_{q_{k}}[\log p_{k}( \bm{x}_{k}|\bm{x}_{k+1},\bm{y})]&=\mathrm{E}_{q_{k}}[\nabla_{\bm{x}_{k}}\log p _{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})]\\ \nabla_{\bm{\Lambda}_{k}}\mathrm{E}_{q_{k}}[\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1} ,\bm{y})]&=-\frac{1}{2}\Lambda_{k}^{-2}\mathrm{E}_{q_{k}}[\text{Tr}(\nabla^{2} _{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})]\end{split}$ (15) $$

where $\nabla_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ and $\nabla^{2}_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ are the gradient and Hessian of $\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ respectively, and $\text{Tr}(\cdot)$ returns the trace of the input matrix.

As a special case of steepest descent, gradient descent updates parameter that lies in the Euclidean space. However, our objective is to optimize parameters that represent a distribution, it makes sense to take the steepest descent direction in the distribution space. As in natural gradient descent [], the parameter to be optimized lies on a Riemannian manifold and we choose the steepest descent direction along that manifold. Thus, we choose KL-divergence as the metric of distribution space and take steepest descent in this space. For KL-divergence metric, the natural gradient of parameter of a loss function $\mathcal{L}=\mathrm{E}_{q_{\phi}}[h(\bm{x})]$ is defined as $\tilde{\nabla}_{\phi}\mathcal{L}=\bm{F}_{\phi}^{-1}\nabla_{\phi}\mathcal{L}$ [], where $\bm{F}_{\phi}$ is the Fisher information matrix of $\phi$ given by the variance of the gradient of log probability of parameter, i.e., $\text{Cov}_{q_{\phi}}[\nabla_{\phi}\log q_{\phi}(\bm{x})]$. For $q(\bm{x})=\mathcal{N}(\bm{x};\bm{\mu},\bm{\Sigma})=\mathcal{N}(\bm{x};\bm{\mu}
,\Lambda^{-1}\bm{I})$, the Fisher information matrices of mean and precision are given respectively by $\bm{F}_{\bm{\mu}}=\Lambda\bm{I}$ and $\bm{F}_{\Lambda}=\frac{1}{2}\Lambda^{-2}\bm{I}$. Thus, the natural gradients of parameters $\{\bm{\mu},\Lambda\}$ have concise forms given by

$$ $\displaystyle\begin{split}\tilde{\nabla}_{\bm{\mu}}\mathcal{L}&=\bm{F}_{\bm{ \mu}}^{-1}\nabla_{\bm{\mu}}\mathrm{E}_{q}[h(\bm{x})]=\Lambda^{-1}\mathrm{E}_{q }[\nabla_{\bm{x}}h(\bm{x})]\\ \tilde{\nabla}_{\Lambda}\mathcal{L}&=\bm{F}_{\Lambda}^{-1}\nabla_{\Lambda} \mathrm{E}_{q}[h(\bm{x})]=-\mathrm{E}_{q}[\text{Tr}(\nabla^{2}_{\bm{x}}h(\bm{x }))].\end{split}$ (16) $$

Following the natural gradient given in (16), we update the variational parameters $\phi_{k}=\{\bm{\mu}_{k},\Lambda_{k}=v_{k}^{-1}\}$ of loss function in (14) using natural gradient descent (NGD) as

$$ $\displaystyle\begin{split}\bm{\mu}_{k}&\leftarrow\bm{\mu}_{k}-s_{1}\Lambda_{k} ^{-1}\nabla_{\mu_{k}}KL(q_{k}||p_{k})=\bm{\mu}_{k}+s_{1}\Lambda_{k}^{-1} \mathrm{E}_{q_{k}}[\nabla_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y })]\\ \Lambda_{k}&\leftarrow\Lambda_{k}-2s_{2}\Lambda_{k}^{2}\nabla_{\Lambda_{k}}KL( q_{k}||p_{k})=\Lambda_{k}-s_{2}(N\Lambda_{k}+\mathrm{E}_{q_{k}}[\text{Tr}( \nabla^{2}_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y}))])\end{split}$ (17) $$

where $s_{1}$ and $s_{2}$ are step sizes. We obtain a stochastic NGD update when the expectations of gradient and Hessian matrix in (17) are approximated by sample mean:

$$ $\displaystyle\begin{split}\bm{\mu}_{k}^{(i+1)}&=\bm{\mu}_{k}^{(i)}+s_{1}( \Lambda_{k}^{(i)})^{-1}\frac{1}{L}\sum_{i=1}^{L}\nabla_{\bm{x}_{k}}\log p_{k}( \bm{x}_{k}|\bm{x}_{k+1},\bm{y})|_{\bm{x}_{k}=\bm{x}_{k}^{(i)}\sim q_{k}^{(i)}} \\ \Lambda_{k}^{(i+1)}&=\Lambda_{k}^{(i)}-s_{2}\left(N\Lambda_{k}^{(i)}+\frac{1}{ L}\sum_{i=1}^{L}\text{Tr}(\nabla^{2}_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_ {k+1},\bm{y}))|_{\bm{x}_{k}=\bm{x}_{k}^{(i)}\sim q_{k}^{(i)}}\right)\end{split}$ (18) $$

where $L$ is the number of samples. The stochastic update of parameters of $q_{k}$ converges to the local minima of KL divergence $KL(q_{k}||p_{k})$ which is a Gaussian approximation of the posterior $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$. It is worth noting that we choose stochastic NGD since it achieves a good performance and parameters involved are easy to tune in our experiments. Other optimization methods may be applied. In stochastic NGD update (18), the gradient and Hessian of $\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ are required. We next introduce some approximations to simplify the calculation.

### 4.3 Score-Based Gradient Calculation

From Bayes’ rule, the score of reverse conditional $p(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\frac{p(\bm{x}_{k+1}|\bm{x}_{k})p(\bm{x}_{k}
|\bm{y})}{p(\bm{x}_{k+1}|\bm{y})}$ involved in (18) is given by $\nabla_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\nabla_{\bm{x}_{
k}}\log p(\bm{x}_{k+1}|\bm{x}_{k})+\nabla_{\bm{x}_{k}}\log p(\bm{y}|\bm{x}_{k}
)+\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k})$, where $\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k})$ is the noisy score function which can be approximated by a well-trained score network $\bm{s}_{\bm{\theta}}(\bm{x}_{k},\sigma_{k})$ and $\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k+1}|\bm{x}_{k})$ can be calculated explicitly for both VE and VP diffusion models. For VE diffusion $\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k+1}|\bm{x}_{k})=\frac{\bm{x}_{k+1}-\bm{x}_{
k}}{\sigma_{k+1}^{2}-\sigma_{k}^{2}}$, and for VP diffusion $\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k+1}|\bm{x}_{k})=\frac{\sqrt{1-\beta_{k+1}}}
{\beta_{k+1}}\bm{x}_{k+1}-\frac{1-\beta_{k+1}}{\beta_{k+1}}\bm{x}_{k}$. However, the likelihood score, i.e., the gradient of logarithm conditional $\nabla_{\bm{x}_{k}}\log p(\bm{y}|\bm{x}_{k})$ is hard to handle in general. For linear inverse problems, several SVD based approximations of $\nabla_{\bm{x}_{k}}\log p(\bm{y}|\bm{x}_{k})$ are proposed in [, , ] for linear measurements and Gaussian approximation for general measurements are discussed in []. In [], the authors propose the following approximation that can be applied for general measurements:

$$ $\displaystyle\log p(\bm{y}|\bm{x}_{k})\approx\log p(\bm{y}|\hat{\bm{x}}_{0}( \bm{x}_{k}))$ (19) $$

where $\hat{\bm{x}}_{0}(\bm{x}_{k})$ is the MMSE estimate of $\bm{x}_{0}$. For VE diffusion, according to the Tweedie formula:

$$ $\displaystyle\hat{\bm{x}}_{0}(\bm{x}_{k})=\mathrm{E}_{p(\bm{x}_{0}|\bm{x}_{k}) }[\bm{x}_{0}]=\bm{x}_{k}+\sigma_{k}^{2}\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k}).$ (20) $$

Similarly, for VP diffusion, the MMSE estimation of $\bm{x}_{0}$ given $\bm{x}_{k}$ is

$$ $\displaystyle\hat{\bm{x}}_{0}(\bm{x}_{k})=\mathrm{E}_{p(\bm{x}_{0}|\bm{x}_{k}) }[\bm{x}_{0}]=\frac{1}{\sqrt{\bar{\alpha}_{k}}}(\bm{x}_{k}+(1-\bar{\alpha}_{k} )\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k})).$ (21) $$

The approximation error of (19) can be quantified with Jenson’s Gap as given in []. We choose the likelihood approximation of (19) in our implementation, other approximation methods can be applied to RMP as discussed in Appendix. As a conclusion, the gradient is calculated as

$$ $\displaystyle\nabla_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y}) \approx\nabla_{\bm{x}_{k}}\log p(\bm{x}_{k+1}|\bm{x}_{k})+\gamma_{k}\nabla_{ \bm{x}_{k}}\log p(\bm{y}|\hat{\bm{x}}_{0}(\bm{x}_{k}))+\bm{s}_{\bm{\theta}}( \bm{x}_{k},\sigma_{k}).$ (22) $$

where the parameter $\gamma_{k}$ is added to balance the approximated likelihood score and prior score. We set $\gamma_{k}=\zeta\frac{\|\bm{s}_{\bm{\theta}}(\bm{x}_{k},\sigma_{k})\|_{2}}{\|
\nabla\log p(\bm{y}|\hat{\bm{x}}_{0}(\bm{x}_{k}))\|_{2}}$ where $\zeta$ is a hyper parameter to be tuned for different problems. The idea behind the strategy is that we should always keep a balance between the data score and the likelihood score.

Figure: Algorithm 2 VE/VP-RMP with Score-based Stochastic NGD

### 4.4 Fixed Precision Update

In the update (18), the Hessian matrix $\nabla^{2}_{\bm{x}_{k}}\log p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ is difficult to acquire in general. Also, the complexity involved in the calculation of Hessian may prevent the application of the algorithm. Thus, we introduce an approximation that does not require the calculation of Hessian. According to the Proposition 1, the update of precision $\Lambda_{k}^{(i+1)}$ converges to the precision of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$. Thus, we can fix the update of $\Lambda_{k}^{(i+1)}$ in Algorithm 1 to the precision of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ and only update $\bm{\mu}_{k}$ at each step. According to Proposition 1, for VE diffusion model, if we set $\bm{C}_{\bm{x}_{0}}=v_{\bm{x}_{0}}\bm{I}$, then the inverse of precision is given by $(\Lambda_{k}^{(i)})^{-1}=\frac{(\sigma_{k}^{2}+v_{\bm{x}_{0}})(\sigma_{k+1}^{2
}-\sigma_{k}^{2})}{\sigma_{k+1}^{2}+v_{\bm{x}_{0}}}$. We cannot calculate $(\Lambda_{k}^{(i)})^{-1}$ directly since $v_{\bm{x}_{0}}$ is unknown. However, for the cases that $k$ is close to $T$, $\sigma_{k}^{2}$ is large enough comparing to $v_{\bm{x}_{0}}$. Thus, $(\Lambda_{k}^{(i)})^{-1}$ can be approximated by $\frac{\sigma_{k}^{2}(\sigma_{k+1}^{2}-\sigma_{k}^{2})}{\sigma_{k+1}^{2}}$. For the case that $k$ is close to 0 0, $\sigma_{k+1}^{2},\sigma_{k}^{2}\rightarrow 0$, we have $(\Lambda_{k}^{(i)})^{-1}\approx\sigma_{k+1}^{2}-\sigma_{k}^{2}$. Similarly, for VP diffusion model, if we set $\bm{C}_{\bm{x}_{0}}=v_{\bm{x}_{0}}\bm{I}$, then $(\Lambda_{k}^{(i)})^{-1}=\frac{\frac{\beta_{k+1}}{1-\beta_{k+1}}(1-\bar{\alpha
}_{k}+\bar{\alpha}_{k}v_{\bm{x}_{0}})}{\frac{\beta_{k+1}}{1-\beta_{k+1}}+(1-
\bar{\alpha}_{k}+\bar{\alpha}_{k}v_{\bm{x}_{0}})}$. Since $\beta_{k}\approx 0$ and $\bar{\alpha}_{k}\approx 0$, we have $(\Lambda_{k}^{(i)})^{-1}\approx\beta_{k+1}$. With stochastic NGD based VI and score-based approximations of gradient and Hessian, we summarize an algorithm given in Algorithm 2.

## 5 Experiments

Figure: 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$.
Refer to caption: x2.png

### 5.1 Gaussian Mixture Model

We generate $x_{0}$ from a Gaussian mixture prior $p(x_{0})=\frac{1}{2}\mathcal{N}(x_{0};\mu_{1},v_{1}^{2})+\frac{1}{2}\mathcal{N
}(x_{0};\mu_{2},v_{2}^{2})$ and get measurement from $y=ax_{0}+v_{0}\varepsilon$, where $\varepsilon\sim\mathcal{N}(0,1)$, to demonstrate the rationale of RMP. In the experiments, we set $\mu_{1}=1$, $\mu_{2}=-1$, $v_{1}=v_{2}=0.2$, $v_{0}=0.5$, $a=1$, and $T=1000$. Since the model is simple, the data and likelihood score can be calculated exactly (given in Appendix D). By running VP-RMP in Algorithm 2 (with exact score), the propagation of the reverse mean is shown in Figs. 1 and 2. In the right part of Fig. 1, $x_{T}$ is initialized randomly from $p_{T}(x_{T})$ for various measurement $y$. We see from the plot that the final output of VP-RMP, i.e., the reverse mean $\mu_{0}$ converges to $\mathrm{E}_{p(x_{0}|y)}[x_{0}]$. In the top-left plot of Fig. 2, for various measurement $y$, we initialize $x_{T}=0$. The final output of VP-RMP is almost the same as that in Fig. 1 which converges to $\mathrm{E}_{p(x_{0}|y)}[x_{0}]$. In the top-right plot of Fig. 2, $x_{T}$ is initialized randomly from $p_{T}(x_{T})$ and $y$ is set to 0.2. The reverse mean converges to the posterior mean $\mathrm{E}_{p(x_{0}|y=0.2)}[x_{0}]$ no matter what $x_{T}$ is. In the bottom-left plot of Fig. 2, $x_{T}$ and $y$ are fixed to three cases. We see that the evolution of reverse mean is almost deterministic when the initial point $x_{T}$ and $y$ are fixed. In the bottom-right plot of Fig. 2, the theoretical results of $\mathrm{E}_{p(x_{0}|y=\bar{y})}[x_{0}]$ and estimations of $x_{0}$ using VP-RMP for various measurements $y$ are shown where points A, B, and C correspond to the outputs of VP-RMP for the three cases given in the bottom-left plot respectively. We see from the plot that the outputs of VP-RMP consists with the posterior mean well.

**Table 1: Quantitative results (PSNR, SSIM, FID, LPIPS) of solving linear inverse problems with Gaussian noise ($\epsilon=0.05$) on FFHQ validation dataset. Bold: best, $\underline{\text{Underline}}$: second best.**
| Methods | SR (4 $\times$) | Inpaint (box) | Inpaint (random) | Deblur (Gauss) | Deblur (motion) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ |  |
| VE-RMP (Ours) | $\bm{29.26}$ | $\bm{0.8421}$ | $25.23$ | ${0.8361}$ | $\bm{34.58}$ | $\underline{0.9315}$ | $\bm{28.28}$ | $\bm{0.8214}$ | $\underline{27.74}$ | $\bm{0.8232}$ |
| VP-RMP (Ours) | $\underline{28.89}$ | $\underline{0.8410}$ | $\bm{25.78}$ | $\bm{0.8716}$ | $\underline{34.27}$ | $\bm{0.9291}$ | $\underline{27.96}$ | $\underline{0.8153}$ | $\bm{28.10}$ | 0.7921 |
| DPS | 24.32 | 0.7016 | 24.68 | 0.8182 | 29.85 | 0.8554 | 25.50 | 0.7210 | 23.42 | ${0.6712}$ |
| MCG | 18.16 | 0.2145 | 11.30 | 0.2450 | 11.06 | 0.0812 | 11.65 | 0.1317 | 11.53 | 0.1065 |
| Methods | SR (4 $\times$) | Inpaint (box) | Inpaint (random) | Deblur(Gauss) | Deblur (motion) |  |  |  |  |  |
| FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ |  |
| VE-RMP (Ours) | 89.04 | 0.2278 | 96.36 | 0.2372 | ${36.37}$ | $\underline{0.0948}$ | 94.43 | 0.2515 | 90.63 | $\underline{0.2438}$ |
| VP-RMP (Ours) | $\bm{57.27}$ | $\bm{0.1890}$ | $\bm{28.36}$ | $\bm{0.1250}$ | $\bm{22.41}$ | $\bm{0.0771}$ | $\bm{59.10}$ | $\bm{0.2024}$ | $\bm{62.80}$ | $\bm{0.2383}$ |
| DPS | ${72.44}$ | ${0.2484}$ | ${53.35}$ | ${0.1905}$ | 57.74 | 0.1887 | $\underline{60.18}$ | ${0.2377}$ | ${68.83}$ | ${0.2576}$ |
| MCG | 227.65 | 0.6232 | 443.24 | 0.7929 | 486.09 | 0.8224 | 354.01 | 0.7760 | 461.38 | 0.7222 |

### 5.2 Image Reconstruction

To show the advantage of RMP over posterior sampling based algorithms, we compare our proposed VE/VP-RMP in Algorithm 2 with baseline algorithms for different image reconstruction tasks on FFHQ $256\times 256$ dataset []. For VP-based methods, the pre-trained score network for FFHQ $256\times 256$ is take from [] and for VE-based methods, the pre-trained score networks for FFHQ $256\times 256$ is taken from []. The algorithms we compare include Diffusion Posterior Sampling (DPS) [], Manifold Constrained Gradients (MCG). Algorithms such as Denoising Diffusion Restoration Models (DDRM) and $\Pi$GDM [] are discussed in the appendix part since these methods have difference likelihood score approximation and DDRM can only be applied for linear measurements. Unless otherwise specified, the measurement noise is set to Gaussian. The following metrics are used for comparisons: Frechet Inception Distance (FID) [], Learned Perceptual Image Patch Similarity (LPIPS) [], Peak Signal-to-Noise-Ratio (PSNR) and Structural Similarity Index Measure (SSIM) [].

For VP-based algorithms i.e., DPS, MCG and VP-RMP, $\beta_{k}$ varies linearly with $k$ in range from 0.0001 to 0.02 as given in []. The timesteps for DPS, MCG, and VP-RMP are all fixed to $T=400$. The inner loop size $T_{in}=1$ for VP-RMP. For VE-RMP, we set $T=30$ and $T_{in}=20$ for all tasks. $\{\sigma_{k}\},k=1,2\cdots,T$ is set as a positive geometric sequence as suggested in [] that satisfies $\frac{\sigma_{k}}{\sigma_{k+1}}>1$. In all experiments, we set $\sigma_{k}=\sigma_{min}(\frac{\sigma_{max}}{\sigma_{min}})^{\frac{k-1}{T-1}}$ where $\sigma_{min}=0.01$ and $\sigma_{max}=100$ that matches the pre-trained score network given in []. The settings of step size $s_{1}$ and $\zeta$ for VE/VP-RMP are given in Appendix E. Other parameters in DPS, MCG are set according to the configurations provided by its’ authors. The code of our algorithms is developed under the framework developed by DPS, and the measurement operator used in our comparisons are the same for all algorithms such that the comparisons are fair enough for all algorithms. Our code is available at [https://github.com/neuripsrmp/rmp.git](https://github.com/neuripsrmp/rmp.git).

Linear Inverse Problems: We compare algorithms on linear inverse problems of image reconstruction including super resolution, inpainting (box, random) and deblur (Gaussian deblur, motion deblur). The evaluation results on different metrics are given in Tables 1. We see from Tables 1 that VE/VP-RMP outperforms other algorithms on almost all tasks and VP-RMP achieves the best performance in both FID and LPIPS metrics in all tasks. We present some results in Fig. 3. More results on FFHQ and ImageNet [] datasets can be found in Appendix F.

Nonlinear Inverse Problems: We compare VE/VP-based RMP with DPS on non-linear image reconstruction tasks including phase retrieval and non-linear deblur. The comparison results is presented in Table 2. We see that similarly to the linear cases, RMP achieves better performance comparing to DPS in most of the metrics.

Figure: Figure 3: Part of the results on solving inverse problems with Gaussian noise ($\epsilon=0.05$).
Refer to caption: extracted/5908909/figures/all.png

**Table 2: Quantitative evaluation (PSNR, SSIM, FID, LPIPS) of solving nonlinear inverse problems on FFHQ validation dataset. Bold: best, $\underline{\text{Underline}}$: second best.**
| Methods | Phase retrieval | Nonlinear deblur |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PSNR | SSIM | FID | LPIPS | PSNR | SSIM | FID | LPIPS |  |
| VE-RMP (Ours) | $\bm{27.98}$ | $\bm{0.7860}$ | $\underline{72.84}$ | $\underline{0.1857}$ | $\bm{25.21}$ | $\bm{0.7222}$ | 143.02 | ${0.3401}$ |
| VP-RMP (Ours) | $\underline{25.32}$ | $\underline{0.7632}$ | $\bm{62.11}$ | $\bm{0.1512}$ | $\underline{24.33}$ | $\underline{0.6872}$ | $\bm{66.47}$ | $\bm{0.2219}$ |
| DPS | 12.75 | ${0.4142}$ | ${218.78}$ | 0.5828 | 23.62 | ${0.6696}$ | $\underline{76.85}$ | $\underline{0.2685}$ |

Figure: Figure 4: Performance comparisons of algorithms with different NFE on SR task.
Refer to caption: x3.png

### 5.3 Complexity vs Performance

The complexity of diffusion-based methods dependent on the number of neural function evaluations (NFE). For our algorithm, RMP consists of $T$ outer loops and each outer loop has $T_{in}$ inner loops. The operations in the outer loops can be ignored since only variable assignments and scalar calculation are involved. In each inner loop, two main operations are involved: denoising with score network $\bm{s}_{\bm{\theta}}(\bm{x}_{k},\sigma_{k})$ and stochastic natural gradient descent step which requires the likelihood score. Thus, the NFE of VE/VP-RMP is $TT_{in}$. To show the performance of RMPs with different NFE, we compare the curves of NFE versus the performance of VP-RMP and DPS on SR task in Fig. 4. From the figure, we see that RMP outperforms DPS significantly on all metrics for all NFEs. We also note that for the PSNR metric, VP-RMP with $50$ NFEs achieves a higher PSNR than DPS with $1000$ NFEs.

## 6 Colculsions

In this paper, we propose a score-based variational inference framework for inverse problems. By optimizing the reverse KL divergence at each step of the reverse process and tracking the evolution of reverse mean at each reverse step, a practical algorithm for general inverse problems is proposed. Extensive experiments validate our theoretical results and demonstrate that our proposed algorithm achieves superior performance on various image reconstruction tasks.

## Appendix A Proof of Proposition 1

### A.1 Gaussian Property

First, the reverse conditional $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ is given by

$$ $\displaystyle p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=\frac{p(\bm{x}_{k+1}|\bm{x }_{k})p(\bm{x}_{k}|\bm{y})}{p(\bm{x}_{k+1}|\bm{y})}$ (23) $$

where $\bm{x}_{k}$ and $\bm{x}_{k+1}$ can be viewed as the discretization of continuous diffusion process of $\bm{x}_{t}$. We relate $\{\bm{x}_{k}\}_{k=0}^{T}$ to continuous stochastic process $\{\bm{x}_{t}\}_{t=0}^{1}$ by letting $\bm{x}_{k}=\bm{x}_{t=k\Delta t}$, where $\Delta t=\frac{1}{T}$. The discrete $\{\bm{x}_{k}\}_{k=0}^{T}$ becomes a continuous process in the limit $\Delta t\rightarrow 0$. Thus, we focus on the continuous form of the reverse conditional, i.e., $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})=p(\bm{x}_{t}|\bm{x}_{t+\Delta t},\bm{y})$. The reverse conditional in a continuous form is given by

$$ $\displaystyle p(\bm{x}_{t}|\bm{x}_{t+\Delta t},\bm{y})=\frac{p(\bm{x}_{t+ \Delta t}|\bm{x}_{t})p(\bm{x}_{t}|\bm{y})}{p(\bm{x}_{t+\Delta t}|\bm{y})}.$ (24) $$

As shown in [], the forward process can be characterized by the following SDE:

$$ $\displaystyle d\bm{x}=\bm{f}_{t}(\bm{x})dt+g_{t}(\bm{x})d\bar{\bm{w}}$ (25) $$

where $\bar{\bm{w}}$ is the standard Wiener process. For VE diffusion, $\bm{f}_{t}(\bm{x})=0$ and $g_{t}(\bm{x})=\sqrt{\frac{d\sigma(t)^{2}}{dt}}$, and for VP diffusion, $\bm{f}_{t}(\bm{x})=-\frac{1}{2}\beta(t)\bm{x}$ and $g_{t}(\bm{x})=\sqrt{\beta(t)}$. When $\Delta t\rightarrow 0$ we have

$$ $\displaystyle\bm{x}_{t+\Delta t}-\bm{x}_{t}=\bm{f}_{t}(\bm{x}_{t})\Delta t+g_{ t}(\bm{x}_{t})\sqrt{\Delta t}\bm{\varepsilon}$ (26) $$

where $\bm{\varepsilon}\sim\mathcal{N}(0,\bm{I})$. Then

$$ $\displaystyle\begin{split}p(\bm{x}_{t+\Delta t}|\bm{x}_{t})&=\mathcal{N}(\bm{x }_{t+\Delta t};\bm{x}_{t}+\bm{f}_{t}(\bm{x}_{t})\Delta t,g_{t}^{2}(\bm{x}_{t}) \Delta t)\\ &\propto\exp{\left(-\frac{\|\bm{x}_{t+\Delta t}-\bm{x}_{t}-\bm{f}_{t}(\bm{x}_{ t})\Delta t\|_{2}^{2}}{2g_{t}^{2}(\bm{x}_{t})\Delta t}\right)}\end{split}$ (27) $$

and

$$ $\displaystyle p(\bm{x}_{t}|\bm{x}_{t+\Delta t},\bm{y})\propto\exp{\left(-\frac {\|\bm{x}_{t+\Delta t}-\bm{x}_{t}-\bm{f}_{t}(\bm{x}_{t})\Delta t\|_{2}^{2}}{2g _{t}^{2}(\bm{x}_{t})\Delta t}\!+\!\log p(\bm{x}_{t}|\bm{y})\!-\!\log p(\bm{x}_ {t+\Delta t}|\bm{y})\right)}.$ (28) $$

When $\Delta t\rightarrow 0$, $\log p(\bm{x}_{t}|\bm{y})$ can be expressed by Taylor expansion:

$$ $\displaystyle\log p(\bm{x}_{t+\Delta t}|\bm{y})\approx\log p(\bm{x}_{t}|\bm{y} )+(\bm{x}_{t+\Delta t}-\bm{x}_{t})^{T}\nabla_{\bm{x}_{t}}\log p(\bm{x}_{t}|\bm {y})+\Delta t\frac{\partial}{\partial t}\log p(\bm{x}_{t}|\bm{y}).$ (29) $$

Thus, we have

$$ $\displaystyle\begin{split}p(\bm{x}_{t}|\bm{x}_{t+\Delta t},\bm{y})\propto\exp{ \left(-\frac{\|\bm{x}_{t+\Delta t}-\bm{x}_{t}-(\bm{f}_{t}(\bm{x}_{t})-g_{t}^{2 }(\bm{x}_{t})\nabla_{\bm{x}_{t}}\log p(\bm{x}_{t}|\bm{y}))\Delta t\|_{2}^{2}}{ 2g_{t}^{2}(\bm{x}_{t})\Delta t}+\mathcal{O}(\Delta t)\right)}\end{split}$ (30) $$

which is Gaussian when $\Delta t\rightarrow 0$. Thus, $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ is Gaussian when $\Delta t\rightarrow 0$.

### A.2 VE Diffusion

For VE diffusion models, the mean of $p(\bm{x}_{k}|\bm{y})$ is given by

$$ $\displaystyle\begin{split}\mathrm{E}_{p(\bm{x}_{k}|\bm{y})}[\bm{x}_{k}]&=\int \bm{x}_{k}p(\bm{x}_{k}|\bm{y})d\bm{x}_{k}\\ &=\int\bm{x}_{k}\int p(\bm{x}_{k}|\bm{x}_{0})p(\bm{x}_{0}|\bm{y})d\bm{x}_{0}d \bm{x}_{k}\\ &=\int\int\bm{x}_{k}p(\bm{x}_{k}|\bm{x}_{0})d\bm{x}_{k}p(\bm{x}_{0}|\bm{y})d \bm{x}_{0}\\ &=\int\bm{x}_{0}p(\bm{x}_{0}|\bm{y})d\bm{x}_{0}=\mathrm{E}_{p(\bm{x}_{0}|\bm{y })}[\bm{x}_{0}]\end{split}$ (31) $$

and its covariance matrix is given by

$$ $\displaystyle\begin{split}&Cov_{p(\bm{x}_{k}|\bm{y})}[\bm{x}_{k}]\\ =\ &\int\bm{x}_{k}\bm{x}_{k}^{T}p(\bm{x}_{k}|\bm{y})d\bm{x}_{k}-\mathrm{E}_{p( \bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{ T}\\ =\ &\int\bm{x}_{k}\bm{x}_{k}^{T}\int p(\bm{x}_{k}|\bm{x}_{0})p(\bm{x}_{0}|\bm{ y})d\bm{x}_{0}d\bm{x}_{k}-\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm {E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ =\ &\int\int\bm{x}_{k}\bm{x}_{k}^{T}p(\bm{x}_{k}|\bm{x}_{0})d\bm{x}_{k}p(\bm{x }_{0}|\bm{y})d\bm{x}_{0}-\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{ E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ =\ &\int(\bm{x}_{0}\bm{x}_{0}^{T}+\sigma_{k}^{2}\bm{I})p(\bm{x}_{0}|\bm{y})d \bm{x}_{0}-\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{ 0}|\bm{y})}[\bm{x}_{0}]^{T}\\ =\ &\sigma_{k}^{2}\bm{I}+Cov_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]+\mathrm{E}_{p( \bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{ T}-\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y })}[\bm{x}_{0}]^{T}\\ =\ &\sigma_{k}^{2}\bm{I}+Cov_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]=\sigma_{k}^{2} \bm{I}+\bm{C}_{\bm{x}_{0}}.\end{split}$ (32) $$

Since $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ can be expanded as a Gaussian around $\bm{x}_{k+1}$ using the second order Taylor approximation when $\Delta t\rightarrow 0$ and $p(\bm{x}_{k+1}|\bm{x}_{k})=\mathcal{N}(\bm{x}_{k+1};\bm{x}_{k},(\sigma_{k+1}^{
2}-\sigma_{k}^{2})\bm{I})$, then $p(\bm{x}_{k}|\bm{y})$ can also be expanded around $\bm{x}_{k+1}$ and approximately expressed in the form of a Gaussian when $\Delta t\rightarrow 0$, i.e., $p(\bm{x}_{k}|\bm{y})\approx\mathcal{N}(\bm{x}_{k};\mathrm{E}_{p(\bm{x}_{0}|\bm
{y})}[\bm{x}_{0}],(\sigma_{k}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}}))$. For two multivariate Gaussian distribution $G_{1}(\bm{x})=\mathcal{N}(\bm{x},\bm{\mu}_{1},\bm{\Sigma}_{1})$ and $G_{2}(\bm{x})=\mathcal{N}(\bm{x},\bm{\mu}_{2},\bm{\Sigma}_{2})$, the product $G_{1}(\bm{x})G_{2}(\bm{x})$ is also Gaussian with mean and covariance given by

$$ $\displaystyle\begin{split}\bm{\mu}_{3}&=\bm{\Sigma}_{2}(\bm{\Sigma}_{1}+\bm{ \Sigma}_{2})^{-1}\bm{\mu}_{1}+\bm{\Sigma}_{1}(\bm{\Sigma}_{1}+\bm{\Sigma}_{2}) ^{-1}\bm{\mu}_{2}\\ \bm{\Sigma}_{3}&=\bm{\Sigma}_{1}(\bm{\Sigma}_{1}+\bm{\Sigma}_{2})^{-1}\bm{ \Sigma}_{2}.\end{split}$ (33) $$

Thus, the mean and covariance of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ can be calculated according to (33) and are given respectively by:

$$ $\displaystyle\begin{split}\bm{\mu}_{k}&=(\sigma_{k}^{2}\bm{I}+\bm{C}_{\bm{x}_{ 0}})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\bm{x}_{k+1}+(\sigma_{k+1 }^{2}-\sigma_{k}^{2})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\mathrm{ E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\\ &=\bm{V}_{k,p_{1}}\bm{x}_{k+1}+\bm{V}_{k,p_{2}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y} )}[\bm{x}_{0}]\\ \bm{C}_{k}&=(\sigma_{k+1}^{2}-\sigma_{k}^{2})(\sigma_{k}^{2}\bm{I}+\bm{C}_{\bm {x}_{0}})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}.\end{split}$ (34) $$

### A.3 VP Diffusion

For VP diffusion models, the mean of $p(\bm{x}_{k}|\bm{y})$ is given by

$$ $\displaystyle\begin{split}\mathrm{E}_{p(\bm{x}_{k}|\bm{y})}[\bm{x}_{k}]&=\int \bm{x}_{k}p(\bm{x}_{k}|\bm{y})d\bm{x}_{k}\\ &=\int\bm{x}_{k}\int p(\bm{x}_{k}|\bm{x}_{0})p(\bm{x}_{0}|\bm{y})d\bm{x}_{0}d \bm{x}_{k}\\ &=\int\int\bm{x}_{k}p(\bm{x}_{k}|\bm{x}_{0})d\bm{x}_{k}p(\bm{x}_{0}|\bm{y})d \bm{x}_{0}\\ &=\int\sqrt{\bar{\alpha}_{k}}\bm{x}_{0}p(\bm{x}_{0}|\bm{y})d\bm{x}_{0}=\sqrt{ \bar{\alpha}_{k}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\end{split}$ (35) $$

and its covariance matrix is given by

$$ $\displaystyle\begin{split}Cov_{p(\bm{x}_{k}|\bm{y})}[\bm{x}_{k}]&=\int\bm{x}_{ k}\bm{x}_{k}^{T}p(\bm{x}_{k}|\bm{y})d\bm{x}_{k}-\bar{\alpha}_{k}\mathrm{E}_{p( \bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{ T}\\ &=\int\bm{x}_{k}\bm{x}_{k}^{T}\int p(\bm{x}_{k}|\bm{x}_{0})p(\bm{x}_{0}|\bm{y} )d\bm{x}_{0}d\bm{x}_{k}-\bar{\alpha}_{k}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{ x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ &=\int\int\bm{x}_{k}\bm{x}_{k}^{T}p(\bm{x}_{k}|\bm{x}_{0})d\bm{x}_{k}p(\bm{x}_ {0}|\bm{y})d\bm{x}_{0}-\bar{\alpha}_{k}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x }_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ &=\int(\bar{\alpha}_{k}\bm{x}_{0}\bm{x}_{0}^{T}+(1-\bar{\alpha}_{k})\bm{I})p( \bm{x}_{0}|\bm{y})d\bm{x}_{0}-\bar{\alpha}_{k}\mathrm{E}_{p(\bm{x}_{0}|\bm{y}) }[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ &=(1-\bar{\alpha}_{k})\bm{I}+\bar{\alpha}_{k}(Cov_{p(\bm{x}_{0}|\bm{y})}[\bm{x }_{0}]+\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm{E}_{p(\bm{x}_{0}| \bm{y})}[\bm{x}_{0}]^{T})\\ &\ \ \ \ -\bar{\alpha}_{k}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\mathrm {E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]^{T}\\ &=(1-\bar{\alpha}_{k})\bm{I}+\bar{\alpha}_{k}Cov_{p(\bm{x}_{0}|\bm{y})}[\bm{x} _{0}]=(1-\bar{\alpha}_{k})\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0}}\end{split}$ (36) $$

Similarly to the case of VE diffusion, since $p(\bm{x}_{k+1}|\bm{x}_{k})=\mathcal{N}(\bm{x}_{k+1};\sqrt{1-\beta_{k+1}}\bm{x}
_{k},\beta_{k+1}\bm{I})$, the mean and covariance of $p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})$ are given respectively by:

$$ $\displaystyle\begin{split}\bm{\mu}_{k}&=((1-\bar{\alpha}_{k})\bm{I}+\bar{ \alpha}_{k}\bm{C}_{\bm{x}_{0}})\left(\left(\frac{\beta_{k+1}}{1-\beta_{k+1}}+1 -\bar{\alpha}_{k}\right)\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0}}\right)^{-1} \frac{\bm{x}_{k+1}}{\sqrt{1-\beta_{k+1}}}\\ &\ \ \ \ +\frac{\beta_{k+1}}{1-\beta_{k+1}}\left(\left(\frac{\beta_{k+1}}{1- \beta_{k+1}}+1-\bar{\alpha}_{k}\right)\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0 }}\right)^{-1}\sqrt{\bar{\alpha}_{k}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_ {0}]\\ &=((1-\bar{\alpha}_{k})\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0}})\left((1- \bar{\alpha}_{k+1})\bm{I}+\bar{\alpha}_{k+1}\bm{C}_{\bm{x}_{0}}\right)^{-1} \sqrt{\alpha_{k+1}}\bm{x}_{k+1}\\ &\ \ \ \ +(1-\alpha_{k+1})\left((1-\bar{\alpha}_{k+1})\bm{I}+\bar{\alpha}_{k+1 }\bm{C}_{\bm{x}_{0}}\right)^{-1}\sqrt{\bar{\alpha}_{k}}\mathrm{E}_{p(\bm{x}_{0 }|\bm{y})}[\bm{x}_{0}]\\ &=\bm{V}_{k,p_{3}}\bm{x}_{k+1}+\bm{V}_{k,p_{4}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y} )}[\bm{x}_{0}]\end{split}$ (37) $$

and

$$ $\displaystyle\bm{C}_{k}=\frac{\beta_{k+1}}{1-\beta_{k+1}}((1-\bar{\alpha}_{k}) \bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0}})\left(\left(\frac{\beta_{k+1}}{1- \beta_{k+1}}+1-\bar{\alpha}_{k}\right)\bm{I}+\bar{\alpha}_{k}\bm{C}_{\bm{x}_{0 }}\right)^{-1}.$ (38) $$

## Appendix B Proof of Theorem 1

### B.1 VE Diffusion

For VE diffusion model, if we set $\bm{x}_{k}=\bm{\mu}_{k}(\bm{x}_{k+1}=\bm{\mu}_{k+1},\bm{y})$, $\forall k=0:T-1$, then from Proposition 1, we have

$$ $\displaystyle\begin{split}&\bm{\mu}_{0}(\bm{x}_{1},\bm{y})\\ =\ &\bm{V}_{0,p_{1}}\bm{x}_{1}+\bm{V}_{0,p_{2}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y} )}[\bm{x}_{0}]\\ =\ &\bm{V}_{0,p_{1}}(\bm{V}_{1,p_{1}}\bm{x}_{2}+\bm{V}_{1,p_{2}}\mathrm{E}_{p( \bm{x}_{0}|\bm{y})}[\bm{x}_{0}])+\bm{V}_{0,p_{2}}\mathrm{E}_{p(\bm{x}_{0}|\bm{ y})}[\bm{x}_{0}]\\ =\ &\bm{V}_{0,p_{1}}\bm{V}_{1,p_{1}}\bm{x}_{2}+\bm{V}_{0,p_{1}}\bm{V}_{1,p_{2} }\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]+\bm{V}_{0,p_{2}}\mathrm{E}_{p( \bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\\ =\ &\cdots\\ =\ &\prod_{i=0}^{T-1}\bm{V}_{i,p_{1}}\bm{\mu}_{T}+(\bm{V}_{0,p_{1}}(\bm{V}_{1, p_{1}}(\cdots\bm{V}_{T-2,p_{1}}(\bm{V}_{T-1,p_{2}})+\bm{V}_{T-2,p_{2}})+\bm{V} _{1,p_{2}})+\bm{V}_{0,p_{2}})\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]). \end{split}$ (39) $$

For the first part, the coefficient of $\bm{\mu}_{T}$ equals

$$ $\displaystyle\prod_{i=0}^{T-1}\bm{V}_{i,p_{1}}=\prod_{i=0}^{T-1}(\sigma_{i}^{2 }\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{i+1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}= (\sigma_{0}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{T}^{2}\bm{I}+\bm{C}_{\bm{x} _{0}})^{-1}.$ (40) $$

For the second part, $\forall k=1:T-1$, we have

$$ $\displaystyle\begin{split}\bm{V}_{k-1,p_{1}}\bm{V}_{k,p_{2}}+\bm{V}_{k-1,p_{2} }&=(\sigma_{k-1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{k}^{2}\bm{I}+\bm{C}_{ \bm{x}_{0}})^{-1}(\sigma_{k+1}^{2}-\sigma_{k}^{2})(\sigma_{k+1}^{2}\bm{I}+\bm{ C}_{\bm{x}_{0}})^{-1}\\ &\ \ \ \ +(\sigma_{k}^{2}-\sigma_{k-1}^{2})(\sigma_{k}^{2}\bm{I}+\bm{C}_{\bm{x }_{0}})^{-1}\\ &=(\sigma_{k+1}^{2}-\sigma_{k-1}^{2})(\sigma_{k+1}^{2}\bm{I}+\bm{C}_{\bm{x}_{0 }})^{-1}.\end{split}$ $$

Similarly, we get

$$ $\displaystyle(\bm{V}_{0,p_{1}}(\bm{V}_{1,p_{1}}(\cdots\bm{V}_{T-2,p_{1}}(\bm{V }_{T-1,p_{2}})+\bm{V}_{T-2,p_{2}})+\bm{V}_{1,p_{2}})+\bm{V}_{0,p_{2}})=(\sigma _{T}^{2}-\sigma_{0}^{2})(\sigma_{T}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}$ (41) $$

Thus, we have

$$ $\displaystyle\bm{\mu}_{0}=(\sigma_{0}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})(\sigma_{T }^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\bm{\mu}_{T}+(\sigma_{T}^{2}-\sigma_{0}^{ 2})(\sigma_{T}^{2}\bm{I}+\bm{C}_{\bm{x}_{0}})^{-1}\mathrm{E}_{p(\bm{x}_{0}|\bm {y})}[\bm{x}_{0}].$ (42) $$

When $\sigma_{T}^{2}\rightarrow\infty$, $\bm{\mu}_{0}(\bm{x}_{1},\bm{y})\rightarrow\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[
\bm{x}_{0}]$, which is the posterior mean.

### B.2 VP Diffusion

Similarly to VE diffusion, for VP diffusion model,

$$ $\displaystyle\begin{split}&\bm{\mu}_{0}(\bm{x}_{1},\bm{y})\\ =\ &\bm{V}_{0,p_{3}}\bm{x}_{1}+\bm{V}_{0,p_{4}}\mathrm{E}_{p(\bm{x}_{0}|\bm{y} )}[\bm{x}_{0}]\\ =\ &\prod_{i=0}^{T-1}\bm{V}_{i,p_{3}}\bm{\mu}_{T}+(\bm{V}_{0,p_{3}}(\bm{V}_{1, p_{3}}(\cdots\bm{V}_{T-2,p_{3}}(\bm{V}_{T-1,p_{4}})+\bm{V}_{T-2,p_{4}})+\bm{V} _{1,p_{4}})+\bm{V}_{0,p_{4}})\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]) \end{split}$ (43) $$

For the first part, the coefficient of $\bm{\mu}_{T}$ equals

$$ $\displaystyle\begin{split}\prod_{i=0}^{T-1}\bm{V}_{i,p_{3}}&=\prod_{i=0}^{T-1} ((1-\bar{\alpha}_{i})\bm{I}+\bar{\alpha}_{i}\bm{C}_{\bm{x}_{0}})\left((1-\bar{ \alpha}_{i+1})\bm{I}+\bar{\alpha}_{i+1}\bm{C}_{\bm{x}_{0}}\right)^{-1}\sqrt{ \alpha_{i+1}}\\ &=\sqrt{\bar{\alpha}_{T}}((1-\bar{\alpha}_{0})\bm{I}+\bar{\alpha}_{0}\bm{C}_{ \bm{x}_{0}})((1-\bar{\alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{x}_{0}})^{ -1}.\end{split}$ (44) $$

For the second part, $\forall i=1:T-1$, we have

$$ $\displaystyle\begin{split}\bm{V}_{i-1,p_{3}}\bm{V}_{i,p_{4}}+\bm{V}_{i-1,p_{4} }&=((1-\bar{\alpha}_{i-1})\bm{I}+\bar{\alpha}_{i-1}\bm{C}_{\bm{x}_{0}})\left(1 -\bar{\alpha}_{i}+\bar{\alpha}_{i}\bm{C}_{\bm{x}_{0}}\right)^{-1}\\ &\ \ \ \ \cdot\sqrt{\alpha_{i}}(1-\alpha_{i+1})\left((1-\bar{\alpha}_{i+1})\bm {I}+\bar{\alpha}_{i+1}\bm{C}_{\bm{x}_{0}}\right)^{-1}\sqrt{\bar{\alpha}_{i}}\\ &\ \ \ \ +(1-\alpha_{i})\left((1-\bar{\alpha}_{i})\bm{I}+\bar{\alpha}_{i}\bm{C }_{\bm{x}_{0}}\right)^{-1}\sqrt{\bar{\alpha}_{i-1}}\\ &=(1-\alpha_{i}\alpha_{i+1})\sqrt{\bar{\alpha}_{i-1}}((1-\bar{\alpha}_{i+1}) \bm{I}+\bar{\alpha}_{i+1}\bm{C}_{\bm{x}_{0}})^{-1}.\end{split}$ (45) $$

Similarly, we obtain

$$ $\displaystyle\begin{split}&\bm{V}_{0,p_{3}}(\bm{V}_{1,p_{3}}(\cdots\bm{V}_{T-2 ,p_{3}}(\bm{V}_{T-1,p_{4}})+\bm{V}_{T-2,p_{4}})+\bm{V}_{1,p_{4}})+\bm{V}_{0,p_ {4}}\\ =\ &(1-\alpha_{0}\alpha_{1}\cdots\alpha_{T})\sqrt{\bar{\alpha}_{0}}((1-\bar{ \alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{x}_{0}})^{-1}\\ =\ &(1-\bar{\alpha}_{T})((1-\bar{\alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{ \bm{x}_{0}})^{-1}.\end{split}$ (46) $$

Thus, we have

$$ $\displaystyle\begin{split}\bm{\mu}_{0}(\bm{x}_{1},\bm{y})=\ &\sqrt{\bar{\alpha }_{T}}((1-\bar{\alpha}_{0})\bm{I}+\bar{\alpha}_{0}\bm{C}_{\bm{x}_{0}})((1-\bar {\alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{x}_{0}})^{-1}\bm{\mu}_{T}\\ &+(1-\bar{\alpha}_{T})((1-\bar{\alpha}_{T})\bm{I}+\bar{\alpha}_{T}\bm{C}_{\bm{ x}_{0}})^{-1}\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]\end{split}$ (47) $$

and $\bm{\mu}_{0}\rightarrow\mathrm{E}_{p(\bm{x}_{0}|\bm{y})}[\bm{x}_{0}]$ as $\bar{\alpha}_{T}\rightarrow 0$.

## Appendix C Proof of Proposition 2

We have

$$ $\displaystyle\begin{split}p(\bm{x}_{k}|\bm{x}_{k+1:T},\bm{y})&=\frac{p(\bm{x}_ {k:T},\bm{y})}{p(\bm{x}_{k+1:T},\bm{y})}\\ &=\frac{\int p(\bm{x}_{0})p(\bm{y}|\bm{x}_{0})p(\bm{x}_{k:T}|\bm{x}_{0})d\bm{x }_{0}}{\int p(\bm{x}_{0})p(\bm{y}|\bm{x}_{0})p(\bm{x}_{k+1:T}|\bm{x}_{0})d\bm{ x}_{0}}\\ &=\frac{p(\bm{x}_{k+1:T}|\bm{x}_{k})\int p(\bm{x}_{0})p(\bm{y}|\bm{x}_{0})p( \bm{x}_{k}|\bm{x}_{0})d\bm{x}_{0}}{p(\bm{x}_{k+2:T}|\bm{x}_{k+1})\int p(\bm{x} _{0})p(\bm{y}|\bm{x}_{0})p(\bm{x}_{k+1}|\bm{x}_{0})d\bm{x}_{0}}\\ &=\frac{p(\bm{x}_{k+1}|\bm{x}_{k})\int p(\bm{x}_{0})p(\bm{y}|\bm{x}_{0})p(\bm{ x}_{k}|\bm{x}_{0})d\bm{x}_{0}}{\int p(\bm{x}_{0})p(\bm{y}|\bm{x}_{0})p(\bm{x}_ {k+1}|\bm{x}_{0})d\bm{x}_{0}}\\ &=\frac{p(\bm{x}_{k},\bm{x}_{k+1},\bm{y})}{p(\bm{x}_{k+1},\bm{y})}=p_{k}(\bm{x }_{k}|\bm{x}_{k+1},\bm{y}).\end{split}$ (48) $$

We minimize the KL divergence between variational joint posterior $q(\bm{x}_{0:T}|\bm{y})=q(\bm{x}_{T}|\bm{y})\prod_{k=T-1}^{0}q_{k}(\bm{x}_{k}|
\bm{x}_{k+1},\bm{y})$ and joint posterior $p(\bm{x}_{0:T}|\bm{y})=p(\bm{x}_{T}|\bm{y})\prod_{k=T-1}^{0}p(\bm{x}_{k}|\bm{x
}_{k+1:T},\bm{y})$ to obtain the optimal variational distribution $q$:

$$ $\displaystyle\begin{split}\mathcal{F}&=\int q(\bm{x}_{0:T}|\bm{y})\log\frac{q( \bm{x}_{0:T}|\bm{y})}{p(\bm{x}_{0:T}|\bm{y})}d\bm{x}_{0:T}\\ &=\sum_{k=T-1}^{0}\int q(\bm{x}_{0:T}|\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x} _{k+1},\bm{y})}{p(\bm{x}_{k}|\bm{x}_{k+1:T},\bm{y})}d\bm{x}_{0:T}\\ &=\sum_{k=T-1}^{0}\int q(\bm{x}_{k:T}|\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x} _{k+1},\bm{y})}{p(\bm{x}_{k}|\bm{x}_{k+1:T},\bm{y})}d\bm{x}_{k:T}\\ &=\sum_{k=T-1}^{0}\int q(\bm{x}_{k:T}|\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x} _{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}d\bm{x}_{k:T}\\ &=\sum_{k=T-1}^{0}\int q(\bm{x}_{k:k+1}|\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{ x}_{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}d\bm{x}_{k}d\bm{x}_{k+ 1}\\ &=\sum_{k=T-1}^{0}\int q(\bm{x}_{k+1}|\bm{y})\int q_{k}(\bm{x}_{k}|\bm{x}_{k+1 },\bm{y})\log\frac{q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm {x}_{k+1},\bm{y})}d\bm{x}_{k}d\bm{x}_{k+1}\end{split}$ (49) $$

where the last line follows from

$$ $\displaystyle\begin{split}q(\bm{x}_{k:k+1}|\bm{y})&=\int q(\bm{x}_{T}|\bm{y}) \prod_{i=k}^{T-1}q_{i}(\bm{x}_{i}|\bm{x}_{i+1},\bm{y})d\bm{x}_{k+2:T}\\ &=\int q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})q_{k+1}(\bm{x}_{k+1}|\bm{x}_{k+2}, \bm{y})q(\bm{x}_{k+2:T}|\bm{y})d\bm{x}_{k+2:T}\\ &=q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})\int q_{k+1}(\bm{x}_{k+1}|\bm{x}_{k+2}, \bm{y})q(\bm{x}_{k+2:T}|\bm{y})d\bm{x}_{k+2:T}\\ &=q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})q(\bm{x}_{k+1}|\bm{y}).\end{split}$ (50) $$

Thus, the minimization of $\mathcal{F}$ is equivalent to the minimization of

$$ $\displaystyle\mathcal{F}_{k}=\int q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})\log \frac{q_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm{y})}{p_{k}(\bm{x}_{k}|\bm{x}_{k+1},\bm {y})}d\bm{x}_{k},\forall k=0,1,\cdots T-1.$ (51) $$

## Appendix D Gaussian Mixture Model

The data prior of $x_{0}$ is given by

$$ $\displaystyle\begin{split}p(x_{0})&=\frac{1}{2}\mathcal{N}(x_{0};\mu_{1},v_{1} ^{2})+\frac{1}{2}\mathcal{N}(x_{0};\mu_{2},v_{2}^{2})\\ &=\frac{1}{2}(p_{1}(x_{0})+p_{2}(x_{0})).\end{split}$ (52) $$

Then, the data score is

$$ $\displaystyle\nabla_{x_{0}}\log p(x_{0})=\frac{1}{2p(x_{0})}\left(-p_{1}(x_{0} )\frac{x_{0}-\mu_{1}}{\sigma_{1}^{2}}-p_{2}(x_{0})\frac{x_{0}-\mu_{2}}{\sigma_ {2}^{2}}\right).$ (53) $$

The measurement $y=ax_{0}+\sigma_{0}\varepsilon_{1}$ and for VP diffusion, $\bm{x}_{k}=\sqrt{\bar{\alpha}_{k}}x_{0}+\sqrt{1-\bar{\alpha}_{k}}\varepsilon_{2}$ where $\varepsilon_{1}$ and $\varepsilon_{2}$ are i.i.d. from $\mathcal{N}(0,1)$. Then,

$$ $\displaystyle x_{0}=\frac{x_{k}-\sqrt{1-\bar{\alpha}_{k}}\varepsilon_{2}}{ \sqrt{\bar{\alpha}_{k}}}$ (54) $$

and

$$ $\displaystyle y=a\frac{x_{k}-\sqrt{1-\bar{\alpha}_{k}}\varepsilon_{2}}{\sqrt{ \bar{\alpha}_{k}}}+\sigma_{0}\varepsilon_{1}.$ (55) $$

Thus, $p(y|x_{k})\sim\mathcal{N}(\frac{ax_{k}}{\sqrt{\bar{\alpha}_{k}}},\frac{a^{2}(1
-\bar{\alpha}_{k})}{\bar{\alpha}_{k}}+\sigma_{0}^{2})$ and the likelihood score

$$ $\displaystyle\nabla_{x_{k}}\log p(y|x_{k})=\frac{a}{\sqrt{\bar{\alpha}_{k}}} \frac{y-\frac{a}{\sqrt{\bar{\alpha}_{k}}}x_{k}}{\frac{a^{2}(1-\bar{\alpha}_{k} )}{\bar{\alpha}_{k}}+\sigma_{0}^{2}}.$ (56) $$

## Appendix E RMP Implementation Details

### E.1 Parameters setting

In our experiments, the step size $s_{1}$ and $\zeta$ for different tasks on FFHQ dataset for VP/VE-RMP are set according to Table 3 and Table 4 respectively.

**Table 3: Hyper-parameters of VP-RMP for different tasks**
| parameter | SR ($4\times$) | box | random | Gauss | motion | nonlinear | phase |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $s_{1}$ | 0.9 | 0.6 | 0.6 | 0.9 | 0.9 | 0.8 | 0.5 |
| $\zeta$ | 0.15 | 0.3 | 0.3 | 0.5 | 0.5 | 0.8 | 0.5 |

**Table 4: Hyper-parameters of VE-RMP for different tasks**
| parameter | SR ($4\times$) | box | random | Gauss | motion | nonlinear | phase |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $s_{1}$ | 0.1 | 0.05 | 0.1 | 0.05 | 0.05 | 0.05 | 0.05 |
| $\zeta$ | 0.07 | 0.1 | 0.25 | 0.15 | 0.15 | 0.15 | 0.35 |

### E.2 Running Time Comparisons

On FFHQ dataset, the running time of RMP and DPS for different image reconstruction tasks are shown in Fig. 5. From Fig. 5, we see that VE/VP-RMP achieves a better performance with less running time for all tasks.

Figure: Figure 5: Running time of RMP and DPS for different image reconstruction tasks.
Refer to caption: x5.png

### E.3 Approximation of Likelihood Score for Natural Images

In addition to the approximation of likelihood score we used in the main results for image reconstruction tasks, other approximation methods [, ] of likelihood score can be employed to RMP. We present the performance of VP-RMP with the approximation method used in $\Pi$GDM in Table 5. It is shown that the performance of VP-RMP with likelihood approximation method in [] achieves similar performance as that of VP-RMP with DPS approximation.

**Table 5: The performance of VP-RMPs with different approximations of likelihood score on FFHQ dataset.**
| Methods | SR (4 $\times$) |  |  |  |
| --- | --- | --- | --- | --- |
| PSNR | SSIM | FID | LPIPS |  |
| VP-RMP (with approximation in (19)) | ${28.86}$ | ${0.8413}$ | ${59.7112}$ | ${0.1825}$ |
| VP-RMP (with approximation in $\Pi$GDM []) | ${28.70}$ | ${0.8477}$ | ${61.8101}$ | ${0.1988}$ |

## Appendix F Image Reconstruction Results

The image reconstruction results for inpainting (random) and Gaussian deblur tasks on FFHQ dataset are shown in Fig. 6 and more results of various image reconstruction tasks on ImageNet are presented in Fig. 7.

Figure: Figure 6: Inpainting (Random) and Gaussian deblur ($\epsilon=0.05$).
Refer to caption: extracted/5908909/figures/inpaint_random_ffhq.png

Figure: Figure 7: SR ($\times 4$) and Gaussian deblur ($\epsilon=0.05$).
Refer to caption: extracted/5908909/figures/sr_imagenet.png

**Table 6: Quantitative evaluation of algorithms on ImageNet dataset.**
| Methods | SR (4 $\times$) | Inpaint (box) | Inpaint (random) | Deblur (Gauss) | Deblur (motion) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ |  |
| VP-RMP (Ours) | $\bm{26.21}$ | $\bm{0.7261}$ | $\bm{19.29}$ | $\bm{0.7830}$ | $\bm{29.71}$ | ${0.8720}$ | $\bm{24.65}$ | ${0.6806}$ | $\bm{25.30}$ | $\bm{0.6533}$ |
| DPS | 21.24 | 0.5655 | 17.96 | 0.7274 | 27.11 | 0.7862 | 22.49 | $0.5996$ | 19.72 | ${0.5118}$ |
| DDRM | ${24.28}$ | ${0.7137}$ | ${19.01}$ | ${0.7801}$ | ${29.01}$ | $\bm{0.8742}$ | 23.81 | $\bm{0.7109}$ | - | - |
| Methods | SR (4 $\times$) | Inpaint (box) | Inpaint (random) | Deblur(Gauss) | Deblur (motion) |  |  |  |  |  |
| FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | LPIPS $\downarrow$ |  |
| VP-RMP (Ours) | $\bm{75.76}$ | $\bm{0.3021}$ | $\bm{106.73}$ | $\bm{0.2156}$ | $\bm{26.45}$ | ${0.1400}$ | $\bm{85.92}$ | ${0.3105}$ | $\bm{71.00}$ | $\bm{0.2933}$ |
| DPS | ${187.14}$ | ${0.4147}$ | ${146.54}$ | ${0.2664}$ | 76.20 | 0.2424 | ${86.28}$ | ${0.3224}$ | ${144.43}$ | ${0.3795}$ |
| DDRM | ${117.82}$ | ${0.3086}$ | ${119.65}$ | ${0.2232}$ | ${27.18}$ | $\bm{0.1066}$ | ${95.07}$ | $\bm{0.2931}$ | - | - |

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: See above
soundness: 2
strengths: The topic of solving inverse problems with score-based models is interesting.
summary: The authors propose Reverse Mean Propagation (RMP) as a new framework that leverages score-based variational inference to solve inverse problems by tracking the posterior mean at each step of the reverse diffusion process. RMP is claimed to be more computationally efficient than traditional diffusion-based methods that rely on generating multiple samples from the posterior, as it can achieve the posterior mean directly.
weaknesses: * The posterior mean is not the only relevant quantity; access to
  samples from the posterior distribution also enables uncertainty
  quantification. This work provides only an approximation to the mean,
  sacrificing sample generation for computational efficiency, which may
  not even be fully realized unless unjustified approximations are made
  (see more below).

* While RMP aims to reduce complexity relative to sampling methods, it still requires nested optimization loops and numerous neural network evaluations.

* The method relies heavily on approximating the likelihood score, $\nabla_{x_k} \log p(y | x_k)$, which is acknowledged as "hard to handle in general." The chosen approximation—replacing the likelihood score with the score at the MMSE estimate of $x_0$—introduces an error quantified by Jensen's Gap.


* The performance of RMP depends on multiple hyperparameters, such as step sizes and the likelihood score balancing parameter. Sensitivity to these hyperparameters and their generalizability across datasets and inverse problems remain unclear. 

* Algorithm 2 uses a fixed-precision update to avoid computing the Hessian further introducing approximation errors.

* The authors note that the reverse conditional probability is strictly Gaussian only as $\Delta t \to 0$, an assumption that may not hold in practice, introducing further errors.

* In high-dimensional settings, computing the MMSE estimate of $x_0$ at each step, as required for likelihood score approximation, can be computationally intensive. The manuscript lacks a detailed analysis of RMP’s computational cost in high-dimensional scenarios.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. I'm trying to understand proposition 1 from a continuous perspective, as it derives the continuous-time limit of the distribution $p({\bf x} _k|{\bf x} _{k+1},{\bf y})$. However, it appears that these two approaches differ significantly.

   Consider the forward process $d{\bf x} _t=\sqrt{\frac{d}{dt}(\sigma _t^2)}d{\bf B} _t$, $t\in[0,1]$, ${\bf x} _0\sim p({\bf x} _0|{\bf y})$. We know that $p({\bf x} _t|{\bf x} _0)={\cal N}({\bf x} _0,\sigma _t^2{\bf I})$, and the marginal distribution of ${\bf x} _t$ is $p _t({\bf x} _t|{\bf y})$. The backward process is $d{\bf x} _t=-\frac{d}{dt}(\sigma _t^2)\nabla\log p _t({\bf x} _t|{\bf y})dt+\sqrt{\frac{d}{dt}(\sigma _t^2)}d{\bf B}^\gets _t$, where ${\bf B}^\gets _t$ is a backward BM. If initialized at ${\bf x} _1\sim p _1({\bf x} _1|{\bf y})$, then ${\bf x} _0\sim p _0({\bf x} _0|{\bf y})$. We can express the score as $\nabla\log p _t({\bf x} _t|{\bf y})=\frac{1}{\sigma _t^2}(\mathbb{E}({\bf x} _0|{\bf x} _t,{\bf y})-{\bf x} _t)$, and thus the discretization scheme can be written as ${\bf x} _{t-\Delta t}\approx {\bf x} _t+\frac{\sigma _t^2-\sigma _{t-\Delta t}^2}{\sigma _t^2}(\mathbb{E}({\bf x} _0|{\bf x} _t,{\bf y})-{\bf x} _t)+{\cal N}(0,(\sigma^2 _t-\sigma^2 _{t-\Delta t}){\bf I})$. (I hope my calculation is correct.) 
   

   Compare this with proposition 1, the posterior variance is independent of ${\bf y}$, while the posterior mean involves $\mathbb{E}({\bf x} _0|{\bf x} _t,{\bf y})$ rather than $\mathbb{E}({\bf x} _0|{\bf y})$. Could you clarify the differences and connections between these two approaches?

2. I noticed an extra term $N\Lambda _k$ in line 2 of equations 17 and 18, which does not appear in equation 16. Is this an error? Additionally, since we are performing gradient descent, the plus and minus signs in equation 17 seem incorrect after the "=" sign.

3. Minor comments:

- Please use parenthetical citation `\citep{}` instead of `\cite{}` when the authors' names aren't part of the text. For example, the first sentence should be "Diffusion models (Sohl-Dickstein et al. 2015; Song & Ermon 2019; Ho et al. 2020; Song et al. 2020a; Rombach et al. 2022) have shown impressive performance for image generation."

- In proposition 1, the reverse conditional should be *approximately* Gaussian when $\Delta t\to0$. $k=0\cdots T-1$ should be $k=0,\cdots,T-1$.

- In line 234: $p({\bf x} _k|{\bf x} _{k+1:T},{\bf y})$ should be $p({\bf x} _k|{\bf x} _{k+1},{\bf y})$.

- Throughout the paper, $KL$ should be ${\rm KL}$, and $Cov$ should be ${\rm Cov}$.

- In equation 15, $\nabla _{{\bm\Lambda} _k}$ should be $\nabla _{{\Lambda} _k}$.

- In line 287, there is no need to mention $\Lambda _k=v _k^{-1}$.

- In algorithm 2, the roles of $T _s$ and $T _{in}$ need to be stated.

- I suggest the author(s) replace figures 1 and 2 with higher resolution versions, especially the bottom-right in figure 2.

- The title of section 6 has a spelling error.
soundness: 3
strengths: The paper focuses on predicting $\mathbb{E}({\bf x} _0|{\bf y})$ rather than sampling from the distribution $p({\bf x} _0|{\bf y})$, which sets it apart from most existing works. The Gaussian VI approach for learning the reverse conditional mean is innovative, and the experimental results presented are also compelling. The writing is clear, and the mathematical derivations are solid. Overall, this paper is of high quality and meets the standards of the ICLR conference.
summary: This paper tackles the challenge of solving inverse problems using diffusion models and introduces a framework called Reverse Mean Propagation (RMP) to recover the posterior mean of the latent variable ${\bf x} _0$ given the measurement ${\bf y}$. RMP is based on the insight that the reverse conditional distributions $p({\bf x} _k|{\bf x} _{k+1},{\bf y})$ in diffusion models are approximately Gaussian, with their mean and variance depending on the unknown quantities $\mathbb{E}({\bf x} _0|{\bf y})$ and ${\rm Cov}({\bf x} _0|{\bf y})$. The paper proposes a Gaussian variational inference approach to learn the reverse conditional mean and approximate the reverse conditional variance.
weaknesses: The introduction section would benefit from a concise summary of the main contributions. Additionally, a detailed literature review of existing methods for inverse problems, particularly those related to VI, should be included to highlight the novelty of this approach.

In the experiments, the proposed method RMP is compared with DPS, MCG, DDRM, and $\Pi$GDM. While RMP demonstrates strong performance against these baselines, it is important to note that most of them (DPS, MCG, and $\Pi$GDM) do not require additional training, and can be used in a plug-and-play manner with any pretrained diffusion model and suitable measurement process ${\bf y}={\cal A}({\bf x} _0)+{\bf w} _0$. However, RMP incurs extra overhead due to the stochastic NGD update for Gaussian VI. Furthermore, in section 5.3, the comparison of NFE between RMP and DPS focuses on inference time, and does not account for training complexity, which may render the comparison a little bit unfair.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1: In line 703, the authors present the estimated form $p(x_k|y) \approx \mathcal{N}(x_k; \mathbb{E}\_{p(x_0|y)}[x_0], (\sigma_k^2 I + C_{x_0}))$. 
However, $x_k|y$ is not generally an asymptotically Gaussian distribution when $\Delta t \rightarrow 0$. Could the authors provide further clarification on this estimation?
soundness: 3
strengths: S1: The paper proposes a novel framework, Reverse Mean Propagation, which significantly reduces the complexity of estimating the posterior mean compared to other methods that rely on generating samples from the posterior distribution.

S2: The details of each component of the RMP algorithm are well presented, including the reverse mean updates and estimation using stochastic natural gradient descent. The discussion on estimating the trace of the Hessian matrix $\nabla^2_{x_k} \log p_k$ is also clear and informative.
summary: This paper introduces a novel framework called Reverse Mean Propagation (RMP) for estimating the posterior mean. The framework involves tracking the mean and covariance in the evolution of the conditional reverse diffusion process.
This is achieved by minimizing the reverse KL divergence at each step through a variational inference with stochastic natural gradient descent.
In experiments, the authors validate their method on image reconstruction tasks. The results show that RMP outperforms than the existing baselines, including DPS and MCG methods.
weaknesses: W1: In line 148, the authors discuss the variance preserving (VP) diffusion model and the variance exploding (VE) diffusion model, mentioning their different training approaches. It would be helpful to explain that the VP scheme, like the VE scheme, is also equivalent to learning the score functions of perturbed data distributions[1]. Clarifying this point could provide more comprehensive background information.

W2: In Table 2, the FID results appear inconsistent with the previously reported results of the DPS method in Table 4 of [2].

W3: The comparison methods in the paper are relatively limited. Including additional baselines, such as PnP-ADMM [3], would be beneficial. 

W4: As a specific type of VI method, there is limited discussion on previous works that directly estimate the posterior mean. 

[1] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." ICLR 2021.

[2] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." ICLR 2023.

[3] Chan, Stanley H., Xiran Wang, and Omar A. Elgendy. "Plug-and-play ADMM for image restoration: Fixed-point convergence and applications." IEEE Transactions on Computational Imaging 3.1 (2016): 84-98.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. It seems that the overall complexity is the product of the number of the outer loops $T$ and the number of the inner loops $T_{in}$. Would different choices of these numbers affect the performance? For example, one can make $T$ larger, in the hope that the correspondng posterior $p(x_k|x_{k+1},y)$ would become easier to approximate so that the convergence would be faster for VI and $T_{in}$ can therefore be smaller.

2. It seems that VE-RMP tends to require smaller $T$, can the author explain it?
soundness: 3
strengths: 1. The paper is clearly written and estimating the posterior mean directly seems new for score-based inverse problems.
2. The authors show that in the limit of continous time, the end point of reverse mean propogation would recover the exact posterior mean, with different conditions for VP and VE diffusions respectively.
3. Natural gradient descent is employed for optimization.
4. The authors have evaluated their algorithms thoroughly, and also conducted an ablation study on the complexity vs performance.
summary: This paper presents a score-based variational inference method for inverse problems. The key ingredient is a framework called reverse mean propogation (RMP) that targets the posterior mean directly. The authors show that RMP can be implemented by solving a variational problem which decomposes into independent sub variational problems at each reverse step. Natural gradient descent is used for learning variational parameters. The authors show that the proposed algorithm outperforms state-of-the-art algorithms in various inverse problems.
weaknesses: 1. Although the derivation of the proposed algorithm seems right, there are lots of approximations that would damage the validity of the methods. For example, proposition 1 and theorem 1 only hold in the continuous limit, not sure how accurate it would be, especially when the step size is large (or equivalently, $T$ is small). Can the authors provide an empirical analysis or theoretical bounds on the approximation error for finite T? For example, can the authors show how the algorithm's performance changes as T varies, to demonstrate how closely it approaches the continuous limit in practice?

2. It seems that the proposed method also uses the same approximation as in DPS for the likelihood score (section 4.3). As this is the main reason that DPS would perform poorly, can the author explain why your algorithm would perform better than DPS? any intuition? Can the authors provide a detailed comparison or ablation study isolating the impact of the likelihood score approximation versus other aspects of their approach?

3. There are some hyperparameters (i.e., $\gamma_k$ in equation (22) and the stepsize $s_1$) that would be crucial for the performance. The paper lacks ablation studies on it. Can the authors conduct and report sensitivity analyses for $\gamma_k$ and $s_1$, showing how performance metrics change across reasonable ranges for these parameters?

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - Is such an expensive two-loop circulation common in this field? If not, the authors should discuss this in their limitations. If so, a comparison with related works is recommended.
soundness: 3
strengths: - This paper is well written and easy to follow.
- Although diffusion-based approaches for solving inverse problems are extensively investigated, the estimating strategy for the posterior mean in equation (7) is novel.
- The gradient estimation in equation (22) with a pre-trained score model is non-trivial. It proves that this method can provide reliable posterior mean estimation in the experiments.
- Figure 2 clearly illustrates how RMP propagates the mean through the diffusion process, which highlights the difference from the models targeting the distributional information.
summary: This paper targets the estimating the posterior mean in the inverse problems.
To do this, the authors first argument the intractable posterior distribution with a diffusion process, and derive the mean and variance propagation rule in the backward diffusion process.
To estimate the intractable transitions incorporated in the above propagation rule, variational inference objectives as well as the natural gradient update are developed.
The above estimation relies on a pretrained score model on the clean data.
The experiments on toy examples and high-resolution images both demonstrate the effectiveness of the proposed method compared to baselines.
weaknesses: - The citation formats, \citep and \citet, need to be more carefully distinguished.
- RMP requires an additional pre-trained score network to estimate the posterior mean. The evaluation in the experiments does not count the time consumption for training the score network. Moreover, the difficulty of obtaining such a score model can depend on the modeling objects, where images are only a special instance. If this is not possible, the authors should discuss this in their limitations. 
- For each fixed measurement $y$, RMP requires two loops to obtain a posterior mean estimation. The outer loop iterates the timesteps in the diffusion process, and the  inner loop optimizes transition means by variational inference. Adding more computational cost comparison in experiments would help better evaluate the proposed method.
- For a different $y$, we need to re-run the expensive RMP to establish a posterior mean estimation. Especially, the $\mu$ and $\Lambda$ need to be re-estimated. Is an amortized inference approach possible here?

Meta-review: This paper proposes a variational inference method for score-based approach to Bayesian inverse problems. The main contribution is a reverse mean propogation method (RMP) that directly approximates the posterior mean, which is implemented by solving multiple variational problems via natural gradient descent. While reviewers and I agree this is an interesting idea, major concerns included compromised accuracy due to numerous approximations and (more thorough) comparison to existing approaches and they seemed to have not completed resolved. Recognizing the potential of the idea, I recommend the authors take the discussions into consideration and submit again.

Decision: Reject