Latent representation learning based model correction and uncertainty quantification for PDEs

1. Problem Statement

The paper addresses the challenge of solving Partial Differential Equations (PDEs) when the governing physics are misspecified due to simplified assumptions, missing physics, or limited observations.

• The Core Issue: In Physics-Informed Neural Networks (PINNs), if the differential operator is incorrect (structural misspecification), standard training minimizes a loss function that forces the solution to satisfy the wrong physics. This leads to biased predictions and miscalibrated uncertainty estimates.
• Limitations of Existing Methods:
  • Parameter Space Correction: Traditional methods often correct models by tuning parameters in parameter space. Uncertainty quantification (UQ) here typically relies on sampling (e.g., Bayesian inference, Ensembles), which is computationally expensive and prohibitive for large-scale PDEs.
  • Additive Correction: While some methods add a discrepancy term, they often treat solution uncertainty and correction uncertainty separately, leading to inconsistent uncertainty propagation.
  • Identifiability: Correcting operators from noisy, sparse data is an ill-posed problem; distinct corrected models may fit data equally well, creating ambiguity that requires robust UQ.

2. Methodology: LVM-GP with Model Correction

The authors propose a Latent Representation Learning framework that extends their previous LVM-GP (Latent Variable Model - Gaussian Process) solver. The core innovation is a shared, confidence-aware stochastic latent representation that drives both the solution and the correction simultaneously.

A. Architecture

The framework consists of three main components:

1. Confidence-Aware Encoder:

   • Constructs a latent variable $z_1(x, \omega)$ by interpolating between a deterministic feature map $\bar{z}(x)$ and a Gaussian Process (GP) prior $z_0(x, \omega)$.
   • Uses a learnable confidence function $m(x) \in [0, 1]$ to control the balance:
     $$z_1 = \text{diag}(m(x))\bar{z}(x) + \text{diag}(1-m(x))z_0(x, \omega)$$
   • This allows the model to balance deterministic features with stochastic uncertainty, preserving spatial structure via the GP kernel.

2. Shared Latent Representation:

   • The same latent variable $z_1$ is fed into two decoders. This ensures that the uncertainty in the solution and the uncertainty in the correction are coherently coupled and propagated from the same source of epistemic uncertainty.

3. Dual Decoders:

   • Solution Decoder: A neural operator (using deep integral layers) that predicts the solution field $u(x)$ conditioned on $z_1$. It outputs a mean $\mu_u$ and aleatoric variance $\sigma_u^2$.
   • Correction Decoder: A network $S_\psi$ that infers a discrepancy term conditioned on the same $z_1$.
   • Physics Constraint: The forcing term $f(x)$ is modeled as:
     $$\mu_f(x) = \tilde{\mathcal{N}}_x[\mu_u(x)] + S_\psi(z_1(x))$$
     where $\tilde{\mathcal{N}}_x$ is the misspecified operator and $S_\psi$ is the learned correction.

B. Training and Inference

• Objective Function: The model is trained by maximizing the data likelihood while regularizing the latent space to prevent degeneracy.
  $$\mathcal{L} = \mathcal{L}_{data} - \beta \mathcal{L}_{reg}$$
  • $\mathcal{L}_{data}$: Negative log-likelihood of observed data ($u, f, b$) under the conditional Gaussian models.
  • $\mathcal{L}_{reg}$: A KL-divergence term ensuring the encoder's distribution stays close to a standard normal prior.
• Inference: Uncertainty is quantified via Monte Carlo sampling of the latent variable $\omega_E$. The predictive mean and variance are computed by averaging over samples, capturing both aleatoric (noise) and epistemic (model) uncertainty without repeated retraining or parameter sampling.

3. Key Contributions

1. Unified Latent-Space Correction: Introduction of a framework where the solution and the discrepancy term share a single stochastic latent field. This enables consistent uncertainty propagation across all quantities of interest (solution, forcing, and correction).
2. Joint Learning Strategy: Development of a unified objective that embeds PDE information directly into the data likelihood, allowing simultaneous learning of the solution and correction from a single coupled objective.
3. Efficiency: The method achieves uncertainty quantification comparable to Ensemble PINNs and Bayesian PINNs (B-PINNs) but avoids the computational cost of posterior sampling or ensemble training.
4. Robustness: The inclusion of latent-space regularization enhances robustness against overfitting and model collapse, particularly in sparse and noisy data regimes.

4. Numerical Results

The method was validated on four benchmark problems:

• ODE System:

  • Demonstrated that the corrected model (S3) recovers accuracy comparable to the correctly specified model (S1) even when the reaction term is structurally wrong (e.g., using $\cos(u)$ instead of $u(1-u)$).
  • Showed that uncertainty bands correctly cover errors in misspecified scenarios, whereas uncorrected models produced overconfident, biased predictions.
  • Performance matched B-PINNs but with significantly lower computational overhead.

• Reaction-Diffusion Equation:

  • Successfully corrected a misspecified reaction term ($u^2$ vs. $u(1-u)$).
  • The corrected model reduced prediction errors significantly compared to the misspecified baseline, with errors falling within the predicted uncertainty bounds.

• 2D Non-Newtonian Channel Flow:

  • Corrected a Newtonian model (constant viscosity) to approximate a power-law fluid.
  • The correction network $s(y)$ successfully captured the systematic discrepancy, yielding smoother and more accurate predictions for the source term $f$ compared to learning a spatially varying viscosity directly.

• 2D Cavity Flow (Navier-Stokes):

  • Corrected a misspecified viscosity model for non-Newtonian flow.
  • Results: The proposed method achieved relative $L_2$ errors of 0.0096 (u) and 0.0119 (v), outperforming the misspecified LVM-GP (0.0981/0.1634) and matching/exceeding Ensemble PINNs (0.0131/0.0223).
  • Visualizations confirmed that the corrected model closely matched reference solutions, while the misspecified model showed significant deviations.

5. Significance

• Reliability in Misspecified Physics: The paper provides a robust solution for scientific computing scenarios where the underlying physics are imperfectly known, a common real-world constraint.
• Computational Efficiency: By moving uncertainty quantification from parameter space to latent representation space, the method offers a scalable alternative to Bayesian inference and ensemble methods, making it feasible for complex, high-dimensional PDEs.
• Coherent Uncertainty: The shared latent structure ensures that the uncertainty in the correction term is naturally linked to the solution, providing a more interpretable and reliable measure of confidence in the model's predictions.
• Generalizability: The framework is applicable to a wide range of problems, from ODEs to complex fluid dynamics, demonstrating versatility in handling both additive and operator-level discrepancies.