Efficient Physics-Constrained Diffusion Models for Solving Inverse Problems

Authors: Seungjun Lee, Shinjae Yoo
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: physics-constraints inverse problem, diffusion model, PDE, generative modeling

Abstract:
Solving inverse problems in scientific and engineering domains often involves complex, nonlinear forward physics and ill-posed conditions. 
Recent advancements in diffusion model have shown promise for general inverse problems, yet their application to scientific domains remains less explored and is hindered by the complexity and high non-linearity of physics constraints. We present a physics-constrained diffusion model (PCDM) designed to solve inverse problems in scientific and engineering domains by efficiently integrating pre-trained diffusion models and physics-constrained objectives.
We leverage accelerated diffusion sampling to enable a practical generation process while strictly adhering to physics constraints by solving optimization problems at each timestep. By decoupling the likelihood optimization from the reverse diffusion steps, we ensure that the solutions remain physically consistent, even when employing fewer sampling steps.
We validate our method on a wide range of challenging physics-constrained inverse problems, including data assimilation, topology optimization, and full-waveform inversion. Experimental results show that our approach significantly outperforms existing methods in efficiency and precision, making it practical for real-world applications.

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

Reviews and discussion:
<thread>
creation date:
1729721688678

thread_type:
official_review

<thread_object>
creator: reviewer_liuf

<content>
summary:
The paper propose a method for the solution of an inverse problem where the forward problem is a solution of some physical simulating. 
The claim is that the result of the algorithm produces solutions that do not only honour the prior but also obey the physics.

strengths:
The results look interesting. The paper may be improved so the results will make sense.

weaknesses:
To be honest I could not understand the paper, even though I have been working in this field for many years.  The authors invented new jargon, "physics constrained" which means, what exactly? What is the constraint they are fulfilling? On which variables? How do you deal with the constraints? Lagrange multipliers? elimination? penalty?
There is a huge branch of inverse problems that treat them as PDE constrained optimization. Clearly, this escaped from the authors. There is a large number of papers that introduce constraints into inverse problems (e.g 0 \le x) but clearly this is not one of these examples. The authors should try to rewrite the paper and be a bit more precise about what they do,

Similarly, in section 3, the equations flow smoothly and I can easily understand how to get from (7) and the way to (10). 
Then you switch to section 3.2 and I cannot see how (11) and on is related to the previous section. 

Finally in ADMM (eq 12) there is another term for Lagrange multiplier that you are missing. The solution of your problem is different than the original problem.

questions:
1. What is physics constrained, please define mathematically

2. How to get from the Langevin dynamics of (8-9) to your optimization problem (11-12)

3. Why and how are the two related

4. Why don't you use Lagrange multipliers for (12), add a term p^T(z-x)

5. Can you clarify the overall algorithm?

6. How do you ensure that you fit the data to some given tolerance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730352926517

thread_type:
official_review

<thread_object>
creator: reviewer_dxjv

<content>
summary:
This paper proposes PCDM (physics-constrained diffusion model), an inverse problem solver that leverages diffusion model as plug-and-play prior. PCDM uses the idea of variable splitting and proposes to solve the underlying optimization problem with implicit diffusion model regularization. The authors demonstrate its application in full-waveform inversion, data assimilation, and topology optimization.

strengths:
1. Applying plug-and-play diffusion model methods to physics-constrained inverse problems is relatively new to the diffusion model community. 
2. The paper is generally easy to follow.

