DiffusionPDE: Generative PDE-Solving Under Partial Observation

We introduce a general framework for solving partial differential equations (PDEs) using generative diffusion models. In particular, we focus on the scenarios where we do not have the full knowledge of the scene necessary to apply classical solvers. Most existing forward or inverse PDE approaches perform poorly when the observations on the data or the underlying coefficients are incomplete, which is a common assumption for real-world measurements. In this work, we propose DiffusionPDE that can simultaneously fill in the missing information and solve a PDE by modeling the joint distribution of the solution and coefficient spaces. We show that the learned generative priors lead to a versatile framework for accurately solving a wide range of PDEs under partial observation, significantly outperforming the state-of-the-art methods for both forward and inverse directions.

Paper

Similar papers

Peer review

Reviewer y14J5/10 · confidence 4/52024-07-08

Summary

This paper introduces diffusion methods to tackle the partially observed PDEs, named DiffusionPDE. By learning the joint distribution of solution and coefficient space, the proposed model can handle both forward and inverse problems. The authors experiment with diverse PDEs and settings to demonstrate the model's effectiveness.

Strengths

- This paper successfully utilizes the diffusion methods in solving PDEs, covering both forward and inverse problems. - The main text and supplementary materials provide diverse experiment settings, which can well support the model’s effectiveness on partial observations. - This paper is overall clear and well-written.

Weaknesses

1. The technical contribution is limited. From a technical view, this paper is an application of the diffusion model in PDE solving. There are also some previous methods that also use diffusion methods and leverage the PDE loss [1]. Thus, I think the technical novelty is limited. [1] A Physics-informed Diffusion Model for High-fidelity Flow Field Reconstruction, JCP 2023 2. Some powerful baselines are missing. - According to Figure 1, I think the base model of DiffusionPDE is U-Net. How about comparing it with a single U-Net? I think U-Net could be a powerful baseline. - There are also some latest models that are good at processing partially observed or irregularly placed PDEs, such as OFormer [1] and Transolver [2]. They should include them as baselines. [1] Transformer for Partial Differential Equations' Operator Learning, TMLR 2023 [2] Transolver: A Fast Transformer Solver for PDEs on General Geometries, ICML 2024 3. Model efficiency comparisons are needed, including GPU memory and running time. 4. I think the proposed model cannot predict the future evolution of a time-dependent PDE. Current tasks are all about “reconstruction” or “imputation”.

Questions

1. Figure 4, do both forward and inverse tasks use the same diffusion model? Or do we need to train two models for these two different tasks? 2. I think the base model is U-Net. So how does DiffusionPDE handle the spatially scattered partial observations? Is the input still in the regular grid, but only the sampled locations have ground-truth values?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

I appreciate that they have discussed the limitations. But I think the mentioned issues about efficiency and limitations on temporal modeling are not trivial. More discussions are expected.

Authorsrebuttal2024-08-13

Clarification on DiffusionPDE’s Robust Performance Using Continuous Coordinates and Bilinear Interpolation

We greatly appreciate your response and recognition of our method. We would like to clarify that DiffusionPDE can utilize continuous coordinates with bilinear interpolation in our prediction space to obtain predicted values for points that are not on the grid. By doing so, for example, the forward problem of the non-bounded Navier-Stokes equation results in a relative error of approximately 6.0%, which is comparable to the scenario without interpolation (approximately 6.9%) and even smaller as more points are considered during the calculation of observation loss, while still remaining significantly lower than OFormer (approximately 16.2%). Overall, the diffusion model is powerful, enabling DiffusionPDE to consistently outperform OFormer. We are happy to mention all these discussions in the revised paper.

Reviewer AaET6/10 · confidence 4/52024-07-08

Summary

The paper proposes to solve PDEs given only sparse measurements by jointly modeling the solution and coefficient space (e.g. the initial conditions) using a diffusion model. By applying diffusion posterior sampling (DPS) the authors obtain samples that are consistent with the sparse measurements and the underlying PDE equations. Several experiments show superior performance of the method compared to standard baselines such as PINNs and FNOs.

Strengths

