Ensemble Kalman Diffusion Guidance: A Derivative-free Method for Inverse Problems

## Abstract

Abstract When solving inverse problems, it is increasingly popular to use pre-trained diffusion models as plug-and-play priors. This framework can accommodate different forward models without re-training while preserving the generative capability of diffusion models. Despite their success in many imaging inverse problems, most existing methods rely on privileged information such as derivative, pseudo-inverse, or full knowledge about the forward model. This reliance poses a substantial limitation that restricts their use in a wide range of problems where such information is unavailable, such as in many scientific applications. To address this issue, we propose Ensemble Kalman Diffusion Guidance (EnKG) for diffusion models, a derivative-free approach that can solve inverse problems by only accessing forward model evaluations and a pre-trained diffusion model prior. We study the empirical effectiveness of our method across various inverse problems, including scientific settings such as inferring fluid flows and astronomical objects, which are highly non-linear inverse problems that often only permit black-box access to the forward model.

## 1 Introduction

The idea of using pre-trained diffusion models (Song et al., ; Ho et al., ) as plug-and-play priors for solving inverse problems has been increasingly popular and successful across various applications including medical imaging (Song et al., ; Sun et al., ), image restoration (Chung et al., ; Wang et al., ), and music generation (Huang et al., ). A key advantage of this approach is its flexibility to accommodate different problems without re-training while maintaining the expressive power of diffusion models to capture complex and high-dimensional prior data distributions. However, most existing algorithms rely on privileged information of the forward model, such as its derivative (Chung et al., ; Song et al., ), pseudo-inverse (Song et al., ), or knowledge of its parameterization (Chung et al., ). This reliance poses a substantial limitation that prevents their use in problems where such information is generally unavailable. For instance, in many scientific applications (Oliver et al., ; Evensen & Van Leeuwen, ; Iglesias, ), the forward model consists of a system of partial differential equations whose derivative or pseudo-inverse is generally unavailable or even undefined.

The goal of this work is to develop an efficient method that only requires black-box access to the forward model and pre-trained diffusion model for solving general inverse problems. Such an approach could significantly extend the range of diffusion-based inverse problems studied in the current literature, unlocking a new class of applications – especially many scientific applications.
The major challenge here arises from the difficulty of approximating the gradient of a general forward model with only black-box access. The standard zero-order gradient estimation methods are known to scale poorly with the problem dimension (Berahas et al., ).

In this paper, we first propose a generic prediction-correction (PC) framework using an optimization perspective that includes existing diffusion guidance-based methods (Chung et al., ; Song et al., ; ; Peng et al., ; Tang et al., ) as special cases.
The proposed PC framework consists of two main steps for each iteration: 1) the prediction step, which is a numerical integration step of the unconditional reverse ordinary differential equation (ODE) or stochastic differential equation (SDE); 2) the correction step, which is a proximal operator that maps the current point towards the high likelihood region while remaining close to the original trajectory. This viewpoint not only provides an alternative view and additional insights of the existing methods but also allows us to introduce new designs to develop a fully derivative-free guidance method. Our approach, called Ensemble Kalman Diffusion Guidance (EnKG), uses an ensemble of particles to estimate the guidance term. A key technical step is to use the empirical covariance matrix of this ensemble instead of relying on a usual scalar weight for the $L^{2}$ regularization in the correction step. Using this step allows us to derive the update rule for each ensemble member without the need for gradient, following the same approach from classic ensemble Kalman methods (Evensen, ; Schillings & Stuart, ).
Consequently, EnKG is a fully derivative-free diffusion guidance method, and can solve challenging inverse problems that only permit black-box access to the forward model.

Contributions

- •
We present a generic prediction-correction (PC) framework that provides an alternative interpretation of guided diffusion, as well as additional insights of the existing methods.
- •
Building upon the PC framework, we propose Ensemble Kalman Diffusion Guidance (EnKG), a fully derivative-free approach that leverages pre-trained model in a plug-and-play manner for solving general inverse problems. EnKG only requires black-box access to the forward model and can accommodate different forward models without any re-training.
- •
We evaluate on various inverse problems including the standard imaging tasks and scientific problems like the Navier-Stokes equation and black-hole imaging. On the standard imaging tasks,
the proposed EnKG outperforms baseline methods on the nonlinear phase retrieval by a large margin. More importantly, EnKG remains effective and efficent on the highly nonlinear inverse problem where the gradient information is genuinely inaccessible such as Navier-Stokes, while the baseline methods fail to generate reasonable predictions.

## 2 Background & Problem Setting

##### Problem setting

Let $G:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ denote the forward model that maps the true unobserved source ${\bm{x}}$ to observations ${\bm{y}}$. We consider the following general inverse problem given by:

$$ ${\bm{y}}=G({\bm{x}})+\xi,\quad{\bm{x}}\in\mathbb{R}^{n},{\bm{y}},\xi\in\mathbb {R}^{m}$ (1) $$

where we only have black-box access to $G$, which we generally assume to be non-linear. $\xi$ represents the observation noise which is typically Gaussian, i.e., $\xi\sim{\mathcal{N}}(0,\Gamma)$, and ${\bm{y}}$ represents the observation which is usually sparse, i.e., $m<n$.
In this paper, solving the inverse problem amounts to finding the ${\bm{x}}$ that has high posterior $p({\bm{x}}|{\bm{y}})$, which is often expressed via Bayes’s rule as $p({\bm{x}}|{\bm{y}})\propto p({\bm{y}}|{\bm{x}})\cdot p({\bm{x}})$. Here $p({\bm{x}})$ is the prior distribution over source signals (which we instantiate using a pre-trained diffusion model), and $p({\bm{y}}|{\bm{x}})$ is defined as (1). Because we only have black-box access to $G$, we can only sample from $p({\bm{y}}|{\bm{x}})$, and do not know its functional form. For simplicity, we focus on finding the max a posteriori estimate: $\arg\max_{\bm{x}}p({\bm{y}}|{\bm{x}})\cdot p({\bm{x}})$.

##### Diffusion models

Diffusion models (Song et al., ; Karras et al., ) capture the prior $p({\bm{x}})$ implicitly using a diffusion process, which includes a forward process and backward process. The forward process transforms a data distribution ${\bm{x}}_{0}\sim p_{\mathrm{data}}$ into a Gaussian distribution ${\bm{x}}_{T}\sim{\mathcal{N}}(0,\sigma^{2}(T){\bm{I}})$ defined by a pre-determined stochastic process. The Gaussian distribution is often referred to as noise, and so the forward process ($t$ going from 0 to $T$) is typically used to create training data (iteratively noisier versions of ${\bm{x}}_{0}\sim p_{\mathrm{data}}$) for the diffusion model. The backward process ($t$ going from $T$ to 0), which is typically learned in a diffusion model, is the standard generative model and operates by sequentially denoising the noisy data into clean data, which can be done by either a probability flow ODE or a reverse-time stochastic process. In this paper, we consider the following probability flow ODE throughout the whole paper for light notation without loss of generality because every other probability flow ODE is equivalent to it up to a simple reparameterization as shown by  Karras et al. ():

$$ ${\mathrm{d}}{\bm{x}}_{t}=-\dot{\sigma}(t)\sigma(t)\nabla_{{\bm{x}}_{t}}\log p_ {t}({\bm{x}}_{t}){\mathrm{d}}t.$ (2) $$

Training a diffusion model amounts to training the so-called score function $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$, which we assume is already completed. Given a trained diffusion model, we can sample $p({\bm{x}})$ by integrating (2) starting from a random noise sample. As such, the focus of this paper is not on training diffusion models, but rather on guiding pre-trained diffusion models for solving inverse problems, discussed next.

##### Diffusion guidance for inverse problems

Arguably the most popular approach to solving inverse problems with a pre-trained diffusion model is guidance-based (Chung et al., ; Wang et al., ; Kawar et al., ; Song et al., ; Zhu et al., ; Rout et al., ; Chung et al., ; Tang et al., ). Guidance-based methods are originally interpreted as the conditional reverse diffusion process targeting the posterior distribution. For ease of notation and clear presentation, we use the probability flow ODE to represent the reverse process and rewrite it with Bayes Theorem.

$$ $\displaystyle{\mathrm{d}}{\bm{x}}_{t}$ $\displaystyle=-\dot{\sigma}(t)\sigma(t)\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x} }_{t}|{\bm{y}}){\mathrm{d}}t,$ $\displaystyle=-\dot{\sigma}(t)\sigma(t)\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x} }_{t}){\mathrm{d}}t-\dot{\sigma}(t)\sigma(t)\nabla_{{\bm{x}}_{t}}\log p_{t}({ \bm{y}}|{\bm{x}}_{t}){\mathrm{d}}t,$ (3) $$

where $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ is the unconditional score and the $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{y}}|{\bm{x}}_{t})$ is the guidance from likelihood.
In practice, the unconditional score is approximated by a pre-trained diffusion model $s_{\theta}({\bm{x}}_{t},t)$ and the weight of the likelihood term is considered as a hyperparameter. The corresponding reverse dynamics is given by

$$ ${\mathrm{d}}{\bm{x}}_{t}=-\dot{\sigma}(t)\sigma(t)s_{\theta}({\bm{x}}_{t},t){ \mathrm{d}}t-w_{t}\nabla_{{\bm{x}}_{t}}\log\hat{p}_{t}({\bm{y}}|{\bm{x}}_{t}){ \mathrm{d}}t,$ (4) $$

where $w_{t}$ is the adaptive time-dependent weight. The design of $w_{t}$ in Eq. (4) varies across different methods but it is typically not related to $\dot{\sigma}(t)\sigma(t)$ that Eq. (2) suggests, which makes it hard to interpret from a posterior sampling perspective. In this paper, we will take an optimization perspective develop a useful interpretation for designing our proposed algorithm.

One key issue with Eq. (4) is that many algorithms for sampling along Eq. (4) assume the ability to access the gradient $\nabla_{{\bm{x}}_{t}}\log\hat{p}_{t}({\bm{y}}|{\bm{x}}_{t})$. When this gradient is unavailable (e.g., when one only has black-box access to $\hat{p}_{t}({\bm{y}})$), then one must develop a derivative-free approach, which is the central technical contribution of our work.

##### Derivative-free optimization

