Multi-variance replica exchange stochastic gradient MCMC for inverse and forward Bayesian physics-informed neural network

Abstract

Physics-informed neural network (PINN) has been successfully applied in solving a variety of nonlinear non-convex forward and inverse problems. However, the training is challenging because of the non-convex loss functions and the multiple optima in the Bayesian inverse problem. In this work, we propose a multi-variance replica exchange stochastic gradient Langevin diffusion method to tackle the challenge of the multiple local optima in the optimization and the challenge of the multiple modal posterior distribution in the inverse problem. Replica exchange methods are capable of escaping from the local traps and accelerating the convergence; two chains with different temperatures are designed where the low temperature chain aims for the local convergence, and the target of the high temperature chain is to travel globally and explore the whole loss function entropy landscape. However, it may not be efficient to solve mathematical inversion problems by using the vanilla replica method directly since the method doubles the computational cost in evaluating the forward solvers (likelihood functions) in the two chains. To address this issue, we propose to make different assumptions on the energy function estimation and this facilities one to use solvers of different fidelities in the likelihood function evaluation. More precisely, one can use a solver with low fidelity in the high temperature chain while use a solver with high fidelity in the low temperature chain. Our proposed method significantly lowers the computational cost in the high temperature chain, meanwhile preserves the accuracy and converges very fast. We give an unbiased estimate of the swapping rate and give an estimation of the discretization error of the scheme. To verify our idea, we design and solve four inverse problems which have multiple modes. The proposed method is also employed to train the Bayesian PINN to solve the forward and inverse problems; faster and more accurate convergence has been observed when compared to the stochastic gradient Langevin diffusion (SGLD) method and vanila replica exchange methods.

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

“…However, compared to vanilla SGLD [26], the computational cost of reSGLD doubles. To reduce the computational cost of reSGLD, we adopt the idea proposed in [24] and develop m-reSGLD (Algorithm 2); that is, an algorithm that uses (1) a full training regime for the low temperature particle θ 1 , which exploits the landscape of the energy function U , and (2) a accelerated training regime for the high temperature particle θ 2 , which explores the landscape of U . To this end, we first fully train both particles for a fixed number of burn-in epochs.…”

“…They demonstrated that the samples generated by their proposed algorithm converge to the target posterior distribution. Replicaexchange MCMC and (its stochastic gradient variants) were proposed [24,5,12] to accelerate the Langevin diffusion. In replica-exchange frameworks, instead of using one particle to sample from the posterior, one uses two particles with different temperature parameters.…”

“…Hence, one needs to balance having fast convergence (i.e., fewer iteration steps) and haven high per-iteration computational cost. To tackle such a problem, in our previous work [24], we proposed the multi-variance stochastic gradient Langevin diffusion (m-reSGLD) framework, which assumes that the estimators for the two particles have different accuracy. Hence, one can use solvers of different accuracy when using replica-exchange frameworks.…”

Theoretical and numerical studies of inverse source problem for the linear parabolic equation with sparse boundary measurements

“…One choice is to estimate each state p k (x) given a sequence of observations up to k in time by the posterior distribution. This task of sequential prediction based on the online observations can then be categorized as the standard Bayesian filtering problems which have been thoroughly studied in the control theory [10,5,24,20]. This motivates us to formulate the inversion problem in the Bayesian framework.…”

NH-PINN: Neural homogenization based physics-informed neural network for multiscale problems

“…Physics-informed neural network (PINN) is a neural network approach of solving the partial differential equations (PDE) [22,17,16]. The idea of the PINN is to approximate the solution of the PDE by a network.…”