Accelerating Convergence of Replica Exchange Stochastic Gradient MCMC via Variance Reduction

1 INTRODUCTION . Stochastic gradient Monte Carlo methods ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Li et al. , 2016 ) are the golden standard for Bayesian inference in deep learning due to their theoretical guarantees in uncertainty quantification ( Vollmer et al. , 2016 ; Chen et al. , 2015 ) and non-convex optimization ( Zhang et al. , 2017 ) . However , despite their scalability with respect to the data size , their mixing rates are often extremely slow for complex deep neural networks with rugged energy landscapes ( Li et al. , 2018 ) . To speed up the convergence , several techniques have been proposed in the literature in order to accelerate their exploration of multiple modes on the energy landscape , for example , dynamic temperatures ( Ye et al. , 2017 ) and cyclic learning rates ( Zhang et al. , 2020 ) , to name a few . However , such strategies only explore contiguously a limited region around a few informative modes . Inspired by the successes of replica exchange , also known as parallel tempering , in traditional Monte Carlo methods ( Swendsen & Wang , 1986 ; Earl & Deem , 2005 ) , reSGLD ( Deng et al. , ∗Equal contribution 2020 ) uses multiple processes based on stochastic gradient Langevin dynamics ( SGLD ) where interactions between different SGLD chains are conducted in a manner that encourages large jumps . In addition to the ideal utilization of parallel computation , the resulting process is able to jump to more informative modes for more robust uncertainty quantification . However , the noisy energy estimators in mini-batch settings lead to a large bias in the naı̈ve swaps , and a large correction is required to reduce the bias , which yields few effective swaps and insignificant accelerations . Therefore , how to reduce the variance of noisy energy estimators becomes essential in speeding up the convergence . A long standing technique for variance reduction is the control variates method . The key to reducing the variance is to properly design correlated control variates so as to counteract some noise . Towards this direction , Dubey et al . ( 2016 ) ; Xu et al . ( 2018 ) proposed to update the control variate periodically for the stochastic gradient estimators and Baker et al . ( 2019 ) studied the construction of control variates using local modes . Despite the advantages in near-convex problems , a natural discrepancy between theory ( Chatterji et al. , 2018 ; Xu et al. , 2018 ; Zou et al. , 2019b ) and practice ( He et al. , 2016 ; Devlin et al. , 2019 ) is whether we should avoid the gradient noise in non-convex problems . To fill in the gap , we only focus on the variance reduction of noisy energy estimators to exploit the theoretical accelerations but no longer consider the variance reduction of the noisy gradients so that the empirical experience from stochastic gradient descents with momentum ( M-SGD ) can be naturally imported . In this paper we propose the variance-reduced replica exchange stochastic gradient Langevin dynamics ( VR-reSGLD ) algorithm to accelerate convergence by reducing the variance of the noisy energy estimators . This algorithm not only shows the potential of exponential acceleration via much more effective swaps in the non-asymptotic analysis but also demonstrates remarkable performance in practical tasks where a limited time is required ; while others ( Xu et al. , 2018 ; Zou et al. , 2019a ) may only work well when the dynamics is sufficiently mixed and the discretization error becomes a major component . Moreover , the existing discretization error of the Langevin-based Markov jump processes ( Chen et al. , 2019 ; Deng et al. , 2020 ; Futami et al. , 2020 ) is exponentially dependent on time due to the limitation of Grönwall ’ s inequality . To avoid such a crude estimate , we consider the generalized Girsanov theorem and a change of Poisson measure . As a result , we obtain a much tighter discretization error only polynomially dependent on time . Empirically , we test the algorithm through extensive experiments and achieve state-of-the-art performance in both optimization and uncertainty estimates . 2 PRELIMINARIES . A common problem , in Bayesian inference , is the simulation from a posterior P ( β|X ) ∝ P ( β ) ∏N i=1 P ( xi|β ) , where P ( β ) is a proper prior , ∏N i=1 P ( xi|β ) is the likelihood function and N is the number of data points . When N is large , the standard Langevin dynamics is too costly in evaluating the gradients . To tackle this issue , stochastic gradient Langevin dynamics ( SGLD ) ( Welling & Teh , 2011 ) was proposed to make the algorithm scalable by approximating the gradient through a mini-batch data B of size n such that βk = βk−1 − ηk N n ∑ i∈Bk ∇L ( xi|βk−1 ) + √ 2ηkτξk , ( 1 ) where βk ∈ Rd , τ denotes the temperature , ηk is the learning rate at iteration k , ξk is a standard Gaussian vector , and L ( · ) : = − log P ( β|X ) is the energy function . SGLD is known to converge weakly to a stationary Gibbs measure πτ ( β ) ∝ exp ( −L ( β ) /τ ) as ηk decays to 0 ( Teh et al. , 2016 ) . The temperature τ is the key to accelerating the computations in multi-modal distributions . On the one hand , a high temperature flattens the Gibbs distribution exp ( −L ( β ) /τ ) ( see the red curve in Fig.1 ( a ) ) and accelerates mixing by facilitating exploration of the whole domain , but the resulting distribution becomes much less concentrated around the global optima . On the other hand , a low temperature exploits the local region rapidly ; however , it may cause the particles to stick in a local region for an exponentially long time , as shown in the blue curve in Fig.1 ( a , b ) . To bridge the gap between global exploration and local exploitation , Deng et al . ( 2020 ) proposed the replica exchange SGLD algorithm ( reSGLD ) , which consists of a low-temperature SGLD to encourage exploitation and a high-temperature SGLD to support exploration β ( 1 ) k = β ( 1 ) k−1 − ηk N n ∑ i∈Bk ∇L ( xi|β ( 1 ) k−1 ) + √ 2ηkτ ( 1 ) ξ ( 1 ) k β ( 2 ) k = β ( 2 ) k−1 − ηk N n ∑ i∈Bk ∇L ( xi|β ( 2 ) k−1 ) + √ 2ηkτ ( 2 ) ξ ( 2 ) k , where the invariant measure is known to be π ( β ( 1 ) , β ( 2 ) ) ∝ exp ( −L ( β ( 1 ) ) τ ( 1 ) − L ( β ( 2 ) ) τ ( 2 ) ) as ηk → 0 and τ ( 1 ) < τ ( 2 ) . Moreover , the two processes may swap the positions to allow tunneling between different modes . To avoid inducing a large bias in mini-batch settings , a corrected swapping rate Ŝ is developed such that Ŝ = exp { ( 1 τ ( 1 ) − 1 τ ( 2 ) ) ( N n ∑ i∈Bk L ( xi|β ( 1 ) k ) − N n ∑ i∈Bk L ( xi|β ( 2 ) k ) − ( 1 τ ( 1 ) − 1 τ ( 2 ) ) σ̂2 F ) } , where σ̂2 is an estimator of the variance of Nn ∑ i∈Bk L ( xi|β ( 1 ) k ) − N n ∑ i∈Bk L ( xi|β ( 2 ) k ) and F is the correction factor to balance between acceleration and bias . In other words , the parameters switch the positions from ( β ( 1 ) k , β ( 2 ) k ) to ( β ( 2 ) k , β ( 1 ) k ) with a probability r ( 1∧ Ŝ ) ηk , where the constant r is the swapping intensity and can set to 1ηk for simplicity . From a probabilistic point of view , reSGLD is a discretization scheme of replica exchange Langevin diffusion ( reLD ) in mini-batch settings . Given a smooth test function f and a swapping-rate function S , the infinitesimal generator LS associated with the continuous-time reLD follows LSf ( β ( 1 ) , β ( 2 ) ) = −〈∇β ( 1 ) f ( β ( 1 ) , β ( 2 ) ) , ∇L ( β ( 1 ) ) 〉 − 〈∇β ( 2 ) f ( β ( 1 ) , β ( 2 ) ) , ∇L ( β ( 2 ) ) 〉 + τ ( 1 ) ∆β ( 1 ) f ( β ( 1 ) , β ( 2 ) ) + τ ( 2 ) ∆β ( 2 ) f ( β ( 1 ) , β ( 2 ) ) + rS ( β ( 1 ) , β ( 2 ) ) · ( f ( β ( 2 ) , β ( 1 ) ) − f ( β ( 1 ) , β ( 2 ) ) ) , where the last term arises from swaps and ∆β ( · ) is the the Laplace operator with respect to β ( · ) . Note that the infinitesimal generator is closely related to Dirichlet forms in characterizing the evolution of a stochastic process . By standard calculations in Markov semigroups ( Chen et al. , 2019 ) , the Dirichlet form ES associated with the infinitesimal generator LS follows ES ( f ) = ∫ ( τ ( 1 ) ‖∇β ( 1 ) f ( β ( 1 ) , β ( 2 ) ) ‖2 + τ ( 2 ) ‖∇β ( 2 ) f ( β ( 1 ) , β ( 2 ) ) ‖2 ) dπ ( β ( 1 ) , β ( 2 ) ) ︸ ︷︷ ︸ vanilla term E ( f ) + r 2 ∫ S ( β ( 1 ) , β ( 2 ) ) · ( f ( β ( 2 ) , β ( 1 ) ) − f ( β ( 1 ) , β ( 2 ) ) ) 2dπ ( β ( 1 ) , β ( 2 ) ) ︸ ︷︷ ︸ acceleration term , ( 2 ) which leads to a strictly positive acceleration under mild conditions and is crucial for the exponentially accelerated convergence in the W2 distance ( see Fig.1 ( c ) ) . However , the acceleration depends on the swapping-rate function S and becomes much smaller given a noisy estimate of N n ∑ i∈B L ( xi|β ) due to the demand of large corrections to reduce the bias . 3 VARIANCE REDUCTION IN REPLICA EXCHANGE STOCHASTIC GRADIENT LANGEVIN DYNAMICS . The desire to obtain more effective swaps and larger accelerations drives us to design more efficient energy estimators . A naı̈ve idea would be to apply a large batch size n , which reduces the variance of the noisy energy estimator proportionally . However , this comes with a significantly increased memory overhead and computations and therefore is inappropriate for big data problems . A natural idea to propose more effective swaps is to reduce the variance of the noisy energy estimator L ( B|β ( h ) ) = Nn ∑ i∈B L ( xi|β ( h ) ) for h ∈ { 1 , 2 } . Considering an unbiased estimator L ( B|β̂ ( h ) ) for ∑N i=1 L ( xi|β̂ ( h ) ) and a constant c , we see that a new estimator L̃ ( B|β ( h ) ) , which follows L̃ ( B|β ( h ) ) = L ( B|β ( h ) ) + c ( L ( B|β̂ ( h ) ) − N∑ i=1 L ( xi|β̂ ( h ) ) ) , ( 3 ) is still the unbiased estimator for ∑N i=1 L ( xi|β ( h ) ) . By decomposing the variance , we have Var ( L̃ ( B|β ( h ) ) ) = Var ( L ( B|β ( h ) ) ) + c2Var ( L ( B|β̂ ( h ) ) ) + 2cCov ( L ( B|β ( h ) ) , L ( B|β̂ ( h ) ) ) . In such a case , Var ( L̃ ( B|β ( h ) ) ) achieves the minimum variance ( 1 − ρ2 ) Var ( L ( B|β ( h ) ) ) given c ? : = −Cov ( L ( B|β ( h ) ) , L ( B|β̂ ( h ) ) ) Var ( L ( B|β̂ ( h ) ) ) , where Cov ( · , · ) denotes the covariance and ρ is the correlation coefficient of L ( B|β ( h ) ) and L ( B|β̂ ( h ) ) . To propose a correlated control variate , we follow Johnson & Zhang ( 2013 ) and update β̂ ( h ) = β ( h ) mb km c every m iterations . Moreover , the optimal c ? is often unknown in practice . To handle this issue , a well-known solution ( Johnson & Zhang , 2013 ) is to fix c = −1 given a high correlation |ρ| of the estimators and then we can present the VR-reSGLD algorithm in Algorithm 1 . Since the exact variance for correcting the stochastic swapping rate is unknown and even time-varying , we follow Deng et al . ( 2020 ) and propose to use stochastic approximation ( Robbins & Monro , 1951 ) to adaptively update the unknown variance . Variants of VR-reSGLD The number of iterations m to update the control variate β̂ ( h ) gives rise to a trade-off in computations and variance reduction . A small m introduces a highly correlated control variate at the cost of expensive computations ; a largem , however , may yield a less correlated control variate and setting c = −1 fails to reduce the variance . In spirit of the adaptive variance in Deng et al . ( 2020 ) to estimate the unknown variance , we explore the idea of the adaptive coefficient c̃k = ( 1− γk ) c̃k−m + γkck such that the unknown optimal c ? is well approximated . We present the adaptive VR-reSGLD in Algorithm 2 in Appendix E.2 and show empirically later that the adaptive VR-reSGLD leads to a significant improvement over VR-reSGLD for the less correlated estimators . A parallel line of research is to exploit the SAGA algorithm ( Defazio et al. , 2014 ) in the study of variance reduction . Despite the most effective performance in variance reduction ( Chatterji et al. , 2018 ) , the SAGA type of sampling algorithms require an excessively memory storage of O ( Nd ) , which is too costly for big data problems . Therefore , we leave the study of the lightweight SAGA algorithm inspired by Harikandeh et al . ( 2015 ) ; Zhou et al . ( 2019 ) for future works . Related work Although our VR-reSGLD is , in spirit , similar to VR-SGLD ( Dubey et al. , 2016 ; Xu et al. , 2018 ) , it differs from VR-SGLD in two aspects : First , VR-SGLD conducts variance reduction on the gradient and only shows promises in the nearly log-concave distributions or when the Markov process is sufficiently converged ; however , our VR-reSGLD solely focuses on the variance reduction of the energy estimator to propose more effective swaps , and therefore we can import the empirical experience in hyper-parameter tuning from M-SGD to our proposed algorithm . Second , VR-SGLD doesn ’ t accelerate the continuous-time Markov process but only focuses on reducing the discretization error ; VR-reSGLD possesses a larger acceleration term in the Dirichlet form ( 2 ) and shows a potential in exponentially speeding up the convergence of the continuous-time process in the early stage , in addition to the improvement on the discretization error . In other words , our algorithm is not only theoretically sound but also more empirically appealing for a wide variety of problems in non-convex learning . Algorithm 1 Variance-reduced replica exchange stochastic gradient Langevin dynamics ( VRreSGLD ) . The learning rate and temperature can be set to dynamic to speed up the computations . A larger smoothing factor γ captures the trend better but becomes less robust . T is the thinning factor to avoid a cumbersome system . Input The initial parameters β ( 1 ) 0 and β ( 2 ) 0 , learning rate η , temperatures τ ( 1 ) and τ ( 2 ) , correction factor F and smoothing factor γ. repeat Parallel sampling Randomly pick a mini-batch set Bk of size n. β ( h ) k = β ( h ) k−1 − η N n ∑ i∈Bk ∇L ( xi|β ( h ) k−1 ) + √ 2ητ ( h ) ξ ( h ) k , for h ∈ { 1 , 2 } . ( 4 ) Variance-reduced energy estimators Update L̂ ( h ) = ∑N i=1 L ( xi ∣∣∣β ( h ) mb k m c ) every m iterations . L̃ ( Bk|β ( h ) k ) = N n ∑ i∈Bk [ L ( xi|β ( h ) k ) − L ( xi ∣∣∣β ( h ) mb km c ) ] + L̂ ( h ) , for h ∈ { 1 , 2 } . ( 5 ) if k mod m = 0 then Update σ̃2k = ( 1− γ ) σ̃2k−m + γσ2k , where σ2k is an estimate for Var ( L̃ ( Bk|β ( 1 ) k ) − L̃ ( Bk|β ( 2 ) k ) ) . end if Bias-reduced swaps Swap β ( 1 ) k+1 and β ( 2 ) k+1 if u < S̃η , m , n , where u ∼ Unif [ 0 , 1 ] , and S̃η , m , n follows S̃η , m , n = exp { ( 1 τ ( 1 ) − 1 τ ( 2 ) ) ( L̃ ( Bk+1|β ( 1 ) k+1 ) − L̃ ( Bk+1|β ( 2 ) k+1 ) − 1 F ( 1 τ ( 1 ) − 1 τ ( 2 ) ) σ̃2 mb k m c ) } . ( 6 ) until k = kmax . Output : The low-temperature process { β ( 1 ) iT } bkmax/Tc i=1 , where T is the thinning factor .

Summary: The paper presents VR-reSGLD, a method to accelerate replica exchange stochastic gradient Langevin diffusion (reSGLD), which has been proposed recently to tackle non-convex learning problems. reSGLD suffers from two major sources of error resulting in low swapping rates: minibatch noise and the discretization error of Langevin diffusion. The idea of the paper is to use control variates to reduce the variance of energy estimators and thereby improve the swapping rate (which should lead to an accelerated convergence). Unlike previous modifications of SGLD, the variance reduction proposed in the paper aims at improving the energy estimators rather than the gradient estimators. The paper presents non-asymptotic results backing the intended acceleration of the Markov jump process. Numerical experiments illustrate the performance gain achieved by VR-reSGLD. These tests include a one-dimensional example (learning the component mean of a bimodal mixture of Gaussians) and Bayesian training of DNNs based on CIFAR imaging data.