Score-Based Variational Inference for Inverse Problems

We thank reviewer 4zYw for careful reading and constructive comments.

W1: The citation formats, \citep and \citet, need to be more carefully distinguished.

We appreciate the reviewer's comment and have addressed this issue in our manuscript.

W2: 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.

RMP is a type of Plug and Play method that requires only a pre-trained score network for inference. The pre-trained score networks used for image reconstruction in our experiments are taken from previous works. It is fair not to include the training time of the score network in our comparisons, since all comparison methods require such a score network. To train such a score network, a large dataset is required, and denoising score matching is used for the training of the score network with the network architecture chosen, such as NCSN. For some applications, the clean dataset may be unavailable, which is a major obstacle. This limitation is common to all score-based methods used for inverse problems, including DPS, DDRM, and Pi-DGM. We have added a discussion on the limitations of RMP in the appendix of our manuscript.

W3: 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.

We add VP-RMP with PiGDM approximation and PnP-ADMM to the running time comparison in Fig. 8 in Appendix E.2 of our manuscript. Also, we add a comparison of the performance of VP-RMP versus loop number in Fig. 7.

W4: 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?

Our proposed RMP algorithm is a non-amortized variational inference method that necessitates the updating of $\mu$ and $\Lambda$ at each reverse step to accommodate distinct measurements $y$. Amortized variational inference learns a shared inference network that connects measurements to the posterior distribution of latent variables [1]. As shown in previous works [2], amortized variational inference incurs an amortization gap for general settings. Non-amortized methods typically perform better than amortized methods since the inference is based on the current measurement rather than relying on a generalized network for all measurements.

To solve inverse problems using score-based data priors, we believe that amortized methods can be applied to applications with loose requirements for the solution, where the focus is more on computational complexity. This is because amortized methods reduce inference complexity and provide direct output estimation. This is an interesting topic that deserves more future research.

[1] Ganguly, Ankush, Sanjana Jain, and Ukrit Watchareeruetai. "Amortized Variational Inference: A Systematic Review." Journal of Artificial Intelligence Research 78 (2023): 167-215.

[2] Margossian, Charles C., and David M. Blei. "Amortized Variational Inference: When and Why?." arXiv preprint arXiv:2307.11018 (2023).

Q1: 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.

To the best of our knowledge, our proposed method, RMP, is the first to employ variational inference with a score-based prior for tracking the density evolution of the conditional reverse process, while the mean simultaneously converging on the posterior mean. The two-loop form is not commonly used in score-based methods for inverse problems. We have addressed this limitation in Appendix G of our manuscript.

We thank reviewer nsm7 for careful reading and constructive comments.

W1. 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?

Our proposed RMP algorithm is essentially a score-based implementation of the Forker-Planck equation [1] for a conditional reverse process that tracks the evolution of the probability density of that process. RMP approximates this evolution using the reverse mean and covariance, as presented in Proposition 1 of our manuscript. The implementation of RMP, based on variational inference, serves as an approximation method for the discretization of the evolution process. This approximation error for finite $T$ is generally an open question since $p(x\_0|y)$ is generally intractable. The RMP algorithm presented in Algorithm 2 is a discrete realization of the RMP framework, meaning that RMP approaches the posterior mean strictly as T tends to infinity. The theoretical performance of RMP when the $T$ is finite can be characterized by the variational loss. We next give an intuition on the theoretical bound of approximation error. From the discrete variational loss F defined in (49), a continuous loss can also be defined similarly. The theoretical performance of RMP when T is finite can be characterized by the variational loss. Next, we will provide an intuition regarding the theoretical bound on the approximation error. Similarly, from the discrete variational loss F defined in (49), a continuous loss can also be defined. The difference between the discrete loss and the continuous loss can be interpreted as the approximation error bound. The continuous loss can be upper bounded by the discrete loss using Jensen's inequality, and Jensen's gap also provides the approximation bound. We have added an experiment in the appendix of our manuscript to demonstrate the performance trend of VP-RMP compared to the outer loop number $T$, as shown in Fig. 7.

[1] Risken, H. "The Fokker-Planck Equation." (1996).

W2. 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?

Different from the previous works that draw samples from the posterior distribution of data using a SDE, our proposed RMP algorithm is essentially a score-based implementation of the Fokker-Planck equation for a conditional reverse process that tracks the evolution of the probability density of that process. RMP approximates this evolution using the reverse mean and covariance. The implementation of RMP, based on variational inference, serves as an approximation method for the discretization of the evolution process. This represents a fundamental difference between our method and previous sampling-based methods. We demonstrate this difference with the following example: We set the parameters of VP-RMP as follows: $T\_{in} = 1$, step size $s1 = 1/\alpha\_{k+1}$, and $x\_k = \mu\_k$ (which is not sample from the variational distribution). In this case, RMP reduces to DPS. Therefore, the reason RMP performs better than DPS is that at each reverse step, DPS samples from the posterior, while RMP tracks the evolution of the density using the mean and covariance (approximated by precision), which leads to better performance in PSNR.

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

