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 Deep-ONets 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.

Uncertainty Quantification in Scientific Machine Learning: Methods, Metrics, and Comparisons

“…The objective of this section is to demonstrate the advantage of U-DeepONet as compared to standard DeepONet, which does not provide epistemic uncertainty estimates. Concurrently with the present work, Lin et al [92] employed replica-exchange MCMC in conjunction with DeepONet for operator learning. Their model is a version of U-DeepONet, with MCMC as the posterior inference algorithm.…”

“…( 8). Note that concurrently with the present work, Markov chain Monte Carlo has been employed for posterior inference in conjunction with DeepONet-based operator learning [92].…”

“…Specifically, the sampler swaps between the two particle chains with a certain ratio and at each step samples from the current chain. Concurrently with the present work, the replica-exchange MCMC technique has been employed in conjunction with DeepONet for operator learning [92]. Additional UQ methods based on MCMC can be found, indicatively, in [35,149,[163][164][165].…”

MIONet: Learning multiple-input operators via tensor product

“…Among these approaches, DeepONet has been applied and demonstrated good performance in diverse applications, such as high-speed boundary layer problems [7], multiphysics and multiscale problems of hypersonics [31] and electroconvection [2], multiscale bubble growth dynamics [22,23], fractional derivative operators [27], stochastic differential equations [27], solar-thermal system [34], and aortic dissection [45]. Several extensions of DeepONet have also been developed, such as Bayesian DeepONet [24], DeepONet with proper orthogonal decomposition (POD-DeepONet) [28], multiscale DeepONet [25], neural operator with coupled attention [13], and physics-informed DeepONet [42,9].…”

Multifidelity deep neural operators for efficient learning of partial differential equations with application to fast inverse design of nanoscale heat transport

“…Different architectures of deep neural operators have been developed, such as deep operator networks (DeepONet) [11,12,14], Fourier neural operators [13,14], nonlocal kernel networks [15], and several others [16,17,18]. Among these deep neural operators, DeepONet was the first one to be proposed (in 2019) [11], and many subsequent extensions and improvements have been developed, such as DeepONet with proper orthogonal decomposition (POD-DeepONet) [14], DeepONet for multiple-input operators (MIONet) [19], DeepONet for multi-physics problems via physics decomposition (DeepM&Mnet) [20,21], DeepONet with uncertainty quantification [22,23,24], multiscale DeepONet [25], POD-DeepONet with causality [26], and physics-informed DeepONet [27,28].…”