Leveraging viscous Hamilton-Jacobi PDEs for uncertainty quantification in scientific machine learning

TL;DR

The paper presents a theoretical and computational framework that rewrites Bayesian inference in SciML as a problem governed by multi-time viscous Hamilton-Jacobi PDEs, enabling posterior moments to be read from PDE derivatives. By specializing to linear models with Gaussian likelihoods and Gaussian priors, the authors derive Riccati ODEs whose solutions yield the posterior mean and covariance, and they enable data-point additions/removals and hyperparameter tuning without retraining on the full dataset. The method supports continual learning and streaming updates, demonstrated across boundary-value ODEs, 1D advection-diffusion, and 2D Helmholtz problems, with domain decompositions and active-learning workflows. This approach offers computational and memory advantages for real-time uncertainty quantification in SciML, and it points to future extensions to broader model classes and PDE solvers.

Abstract

Uncertainty quantification (UQ) in scientific machine learning (SciML) combines the powerful predictive power of SciML with methods for quantifying the reliability of the learned models. However, two major challenges remain: limited interpretability and expensive training procedures. We provide a new interpretation for UQ problems by establishing a new theoretical connection between some Bayesian inference problems arising in SciML and viscous Hamilton-Jacobi partial differential equations (HJ PDEs). Namely, we show that the posterior mean and covariance can be recovered from the spatial gradient and Hessian of the solution to a viscous HJ PDE. As a first exploration of this connection, we specialize to Bayesian inference problems with linear models, Gaussian likelihoods, and Gaussian priors. In this case, the associated viscous HJ PDEs can be solved using Riccati ODEs, and we develop a new Riccati-based methodology that provides computational advantages when continuously updating the model predictions. Specifically, our Riccati-based approach can efficiently add or remove data points to the training set invariant to the order of the data and continuously tune hyperparameters. Moreover, neither update requires retraining on or access to previously incorporated data. We provide several examples from SciML involving noisy data and \textit{epistemic uncertainty} to illustrate the potential advantages of our approach. In particular, this approach's amenability to data streaming applications demonstrates its potential for real-time inferences, which, in turn, allows for applications in which the predicted uncertainty is used to dynamically alter the learning process.

Figures (13)

• Figure 1: (See Section \ref{['sec:general_connection']}) Illustration of a connection between a Bayesian inference problem in scientific machine learning (top) and the solution to a multi-time viscous HJ PDE (bottom). The colors indicate the associated quantities between problems. This color scheme is reused in the subsequent illustrations of our connection. The arrow labels describe how the boxed quantities are related. For example, the posterior mean in the learning problem is equivalent to the spatial gradient of the solution to the multi-time viscous HJ PDE (red).
• Figure 1: (See Section \ref{['sec:general_connection']}) Mathematical formulation of the connection between a Bayesian inference problem with linear model and Gaussian likelihood (top) and the solution to a multi-time viscous HJ PDE with quadratic Hamiltonian (bottom). The content of this illustration matches that of Figure \ref{['fig:intro_connection_in_words']} by replacing each term in Figure \ref{['fig:intro_connection_in_words']} with its corresponding mathematical expression. The colors indicate the associated quantities. The arrow labels describe how the boxed quantities are related.
• Figure 1: (See Section \ref{['sec:quadIC']}) Mathematical formulation of the connection between a Bayesian inference problem with linear model, Gaussian likelihood, and Gaussian prior (top) and the solution to a multi-time viscous HJ PDE with quadratic Hamiltonian and quadratic initial condition (bottom). The content of this illustration is a special case of the connection in Figure \ref{['fig:connection_in_math']}, where the prior is Gaussian (set $g({\boldsymbol{\theta}}) = \frac{1}{2}\langle {\boldsymbol{\theta}}, \Lambda^{-1} {\boldsymbol{\theta}}\rangle$ in Figure \ref{['fig:connection_in_math']}). The colors indicate the associated quantities. The arrow labels describe how the boxed quantities are related.
• Figure 1: Results of solving \ref{['eq:example_1']} using continual learning and our Riccati-based approach. (a) and (b) show the predicted mean and uncertainty of $u$ and $f$, respectively, after the 101th, 151th, and 201th noisy data point of $f$ becomes available. Our Riccati-based approach naturally coincides with the continual learning framework while inherently avoiding catastrophic forgetting. Our approach allows us to incrementally update the learned models without accessing the historical data ($\boldsymbol{\cdot}$), while also providing a quantitative metric for our confidence in the learned models in the form of the predicted uncertainties ($\blacksquare$). We observe that regions of low predicted uncertainty generally coincide with regions of high inference accuracy, which implies that this confidence metric is a good indicator of the reliability of the model.
• Figure 1: Deleting two outliers ($\boldsymbol{\times}$) sequentially from the trained model using our Riccati-based approach to solve \ref{['eq:example_1']}. The left column of (a) and (b) shows the inferences of $u$ and $f$, respectively, after the initial training with all data points. The middle column shows the inferences after the first outlier is deleted, and the right column shows the inferences after two outliers are deleted. The outliers are deleted by solving the associated Riccati ODEs backward in time. This process only uses the results of the previous training step and information of the point to be deleted, which provides potential computational advantages over more standard SciML approaches that would otherwise require retraining on the entire remaining dataset.

Theorems & Definitions (1)

• Remark 2.1