Scalable Uncertainty Quantification for Deep Operator Networks using Randomized Priors

Abstract

We present a simple and effective approach for posterior uncertainty quantification in deep operator networks (DeepONets); an emerging paradigm for supervised learning in function spaces. We adopt a frequentist approach based on randomized prior ensembles, and put forth an efficient vectorized implementation for fast parallel inference on accelerated hardware. Through a collection of representative examples in computational mechanics and climate modeling, we show that the merits of the proposed approach are fourfold. (1) It can provide more robust and accurate predictions when compared against deterministic DeepONets. (2) It shows great capability in providing reliable uncertainty estimates on scarce data-sets with multi-scale function pairs. (3) It can effectively detect out-of-distribution and adversarial examples. (4) It can seamlessly quantify uncertainty due to model bias, as well as noise corruption in the data. Finally, we provide an optimized JAX library called {\em UQDeepONet} that can accommodate large model architectures, large ensemble sizes, as well as large data-sets with excellent parallel performance on accelerated hardware, thereby enabling uncertainty quantification for DeepONets in realistic large-scale applications.

23 pages, 11 figures

References in corpus (13)

• Array Programming with NumPy
• Hidden Physics Models: Machine Learning of Nonlinear Partial Differential Equations
• Weight Uncertainty in Neural Networks
• Fourier Neural Operator for Parametric Partial Differential Equations
• On the eigenvector bias of Fourier feature networks: From regression to solving multi-scale PDEs with physics-informed neural networks
• Uncertainty Quantification in Scientific Machine Learning: Methods, Metrics, and Comparisons
• Operator learning for predicting multiscale bubble growth dynamics
• Learning Operators with Coupled Attention
• DeepONet prediction of linear instability waves in high-speed boundary layers
• DeepM&Mnet for hypersonics: Predicting the coupled flow and finite-rate chemistry behind a normal shock using neural-network approximation of operators
• Accelerated replica exchange stochastic gradient Langevin diffusion enhanced Bayesian DeepONet for solving noisy parametric PDEs
• Output-Weighted Sampling for Multi-Armed Bandits with Extreme Payoffs
• Improved architectures and training algorithms for deep operator networks

Cited by in corpus (6)

• Reliable extrapolation of deep neural operators informed by physics or sparse observations
• RiemannONets: Interpretable Neural Operators for Riemann Problems
• Empowering Bayesian Neural Networks with Functional Priors through Anchored Ensembling for Mechanics Surrogate Modeling Applications
• Composite Bayesian Optimization In Function Spaces Using NEON -- Neural Epistemic Operator Networks
• Uncertainty-Aware Data-Based Method for Fast and Reliable Shape Optimization
• Alpha-VI DeepONet: A prior-robust variational Bayesian approach for enhancing DeepONets with uncertainty quantification