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

arXiv: https://arxiv.org/abs/2111.02484 · HTML: https://mlsanity.com/paper/2111.02484/

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 29 references that are in the map.
- First posted on arXiv 3 Nov 2021.

## Builds on (most cited first, 15 of 15)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Model Reduction and Neural Networks for Parametric PDEs](https://mlsanity.com/paper/2005.03180/) (2020)
- [Quantifying total uncertainty in physics-informed neural networks for solving forward and inverse stochastic problems](https://mlsanity.com/paper/1809.08327/) (2018)
- [A comprehensive deep learning-based approach to reduced order modeling of nonlinear time-dependent parametrized PDEs](https://mlsanity.com/paper/2001.04001/) (2020)
- [A Hitting Time Analysis of Stochastic Gradient Langevin Dynamics](https://mlsanity.com/paper/1702.05575/) (2017)
- [Non-convex Learning via Replica Exchange Stochastic Gradient MCMC](https://mlsanity.com/paper/2008.05367/) (2020)
- [Accelerating Nonconvex Learning via Replica Exchange Langevin Diffusion](https://mlsanity.com/paper/2007.01990/) (2020)
- [Multi-agent Reinforcement Learning Accelerated MCMC on Multiscale Inversion Problem](https://mlsanity.com/paper/2011.08954/) (2020)

## Built on by (most cited first, 22 of 22)

- [MIONet: Learning multiple-input operators via tensor product](https://mlsanity.com/paper/2202.06137/) (2022)
- [Multifidelity deep neural operators for efficient learning of partial differential equations with application to fast inverse design of nanoscale heat transport](https://mlsanity.com/paper/2204.06684/) (2022)
- [Physics-Informed Machine Learning: A Survey on Problems, Methods and Applications](https://mlsanity.com/paper/2211.08064/) (2022)
- [Reliable extrapolation of deep neural operators informed by physics or sparse observations](https://mlsanity.com/paper/2212.06347/) (2022)
- [Scalable Uncertainty Quantification for Deep Operator Networks using Randomized Priors](https://mlsanity.com/paper/2203.03048/) (2022)
- [Deep Learning Methods for Partial Differential Equations and Related Parameter Identification Problems](https://mlsanity.com/paper/2212.03130/) (2022)
- [NeuralUQ: A comprehensive library for uncertainty quantification in neural differential equations and operators](https://mlsanity.com/paper/2208.11866/) (2022)
- [Uncertainty quantification for noisy inputs-outputs in physics-informed neural networks and neural operators](https://mlsanity.com/paper/2311.11262/) (2023)
- [Sparse identification of nonlinear dynamics and Koopman operators with Shallow Recurrent Decoder Networks](https://mlsanity.com/paper/2501.13329/) (2025)
- [IB-UQ: Information bottleneck based uncertainty quantification for neural function regression and neural operator learning](https://mlsanity.com/paper/2302.03271/) (2023)
- [D2NO: Efficient Handling of Heterogeneous Input Function Spaces with Distributed Deep Neural Operators](https://mlsanity.com/paper/2310.18888/) (2023)
- [PROSE: Predicting Operators and Symbolic Expressions using Multimodal Transformers](https://mlsanity.com/paper/2309.16816/) (2023)
- [On Approximating the Dynamic Response of Synchronous Generators via Operator Learning: A Step Towards Building Deep Operator-based Power Grid Simulators](https://mlsanity.com/paper/2301.12538/) (2023)
- [Variational Bayes Deep Operator Network: A data-driven Bayesian solver for parametric differential equations](https://mlsanity.com/paper/2206.05655/) (2022)
- [A Multimodal PDE Foundation Model for Prediction and Scientific Text Descriptions](https://mlsanity.com/paper/2502.06026/) (2025)
- [A hybrid Decoder-DeepONet operator regression framework for unaligned observation data](https://mlsanity.com/paper/2308.09274/) (2023)
- [Uncertainty quantification for deeponets with ensemble kalman inversion](https://mlsanity.com/paper/2403.03444/) (2024)
- [Alpha-VI DeepONet: A prior-robust variational Bayesian approach for enhancing DeepONets with uncertainty quantification](https://mlsanity.com/paper/2408.00681/) (2024)
- [Coefficient-to-Basis Network: A Fine-Tunable Operator Learning Framework for Inverse Problems with Adaptive Discretizations and Theoretical Guarantees](https://mlsanity.com/paper/2503.08642/) (2025)
- [B-LSTM-MIONet: Bayesian LSTM-based Neural Operators for Learning the Response of Complex Dynamical Systems to Length-Variant Multiple Input Functions](https://mlsanity.com/paper/2311.16519/) (2023)
- [A novel data generation scheme for surrogate modelling with deep operator networks](https://mlsanity.com/paper/2402.16903/) (2024)
- [Physics-informed Multiple-Input Operators for efficient dynamic response prediction of structures](https://mlsanity.com/paper/2505.07090/) (2025)

## Abstract (arXiv)

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.