Our approach relies on derivative-free optimization in order to solve the inverse problem
Derivative-free optimization (DFO) refers to the optimization of a function where access to its derivative is not available, and only black-box access is available.
Traditional DFO algorithms are broadly classified as direct search, which includes the coordinate search (Fermi, ), stochastic finite-difference approximations of the gradient (Spall, ), and Nelder-Mead simplex methods (Nelder & Mead, ), and model-based methods which include descent and trust region methods (Conn et al., ; Bortz & Kelley, ). Recent stochastic zero-order optimization techniques involve approximating the gradient via Gaussian smoothing (Nesterov & Spokoiny, ); these gradient estimates can be plugged into algorithms like stochastic gradient descent (SGD) and directly parallel their first-order counterparts (LeCun et al., ). Our approach builds on top of Ensemble Kalman Inversion (EKI) which is a popular scientific computing method used primarily for solving physical inverse problems (Iglesias et al., ; Calvello et al., ). From an optimization perspective, the method can be seen as performing gradient decent with a particle-based approximation to the derivative of the forward operator (Schillings & Stuart, ; Kovachki & Stuart, ).

## 3 Method

To develop our Ensemble Kalman Diffusion Guidance (EnKG) method, we first provide an interpretation of diffusion guidance through the lens of the prediction-correction framework. EnKG can be viewed as an instantiation which enables derivative-free approximation of the guidance term.

### 3.1 Prediction-correction interpretation of guidance-based methods

Inspired by the idea of the Predictor-Corrector continuation method in numerical analysis (Allgower & Georg, ), we show how to express the guidance-based methods within the following prediction-correction framework.
Suppose the time discretization scheme is $T=t_{0}>t_{1}\dots>t_{N}=0$. Let $w_{i}=w_{t_{i}}$ for light notation.
As illustrated in Algorithm 1, guidance-based methods for inverse problems can be summarized into the following steps.

##### Prior prediction step

This step is simply a numerical integration step of the unconditional probability flow ODE, i.e., by moving one step along the unconditional ODE trajectory:

$$ ${\bm{x}}^{\prime}_{i}={\bm{x}}_{i}-\dot{\sigma}(t_{i})\sigma(t_{i})s_{\theta} \left({\bm{x}}_{i},t_{i}\right)(t_{i+1}-t_{i}).$ (5) $$

Figure: Algorithm 1 Generic Guidance-based Method (ODE version)

Figure: Figure 1: Illustration of the prediction-correction interpretation for guidance-based methods on a 1D Gaussian mixture example. From left to right, the probability flow ODE gradually transforms $p_{t}({\bm{x}}_{t})$ from a Gaussian into a mixture of two Gaussians. The grey lines indicate the vector field of the probability flow. The prediction step is simply a numerical integration step over the probability flow trajectory. The correction step moves towards the MAP estimator while staying near to the initial prediction point.
Refer to caption: extracted/5889642/figs/pc_demo.png

##### Log-likelihood estimation step

This step estimates the log-likelihood $\log p({\bm{y}}|{\bm{x}}_{t})$:

$$ $\log\hat{p}({\bm{y}}|{\bm{x}}_{i})\approx\log p({\bm{y}}|{\bm{x}}_{i}).$ $$

##### Guidance correction step

This step solves the following optimization problem that formulates a compromise between maximizing the log-likelihood and being near ${\bm{x}}\prime_{i}$:

$$ ${\bm{x}}_{i+1}=\arg\min_{{\bm{x}}_{i+1}}\frac{\|{\bm{x}}_{i+1}-{\bm{x}}^{ \prime}_{i}\|^{2}_{2}}{2w_{i}}-\log\hat{p}({\bm{y}}|{\bm{x}}_{i+1}),$ (6) $$

where the larger guidance scale $w_{i}$ gives the solution point near the MAP estimator and smaller value leads to small movement towards the MAP estimator. Eq. (6) is essentially a proximal operator (Parikh et al., ) if $w_{i}$ is lower bounded by a positive number. Thus it is firmly non-expansive in the convex setting and will converge to MAP estimator when applied recursively.

To solve Eq. (6) efficiently in practice, we may use first-order Taylor approximation of $\log\hat{p}({\bm{y}}|{\bm{x}}_{i+1})$ at ${\bm{x}}_{i}^{\prime}$, which gives:

$$ $\log\hat{p}({\bm{y}}|{\bm{x}}_{i+1})=\log\hat{p}({\bm{y}}|{\bm{x}}_{i}^{\prime })+\nabla^{\top}\log\hat{p}({\bm{y}}|{\bm{x}}_{i}^{\prime})\left({\bm{x}}_{i+1 }-{\bm{x}}_{i}^{\prime}\right)+O(\Delta{\bm{x}}_{i}^{2}).$ (7) $$

Substituting the approximation Eq. (7) into the correction step (6) gives:

$$ $\displaystyle{\bm{x}}_{i+1}$ $\displaystyle\approx\arg\min_{{\bm{x}}_{i+1}}\frac{\|{\bm{x}}_{i+1}-{\bm{x}}^{ \prime}_{i}\|^{2}_{2}}{2w_{i}}-\log\hat{p}({\bm{y}}|{\bm{x}}_{i}^{\prime})- \nabla^{\top}\log\hat{p}({\bm{y}}|{\bm{x}}_{i}^{\prime})\left({\bm{x}}_{i+1}-{ \bm{x}}_{i}^{\prime}\right)$ (8) $\displaystyle={\bm{x}}_{i}^{\prime}+w_{i}\nabla\log\hat{p}({\bm{y}}|{\bm{x}}_{ i}^{\prime}),$ (9) $$

which can recover the gradient step structure of most existing guidance-based methods (Chung et al., ; Song et al., ; ; Mardani et al., ).

##### Putting it together.

Figure 1 visually demonstrates the Prediction-Correction interpretation in a 1D Gaussian mixture example. This depiction shows how the guidance-based methods iteratively step towards the MAP estimator while staying close to the initial unconditional generation trajectory defined by the prediction step. More importantly, this PC framework allows more degrees of freedom in design. In the next section, we introduce ensemble covariance matrix to the step size and derive the update rule that does not require gradient information. To stimulate further research, Section 6 discusses other potential design innovations within the PC framework.

### 3.2 Our approach: Ensemble Kalman Diffusion Guidance

Figure: Algorithm 2 Our method: Ensemble Kalman Diffusion Guidance (EnKG).

We now demonstrate how the correction step can be performed in a derivative-free manner using the idea of statistical linearization. Our overall approach is described in Algorithm 2.

##### Likelihood estimation.

The likelihood term can be factorized as follows:

$$ $p({\bm{y}}|{\bm{x}}_{i})=\int p({\bm{y}}|{\bm{x}}_{N})p({\bm{x}}_{N}|{\bm{x}}_ {i}){\mathrm{d}}{\bm{x}}_{N}=\mathbb{E}_{{\bm{x}}_{N}\sim p({\bm{x}}_{N}|{\bm{ x}}_{i})}p({\bm{y}}|{\bm{x}}_{N}),$ (10) $$

which is intractable in general. We use the following Monte Carlo approximation:

$$ $p({\bm{y}}|{\bm{x}}_{i})=\mathbb{E}_{{\bm{x}}_{N}\sim p({\bm{x}}_{N}|{\bm{x}}_ {i})}p({\bm{y}}|{\bm{x}}_{N})\approx p({\bm{y}}|\hat{{\bm{x}}}_{N}),$ (11) $$

where $\hat{{\bm{x}}}_{N}$ is obtained by running the Probability Flow ODE solver $\phi$ starting at ${\bm{x}}_{i}$. One attractive property of this estimate compared to popular ones based on $\mathbb{E}[{\bm{x}}_{N}|{\bm{x}}_{i}]$ and isotropic Gaussian approximations in previous works Chung et al. (); Song et al. (; ) is that our approximation stays on data manifold but the Gaussian approximations include additive noise that do not live on data manifold.
This aspect is particularly important for scientific inverse problems based on partial differential equations (PDEs), where staying on the data manifold is important for reliably solving the forward model $p({\bm{y}}|{\bm{x}})$. For instance, we observe that Gaussian approximations frequently violate the stability conditions of numerical PDE solvers, leading to meaningless estimates.

##### Derivative-free correction step.

Consider an ensemble of particles $\{{\bm{x}}_{i}^{(j)}\}_{j=1}^{J}$. Let $\bar{{\bm{x}}}_{i}$ denote their empirical mean and ${\bm{C}}_{xx}^{(i)}$ denote their empirical covariance matrix, at the $i$-th iteration.

$$ $\bar{{\bm{x}}}_{i}=\frac{1}{J}\sum_{j=1}^{J}{\bm{x}}_{i}^{(j)},{\bm{C}}_{xx}^{ (i)}=\frac{1}{J}\sum_{j=1}^{J}\left({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i} \right)\left({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i}\right)^{\top}.$ $$

Instead of the commonly used scalar weight $w_{i}$, we use a weighting matrix $w_{i}{\bm{C}}_{xx}^{(i)}$ in Eq. (8):

$$ $\displaystyle{\bm{x}}_{i+1}^{(j)}\approx\arg\min_{{\bm{x}}_{i+1}^{(j)}}$ $\displaystyle\frac{1}{2}\left({\bm{x}}_{i+1}-{\bm{x}}^{\prime}_{i}\right)^{ \top}\left(w_{i}{\bm{C}}_{xx}^{(i)}\right)^{-1}\left({\bm{x}}_{i+1}-{\bm{x}}^{ \prime}_{i}\right)$ (12) $\displaystyle-\nabla^{\top}\log\hat{p}\left({\bm{y}}|{\bm{x}}_{i}^{\prime(j)} \right)\left({\bm{x}}_{i+1}-{\bm{x}}_{i}^{\prime(j)}\right)$ (13) $\displaystyle={\bm{x}}_{i}^{\prime(j)}+$ $\displaystyle w_{i}{\bm{C}}_{xx}^{(i)}\nabla\log\hat{p}\left({\bm{y}}|{\bm{x}} _{i}^{\prime(j)}\right),$ (14) $$

which can be viewed as a gradient step preconditioned by ${\bm{C}}_{xx}^{(i)}$. Leveraging the idea of statistical linearization in the ensemble Kalman methods (Bergemann & Reich, ; Schillings & Stuart, ), we can perform the preconditioned gradient step without explicit derivative.

###### Assumption 1 .

