Title: Double Randomized Underdamped Langevin with Dimension-Independent Convergence Guarantee
Venue: NeurIPS 2023
Accepted: yes
Keywords: Langevin

Dimension dependence

Acceleration
Abstract: This paper focuses on the high-dimensional sampling of log-concave distributions with composite structures: $p^*(\mathrm{d}x)\propto \exp(-g(x)-f(x))\mathrm{d}x$. We develop a double randomization technique, which leads to a fast underdamped Langevin algorithm with a dimension-independent convergence guarantee. We prove that the algorithm enjoys an overall $\tilde{\mathcal{O}}\left(\frac{\left(\mathrm{tr}(H)\right)^{1/3}}{\epsilon^{2/3}}\right)$ iteration complexity to reach an $\epsilon$-tolerated sample whose distribution $p$ admits $W_2(p,p^*)\leq \epsilon$.  Here,  $H$ is an upper bound of the Hessian matrices for $f$ and does not explicitly depend on dimension $d$. For the posterior sampling over linear models with normalized data, we show a clear superiority of convergence rate which is dimension-free and outperforms the previous best-known results by a $d^{1/3}$ factor. The analysis to achieve a faster convergence rate brings new insights into high-dimensional sampling.

Double Randomized Underdamped Langevin with Dimension-Independent Convergence Guarantee

Abstract

This paper focuses on the high-dimensional sampling of log-concave distributions with composite structures: p ∗ (d x ) ∝ exp( -g ( x ) -f ( x ))d x . We develop a double randomization technique, which leads to a fast underdamped Langevin algorithm with a dimension-independent convergence guarantee. We prove that the algorithm enjoys an overall ˜ O ( (tr( H )) 1 / 3 ϵ 2 / 3 ) iteration complexity to reach an ϵ -tolerated sample whose distribution p admits W 2 ( p, p ∗ ) ≤ ϵ . Here, H is an upper bound of the Hessian matrices for f and does not explicitly depend on dimension d . For the posterior sampling over linear models with normalized data, we show a clear superiority of convergence rate which is dimension-free and outperforms the previous best-known results by a d 1 / 3 factor. The analysis to achieve a faster convergence rate brings new insights into high-dimensional sampling.

1 Introduction

Sampling from a high-dimensional distribution serves as one of the key components in statistics, machine learning, and scientific computing, and constitutes the foundation of the fields including Bayesian statistics and generative models [Liu and Liu, 2001, Brooks et al., 2011, Song et al., 2020]. Recently, there is an emerging trend in designing provably faster Markov Chain Monte Carlo (MCMC) algorithms using techniques from first-order optimization [Dalalyan, 2017, Durmus et al., 2019, Cheng and Bartlett, 2018, Vempala and Wibisono, 2019, Chewi et al., 2021]. One typical MCMCalgorithm that allows following the idea is the Langevin-type algorithms, which are of the central interest of this paper.

Langevin-type algorithms originated in statistical physics discretize a stochastic differential equation with stationary distribution corresponding to the target. For Gibbs distribution p ( x ) ∝ e -U ( x ) , the standard overdamped Langevin algorithm uses Euler Maruyama scheme to discretize the diffusion process d x t = -∇ U ( x t )d t + √ 2d B t . The algorithm iteratively performs the updates as for n = 0 , 1 , 2 , . . . ,

where ϵ n ∼ N (0 , I ) follows the normal distribution and h > 0 is the step size.

This paper focuses on the dimension dependence of the convergence behavior of Langevin-type algorithms. Specifically, we consider sampling problems over potentially high-dimensional and strongly log-concave distributions with a composite structure: p ∝ e -U ( x ) = e -g ( x ) -f ( x ) where g ( x ) = m 2 ∥ x ∥ 2 . Regarding p as a posterior distribution, e -g ( x ) corresponds to a Gaussian prior and e -f ( x ) corresponds to the likelihood. This structure also includes general m -strongly convex

potential functions U ( x ) , which can be split into a strongly convex term g ( x ) = m 2 ∥ x ∥ 2 and a weakly convex one f ( x ) = U ( x ) -m 2 ∥ x ∥ 2 .

The analysis of Langevin sampling from an optimization viewpoint may date back to Jordan et al. [1998]. The viewpoint has led to a surge of works that establish quantitative convergence guarantees for overdamped Langevin algorithms to sample log-concave distributions [Dalalyan, 2017, Durmus and Moulines, 2016]. It should be noted that these convergence guarantees always involve an extra d dimension dependence due to the injection of non-negligible Gaussian noise, whereas the convergence rates of first-order optimization are often dimension-free [Nesterov, 2003]. Various accelerated methods have been proposed that can mitigate the dimension dependence and also achieve faster convergence. First, as the momentum acceleration of Langevin algorithms [Ma et al., 2021], underdamped Langevin Monte Carlo has a more stable trajectory and is known to exhibit a faster convergence rate [Cheng et al., 2018, Ma et al., 2021, Zhang et al., 2023]. Apart from using a different trajectory, more stable discretization schemes of the same diffusion process also lead to better convergence [Shen and Lee, 2019, Wibisono, 2019, He et al., 2020, Li et al., 2019]. However, the dimension dependence remains. To the best of our knowledge, the fastest randomized midpoint method for underdamped Langevin under Wasserstein distance achieves a O ( d 1 / 3 ϵ 2 / 3 ) convergence rate and still has a d 1 / 3 dimension dependence [Shen and Lee, 2019].

Recently another thread of works studies the convergence with weaker dimension dependence. Some researchers explore further general assumptions on the target distribution. The idea is by observing that some smoothness conditions can average out the dimension-dependent errors brought by the noise using Ito's formula. For example, Li et al. [2021] achieves a √ d dimension dependence with a 3-rd order growth condition. Another approach investigates the curvature of the target distribution and illustrates that dimension often does not determine the complexity of the sampling problem. Vono et al. [2022] propose an ADMM-type splitting algorithm with a dimension-free convergence rate when the likelihood is separable. Along the same thread, Freund et al. [2022] study the convergence rate of Langevin algorithms with no explicit dependence on dimension and propose to characterize the convergence rate by the upper bound of Hessian matrices of f . Specifically, by letting H be the upper bound of the Hessian matrices of f , they show a variant of the overdamped Langevin algorithm achieves a convergence rate of O ( tr( H ) ϵ ) in KL divergence. This result improves the rate under wide conditions because many high-dimensional sampling problems are intrinsically low-dimensional in the sense that tr( H ) = o ( d ) , which frequently appears in machine learning. For example, when the potential function has a ridge separable structure with mild conditions, tr( H ) can be dimension-free (see Section 3.3 for more details).

Though these works succeed in obtaining a convergence rate of the Langevin algorithm with a weak dependence on the dimension, some important questions remain:

' How to design provably faster algorithms with weak dependencies on dimension even for the log-concave sampling?'

Such a question is significant for understanding high-dimensional sampling and already includes lots of applications in real practice.

In this paper, we follow the regime of Freund et al. [2022] to design provably faster Langevin algorithms, which answers the above question affirmatively. We propose a double-randomized algorithm showing that variants of underdamped algorithms inherit the lower dimensional dependence in overdamped Langevin and a special design of random stepsize can still keep the property when using more stable discretization.

Specifically, we develop a double randomization technique and obtain the Double-Randomized Underdamped Langevin (DRUL) algorithm. We consider an averaged contraction effect to avoid dimension dependence. We design two distributions to achieve the mid-point acceleration with a randomized stepsize. DRUL is proven to enjoy an overall ˜ O ( (tr( H )) 1 / 3 ϵ 2 / 3 ) iteration and gradient complexity. For the posterior sampling over linear models, we show a clear superiority of DRUL, which achieves a dimension-free ˜ O ( 1 ϵ 2 / 3 ) convergence rate. The novel perspective of DRUL is introducing the random step size at each update. Such a random step size combined with a random midpoint point proposed by Shen and Lee [2019] reduces the discretization error and provides new insights into designing provably faster sampling algorithms.

In summary, the contributions of the paper are listed below:

(A) We propose the Double-Randomized technique and design the DRUL algorithm.

(B) We show the DRUL converges to the target distribution in Wasserstein distance in ˜ O ( (tr( H )+ ∥ x ∗ ∥ 2 ) 1 / 3 ϵ 2 / 3 ) iterations. For posterior sampling over generalized linear models, a dimension-free ˜ O ( 1 ϵ 2 / 3 ) complexity can be achieved.