We conduct a parameter sensitivity analysis for the hyperparameters stepsize $s_1$ and $\gamma_k$ involved in the RMP algorithm. In Figs. 5 and 6 of Appendix E.1, the performance of VP-RMP and VE-RMP versus hyperparameters, in terms of four metrics, is shown in the form of a 2D colormesh. From the colormesh, we see that all four metrics exhibit similar variation when the hyperparameters change, and these metrics vary smoothly with respect to the hyperparameters. Therefore, we can choose the hyperparameters based on the colormesh.

Thanks for the response and additional results. I will maintain my score.

Regarding your discussion on the connection to DPS, the coefficient before term $\nabla \log p(y|x_k)$ is set to 1 in DPS, while in your case it is set as a hyperparameter $\gamma_k$. What would DPS perform if one also treat this coefficient as a hyperparameter (just as you did in RMP), instead of fixing at 1?

We thank reviewer k8qn for carefully reading and constructive comments

W1: 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.

We thank reviewer k8qn for the constructive comments. We have added a discussion of our contributions and a literature review of existing methods for inverse problems and their relation to VI in the introduction. Please see our manuscript.

W2: In the experiments, the proposed method RMP is compared with DPS, MCG, DDRM, and ΠGDM. While RMP demonstrates strong performance against these baselines, it is important to note that most of them (DPS, MCG, and Π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 𝑦=𝐴(𝑥0)+𝑤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.

Similar to DPS, MCG, and ΠGDM, the proposed method RMP is also a plug and play method that does not require additional training if a pretrained score network is available. With a pretrained score network, RMP works well for suitable measurement process since the stochastic NGD involved in RMP updates the mean based on the natural gradient approximated by score networks. Therefore, the comparison of inference times between our method and DPS is fair, as no extra training is required for our method.

Q1:

We thank reviewer k8qn for careful reading and constructive comment. Your calculation is completely correct and can lead to an alternative way to understand Proposition 1 if we take one step further. Specifically, if we replace the score ∇log⁡𝑝𝑡(𝑥𝑡|𝑦) using the Gaussian approximation with the mean and variance given by (31) and (32), i.e., $\nabla log⁡ 𝑝\_𝑡(𝑥\_𝑡|𝑦) =(\sigma\_k^2 I + C\_{x\_0})^{-1}(E[x\_0|y] - x\_t)$ and take expactation, we arrive at the VE case of Proposition 1. It worth noting that the Gaussian approximation holds when $\Delta t\rightarrow 0$, which comes from the Gaussian property of the reverse diffusion process [1,2].

[1]Anderson, Brian DO. "Reverse-time diffusion equation models." Stochastic Processes and their Applications 12.3 (1982): 313-326.

[2] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

Q2: I noticed an extra term $𝑁\Lambda\_𝑘$ 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.

Dear Reviewer k8qn, this is not an error. Equation (16) represents the natural gradient expression of the loss function defined in Line 278, and Equations (17) and (18) represent the natural gradient updates of the parameters based on the loss function defined in (14). Additionally, the $𝑁\Lambda_𝑘$ term comes from the natural gradient of the first part of loss function in equation (14), specifically the gradient of $-N/2\log(2\pi/\Lambda_𝑘)$. Furthermore, the minus signs in equation 17 arise from the loss function in equation (14).

Q3: We thank the reviewer for carefully reading. We have correct these problems in our manuscript.

We thank reviewer ZpSm for careful reading of our work and the positive comments.

Reply to reviewer's comments:

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.

We thank Reviewer ZpSm for the valuable advice. We have clarified this point in Section 2.3.

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

The inconsistency arises from the fact that most of the recovery results of DPS differ significantly from the original image, which indicates a high rate of unsuccessful recovery, as discussed in the Limitations section of the DPS paper. The metrics reported by DPS for phase retrieval in their paper only account for successful recovery. Additionally, we note that the comparison results of DPS were obtained by running the code provided by the authors, and the parameters of DPS were set according to their code instructions.

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

We have added PnP-ADMM for comparison in our manuscript. The implementation of PnP-ADMM is from the scico library. The parameters for PnP-ADMM are set to $\rho=0.2$ and $\text{maxiter}=12$, and the image denoiser is chosen as the pretrained DnCNN denoiser. From the comparison results, we observe that RMP significantly outperforms PnP-ADMM in all metrics.

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

In this work, we present an evolution analysis of the transition probability density of the conditional reverse diffusion process, which is described by the Fokker-Planck Equation [1,2]. Different from the previous works that draw samples from the posterior distribution of data using a stochastic differential equation (SDE), our proposed RMP algorithm is essentially a score-based implementation of the Fokker-Planck equation for a conditional reverse process that tracks the evolution of the probability density of that process. RMP approximates this evolution using the reverse mean and covariance. The implementation, based on variational inference, serves as an approximation method for the discretization of the evolution process. This represents a fundamental difference between our method and previous sampling-based methods.

We have added a discussion on related works in the introduction of our manuscript. Please see the highlighted part in the introduction.

[1] Risken, H. "The Fokker-Planck Equation." (1996)

[2] Jordan, Richard, David Kinderlehrer, and Felix Otto. "The variational formulation of the Fokker-Planck equation." SIAM journal on mathematical analysis 29.1 (1998): 1-17.

Q1: In line 703, the authors present the estimated form $𝑝(𝑥\_𝑘|𝑦) \approx \mathcal{N}(𝑥\_𝑘;𝐸 𝑝(𝑥\_0|𝑦)[𝑥\_0],(\sigma\_𝑘^2 𝐼+𝐶\_{𝑥\_0}))$. However, $p(𝑥\_𝑘|𝑦)$ is not generally an asymptotically Gaussian distribution when $\Delta 𝑡\rightarrow 0$. Could the authors provide further clarification on this estimation?

It is true that $p(x_k|y)$ is generally not a Gaussian distribution as $\Delta t \to 0$. In Appendix A.2, we aim to calculate the mean and covariance of $p(x_k|x_{k+1},y)$, which is asymptotically Gaussian as $\Delta t \to 0$ [1,2]. Since $p(x_k|x_{k+1})$ is also Gaussian, we can approximate $p(x_k|y)$ as Gaussian when $\Delta t \to 0$ in the calculation.

[1]Anderson, Brian DO. "Reverse-time diffusion equation models." Stochastic Processes and their Applications 12.3 (1982): 313-326.

[2] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

We thank the reviewer for careful reading and constructive comments. We summarize our work as follows:

The Fokker-Planck equation [1,2] describes the evolution of the probability density of a stochastic process which is deterministic. In this work, we present an evolution analysis of the probability density for the conditional reverse diffusion process, as described by the Fokker-Planck equation. Based on the analysis, we propose a variational inference framework that minimizes the reverse KL divergence, referred to as RMP, since it propagates the mean at each reverse step. Different from the previous works that draw samples from the posterior with a reverse diffusion process, our proposed RMP algorithm is essentially an approximation of the Fokker-Planck equation for a conditional reverse process that tracks the deterministic evolution of the probability density. This represents a fundamental difference between our method and previous methods. Unlike sampling-based methods that have to sample from the posterior multiple times and average out to get the posterior mean which is time-consuming, RMP converges to the posterior mean.

The main contributions of our work are summarized as follows:

• We characterize the evolution of the transaction probability density of the conditional reverse diffusion process in terms of its mean and covariance. Based on this characterization, we propose the RMP framework, which approaches the posterior mean.
• We connect the variational inference with conditional diffusion process and propose an implementation of the RMP framework based on stochastic natural gradient descent with score-based prior and suitable approximation that reduces complexity.
• We conduct extensive experiments to demonstrate the validity of the theory of RMP framework and shows that RMP outperforms state-of-the-art algorithms on reconstruction performance with lower computational complexity in various problems.

[1] Risken, H. "The Fokker-Planck Equation." (1996)

[2] Jordan, Richard, David Kinderlehrer, and Felix Otto. "The variational formulation of the Fokker-Planck equation." SIAM journal on mathematical analysis 29.1 (1998): 1-17.

Reply to reviewer's comments:

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

  This is not an issue with RMP. For example, in our image reconstruction experiments, we set the inner loop number of VP-RMP to 1, resulting in an algorithm without nested loops. The comparison of VP-RMP with existing sampling-based methods, such as DPS, DDRM, and PnP-ADMM, demonstrates that VP-RMP outperforms these methods with fewer NFEs and reduced running time.

• W2: The method relies heavily on approximating the likelihood score, ∇𝑥𝑘log⁡𝑝(𝑦|𝑥𝑘), which is acknowledged as "hard to handle in general." The chosen approximation—replacing the likelihood score with the score at the MMSE estimate of 𝑥0—introduces an error quantified by Jensen's Gap.

There are many existing methods for approximating likelihood scores, and the approximation introduced by DPS is just one of them. In Appendix E.3, we compare our approach with the likelihood approximation methods introduced by DPS and PiGDM, demonstrating that RMP performs well with these approximations.

• W3: 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.

  We perform a parameter sensitivity analysis for the hyperparameters (stepsize s1 and 𝛾𝑘) involved in the RMP algorithm, as detailed in the Appendix. Figures 5 and 6 in Appendix E.1 present the performance of VP-RMP and VE-RMP concerning these hyperparameters, measured by four metrics, in the form of a 2D colormesh. From the colormesh, we observe that all four metrics show similar variations and change smoothly with respect to the hyperparameters. Therefore, we can select the hyperparameters based on the colormesh, as shown in Tables 3 and 4.

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

  Our proposed RMP is the first method to approximate the evolution of the density of the conditional reverse diffusion process. Unlike previous methods that sample from the posterior distribution of data using a SDE, which is a stochastic sampling method, our method tracks the density evolution, which is deterministic. It true that the fixed-precision update to avoid computing complexity introduces approximation error when T is small. We show that the approximation involved in fixed-precision update is robust to many applications as shown in our experiments of Gaussian mixture model and various image reconstruction tasks.