Replica exchange enhanced adaptively weighted stochastic gradient Langevin dynamics for Bayesian sampling and optimization

We propose the replica exchange adaptively weighted stochastic gradient Langevin dynamics (REAWSGLD) algorithm, designed explicitly for Bayesian learning with complex energy landscapes encountered in big data problems. By merging the 1/k-ensemble and replica exchange methods, this algorithm effectively escapes local traps in Monte Carlo simulation and non-convex optimization. It operates by running two Langevin dynamics processes concurrently at different temperatures, enabling position swaps between them. The lower temperature process, influenced by the 1/k-ensemble method, focuses on exploiting local geometry by protruding low-energy regions and biasing the sampling towards them. Meanwhile, the higher temperature process, influenced by larger noises, facilitates global exploration across the entire domain. The 1/k-ensemble and replica exchange methods are complementary: the 1/k-ensemble method mitigates the risk of the replica exchange method excessively exploring distribution tails, while the replica exchange method enhances the global exploration capability of the 1/k-ensemble method. The proposed algorithm has been empirically evaluated across various experiments, demonstrating its efficacy in navigating complex energy landscapes. The numerical results highlight its potential for Monte Carlo simulation and non-convex optimization in contemporary machine learning tasks.