2 Related works

There has been a surge of works investigating the asymptotic guarantees for the Langevin-type algorithms [Roberts and Tweedie, 1996, Mattingly et al., 2002]. And a series of recent works establish the non-asymptotic quantitative analysis framework of the Langevin-type algorithms. For performance on strongly log-concave distributions, early works [Dalalyan, 2017, Durmus and Moulines, 2016] establish the non-asymptotic convergence rate in TV distance with Lipschitz gradients assumption for overdamped dynamics. Along the same setting, similar convergence rates are achieved in KL divergence using the Wasserstein gradient flow [Cheng and Bartlett, 2018, Durmus et al., 2019] or Wasserstein distance using the contractive property for overdamped Langevin. Underdamped Langevin accelerates the vanilla Langevin algorithms [Ma et al., 2021]. With the same setting mentioned above, a faster convergence rate can be established [Cheng et al., 2018, Dalalyan and Riou-Durand, 2020, Shen and Lee, 2019, Zhang et al., 2023] for underdamped Langevin algorithms. Beyond the log Lipschitz-smooth setting, other examples also show better results can be achieved, such as Durmus and Moulines [2019], Li et al. [2019] for overdamped Langevin algorithms. Previous works also investigate the convergence rate without strongly log-convex assumptions, such as the log-Sobolev inequality (LSI) condition which is similar to the Polyak-Łojasiewicz condition in optimization. With target distributions satisfying LSI and log-Lipschitz-smooth, the results can be extended for both overdamped Langevin [Vempala and Wibisono, 2019] and underdamped Langevin [Ma et al., 2021, Zhang et al., 2023].

Another thread of works explores the dimension dependence of Langevin-type algorithms. Despite the great similarity between the analysis in Langevin algorithms and optimization, the convergence of Langevin-type algorithms depends on dimension in the log-concave setting whereas convex optimization algorithms can often achieve dimension-independent results. With general strongly log-

concave (or LSI) and log-Lipschitz-smooth condition, previous works establish a O ( d ϵ 2 ) convergence rate for overdamped algorithms [Durmus et al., 2019] and a ˜ O ( √ d ϵ ) convergence rate for underdamped algorithms [Cheng et al., 2018] to ensure finding a solution x n with W 2 (Law( x n ) , p ) ≤ ϵ or √ KL (Law( x n ) , p ) ≤ ϵ . With an additional linear growth condition on third-order derivative, Li et al. [2021] achieve a ˜ O ( √ d ϵ ) convergence rate for overdamped Langevin. And for underdamped Langevin algorithms, Shen and Lee [2019] improve the dependence of dimension and obtains a ˜ O ( d 1 / 3 ϵ 2 / 3 ) convergence rate. Recent work of Freund et al. [2022] characterizes the dimensional dependence of convergence rate by the upper bound of the Hessian matrix of f . As discussed in Section 1, Freund et al. [2022] establish a O ( tr( H )+ ∥ x ∗ ∥ 2 ϵ ) convergence rate in KL divergence, which implies a O ( tr( H )+ ∥ x ∗ ∥ 2 ϵ 2 ) convergence rate in Wasserstein distance. And when f admits a ridge-separable formula, Freund et al. [2022] obtain a dimension-free O ( 1 ϵ 2 ) convergence rate with some mild assumptions. We summarize the comparison of these most related works in Table 2.

3 Preliminary and problem setup

3.1 Notations

We use the convention O ( · ) and Ω( · ) to denote lower and upper bounds with a universal constant. ˜ O ( · ) ignores the polylogarithmic dependence. And use f ≲ g to denote f = O ( g ) . Use W 2 ( µ, ν ) to denote the 2 -Wasserstein distance of distribution µ and ν . Use Law( X ) to denote the distribution of the random variable X . The Frobenius norm is denoted by ∥ · ∥ F while ∥ · ∥ 2 stands for operator 2 -norm for matrices.

3.2 Sampling problem

We consider the distributions with a composite structure:

where g ( x ) = m 2 ∥ x ∥ 2 is a quadratic function. In the context of posterior sampling, the associated task is sampling from a distribution with a Gaussian prior. The composite structure also includes the general m -strongly convex function, which can be divided into g ( x ) = m 2 ∥ x ∥ 2 and the weakly convex function f ( x ) = U ( x ) -m 2 ∥ x ∥ 2 . We make the following assumption on f .

Assumption 3.1. f ∈ C 2 is convex and has L -Lipschitz continuous gradients, i.e. 0 ⪯ ∇ 2 f ⪯ LI .

It corresponds to making m -strongly convex and L + m -Lipschitz smooth assumptions on the potential function U , which is a basic setting and widely studied in the Langevin sampling literature (see e.g. Dalalyan [2017], Cheng and Bartlett [2018], Cheng et al. [2018], Shen and Lee [2019]).

In addition to the previous assumption, we follow Freund et al. [2022] to characterize the convergence rate by a new factor H to avoid explicit dimension dependence.

Definition 3.2. Let H be an upper bound of the Hessian matrices of f , i.e. H ⪰ ∇ 2 f ( x ) .

Note that the Lipschitz smooth condition in Assumption 3.1 provides a loose bound for H since we always have H ⪯ LI . This implies that tr( H ) will reach dL in the worst case, whereas this quantity can be much smaller than dL under wide conditions. One typical example is when f admits a so-called ridge separable structure shown below.

3.3 Example: ridge separable functions

Definition 3.3. f is said to admit the ridge separable form if

where σ i are all univariate functions, and a i are given vectors in R d .

Ridge separable functions contain many applications in machine learning, such as regression or classification over generalized linear models as well as (deep) neural networks in the neural tangent kernel regime [Gelman and Hill, 2006, Jacot et al., 2018]. We follow the argument of Freund et al. [2022], showing realizable conditions that ensure tr( H ) to be dimension-free.

Assumption 3.4. The function σ i ∈ C 2 has a bounded second derivative, i.e. σ ′′ i ≤ L 0 for all i ∈ [ n ] . Assumption 3.5. For all i ∈ [ n ] , then norm of a i is bounded by R , i.e. ∥ a i ∥ 2 ≤ R 2 .

When Assumptions 3.4 and 3.5 hold, the Hessian has the upper bound ∇ 2 f ( x ) = ∑ n i =1 1 n σ ′′ ( a T i x ) a i a T i ⪯ L 0 n ∑ n i =1 a i a T i . Let H = L 0 n ∑ n i =1 a i a T i and

thus illustrates that tr( H ) has a dimension-free L 0 R 2 upper bound. Meanwhile the Lipschitz constant of ∇ f can be bounded by

It indicates that worst-case upper bounds of tr( H ) and Lip( ∇ f ) are the same.

Note that Assumptions 3.4 and 3.5 are easy to achieve in practice. For example, consider using posterior sampling to compute the Bayes estimator over a linear model, whose advantages against maximum a posterior estimator have been discussed broadly (see e.g. Audibert [2009]). Here, a i is associated with the data. So Assumption 3.5 can be realized by simply normalizing the data. Moreover, σ i corresponds to the loss function and is only required to have a bounded second derivative by Assumption 3.4. Note that sampling from a ridge separable potential functions are extensively studied in related literature (see [Mou et al., 2021, V ono et al., 2022, Lee et al., 2018] as examples).

3.4 Preliminary

Overdamped Langevin. The overdamped Langevin dynamics for target distribution (3.1) is a diffusion process that evolves along the following SDE

where B t is the standard Brownian motion. The overdamped Langevin algorithms simulate and discretize the SDE (3.3). Different algorithms vary mainly by different discretization methods.

Underdamped Langevin Underdamped Langevin algorithms accelerate the convergence rate of overdamped Langevin algorithms and instead discretize the following dynamics

The diffusion process has the stationary distribution p ( x , v ) ∝ exp { -1 2 u ∥ v ∥ 2 -U ( x ) } .

Contractive Property One notable property that yields the convergence of Langevin dynamics is the contractive property [Cheng and Bartlett, 2018, Cheng et al., 2018], given as follows.

Definition 3.6. A stochastic differential equation has contractive property if there exists a positive constant m and

for any pair of solutions x ( t ) and y ( t ) driven by the same Brownian motion.

Remark 3.7. E ∥ x ( t ) -y ( t ) ∥ 2 is an upper bound of squared Wasserstein distance of Law( x ( t )) and Law( y ( t )) , and thus (3.5) implies a geometric convergence of the Wasserstein distance.

Contractive property can be established for both underdamped and overdamped Langevin dynamics under suitable conditions. The convergence analysis for our proposed algorithms follows a similar argument as the contractive property.

Algorithm 1 Double-Randomized Underdamped Langevin (DRUL)

Require: Iteration N , target function U = m 2 ∥ x ∥ 2 + f ( x ) , initial point ( x 0 , v 0 ) , max step size h , u = 1 and κ = L + m .

4 Double-randomized underdamped Langevin algorithm

In this section, we introduce our double-randomized sampling method and present our main result. We will illustrate our intuition in Section 5.

The proposed algorithm is built upon the following discretization scheme given a fixed point ˜ x n

The purpose of splitting the strongly convex part of U ( x ) is to avoid the md dimension dependence of tr( H ) . Although m is reasonably small in practice, one cannot obtain a fully dimensional-free convergence rate if the term with respect to g is discretized.

Denote the solution of process (4.1) at time t given starting point ( x n , v n ) , Brownian motion { B t } 0 ≤ t ≤ β , step size β and point ˜ x n by J ( β, ˜ x n ; { B t } 0 ≤ t ≤ β , ( x n , v n )) . The double-randomized algorithm performs the following one-step update:

The algorithm introduces a random step size β ∼ ρ ∝ ( 1 h -t hκ ) defined on [0 , h ] other than deterministic h . And let α follows the distribution ρ ′ ∝ ( e t -h κ -t h ) on [0 , h ] .

The analysis focuses on a Gaussian prior setting. Under a Gaussian prior, the following Lemma points out that J is linear in starting point ( x n , v n ) and has a decoupled Brownian motion, and thus the update (4.2) can be easily implemented.

Lemma 4.1. Assume g ( x ) = m 2 ∥ x ∥ 2 . If x n ( t ) follows the discretized Langevin diffusion process (4.1) with starting point ( x n , v n ) and a given ˜ x n , for any t > 0 , it satisfies the integral equation

where

are deterministic functions of t . Moreover, if x ∗ n ( t ) follows the exact Langevin process (3.4) with starting point ( x n , v n ) , then for any t > 0

Given the integral formula (4.3), the algorithm can be summarized as Algorithm 1. In Algorithm 1, A 21 , A 22 is deterministic scalar functions, and Σ : R 2 → R 3 d × 3 d is also deterministic. For the explicit formula, please refer to Appendix A. The distribution of ( H,G,W ) is induced by the coupling between { B t } 0 ≤ t ≤ α and { B t } 0 ≤ t ≤ β in (4.2).

4.1 Main theorem

The convergence guarantee of DRUL can be stated as the following theorem.

Theorem 4.2 ( Main theorem , convergence of DRUL) . For any tolerance ϵ ∈ (0 , 1) , denote the minimizer of U ( x ) by x ∗ and set the step size

where C 2 ≥ 1 is a universal constant. With initial point ( x 0 , v 0 ) , define Ω 0 = E x ∼ p, v ∼N (0 ,u ) ( ∥ x 0 -x ∥ 2 + ∥ x 0 + v 0 -v -x ∥ 2 ) . Then under Assumptions 3.1, when

Algorithm 1 outputs x n such that W 2 (Law( x n ) , p ) ≤ ϵ .

Theorem 4.2 shows an overall ˜ O ( ( tr( H )+ ∥ x ∗ ∥ 2 ) 1 / 3 ϵ 2 / 3 ) iteration and gradient complexity to find an ϵ -approximate sample in 2 -Wasserstein distance. By pre-finding x ∗ using a convex optimization algorithm and linearly drifting the coordinates, we can assume x ∗ = 0 without loss of generality. DRUL admits the stepsize reaching O ( ϵ 2 / 3 (tr( H )) 1 / 3 ) and the overall complexity is ˜ O ( (tr( H )) 1 / 3 ϵ 2 / 3 ) .

Now we consider a realizable case where our result achieves better complexity on d . The concrete example is that f admits a separable structure as discussed in Section 3.3. Corollary 4.3 below shows that when f admits a ridge separable structure, Algorithm 1 enjoys a dimension-free iteration and gradient complexity.

Corollary 4.3. Follow the notations in Theorem 4.2. Further, if f ( x ) admits a ridge-separable form and satisfies Assumptions 3.4 and 3.5. Then if we set the step size

With initial point ( x 0 , v 0 ) and Ω 0 = E x ∼ p, v ∼N (0 ,u ) ( ∥ x 0 -x ∥ 2 + ∥ x 0 + v 0 -v -x ∥ 2 ) . Then when

Algorithm 1 outputs x n such that W 2 (Law( x n ) , p ) ≤ ϵ .

Discussion. Theorem 4.2 establishes the convergence rate of DRUL with a weak dimension dependence. We shall note that for lots of high-dimensional sampling problems, the trace of the Hessian matrices for the potential function is much smaller than d times the largest eigenvalue because the eigenvalues often drop rapidly. In practice, this situation has been considered in lots of topics. For

(a)

(b)

Eigenvalue spectrum of the Hessian

10k

140

2

120

1000

5

10hiddenunits

30hidden units

100

2

50hidden units

100

70hidden units

80

5

2

60

10

40

2

20

1

0

0

0.02

0.04

0.06

0.08

0

100

200

300

400

500

600

700

800

example, one defines the effective rank (e.g. Hsu et al. [2012]) to represent the acting dimension of the data on linear models. In distributed machine learning, one can show fewer bits are needed to be transmitted between machines when the eigenvalues decrease fast (e.g. Hanzely et al. [2018]). As shown in Figure 1(a), the Gram matrix of the MNIST dataset, which is also the Hessian of Bayesian ridge linear regression, has rapidly decreasing eigenvalues. Similar empirical results are observed on deep neural network models by Sagun et al. [2016], as in Figure 1(b). In fact, beyond the ridge separable case, there are many problems admitting tr( H ) = o ( d ) , such as the concentrated posterior distributions with bounded gradients and neural networks with regularization. However, one may notice that the result of Theorem 4.2 is based on a uniform upper bound of the Hessian matrices. We think this is the basic case to show that our algorithm achieves convergence with an intrinsically low-dimension dependence. It is possible to relax the condition to the upper bound of the traces for local Hessian matrices with an additional Hessian smoothness assumption. Please see more discussions in Appendix E.

5 Intuition

In this section, we provide the intuitions of the algorithm design. For the simplicity of the notation, let ̂ x n ( s ) = J ( s, x n ; { B t } 0 ≤ t ≤ s , ( x n , v n )) and x n ( s ) = J ( s, J ( α, x n ; { B t } 0 ≤ t ≤ α , ( x n , v n )); { B t } 0 ≤ t ≤ s , ( x n , v n )) . Then given α and β in Algorithm 1, x n ( s ) = J ( s, ̂ x n ( α ); { B t } 0 ≤ t ≤ β , ( x n , v n )) and x n ( β ) = x n +1 .

Our analysis tracks the dynamics of the distance to the stationary distribution. Specifically, we analysis the dynamics of E Ω n ( t ) = E ∥ x n ( t ) -x ∗ n ( t ) + v n ( t ) -v ∗ n ( t ) ∥ 2 + E ∥ x n ( t ) -x ∗ n ( t ) ∥ 2 and E Ω n = E Ω n (0) , which upper bound the squared 2-Wasserstein distance W 2 (Law( x n ( t )) , p ) 2 and W 2 (Law( x n ) , p ) 2 , respectively. Here, ( x ∗ n ( t ) , v ∗ n ( t )) evolving along the exact Langevin diffusion (3.4) follows the stationary distribution and is synchronously coupled with ( x n ( t ) , v n ( t )) . One can find its formal definition in Appendix B. Via tracking the flow and using the contraction property of the process, the following Lemma characterizes the one-step discretization.

Lemma 5.1. Let x n be the n -step output of Algorithm 1 and x ∗ n ( t ) be defined as above. Set u = 1 L + m . Under Assumptions 3.1, given that x n and x ∗ n are coupled synchronously, we have for any h > 0

where E is

Lemma 5.1 indicates at each step, E Ω n contracts with a local discretization error. Telescoping (5.1) and upper bounding the term E yields the Theorem 4.2. To obtain the convergence rate in the

