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

Q1: We appologize for the confusion in Figure 3(b). We clarify as follows:

• Explain metrics Figure 3(b) compares the computation complexity of different algorithms. As defined in "evaluation metrics" paragrpah in Section 4.2, Seq # DM is the number of sequential diffusion model evaluations, representing the count of diffusion model evaluations that must be performed in sequence where each query depends on the previous ones. Seq # DM grad is similar to Seq # DM but focuses on the gradient evaluation of diffusion model.
• Why we measure number of sequential evaluations? This is because the sequential evaluations determine the minimum runtime irrespective of available computational resources. By isolating the sequential evaluation, we can better assess how much of the algorithm can potentially benefit from parallelization. Similar to the concept of "critical path" for algorithm analysis in the computer science literature [5].
• "What does 0.2,..,0.8 mean?": to compare the relative computation complexity of different algorithms, each axis is normalized by diving by the maximum across all the algorithms. For example, a value of 0.8 indicates that the corresponding algorithm uses 80% of the computational resources compared to the most resource-intensive algorithm for that metric.

Q2: Runtime comparison: To provide a clearer perspective, we report the runtime of different algorithms on the Navier-Stokes equation, tested on a single A100 GPU. These results are now incorporated into Table 2 of the revised manuscript. As shown below, EnKG achieves the best performance while maintaining a runtime comparable to or even shorter than other algorithms.

All the algorithms take a while to finish. We believe this reflects the intrinsic difficulty of the problem rather than a limitation of our algorithm. Solving inverse problems where the forward model involves a numerical PDE solver is computationally expensive across all methods. For instance, Figure 13 in [14] reports runtimes ranging from 3–16 hours for a 2D spatio-temporal PDE problem. It is normal in these type of PDE-based problems to take hours or even days.
Practical Applicability Given its inherently parallel nature, EnKG can take advantage of recent advancements in large-scale distributed computing infrastructure to mitigate these challenges. With appropriate optimization and resource allocation, we believe that EnKG remains scalable even for high-dimensional problems, especially as computational hardware and distributed frameworks continue to improve. Its derivative-free nature and ability to operate with black-box forward models make it a practical option for challenging scientific problems where gradient-based approaches may be infeasible such as numerical weather model calibration [16,17].

Comparison to conditional finetuning Thanks for the references and the interesting question. To investigate how the conditional diffusion models work, we implement the idea of provided reference DEFT [4] on our Navier-Stokes equation problem. Note that the original DEFT requires the gradient information as the input (Eq. (12) in their paper). Since our focus is on the black-box forward model setting, where gradients are unavailable, we replace the gradient input with the meansurement y. The table below compares the relative L2 error of DEFT and EnKG.

As shown in the table, DEFT is not effective in the black-box setting. We hypothesize that this is because DEFT, as a hybrid method, relies critically on gradient information to guide the generation process. Without access to gradients, its performance degrades significantly. In contrast, EnKG is designed to operate without gradient information, making it well-suited for problems where the forward model is a black box. As the results indicate, EnKG achieves significantly better performance than DEFT in this scenario.

[5]: Kohler, Walter H. "A preliminary evaluation of the critical path method for scheduling tasks on multiprocessor systems." IEEE Transactions on Computers 100.12 (1975): 1235-1238.

[6]: Booton, Richard C. "Nonlinear control systems with random inputs." IRE Transactions on Circuit Theory 1.1 (1954): 9-18.

[7]: Kazakov IE (1954) Approximate method of statistical investigations of nonlinear systems. In: Proceedings, Voenno-Vozdushnaya Akademia imeni N.I. Zhukuvskogo, vol 394, pp 1–52

[8]: Evensen, Geir. "The ensemble Kalman filter: Theoretical formulation and practical implementation." Ocean dynamics 53 (2003): 343-367.

[9]: Iglesias, Marco A., Kody JH Law, and Andrew M. Stuart. "Ensemble Kalman methods for inverse problems." Inverse Problems 29.4 (2013): 045001.

[10]: Garbuno-Inigo, Alfredo, et al. "Interacting Langevin diffusions: Gradient structure and ensemble Kalman sampler." SIAM Journal on Applied Dynamical Systems 19.1 (2020): 412-441.

[11]: Calvello, Edoardo, Sebastian Reich, and Andrew M. Stuart. "Ensemble Kalman methods: a mean field perspective." arXiv preprint arXiv:2209.11371 (2022).