- Solving PDEs under partial observation is an important problem in real-world applications - The proposed method is technically sound and improves upon existing baseline methods (PINN, FNO) that do not work well for sparse measurements - Leveraging a pretrained diffusion model as a generative prior to model the joint distribution of solution and coefficient space is a good idea - The presentation of the method is clear and supported by concise algorithms and equations. The paper is well written overall - Experiments consider standard baseline methods for PDEs and cover a sufficient range of different dynamics ---- Post-rebuttal: the authors have addressed quite a few of the initial concerns, and while some concerns (e.g. about the magnitude of the contributions remain), I'd be happy to support an accept. I've raised my score accordingly.

Weaknesses

- The main weakness of the method is the limited novelty. Both sparse measurements and physics-based losses have been considered together with diffusion models, see e.g. Shu et al. (2023). So it seems to me that the main technical novelty is to apply diffusion models to model the joint distribution of two simulation states at different points in time and apply DPS during inference for consistency with the sparse measurements and PDE constraints. - The experiments do not take into account any stochasticity or uncertainty. In principle, DPS will give a distribution of solutions, which is not the case for the other baseline methods, but this is not explored further in the paper. - Since the joint distribution models two states at time 0 and time T (for all experiments except Burgers' equation) and $0 \ll T$, the authors need to simplify the PDE loss $\mathcal{L}_{pde}$ to drop any time derivatives. This is a serious limitation. - It is not clear if DPS works better than classifier-free guidance, as used e.g. in Shu et al. (2023), or other methods for solving inverse problems with diffusion models. - DPS requires a lot of compute during inference for calculating $\mathcal{L}_{pde}$. For a fair comparison, it would be important to show the number of parameters, training time and inference time for all methods.

Questions

- Algorithm 1 shows an adaption of DPS to EDM (Karras et al. 2022). Is this adaptation novel? Can the authors give some intuition why they apply the DPS losses in line 12 and 13 to the 2nd order correction (line 8) and not apply any trapezoidal rules in this case? - Are sparse measurements located on a grid that matches the resolution of the diffusion model or do they have continuous coordinates? In the second case, how are they interpolated to match the data resolution of the diffusion model? Does that make classifier-free guidance difficult to apply? - As noted in the weaknesses: why not use classifier-free guidance? I would like to see a discussion of different methods for inverse problems and diffusion models that can be used here instead of DPS and what are the advantages of using DPS. Reconstructing the solution/coefficient space from sparse measurements alone is a linear inverse problem with a number of different methods that can be used (e.g. Denoising Diffusion Restoration Models; Kawar et al. 2022, among many others) which oftentimes have much nicer theoretical guarantees/higher quality reconstructions and faster sampling speed. When considering these methods, is adding the PDE loss $\mathcal{L}_{pde}$ and thus making the problem a non-linear inverse problem really beneficial? - Likewise, as mentioned above: what are the parameter counts and runtimes of the method and the baselines?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have mentioned slow sampling speed as a limitation in the conclusion, but I think an extended discussion of this would be important to include.

Reviewer 4PC94/10 · confidence 3/52024-07-10

Summary

The work uses a guided diffusion process to solve the PDE forward and inverse problems with partial observations. Instead of learning the parameter-to-solution map ($a\rightarrow u$) as in Neural Operators, the method learns the diffusion process on the joint distribution $(a,u)$, and use guided diffusion for inference under sparse observations. Compared with several baseline method, the proposed method shows improved performance for solving forward and inverse problem with sparse observations.

Strengths

The work uses a guided diffusion process to solve the PDE forward and inverse problems with partial observations.The authors compare with several baseline methods. The idea is clearly presented, and might be useful for the community.

Weaknesses

The paper presents an interesting approach to solving PDE forward or inverse problems with sparse observations, which is an appealing concept given the minimal data requirement. However, this approach raises some concerns about the well-posedness of the problem. For example, in forward problems where sparse observations of the parameter $a(x_i)$ are available, there are infinitely many ways to interpolate $a$ and solve the PDE to obtain $u$. They are all valid solutions that satisfy both the PDE and the observations. This suggests that the method's ability to achieve good recovery might heavily rely on the strong regularization imposed by the training dataset, potentially limiting its practical utility as it may only favor solutions resembling those in the training set. Additionally, in Appendix C, Table 2, the weightings for observation and PDE loss are significantly higher (by two to six orders of magnitude) than those for $\nabla_x \log(p(x))$ as described in Equation 8, which might indicate a predominance of data fitting over the diffusion process. It would be beneficial if the authors could provide more guidance on how these weights were chosen and discuss the implications of using smaller weights. Understanding the rationale behind these choices could help clarify the model's dependency on these parameters and their impact on the solution's behavior.

Questions

(1) The results from the baseline methods (PINO, DeepONet, PINNs, and FNO) are so bad. The claim that these methods are "not easily extendable" invites further scrutiny: (a) All the baseline models are supposed to represent smooth functions. However, in Figure 4, 7, 8, they look discontinuous at the training points. An explanation of how these models were trained and how inference was conducted could clarify why these discrepancies appear. (b) Taking PINNs as an example in the Darcy flow problem. Let $\hat{a}(x)$ and $\hat{u}(x)$ be the (potentially noisy) observation at $x$. We can represent $a(x)$ by a neural network, $a_V(x)$ (use neural net for convenience, could be other representations), and the PDE solution $u(x)$ by a neural net $u_W(x)$, where $V$ and $W$ are the weights of the neural network. We can solve the following optimization problem: $$\min_{V,W} \sum_{x\in T_d} (u_W(x) - \hat{u}(x))^2 + (a_V(x) - \hat{a}(x))^2 + \sum_{x \in T_r} (\nabla \cdot (a_V(x) \nabla u_W(x)) - q(x))^2$$ where $T_d$ is the observational point, and $T_r$ is the residual points, which does not need to be the same as $T_d$. (c) Similar, for a trained neural operator parametrized by W, $G_W[a] = u$. We can solve the following optimization problem: $$\min_V \sum_{x\in T_d} (a_V(x) - \hat{a}(x))^2 + (G\[a_V\](x)-\hat{u}(x))^2$$ where $G[a_V]$ is the solution of the PDE with parameter $a_V$. It seems that all the baseline methods can be used for forward and inverse problems with sparse observation. It is unclear why the proposed method would offer superior performance compared with the baselines. In contexts where full observation are available, as shown in Table 4, one might intuitively expect methods like PINNs—which utilize residual losses to ensure adherence to PDE constraints—or Neural Operators—which establish a direct parameter-to-solution mapping—to deliver more accurate results compared to a method that relies on a diffusion process. This leads to a critical inquiry on why diffusion process gives better accuracy for PDE problems. (2) How is the PDE loss computed? Is it by finite difference on a regular grid? Detailing this in the main text could help readers assess the accuracy and applicability of the PDE loss in different scenarios.

Rating

4

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The author mention several limitation in the conclusion.

Reviewer nBsW8/10 · confidence 5/52024-07-10

Summary

The paper uses score based generative diffusion models to find the forward and backwards solution of a set of PDEs given partial observations of the solution and/or incomplete knowledge of the coefficients. The method performs well, and outperforms other ML methods such as FNO, as well as 'standard' FE type methods, for a range of standard test problems. The method reconstructs This is a novel approach, which delivers good performance, with low errors at a competitive speed. Extensive tests are given, with careful analysis of the results.

Strengths

The use of score based generative methods in this context, where both the solution and the parameter estimates are updated, is novel. The method is clearly effective for the problems considered and should have good applications to real world examples. Extensive tests on a series of standard test problems show that the errors of the method are much lower than other ML based methods such as FNO.

Weaknesses

This paper suffers as do many similar papers from a limited range of examples. It concentrates on the usual examples of PDEs such as NS and Bergers, and in both cases of these it looks at problems with quite moderate viscosity, which are realtively east to solve. This is more or less inevitable for such a short paper as this, especially as comparisons are needed with other method. But I would have liked to have seen more novel examples than the usual ones. This is not really a criticism of this paper, but is something to consider for future work. It would be imporoved by a fairer comparison with other methods which work with incomplete data and measurements. A clear exanple of this being the data assimilation widely used in physical modelling for just this range of problems. These should be descibed somewhere in the introduction and in Section 2. (Although of course these latter methods are slow in comparison.) The method is also limited (see later) to looking at certain slices of the solution.

Questions

1. How does this method compare with a data assimilation approach 2. How easy would it be to extend the method to full time intervals 3. How easy is it to extend the method to higher dimensions 4. Have the authors tried out the method on more challenging PDE examples. 5. Also consider tests on NS and Bergers' eqn with much smaller viscosity.

Rating

8

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

The model as described only looks at slices of solutions of 2D problems. This has been clearly identified by the authors. In this sense it is vrather limited when compared to other ML based approaches, and of course traditional FE based methods. I am pleased that this is recognised and that the authors plan to address this. The DiffusionPDE method will only be truly competitive when this is done, but this paper is a good step in this direction.

Reviewer y14J2024-08-13

I would like to thank the authors' response. After rebuttal, I am clear about the paper's technical contribution. There are still some questions remain. I think the extension of DiffusionPDE to a prediction model is non-trivial. The current framework cannot support this. Also, as the author mentioned, they treat the sparse observation in uniform grids, which is also not a generalizable design since the partial observations can be placed in any position of the domain. Some techniques, such as OFormer, may be more applicable than U-Net. After rebuttal, I think the application scope of DiffusionPDE is limited. The authors should discuss more about the above-mentioned limitations in the revised paper. However, I appreciate the design in the joint modeling of forward-inverse problems. Thus, I improve my score to 5.

Reviewer 4PC92024-08-13

Thanks for the the hard work and for conducting additional experiments. I appreciate the comprehensive benchmarking and the detailed responses to the reviewers' comments. > we know of no existing method that performs well under sparse observations without assumptions about PDE coefficients to the best of our knowledge. I agree. But I'm not asking for out-of-sample distribution. My questions is on the well-posedness of the problem. Using the Darcy flow as an example. Suppose $a$ is sampled from some Gaussian random field with certain correlation structure. If we only have partial observation of $\{x_i, a(x_i)\}$, for i = 1,...,N, and N is relatively small. Suppose the ground truth is $a_{GT}$. However, there could be $a_1$, ..., $a_M$ that are all consistent with the observation. That is, $a_j(x_i) = a_{GT}(x_i)$ for i = 1,...,N. and j = 1,...,M. In this case, DiffusePDE might recover $a_{GT}$ (hypothetically), while FNO, DeepOnet, PINN, etc, might recover $a_1$, ..., $a_M$, which will certainly have error with respect to $a_GT$, but they are all consistent with the observation data and the PDE. Therefore, the low accuracy of other baseline methods does not necessarily mean they are worse. They might be different but still valid solutions to the problem. As an even more straightforward approach in this case, we can sample the Gaussian random field conditional on the observation, and solve the PDE numerically, this would give infinite number of valid solutions to the problem. > During inference, we apply a mask to set non-observed input values to zero. Why do we mask the non-observed input values to zero? It seems a more reasonable approach is to exclude those in the loss function. > For partial observations, PINNs [1] does not learn any knowledge about the PDE, making it difficult to inpaint the missing parts My understanding is that PINN, or any method employing PDE loss, should be able to inpaint missing parts in ways that remain consistent with the PDE. As the PDE loss can be computed and minimized without any data. > Neural operators [2] learn to map the entire coefficient space to the entire solution space and hence have difficulty taking partial observations as input This statement appears to conflict with the typical characterization of neural operators (FNO or DeepONet) as being "resolution invariant" -- the neural operator can be evaluated at any point, even those not in the training points. Could you provide more details on the difficulties encountered? > For full observations, PINNs require iterations of optimization to converge to the solution, which can be vulnerable to local minimal or failure to convergence, leading to higher error. This seems to suggest that PINN fail to train in these cases, which invite questions on how is PINN trained, As the problems (Darcy flow, Poisson, etc) are standard examples in PINN literatures [1,2,3,4]. > Iterative generative models, such as diffusion models, tend to beat feedforward models like GANs. I'm not concerned about GANs. To my understanding, in your approach, the joint distribution between the parameter $a$ and the solution $u$ is modeled. That is, for any particular $a$, there is a conditional distribution of $u$. However, for PDEs, the solution $u$ should be uniquely determined by the parameter $a$. And neural operators aim to learn this mapping. It seems somewhat counter-intuitive that learning a joint distribution could be more effective than learning this mapping directly. As a summary, due to ongoing concerns regarding the problem formulation, implementation of some of the baseline method, and an lack of understanding of advantages of the diffusion-based approach, I will maintain my current score for this review. [1] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations,” [2] DeepXDE(https://deepxde.readthedocs.io/en/latest/demos/pinn_forward.html) [3] Q. He and A. M. Tartakovsky, “Physics-Informed Neural Network Method for Forward and Backward Advection-Dispersion Equations,” [4] S. Wang, S. Sankaran, H. Wang, and P. Perdikaris, “An Expert’s Guide to Training Physics-informed Neural Networks,”

Authorsrebuttal2024-08-13

We greatly appreciate your response. We would like to clarify a few points: > "the well-posedness of the problem" We agree that there are infinitely many valid solutions to the partial observations. However, we assume that there are statistical patterns in the coefficient and solution spaces, which means that each valid solution has its own 'probability' of actually being the solution. We aim to learn this probability distribution with our diffusion model, which excels at sampling a highly likely full state given partial observations, as reflected by the much lower error compared to all prior methods suggested by ourselves and other reviewers, including a simple baseline (suggested by reviewer nBsW) of completing the coefficient space using an RBF kernel. Such problems of statistically inferring unobserved data have been widely studied across domains, such as image inpainting in computer vision and matrix completion in machine learning, which are deemed well-posed problems of maximizing probability. > "Why do we mask the non-observed input values to zero?" For our method and other baselines that require optimization during inference, we indeed exclude non-observed values. For baseline methods that expect a complete input space, e.g., FNO, we have to fill in those values. We have tried filling in the missing values with constant zero values (default) or using RBF kernel or nearest neighbor interpolation, all resulting in lower performance than DiffusionPDE. > "As the PDE loss can be computed and minimized without any data. / This seems to suggest that PINN failed to train in these cases..." PINN is trained well according to the original paper, and the reviewer is correct in noting that PINN can handle partial observations and automatically complete the missing data. However, it’s generally understood that PINN tends to have a test error between 10% and 30%, even in fully observed situations, because the training loss never perfectly converges to 0. This high error is one of the reasons why there is growing interest in neural operators [1]. > "This statement appears to conflict with the typical characterization of neural operators (FNO or DeepONet) as being 'resolution invariant'... Could you provide more details on the difficulties encountered?" FNO can solve problems across different resolutions, and the reviewer is correct in noting that the output of neural operators can be evaluated at any point. However, we cannot assume that the input to neural operators can be sparse like ours. In fact, neural operators require a complete continuous grid as input, whereas we only have very sparse observation points (approximately 3%), which are highly discrete and do not meet this requirement. As a result, the model struggles to learn higher-frequency features. To the best of our knowledge, neural operators have not been shown to be effective on highly sparse inputs. > "It seems somewhat counter-intuitive that learning a joint distribution could be more effective than learning this mapping directly." The neural operators that learn the mapping directly indeed perform comparably to DiffusionPDE given "full" observation of $a$. However, when the observation of $a$ is not complete, the solution $u$ is not uniquely determined. Hence, we need to sample from the conditional distribution of $u$, a task in which DiffusionPDE outperforms neural operators and their variants suggested by the reviewers. [1] Physics-Informed Neural Operator for Learning Partial Differential Equations. Li et. al. ACM/JMS Journal of Data Science 2024.

Reviewer AaET2024-08-13

Rebuttal

Thank you for the updates and the interesting new results in the PDF. It's definitely an interesting direction, and I'm open to supporting acceptance by raising my score. I still find it difficult to strongly argue for acceptance given the limited scope of the technical contributions, though.

Authorsrebuttal2024-08-13

We greatly appreciate your response. We would appreciate it if you could let us know of any specific questions we can address to help raise the score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC