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: In the context of sum-decomposable, smooth, and strongly convex potential functions the authors propose algorithms to improve the upper bound on the approximation error for estimating a Underdamped Langevin Diffusion (ULD) process and on the sampling error of the strongly-log-concave sampling problem. The continuous time ULD process has the strongly-log-concave distribution as its invariant distribution so approximating this process and sampling from this distribution are closely related. In the full gradient setting, the authors present a novel algorithm, ALUM which has optimal iteration complexity in $d$, $N$ and $\epsilon$ for the ULD approximation problem. Here, $N$ is the number of terms in the sum-decomposable function, $d$ is the input dimension and $\epsilon$ is the accuracy parameter. In comparison to the RMM algorithm (Lee et al), the dependence on $d$, $N$ and $\epsilon$ in the iteration complexity are matched, but the number of gradient evaluations are reduced by a factor of 2. Next the stochastic gradient oracle setting is considered: here the authors propose algorithms (SVRG-ALUM and SAGA-ALUM) and analyze the corresponding sampling/discretization error. Finally an improvement in the lower bound for the gradient complexity of the ULD approximation problem to $\Omega(N + d^{1/3}N^{2/3}/\epsilon^{2/3})$ is presented. This shows that the gradient complexity of the presented Variance reduction modifications of ALUM match the lower bound in $d$, $N$ and $n$ in the sum-decomposable setting with a stochastic gradient oracle.