$G\circ\phi$ has bounded first and second order derivatives.
Let $\psi$ denote $G\circ\phi$. There exist constants $M_{1},M_{2}$ such that for all ${\bm{u}},{\bm{u}}^{\prime},{\bm{v}},{\bm{v}}^{\prime}\in\mathbb{R}^{d}$,

$$ $\|\psi({\bm{u}})-\psi({\bm{u}}^{\prime})\|\leq M_{1}\|{\bm{u}}-{\bm{u}}^{ \prime}\|,{\bm{v}}^{\top}H_{\psi}({\bm{v}}^{\prime}){\bm{v}}\leq M_{2}\|{\bm{v }}\|^{2}.$ $$

where $H_{\psi}$ denotes the Hessian matrix of $\psi$.

###### Assumption 2 .

The distance between ensemble particles is bounded. There exists a constant $M_{3}$ such that $\|{\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i}\|<M_{3},j=1,\dots,J$.

###### Assumption 3 .

The observation empirical covariance matrix does not degenerate to zero unless the covariance matrix collapses to zero. In other words, $tr\left(C_{yy}^{(i)}\right)=0$ if and only if $C_{xx}^{(i)}=0$, and

$$ $C_{xx}^{(i)}\neq 0\rightarrow tr\left(C_{yy}^{(i)}\right)>M_{4},M_{4}>0,$ $$

where

$$ $C_{yy}^{(i)}=\frac{1}{J}\sum_{j=1}^{J}\left(\psi({\bm{x}}_{i}^{(j)})-\bar{\psi }_{i}\right)\left(\psi({\bm{x}}_{i}^{(j)})-\bar{\psi}_{i}\right)^{\top},\bar{ \psi}_{i}=\frac{1}{J}\sum_{j=1}^{J}\psi({\bm{x}}_{i}^{(j)}).$ (15) $$

###### Proposition 1 .

Under Assumption 1, 2 and 3, suppose the algorithm is implemented with $w_{i}=1/\left(tr\left(C_{yy}^{(i)}\right)\right)$, then we have the following approximation:

$$ ${\bm{C}}_{xx}^{(i)}\nabla\log\hat{p}\left({\bm{y}}|{\bm{x}}_{i}^{\prime(j)} \right)\approx\frac{1}{J}\sum_{k=1}^{J}\left\langle\psi\left({\bm{x}}^{\prime( k)}_{i}\right)-\bar{G},{\bm{y}}-\psi\left({\bm{x}}^{\prime(j)}_{i}\right) \right\rangle_{\Gamma}\left({\bm{x}}^{\prime(j)}_{i}-\bar{{\bm{x}}_{i}}\right),$ (16) $$

where

$$ $\bar{G}=\frac{1}{J}\sum_{j=1}^{J}G\left(\hat{{\bm{x}}}^{(j)}_{N}\right)=\frac{ 1}{J}\sum_{j=1}^{J}\psi({\bm{x}}_{i}^{\prime(j)}).$ $$

The detailed derivation can be found in Appendix A.2.
Therefore, we can update each particle with the ensemble difference to move toward the high-likelihood region without evaluating the derivative. A full description of our method is given in Algorithm 2. We name it Ensemble Kalman Diffusion Guidance as it is a guidance for diffusion models based on ensemble Kalman approximation. Implementation details are provided in Appendix A.4.

###### Assumption 1 .

###### Assumption 2 .

###### Assumption 3 .

###### Proposition 1 .

## 4 Experiments

We empirically study our method on the classic image restoration problems and two scientific inverse problems. We view the scientific inverse problems as the more interesting domains for evaluating our method, particularly the Navier-Stokes equation where it is impractical to accurately compute the gradient of the forward model.

##### Baselines

We focus on comparing against methods that only use black-box access to the forward model. Our first two baselines incorporate numerical estimation methods to substitute automatic differentiation, named Forward-GSG and Central-GSG given in Algorithm 3.
Specifically, we use forward Gaussian smoothed gradient (Forward GSG in Eq. 28) and central Gaussian smoothed gradient (Central GSG in Eq. 29) to approximate the gradient of the log-likelihood, and plug it into the standard gradient-based method, Diffusion Posterior Sampling (DPS)(Chung et al., ). more details can be found in Appendix A.3. The third baseline is Stochastic Control Guidance (SCG) proposed in Huang et al. (), which is designed to deal with non-differentiable guidance. The last baseline in Diffusion Policy Gradient (DPG) proposed by Tang et al. (). Additionally, for Navier-Stokes equation, we add the traditional method Ensemble Kalman Inversion (EKI) proposed in Iglesias et al. ().

**Table 1: Quantitative evaluation on FFHQ 256x256 dataset. We report average metrics for image quality and samples consistency on four tasks. Measurement noise level $\sigma=0.05$ is used if not otherwise stated.**
|  | Inpaint (box) | SR ($\times$4) | Deblur (Gauss) | Phase retrieval |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ |
| Gradient access |  |  |  |  |  |  |  |  |  |  |  |  |
| DPS^† | 21.77 | 0.767 | 0.213 | 24.90 | 0.710 | 0.265 | 25.46 | 0.708 | 0.212 | 14.14 | 0.401 | 0.486 |
| Black-box access |  |  |  |  |  |  |  |  |  |  |  |  |
| Forward-GSG | 17.82 | 0.562 | 0.302 | 18.08 | 0.469 | 0.384 | 24.43 | 0.704 | 0.206 | 7.88 | 0.070 | 0.838 |
| Central-GSG | 18.76 | 0.720 | 0.229 | 26.55 | 0.740 | 0.169 | 25.39 | 0.775 | 0.180 | 10.10 | 0.353 | 0.691 |
| DPG | 20.89 | 0.752 | 0.184 | 28.12 | 0.831 | 0.126 | 26.42 | 0.798 | 0.143 | 15.47 | 0.486 | 0.495 |
| SCG | 4.71 | 0.302 | 0.763 | 4.71 | 0.302 | 0.760 | 4.69 | 0.300 | 0.759 | 4.623 | 0.294 | 0.764 |
| EnKG(Ours) | 21.70 | 0.727 | 0.286 | 27.17 | 0.773 | 0.237 | 26.13 | 0.723 | 0.224 | 20.06 | 0.584 | 0.393 |

**Table 2: Comparison of different derivative-free algorithms on Navier-Stokes inverse problem. Numbers in parentheses represent the sample standard deviation.**
|  | $\sigma_{\mathrm{noise}}=0$ | $\sigma_{\mathrm{noise}}=1.0$ | $\sigma_{\mathrm{noise}}=2.0$ | Computation cost |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | Relative L2 | Relative L2 | Relative L2 | Total # FME | Total # DME | Seq # FME | Seq # DME |
| EKI | 0.577 (0.138) | 0.609 (0.119) | 0.673 (0.107) | 1024k | 0 | 0.50k | 0 |
| Forward-GSG | 1.687 (0.156) | 1.612 (0.173) | 1.454 (0.154) | 2049k | 1k | 1k | 1k |
| Central-GSG | 2.203* (0.314) | 2.117 (0.295) | 1.746 (0.191) | 2048k | 1k | 1k | 1k |
| DPG | 0.325 (0.188) | 0.408* (0.173) | 0.466 (0.171) | 4000k | 1k | 1k | 1k |
| SCG | 0.908 (0.600) | 0.928 (0.557) | 0.966 (0.546) | 512k | 384k | 0.75k | 1k |
| EnKG(Ours) | 0.120 (0.085) | 0.191 (0.057) | 0.294 (0.061) | 295k | 2695k | 0.14k | 1.3k |

Figure: Figure 2: Visualization of results on Navier-Stokes inverse problem with different levels of observation noise. Each observation is subsampled by a factor of 2, representing a sparse measurement. Note that the results of Central-GSG are here for demonstration purpose because neither Central-GSG nor Forward-GSG is able to produce reasonable results.
Refer to caption: extracted/5889642/figs/sample-demo.png

### 4.1 Image inverse problems

Tackling image inverse problems (e.g., deblurring) is common in the literature and serves as a reasonable reference domain for evaluation. We note that we consider a harder version of the problem where we do not use the gradient of the forward model. Moreover, most imaging problems use a linear forward model (except for phase retrieval), which is significantly simpler than the nonlinear forward models more often used in scientific domains. We do note that, perhaps surprisingly, for the nonlinear phase retrieval task, our approach outperforms the standard strong DPS baseline which actually uses the gradient of the forward model.

##### Problem setting

We evaluate our algorithm on the standard image inpainting, superresolution, deblurring (Gaussian), and phase retrieval problems. For inpainting, the forward model is a box mask with randomized location. For superresolution, we employ bicubic downsampling (either $\times 2$ or $\times 4$) and for Gaussian deblurring, a blurring kernel of size $61\times 61$ with standard deviation 3.0. Finally, phase retrieval takes the magnitude of the Fourier transform of the image as the observation. We use measurement noise $\sigma=0.05$ in all experiments except for superresolution on $64\times 64$ images, where we set $\sigma=0.01$.
The pre-trained diffusion model for FFHQ $64\times 64$ is taken unmodified from Karras et al. (). The model for FFHQ $256\times 256$ is taken from Chung et al. () and converted to the EDM framework using their Variance-Preserving (VP) preconditioning.

##### Evaluation metrics

We evaluate the sample quality of all methods using peak signal signal-to-noise-ratio (PSNR), structural similarity (SSIM) index (Wang et al., ), and learned perceptual image patch similarity (LPIPS) score (Zhang et al., ).

##### Results

We demonstrate the qualitative results for image inverse problems in Figure 7, and show the quantitative results in Table 5 and 1. EnKG achieves comparable or even better performance compared to DPS (Chung et al., ), which has gradient access to the forward function and the diffusion model. EnKG consistently outperforms all baseline methods with black-box access on two data settings and four different tasks.

### 4.2 Navier Stokes

The Navier-Stokes problem is representative of the key class of scientific inverse problems (Iglesias et al., ) that our approach aims to tackle. The gradient of the forward model is impractical to reliably compute via auto-differentiation, as it requires differentiating through a PDE solver. Having effective derivative-free methods would be highly desirable here.

##### Problem setting