[12]: Okulicka-Dłużewska, Felicja. "Pseudoinverse preconditioner for saddle point problem." AIP Conference Proceedings. Vol. 1504. No. 1. American Institute of Physics, 2012.

[13]: Cahueñas, Oskar, Luis M. Hernández-Ramos, and Marcos Raydan. "Pseudoinverse preconditioners and iterative methods for large dense linear least-squares problems." Bull. Comput. Appl. Math 1 (2013): 69-91.

[14]: 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.

[15]: Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.

[16]: Buehner, Mark, Ron McTaggart-Cowan, and Sylvain Heilliette. "An ensemble Kalman filter for numerical weather prediction based on variational data assimilation: VarEnKF." Monthly Weather Review 145.2 (2017): 617-635.

[17]: 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): 321-333.

We truly appreciate the reviewer’s detailed suggestions and careful reading of our manuscript. We address the reviewer’s concerns as follows:

• We would like to clarify that first part of the Assumption 3 is not redundant. Specifically, $C_{yy}$ is the observation empirical covariance matrix. In general, $tr(C_{yy})=0$ does not necessarily imply that $C_{xx}=0$. For example, if the forward map $\phi$ maps everything to a constant, $tr(C_{yy})=0$ but $C_{xx}\neq 0$. Assumption 3 baiscally ensures that the forward model does not map all different particles into identical observation.
• We have corrected the typos mentioned by the reviewer in the revised version of the manuscript.

We thank the reviewer for the detailed comments and constructive suggestions, and we apologize for any confusion caused by the presentation. We address the reviewer's concerns and questions below.

Methodological concerns

W1: Regarding the covariance matrix $C_{xx}$, we first note that our algorithm does not explicitly use $C_{xx}$. It is only used for derivation. When $C_{xx}$ is singular, the correction step can be interpreted as a projected gradient step onto the subspace spanned by the particles. Below, we address the reviewer's concerns in detail.

• Role of the covariance matrix: intuitively, the matrix $C_{xx}$ is an important component for the derivation that facilitates interaction among particles in our framework. In our paper, as we shown in Proposition 1, introducing $C_{xx}$ enables us to derive the update rule of the particles for solving the inverse problem without explicit derivatives. The covariance matrix naturally appears in the expression of statistical linearization [6,7]. The combination of covariance matrix and statistical linearization is a well-established idea and central feature in the family of ensemble Kalman methods [8,9,10] dated back to the 1990s by Evensen et al [8]. To provide more context for this step, We add additional explanation after Eq.(12) in the revised version.
• Singular covariance matrix is valid and does not affect the drivation and proof: we appreciate the reviewer’s observation regarding the singularity the covaraince matrix when the number of particles is less than the problem dimension. While we agree that this case requires further discussion, we emphasize that singular covariance matrix does not affect the derivation or proof as they do not rely on $C_{xx}$ being full-rank. Because the covariance matrix is always positive semi-definite (PSD), its inversion in Eq. (12) can be generalized to Moore–Penrose inverse (aka. pseudoinverse). The use of singular PSD matrix as preconditioner is valid and commonly adopted in the numerical iterative algorithms [12,13].
• Correction step is still effective even with singular covariance matrix In our case, when $C_{xx}$ is singular, Eq.(14) effectively becomes a gradient step projected onto the subspace spanned by the particles. This is effective when the intrinsic problem dimension is smaller than the ambient space dimension. As shown in our experiments, 1~2k particles is already sufficient to solve various inverse problems effectively. As further studied in Figure 6, further increasing the number of particles yeilds diminishing return, supporting the practical utility of our approach even in high-dimensional problems.
• We agree this is an important case to consider and have added the discussion to lines 264-269 to clarify.

W2: Thanks for the constructive feedback. We apologize for any confusion caused by the current presentation. To address the reviewer’s concerns and ensure clarity, we have revised Proposition 1 and Lemma 1 in the updated version. The reasoning path is as follows:

• Starting point We begin with the ensemble update rule as defined in Eq.(16) (revised version).
• Convergence result Based on this update rule, we establish the convergence results presented in Lemma 1.
• Link Lemma 1 to Proposition 1 Lemma 1 is then used to prove Proposition 1.
• Proposition 1 connects the update rule to Eq.(14) Proposition 1 shows that our ensemble update rule given by Eq.(16) approximates the gradient correction step defined in Eq.(12) without explicit derivative.