weaknesses:
1. The proposed PCDM appears to be mathematically equivalent to a special case of the algorithm in Li et al. [1] (specifically, the case using Tweedie's formula). The claim of algorithmic novelty is questionable (line 100). 
2. The "physics-constrained" aspect really comes from the inverse problem itself instead of the novel algorithmic design. Most existing gradient-based plug-and-play diffusion model methods can incorporate that physics loss such as DiffPIR [2], DPS,DAPS [3], RED-diff [4], [5]. These methods are not compared or  discussed in the paper. 
3. The experimental comparison excludes many recent and relevant algorithms. For example DiffPIR [2] and DAPS [3], RED-diff [4].  
4. Reproducibility concerns: important experimental and implementation details are insufficiently documented. See more concrete questions in the next section.
5. There is a lack of ablation studies on important algorithm design parameters, such as the number of likelihood steps per iteration, the optimization threshold $t_s$, and sensitivity to the optimizer configurations. 


[1] : Li, Xiang, et al. "Decoupled data consistency with diffusion purification for image restoration." _arXiv preprint arXiv:2403.06054_ (2024).
[2] : Zhu, Yuanzhi, et al. "Denoising Diffusion Models for Plug-and-Play Image Restoration." _arXiv preprint arXiv:2305.08995_ (2023).
[3] : Zhang, Bingliang, et al. "Improving diffusion inverse problem solving with decoupled noise annealing." _arXiv preprint arXiv:2407.01521_ (2024).
[4] : Mardani, Morteza, et al. "A Variational Perspective on Solving Inverse Problems with Diffusion Models." _The Twelfth International Conference on Learning Representations_.
[5] : Peng, Xinyu, et al. "Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance." _Forty-first International Conference on Machine Learning_. 2024.

questions:
1. I'm a bit surprised at how well the Opt w/o diff baseline can recover the large structure of the ground truth, as shown in Figure 2 and Table 1. This contrasts with traditional FWI literature findings [1] and my own experimental validation on OpenFWI dataset. I'm curious how the authors implement the FWI problem and the corresponding baselines. More specifically, 
	1. What is exactly the Opt w/o diff baseline in Table 1? Is that the Adam optimizer? What initialization strategy was employed? What are the specific hyperparameters used to report the results? 
	3. Why are the residuals of InversionNet and VelocityGAN omitted from Table 1?  
	4. Given that OpenFWI paper does not provide the gradient implementation of the forward model, how did the authors implement the gradient? 
2. What are the hyperparameter selection criteria across compared methods? 
3. Is there any supplementary material or code to facilitate the reproducibility?

[1] : Virieux, Jean, and Stéphane Operto. "An overview of full-waveform inversion in exploration geophysics." _Geophysics_ 74.6 (2009): WCC1-WCC26.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730380756025

thread_type:
official_review

<thread_object>
creator: reviewer_rymp

<content>
summary:
The paper proposes doing MAP estimation using a diffusion plug and play prior (PnP). Namely, the paper aims at solving
$$argmin \|y - \mathcal{A}(x)\| + \lambda \mathcal{R}(x).$$

To do so, it follows the traditional PnP route by using ADMM to split this into solving two proximal problems:

$$ z_{i+1} = argmin_{z} \mathcal{L}_\mu(z, x_i) $$

$$ x_{i+1} = argmin_{x} \mathcal{L}_\mu(z_{i+1}, x) $$

Finally, it replaces the prior proximal problem by a forward backward (with one step) sampling, namely equation (15).
It then evaluates the approach in non-linear problems coming from physics.

strengths:
The paper proposes an adapted method to solving several physics problems where adapting to a physical constraint is cast as having high likelihood. The numerical applications are relevant.

weaknesses:
My main concern is the novelty aspect of the paper. Indeed, several papers have investigated the applications of pretrained diffusion generative models as PnP priors. In particular, Algorithm 1 of [1] is essentially the same as the one proposed in this paper. Unless I'm mistaken, this makes the only novelty in this paper w.r.t. [1] to be the physical applications, which are indeed interesting. But I do not reckon it is worth being accepted to ICLR.

Furthermore, even if the proposed algorithm is conceptually different, it is still part of the broad Plug and Play family and I would expect at least a comparison with [1] or any other Plug and Play with diffusion paper.


[1]  Denoising Diffusion Models for Plug-and-Play Image Restoration
Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, Luc Van Gool; Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, 2023, pp. 1219-1229 

[2] Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction.
 Xu, Xingyu, and Yuejie Chi. arXiv preprint arXiv:2403.17042 (2024).

[3]Graikos, Alexandros, et al. "Diffusion models as plug-and-play priors." Advances in Neural Information Processing Systems 35 (2022): 14715-14728.

[4] F. Coeurdoux, N. Dobigeon and P. Chainais, "Plug-and-Play Split Gibbs Sampler: Embedding Deep Generative Priors in Bayesian Inference," in IEEE Transactions on Image Processing, vol. 33, pp. 3496-3507, 2024, doi: 10.1109/TIP.2024.3404338.

[5] Wu, Zihui, et al. "Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors." arXiv preprint arXiv:2405.18782 (2024).

[6] Wang, Hengkang, et al. "DMPlug: A Plug-in Method for Solving Inverse Problems with Diffusion Models." arXiv preprint arXiv:2405.16749 (2024).

questions:
For the major point, see weaknesses.

Minor questions and remarks.

* Is the left term in eq(6) $x_{t-1}$ ? Otherwise it is not a sampling process, as it does not evolve through time.
* What is $ \hat{\epsilon}_{t}$ in equation (15) ? 
Is it equation (7) with $x_{t_k}$?
* Equation (15) mixes indexes between $t_k$ and $t$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730387104735

thread_type:
official_review

<thread_object>
creator: reviewer_ucwn

<content>
summary:
This paper proposes a variable splitting method for solving physics-constrained inverse problems with diffusion model priors. Throughout the reverse diffusion process, the method alternates between two optimization problems: one to update the noisy estimated image with the diffusion model as a regularizer, and one to enforce data/physics constraints. The authors present experiments on full-waveform inversion, data assimilation, and topology optimization, showing quantitative and qualitative improvement upon baselines in all three applications.

strengths:
* The proposed method is simple and intuitive. Although the methodology lacks technical novelty (see weaknesses), it’s helpful to see that such a simple extension of DPS and other plug-and-play diffusion-based inverse solvers may already go a long way in handling physics inverse problems.
* Validation is done on three very different tasks, and task-specific baselines and the DPS baseline are compared against for each task.

weaknesses:
* The technical contribution is marginal. The idea of variable splitting for diffusion-based inverse solving is not new (Equation 16 is similar to the proximal optimization step in Equation 8 of Song and Shen et al. 2022). The main difference in this work is that there may not be a closed-form solution to Equation 16, so iterative gradient-based optimization is used at each likelihood step.
* In Tables 1 and 2, the smaller number of reverse steps used by PCDM is touted. This is a little misleading, as the Table 1 caption says that PCDM involves 1000 likelihood iterations in addition to 200 reverse steps. For expensive forward models, it may be the case that these 1000 likelihood iterations are very costly. Also, a clarifying question: does one likelihood iteration mean an entire optimization round of Equation 16, or do the 1000 likelihood iterations account for all the gradient steps needed to solve Equation 16 throughout the algorithm?
* I have concerns about how fair the comparison to DPS is. In Figure 3(b), the DPS reconstruction clearly doesn’t match the visual statistics of Kolmogorov flow. I would expect DPS to at least produce something that appears visually plausible even if it doesn’t agree with the physical model. For example, in Figure 4 of SDA (Rozet and Louppe 2023) and Figure 5 of Feng et al. 2024, the DPS reconstructions at least look qualitatively reasonable. I would also be curious how hyperparameters for DPS were chosen.

---
References:

Song and Shen et al. “Solving Inverse Problems in Medical Imaging with Score-Based Generative Models.” ICLR 2022.

Rozet and Loupe. “Score-based Data Assimilation.” NeurIPS 2023.

Feng et al. “Neural Approximate Mirror Maps for Constrained Diffusion Models.” arXiv 2024.

questions:
* Please comment on how hyperparameters for baselines, including DPS, were chosen. I recommend making an appendix to include such details.
* Often it makes more sense to think of physics constraints as priors (i.e., checking whether a solution satisfies a physical model doesn’t involve the observed measurements). Does it make sense in that case to move the physics-consistency term into Equation 12 as an additional regularizer?
* It’s surprising that “Opt w/o diff” in Table 1 has the worst data residual, given that it only optimizes the likelihood term. The authors suggest that this is because it struggles with local minima, but I was under the impression that adding a diffusion regularizer would only complicate the optimization landscape. I would appreciate comments from the authors on why they believe “opt w/o diff” struggles to fit the data and whether they observed the same trend with the other tasks (why wasn’t opt w/o diff included as a baseline for the other tasks?).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680525157

thread_type:
official_review

<thread_object>
creator: reviewer_lnsc

<content>
summary:
The authors address inverse problems by leveraging diffusion models as priors.
They formulate the problem as minimizing a composite objective function comprising a likelihood term, which enforces physics constraints, and a regularization term defined by a diffusion model.
As the resulting problem is difficult to solve directly, the authors utilizes a variable splitting scheme that alternates between minimization over the regularizer and the likelihood.
The regularization step is handled through a backward diffusion step, while the likelihood step is performed by minimizing and L2-regularized inverse problem.
The authors validate their approach on a set of three problems.

strengths:
Solve inverse problems that arise in physics-constrained setups using a variable splitting scheme.

weaknesses:
**Insufficient coverage of the related work**

The authors provide a high-level overview of two lines of research: end-to-end supervised approaches and unsupervised approaches.
While, there is a wealth of methods in inverse problems with diffusion models priors, few are mentioned.
Notably, related works that leverage variable splitting schemes, also known as Split Gibbs sampling, are not discussed; for reference, see [1, 2, 3] and the corresponding Related Work sections.

**Methodological ambiguities**

Section 3.3 introduces the regularization step without a clear justification for its formulation. Specifically, _why it has this form?_.
Furthermore, the method employs two regularization hyperparameters, $\lambda$ and $\mu$, yet only $\mu$ appears in the update equations.
Besides, the regularization step is independent of these hyperparameters.

**Lack of implementation details**

- The paper does not address the sensitivity of the method to its hyperparameters, namely the early stopping criterion and the timing of triggering the optimization (the parameter $t_s$ in line 256).
- The experimental section lacks specific implementation details, such as the hyperparameters for DPS and SDA; details regarding the used pre-trained diffusion models.
- The reported results raises some concerns In Table 1, DPS performance appears almost identical to the method that omits the prior (Opt w/o diff), which warrants further clarification as inverse problems are severely ill-posed hence pure optimization often yields an inconsistent solutions

---
.. [1] Zhu, Yuanzhi, et al. "Denoising diffusion models for plug-and-play image restoration." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

.. [2] Wu, Zihui, et al. "Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors." arXiv preprint arXiv:2405.18782 (2024).

.. [3] Xu, Xingyu, and Yuejie Chi. "Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction." arXiv preprint arXiv:2403.17042 (2024).
 
.. [4] Rozet, François, and Gilles Louppe. "Score-based data assimilation." Advances in Neural Information Processing Systems 36 (2023): 40521-40541.

questions:
**Specific questions**

In the experiments, the formulation of the inverse problem in Experiments 4.1 and 4.3 is unclear, namely
 
- Experiment 4.1: is the operator $A$ a discretization of the d’Alembert operator? Additionally, is $s(r,t)$ provided within the dataset?
 - Experiment 4.3: Given that the problem is defined as a constrained optimization, how does the operator $A$ transform $x$ to yield the observation $y$?

Why was SDA excluded from Experiments 4.1 and 4.3? Although originally developed for data assimilation, it remains applicable as an inverse problem method.
Similarly, why was DPS omitted from Experiment 4.3?


**Broader questions**

- Could this method be applied to inverse problems in image restoration, and how would it compare to existing algorithms in the literature?
- The paper addresses problems of moderate dimensionality, approximately $5000$; have the authors considered Sequential Monte Carlo methods [1, 2, 3], which offer stronger theoretical guarantees?
Given this dimensionality, propagating multiple particles in parallel is feasible and would overcome mode-collapse.

---
.. [1] Dou, Zehao, and Yang Song. "Diffusion posterior sampling for linear inverse problem solving: A filtering perspective." The Twelfth International Conference on Learning Representations. 2024.

.. [2] Cardoso, Gabriel, Janati Yazid,, Sylvain Le Corff, and Eric Moulines. "Monte Carlo guided Denoising Diffusion models for Bayesian linear inverse problems." The Twelfth International Conference on Learning Representations. 2023.

.. [3] Wu, Luhuan, et al. "Practical and asymptotically exact conditional sampling in diffusion models." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732508793130

thread_type:
official_comment

<thread_object>
creator: area_chair_jvki

<content>
title:
Authors' Rebuttal

comment:
Dear Authors,

As the author-reviewer discussion period is approaching its end, I strongly encourage you to read the reviews and engage with the reviewers to ensure the message of your paper has been appropriately conveyed and any outstanding questions have been resolved.

This is a crucial step, as it ensures that both reviewers and authors are on the same page regarding the paper's strengths and areas for improvement.

Thank you again for your submission.

Best regards,

AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1732645471355

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We would like to appreciate all the reviewers for their constructive comments which have led to the revision, and we believe without a doubt have improved the quality of the manuscript. The revised manuscript has been uploaded. We have addressed the individual reviewer's comments below respectively, and here the summary of the major changes is as follows:

1.	We added ablation studies in Figure 5 and Section 4.4 in page 10. 

2.	We added details of the architectures and training procedures, and implementation details for our methods in Appendix A.1 and 2. 

3.	We have provided a brief explanation of the state-of-the-art methods and discussed their difference from our approach in Appendix A.3. Additionally, we conducted additional experiments comparing these state-of-the-art algorithms with our method in Table 5, and Figure 6, 7 in Appendix B.

4. We added the overall algorithm of our method in Algorithm 1 (Appendix A.3).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces Physics-Constrained Diffusion Model (PCDM), for solving inverse problems in physics by leveraging diffusion models as priors. PCDM employs a variable splitting technique, similar to ADMM and plug-and-play methods, to minimize a composite objective function that balances a likelihood term (enforcing physical constraints) with a regularization term defined by a diffusion model. This is achieved by alternating between a step that updates the solution using the diffusion model as a regularizer and a step that enforces data and physics constraints. The authors demonstrate PCDM's effectiveness on three applications: full-waveform inversion, data assimilation, and topology optimization, showing improved results compared to baseline methods. The claimed contribution lies in the integration of diffusion model priors with physical constraints to achieve solutions that are both realistic and physically consistent.

All the reviewers agree that the novelty of this approach is limited, and that the authors seem oblivious to the large amount of literature in the topic. As such the benchmarks are not really useful as they are not compared with state-of-the-art related (or in fact very similar) methods. As such I recommend rejection.