We consider the 2-d Navier-Stokes equation for a viscous, incompressible fluid in vorticity form on a torus,
where ${\bm{u}}\in C\left([0,T];H^{r}_{\text{per}}((0,2\pi)^{2};\mathbb{R}^{2})\right)$ for any $r>0$ is the velocity field,
${\bm{w}}=\nabla\times{\bm{u}}$ is the vorticity, ${\bm{w}}_{0}\in L^{2}_{\text{per}}\left((0,2\pi)^{2};\mathbb{R}\right)$ is the initial vorticity, $\nu\in\mathbb{R}_{+}$ is the viscosity coefficient, and $f\in L^{2}_{\text{per}}\left((0,2\pi)^{2};\mathbb{R}\right)$ is the forcing function.
The solution operator ${\mathcal{G}}$ is defined as the operator mapping the vorticity from the initial vorticity to the vorticity at time $T$.
${\mathcal{G}}:{\bm{w}}_{0}\rightarrow{\bm{w}}_{T}$. Our experiments implement it as a pseudo-spectral solver (He & Sun, ).

$$ $\displaystyle\begin{split}\partial_{t}{\bm{w}}({\bm{x}},t)+{\bm{u}}({\bm{x}},t )\cdot\nabla{\bm{w}}({\bm{x}},t)&=\nu\Delta{\bm{w}}({\bm{x}},t)+f({\bm{x}}), \qquad{\bm{x}}\in(0,2\pi)^{2},t\in(0,T]\\ \nabla\cdot{\bm{u}}({\bm{x}},t)&=0,\qquad\qquad\qquad\qquad\quad\,{\bm{x}}\in( 0,2\pi)^{2},t\in[0,T]\\ {\bm{w}}({\bm{x}},0)&={\bm{w}}_{0}({\bm{x}}),\qquad\qquad\qquad\quad{\bm{x}} \in(0,2\pi)^{2}\end{split}$ (17) $$

The goal is to recover the initial vorticity field from a noisy sparse observation of the vorticity field at time $T=1$. The Navier-Stokes equation  17 does not admit a closed form solution and thus there is no closed form derivative available for the solution operator.
Moreover, obtaining an accurate numerical derivative via automatic differentiation through the numerical solver is challenging due to the extensive computation graph expanded after thousands of discrete time steps.

We first solve the equation up to time $T=5$ using initial conditions from a Gaussian random field, which is highly complicated due to the nonlinearity of Navier-Stokes equation. We sample 20,000 vorticity fields from this empirical measure to train our diffusion model. Then, we independently sample 10 random vorticity fields as the test set.

##### Evaluation metrics

We report the relative $L^{2}$ error to evaluate the accuracy of the algorithm, which is given by $\frac{\left\|\hat{{\bm{w}}}_{0}-{\bm{w}}^{*}_{0}\right\|_{L^{2}}}{\left\|{\bm{
w}}^{*}_{0}\right\|_{L^{2}}}$ where $\hat{{\bm{w}}}_{0}$ is the predicted vorticity and ${\bm{w}}^{*}_{0}$ is the ground truth. To comprehensively capture the computational characteristics of different algorithms, we use the following metrics:
1) Total # Fwd: total number of forward model evaluations; 2) Seq. # Fwd: number of sequential forward model evaluations; 3) Total # DM: total number of diffusion model evaluations; 4) Seq. # DM: number of sequential diffusion model evaluations; 5) Total # DM grad: total number of diffusion model gradient evaluations; 6) Seq. # DM grad: number of sequential diffusion model gradient evaluations.

##### Results

In Table 2, we show the average relative $L^{2}$ error of the recovered ground truth at different noise levels of the observations. Our approach, by far, outperforms methods which use other forms of approximate derivative information. Indeed, we show in Figure 2 that our method give solutions
which qualitatively preserve important features of the flow while other methods completely fail to give a reasonable answer. Additionally, we compare the results in the radial power spectrum as shown in Figure 3(a). The traditional baseline EKI produces an over-smoothed solution, while DPG,SCG, and GSG tend to generate noisy solutions. In Figure 3 (b), we visualize the main source of computation costs, showing that the forward model evaluation of Navier-Stokes equation is significantly slower than other parts as the numerical solver has to take thousands of iterations. Figure 3 (c) compares different algorithms with our proposed efficiency metrics, showing that EnKG requires the least number of forward model evaluations and is highly parallelizable. Therefore, for problems where the forward model evaluation is slow, EnKG stands out to be the most efficient algorithm compared to the baseline methods.

Figure: Figure 3: Panel (a): radial power spectrum of the solutions from different algorithms. Panel (b): runtime of single evaluation of the forward model, diffusion model, and diffusion model gradient (tested on a single A100). Panel (c): comparison of computational characteristics of different algorithms on Navier-Stokes problem.
Refer to caption: extracted/5889642/figs/spec-complexity.png

### 4.3 Black-hole imaging inverse problem

The black-hole imaging problem is interesting due to its highly non-linear and ill-posed forward model (i.e., the sparse observations captured by telescopes on Earth). For evaluation purposes, we assume only black-box access to the forward model.

##### Problem setting

The black hole interferometric imaging system aims to reconstruct image of black holes using a set of telescopes distributed on the Earth. Each pair of telescopes produces a measurement $V_{a,b}^{t}$ called visibility, where $(a,b)$ is a pair of telescopes and $t$ is the measuring time. To mitigate the effect of measurement noise caused by atmosphere turbulence and thermal noise, multiple visibilities can be grouped together to cancel out noise (Chael et al., ). This produces noise-invariant measurements, termed closure phases ${\bm{y}}^{\mathrm{cph}}_{t,(a,b,c)}$ and log closure amplitudes ${\bm{y}}^{\mathrm{camp}}_{t,(a,b,c,d)}$. We specify the likelihood of this measuring function similar as Sun & Bouman (), i.e.,

$$ $\ell({\bm{y}}|{\bm{x}})=\sum_{t}\frac{\|\mathcal{A}^{\mathrm{cph}}_{t}({\bm{x} })-{\bm{y}}^{\mathrm{cph}}_{t}\|_{2}^{2}}{2\beta^{2}_{\mathrm{cph}}}+\sum_{t} \frac{\|\mathcal{A}_{t}^{\mathrm{camp}}({\bm{x}})-{\bm{y}}_{t}^{\mathrm{camp}} \|_{2}^{2}}{2\beta^{2}_{\mathrm{camp}}}+\rho\frac{\|\sum{\bm{x}}_{i}-{\bm{y}}^ {\mathrm{flux}}\|_{2}^{2}}{2},$ (18) $$

where $\mathcal{A}_{t}^{\mathrm{cph}}$ and $\mathcal{A}_{t}^{\mathrm{camp}}$ measures the closure phase and log closure amplitude of black hole images ${\bm{x}}$. $\beta_{\mathrm{cph}}$ and $\beta_{\mathrm{camp}}$ are some known parameters given by the telescope system. The first two terms are the sum of chi-square values for closure phases and log closure amplitudes, and the last term constrain the total flux of the black-hole image.
We trained a diffusion model on GRMHD dataset (Wong et al., ) with resolution 64$\times$ 64 to generate black hole images from telescope measurements.

##### Evaluation metrics

We report the chi-square errors of closure phases $\chi_{\mathrm{cph}}^{2}$ and closure amplitudes $\chi_{\mathrm{camp}}^{2}$ to measure how the generated samples fit the given measurement. We calculate the peak signal-to-noise ratio (PSNR) between reconstructed images and the ground truth. Moreover, since the black-hole imaging system provides only information for low spatial frequencies, following conventional evaluation methodology (EHTC, ), we blur images with a circular Gaussian filter and compute their PSNR on the intrinsic resolution of the imaging system.

##### Results

Figure 4 shows the reconstructed images of the black-hole using our method and the baseline methods with black box access.
EnKG is able to generate black hole images with visual features consistent with the ground truth. Moreover, we report quantitative results of each method in Table 3. EnKG achieves relatively low measurement error and the best (blurred) PSNR compared with baseline methods. SCG achieves slightly better data fitting metrics measured by chi-square value, but produces images that are much noisier than those of EnKG.

Figure: Figure 4: Visualization of generated samples on the black-hole imaging inverse problem. The first row shows the results on the original resolution, while the second row shows the blurred images in the intrinsic resolution of the imaging system.
Refer to caption: extracted/5889642/figs/blackhole-visual.png

**Table 3: Quantitative evaluation of the reconstructed black-hole images.**
|  | PSNR $\uparrow$ | Blurred PSNR $\uparrow$ | $\chi^{2}_{\mathrm{cph}}$ $\downarrow$ | $\chi^{2}_{\mathrm{camp}}$ $\downarrow$ |
| --- | --- | --- | --- | --- |
| Central-GSG | 24.700 | 30.011 | 4.616 | 79.669 |
| SCG | 20.201 | 20.976 | 1.103 | 1.134 |
| DPG | 13.222 | 14.281 | 5.116 | 15.679 |
| EnKG (Ours) | 29.093 | 32.803 | 1.426 | 1.270 |

## 5 Limitations and discussion

The first limitation of the proposed EnKG is that as a optimization-based method, it cannot capture the exact posterior distribution and thus cannot provide reliable uncertainty quantification, which might be important in some applications. A potential future improvement can involve adapting ideas from the sequential Monte Carlo approach to EnKG to enable provable posterior sampling. Secondly, while the per-sample time cost of EnKG is smaller than the standard gradient-based approach, the total runtime is much longer because EnKG maintains a whole ensemble of interacting particles at each iteration. However, as we shown in Figure 6, even a small number of particles can achieve 20-30% relative $L^{2}$ error. A potential improvement could be introducing a scheme that adaptively controls the number of particles according to the optimization landscape to improve efficiency.

## 6 Conclusion

In this work, we propose EnKG, a fully derivative-free approach to solve general inverse problems that only permit black-box access. EnKG can accommodate different forward models without any re-training while maintaining the expressive ability of diffusion models to capture complex distribution. The experiments on various inverse problems arising from imaging and partial differential equations demonstrate the robustness and effectiveness of our methodology.

#### Acknowledgments

