Generative diffusion learning for parametric partial differential equations

1. Problem Statement

The paper addresses the challenge of learning solution operators for parameter-dependent Partial Differential Equations (PDEs). In many scientific applications, solving a PDE $G(u; a) = 0$ for various input parameters $a$ (e.g., random fields, boundary conditions) is computationally expensive due to the need for repeated numerical solves.

Existing "Operator Learning" methods (such as DeepONets and Fourier Neural Operators - FNO) typically learn a deterministic mapping $S_\theta: a \to u$. These methods face two main limitations:

1. Data Sensitivity: They struggle when training data is corrupted by noise (numerical or observational), often treating noise as an error to be minimized rather than a distribution to be modeled.
2. Lack of Uncertainty Quantification: Deterministic models provide a single output for a given input, offering no inherent measure of confidence or predictive uncertainty.

The authors propose a framework to learn the solution operator as a probabilistic conditional distribution $u|a$, enabling robust learning from noisy data and providing uncertainty quantification.

2. Methodology: Probabilistic Diffusion Neural Operator (PDNO)

The core contribution is the Probabilistic Diffusion Neural Operator (PDNO), which reformulates operator learning as a conditional generative modeling task using Denoising Diffusion Probabilistic Models (DDPM).

A. Probabilistic Formulation

Instead of minimizing a deterministic loss (like $L_2$ error), the authors maximize the conditional data likelihood. They assume the observed output $u_\eta$ is generated by:
$$u_\eta = s_\theta(a) + \eta$$
where $\eta$ is random noise. The goal is to learn the conditional density $q_\theta(u|a)$.

• Noise-Free Case: The distribution collapses to a Dirac measure $\delta_{s_\theta(a)}$.
• Noisy Case: The distribution is non-degenerate, allowing the model to learn both the mean solution and the noise characteristics.

B. Conditional DDPM Framework

The method generalizes DDPM to a supervised setting where the input $a$ conditions the generation process.

1. Forward Process: A stochastic differential equation (SDE) gradually adds noise to the true solution $u$ conditioned on $a$, transforming the data distribution into a standard Gaussian.
   $$dU(t; a) = -\frac{1}{2}\alpha(t)U(t; a)dt + \sqrt{\alpha(t)}dW(t)$$
2. Reverse Process: A learned neural network (the score function $\Psi_\theta$) approximates the reverse SDE to transform Gaussian noise back into the solution distribution $u|a$.
   $$dV(t; a) = \left(\frac{1}{2}\alpha(T-t)V(t; a) + \alpha(T-t)\Psi_\theta(T-t, V(t; a), a)\right)dt + \sqrt{\alpha(T-t)}d\tilde{W}(t)$$
3. Training Objective: The model is trained using Denoising Score Matching (DSM). The neural network (typically a UNet) predicts the noise added at each time step $n$. The loss function minimizes the difference between the predicted noise and the actual noise:
   $$L_{DDSM}(\theta) = \mathbb{E} \left[ \| \epsilon_\theta(n, \sqrt{\gamma_n}u_0 + \sqrt{1-\gamma_n}\epsilon_n, a) - \epsilon_n \|^2 \right]$$

C. Functional Space Extension

The paper discusses extending PDNO to infinite-dimensional functional spaces (Hilbert spaces) to achieve mesh-invariance. While the current implementation uses finite-dimensional discretizations, the theoretical framework defines a score operator in Hilbert spaces, addressing the non-existence of Lebesgue measures in infinite dimensions by using conditional expectations (Ornstein-Uhlenbeck processes).

3. Key Contributions

1. Probabilistic Operator Learning: Introduces a framework where the solution operator is a family of conditional distributions rather than a deterministic map. This naturally handles noisy datasets.
2. Uncertainty Quantification: By sampling multiple times from the learned distribution $q_\theta(u|a)$, the method provides confidence intervals and quantifies predictive uncertainty (e.g., distinguishing between aleatoric noise and model uncertainty).
3. Noise Recovery: The model can recover the magnitude and distribution of additive noise in the training data, effectively "denoising" the solution during inference.
4. Robustness to Non-Smooth Solutions: Demonstrates superior performance over deterministic methods (like FNO) when learning solutions with high discontinuities or non-smooth features (e.g., advection equations with shocks).

4. Experimental Results

The authors evaluated PDNO against Fourier Neural Operators (FNO) on several benchmarks:

• 1D Elliptic Equation (Noisy Data):
  • PDNO successfully learned the mean solution and the standard deviation from data corrupted by Gaussian noise (up to 100% noise level).
  • The predicted confidence intervals matched theoretical bounds, whereas deterministic methods failed to characterize uncertainty.
• 2D Elliptic Equation:
  • In noise-free settings, PDNO achieved lower Mean Relative $L_2$ Error (MRLE) than FNO across various mesh sizes.
  • In noisy settings with Gaussian mixture noise, PDNO successfully recovered the multi-modal noise distribution (in simpler cases) and its first two moments.
• Burgers' Equation:
  • FNO outperformed PDNO here. The authors attribute this to the periodic boundary conditions and the smooth nature of the solution, which aligns well with FNO's spectral bias.
• Advection Equation (Highly Non-Smooth):
  • Key Finding: PDNO significantly outperformed FNO (MRLE of 0.0445 vs 0.1325 at $d_U=1024$).
  • PDNO captured fine details of the discontinuous solution (shocks) much better than FNO, which suffered from Gibbs phenomena and smoothing artifacts.
• Inference Speed:
  • PDNO is slower than FNO because it requires sequential sampling (N steps) rather than a single forward pass. However, the accuracy gain in non-smooth and noisy regimes justifies the cost in specific applications.

5. Significance and Conclusion

This work bridges the gap between generative modeling and scientific machine learning.

• Practical Impact: It offers a robust solution for real-world scenarios where data is inherently noisy or incomplete, providing not just a prediction but a measure of trust (uncertainty).
• Theoretical Advancement: It provides a rigorous probabilistic formulation for operator learning, generalizing standard $L_2$ loss minimization to likelihood maximization.
• Limitations & Future Work: The current implementation is mesh-dependent (requiring retraining for different resolutions) and computationally expensive during inference. Future directions include extending the model to infinite-dimensional functional spaces for mesh-invariance and using cascaded diffusion models or latent space diffusion to improve sampling efficiency.

In summary, PDNO represents a paradigm shift from learning "the answer" to learning "the distribution of answers," making it particularly valuable for uncertainty-aware scientific computing and inverse problems.