main theorem, we will show that Algorithm 1 guarantees an upper bound of E which (1) achieves the state-of-the-art dependence on stepsize and (2) satisfies that the dimension dependence can be controlled by the contraction.

Improved discretization error. We accomplish the first goal by matching the expectation, which is detailed by Lemma 5.2.

Lemma 5.2. Let ρ be probability measure defined on [0 , h ] satisfying ρ (d t ) ∝ ( 1 h -t hκ ) d t , and the probability measure ρ ′ defined on [0 , h ] satisfies ρ ′ (d t ) ∝ ( e t -h κ -t h ) d t on [0 , h ] . Then for measurable function F ( t ) , ρ and ρ ′ satisfy that

(A) There exists positive constant C 1 such that E t ∼ ρ ∫ t 0 e s -t κ F ( t )d s = C 1 h E t ∼ ρ ′ F ( t ) .

(B) There exists positive constant C 2 such that E t ∼ ρ ′ | F ( t ) | ≤ C E t ∼ ρ | F ( t ) | .

Claim (A) in Lemma 5.2 states that by choosing a random step size β ∼ ρ ′ , we can leverage the low discretization error of the randomized midpoint method. Denote the random weight w n ( s, α ) = x n ( s ) -x ∗ n ( s ) + v n ( s ) -v ∗ n ( s ) . One can split w n ( s, α ) into ( w n ( s, α ) -w n (0 , 0)) + w n (0 , 0) . The former term can be bounded using the one-step move, and by claim (A), the dominating latter one is

which attains a low discretization error given that E s ∼ ρ ′ ∇ f ( ̂ x n ( s )) is a low biased approximation to E s ∼ ρ ′ ∇ f ( x n ( s )) .

Averaged contraction can control the weight variance. Then we consider the variation of the weight w n ( s, α ) . The time difference of w n ( s, α ) writes

(5.2) indicates the variance Var s ∼ µ ( w ( s, α )) is dimension dependent for nondegenerated distributions µ on [0 , h ] , since it will introduce the v n ( r ) whose difference to v n ( β ) or v n (0) is dimension dependent. We control the dimension dependence via the averaged contraction. And the randomized step size makes it possible to consider the averaged effect. We have the following Lemma to bound the variation of the weight.

Lemma 5.3. Let x n ( t ) , v n ( t ) , v ∗ n ( t ) and v ∗ n ( t ) be defined as above and u = 1 L + m . Under Assumption 3.1, for any t, α ≤ h , we have

Now the dimension dependence of Var( w ( s, α )) is contained in E t ∼ ρ Ω n ( t ) , which can be controlled using the averaged contraction under the stochastic step size.

6 Conclusion

This paper proposes a double-randomized technique and designs the DRUL algorithm. We prove that with strongly convex and Lipschitz smooth assumptions potentials, the algorithm converges to the target distribution in Wasserstein distance in ˜ O ( (tr( H )+ ∥ x ∗ ∥ 2 ) 1 / 3 ϵ 2 / 3 ) iterations. The result illustrates that many sampling tasks in machine learning can achieve a dimension-independent complexity. The proposed DRUL algorithm can be potentially much faster than existing algorithms for highdimensional problems. As a concrete example, when the negative log-likelihood function admits a ridge-separable structure, under mild conditions, a dimension-free ˜ O ( 1 ϵ 2 / 3 ) iteration complexities can be obtained by DRUL. We hope our technique brings new insights for designing dimensionindependent algorithms for high-dimensional sampling.

References

Jun S Liu and Jun S Liu. Monte Carlo strategies in scientific computing , volume 75. Springer, 2001.

Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of markov chain monte carlo . CRC press, 2011.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Arnak S Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 79(3): 651-676, 2017.

Alain Durmus, Szymon Majewski, and Bła˙ zej Miasojedow. Analysis of langevin monte carlo via convex optimization. The Journal of Machine Learning Research , 20(1):2666-2711, 2019.

Xiang Cheng and Peter Bartlett. Convergence of langevin mcmc in kl-divergence. In Algorithmic Learning Theory , pages 186-211. PMLR, 2018.

Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. Advances in neural information processing systems , 32, 2019.

Sinho Chewi, Murat A Erdogdu, Mufan Bill Li, Ruoqi Shen, and Matthew Zhang. Analysis of langevin monte carlo from poincar \ 'e to log-sobolev. arXiv preprint arXiv:2112.12662 , 2021.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. SIAM journal on mathematical analysis , 29(1):1-17, 1998.

Alain Durmus and Eric Moulines. Sampling from a strongly log-concave distribution with the unadjusted langevin algorithm. 2016.

Yurii Nesterov. Introductory lectures on convex optimization: A basic course , volume 87. Springer Science & Business Media, 2003.

Yi-An Ma, Niladri S Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L Bartlett, and Michael I Jordan. Is there an analog of nesterov acceleration for gradient-based mcmc? Bernoulli , 27(3): 1942-1992, 2021.

Xiang Cheng, Niladri S Chatterji, Peter L Bartlett, and Michael I Jordan. Underdamped langevin mcmc: A non-asymptotic analysis. In Conference on learning theory , pages 300-323. PMLR, 2018.

Matthew Zhang, Sinho Chewi, Mufan Bill Li, Krishnakumar Balasubramanian, and Murat A Erdogdu. Improved discretization analysis for underdamped langevin monte carlo. arXiv preprint arXiv:2302.08049 , 2023.

Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems , 32, 2019.

Andre Wibisono. Proximal langevin algorithm: Rapid convergence under isoperimetry. arXiv preprint arXiv:1911.01469 , 2019.

Ye He, Krishnakumar Balasubramanian, and Murat A Erdogdu. On the ergodicity, bias and asymptotic normality of randomized midpoint sampling method. Advances in Neural Information Processing Systems , 33:7366-7376, 2020.

Xuechen Li, Yi Wu, Lester Mackey, and Murat A Erdogdu. Stochastic runge-kutta accelerates langevin monte carlo and beyond. Advances in neural information processing systems , 32, 2019.

Ruilin Li, Hongyuan Zha, and Molei Tao. Sqrt (d) dimension dependence of langevin monte carlo. arXiv preprint arXiv:2109.03839 , 2021.

Maxime Vono, Daniel Paulin, and Arnaud Doucet. Efficient mcmc sampling with dimension-free convergence rate using admm-type splitting. The Journal of Machine Learning Research , 23(1): 1100-1168, 2022.

Yoav Freund, Yi-An Ma, and Tong Zhang. When is the convergence time of langevin algorithms dimension independent? a composite optimization viewpoint. Journal of Machine Learning Research , 23(214):1-32, 2022.

Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli , pages 341-363, 1996.

Jonathan C Mattingly, Andrew M Stuart, and Desmond J Higham. Ergodicity for sdes and approximations: locally lipschitz vector fields and degenerate noise. Stochastic processes and their applications , 101(2):185-232, 2002.

Arnak S Dalalyan and Lionel Riou-Durand. On sampling from a log-concave density using kinetic langevin diffusions. Bernoulli , 26(3):1956-1988, 2020.

Alain Durmus and Eric Moulines. High-dimensional bayesian inference via the unadjusted langevin algorithm. Bernoulli , 25(4A):2854-2882, 2019.

Andrew Gelman and Jennifer Hill. Data analysis using regression and multilevel/hierarchical models . Cambridge university press, 2006.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Jean-Yves Audibert. Fast learning rates in statistical inference through aggregation. 2009.

Wenlong Mou, Yi-An Ma, Martin J Wainwright, Peter L Bartlett, and Michael I Jordan. High-order langevin diffusion yields an accelerated mcmc algorithm. The Journal of Machine Learning Research , 22(1):1919-1959, 2021.

Yin Tat Lee, Zhao Song, and Santosh S Vempala. Algorithmic theory of odes and sampling from well-conditioned logconcave densities. arXiv preprint arXiv:1812.06243 , 2018.

Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476 , 2016.

Daniel Hsu, Sham M Kakade, and Tong Zhang. Random design analysis of ridge regression. In Conference on learning theory , pages 9-1. JMLR Workshop and Conference Proceedings, 2012.

Filip Hanzely, Konstantin Mishchenko, Peter Richtárik, and A. Sega: Variance reduction via gradient sketching. Advances in Neural Information Processing Systems , 31, 2018.

