Accelerated replica exchange stochastic gradient Langevin diffusion enhanced Bayesian DeepONet for solving noisy parametric PDEs

Abstract

The Deep Operator Networks~(DeepONet) is a fundamentally different class of neural networks that we train to approximate nonlinear operators, including the solution operator of parametric partial differential equations (PDE). DeepONets have shown remarkable approximation and generalization capabilities even when trained with relatively small datasets. However, the performance of DeepONets deteriorates when the training data is polluted with noise, a scenario that occurs very often in practice. To enable DeepONets training with noisy data, we propose using the Bayesian framework of replica-exchange Langevin diffusion. Such a framework uses two particles, one for exploring and another for exploiting the loss function landscape of DeepONets. We show that the proposed framework's exploration and exploitation capabilities enable (1) improved training convergence for DeepONets in noisy scenarios and (2) attaching an uncertainty estimate for the predicted solutions of parametric PDEs. In addition, we show that replica-exchange Langeving Diffusion (remarkably) also improves the DeepONet's mean prediction accuracy in noisy scenarios compared with vanilla DeepONets trained with state-of-the-art gradient-based optimization algorithms (e.g. Adam). To reduce the potentially high computational cost of replica, in this work, we propose an accelerated training framework for replica-exchange Langevin diffusion that exploits the neural network architecture of DeepONets to reduce its computational cost up to 25% without compromising the proposed framework's performance. Finally, we illustrate the effectiveness of the proposed Bayesian framework using a series of experiments on four parametric PDE problems.

References in corpus (12)

• Quantifying total uncertainty in physics-informed neural networks for solving forward and inverse stochastic problems
• Stochastic Gradient Hamiltonian Monte Carlo
• Gradient Descent Provably Optimizes Over-parameterized Neural Networks
• On Lazy Training in Differentiable Programming
• DeepM&Mnet: Inferring the electroconvection multiphysics fields based on operator approximation by neural networks
• Computational multiscale methods for quasi-gas dynamic equations
• Non-convex Learning via Replica Exchange Stochastic Gradient MCMC
• Accelerating Nonconvex Learning via Replica Exchange Langevin Diffusion
• Multi-agent Reinforcement Learning Accelerated MCMC on Multiscale Inversion Problem
• Computational multiscale methods for parabolic wave approximations in heterogeneous media
• Multi-variance replica exchange stochastic gradient MCMC for inverse and forward Bayesian physics-informed neural network
• Learning Algorithms for Coarsening Uncertainty Space and Applications to Multiscale Simulations

Cited by in corpus (4)

• Reliable extrapolation of deep neural operators informed by physics or sparse observations
• Deep Learning Methods for Partial Differential Equations and Related Parameter Identification Problems
• Scalable Uncertainty Quantification for Deep Operator Networks using Randomized Priors
• Hybrid explicit-implicit learning for multiscale problems with time dependent source