We hope this resolves the reviewer’s concerns, and we are happy to provide further clarifications if needed.

Technical concerns W3: Thanks for sharing your comment. What we intended to convey is that the guidance term should theoretically depend on the noise schedule in the reverse process if the goal is posterior sampling. To see this, Eq.(4) in the DPS paper expresses the reverse SDE for posterior sampling, where the coefficient of the guidance $\nabla_{x_t} \log p(y|x_t)$ is $-\beta(t)$, directly tied to the noise schedule. However, in DPS’s algorithm (Algorithm 1 in their paper), the coefficient is implemented as a heuristic that is unrelated to $\beta(t)$ (see footnote 5 on page 6 of DPS paper). This suggests that posterior sampling interpretation does not align with how DPS actually works in practice. This is one reason why we want to propose the Prediction-Correction framework as an alternative interpretation.

W4: Thank you for pointing out this issue. We agree that the underlying optimization problem is generally non-convex, which means the algorithm may converge to a local maximum rather than a global one. To address this, we have revised the statement to explicitly clarify that.

Highlight the modification of the revision We apologize for any inconvenience caused by the lack of highlights of the revised text. To facilitate the reviewer’s evaluation, we have highlighted all revisions in the updated manuscript. Additionally, we provide a summary of major updates in the meta response for ease of reference.

Misuse of terminology "pre-conditioner" We thank the reviewer for the reference and detailed explanation of the concern. We agree that describing $C_{xx}$ as the term "pre-conditioner" could be misleading as as it traditionally applies to invertible matrices. To fully address the reviewer's concern, 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. This change has been reflected in lines 264–267 of the updated manuscript.

Flaws in Proposition 1

Issues with Eq.(18): we thank the reviewer for pointing out the issue and apologize for the confusion. There was indeed a typo in the superscript on the right-hand side of Eq.(18). The corrected right-hand side is as follows, where the index $k$ appears in both terms. $\frac{1}{J}\sum_{k=1}^{J}\left\langle \psi\left(x^{\prime(k)}\_i\right) - \bar{G}, y - \psi\left(x^{\prime(j)}\_i\right) \right\rangle_{\Gamma} \left(x^{\prime(k)}\_i -\bar{x\_i} \right)$

We appreciate the reviewer’s careful review and have corrected this typo in the updated manuscript. Additionally, we have reviewed the entire manuscript to ensure these indexing typos have been fixed.

Concerns about derivation of $C^{(i+1)}\_{xx}$ in Lemma 1 We appreciate the reviewer's feedback and agree that this step requires explicit justification. To clarify, we have added the following detailed explanation to explicitly show the equality: we first recall the definition of $C\_{xy}^{(i)}$:

Using this definition, we rewrite one of the cross terms as: $\frac{1}{J}\sum\_{j=1}^{J} \left(w\_i C\_{xy}^{(i)}\left(\bar{\psi}\_i-\psi(x\_i^{(j)})\right)\right)\left(x\_i^{(j)}-\bar{x}\_{i}\right)^\top = - w\_i C\_{xy}^{(i)}C\_{xy}^{(i)\top}$

Similarly, the other cross term can be simplified as $\frac{1}{J}\sum\_{j=1}^{J} \left(x\_i^{(j)}-\bar{x}\_{i}\right)\left(\left(\bar{\psi}\_i-\psi(x\_i^{(j)})\right)^\top w\_iC\_{xy}^{(i)\top}\right) = - w\_i C\_{xy}^{(i)}C\_{xy}^{(i)\top}.$ Therefore, the sume of the cross terms simplifies to $-2w_iC\_{xy}^{(i)}C\_{xy}^{(i)\top}$. The full derivation is provided in lines 758-776 of the updated manuscript.

We hope these clarifications address the reviewer’s concerns, and we are happy to provide further elaboration if needed.

Q1: Thanks for pointing it out. We have removed that sentence from the main text to avoid confusion in the revised version. For a more comprehensive evaluation, we have added a comparison against a few gradient-based methods including DPS as additional reference points in Appendix A.8.