Gilles Hargé. A convex/log-concave correlation inequality for gaussian measure and an application to abstract wiener spaces. Probability theory and related fields , 130(3):415-440, 2004.

Sanjib Basu and Anirban DasGupta. The mean, median, and mode of unimodal distributions: a characterization. Theory of Probability & Its Applications , 41(2):210-223, 1997.

A Simulation of the discretization process

A.1 Integral form of the discretized process

Proof of Lemma 4.1

Proof. Define the aligned vector z n ( s ) = ( x n ( s ) , v n ( s )) and write the diffusion process (4.1) as

where

And ˜ B s is a Brownian Motion in R 2 d . By a variation of constant method, the solution of (A.1) is

To obtain the explicit expression of x n ( t ) and v n ( t ) , respectively, it suffices to present a block-wise form

of the exponential e At . And a direct decomposition yields that

Since z n ( t ) = ( x n ( t ) , v n ( t )) , we obtain the close-form solution of the discretized Langevin process of x n ( t )

A similar proof can be applied to the exact diffusion process and thus x ∗ n ( t ) satisfies

Thus we can obtain the result in Lemma 4.1.

A.2 Controls of the integral forms

Lemma A.1. The functions A 11 ( t ) and A 12 ( t ) in Lemma 4.1 satisfy

for t ∈ [0 , h ] .

Proof. In the proof, we will repeatedly use the inequality e t -1 = O ( t ) . For notational simplicity, denote √ 1 -um by k . For the first claim,

As for the second claim, we have

Note that k = √ 1 -um ≤ 1 and h ≤ 1 κ ≤ 1 . When t ≤ h our claims hold true.

A.3 Simulation of the Brownian motion

Lemma A.2. Let B t be a standard Brownian Motion. In Algorithm 1, we have G = ∫ t 0 A 12 ( s -t )d B s , H = ∫ α 0 A 12 ( s -α )d B s and W = ∫ t 0 A 22 ( s -t )d B s . Then conditioned on the randomness of α and β , ( G,H,W ) follows a mean-zero Gaussian distribution with covariance

Proof. By the properties of Brownian motion, ( G,H,W ) is a zero-mean Gaussian variable. Its variance is given by

The rest of the claims can be proved similarly.

B Proof of the main theorem

In this section, we present the proof of Theorem 4.2.

B.1 Contractive lemma

We begin by defining the n-step exact process ( x ∗ n , v ∗ n ) and ( x ∗ n ( t ) , v ∗ n ( t )) . Denote d p ∗ ∝ e -U ( x ) -1 2 u ∥ v ∥ 2 d x d v . Let ( x ∗ 0 , v ∗ 0 ) ∼ p ∗ and

For the simplicity of notation, we denote ( x ∗ n , v ∗ n ) = ( x ∗ n (0) , v ∗ n (0)) . And we have for any n and t , ( x ∗ n ( t ) , v ∗ n ( t )) ∼ p ∗ . With slight abuse of notation, let ̂ x n ( t ) and x n ( t ) denote inaccurate starting point discretization process and midpoint process in Algorithm 1, respectively. And recall that in Algorithm 1 we have

and

And we assume ̂ x n ( t ) is driven by the same Brownian Motion as x n ( t ) and x ∗ n ( t ) at the n -th step.

Define Ω n ( t ) = ∥ x n ( t ) -x ∗ n ( t ) ∥ 2 + ∥ x n ( t ) -x ∗ n ( t ) + v n ( t ) -v ∗ n ( t ) ∥ 2 and Ω n = ∥ x n -x ∗ n ∥ 2 + ∥ x n -x ∗ n + v n -v ∗ n ∥ 2 . The proof concentrates on the dynamics of the Ω n . Note that Wasserstein distance minimizes all the couplings and Ω n ( t ) ≥ ∥ x n ( t ) -x ∗ n ( t ) ∥ 2 . Hence E Ω n upper bounds of the Wasserstein distance W 2 (Law( x n ) , p ) . Given a synchronously coupled assumption, the dynamics of Ω n ( t ) can be characterized as follow.

B.1.1 Proof of Lemma 5.1

Proof. Let the random process B α t given B α be

where B ′ t is an independent Brownian motion. Note that B α t is a Brownian bridge and B α α = B α . Then when ̂ x n ( α ) is given and x ∗ n ( t ) , x n ( t ) are coupled synchronously, given B α , we have

and

For simplicity, denote z n ( t ) = x n ( t ) -x ∗ n ( t ) and ψ n ( t ) = v n ( t ) -v ∗ n ( t ) . Then by Taylor's Theorem

Define H t = ∫ 1 0 ∇ 2 U ( s x n ( t ) + (1 -s ) x ∗ n ( t ))d s , then

where in the second equality we use (B.4) and divide ∇ f ( ̂ x n ( α )) into ∇ f ( x n ( t )) and the difference ∇ f ( ̂ x n ( α )) -∇ f ( x n ( t )) . Note that A can be converted to a quadratic form

whose eigenvalues λ i are given by characteristic function (1 -λ ) 2 -1 4 ( uλ i ( H t ) + um -2) 2 = 0 . Given the strong-convexity and Lipschitz smoothness of the potential U ( x ) , when u is set to be 1 L + m the eigenvalue of A is greater than 1 2 κ . Therefore,

Then by multiplying e t κ and taking the integral form 0 to β , we obtain

In Algorithm 1 the random step size β ∼ ρ and thus we reach that

B.2 Proof of Theorem 4.2

We devote the rest of this section to the proof of the main theorem. Lemma 5.1 shows that the dynamics of the E Ω n is driven by a shrinkage term and a local discretization error. To achieve the claimed convergence rate, the proof aims to establish an upper bound of local error which is: (1) dimensional-independent, (2) optimal in the sense of max step size h .

Proof. Given that E Ω n is an upper bound of the squared 2 -Wasserstein distance, it suffices to prove E Ω n ≤ ϵ 2 . By Lemma 5.1 and Lemma 5.2, we have

Next, we bound the error term A by

where in inequality a ≲ , we apply Lemma D.1, Lemma 5.3 and Lemma D.5; inequality b ≲ follows by the observation uL = L L + m ≤ 1 . Then by plugging the upper bound of A into (B.5), there exists constant C 2 such that

(B.6) can be simplified as

Note that when β ∼ ρ , E Ω n ( β ) = E Ω n +1 . And E Ω n is an upper bound of Wasserstein distance. Equation (B.7) characterizes the one-step discretization of the Langevin process, where the first term indicates a linear convergence if without discretization error, and the second term is the one-step discretization error.

Without loss of generality, let C 2 ≥ 1 . Since h ≤ 1 ≤ κ , we have E β ∼ ρ e -β κ = κ h ( 1 -e h κ ) ≤ 1 -h 3 κ . Since h ≤ 1 12 C 2 κ , we have h ≤ 1 2 C 2 κ and h ≤ 1 √ C 2 and therefore 1 -C 2 h 2 ≥ 0 , 2 C 2 h 2 ≤ h 6 κ . Hence the first term of (B.7) can be bounded by

As for error term in (B.7), since h ≤ 1 12 C 2 κ ≤ 1 √ 2 C 2 , we have C 2 h 4 1 -C 2 h 2 ≤ 2 C 2 h 4 . Denote r = 1 -h 6 κ and E = 2 C 2 h 4 ( u 2 L m tr( H ) + ∥ x ∗ ∥ 2 ) and therefore

When N ≥ 6 κ log ( 2 E Ω 0 ϵ 2 ) , we have and therefore completes the proof.

C Supporting lemmas

We assume that Assumption 3.1 holds in this section and Section D.

And when h ≤ ϵ 2 / 3 ( 24 C 2 κ ( u 2 L m tr( H )+ ∥ x ∗ ∥ 2 ) ) 1 / 3 , we get

Plugging (B.10) and (B.11) into (B.9) yields

C.1 Proof of Lemma 5.2

We would like to point out that how to match the random step size α and ρ is critical for obtaining a low-bias estimator. And the control for local discretization error in Lemma 5.1 involves a challenging multiple integral. These considerations inspire our choice of ρ and ρ ′ . And in the following lemma, we summarize the main properties of ρ and ρ ′ .

