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

Authors: Hongkai Zheng, Wenda Chu, Austin Wang, Nikola Borislavov Kovachki, Ricardo Baptista, Yisong Yue
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: inverse problem, diffusion model, derivative-free

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 many scientific applications. To address this, 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. 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.

Decision: reject
Reviewer scores: 5, 6, 5, 3

Reviews and discussion:
<thread>
creation date:
1730681993398

thread_type:
official_review

<thread_object>
creator: reviewer_fiay

<content>
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.

strengths:
Introduction of an algorithm that solves inverse problems with diffusion models prior that only requires point-wise access to the forward model

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.

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).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691602942

thread_type:
official_review

<thread_object>
creator: reviewer_g5qz

<content>
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.

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.

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.

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$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721268131

thread_type:
official_review

<thread_object>
creator: reviewer_x3ty

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731021420541

thread_type:
official_review

<thread_object>
creator: reviewer_8wu4

<content>
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.

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.

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**).

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732221977270

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Meta response

comment:
We sincerely thank the reviewers for their constructive feedback, insightful questions, and recognition of our contributions. We are pleased that the reviewers highlight the proposed PC framework as “providing deeper understanding” (x3Ty) of existing guidance-based algorithms and recognize EnKG as an "interesting" (8WU4) and "novel" (G5qz) approach to solve inverse problems with only black-box access. The reviewers appreciate the "strong performance" (x3Ty) of EnKG across various problems as well as its "comprehensive validation" (G5qz). We address the reviewers’ questions, concerns, and suggestions in our individual responses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733256390037

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Summarize the Discussion

comment:
We sincerely thank the reviewers for their constructive feedback and engagement during the discussion. As the discussion period has now closed, we would like to summarize the key points of discussion where reviewers have not responded or confirmed that their concerns were addressed:
- **Derivative-free vs gradient-based** (8WU4) The reviewer raised concerns regarding the trade-offs between derivative-free and gradient-based methods. We argue that this trade-off is inherently challenging to quantify, especially for PDEs, as developing reliable gradient estimation methods remains an open research question [1,2]. We also emphasized the intellectual merit of studying derivative-free methods for inverse problems, which includes:
    - Applicability to non-differentiable problems (e.g., symbolic music generation [3], stochastic weather models [4]).
    - Handling scenarios where the forward model’s internal structure is unknown, such as black-box APIs
    - High parallelizability that allows them to benefit from the recent advancements in large-scale distributed computing infrastructure, as discussed in [5].
    - Allowing for more flexible forward modeling without requiring structural constraints for gradient computation.
- **Question about Assumption 3** (G5qz) We have provided further clarification on Assumption 3 and its practical implication. Empirical evidence is also provided in Figure 5. 
- **Misuse of terminology "pre-conditioner"** (Fiay) We have removed the term "pre-conditioner" and revised the dicussion to focus on the practical role of $C_{xx}$, which is to project the gradient onto the subspace spanned by the particles.
- **Typos in the superscript** (Fiay) We appreciate the reviewer’s careful attention to detail and have corrected this typo in the updated manuscript. 
- **Question about Lemma 1 proof step** (Fiay) We have provided a detailed explanation of the proof step both in rebuttal and the revised manuscript (lines 758-776) to address the reviewer’s concerns.

We hope this summary demonstrates our efforts to address all potential concerns. We kindly ask the reviewer to reassess the score in light of the key contributions and value of our work, as highlighted in our meta response. We feel that our study of derivative-free methods for diffusion guidance is a sufficiently novel intellectual contribution in its own right and worthy of publication at ICLR.

[1]:  Gaggioli, E. L., and Oscar P. Bruno. "Parallel inverse-problem solver for time-domain optical tomography with perfect parallel scaling." Journal of Quantitative Spectroscopy and Radiative Transfer 290 (2022): 108300.
[2]: Fang, Lean, and Ping He. "A duality-preserving adjoint method for segregated Navier–Stokes solvers." Journal of Computational Physics 503 (2024): 112860.

[3]: Huang, Yujia, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli Shama Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. "Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion." In Forty-first International Conference on Machine Learning.

[4]: Miyoshi, Takemasa, and Masaru Kunii. "The local ensemble transform Kalman filter with the Weather Research and Forecasting model: Experiments with real observations." Pure and applied geophysics 169.3 (2012)

[5]: Salimans, Tim, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. "Evolution strategies as a scalable alternative to reinforcement learning." arXiv preprint arXiv:1703.03864 (2017).
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.