We thank Ben Prather, Abhishek Joshi, Vedant Dhruv, C.K. Chan, and Charles Gammie for
the synthetic images used here, generated under NSF grant AST 20-34306.
This research used resources of the Oak Ridge Leadership Computing Facility at the Oak
Ridge National Laboratory, which is supported by the Office of Science of the U.S.
Department of Energy under Contract No. DE-AC05-00OR22725. This research used resources of
the Argonne Leadership Computing Facility, which is a DOE Office of Science User Facility
supported under Contract DE-AC02-06CH11357. This research was done using services provided
by the OSG Consortium, which is supported by the National Science Foundation awards
#2030508 and #1836650. This research is part of the Delta research computing project,
which is supported by the National Science Foundation (award OCI 2005572), and the State
of Illinois. Delta is a joint effort of the University of Illinois at Urbana-Champaign and
its National Center for Supercomputing Applications.

## Appendix A Appendix / supplemental material

### A.1 Notation

**Table 4: Table of notations.**
| Notation | Description |
| --- | --- |
| $G$ | the forward model of the inverse problem |
| $\phi$ | Probability ODE solver |
| $\psi$ | Composition of $G$ and $\phi$ |
| ${\bm{D}}f$ | Jacobian matrix of function $f$ |
| $L^{r}_{\text{per}}$ | Lebesgue space of periodic $r$-integrable functions |
| $H^{r}_{\text{per}}$ | Sobolev space of $r$-times weakly differentiable periodic functions |
| $\hat{\nabla}f$ | approximate gradient of $f$ |
| $\mu$ | Gaussian smoothing factor |
| $Q$ | number of gradient estimation queries |
| $w_{i}$ | log-likelihood gradient scale at step $i$ |
| $N$ | number of sampling steps |
| $E_{\mu,Q}(f({\bm{x}}))$ | gradient estimator of $f({\bm{x}})$ using smoothing factor $\mu$ and $Q$ queries |

### A.2 Proofs

###### Lemma 1 .

Under Assumption 1, 2 and 3, suppose the algorithm is implemented with $w_{i}=1/\left(tr\left(C_{yy}^{(i)}\right)\right)$, then

$$ $tr\left(C_{xx}^{(i+1)}\right)\leq tr\left((C_{xx}^{(i)}\right)-\frac{1}{M_{3}} tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right),$ (19) $$

where

$$ $C_{xy}^{(i)}=\frac{1}{J}\sum_{j=1}^{J}\left({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_ {i}\right)\left(\psi({\bm{x}}_{i}^{(j)})-\bar{\psi}_{i}\right)^{\top}.$ $$

Note that $tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right)$ is zero if and only if all ensemble members collapse to a single point, i.e., $tr\left(C_{xx}^{(i)}\right)=0$. Therefore, $tr\left(C_{xx}^{(i)}\right)$ monotonically decreases to zero in the limit as $i$ goes to infinity.

###### Proof.

We first rewrite the ensemble update at iteration $i$ as follows

$$ ${\bm{x}}_{i+1}^{(j)}={\bm{x}}_{i}^{(j)}+w_{i}C_{xy}^{(i)}\left(y-\psi({\bm{x}} _{i}^{(j)})\right),$ (20) $$

where $j\in\{1,\dots,J\}$. The covariance matrix at the next iteration is given by

$$ $C_{xx}^{(i+1)}=\frac{1}{J}\sum_{j=1}^{J}({\bm{x}}_{i+1}^{(j)}-\bar{{\bm{x}}}_{ i+1})({\bm{x}}_{i+1}^{(j)}-\bar{{\bm{x}}}_{i+1})^{\top}.$ (21) $$

Plugging the update rule in Eq. (20) into Eq. (21), we have

$$ $\displaystyle C_{xx}^{(i+1)}$ $\displaystyle=\frac{1}{J}\sum_{j}\left[({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i}) +w_{i}C_{xy}^{(i)}\left(\bar{\psi}_{i}-\psi({\bm{x}}_{i}^{(j)})\right)\right] \left[({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i})+w_{i}C_{xy}^{(i)}(\bar{\psi}_{i} -\psi({\bm{x}}_{i}^{(j)}))\right]^{\top}$ $\displaystyle=C_{xx}^{(i)}-2w_{i}C_{xy}^{(i)}C_{xy}^{(i)\top}+w_{i}^{2}C_{xy}^ {(i)}C_{yy}^{(i)}C_{xy}^{(i)\top}.$ $$

By linearity of trace, we have

$$ $tr\left(C_{xx}^{(i+1)}\right)=tr\left(C_{xx}^{(i)}\right)-2w_{i}tr\left(C_{xy} ^{(i)}C_{xy}^{(i)\top}\right)+w_{i}^{2}tr\left(C_{xy}^{(i)}C_{yy}^{(i)}C_{xy}^ {(i)\top}\right).$ $$

By cyclic and submultiplicative properties, we have

$$ $w_{i}^{2}tr\left(C_{xy}^{(i)}C_{yy}^{(i)}C_{xy}^{(i)\top}\right)=w_{i}^{2}tr \left(C_{yy}^{(i)}C_{xy}^{(i)\top}C_{xy}^{(i)}\right)\leq w_{i}^{2}tr\left(C_{ yy}^{(i)}\right)tr\left(C_{xy}^{(i)\top}C_{xy}^{(i)}\right).$ $$

Since $w_{i}=1/\left(tr\left(C_{yy}^{(i)}\right)\right)$, we have

$$ $\displaystyle tr\left(C_{xx}^{(i+1)}\right)$ $\displaystyle\leq tr\left(C_{xx}^{(i)}\right)-\frac{2}{tr\left(C_{yy}^{(i)} \right)}tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right)+\frac{1}{tr\left(C_{yy}^{( i)}\right)}tr\left(C_{xy}^{(i)\top}C_{xy}^{(i)}\right)$ $\displaystyle=tr\left(C_{xx}^{(i)}\right)-\frac{1}{tr\left(C_{yy}^{(i)}\right) }tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right).$ $$

By Assumption 1 and 2, we know there exists a constant $M_{4}>0$ such that $tr\left(C_{yy}^{(i)}\right)\leq M_{4}$. Therefore,

$$ $tr\left(C_{xx}^{(i+1)}\right)\leq tr\left(C_{xx}^{(i)}\right)-\frac{1}{tr\left (C_{yy}^{(i)}\right)}tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right)\leq tr\left(C _{xx}^{(i)}\right)-\frac{1}{M_{4}}tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right).$ $$

By Assumption 3, $tr\left(C_{xy}^{(i)}C_{xy}^{(i)\top}\right)$ is zero if and only if all the ensemble members collapse to a single point, which is $tr\left(C_{xx}^{(i)}\right)=0$. Therefore, $tr\left(C_{xx}^{(i)}\right)$ monotonically decreases to zero. Additionally, we empirically check how quickly the average distance decays as we iterate in our experiments as shown in Figure 5.
∎

Figure: Figure 5: Distance of ensemble members quickly decays over update steps. Empirical verification of Lemma 1.
Refer to caption: extracted/5889642/figs/avg_dist.png

###### Proposition 1 .

Under Assumption 1, 2 and 3, suppose the algorithm is implemented with $w_{i}<1/\left(tr\left(C_{yy}^{(i)}\right)\right)$, then we have the following approximation:

$$ $C_{xx}^{(i)}\nabla\log\hat{p}\left({\bm{y}}|{\bm{x}}_{i}^{\prime(j)}\right) \approx\frac{1}{J}\sum_{k=1}^{J}\left\langle G(\hat{{\bm{x}}}^{\prime(k)}_{N}) -\bar{G},{\bm{y}}-G(\hat{{\bm{x}}}^{(j)}_{N})\right\rangle_{\Gamma}\left({\bm{ x}}^{(j)}_{i}-\bar{{\bm{x}}_{i}}\right),$ (22) $$

where

$$ $\bar{G}=\frac{1}{J}\sum_{j=1}^{J}G(\hat{{\bm{x}}}^{(j)}_{N})$ $$

###### Proof.

Note that we can always normalize the problem so that $\Gamma$ is identity. Therefore, without loss of generality and for the ease of notation, we assume $\Gamma=\mathbf{I}$ throughout the whole proof.
Given the inverse problem setting in Eq. 1 where the observation noise is Gaussian, we can rewrite the preconditioned gradient w.r.t ${\bm{x}}_{i}^{\prime(j)}$ as

$$ $\displaystyle~{}C_{xx}^{(i)}\nabla\log\hat{p}\left({\bm{y}}|{\bm{x}}_{i}^{ \prime(j)}\right)$ (23) $\displaystyle=-\frac{1}{J}\sum_{k=1}^{J}\left({\bm{x}}_{i}^{\prime(k)}-\bar{{ \bm{x}}}_{i}\right)\left({\bm{x}}_{i}^{\prime(k)}-\bar{{\bm{x}}}_{i}\right)^{ \top}\nabla\frac{1}{2}\left\|\psi\left({\bm{x}}_{i}^{\prime(j)}\right)-{\bm{y} }\right\|^{2}$ (24) $\displaystyle=-\frac{1}{J}\sum_{k=1}^{J}\left({\bm{x}}_{i}^{\prime(k)}-\bar{{ \bm{x}}}_{i}\right)\left({\bm{x}}_{i}^{\prime(k)}-\bar{{\bm{x}}}_{i}\right)^{ \top}{\bm{D}}^{\top}\psi\left({\bm{x}}_{i}^{\prime(j)}\right)\left(\psi\left({ \bm{x}}_{i}^{\prime(j)}\right)-{\bm{y}}\right)$ (25) $\displaystyle=-\frac{1}{J}\sum_{k=1}^{J}\left({\bm{x}}_{i}^{\prime(k)}-\bar{{ \bm{x}}}_{i}\right)\left({\bm{D}}\psi\left({\bm{x}}_{i}^{\prime(j)}\right){\bm {x}}_{i}^{\prime(k)}-{\bm{D}}\psi\left({\bm{x}}_{i}^{\prime(j)}\right)\bar{{ \bm{x}}}_{i}\right)^{\top}\left(\psi\left({\bm{x}}_{i}^{\prime(j)}\right)-{\bm {y}}\right)$ (26) $\displaystyle=-\frac{1}{J^{2}}\sum_{k=1}^{J}\sum_{l=1}^{J}\left({\bm{x}}_{i}^{ \prime(j)}-\bar{{\bm{x}}}_{i}\right)\left({\bm{D}}\psi\left({\bm{x}}_{i}^{ \prime(j)}\right)\left({\bm{x}}_{i}^{\prime(k)}-{\bm{x}}_{i}^{(l)}\right) \right)^{\top}\left(\psi\left({\bm{x}}_{i}^{\prime(j)}\right)-{\bm{y}}\right).$ (27) $$

