Improving Sampling Accuracy of Stochastic Gradient MCMC Methods via Non-uniform Subsampling of Gradients

1 INTRODUCTION . Many MCMC methods use physics-inspired evolution such as Langevin dynamics ( Brooks et al. , 2011 ) to utilize gradient information for exploring posterior distributions over continuous parameter space efficiently . However , gradient-based MCMC methods are often limited by the computational cost of computing the gradient on large data sets . Motivated by the great success of stochastic gradient methods for optimization , stochastic gradient MCMC methods ( SG-MCMC ) for sampling have also been gaining increasing attention . When the accurate but expensive-to-evaluate batch gradients in a MCMC method are replaced by computationally cheaper estimates based on a subset of the data , the method is turned to a stochastic gradient version . Classical examples include SG ( overdamped ) Langevin Dynamics ( Welling & Teh , 2011 ) and SG Hamiltonian Monte Carlo ( Chen et al. , 2014 ) , all of which were designed for scalability suitable for machine learning tasks . However , directly replacing the batch gradient by a ( uniform ) stochastic one without additional mitigation will generally cause a MCMC method to sample from a statistical distribution different from the target , because the transition kernel of the MCMC method gets corrupted by the noise of subsampled gradient . In general , the additional noise is tolerable if the learning rate/step size is tiny or decreasing . However , when large steps are used for better efficiency , the extra noise is nonnegligible and undermines the performance of downstream applications such as Bayesian inference . In this paper , we present a state-dependent non-uniform SG-MCMC algorithm termed Exponentially Weighted Stochastic Gradients method ( EWSG ) , which continues the efforts of uniform SG- MCMC methods for better scalability . Our approach is based on designing the transition kernel of a SG-MCMC method to match the transition kernel of a full-gradient-based MCMC method . This matching leads to non-uniform ( in fact , exponential ) weights that aim at capturing the entire statevariable distribution of the full-gradient-based MCMC method , rather than just providing unbiased gradient estimator or reducing its variance . When focusing on the variance , the advantage of EWSG is the following : recall the stochasticity of a SG-MCMC method can be decomposed into the intrinsic randomness of MCMC and the randomness introduced by gradient subsampling ; in conventional uniform subsampling treatments , the latter randomness is independent of the former , and thus when they are coupled together , variances add up ; EWSG , on the other hand , dynamically chooses the weight of each datum according to the current state of the MCMC , and thus the variances do not add up due to dependence . However , the gained accuracy is beyond reduced variance , as EWSG , when converged , samples from a distribution close to the invariant distribution of the full-gradient MCMC method ( which has no variance of the 2nd type ) , because its transition kernel ( of the corresponding Markov process ) is close to that of the full-gradient-MCMC method . This is how better sampling accuracy can be achieved . Our main demonstration of EWSG is based on 2nd-order Langevin equations ( a.k.a . inertial , kinetic , or underdamped Langevin ) , although it works for other MCMC methods too ( e.g. , Sec.F , G ) . To concentrate on the role of non-uniform SG weights , we will work with constant step sizes only . The fact that EWSG has locally reduced variance than its uniform counterpart is rigorously shown in Theorem 3 , and a global non-asymptotic analysis of EWSG is given in Theorem 4 to quantify its convergence properties and demonstrate the advantage over its uniform SG counterpart . A number of experiments on synthetic and real world data sets , across downstream tasks including Bayesian logistic regression and Bayesian neural networks , are conducted to validate our theoretical results and demonstrate the effectiveness of EWSG . In addition to improved accuracy , the convergence speed was empirically observed , in a fair comparison setup based on the same data pass , to be comparable to its uniform counterpart when hyper-parameters are appropriately chosen . The convergence ( per data pass ) was also seen to be clearly faster than a classical Variance Reduction ( VR ) approach ( note : for sampling , not optimization ) , and EWSG hence provides a useful alternative to VR . Additional theoretical investigation of EWSG convergence speed is provided in Sec . I. Terminology-wise , ∇V will be called the full/batch-gradient , n∇VI with random I will be called stochastic gradient ( SG ) , and when I is uniform distributed it will be called a uniform SG/subsampling , otherwise non-uniform . When uniform SG is used to approximate the batchgradient in underdamped Langevin , the method will be referred to as ( vanilla ) stochastic gradient underdamped Langevin dynamics ( SGULD/SGHMC ) 1 , and it serves as a baseline in experiments . 2 RELATED WORK . Stochastic Gradient MCMC Methods Since the seminal work of SGLD ( Welling & Teh , 2011 ) , much progress ( Ahn et al. , 2012 ; Patterson & Teh , 2013 ) has been made in the field of SG-MCMC . Teh et al . ( 2016 ) theoretically justified the convergence of SGLD and offered practical guidance on tuning step size . Li et al . ( 2016 ) introduced a preconditioner and improved stability of SGLD . We also refer to Maclaurin & Adams ( 2015 ) and Fu & Zhang ( 2017 ) which will be discussed in Sec.5 . While these work were mostly based on 1st-order ( overdamped ) Langevin , other dynamics were considered too . For instance , Chen et al . ( 2014 ) proposed SGHMC , which is closely related to 2ndorder Langevin dynamics ( Bou-Rabee & Sanz-Serna , 2018 ; Bou-Rabee et al. , 2018 ) , and Ma et al . ( 2015 ) put it in a more general framework . 2nd-order Langevin was recently shown to be faster than the 1st-order version in appropriate setups ( Cheng et al. , 2018b ; a ) and began to gain more attention . Variance Reduction For optimization , vanilla SG methods usually find approximate solutions quickly but the convergence slows down when an accurate solution is needed ( Bach , 2013 ; Johnson & Zhang , 2013 ) . SAG ( Schmidt et al. , 2017 ) improved the convergence speed of stochastic gradient methods to linear , which is the same as gradient descent methods with full gradient , at the expense of large memory overhead . SVRG ( Johnson & Zhang , 2013 ) successfully reduced this memory overhead . SAGA ( Defazio et al. , 2014 ) furthers improved convergence speed over SAG and SVRG . For 1SGULD is the same as the well-known SGHMC with B̂ = 0 , see ( Chen et al. , 2014 , Eq ( 13 ) and section 3.3 ) for details . To be consistent with existing literature , we will refer SGULD as SGHMC in the sequel . sampling , Dubey et al . ( 2016 ) applied VR techniques to SGLD ( see also ( Baker et al. , 2019 ; Chatterji et al. , 2018 ) ) . However , many VR methods have large memory overhead and/or periodically use the whole data set for gradient estimation calibration , and hence can be resource-demanding . EWSG is derived based on matching transition kernels of MCMC and improves the accuracy of the entire distribution rather than just the variance . However , it does have a consequence of variance reduction and thus can be implicitly regarded as a VR method . When compared to the classic work on VR for SG-MCMC ( Dubey et al. , 2016 ) , EWSG converges faster when the same amount of data pass is used , although its sampling accuracy is below that of VR for Gaussian targets ( but well above vanilla SG ; Sec.5.1 ) . In this sense , EWSG and VR suit different application domains : EWSG can replace vanilla SG for tasks in which the priority is speed and then accuracy , as it keeps the speed but improves the accuracy ; on the other hand , VR remains to be the heavy weapon for accuracydemanding scenarios . Importantly , EWSG , as a generic way to improve SG-MCMC methods , can be combined with VR too ( e.g. , Sec.G ) ; thus , they are not exclusive or competitors . Importance Sampling ( IS ) IS employs nonuniform weights to improve SG methods for optimization . Traditional IS uses fixes weights that do not change along iterations , and the weight computation requires prior information of gradient terms , e.g. , Lipschitz constants of gradient ( Needell et al. , 2014 ; Schmidt et al. , 2015 ; Csiba & Richtárik , 2018 ) , which are usually unknown or difficult to estimate . Adaptive IS was also proposed in which the importance was re-evaluated at each iteration , whose computation usually required the entire data set per iteration and may also require information like the upper bound of gradient ( Zhao & Zhang , 2015 ; Zhu , 2016 ) . For sampling , it is not easy to combine IS with SG ( Fu & Zhang , 2017 ) ; the same paper is , to our knowledge , the closest to this goal and will be compared with in Sec.5.3 . EWSG can be viewed as a way to combine ( adaptive ) IS with SG for efficient sampling . It require no oracle about the gradient , nor any evaluation over the full data set . Instead , an inner-loop Metropolis chain maintains a random index that approximates a state-dependent non-uniform distribution ( i.e . the weights/importance ) . 3 UNDERDAMPED LANGEVIN : THE BACKGROUND OF A MCMC METHOD . Underdamped Langevin Dynamics ( ULD ) is { dθ = rdt dr = − ( ∇V ( θ ) + γr ) dt+ σdW ( 1 ) where θ , r ∈ Rd are state and momentum variables , V is a potential energy function which in our context ( originated from cost minimization or Bayesian inference over many data ) is the sum of many terms V ( θ ) = ∑n i=1 Vi ( θ ) , γ is a friction coefficient , σ is intrinsic noise amplitude , and W is a standard d-dimensional Wiener process . Under mild assumptions on the potential V ( Pavliotis , 2014 ) , Langevin dynamics admits a unique invariant distribution π ( θ , r ) ∼ exp ( − 1T ( V ( θ ) + ‖r‖2 2 ) ) and is in many cases geometric ergodic . T is the temperature of system determined via the fluctuation dissipation theorem σ2 = 2γT ( Kubo , 1966 ) . The main reason for considering ULD rather than overdamped one is that ULD can converge faster than overdamped Langevin , in particular in high-dimension space ( e.g. , Cheng et al . ( 2018b ; a ) ; Tao & Ohsawa ( 2020 ) ) . Like the overdamped version , numerical integrators for ULD with well captured statistical properties of the continuous process have been extensively investigated ( e.g , Roberts et al . ( 1996 ) ; Bou-Rabee & Owhadi ( 2010 ) ) , and both the overdamped and underdamped integrators are friendly to derivations that will allow us to obtain explicit expressions of the non-uniform weights . 4 MAIN WORK . 4.1 AN ILLUSTRATION OF NON-OPTIMALITY OF UNIFORM SUBSAMPLING . In many applications , cases where data size n is larger than dimension d are not uncommon . In such cases , { ∇Vi } i=1,2 , ··· , n ⊂ Rd are linearly dependent and hence it is likely that there exist probability distributions { pi } i=1,2 , ··· , n other than the uniform one such that the gradient estimate is unbiased . This opens up the door to develop non-uniform subsampling schemes ( weights may be θ dependent ) , which can help reduce introduced additional variance while maintaining unbiasedness . In fact , in a reasonable setup , it turns out an optimal way of subsampling gradients , is far from being uniform : Theorem 1 Suppose given θ ∈ Rd , the errors of SG approximation bi = n∇Vi ( θ ) −∇V ( θ ) , 1 ≤ i ≤ n are i.i.d . absolutely continuous random vectors with possibly-θ-dependent density p ( x|θ ) . Define p ∈ Rn as a sparse vector if the number of non-zero entries in p is no greater than d + 1 . Then with probability 1 , the optimal probability distribution p ? that is unbiased and minimizes the trace of the covariance of n∇VI ( θ ) , i.e . p ? which solves the following , is a sparse vector . min p Tr ( EI∼p [ bIbTI ] ) s.t . EI∼p [ bI ] = 0 , ( 2 ) Despite the sparsity of p ? , which seemingly suggests one only needs at most d+1 gradient terms per iteration when using SG methods , it is not practical because p ? requires solving the linear programming problem ( 2 ) in Theorem 1 , for which an entire data pass is needed . Nevertheless , this result shows uniform SG can be far from optimal and motivates us to propose an exponentially weighted stochastic gradient method , which has reduced local variance with high probability and at the same time remains efficiently implementable without necessarily using all the data per parameter update .

Summary: This paper proposes a non-uniform sampling method for stochastic gradient minibatches for SG-MCMC. By sampling the indices of the stochastic gradients according to a parameter-specific (exponentially weighted) non-uniform distribution, the paper shows that it exactly matches the transition kernel of full batch gradient MCMC (for underdamped langevin dynamics) in Theorem 2. Although this exponentially weight distribution is intractable, the paper presents an approximate sampler in Algorithm 1 (that both uses a 1-step Metropolis Hastings approximation and a deterministic approximation for the $x$ term $r_{k+1} = r_k$). Finally the paper compares Algorithm 1 with other SGMCMC methods in a small synthetic Gaussian example, bayesian logistic regression on Covertype data, and a Bayesian neural network on MNIST.