Q2: Thanks for the question. We assume G is black-box function, meaning that for any input $x$, we can evaluate $G(x)$ without having access to the internal structure or parameters of $G$ - similar to scenarios where software provides an API to evaluate $G(x)$ but does not expose its source code. If this does not align with the reviewer’s interpretation of “unknown,” we would appreciate clarification. In the black hole imaging expeirment, we evaluate $G(x)$ with the forward model inImage.observe from ethim library, which is a function of $x$. As defined in Eq.(1), $x$ denotes the input to the forward model where the observed data is $y$.

Q3: Thanks for the questions.

• Proof clarification We agree that the last step in the proof could benefit from additional clarification. In response, we have revised the last step of the proof as follows: By assumption 1 and 2, we know that both $tr\left(C^{(i)}\_{xx}\right)$ and $tr\left(C^{(i)}\_{yy}\right)$ are upper bounded. Further by Assumption 3, $tr\left(C^{(i)}\_{xy}C^{(i)\top}\_{xy}\right)$ is lower bounded unless all the particles collapse in a single point. Therefore, there exists an $\alpha>0$ such that $tr\left(C^{(i)}\_{xy}C^{(i)\top}\_{xy}\right)\geq \alpha \cdot tr\left(C^{(i)}\_{xx}\right)tr\left(C^{(i)}\_{yy}\right).$ Theforefore, plugging this into Eq.(23) gives $tr\left(C^{(i+1)}\_{xx}\right) \leq (1-\alpha) \cdot tr\left(C^{(i)}\_{xx}\right),$ showing that $tr\left(C^{(i)}\_{xx}\right)$ monotonically decreases to zero.
• Ill-posedness
  • Different initialization may lead to different solutions The convergence result indicates that the algorithm returns a point estimate (an ensemble of particles that are close to each other) for each run. However, the randomness in the initialization (prior step) can still lead to different convergent points across runs just like other optimization methods. Thus, the convergence does not imply the algorithm cannot find different solutions.
  • Clarification on ill-posedness The ill-posedness suggests that potentially many $x$ satisfy $y=G(x)$. However, when equipped with a proper prior, the potential solution space will be largely constrained, which is the motivation of using strong diffusion prior. In such cases, the set of possible solutions may reduce to a unimodal distribution.

Possible Errata:

• We apologize for the confusion. The inner product with subscript $\Gamma$ represents the weighted Euclidean inner product, where $\Gamma$ is the covariance matrix of the Gaussian noise in the observation, as defined in Eq.(1). We have added this notation and its definition to Table 4 for clarity.
• Other issues have been fixed in the revised version. Thanks again for the detailed review.

We thank the reviewer for the constructive feedback and detailed review, and we appologize for any confusion caused by the presentation. In response, we address the reviewer's concerns and questions below.

W1: Role of the covariance matrix: intuitively, the matrix $C_{xx}$ is an important component for the derivation that facilitates interaction among particles in our framework. In our paper, as we shown in Proposition 1, introducing $C_{xx}$ enables us to derive the update rule of the particles for solving the inverse problem without explicit derivatives. The covariance matrix naturally appears in the expression of statistical linearization [1,2]. The combination of covariance matrix and statistical linearization [1,2] is a well-established idea and central feature in the family of ensemble Kalman methods [3,4,5,6] dated back to the 1990s by Evensen et al [3]. To provide more context for this step, We add additional explanation after Eq.(12) in the revised version.

Clarification on contribution

• In a nutshell, the proposed EnKG is a novel algorithm combines the expressive capabilities of data-driven priors with the ability to solve inverse problems using only black-box access to the forward model. EnKG introduces the core idea of statistical linearization from ensemble Kalman methods [3,4,5] into the Prediction-Correction framework for guided diffusion. At a high level, EnKG can also be viewed as an ensemble Kalman Inversion method regularized by the diffusion model prior.
• Relation to existing literature
  • The use of covariance matrix stems from statistical linearization. This idea originates from [1,2] and is the fundamental building block of the family of ensemble Kalman method (See [6,7] for overview). However, its integration into the proximal operator defined in Eq.(12) is proposed by us.
  • The convergence result in Lemma 1 is also developed by us. While prior works [9,10] establish the convergence results for other ensemble Kalman method in the non-linear setting, their proof do not directly apply to our case due to the difference in the update rule. That said, they share similar intuition and assumptions (e.g., regularity assumptions and bounded particles).
  • The use of the local linearity of the operator is mentioned in [8,10]. However, [8] does not provide a formal proof and focuses on a different objective function and a momentum-based variant while [10] examines a different variant of EKI with Tikhonov regularization and covariance inflation.