By definition, we have

$$ $\displaystyle tr\left(C_{xx}^{(i)}\right)$ $\displaystyle=tr\left(\frac{1}{J}\sum_{j=1}^{J}\left({\bm{x}}_{i}^{(j)}-\bar{{ \bm{x}}}_{i}\right)\left({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i}\right)^{\top}\right)$ $\displaystyle=\frac{1}{J}\sum_{j=1}^{J}tr\left(\left({\bm{x}}_{i}^{(j)}-\bar{{ \bm{x}}}_{i}\right)^{\top}\left({\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i}\right)\right)$ $\displaystyle=\frac{1}{J}\sum_{j=1}^{J}\|{\bm{x}}_{i}^{(j)}-\bar{{\bm{x}}}_{i} \|_{2}^{2},$ $$

which represents the average distance between ensemble members. By Lemma 1, we know that $tr\left(C_{xx}^{(i)}\right)$ monotonically decreases to zero in the limit. Therefore, the ensemble members will get sufficiently close as we iterate. Therefore, we can apply first-order Taylor approximation to $h$ at ${\bm{x}}_{i}^{\prime(j)}$ under Assumption 1 and obtain

$$ $\displaystyle\psi\left({\bm{x}}_{i}^{\prime(k)}\right)$ $\displaystyle=\psi\left({\bm{x}}_{i}^{\prime(j)}+{\bm{x}}_{i}^{\prime(k)}-{\bm {x}}_{i}^{\prime(j)}\right)$ $\displaystyle=\psi\left({\bm{x}}_{i}^{\prime(j)}\right)+{\bm{D}}\psi\left({\bm {x}}_{i}^{\prime(j)}\right)\left({\bm{x}}_{i}^{\prime(k)}-{\bm{x}}_{i}^{\prime (j)}\right)+o\left(\|{\bm{x}}_{i}^{\prime(k)}-{\bm{x}}_{i}^{\prime(j)}\|_{2}^{ 2}\right),$ $$

where $k\in\{1,\dots,J\}$.
Therefore for any $k,l\in\{1,\dots,J\}$, we have

$$ $\psi\left({\bm{x}}_{i}^{\prime(k)}\right)-\psi\left({\bm{x}}_{i}^{\prime(l)} \right)={\bm{D}}\psi\left({\bm{x}}_{i}^{\prime(j)}\right)\left({\bm{x}}_{i}^{ \prime(k)}-{\bm{x}}_{i}^{\prime(l)}\right)$ $$

We then plug it into Eq. 27

$$ $\displaystyle~{}C_{xx}^{(i)}\nabla\log\hat{p}\left({\bm{y}}|{\bm{x}}_{i}^{ \prime(j)}\right)$ $\displaystyle\approx-\frac{1}{J^{2}}\sum_{k=1}^{J}\sum_{l=1}^{J}\left({\bm{x}} _{i}^{(j)}-\bar{{\bm{x}}}_{i}\right)\left(\psi\left({\bm{x}}_{i}^{\prime(k)} \right)-\psi\left({\bm{x}}_{i}^{(l)}\right)\right)^{\top}\left(\psi\left({\bm{ x}}_{i}^{\prime(j)}\right)-{\bm{y}}\right)$ $\displaystyle=-\frac{1}{J}\sum_{k=1}^{J}\left({\bm{x}}_{i}^{\prime(j)}-\bar{{ \bm{x}}}_{i}\right)\left(\psi\left({\bm{x}}_{i}^{\prime(k)}\right)-\bar{\psi}_ {i}\right)^{\top}\left(\psi\left({\bm{x}}_{i}^{\prime(j)}\right)-{\bm{y}}\right)$ $\displaystyle=-\frac{1}{J}\sum_{k=1}^{J}\left\langle\psi\left({\bm{x}}_{i}^{ \prime(k)}\right)-\bar{\psi}_{i},\psi\left({\bm{x}}_{i}^{\prime(j)}\right)-{ \bm{y}}\right\rangle\left({\bm{x}}_{i}^{\prime(j)}-\bar{{\bm{x}}}_{i}\right)$ $\displaystyle=\frac{1}{J}\sum_{k=1}^{J}\left\langle G(\hat{{\bm{x}}}^{\prime(k )}_{N})-\bar{G},{\bm{y}}-G(\hat{{\bm{x}}}^{(j)}_{N})\right\rangle\left({\bm{x} }^{(j)}_{i}-\bar{{\bm{x}}_{i}}\right),$ $$

concluding the proof.
∎

###### Lemma 1 .

###### Proof.

###### Proposition 1 .

###### Proof.

### A.3 Zero-order gradient estimation baseline

We use the forward Gaussian smoothing and central Gaussian smoothing gradient estimation methods to establish a baseline to compare against. These methods approximate the gradient of a function using only function evaluations and can be expressed in the following (forward-GSG) form :

$$ $\hat{\nabla}f({\bm{x}})=\sum_{i}^{Q}\frac{f({\bm{x}}+\mu{\bm{u}}_{i})-f({\bm{x }})}{\mu}\tilde{{\bm{u}}}_{i}$ (28) $$

And Central-GSG:

$$ $\hat{\nabla}f({\bm{x}})=\sum_{i}^{Q}\frac{f({\bm{x}}+\mu{\bm{u}}_{i})-f({\bm{x }}-\mu{\bm{u}}_{i})}{2\mu}\tilde{{\bm{u}}}_{i}$ (29) $$

For Gaussian smoothing, ${\bm{u}}_{i}$ follows the standard normal distribution and $\tilde{{\bm{u}}}_{i}=\frac{1}{Q}{\bm{u}}_{i}$. The smoothing factor $\mu$ and number of queries $Q$ are both tunable hyperparameters.

Posterior sampling requires computation of the scores $\nabla_{{\bm{x}}_{t}}p({\bm{x}}_{t})$ and $\nabla_{{\bm{x}}_{t}}p({\bm{y}}\mid{\bm{x}}_{t})$; the former is learned by the pre-trained diffusion model, and the latter can be estimated by various approximation methods. In our baseline derivative-free inverse problem solver, we substitute the explicit automatic differentiation used in algorithms such as DPS with (28) and (29). We estimate this gradient by leveraging the fact that a probability flow ODE deterministically maps every ${\bm{x}}_{t}$ to ${\bm{x}}_{0}$; $\hat{\nabla}_{\hat{{\bm{x}}}_{0}}p({\bm{y}}\mid\hat{{\bm{x}}}_{0})$ is approximated with Gaussian smoothing, and a vector-Jacobian product (VJP) is used to then calculate $\hat{\nabla}_{{\bm{x}}_{t}}p({\bm{y}}\mid{\bm{x}}_{t})$. Our gradient estimate is defined as follows:

$$ $\hat{\nabla}_{{\bm{x}}_{t}}\log{p({\bm{y}}\mid{\bm{x}}_{t})}=\hat{\nabla}_{{ \bm{x}}_{t}}\log p({\bm{y}}\mid\hat{{\bm{x}}}_{0})={\bm{D}}_{{\bm{x}}_{t}}^{ \top}\hat{{\bm{x}}}_{0}\hat{\nabla}_{\hat{{\bm{x}}}_{0}}\log{p({\bm{y}}\mid \hat{{\bm{x}}}_{0})}$ (30) $$

${\bm{D}}_{{\bm{x}}_{t}}^{\top}$ is the transpose of the Jacobian matrix; (30) can be efficiently computed using automatic differentiation. Note that although automatic differentiation is used, differentiation through the forward model does not occur. Thus, this method is still applicable to non-differentiable inverse problems. Furthermore, we choose to perturb $\hat{{\bm{x}}}_{0}$ and use a VJP rather than directly perturb ${\bm{x}}_{t}$ so that we can avoid repeated forward passes through the pre-trained network, which is very expensive. Pseudocode for these algorithms is provided in Algorithm LABEL:alg:dps-zo.

Figure: Algorithm 3 Central/Forward-GSG baseline with $\sigma(t)=t$ and $s(t)=1$

### A.4 EnKG Implementation Details

There are mainly two design choices in our algorithm 2 to be made. The first is the step size $w_{i}$ which controls the extent to which the correction step moves towards the MAP estimator. In the ensemble Kalman literature (Kovachki & Stuart, ), the following adaptive step size is widely used and we adopt it for our experiments as well.

$$ $w_{i}^{-1}=\frac{1}{J^{2}}\sqrt{\sum_{k=1}^{J}\left\|G(\hat{{\bm{x}}}^{\prime( k)}_{N})-\bar{G}\right\|^{2}\left\|{\bm{y}}-G(\hat{{\bm{x}}}^{(j)}_{N})\right \|^{2}}$ (31) $$

Secondly, we find it useful to perform two correction steps in Eq. (6) when solving highly nonlinear and high-dimensional problems such as Navier Stokes. Therefore, we perform two correction steps at each iteration when running experiments on Navier Stokes.

### A.5 Baseline Details

### A.6 Additional results

Figure: Figure 6: Ablation study on the number of particles for Navier-Stokes. The shaded region represents best and worst particle.
Refer to caption: x1.png

**Table 5: Qualitative evaluation on FFHQ 64x64 dataset. We report average metrics for image quality and samples consistency on four tasks. Measurement noise level $\sigma=0.05$ is used if not otherwise stated.**
|  | Inpaint (box) | SR ($\times$2, $\sigma=0.01$) | Deblur (Gauss) | Phase retrieval |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ |
| Gradient access |  |  |  |  |  |  |  |  |  |  |  |  |
| DPS^† | 22.58 | 0.767 | 0.087 | 29.19 | 0.924 | 0.032 | 21.50 | 0.663 | 0.157 | 12.19 | 0.245 | 0.592 |
| Black-box access |  |  |  |  |  |  |  |  |  |  |  |  |
| Forward-GSG | 19.62 | 0.612 | 0.189 | 25.25 | 0.836 | 0.093 | 20.27 | 0.606 | 0.170 | 10.307 | 0.170 | 0.493 |
| Central-GSG | 21.37 | 0.764 | 0.095 | 27.41 | 0.916 | 0.030 | 20.88 | 0.729 | 0.123 | 11.36 | 0.283 | 0.619 |
| DPG | 21.92 | 0.799 | 0.088 | 26.86 | 0.917 | 0.027 | 20.00 | 0.734 | 0.114 | 15.56 | 0.438 | 0.446 |
| SCG | 20.27 | 0.734 | 0.098 | 27.02 | 0.910 | 0.036 | 20.73 | 0.754 | 0.100 | 10.59 | 0.233 | 0.617 |
| EnKG(Ours) | 23.53 | 0.822 | 0.067 | 29.52 | 0.930 | 0.036 | 22.02 | 0.698 | 0.136 | 26.14 | 0.840 | 0.122 |