Lemma C.1. Let ρ be probability measure defined on [0 , h ] satisfying ρ (d t ) ∝ ( 1 h -t hκ ) d t , and the probability measure ρ ′ defined on [0 , h ] satisfies ρ ′ (d t ) ∝ ( e t -h κ -t h ) d t on [0 , h ] . Then for measurable function F ( t ) , ρ and ρ ′ satisfy that

(A) There exists positive constant C 1 such that E t ∼ ρ ∫ t 0 e s -t κ F ( t ) = C 1 h E t ∼ ρ ′ F ( t ) .

(B) There exists positive constant C 2 such that E t ∼ ρ ′ | F ( t ) | ≤ C E t ∼ ρ | F ( t ) | .

Proof. We begin with the first claim. By the definition of ρ and integration by parts, we obtain

which indicates that

ρ ∝ 1 h -t hκ has a normalizing constant which is of order Θ(1) . Denote the normalizing constant of ∫ h 0 ( e t -h κ -t h ) d s by a scalar function Z ′ ( h ) . To obtain the first claim, it suffices to show Z ′ ( h ) is of order O ( h ) . Then given that e -h κ ≥ 1 -h κ , Z ′ ( h ) satisfies

which proves the first claim.

Then we prove the second claim of the proposition. Note that when h ≤ κ

Thus

As for ρ , denote the normalizing constant of ρ by Z ( h ) = ∫ h 0 1 h -t hκ . Since h = o (1) and κ > 1 ,

where the last inequality follows by that t ≤ h = 0(1) and κ > 1 . Combining (C.2) and (C.3) yields Claim (B).

Then we have

C.2 Convergence lemma

Lemma C.2. Let x ∗ be the minimizer of the potential function, that is, x ∗ = arg min { g ( x ) + f ( x ) } , and A is a positive definite matrix that mI ⪯ A ⪯ LI . Let x n ( t ) and x ∗ n ( t ) be defined as (B.3) and (B.1). Assume that u = 1 L + m . Then we have the following bound

and

Proof. We begin with the first bound. By Young's inequality and Lipschitz smoothness of f ,

Since U ( · ) is m strongly convex and ∥ H x ∥ is a convex function with respect to x , Theorem 1.1 of Hargé [2004] implies that

On the other hand, by Theorem 7 of Basu and DasGupta [1997],

where Σ is the covariance of p and therefore bounded by 1 m . Hence

Combining equation (C.5) and (C.6), we obtain that

Combining (C.4) and (C.7) and then we achieve the first claim.

To bound E ∥ A ( x n ( t ) + v n ( t )) ∥ 2 , we have

where the second inequality follows by the Lipschitz smoothness of f ( x ) and Young's inequality; in the third inequality we use (C.7) and v ∼ N (0 , 1 m ) .

To prove the third claim, we have

D Technical bounds

In this section, we present some controls that are useful in the proof of Theorem 4.2.

Lemma D.1. Let x n ( t ) , ̂ x n ( t ) be defined as above and u = 1 L + m . For any fixed t, α ≤ h .

(i) When Assumption 3.1 hold and H is a uniform upper bound of ∇ 2 f , i.e. ∇ 2 f ( x ) ⪯ H for all x ∈ R d , x ∗ n ( t ) and ̂ x n ( α ) satisfies

(ii) When Assumptions 3.1, E.1 and E.2 hold, we have

Proof. For any x , y ∈ R d , we have

Denote ¯ H = ∫ 1 0 ∇ 2 f ( x + k ( y -x ))d k and ¯ H depends on x and y . Substitute x , y by x n ( t ) , ̂ x n ( α ) , and then we obtain ∥∇ f ( x n ( t )) -∇ f ( ̂ x n ( α )) ∥ 2 = ∥ ¯ H ( x n ( t ) -̂ x n ( α )) ∥ 2 . Then by plugging in the close-form solution into ∥ ¯ H ( x n ( t ) -̂ x n ( α )) ∥ 2 ,

where in a = we use Lemma 4.1; b ≲ follows by the Jensen's inequality. Besides Lemma A.1 implies that ( A 11 ( t ) -A 11 ( α ) -A 12 ( t ) + A 12 ( α )) 2 ∥ H x n ∥ 2 ≲ h 2 ∥ H x n ∥ 2 and ( A 12 ( t ) -A 12 ( α )) 2 ∥ H ( x n + v n ) ∥ 2 ≲ h 2 ∥ H ( x n + v n ) ∥ 2 .

For the expectation of term ① , we have

where in the last inequality we use that t, α < h and A 12 ( t ) = O ( t ) . Let ¯ H ′ = ∫ 1 0 ∇ 2 f ( ̂ x n ( α ) + k ( ̂ x n ( α ) -x n ))d k . By the definition of ̂ x n ( α ) , E ∥ ¯ H ( ∇ f ( ̂ x n ( α )) -∇ f ( x n )) ∥ 2 can be upper bounded by

where in the last inequality we apply Lemma A.1. Note that x n (0) = x n and v n (0) = v n . Then by Lemma C.2,

where in the last inequality we use u = 1 L + m ≤ min { 1 L , 1 m } . Plugging (D.4) into (D.3) and then we have an upper bound of the expectation of term ① . Plugging the upper bound of term ① into (D.2),

and thus we have the following inequality

where in the last inequality we use Lemma C.2 and the definition of Ω n . When we assume ∇ 2 f ( x ) has a uniform upper bound H , we have

and

Besides we have tr ( ( ¯ H ′ ) 2 ) +tr ( ¯ H 2 ) ≤ L tr ( ¯ H ′ ) + L tr ( ¯ H ) ≲ L tr( H ) . These bounds yield the control

And when f is L 2 -Hessian smooth, by Lemmas D.2 and D.3, we have

Lemma D.2. Let ¯ H be defined as above. When Assumptions E.1 and E.2 hold, we have

Proof. First, we have

where a ≲ follows by dividing ¯ H into E [ ¯ H|F n ] and ¯ HE [ ¯ H|F n ] . In the last inequality, E [ ¯ H|F n ] is independent of B s , and thus

As for the firs term, let { B ′ s } 0 ≤ s ≤ h be a standard Brownian motion which is independent of { B s } 0 ≤ s ≤ h . Define ̂ x ′ n ( α ) = J ( α, x n ; { B ′ s } 0 ≤ s ≤ α , ( x n , v n )) and x ′ n ( t ) = J ( α, ̂ x ′ n ( α ); { B ′ s } 0 ≤ s ≤ t , ( x n , v n )) . Then we have

where a ≤ follows by the definition of ¯ H and Jensen's inequality; in b ≤ we use the Hessian smoothness of f ; c ≲ follows by the definition of x n ( t ) , x ′ n ( t ) , ̂ x n ( α ) and ̂ x ′ n ( α ) . For the first term of the last inequality, by the definition of ̂ x n ( α ) and ̂ x ′ n ( α ) , we have

And thus we have

Plugging (D.11) and (D.9) in to control (D.8) yields

where the last inequality follows by Assumption E.2 and the observation that E tr ( ( E [ ¯ H|F n ] ) 2 ) ≤ L EE [tr ( ¯ H ) |F n ] ≤ LM .

Lemma D.3. Let ¯ H and ¯ H ′ be defined as above. When Assumptions E.2 and E.1 hold, we have

Proof. Similar to the proof of Lemma D.2, we have

We follow the notation of B ′ s and ̂ x ′ n ( α ) in (D.10). Then

Combining (D.12) and (D.13) and using tr ( E [ ¯ H ′ |F n ] ) ≤ LM yields

To prove Lemma 5.3, we consider the following Lemma which also includes the Hessian smooth case.

Lemma D.4. Let x n ( t ) , v n ( t ) , v ∗ n ( t ) and v ∗ n ( t ) be defined as above and u = 1 L + m .

(i) Under Assumption 3.1, for any t, α ≤ h , we have

(ii) Under Assumptions 3.1, E.1 and E.2, for any t, α ≤ h , we have

Proof. Since the discretized process and the continuous process are synchronously coupled,

Hence the square norm satisfies

where a ≤ follows by equation [D.15] and Jensen's inequality; in b ≤ we use | 1 -um | ≤ 1 and substitute the square norm by Ω n ( s ) . To bound the second term in the last inequality, we have

where in the second inequality, we use the Lipschitz smoothness of f and ∥ x n ( t ) -x ∗ n ( t ) ∥ ≤ Ω n ( t ) . By combining bound (D.16) and (D.17), we arrive at that