• To address the reviewer's concerns and improve the clarity, we have also updated the Related Work section to extend the dicussion on the literature.

W2: Thank you for your feedback.

• To the best of our knowledge, no prior work explicitly provides the Predictor-Corrector interpretation for guidance-based methods as presented in our paper. We did just find a concurrent ICLR 2025 submission [11] that proposes a similar concept but focus on classifier-free guidance.
• The primary contribution of our work is the development of the EnKG algorithm, which integrates the Predictor-Corrector framework with ensemble Kalman methods. We are open to revising our contribution statement and softening the wording around the Predictor-Corrector interpretation if the reviewer can provide specific references to prior work that aligns with our approach.

Dear Reviewer G5qz,

We want to thank you for your careful review and valuable feedback. We have carefully revised the manuscript based on your feedback and have provided further clarifications in our rebuttal. As the discussion period concludes today, we would like to check if our responses have addressed your concerns. If so, we would greatly appreciate it if you could update your score to reflect this.

We sincerely value your input and are happy to provide further clarifications if needed.

Best regards

Apologies for the delayed response.

I have reflected further on the feasibility of the proposed theory.

Since $N$ is finite, $\text{tr}(C^{(i+1)}) \leq (1-\alpha) \text{tr}(C^{(i)})$ does not guarantee that the quantity converges to zero. Other types of inequalities, such as $\text{tr}(C^{(i+1)}) \leq \text{tr}(C^{(i)}) - \beta$, would be required.

Unfortunately, I do not believe such an inequality is achievable under the given assumptions.

Additionally, I would like to highlight some points that were missed during the previous review process: If a matrix $A = \sum_i a_i a_i^T$, then $\text{tr}(A) = 0 \iff a_i = 0 \, \forall i \iff A = 0$.

The assumptions imply that $\psi(x^{(j)}) - \bar{\psi} = 0 \implies x^{(j)} = \bar{x}$, which is a rather strong condition, especially when $\psi$ is ill-posed.

I recommend that the authors provide empirical evidence to demonstrate that the proposed derivative-free correction adequately approximates the actual correction term.

Q1: Thanks for the good question.

• As discussed in our response to W3, the proposed EnKG is well-suited for large-scale problem when the forward model is the major cost and can be evaluated in parallel. In such cases, EnKG’s reliance on parallelizable updates ensures scalability when equipped with the modern distributed computing infrastructure. However, if the forward model is inherently slow and cannot be parallelized, EnKG’s runtime may become a limiting factor.
• Exploring even higher-dimensional problems to benchmark EnKG and other derivative-free methods is an exciting direction for future research. However, such an effort is beyond the scope of this paper, which primarily focuses on introducing and analyzing the new algorithmic design.

Q2: This is an interesting and insightful question.

Additional controlled experiment To investigate the dependence on the quality of pre-trained diffusion models, we conducted a controlled experiment on Navier-Stokes equation. Specifically, we trained a diffusion model prior using only 1/10 of the original training set and limited the training to 15k steps to simulate a lower-quality model. We evaluate the top two algorithms EnKG and DPG, and report the relative L2 error below.

Robust performance We observe that while both algorithms experienced a performance drop due to the reduced quality of the diffusion model, the decline was relatively small compared to the significant reduction in training data. Notably, our EnKG demonstrated greater robustness, with a smaller performance drop than the best baseline method, DPG. These results indicate that while EnKG benefits from high-quality diffusion models, it is not overly sensitive to their quality. It maintains strong performance even with reduced model capabilities. We have included this interesting experiment and results in Appendix A.9.

[1]: Kohler, Walter H. "A preliminary evaluation of the critical path method for scheduling tasks on multiprocessor systems." IEEE Transactions on Computers 100.12 (1975): 1235-1238.

W1: Thank you for pointing out this important consideration. We agree that EnKG may face high overall computational costs when scaling to very high-dimensional problems. However, due to its inherently parallel nature, EnKG can take advantage of recent advancements in large-scale distributed computing infrastructure to mitigate these challenges. With appropriate optimization and resource allocation, we believe that EnKG remains scalable even for high-dimensional problems, especially as computational hardware and distributed frameworks continue to improve. Regarding the reliance on the pretrained model, we address the concern in the response to Q2 below.

