Non-convex Learning via Replica Exchange Stochastic Gradient MCMC

Abstract

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ïve 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.

Accepted by ICML 2020

Cited by in corpus (9)

• Accelerated replica exchange stochastic gradient Langevin diffusion enhanced Bayesian DeepONet for solving noisy parametric PDEs
• flexgrid2vec: Learning Efficient Visual Representations Vectors
• Faster Convergence of Stochastic Gradient Langevin Dynamics for Non-Log-Concave Sampling
• A Contour Stochastic Gradient Langevin Dynamics Algorithm for Simulations of Multi-modal Distributions
• Multi-variance replica exchange stochastic gradient MCMC for inverse and forward Bayesian physics-informed neural network
• Stein Self-Repulsive Dynamics: Benefits From Past Samples
• UPANets: Learning from the Universal Pixel Attention Networks
• Accelerating Convergence of Replica Exchange Stochastic Gradient MCMC via Variance Reduction
• A New Framework for Variance-Reduced Hamiltonian Monte Carlo