Active operator learning with predictive uncertainty quantification for partial differential equations

TL;DR

This work proposes a lightweight predictive UQ method tailored for Deep operator networks (DeepONets) that also generalizes to other operator networks and demonstrates how predictive uncertainties can be used in the context of Bayesian optimization and active learning problems to yield improvements in accuracy and data-efficiency for outer-loop optimization procedures.

Abstract

With the increased prevalence of neural operators being used to provide rapid solutions to partial differential equations (PDEs), understanding the accuracy of model predictions and the associated error levels is necessary for deploying reliable surrogate models in scientific applications. Existing uncertainty quantification (UQ) frameworks employ ensembles or Bayesian methods, which can incur substantial computational costs during both training and inference. We propose a lightweight predictive UQ method tailored for Deep operator networks (DeepONets) that also generalizes to other operator networks. Numerical experiments on linear and nonlinear PDEs demonstrate that the framework's uncertainty estimates are unbiased and provide accurate out-of-distribution uncertainty predictions with a sufficiently large training dataset. Our framework provides fast inference and uncertainty estimates that can efficiently drive outer-loop analyses that would be prohibitively expensive with conventional solvers. We demonstrate how predictive uncertainties can be used in the context of Bayesian optimization and active learning problems to yield improvements in accuracy and data-efficiency for outer-loop optimization procedures. In the active learning setup, we extend the framework to Fourier Neural Operators (FNO) and describe a generalized method for other operator networks. To enable real-time deployment, we introduce an inference strategy based on precomputed trunk outputs and a sparse placement matrix, reducing evaluation time by more than a factor of five. Our method provides a practical route to uncertainty-aware operator learning in time-sensitive settings.

Figures (19)

• Figure 1: (a) The uncertainty-equipped operator network architecture splits predictions into mean and variance estimates. Depending on the operator network used, the uncertainty predictions could be integrated into the network like for our DeepONet, or an external (secondary) network such as for FNO. Which framework is used will depend on network architecture, but generalizes to any operator network. We interpret network outputs as parameters for a predictive probability distribution, which are calibrated to the observed error distributions during training. (b) The uncertainty estimates provided by the operator networks are employed to help guide outer-loop data aquisition procedures in function spaces. The quantity of interest (QoI) and variance estimates are derived from the operator network predictions, and this information is used to guide data acquisition by balancing exploitation and exploration.
• Figure 1: DeepONet architecture variants for the incorporation of boundary conditions (left) and the prediction of uncertainty parameters (right). The simplest variant for handling boundary conditions (top-left) was found to match the performance of more structured architectures. For uncertainty-equipped operator networks, we found that decoupling the trunk predictions for mean and uncertainty parameters consistently improved the performance of the network. The second variant on the right was selected since it introduces far less trainable parameters than independent uncertainty networks and was observed to achieve similar levels of performance.
• Figure 2: Observed network errors during training (top left) and interpretation of network outputs as parameters for a normal distribution (bottom left). The means and standard deviations provided by the network are used to assign likelihoods to the observed target values, and the network aims to increase these likelihoods during the training process. Histograms of the prediction errors for two input examples are shown to the right with predictive uncertainties overlaid to illustrate how the network uncertainty matches the observed errors near the end of training.
• Figure 3: Inference is optimized for fast evaluations on a fixed grid by precomputing the trunk outputs for each grid location and constructing a sparse placement matrix to map vectorized network outputs to the correct array locations. The timings reported in the top row are one-time computations performed after training. Red timings denote unoptimized speeds using loops and manual entry placement, while black timings reflect the average runtime per input data pair $(f, g)$ processed in batches. We note that the speedups reported use NumPy's $\texttt{einsum}$ operations applied to batched inputs and are not representative of single-example evaluations. The $\log \sigma$ calculation has been omitted from the diagram for brevity.
• Figure 4: Example input data consisting of 2-dimensional interior data and 1-dimensional boundary data. DeepONet sensor locations for the interior data (blue) and boundary data (red) are shown on the right.