W2: We thank the reviewer's comment on the scope of our experiments. Establishing a proper set of benchmarks for this style of problem is a major project in its own right and we are excited to pursue it in the future.

W3: We would like to first clarify that EnKG does not explicitly calculate or store the covariance matrix, as shown in Algorithm 2. Instead, the ensemble updates are performed using vector inner products, which are computationally efficient. Consequently, the primary computational cost arises from the forward model queries and diffusion model queries, which are reflected in the metrics for computational complexity defined in the “Evaluation Metrics” section of Sec. 4.2. For the reviewer’s convenience, we explain these metrics below.

Metrics for evaluating computational complexity

• Total # Fwd is the total number of forward model queries i.e., $G(x)$.
• Seq # Fwd is the number of sequential forward model queries, representing the count of forward model queries that must be performed in sequence where each query depends on the previous ones.
• Total # DM and Seq # DM: similar to Total # Fwd and Seq # Fwd but focuses on the diffusion model evaluations.
• Total # DM grad and Seq # DM grad: similar to Total # Fwd and Seq # Fwd but focuses on the gradient evaluation of the diffusion model.

The sequential metrics are important as they determine the minimum runtime irrespective of available computational resources. By isolating the sequential evaluation, we can better assess how much of the algorithm can potentially benefit from parallelization. Similar to the concept of "critical path" for algorithm analysis in the computer science literature [1].

Analysis of Computational Complexity Figure 3(b) compares different algorithm through the len of the above metrics. As we can see, EnKG requires the least sequential forward model calls, which indicates that EnkG can scale up to large-scale problem when the forward model is the major cost and can be evaluated in parallel.

[1]: Deng, Chengyuan, et al. "OpenFWI: Large-scale multi-structural benchmark datasets for full waveform inversion." Advances in Neural Information Processing Systems 35 (2022): 6007-6020.

[2]: Allgower, Eugene L., and Kurt Georg. Numerical continuation methods: an introduction. Vol. 13. Springer Science & Business Media, 2012.

[3]: Kochkov, Dmitrii, et al. "Machine learning–accelerated computational fluid dynamics." Proceedings of the National Academy of Sciences 118.21 (2021): e2101784118.

[4]: Chen, Ricky TQ, et al. "Neural ordinary differential equations." Advances in neural information processing systems 31 (2018).

[5]: 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.

[6]: Buehner, Mark, Ron McTaggart-Cowan, and Sylvain Heilliette. "An ensemble Kalman filter for numerical weather prediction based on variational data assimilation: VarEnKF." Monthly Weather Review 145.2 (2017): 617-635.

[7]: 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): 321-333.

[8]: Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.

[9]: Kohler, Walter H. "A preliminary evaluation of the critical path method for scheduling tasks on multiprocessor systems." IEEE Transactions on Computers 100.12 (1975): 1235-1238.

[10]: Chung, Hyungjin, et al. "Diffusion Posterior Sampling for General Noisy Inverse Problems." The Eleventh International Conference on Learning Representations.

[11]: Peng, Xinyu, et al. "Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance." Forty-first International Conference on Machine Learning. 2024.

[12]: Tang, Haoyue, et al. "Solving General Noisy Inverse Problem via Posterior Sampling: A Policy Gradient Viewpoint." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

Q4: Parameter choices: We apologize for the confusion and clarify the parameters choices as follows. First, $Q$ is the number of forward model queries for our baselines (Forward-GSG and Central-GSG). We set $Q=2048$ and it is chosen so that the overall runtime cost of Forward-GSG and Central-GSG is similar to the proposed EnKG. For the choice of $J$, we conducted an ablation study on Navier-Stokes equation, as shown in the figures 6 and 8. The performance improves as we increase the number of particles, with $J=1024, 2048$ offering the best trade-off between accuracy and runtime. For the numer of steps, we simply fixed it to 144 as we follows the EDM framework [8]. Figure 2 in [8] suggests that the EDM sampler yields marginal improvement beyond this point. The parameter choices of EnKG for each type of problem are specified as below.