Figure: Figure 7: Qualitative results on FFHQ 256.
Refer to caption: extracted/5889642/figs/ffhq256-demo.png

We include more qualitative results for inverse problems on FFHQ 256x256 dataset in Figure 7.

**Table 6: Hyperparameter choices for DPS with zero-order gradient estimation ($64\times 64$).**
|  | Inpaint (box) | SR ($\times$2, $\sigma=0.01$) | Deblur (Gauss) | Phase retrieval |
| --- | --- | --- | --- | --- |
| Forward GSG |  |  |  |  |
| $\mu$ | 0.001 | 0.001 | 0.001 | 0.001 |
| $Q$ | 10000 | 10000 | 10000 | 10000 |
| $w_{i}$ | 1.0 | 1.0 | 1.0 | 1.0 |
| $N$ | 1000 | 1000 | 1000 | 1000 |
| Central GSG |  |  |  |  |
| $\mu$ | 0.001 | 0.001 | 0.001 | 0.001 |
| $Q$ | 10000 | 10000 | 10000 | 10000 |
| $w_{i}$ | 1.0 | 1.0 | 1.0 | 1.0 |
| $N$ | 1000 | 1000 | 1000 | 1000 |

**Table 7: Hyperparameter choices for baselines based on zero-order gradient estimation ($256\times 256$).**
|  | Inpaint (box) | SR ($\times$4, $\sigma=0.05$) | Deblur (Gauss) | Phase retrieval |
| --- | --- | --- | --- | --- |
| Forward-GSG |  |  |  |  |
| $\mu$ | 0.01 | 0.01 | 0.01 | 0.01 |
| $Q$ | 10000 | 10000 | 20000 | 20000 |
| $w_{i}$ | 1.0 | 1.0 | 3.0 | 0.7 |
| $N$ | 1000 | 1000 | 1000 | 1000 |
| Central-GSG |  |  |  |  |
| $\mu$ | 0.01 | 0.01 | 0.01 | 0.01 |
| $Q$ | 10000 | 10000 | 20000 | 20000 |
| $w_{i}$ | 1.0 | 1.0 | 3.0 | 0.7 |
| $N$ | 1000 | 1000 | 1000 | 1000 |

### A.7 Details of black hole imaging

The measurement of black hole imaging is defined as (Sun & Bouman, )

$$ $\displaystyle{\bm{y}}_{t,(a,b,c)}^{\text{cph}}=\angle(V_{a,b}^{t}V_{b,c}^{t}V_ {a,c}^{t}):=\mathcal{A}_{t,(a,b,c)}^{\text{cph}}(\mathbf{x})$ (32) $\displaystyle{\bm{y}}_{t,(a,b,c,d)}^{\text{camp}}=\log\left(\frac{|V_{a,b}^{t} ||V_{c,d}^{t}|}{|V_{a,c}^{t}||V_{b,d}^{t}|}\right):=\mathcal{A}_{t,(a,b,c,d)}^ {\text{camp}}({\bm{x}})$ (33) $$

where $V_{a,b}$ is the visibility defined by

$$ $V_{a,b}^{t}({\bm{x}})=g_{a}^{t}g_{b}^{t}\exp(-i(\phi_{a}^{t}-\phi_{b}^{t})) \cdot\tilde{\bm{I}}^{t}_{a,b}({\bm{x}})+\eta_{a,b}.$ (34) $$

$g_{a},g_{b}$ are telescope-based gain errors, $\phi_{a}^{t},\phi_{b}^{t}$ are phase errors, and $\eta_{a,b}$ is baseline-based Gaussian noise. The measurements consist of $(M-1)(M-2)/2$ closure phases ${\bm{y}}^{\text{cph}}$ and $M(M-3)/2$ log closure amplitudes ${\bm{y}}^{\text{camp}}$ for an array of $M$ telescopes. Our experiments use $M=9$ telescopes from Event Horizon Telescope.

Figure: Figure 8: Vorticity field predicted by EnKG with different number of particles. From left to right, the result gets better as we increase the number of particles.
Refer to caption: extracted/5889642/figs/ablation-noise0.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - **Q1:** Use of a `pretrained diffusion model`  is mentionned. I can see this type of models for images, but for complex data with varying domain of definition, like in many scientific applications, the diffusion model needs to be retrained for each type of data no? IT is not completely clear how reusable these models are outside computational imaging.

- **Q2:** The prediction-correction scheme strongly relate to the usual Plug-and-Play methods classicaly used in inverse problem resolution. The proposed scheme is related to the Forward-Backward Splitting (FBS) method, which is typically used in the DiffPIR [B]. From what I understand, equation (9) in this paper corresponds to equation (13) in [A]. Could you elaborate on the main difference of between the proposed method and the DiffPIR framework? (except the derivation free aspect). Other missing references are [B], a survey on using diffusion models for inverse problems, and [C], which also consider diffusion prior for inverse problems, with gradient based method but applied to the Black-hole imaging problem. Adding comparison with their method would be very useful.

- **Q3:** To make is possible to evaluate how good the results are compared to methods that are able to access the gradient of the forward operator, it is necessary to add a few methods that have access to the forward operator's gradients. Indeed, even though ODE solvers are not always natively differentiable, there are more and more works that consider making them differentiable, for instance using `jax` for Navier stokes using pseudo-spectral solver  [here](https://github.com/google/jax-cfd). The interesting question is: should we spend some time making them differentiable or do we not gain much by doing so. Therefor, quantifying how much is lost on simple cases as the ones presented here is necessary to make the case of derivative free methods. In particular, adding the results of the DPS method and the DiffPIR method would be very useful at least for the imaging task. Note that these models are both implemeted in the [`deepinv`](https://deepinv.github.io) library (see [here](https://deepinv.github.io/deepinv/deepinv.sampling.html)). Also, adding the DPS base line and PnP-DM form [C] for black hole imaging would better illustrate how much we loose by note considering the gradient of this differentiable operator. Adding them for Navier-Stokes would also be very interesting but probably more challenging.

- **Q4**: The value of $J$ and $Q$ in the experiments are not reported. Could the author provide them? From equation (16), I understand that we need to compute the forward operator $J$ or $Q$ times at each iteration. From the NAvier stokes experiment, assuming that the procedure is run for 1k steps, I guess $J = 295$ and $Q=2048$? How were these value chosen? What are their impact on the results? Also, would having access to the gradient mean going roughly 100 times faster than EnKG? (The computational cost of computing the gradient through autodiff is approximately x2/3 times the cost of evaluating the forward operator). Note that the metrics chosen (# of evaluation of Fwd/DM, Seq) are not clear. Better definition of what they represent mean would be useful. Adding the total runtime of the method would help a lot to assess the computational cost.

- **Q5:** In the black-hole experiment, how many simulation are used to train the diffusion model?

### Minor comments, nitpicks and typos

- Missing ref: 
- l.066: "One more challenging" -> "On"
- l.069: "More computationally efficient" -> than what?
- l.078: "often Gaussian" -> They are almost never Gaussian, but they are modeled as such and this gives reasonable results.
- l.141: "and Nelder-Mead simplex methods" -> Extra `,`.
- l.198: The drop of the subscript $x$ for the gradient is confusing
- Eq (7) -> The notation $\Delta x_i$ is not defined. As it is not used anywhere else, I would recomment using $\|x_{i+1} - x_i\|$ instead.
- Eq (12) -> $x_i'$ should have a super script $(j)$. It is also not completely consistent for $x_{i+1}$.
- l.302: "instead"
- l.316: `our approach outperforms the standard strong DPS baseline` -> The DPS result are not present in the table, so I think they are missing?
- l.335: What is the `EDM` framework?


- l.847: The change from  "l" to "j" should be made explicit as it is not immediately clear and can appear as a typo. 

### References

[A] : Daras, Giannis, et al. [A survey on diffusion models for inverse problems.](https://arxiv.org/pdf/2410.00083) arXiv preprint, 2024.  
[B] : Zhu, Yuanzhi, et al. [Denoising diffusion models for plug-and-play image restoration.](https://arxiv.org/pdf/2305.08995) Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.  
[C] : Wu, Zihui, et al. [Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors.](https://arxiv.org/pdf/2405.18782) arXiv preprint, 2024.
soundness: 3
strengths: - Interesting approach to solve inverse problems.
- Derivative-free approaches can be useful in many cases and have received much less attention than gradient-based methods.
summary: This paper proposes a new approach to solve inverse problems using a derivative-free optimization method based on the Ensemble Kalman Filter. The core idea is to approximate the data-fidelity term gradient with a statistical
linearization from the ensemble Kalman methods. The method is applied to three  types of inverse problems: computational imaging problems, the Navier-Stokes equation, and the black-hole imaging problem. The method is compared other derivative-free baselines.
weaknesses: - Positioning relative to the state-of-the-art is not clear, in particular with respect to the proposed framework, which seems to be a variant of the existing ones (see **Q2**).
- The evaluation is not completely satisfactory (see **Q3, Q4**).

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1) How does EnKG perform when applied to larger-scale problems or high-dimensional data? Are there specific limitations to its computational efficiency?
2) Is there any dependence on the performance and the pertrained model? How sensitive is the method to the quality and type of pre-trained diffusion model used? What are the implications if a suitable model is not available?
soundness: 3
strengths: 1)	EnKG operates without gradients, needing only black-box access to forward models, making it highly applicable to complex inverse problems with unknown or undefined derivatives. The proposed PC framework generalizes existing methods, enabling adaptability across various inverse problems without retraining.
2)	The current work demonstrates strong performance, notably in complex tasks like the Navier-Stokes equation, outperforming gradient-based solutions
3)	Provides deeper understanding and new interpretations of diffusion-based approaches, contributing to the field of inverse problem-solving.
summary: The paper presents Ensemble Kalman Diffusion Guidance (EnKG), a novel derivative-free method for solving inverse problems using pre-trained diffusion models. Traditional approaches often require detailed knowledge of the forward model, such as derivatives, which limits their applicability. EnKG overcomes this by relying solely on black-box evaluations of the forward model and a pre-trained diffusion model.
Key contributions are twofold: 1) EnKG operates solely with black-box access to forward model evaluations and a pre-trained diffusion model, making it particularly useful in scenarios where derivative information is inaccessible. 2) The authors introduce a prediction-correction (PC) framework that utilizes the empirical covariance matrix of ensemble particles during the correction step. This innovation allows EnKG to effectively bypass reliance on gradients, enhancing its applicability in non-linear inverse problems.
The paper demonstrates the effectiveness of EnKG across various inverse problems, including applications in fluid dynamics. These examples highlight the method's capability to handle complex, non-linear scenarios that are common in scientific research.
In summary, this work expands the toolkit for addressing inverse problems in machine learning by introducing a flexible and robust approach that maintains the generative power of diffusion models.
weaknesses: 1) The method may face challenges when scaling to very large models or high-dimensional data, as ensemble-based approaches can become computationally expensive. A further insight along this line would be useful. Also, it relies on pre-trained diffusion models, which might limit effectiveness if high-quality models are not available for certain tasks.
2) The empirical validation focuses on specific problem sets; broader testing across diverse applications would strengthen the generalizability claim. 
3) A further analysis on the algorithm complexity would be beneficial as the combination of the prediction and correction steps might introduce additional computational and implementation complexity due to ensemble covariance estimation.

