Non-convex Learning via Replica Exchange Stochastic Gradient MCMC

TL;DR

This paper introduces an adaptive replica exchange stochastic gradient MCMC method that corrects biases in mini-batch settings, enabling scalable and accelerated sampling for deep neural network training, achieving state-of-the-art results.

Contribution

It proposes an adaptive reSGMCMC algorithm that corrects biases in mini-batch replica exchange MCMC, improving scalability and convergence in deep learning applications.

Findings

Achieves state-of-the-art results on CIFAR10, CIFAR100, and SVHN datasets.

Demonstrates an acceleration-accuracy trade-off in the numerical discretization process.

Validates effectiveness through extensive experiments in supervised and semi-supervised learning.

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\"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.