Runtime comparison We also report the runtime for different algorithms on the Navier-Stokes equation, measured on a single A100 GPU. Table 2 has been updated to include these runtime results. For PnP-DM, we explored all the hyperparameter combinations in their paper, but all of them crashed with numerical instability in the solver. We found that the Langevin Monte Carlo generates noisy vorticity field that violates the stability condition of the PDE solver. So we had to use very small time step for statbility, which significantly increased runtime (estimated to exceed 100 hours). Thus, we marked these results as "timeout" below. On the other hand, DPS is about 3.5x-4x faster but produces worse results than EnKG as reported in our earlier comparisons.

Elaboration on the computation complexity metrics

• We identify that the primary computation resource requirements of inverse problem solvers are the forward model query and diffusion model query. Our chosen metrics aim to comprehensively represent these resource demands. More specifically,
  • Total # Fwd is the total number of forward model queries i.e., $G(x)$.
  • Seq # Fwd is the number of sequential forward model queries, representing the count of forward model queries that must be performed in sequence where each query depends on the previous ones.
  • Total # DM and Seq # DM: similar to Total # Fwd and Seq # Fwd but focuses on the diffusion model evaluations.
  • Total # DM grad and Seq # DM grad: similar to Total # Fwd and Seq # Fwd but focuses on the gradient evaluation of the diffusion model.
• Why we measure number of sequential evaluations? This is because the sequential evaluations determine the minimum runtime irrespective of available computational resources. By isolating the sequential evaluation, we can better assess how much of the algorithm can potentially benefit from parallelization. Similar to the concept of "critical path" for algorithm analysis in the computer science literature [9].

Q5: 50k black-hole GRMHD simulations are used for training diffusion model.

Minor comments, nitpicks and typos: thanks for the reviewer's detailed comments. We address your concerns below:

• Update "more computationally efficient than other derivative-free methods" for clarity.
• EDM is proposed by [8] to express the diffusion models in a unified and clean framework (introduced in Line 96). We have added the citation to the word EDM to avoid confusion.
• We have fixed the other issues in the revised version.

Q3: We appreciate the insightful question and constructive suggestions.

Part 1: additional comparison against gradient-based methods as reference

• For a more comprehensive evaluation, we have added results of gradient-based methods in Table 8, Table 9, and Table 10 in Appendix A.8 (last page). We also present them right below for the reviewer's convenience. Specifically, we have added DPS and DiffPIR for imaging tasks, DPS and PnP-DM to Navier-stokes equation and black-hole experiments as requested by the reviewer.
• The results show that EnKG clearly outperforms gradient-based methods like DPS and PnP-DM on Navier-Stokes equation and achieves comparable results on image restoration and black hole imaging. The details about the comparison are added to Appendix A.8.

FFHQ256 image restoration tasks:

Relative L2 error on Navier-Stokes equation:

Black hole imaging

Part 2: response to the question on differentiable solvers

• We acknowledge the reviewer’s point about differentiable solvers. While differentiable solvers are valuable in certain contexts, we argue that developing derivative-free methods remains a valuable direction for the following reasons:
  • Autograd is not always efficient or reliable Automatic differentiation frameworks like JAX or PyTorch provide a general and straightforward way to estimate gradient. However, it often becomes computationally prohibitive for PDEs with high iteration counts due tracking a huge computation graph caused by unrolling the iterative numerical solvers. For example, in our Navier-Stokes equation problem, the autograd encounter out-of-memory issue when the pseudospectral solver unrolls for more than around 6k steps (on a A100-40GB GPU). Additionally, the paper [3] associated with aforementioned jax-cfd package explicitly mentioned this issue in the first paragraph of page 4 in their paper. Therefore, autodiff is only practical for a limited range of problems. Moreover, autodiff can introduce additional numerical error, as evidenced in neural ODE work [4], which prompted the authors to derive an adjoint method to control the error.
  • Efficient and robust adjoint methods are not always easy to design While adjoint methods exploit PDE structure for efficient gradient estimation, they often require problem-specific derivations. Robust and efficient designs for adjoint methods remain an active area of research [5].
  • Gradient could be inaccessible for many other reasons For instance, the forward model could be stochastic, partially unknown, or implemented as a black-box software, making gradient computation infeasible. We would be happy to provide further elaboration if the reviewer wishes.
• Given these challenges, we believe developing derivative-free methods is also a valuable direction. For example, derivative-free methods have been widely used in operational weather forecasting centers [6,7].