By (C.3), we have h 2 ∫ h 0 1 h E Ω n ( t )d t = h 2 E t ∼ unif[0 ,h ] E Ω n ( t ) ≲ h 2 E t ∼ ρ E Ω n ( t ) . When the Hessian of f is upper bounded by H , by (D.6), we have

which completes the first claim.

When the f is L 2 -Hessian smooth, by (D.7), we have

which establishes the second claim.

Lemma D.5. Let x n ( t ) and ̂ x n ( t ) be defined as above and u = 1 L + m . Then,

(i) When Assumption 3.1 holds, and H be defined as Definition 3.2. Then

(ii) When Assumptions 3.1, E.1 and E.2 hold, we have

Proof. Note that x n ( t ) depends on α . First, we have

Define ¯ H ′′ = ∫ 1 0 ∇ 2 f ( x n ( s ) + k ( ̂ x n ( s ) -x n ( s ))) and recall that ¯ H ′ = ∫ 1 0 ∇ 2 f ( ̂ x n ( α ) + k ( x n -̂ x n ( α )))d k . And then E [ ∥∇ f ( x n ( s )) -∇ f ( ̂ x n ( s )) ∥ 2 ∣ ∣ α, s ] can be bounded by

where a = is by plugging in the close-form solution of x n ( s ) and ̂ x n ( s ) ; b ≤ follows by A 12 ( t -s ) ≲ h and the Lipschitz continuous of ∇ f .

where a ≲ follows by (D.4). When ∇ 2 f has a uniform upper bound H , we have

where in the last inequality, we use that E tr( ¯ H ′ 2 ) ≤ L E tr( ¯ H ′ ) ≤ L tr( H ) and tr( H 2 ) ≤ L tr( H ) . And when Assumptions E.1 and E.2 hold, by Lemma D.3,

E Local Hessian bound results

As discussed in Section 4.1, here we show that the uniform Hessian upper bound can be relaxed to local trace control with an additional Hessian smooth assumption. Specifically, we make the following assumptions on f .

Assumption E.1. f has a L 2 -Lipschitz Hessian.

Assumption E.1 is a standard assumption that frequently appears in the literature of optimization and sampling [Ma et al., 2021, Dalalyan and Riou-Durand, 2020, Durmus and Moulines, 2019].

Assumption E.2. There exist M > 0 such that for all x ∈ R , tr ( ∇ 2 f ( x )) ≤ M .

Assumption E.2 relax the uniform Hessian bound ∇ 2 f ( x ) ⪯ H for x ∈ R d under which we obtain a convergence rate with tr( H ) . Assumption E.2 only requires that the trace of Hessian can be controlled locally. Under Assumptions E.1 and E.2, the convergence of DRUL only has a weak dimension dependency, as stated in Theorem E.3.

Theorem E.3. For any tolerance ϵ ∈ (0 , 1) , denote the minimizer of U ( x ) by x ∗ . Assume that Assumptions 3.1, E.1 and E.2 hold. Set the step size

where C 2 ≥ 1 is a universal constant. With initial point ( x 0 , v 0 ) , define Ω 0 = E x ∼ p, v ∼N (0 ,u ) ( ∥ x 0 -x ∥ 2 + ∥ x 0 + v 0 -v -x ∥ 2 ) . Then under Assumptions 3.1, when

Algorithm 1 outputs x n such that W 2 (Law( x n ) , p ) ≤ ϵ .

Without loss of generality, assume x ∗ = 0 . Then Theorem E.3 indicates a convergence rate of ˜ O ( M 1 / 3 ϵ -2 / 3 + d 2 / 7 ϵ -2 / 7 ) , which has a lower dimension dependency both in the sense of the order of d and the relevant ϵ coefficient. When the Hessian of any x ∈ R d has a dimension-free trace, we arrive at a convergence rate of O ( ϵ -2 / 3 + d 2 / 7 ϵ -2 / 7 ) .

E.1 Proof of Theorem E.3

In this section, we prove the convergence guarantee using the upper bound of the local Hessian trace.

Proof. We follow the notations in the proof of Theorem 4.2 in Appendix B.2. Similarly, denote the error in (B.5) by

Then

where a ≲ follows by Lemmas D.1, D.4 and D.5. We shall note that E Ω n +1 = E β ∼ ρ Ω n ( β ) . Combining the above control with (B.5), we have

where the last inequality follows by the control (B.8) and the choice of h which satisfies 1 -C 2 h 2 ≥ 1 2 . Denote r = 1 -h 6 κ and E ′ = 2 C 2 h 4 ( u 2 L m M + ∥ x ∗ ∥ 2 + L 2 2 u 4 h 4 d 2 ) and therefore

When N ≥ 6 κ log ( 2 E Ω 0 ϵ 2 ) , we have

And when

we get E 1 -r ≤ ϵ 2 2 . Together with the (E.1), we obtain the convergence rate in Theorem E.3.

Finally, we would say that the trace of Hessian should be a more realistic factor to characterize the complexity of Langevin algorithms. Beyond the ridge separable functions described in Section 3.3, we illustrate the Bayesian sampling over the two-layer neural network model and show that the neural network also has a bounded trace under suitable conditions. Note that our convergence analysis cannot be applied to this task since neural networks are non-convex models, and we leave the extensions of analysis to the general non-log-concave case as future work.

Proposition E.4. Define f ( W , w ) = w ⊤ σ ( W ⊤ x ) , where σ is the activation function. When ∥ x ∥ 1 ≤ r 1 , ∥ w ∥ ≤ r 2 and σ ′′ ( x ) ≤ α , we have tr ( ∇ 2 f ( W , w ) ) ≤ αr 1 r 2 .

Proof. By direct computation, we have

