# Non-convex Learning via Replica Exchange Stochastic Gradient MCMC

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 12 Aug 2020.

## Built on by (most cited first, 30 of 34)

- [Conformalized-DeepONet: A Distribution-Free Framework for Uncertainty Quantification in Deep Operator Networks](https://mlsanity.com/paper/2402.15406/) (2024)
- [Bayesian autoencoders for data-driven discovery of coordinates, governing equations and fundamental constants](https://mlsanity.com/paper/2211.10575/) (2022)
- [Accelerated replica exchange stochastic gradient Langevin diffusion enhanced Bayesian DeepONet for solving noisy parametric PDEs](https://mlsanity.com/paper/2111.02484/) (2021)
- [On Convergence of Federated Averaging Langevin Dynamics](https://mlsanity.com/paper/2112.05120/) (2021)
- [Polygonal Unadjusted Langevin Algorithms: Creating stable and efficient adaptive algorithms for neural networks](https://mlsanity.com/paper/2105.13937/) (2021)
- [Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfalls](https://mlsanity.com/paper/2212.07959/) (2022)
- [Deep Operator Learning-based Surrogate Models with Uncertainty Quantification for Optimizing Internal Cooling Channel Rib Profiles](https://mlsanity.com/paper/2306.00810/) (2023)
- [Accelerating Approximate Thompson Sampling with Underdamped Langevin Monte Carlo](https://mlsanity.com/paper/2401.11665/) (2024)
- [Interacting Contour Stochastic Gradient Langevin Dynamics](https://mlsanity.com/paper/2202.09867/) (2022)
- [Gradient-based Discrete Sampling with Automatic Cyclical Scheduling](https://mlsanity.com/paper/2402.17699/) (2024)
- [UPANets: Learning from the Universal Pixel Attention Networks](https://mlsanity.com/paper/2103.08640/) (2021)
- [Non-reversible Parallel Tempering for Deep Posterior Approximation](https://mlsanity.com/paper/2211.10837/) (2022)
- [Constrained Exploration via Reflected Replica Exchange Stochastic Gradient Langevin Dynamics](https://mlsanity.com/paper/2405.07839/) (2024)
- [Accelerating Langevin Monte Carlo Sampling: A Large Deviations Analysis](https://mlsanity.com/paper/2503.19066/) (2025)
- [Hard-Attention Gates with Gradient Routing for Endoscopic Image Computing](https://mlsanity.com/paper/2407.04400/) (2024)
- [Morephy-Net: An Evolutionary Multi-objective Optimization for Replica-Exchange-based Physics-informed Neural Operator Learning Networks](https://mlsanity.com/paper/2509.00663/) (2025)
- [Active Learning with Bayesian Multi-Fidelity Laplace Neural Operators for Oscillatory Parametric PDEs](https://mlsanity.com/paper/2502.00550/) (2025)
- [Parameter Expanded Stochastic Gradient Markov Chain Monte Carlo](https://mlsanity.com/paper/2503.00699/) (2025)
- [Subsampling Error in Stochastic Gradient Langevin Diffusions](https://mlsanity.com/paper/2305.13882/) (2023)
- [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)
- [Bayesian Computation in Deep Learning](https://mlsanity.com/paper/2502.18300/) (2025)
- [Enhancing Gradient-based Discrete Sampling via Parallel Tempering](https://mlsanity.com/paper/2502.19240/) (2025)
- [JaxSGMC: Modular stochastic gradient MCMC in JAX](https://mlsanity.com/paper/2505.11190/) (2025)
- [Continuous Policy and Value Iteration for Relaxed Stochastic Control Problems and Its Convergence](https://mlsanity.com/paper/2506.08121/) (2025)
- [One Period to Rule Them All: Identifying Critical Learning Periods in Deep Networks](https://mlsanity.com/paper/2506.15954/) (2025)
- [A New Framework for Variance-Reduced Hamiltonian Monte Carlo](https://mlsanity.com/paper/2102.04613/) (2021)
- [The surrogate Gibbs-posterior of a corrected stochastic MALA: Towards uncertainty quantification for neural networks](https://mlsanity.com/paper/2310.09335/) (2023)
- [Feynman-Kac Operator Expectation Estimator](https://mlsanity.com/paper/2407.02010/) (2024)
- [Exploring Non-Convex Discrete Energy Landscapes: An Efficient Langevin-Like Sampler with Replica Exchange](https://mlsanity.com/paper/2501.17323/) (2025)
- [A hierarchical resource-efficient deep policy gradient method for continuous-time optimal control problems](https://mlsanity.com/paper/2502.14141/) (2025)

## Abstract (arXiv)

Replica exchange Monte Carlo (reMC), also known as parallel tempering, is an important technique for accelerating the convergence of the conventional Markov Chain Monte Carlo (MCMC) algorithms. However, such a method requires the evaluation of the energy function based on the full dataset and is not scalable to big data. The na"ive implementation of reMC in mini-batch settings introduces large biases, which cannot be directly extended to the stochastic gradient MCMC (SGMCMC), the standard sampling method for simulating from deep neural networks (DNNs). In this paper, we propose an adaptive replica exchange SGMCMC (reSGMCMC) to automatically correct the bias and study the corresponding properties. The analysis implies an acceleration-accuracy trade-off in the numerical discretization of a Markov jump process in a stochastic environment. Empirically, we test the algorithm through extensive experiments on various setups and obtain the state-of-the-art results on CIFAR10, CIFAR100, and SVHN in both supervised learning and semi-supervised learning tasks.