We appreciate the constructive suggestions and detailed review, and we apologize for any confusion caused by our presentation. We address the reviewer's questions and concerns below.

Q1: Thanks for the good question. The prior diffusion models typically need to be trained on domain-specific data to fully leverage domain characteristics, as we did with Navier-Stokes equation and black-hole data. Once trained, it can reused for different problems (various forward models) within that domain. For example, the same black-hole prior can be reused for different telescope configurations, such as varying the number of telescopes, their positions, or the level of measurement noise. Similarly, this adaptability extends to other domains, such as seismology and ocean dynamics, where the same prior can support different problem configurations, which is practically useful.

Exploring diffusion priors that generalize across multiple domains is an exciting direction for future research. However, such an effort would require studying the interplay between prior transferability and the underlying problem characteristics, which extends beyond the scope of this paper.

Q2: We thank the reviewer for the referrences and questions, and we applologize for any confusion caused by our presentation.

• Correction step is general and requires algorithmic design We would like to first clarify that Eq.(9) is a general expression for showing how existing guidance-based methods fit within our Prediction-Correction framework. This feature is shared by many other algorithms [10,11,12], including DiffPIR. The specific design of the correction step is an important algorithmic contribution and differentiates various methods.
• EnKG vs DiffPIR The correction step of EnKG is conceptualized in Eq.(14) and implemented as Eq.(16) or Eq.(17). If the reviewer is referring to a comparison between Eq.(14) or Eq.(16) in our paper and Eq.(13) in DiffPIR, we elaborate on the major differences below:
  • Optimization space In Eq.(13) of DiffPIR, they optimize over $x_0^{(t)}$, the clean data space $p_0$ (implemented with Tweedie's formula in their paper). In contrast, as shown in Eq.(12), EnKG tackles the subproblem in $x_{t}$, which resides in the $p_{t}$ space.
  • Update rule EnKG solves the problem in Eq.(12) in a completely different way. EnKG runs an ensemble of interacting particles to solve it without gradient information. In comparison, DiffPIR performs the standard gradient step. These distinct approaches yield entirely different algorithms, as evident from comparing Algorithm 2 in the proposed EnKG and Algorithm 1 in the DiffPIR paper.
  • Derivative-Free Optimization The differences in algorithm lead to a key different property: derivative-free nature of EnKG. EnKG can solve problems that DiffPIR cannot handle as the gradient is not always avaliable. We discuss the broader value of this property further in our response to Q3.
  • Frameworks for Interpreting Guidance Our prediction-correction (PC) scheme also offers a different intepretation of the guidance-based methods, distinct from the DiffPIR's HQS framework. Inspired by numerical continuation [2], the PC scheme progresses along a prior-defined path while optimizing for high-likelihood regions, which enables a clean development of our EnKG method. In contrast, DiffPIR views the generic guidance-based methods through the len of a series of MAP estimation problems (Appendix A.1 in their paper).
• We have incorporated the reference [A] at line 103 and added the results of PnP-DM as requested. These details are further elaborated in our response to Q3.
• We hope this comparison clarifies the distinctions between EnKG and DiffPIR. Please let us know if further elaboration is needed.

Main text

• Lines 141-143: expanded the discussion of related work. (G5qZ)
• Lines 198-200: revised the discussion to focus on the practical setting of non-convex optimization. (non-convex). (Fiay)
• Lines 264-269: revised the discussion about the role of the covariance matrix. (Fiay)
• Lines 287-299: revise Proposition 1 for clarity and corrected typos. (Fiay)
• Table 2: added a column to report runtime for different algorithms. (Fiay)
• Lines 430-458: revised the explanation of evaluations metrics. (8WU4, Fiay, x3Ty)
• Figure 3: updated the caption for clarity. (Fiay)

Appendix

• Table 4: included the definition of weighted Euclidean inner product. (G5qZ)
• Lines 731-742: revised the statement of Lemma 1 for clarity. (Fiay)
• Lines 758-776: added an intermediate reasoning step in the proof for better readability. (Fiay)
• Lines 791-799: revise the last step for clarity. (G5qZ)
• Appendix A.8: added additional comparisons to gradient-based methods for reference. (8WU4)
• Appendix A.9: additional experiment on the dependence on the quality of pre-trained diffusion models. (x3Ty)