Therefore,

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of sampling from a Gibbs distribution $p(x) \\propto e^{-U(x)}$ using discretized Langevin dynamics. Since the approximation error of such methods usually depends on $\\mathrm{Tr}(\\nabla^2 U)$, the proposed algorithm splits $U$ into a quadratic part $g(x) = \\frac m2 ||x||^2$ and a remainder $f$, and only discretizes the dynamics according to $f$.\n\nSuch a split was already considered in [Freund et al. '21]; however, the main novelty of this article is that the discretization on $f$ is performed through a two-step method with random step-sizes, instead of a simple gradient update. This scheme shaves off a factor of $\\epsilon^{-1/3}$ from the required time complexity to reach a Wasserstein error of $\\epsilon$.\n\nThe proof is based on an approximate contraction bound for a quantity $\\Omega_n$ that bounds the desired Wasserstein distance, followed by a precise analysis and optimization of the error in the aforementioned bound.\n\nSTRENGTHS:\nThis paper presents a novel algorithm for Langevin simulation, that achieves state-of-the-art performance for the dependency on both the precision requirement $\\epsilon$ and the ambient dimension $d$. The proposed algorithm is fairly simple (at least for quadratic $g$) and easy to implement, and the main ideas behind it are clearly explained. Overall, the paper is fairly well-written, with only a few typos here and there.\n\nWEAKNESSES:\nClarity/soundness: some of the proofs are very hard to parse due to the amount of simplifications made at once from one line to the next. This is especially felt in B.2, where the first two inequalities below l.447 contain around 5 sub-inequalities to check each. The specifications $\\alpha \\sim \\rho'$ and $\\beta \\sim \\rho$ are also used inconsistently, which makes it hard to understand with respect to what quantities each expectation is taken.\n\nNovelty/significance: The relationship to [Shen and Lee '19] and [Freund et al. '22] should probably be expanded. From what I understand, this paper unites the randomized midpoint method of the former with the combined optimization viewpoint of the latter, but it is unclear if there are challenges other than computational to this endeavor. Namely, neither of those methods require such a complicated step-size scheme, which seems to be the main novelty of the paper, but the need for it is unclear.\n\nMinor remarks:\n- eq. (3.5): what is A?\n- l.185: \"squre\"\n- the equation below l.279 should be a scalar product.\n- in (B.4), shouldn't $z_n(t)$ be $\\hat x_n(\\alpha) - x_n^*(t)$ instead ? This change does ripple through the proof of Lemma 5.1, so I'm not actually sure of how minor it is.\n\nQUESTIONS:\n- How important is the choice of $\\rho$ ? I understand from Lemma 5.2 and its implications that for a given $\\rho$, the choice of $\\rho'$ is important to ensure these conditions, but why can't I, for example, choose a uniform prior for $\\beta$?\n- In section 5.2, it seems like you take $\\bar w_n = x_n - x_n^* + v_n - v_n^*$ as the expectation of $w_n(s, \\alpha)$; why is it the case ?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper suggests a novel version of the Unadjusted Langevin Algorithm with sample complexity in Wasserstein-2 distance scaling with the effective dimension of the problem (trace of the potential's Hessian) instead of the ambient space dimension in case of strongly log-concave distributions. This result completes and generalizes the results of [Shen and Lee, 2019].\n\nSTRENGTHS:\nThe research direction towards studying the sample complexity rates in terms of the effective dimension is interesting and potentially allows to explain the successful behaviour of the ULA-type algorithms in the high-dimensional problems.\n\nWEAKNESSES:\nFirst of all, the suggested Algorithm 1 (DRUL) does not seem to be really an implementable one, since the discretization error in $x_{n+1}$ and $v_{n+1}$ is to appear in the practical implementations. It is not clear if the control of this discretization error would not yield an explicit dimension dependence in the stepsize $h$ in Theorem 4.2. \n\nSecond, the sample complexity scaling as $\\varepsilon^{-2/3}$ is not completely convincing. For example, for the ridge separable potentials, which are the main motivating example towards the paper, the Hamiltonian Monte-Carlo method is known to obtain a sample complexity of order $(d/\\varepsilon)^{1/4}$, see e.g. [Mangoubi et al, 2017]. Thus there is a natural question if the $\\varepsilon^{-2/3}$ complexity optimal for Langevin-type algorithms? Again it seems that the particular rate can degrade after the intergral discretization in Alg. 1 taken into account.\n\nReferences:\nMangoubi, O., & Smith, A. (2017). Rapid mixing of Hamiltonian Monte Carlo on strongly log-concave distributions. arXiv preprint arXiv:1708.07114.\n\nQUESTIONS:\n1. Is it possible to add any numerical findings illustrating the superiority of the doubly randomized ULA (with random step size) against the one with constant or decreasing step size? If one could trace the precise dependence upon the $\\operatorname{trace}{H}$ even in the toyish setup, I would lean towards increasing my score.\n \n2. Are there any novel technical contributions developed to prove the result of Theorem 4.2? If yes, please add the corresponding discussion to the main text.\n\nLIMITATIONS:\nThe paper is theoretical and no negative societal impact is expected.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper adapts the Randomized Midpoint Method to the composite optimization context considered in Freund et al, and consequently improves the dependence from $O(tr(H)/\\epsilon)$ to $O((tr(H)/\\epsilon)^{1/3})$.\n\nSTRENGTHS:\nThe application of randomized midpoint in this composite sampling is novel, and there is genuine improvement in the rate estimate when compared to Freund et al.\n\nThe technique of double randomization is new and requires some novel analysis when contrasted with prior works.\n\nThe authors do a decent job of illustrating that the trace of the Hessian is $o(d)$ through some figures and discussion.\n\nWEAKNESSES:\nThe primary contributions of this paper are not particularly original and mostly stem from combining the framework in Freund et al. with the known analysis for randomized midpoint in Shen and Lee. This in my view is the primary weakness of the paper.\n\nIn general, claims about the “dimension-free” nature of the convergence guarantees need to be careful since the composite structure assumption is quite strong, although the authors in general do a good job of qualifying their claims.\n\nOverall, while this work contains some novel claims and results, and is a bona fide improvement on prior work. However, the technical novelty is not significant, and I am borderline on this paper as a result.\n\nQUESTIONS:\nIt seems inappropriate to compare this to the original randomized midpoint work/other work for the standard Langevin Monte Carlo, since they do not assume the composite structure of the problem. The primary highlighted comparison is with respect to Freund et al. (2022), in which case the gain is more like $tr(H)^{1/3}/\\epsilon^{4/3}$. If $tr(H)$ is $O(1)$ then this is only a gain in epsilon, which is usually smaller than $d$.\n\nWhat is the previous proof referred to in L. 253?\n\nTerminology of “acceleration” should probably be avoided since this is classically used only to refer to sqrt(kappa) rates. A better term might be “improved discretization error”.\n\nTypos:\n\nL. 140 What is being made strongly convex?\n\nL. 185 squred -> squared\n\nL. 197 denote solution -> denote the solution\n\nLIMITATIONS:\nI have outlined my concerns already in the previous sections.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose a Langevin-type algorithm for sampling a strongly log-concave distribution with a composite structure. Their method can be viewed as a variant of the randomized midpoint method, with two key modifications: (i) they only discretize the smooth convex part of the negative log likelihood but retain the strongly convex part; (ii) they draw both step sizes in the algorithm randomly according to carefully crafted distributions. It is shown that the algorithm achieves an accelerated rate without explicit dependence on the dimension.\n\nSTRENGTHS:\n- The result is interesting and noteworthy. To sample a strongly log-concave distribution, the best-known iteration complexity bound either has an undesirable dimension dependence, such as $\\tilde{O}(\\frac{d^{1/3}}{\\epsilon^{2/3}})$ in [Shen and Lee, 2019], or a worse dependence on $\\epsilon$, such as $\\tilde{O}(\\frac{\\mathrm{tr}(H)}{\\epsilon})$ in [Freund et al., 2022]. In this work, the authors manage to achieve the best of both worlds and prove a complexity bound of $\\tilde{O}(\\frac{(\\mathrm{tr}(H))^{1/3}}{\\epsilon^{2/3}})$. \n- The authors introduce the double randomized technique to reduce the discretization error, which seems novel to me.\n\nWEAKNESSES:\nI think the presentation of the paper can be improved.\n\n- In particular, the explanation in Section 5 is not very helpful, and it remains unclear to me why the randomized step size helps reduce the discretization error, and why the authors choose the specific distribution in Lemma 5.2. Moreover, it would be helpful if the authors can compare their analysis with the one in [Shen and Lee, 2019] to better explain how they remove the dimension dependence.  \n\n-  Also, there are numerous typos in the main text and the proofs in the appendix, which sometimes make it hard to understand. Please see the \"Questions\" section for more details.\n\nQUESTIONS:\n- While this work focuses on the dimension dependence, the condition number $\\kappa = L/\\mu$ can also impact the convergence rate greatly. How is the result in this paper compared with the existing works in terms of the dependence on $\\kappa$?\n- In Lemma 5.1, it is unclear to me what it means that \"$x_n$ and $x_n^*$ are coupled synchronously\". Do you mean that they are driven by the same Brownian process?\n- Page 9, Lines 277-281: I am confused by this paragraph. By definition, isn't the random weight $w_n(s,\\alpha)$ a $d$-dimension vector? If so, how can you apply Claim (A) in Lemma 5.2?\n- Page 9, Lines 282-292: It is also unclear to me why \"the randomized step size make it possible to consider the averaged effect\".\n- Page 12, (A.1): I am not sure why the authors introduce the extra parameter $\\gamma$. As far as I can see, $\\gamma$ is fixed as 2 in the rest of the proof.  \n- Page 14, Section B.1.1: I don't see why the random process $B_t^{\\alpha}$ is a Brownian bridge. Is it supposed to be the random process $B_t$ conditioned on $B_{\\alpha}$? And why do we need to introduce this process in the first place?\n- Page 14, the equations under Line 431: Here the authors exchange the order of differentiation and expectation, which is not justified. Indeed, it is not even clear if $\\Omega(t)$ is differentiable since it involves the solution trajectories of SDEs. \n\n\nTypos in the paper:\n\n- The convergence rates reported in the introduction are inconsistent with the ones in Table 1. Specifically, the rate by [Shen and Lee, 2019] should be $\\tilde{O}(\\frac{d^{1/3}}{\\epsilon^{2/3}})$ (Page 2, Line 47), and the rate by [Freund et al., 2022] should scale linearly with $O(\\frac{1}{\\epsilon})$ (Table 1, Row 5). \n- Page 1, Line 49: \"convergence dependence\" -> \"dimension dependence\"\n- Page 5, Definition 3.6: $A$ is undefined. \n- Page 9, Line 275: the integral should be $\\int_{0}^t e^{\\frac{s-t}{\\kappa}} F(s) ds$.\n\nLIMITATIONS:\nYes, the authors addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}