Optimal Underdamped Langevin MCMC Method

1 3N 2 3 /ε 2 3 ) gradient evaluations to achieve ε-error ( in √ E∥·∥22 distance ) for approximating d-dimensional ULD . Moreover , we prove a lower bound of gradient complexity as Ω ( N + d 1 3N 2 3 /ε 2 3 ) , which indicates that our method is optimal in dependence of N , ε , and d. In particular , we apply our method to sample the strongly-log-concave distribution and obtain gradient complexity better than all existing gradient based sampling algorithms . Experimental results on both synthetic and real-world data show that our new method consistently outperforms the existing ULD approaches . 1 Introduction . Sampling is an important research problem in statistics learning with many applications such as Bayesian inference [ 1 ] , multi-arm bandit optimization [ 2 ] , and reinforcement learning [ 3 ] . One of the fundamental problems in these applications is to sample from a high-dimensional stronglylog-concave distribution . Recently , several Markov chain Monte Carlo ( MCMC ) based methods were proposed to solve this problem based on underdamped Langevin diffusion ( ULD ) . This continuous diffusion process converges to the target distribution exponentially fast . Thus , the methods approximating a ULD process could be used to sample from the target distribution within certain accuracy . Multiple discretization methods have been proposed for approximating ULD . Among them , the Euler-Maruyama discretization [ 4 ] is the simplest one but generates the largest error . Recently the left point method ( LPM ) 1 [ 5 ] was introduced to fix the gradient term in ULD to be the gradient at k-th iteration , and then integrate the new linear stochastic differential equation ( SDE ) with a small time-interval . Subsequently , Shen and Lee [ 6 ] proposed randomized midpoint method ( RMM ) with smaller error . There are also discretization schemes based on splitting [ 7 ] or Runge-Kutta method [ 8 , 9 ] . More recently , Cao et al . [ 10 ] derived an information-based complexity lower bound for simulating a d-dimensional ULD . Under the assumption that the full gradient oracle ∇f ( x ) is evaluated at most n times , they show a lower bound for worst-case error by perturbation analysis , which matches the discretization error upper bound of RMM in the dependence of d and n. Although the ULD-MCMC methods with full gradient oracle are largely understood , many real-world applications involve summation form of potential function and large-scale data , which leads to the need of stochastic gradient methods . The vanilla stochastic gradient methods have been used to replace full gradient [ 5 ] . Albeit the computational cost for each iteration is reduced , the variance of 1Although this method is mostly just denoted as ULD-MCMC , we adopt the name LPM to distinguish it from other discretization methods . The name comes from the fact that gradient is evaluated at the left point of the time interval . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . stochastic gradient is much larger than the discretization error and therefore degenerates the overall performance . Previous works [ 11 , 12 ] used stochastic variance reduced gradient ( SVRG ) [ 13 ] and SAGA [ 14 ] instead , but the gradient complexities of these methods are still worse than the full gradient RMM in terms of dependence on accuracy ε . Thus , there exists a natural question : What is the optimal ULD-MCMC method with sum-decomposable potential ? In this paper , we focus on optimal dependence of dimension d , components number N and accuracy ε in gradient complexity for estimating a ULD process . We answer this question by two parts . We first provide a novel ULD-MCMC method and derive the corresponding complexity upper bound in Sections 4 and 5 . After that , we analyze the worse case error and show that the lower bound matches the upper bound in Section 6 . The major contributions of our paper can be summarized as follows . 1 . We propose a new full gradient ULD-MCMC method , called as AcceLerated ULD method ( ALUM ) , whose discretization error has the same order dependence on dimension d , step size h as RMM . Although RMM already has optimal asymptotic complexity in full gradient setting , ALUM is still of practical interest . Compared with RMM , which uses two gradient evaluations at each iteration , ALUM uses gradient less frequently and only requires one gradient at each iteration to achieve constant speedup . 2 . We further propose VR-ALUM methods , including SVRG-ALUM and SAGA-ALUM , which utilize the unbiased variance reduction techniques in ALUM under sum-decomposable setting . We show that these methods achieve better gradient complexity than all existing gradient based MCMC approaches . These gradient complexities for sampling from a strongly-log-concave distribution are compared in Table 1 . 3 . We derive an information-based lower bound on worst-case error for estimating a ULD process with only gradient oracle and weighted Brownian motion oracle . We show that in order to achieve ε approximation accuracy , Ω ( N + d 1 3N 2 3 ε− 2 3 ) single component gradient evaluations are needed . This lower bound matches the upper bound for VR-ALUM in terms of dependence of dimension d , sample size N , and accuracy ε . Therefore , our VR-ALUM methods are indeed optimal for estimating a ULD process under sum-decomposable setting . 2In Zou et al . [ 11 ] , the authors call their method SVR-HMC . However , their method is not based on Hamiltonian Monte Carlo ( HMC ) , but based on ULD . Their method is just applying SVRG to replace full gradient in LPM . 3Further explanation and comparison is shown in Appendix A.2 . 2 Related work . Though we mainly use the gradient oracle in this paper , there is also a large body of sampling algorithms which leverage zeroth order potential value oracle . For example , Metropolis-Hastings accept-reject step could be used to ensure an MCMC converges to a stationary distribution equal to the target distribution [ 19 , 20 , 21 , 22 , 23 ] and the linear convergence can be obtained . For example , MALA [ 24 , 21 , 25 ] , as Metropolis-adjusted ULA , can achieve ε error in total variation distance within O ( d2 log ( 1/ε ) ) steps under certain initialization [ 26 ] . Note that purely gradient based sampling algorithms typically only converge sub-linearly because the stationary distribution is different from the target distribution . Other diffusion processes could also be used for constructing MCMC . ULA , a discretization of Langevin diffusion ( LD ) [ 27 , 24 ] , is the first commonly used gradient-based MCMC . We will provide an overview for some LD-based and ULD-based methods in Appendix A.1 . Later , higher order diffusion process is also discussed . Although the path is even smoother than ULD , currently it is not clear whether such higher order smoothness of path could be leveraged to accelerate the convergence without extra assumptions on potential function . Mou et al . [ 28 ] studied third-order diffusion , but their acceleration requires special structure or high order smoothness of potential . The more general diffusion process , which converges to the target distribution , was studied in Ma et al . [ 29 ] . 3 Preliminary . The problem of sampling from a strongly-log-concave distribution involves a probability density function p∗ ( x ) defined on a real vector space Rd . A corresponding potential function f ( x ) = − log ( p∗ ( x ) ) can be defined such that f ( x ) is strongly convex . One way to solve this sampling problem is constructing a Markov chain that converges to a stationary distribution that is the same as or similar to the target distribution . In Appendix A.1 , we introduce several such Markov chains as discretization of certain continuous stochastic processes , and we roughly analyze discretization error of these methods with pointing out the bottleneck . We use Õ ( f ) = O ( f ) logO ( 1 ) ( f ) to omit logarithm factor . ∥·∥2 means the Euclidean norm . We define a norm of random vector as ∥·∥L2 = √ E∥·∥22 . Next , we list our assumptions on potential f ( x ) . Assumption 1 ( Sum-decomposable ) . f ( x ) = ∑N i=1 fi ( x ) , where integer N is the sample size . Assumption 2 ( Smoothness ) . Each function fi is twice differentiable on Rd and there exists a constant L > 0 , such that ∇2fi ( x ) ≼ LN I for any x ∈ R d where I is the identity matrix . It can be easily verified that f ( x ) is L-smooth . Assumption 3 ( Strong Convexity ) . There exists a constant m > 0 such that : f ( x ) − f ( y ) ≥ ⟨∇f ( y ) , x− y⟩+ m2 ∥x− y∥ 2 2 . We define the condition number κ : = L/m . We finally define the 2-Wasserstein distance between distributions . For any pair of probability measures µ and ν on the same parameter space , a transference plan ζ between µ and ν is a joint distribution such that the marginal distributions on two sets of coordinates are µ and ν , respectively . We denote Γ ( µ , ν ) as the set of all transference plans , and define the 2-Wasserstein distance between µ and ν as follows : W 22 ( µ , ν ) = inf ζ∈Γ ( µ , ν ) ∫ ∥x− y∥22dζ ( x , y ) . 4 AcceLerated ULD-MCMC ( ALUM ) methods . In this section , we propose a class of AcceLerated ULD-MCMC ( ALUM ) methods based on approximation of the continuous ULD process : dXt = Vtdt , dVt = −∇f ( Xt ) dt− γVtdt+ √ 2γdBt . ( 1 ) The gradient term in the above SDE could be highly non-linear such that closed form solution is not available . Thus , we propose an estimation of SDE solution at time point h which only requires gradient evaluation at one single point : X ( o ) h = X0 + ψ1 ( h ) V0 − hψ1 ( h− ah ) ∇f ( X ( e ) ah ) + ex , [ 0 , h ] , V ( o ) h = ψ0 ( h ) V0 − hψ0 ( h− ah ) ∇f ( X ( e ) ah ) + ev , [ 0 , h ] , X ( e ) ah = X0 + ψ1 ( ah ) V0 + ex , [ 0 , ah ] , ( 2 ) where the random variable a is uniform random variable in [ 0 , 1 ] , ψi ( · ) and noise terms ex , [ 0 , h ] , ev , [ 0 , h ] , ex , [ 0 , ah ] are defined as follows 4. ψ0 ( h ) = e −γh , ψ1 ( h ) = 1 γ ( 1− e−γh ) , ex , [ 0 , h ] = √ 2γ ∫ h 0 ψ1 ( h− s ) dBs , ev , [ 0 , h ] = √ 2γ ∫ h 0 ψ0 ( h− s ) dBs . ( 3 ) 4.1 Full gradient ALUM . In this subsection , we propose a full gradient ALUM method , which is shown in Algorithm 1 . Specifically , the update rule is obtained by setting x ( o ) k+1 , v ( o ) k+1 to X ( o ) h , V ( o ) h where the ULD starts from X0 = x ( o ) k , V0 = v ( o ) k and h is a small step size . Algorithm 1 : Full gradient ALUM Method Input : Initial point ( x ( o ) 0 , v ( o ) 0 ) , parameter γ , iteration number K , and step size h > 0. for k = 0 to K − 1 do Randomly sample ak uniformly from [ 0 , 1 ] ; Generate ex , [ 0 , h ] , k , ev , [ 0 , h ] , k , ex , [ 0 , akh ] , k according to Appendix A.5 ; x ( e ) k = x ( o ) k + ψ1 ( akh ) v ( o ) k + ex , [ 0 , akh ] , k ; Calculate full gradient ∇f ( x ( e ) k ) ; x ( o ) k+1 = x ( o ) k +ψ1 ( h ) v ( o ) k −hψ1 ( h−akh ) ∇f ( x ( e ) k ) +ex , [ 0 , h ] , k ; v ( o ) k+1 = ψ0 ( h ) v ( o ) k −hψ0 ( h−akh ) ∇f ( x ( e ) k ) +ev , [ 0 , h ] , k ; end for Output : x ( o ) K . Next , we compare our method with both RMM and Nesterov ’ s accelerated gradient method ( NAG ) which has been shown optimal for gradient based optimization . Compared to RMM : By comparing the formula in ( 19 ) and ( 2 ) , it is easy to see that ALUM and RMM differ in one single term −ψ2 ( ah ) ∇f ( X0 ) . Our motivation of dropping this term is to reduce the computations . By deleting this term from the algorithm , only one gradient evaluation at a randomized midpoint is needed at each iteration . This compares favorably to the RMM , which requires two gradient evaluations at each iteration . However , could this method still estimate ULD accurately with only half gradient evaluations ? We firmly answer this question with rigorous error analysis in Theorem 2 . Roughly speaking , if we only consider the error ’ s dependence on step size h , the answer is yes . In this section , we only give an intuition of how that is possible . In Appendix A.1 , we show that RMM has very low bias and high variance . Therefore , the bottleneck is the variance , and increasing the bias slightly would not degenerate the overall performance much if the bias is not larger than the error introduced by variance . The dropped term −ψ2 ( ah ) ∇f ( X0 ) has norm O ( h2 ) , and the coefficient of gradient in Vh is O ( h ) . Thus , the bias introduced in single step is O ( h3 ) when we use step size h. After accumulating for T/h = O ( h−1 ) iterations , the bias is O ( h2 ) , which still has better dependence on h than the square root of variance O ( h3/2 ) . With increased bias , the complexities of ALUM and some other variants5 still have the same dependence on ε as RMM , the highest order dependence on κ could deteriorate as discussed in Section 7 . Compared to Nesterov ’ s accelerated gradient ( NAG ) method : NAG is a gradient based optimization method [ 30 ] . Based on [ 31 ] , NAG could be formulated as the following momentum method : 4ex , [ 0 , ah ] is obtained by simply substituting h with ah in the definition of ex , [ 0 , h ] . 5We can drop ex , [ 0 , h ] , k and ex , [ 0 , akh ] , k in ( 2 ) to derive other variants of ALUM . The bias of these variants increases to O ( h3/2 ) , which is still no larger than the square root of variance . The maximum order dependence on h is the same , thus the final iteration complexity has the same dependence on ε as RMM . xk+1 = xk + ckxk + c ′ k∇f ( x ( e ) k ) , vk+1 = ckxk + c ′ k∇f ( x ( e ) k ) , x ( e ) k = xk + ckxk . ( 4 ) We can see that the main differences between Algorithm 1 and ( 4 ) are coefficients and additional Gaussian noise terms . The coefficients in NAG are set in a deterministic way . However , the coefficients in ALUM could be random values from fixed distributions . The additional noise terms in ALUM come from the Brownian motion term in ULD and are necessary for sampling methods . Despite the differences , both NAG and ALUM use one gradient at each iteration , and both of them take a big jump along current momentum direction to calculate the gradient instead of directly computing the gradient at current iterate .

Summary: The paper makes new contributions to the literature on sampling from a strongly log-concave smooth measure. This is done by approximating the Underdamped Langevin Diffusion (ULD), and so, a fortiori, the paper also contributes to this approximation problem.  The main result is a new algorithm to approximate the ULD, which is well suited to a stochastic setting, where the gradient of the potential is replaced by certain estimators.  In the full gradient (i.e. without using an estimator for the gradient) setting the new algorithm matches the asymptotic performance of previous works while cutting the number of Gradient queries by about half. In the stochastic setting, if the potential can be represented as a sum of $N$ smooth functions, then the authors also show how to modify their algorithm to accommodate estimators to the gradient and reduce the overall query and iteration complexity.  The results of the stochastic setting are also complemented by a lower bound which essentially show that the dependence on the parameters $N$, and $d$, the dimension are essentially optimal.