---

confidence: 4
contribution: 4
initial_rating: 5
presentation: 3
questions: - Line 316: Is the DPS baseline included in the table?
- Black-Hole Imaging Problem: In the black-hole imaging problem, how is $G(\phi(x_i, t_i))$ computed if $G$ is unknown and $\phi(x_i, t_i)$ differs from the observed data? Could you provide a step-by-step explanation of how the black-box forward model is handled in the case? Additionally, please clarify if any assumptions are made about $G$ or generated samples in this scenario.
- Proof of Lemma 1: The proof shows a monotonic decrease in $\text{tr}(C_{xx}^{(i)})$. Why does this quantity converge to zero? From another perspective, a vanishing covariance implies the trajectories converge to a single point. Does this implication contradict the ill-posed nature of inverse problems, which typically have many possible solutions?

Possible Errata
- Lines 225 and 284: Remove  \Gamma .
- Line 731: Remove the extraneous ‘(‘.
- Line 847: Replace with $\approx$.
soundness: 2
strengths: 1.	Novel Approach: The paper introduces statistical linearization within ensemble Kalman methods to diffusion-based inverse problems, a novel concept in this context.
2.	Innovative Guidance Term Formulation: The authors present a unique formulation for the guidance term, with a clever trick that replaces the derivative of the forward model with covariance from forward evaluations.
3.	Comprehensive Validation: The effectiveness of EnKG is demonstrated across three different scenarios: (1) cases where the forward model is known and differentiable, (2) cases where the forward model is known but differentiating it is impractical (e.g., PDE-based models), and (3) cases where the forward model is a black box, with observations as the only available information.
summary: Diffusion models have been used to address inverse problems, with numerous diffusion-based solvers that avoid retraining existing diffusion models. These approaches typically rely on pseudo-inverses or derivatives tied to the forward model. This paper introduces a novel diffusion-based inverse solver designed for cases where the forward model is unknown.

The authors propose Ensemble Kalman Diffusion Guidance (EnKG), a derivative-free method that utilizes only forward model evaluations along with a pre-trained diffusion model.
In the proposed method, guidance term is computed as follows:
1. Particles are initialized to compute covariance in following steps.
2. During the diffusion trajectory, the particles are pushed by ODE solver.
3. Then, synthesized samples are applied to the forward model and compute covariance of them.
4. Diffusion trajectory is updated by the formula given the EnKG.
The proposed method replace the derivative of the forward model by computation of covariance and ODE solving.

The empirical results demonstrate the effectiveness of EnKG across diverse inverse problems, including cases where the measurement operator is treated as a black box. Specifically, the method is applied to image inversion problems with explicit forward models, Navier-Stokes inverse problems where the forward model is computed by solving PDEs, and black hole imaging, where the forward model is a black box.
weaknesses: 1. Limited Discussion of Related Work: The paper lacks depth in interpreting and explaining related works.
- The motivation behind the weighting matrix $w_i C_{xx}^{(i)}$ is unclear. Could you provide further explanation on the intuition and reasoning behind the choice of weighting matrix?
- Although the Kalman method is a core component, the paper does not provide a thorough explanation of its role and mechanics in this context. Specifically, which parts of the method are directly applied from existing literature and which represent novel contributions of this paper? For example, the introduction of weighting matrix, the derivation using local linearity of the operator in the proof, and the convergence claim. Related to above questions, please clarify the intuition and motivation provided by the literature versus that introduced by the authors.
2. Overstated Contribution: The claimed contributions seem somewhat overstated. The concept of Predictor Corrector interpretation in guidance-based methods is not entirely new.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: - Plot (b) in Figure 3 is difficult to interpret. Could you clarify what the values 0.2, 0.4, ..., 0.8 represent? Additionally, what do "Seq # DM" and "Seq # DM grad" refer to in this context?
- Plot (c) reports the runtime of the proposed algorithm, with an average of  approximately 140 minutes. This is a considerable computational cost (about 2 hours) for solving one inverse problem.
    * How does this runtime compare to other algorithms (aside from EKI) ?
    * Can you comment on the practical applicability of the method given this runtime?
    * Considering the high computational cost, how does the algorithm perform relative to methods that fine-tune or train smaller network components for the guidance term, as seen in [1, 2, 3, 4]?


---
.. [1] Black, Kevin, et al. "Training diffusion models with reinforcement learning." arXiv preprint arXiv:2305.13301 (2023).

.. [2] Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L., ... & Levine, S. (2024). Fine-tuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194.

.. [3] Fan, Ying, et al. "Reinforcement learning for fine-tuning text-to-image diffusion models." Advances in Neural Information Processing Systems 36 (2024).

.. [4] Denker, Alexander, et al. "DEFT: Efficient Finetuning of Conditional Diffusion Models by Learning the Generalised $ h $-transform." arXiv preprint arXiv:2406.01781 (2024).
soundness: 1
strengths: Introduction of an algorithm that solves inverse problems with diffusion models prior that only requires point-wise access to the forward model
summary: The authors address inverse problems using diffusion models as priors within a restrictive setup where the forward model is accessible only point-wise.
They employ a Predictor-Correct framework, where in the prediction stage, samples are drawn from the prior using the ODE describing the diffusion model.
In the correction stage, a MAP problem involving the forward model at the current diffusion step is solved through gradient-based updates.
Unlike previous approaches, the authors approximate the intractable term in this MAP formulation by evaluating it on samples generated via the ODE.
Since only point-wise access to the forward model is available, the authors estimate the gradient through statistical linearization and ensemble Kalman methods.
This approach maintains a set of particles and uses them, along with their centroids, to approximate the gradient.
The authors validate their algorithm on three different inverse problems.
weaknesses: **Methodological concerns**

- The authors’ motivation for their "Derivative-free correction step" in Lines 246–263. The matrix $C_{xx}$ appears without adequate justification, and in Equation (12), the authors invert $C_{xx}$ even though it is a singular matrix, as the number of samples used to compute it is less than the dimensionality of the problem. Consequently, $C_{xx}$ cannot be treated as a preconditioning matrix in this context.
- In the appendix, the authors base their proofs on Equation (20), which is introduced without sufficient explanation. This equation appears to correspond to the ensemble update, which assumes an estimation of the gradient, the very objective of the lemmas and propositions that follow. This circular reasoning raises concerns about the validity of the proofs.


**Technical concerns**

- In Lines 125–127 (paragraph following Equation (4)), the authors suggest that the guidance term depends on the noise scheduler $\dot{\sigma} \sigma$. However, this dependence result from their specific formulation. As verification, the authors can review Equation (5) in [1] or write DPS's algorithm in terms of the score. The guidance term is not scaled by the noise scheduler.
- The statement in Lines 194–195 is misplaced. Specifically, $\log \hat{p}(y | x_{i+1})$ is a composition of the simulated ODE and the forward model, making it highly non-convex. Therefore, the hypothesis of convexity is unrealistic. Besides, this term varies at each diffusion step, as it is composed with the ODE at different time steps, which diverges from the requirements outlined in [2], Chapter 4. Hence, the iterative updates may not converge to a true MAP estimate within this setups.


**Errors and clarifications**

- Equations (12)–(14) lack clarity. The variable $x_{i+1}'$ is undefined, and while the argmin is specified with respect to $x_{i+1}^{(j)}$, this variable does not appear in the equations.
- In Lines 897–903, the gradients of $p$ are missing a logarithmic term.
- The first part of Assumption (3) regarding $C_{xx}$ is redundant. Since $C_{xx}$ is a positive semi-definite matrix, its trace is non-negative, being the sum of positive eigenvalues. Therefore, if the trace of $C_{xx}$ is zero, $C_{xx}$ must be the zero matrix.
- In Lines 316–317, "DPG" should replace "DPS," as DPS is not included in the experiments.
- The term "guidance" is repeated twice in Line 52.


---

.. [1] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

.. [2] Parikh, Neal, and Stephen Boyd. "Proximal algorithms." Foundations and trends® in Optimization 1.3 (2014): 127-239.

Meta-review: The paper proposes a method for solving inverse problems with diffusion models that only requires forward model evaluations as opposed to being able to differentiate through the forward model. The paper approximates the gradient corresponding to a data consistency term. 

The reviewers and the AC agree that the problem considered is interesting and important and the approach promising.

However, the majority of the reviewers recommend rejecting the paper. The key issues that remain after the rebuttal and discussion are:

- Reviewer Fiay identified methodological and technical issues, and as pointed out by reviewer Fiay some of the technical issues remain in the revision, although several have been fixed. 

- Reviewers 8WU4 and G5qZ both find that the positioning relative to gradient-based methods and tradeoffs are insufficient, and it remains questionable how gradient-free methods can outperform gradient based ones. Moreover the reviewers oversells the contributions.

The paper has potential and the problem and approach are interesting, but for the two remaining key issues, I recommend rejecting the paper.

Decision: Reject