A CCELERATING H AMILTONIAN M ONTE C ARLO VIA C HEBYSHEV I NTEGRATION T IME
A BSTRACT
Hamiltonian Monte Carlo (HMC) is a popular method in sampling.
While there are quite a few works of studying this method on various aspects, an interesting question is how to choose its integration time to achieve acceleration.
In this work, we consider accelerating the process of sampling from a distribution π ( x ) ∝ exp( − f ( x )) via HMC via time-varying integration time.
When the potential f is
L -smooth and m -strongly convex, i e. for sampling from a log-smooth and strongly log-concave target distribution π , it is known that under a constant integration time, the number of iterations that ideal HMC takes to get an (cid:15) Wassersteindistance to the target π is O ( κ log 1 (cid:15) ) where κ :=
Lm is the condition number.
We propose a scheme of time-varying integration time based on the roots of Chebyshev polynomials.
We show that in the case of quadratic potential f , i.e. when the target π is a Gaussian distribution, ideal HMC with this choice of integration time only takes O ( √ κ log 1 (cid:15) ) number of iterations to reach Wasserstein-2 distance less than (cid:15) ; this improvement on the dependence on condition number is akin to acceleration in optimization.
The design and analysis of HMC with the proposed integration time is built on the tools of Chebyshev polynomials.
Experiments ﬁnd the advantage of adopting our scheme of time-varying integration time even for sampling from distributions with smooth strongly convex potentials that are not quadratic
I NTRODUCTION
Markov chain Monte Carlo (MCMC) algorithms are fundamental techniques for sampling from probability distributions, which is a task that naturally arises in statistics (Duane et al., 1987;
Girolami & Calderhead, 2011), optimization (Flaxman et al., 2005; Duchi et al., 2012; Jin et al., 2017), machine learning and others (Wenzel et al., 2020; Salakhutdinov & Mnih, 2008; Koller & Friedman, 2009;
Welling & Teh, 2011).
Among all the MCMC algorithms, the most popular ones perhaps are Langevin methods (Li et al., 2022; Dalalyan, 2017; Durmus et al., 2019; Vempala & Wibisono, 2019; Lee et al., 2021b;
Chewi et al., 2020) and Hamiltonian Monte Carlo (HMC) (Neal, 2012; Betancourt, 2017; Hoffman & Gelman, 2014; Levy et al., 2018).
For the former, recently there have been a sequence of works leveraging some techniques in optimization to design Langevin methods, which include borrowing the idea of momentum methods like Nesterov acceleration (Nesterov, 2013) to design fast methods, e.g., (Ma et al., 2021; Dalalyan & Riou-Durand, 2020).
Speciﬁcally, Ma et al.
(2021) show that for sampling from distributions satisfying the log-Sobolev inequality, under-damped
Langevin improves the iteration complexity of over-damped Langevin from O ( d(cid:15) ) to O ( (cid:113) d(cid:15) ) , where d is the dimension and (cid:15) is the error in KL divergence, though whether their result has an optimal dependency on the condition number is not clear.
On the other hand, compared to Langevin methods, the connection between HMCs and techniques in optimization seems rather loose.
Moreover, to our knowledge, little is known about how to accelerate HMCs with a provable acceleration guarantee for converging to a target distribution.
Speciﬁcally, Chen & Vempala (2019) show that for sampling from strongly log-concave distributions, the iteration complexity of ideal HMC is O ( κ log 1 (cid:15) ) , and Vishnoi (2021) shows the same rate of ideal HMC when the potential is strongly convex quadratic in a nice tutorial.
In contrast, there are a few methods that exhibit acceleration when minimizing strongly convex quadratic functions in optimization.
For example, while Heavy Ball (Polyak, 1964) does not have an accelerated linear rate globally for minimizing general smooth strongly convex functions, it does show acceleration when minimizing strongly convex quadratic functions (Wang et al., 2020;
This observation makes us wonder whether one can get an accelerated linear rate of ideal HMC for sampling, i.e., O ( √ κ log 1 (cid:15) ) , akin to acceleration in optimization.
We answer this question afﬁrmatively, at least in the Gaussian case.
We propose a time-varying integration time for HMC, and we show that ideal HMC with this time-varying integration time exhibits acceleration when the potential is a strongly convex quadratic (i.e. the target π is a Gaussian), compared to what is established in Chen & Vempala (2019) and Vishnoi (2021) for using a constant integration time.
Our proposed time-varying integration time at each iteration of HMC depends on the total number of iterations K , the current iteration index k , the strong convexity constant m , and the smoothness constant L of the potential; therefore, the integration time at each iteration is simple to compute and is set before executing HMC.
Our proposed integration time is based on the roots of Chebysev polynomials, which we will describe in details in the next section.
In optimization,
Chebyshev polynomials have been used to help design accelerated algorithms for minimizing strongly convex quadratic functions, i.e., Chebyshev iteration (see e.g., Section 2.3 in d’Aspremont et al.
Our result of accelerating HMC via using the proposed Chebyshev integration time can be viewed as the sampling counterpart of acceleration from optimization.
Interestingly, for minimizing strongly convex quadratic functions, acceleration of vanilla gradient descent can be achieved via a scheme of step sizes that is based on a Chebyshev polynomial, see e.g., Agarwal et al.
and our work is inspired by a nice blog article by Pedregosa (2021).
Hence, our acceleration result of
HMC can also be viewed as a counterpart in this sense.
In addition to our theoretical ﬁndings, we conduct experiments of sampling from a Gaussian as well as sampling from distributions whose potentials are not quadratics, which include sampling from a mixture of two Gaussians, Bayesian logistic regression, and sampling from a hard distribution that was proposed in Lee et al.
(2021a)
for establishing some lower-bound results of certain Metropolized sampling methods.
Experimental results show that our proposed time-varying integration time also leads to a better performance compared to using the constant integration time of Chen & Vempala (2019) and Vishnoi (2021) for sampling from the distributions whose potential functions are not quadratic.
We conjecture that our proposed time-varying integration time also helps accelerate HMC for sampling from log-smooth and strongly log-concave distributions, and we leave the analysis of such cases for future work.
P RELIMINARIES 2. H AMILTONIAN M ONTE C ARLO (HMC)
Suppose we want to sample from a target probability distribution ν ( x ) ∝ exp( − f ( x )) on R d , where f : R d → R is a continuous function which we refer to as the potential.
Denote x ∈ R
d
the position and v ∈ R
d the velocity of a particle.
In this paper, we consider the standard Hamiltonian of the particle (Chen & Vempala, 2019; Neal, 2012), which is deﬁned as
[Equation]
(1) while we refer the readers to Girolami & Calderhead (2011); Hirt et al.
; Brofos & Lederman (2021) and the references therein for other notions of the Hamiltonian.
The Hamiltonian ﬂow generated by H is the ﬂow of the particle which evolves according to the following differential equations:
[Equation]
For the standard Hamiltonian deﬁned in (1), the Hamiltonian ﬂow becomes dx dt = v
and dv dt = −∇ f ( x ) .
We will write ( x t , v t ) =
HMC t ( x 0 , v 0 ) as the position x and the velocity v of the Hamiltonian ﬂow after integration time t starting from ( x 0 , v 0 ) .
There are many important properties of the Hamiltonian ﬂow including that the Hamiltonian is conserved along the ﬂow, the vector ﬁeld associated with the ﬂow is divergence free, and the Hamiltonian dynamic is time reversible, see e.g., Section 3 in Vishnoi (2021).
The Ideal HMC algorithm (see Algorithm 1) proceeds as follows: in each iteration k , sample an initial velocity from the normal distribution, and then ﬂow following the Hamiltonian ﬂow with a pre-speciﬁed integration time η k .
It is well-known that ideal HMC preserves the target density π ( x ) ∝ exp( − f ( x )) ; see e.g., Theorem 5.1 in Vishnoi (2021).
Furthermore, in each iteration, HMC brings the density of the iterates x k ∼ ρ k closer to the target π .
However, the Hamiltonian ﬂow
HMC t ( x 0 , v 0 ) is in general difﬁcult to simulate exactly, except for some special potentials.
In practice, the Verlet integrator is commonly used to approximate the ﬂow and a Metropolis-Hastings ﬁlter is applied to correct the induced bias arises from the use of the integrator (Tripuraneni et al., 2017; Brofos & Lederman, 2021; Hoffman et al., 2021; Lee et al., 2021a; Chen et al., 2020).
In recent years, there have been some progress on showing some rigorous theoretical guarantees of HMCs for converging to a target distribution, e.g., Chen et al.
; Durmus et al.
(2017); Bou-Rabee & Eberle (2021); Mangoubi & Smith (2019; 2021); Mangoubi & Vishnoi (2018).
There are also other variants of HMCs proposed in the literature, e.g., Riou-Durand & Vogrinc (2022); Bou-Rabee & Sanz-Serna (2017); Zou & Gu (2021); Steeg & Galstyan (2021); Hoffman & Gelman (2014);
Tripuraneni et al.
; Chen et al.
(2014), to name just a few.
Recall that the 2-Wasserstein distance between probability distributions ν
and ν 2 is
[Equation]
where Γ( ν 1 , ν 2 ) represents the set of all couplings of ν 1 and ν 2 .
2. A NALYSIS OF HMC IN QUADRATIC CASE WITH CONSTANT INTEGRATION TIME
In the following, we replicate the analysis of ideal HMC with a constant integration time for quadratic potentials (Vishnoi, 2021), which provides the necessary ingredients for introducing our method in the next section.
Speciﬁcally, we consider the following quadratic potential:
[Equation]
which means the target density is the Gaussian distribution
π = N (0 , Λ − 1 ) , where Λ the diagonal matrix whose j th diagonal entry is λ j .
We note for a general Gaussian target N ( µ, Σ) for some µ ∈ R d and Σ (cid:31) 0 , we can shift and rotate the coordinates to make µ = 0 and Σ a diagonal matrix, and our analysis below applies.
So without loss of generality, we may assume the quadratic potential is separable, as in (3).
In this quadratic case, the Hamiltonian ﬂow (2) becomes a linear system of differential equations, and we have an exact solution given by sinusoidal functions, which are
[Equation]
In particular, we recall the following result on the deviation between two co-evolving particles with the same initial velocity.
Lemma 1.
(Vishnoi, 2021)
Let x 0 , y 0 ∈ R d .
Consider the following coupling: ( x t , v t ) =
HMC t ( x 0 , ξ ) and ( y t , u t ) =
HMC t ( y 0 , ξ ) for some ξ ∈ R d .
Then for all t ≥ 0 and for all j ∈ [ d ] , it holds that
[Equation]
Using Lemma 1, we can derive the convergence rate of ideal HMC for the quadratic potential as follows.
Lemma 2.
(Vishnoi, 2021) Let π ∝ exp( − f ) =
N (0 , Λ − 1 ) be the target distribution, where f ( x ) is deﬁned on (3).
Let ρ K be the distribution of x K generated by Algorithm 1 at the ﬁnal iteration K .
Then for any ρ 0 and any K ≥ 1 , we have
[Equation]
We replicate the proof of Lemma 1 and Lemma 2 in Appendix B for the reader’s convenience.
Vishnoi (2021) shows that by choosing ( Constant integration time )
[Equation]
one has that cos (cid:16)(cid:112) 2 λ j η ( K ) k
(cid:17) ≤ 1 − Θ (cid:0) mL (cid:1) for all the iterations k ∈
[ K ] and dimensions j ∈ [ d ] .
Hence, by Lemma 2, the distance satisﬁes
[Equation]
after K iterations of ideal HMC with the constant integration time.
On the other hand, for general smooth strongly convex potentials f ( · ) , Chen & Vempala (2019) show the same convergence rate 1 − Θ (cid:0) mL (cid:1) of HMC using a constant integration time
η ( K ) k = c √ L , where c > 0 is a universal constant.
Therefore, under the constant integration time, HMC needs O ( κ log 1 (cid:15) ) iterations to reach error W 2 ( ρ K , π ) ≤ (cid:15) ,
where κ =
Lm is condition number.
Furthermore, they also show that the relaxation time of ideal HMC with a constant integration time is Ω( κ ) for the Gaussian case.
2. C HEBYSHEV POLYNOMIALS
We denote Φ K ( · ) the degree- K Chebyshev polynomial of the ﬁrst kind , which is deﬁned by:
[Equation]
Our proposed integration time is built on a scaled-and-shifted Chebyshev polynomial, deﬁned as:
[Equation]
where h ( · ) is the mapping h ( λ ) :
=
L + m − 2 λ
L − m .
Observe that the mapping h ( · ) maps all λ ∈ [ m, L ] into the interval [ − 1 , 1] .
The roots of the degree- K scaled-and-shifted Chebyshev polynomial ¯Φ K ( λ ) are (Chebyshev roots)
[Equation]
where k
= 1 , 2 , . . .
, K , i.e., ¯Φ K ( r ( K ) k )
= 0 .
We now recall the following key result regarding the scaled-and-shifted Chebyshev polynomial ¯Φ K .
Lemma 3.
(e.g., Section 2.3 in d’Aspremont et al.
For any positive integer K , we have
[Equation]
The proof of Lemma 3 is in Appendix B.
C HEBYSHEV INTEGRATION TIME
We are now ready to introduce our scheme of time-varying integration time.
Let K be the pre-speciﬁed total number of iterations of HMC.
Our proposed method will ﬁrst permute the array [1 , 2 , . . .
, K ] before executing HMC for K iterations.
Denote σ ( k ) the k th element of the array [1 , 2 , . .
, K ] after an arbitrary permutation σ .
Then, we propose to set the integration time of HMC at iteration k , i.e., set η ( K ) k , as follows: [Equation] [Equation] [Equation]
[Equation]
[Equation]
We note the usage of the permutation σ is not needed in our analysis below; however, it seems to help improve performance in practice.
Speciﬁcally, though the guarantees of HMC at the ﬁnal iteration K provided in Theorem 1 and Lemma 4 below is the same regardless of the permutation, the progress of HMC varies under different permutations of the integration time, which is why we recommend an arbitrary permutation of the integration time in practice.
Our main result is the following improved convergence rate of HMC under the Chebyshev integration time, for quadratic potentials.
Theorem 1.
Denote the target distribution π ∝ exp( − f ( x )) =
N (0 , Λ − 1 ) , where f ( x ) is deﬁned on (3), and denote the condition number κ :=
Lm .
Let ρ K be the distribution of x K generated by
Algorithm 1 at the ﬁnal iteration K .
Then, we have
[Equation]
Consequently, the total number of iterations K such that the Wasserstein-2 distance satisﬁes
W
[Equation]
K , [Equation] ≤
[Equation]
is
[Equation]
Theorem 1 shows an accelerated linear rate 1 − Θ (cid:16) 1 √ κ (cid:17) using Chebyshev integration time, and hence improves the previous result of 1 − Θ (cid:0) 1 κ (cid:1) as discussed above.
The proof of Theorem 1 relies on the following lemma, which upper-bounds the cosine products that appear in the bound of the Wdistance in Lemma 2 by the scaled-and-shifted Chebyshev polynomial ¯Φ K ( λ ) on (7).
Lemma 4.
Denote | P
Cos K ( λ ) |
[Equation]
(cid:12)(cid:12) .
Suppose λ ∈ [ m, L ] .
Then, we have for any positive integer K ,
[Equation]
The proof of Lemma 4 is available in Appendix C. Figure 1 compares the cosine product max λ ∈ [ m,L ]
[Equation]
[-∇ ¯ x s +
-]Π ks = [Equation]
[-¯ x s + θ ¯ v s + 12 ; ¯ v s +
 = ¯ v-](cid:16) √ 2 λη ( K ) s (cid:17)(cid:12)(cid:12)(cid:12) in Lemma 2 of using the proposed integration time and that
[Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
Compute the ratio α
[Table]
[Table]
[Table]
Else 18:
[Table]
[Table]
end for of using the constant integration time, which illustrates acceleration via the proposed Chebyshev integration time.
We now provide the proof of Theorem 1.
Proof.
(of Theorem 1)
From Lemma 2, we have
[Equation]
[Equation] We can upper-bound the cosine product of any j ∈ [ d ] as,
[Equation]
where (a) is due to the use of Chebyshev integration time (10), (b) is by Lemma 4, and (c) is by
Lemma 3.
Combining (12) and (13) leads to the result.
HMC with Chebyshev Integration Time for General Distributions To sample from general strongly log-concave distributions, we propose Algorithm 2, which adopts the Verlet integrator (a.k.a.
the leapfrog integrator) to simulate the Hamiltonian ﬂow HMC η ( · , ξ ) and uses Metropolis ﬁlter to correct the bias.
It is noted that the number of leapfrog steps S k in each iteration k is equal to the integration time η ( K ) k divided by the step size θ used in the leapfrog steps.
More precisely, we have S k
= (cid:98) η
( K ) k
θ (cid:99) in iteration k of HMC.
E XPERIMENTS
We now evaluate HMC with the proposed Chebyshev integration time (Algorithm 2) and HMC with the constant integration time (Algorithm 2 with line 7 replaced by the constant integration time (5)) in several tasks.
For all the tasks in the experiments, the total number of iterations of HMCs is set to be K = 10 , 000 , and hence we collect K = 10 , 000 samples along the trajectory.
For the step size θ in the leapfrog steps, we let θ ∈ { 0 .
To evaluate the methods,
we
[Table] [Table] [Table]
compute effective sample size (ESS), which is a common performance metric of HMCs (Girolami & Calderhead, 2011; Brofos & Lederman, 2021;
Hirt et al., 2021; Riou-Durand & Vogrinc, 2022;
Hoffman et al., 2021; Hoffman & Gelman, 2014; Steeg & Galstyan, 2021), by using the toolkit ArViz (Kumar et al., 2019).
The ESS of a sequence of N dependent samples is computed based on the autocorrelations within the sequence at different lags: ESS :=
N/ (1 + 2 (cid:80) k γ ( k )) , where γ ( k ) is an estimate of the autocorrelation at lag k .
We consider 4 metrics, which are (1) Mean ESS: the average of ESS of all variables.
That is, ESS is computed for each variable/dimension, and Mean ESS is the average of them.
Min ESS: the lowest value of ESS among the ESSs of all variables; (3)
Mean ESS/Sec.
Mean ESS normalized by the CPU time in seconds; (4) Min ESS/Sec.:
Minimum
ESS normalized by the CPU time in seconds.
In the following tables, we denote “Cheby.”
as our proposed method, and “Const.”
as HMC with the the constant integration time (Vishnoi, 2021;
Chen & Vempala, 2019).
Each of the conﬁgurations is repeated 10 times, and we report the average and the standard deviation of the results.
We also report the acceptance rate of the Metropolis ﬁlter (Acc. Prob) on the tables.
Our implementation of the experiments is done by modifying a publicly available code of HMCs by Brofos & Lederman (2021).
Code for our experiments can be found in the supplementary.
4. I DEAL HMC FLOW FOR SAMPLING FROM A G USSIAN WITH A DIAGONAL COVARIANCE
Before evaluating the empirical performance of Algorithm 2 in the following subsections, here we discuss and compare the use of a arbitrary permutation of the Chebyshev integration time and that without permutation (as well as that of using a constant integration time).
We simulate ideal HMC for sampling from a Gaussian N ( µ, Σ) , where µ =
[Equation]
and
[Equation]
(cid:21) .
It is noted that ideal
HMC ﬂow for this case has a closed-form solution as (4) shows.
The result are reported on Table 1.
From the table, the use of a Chebyshev integration time allows to obtain a larger ESS than that from using a constant integration time, and a arbitrary permutation helps get a better result.
An explanation is that the ESS is a quantity that is computed along the trajectory of a chain, and therefore a permutation of the integration time could make a difference.
We remark that the observation here (a arbitrary permutation of time generates a larger ESS) does not contradict to Theorem 1, since
Theorem 1 is about the guarantee in W 2 distance at the last iteration K .
4. S AMPLING FROM A G AUSSIAN
We sample N ( µ, Σ) , where µ = (cid:20) 01 (cid:21) and Σ = (cid:20)0 .0 .
5(cid:21) .
Therefore, the strong convexity constant m is approximately 0 .
01 and the smoothness constant L is approximately 1 .
Table 2 shows the results.
HMC with Chebyshev integration time consistently outperforms that of using the constant integration time in terms of all the metrics: Mean ESS, Min ESS, Mean ESS/Sec, and Min ESS/Sec.
We also plot two quantities throughout the iterations of HMCs on Figure 2.
Speciﬁcally, Sub-ﬁgure (a) on Figure 2 plots the size of the difference between the targeted covariance Σ and an estimated covariance ˆΣ k at each iteration k of HMC, where ˆΣ k is the sample covariance of 10 , 000 samples collected from a number of 10 , 000 HMC chains at their k th iteration.
Sub-ﬁgure (b) plots a discrete
TV distance that is computed as follows.
We use a built-in function of Numpy to sample 10 ,samples from the target distribution, while we also have 10 , 000 samples collected from a number
[Table]
[Figure]
of 10 , 000 HMC chains at each iteration k .
Using these two sets of samples, we construct two histograms with 30 number of bins for each dimension, we denote them as ˆ π and ˆ ρ k .
The discrete
TV(ˆ π, ˆ ρ k ) at iteration k is 0 .
5 times the sum of the absolute value of the difference between the number of counts of all the pairs of the bins divided by 10 , 000 , which serves as a surrogate of the
Wasserstein-2 distance between the true target π and ρ k from HMC, since computing or estimating the true Wasserstein distance is challenging.
4. S AMPLING FROM A MIXTURE OF TWO G AUSSIANS
For a vector a ∈ R d and a positive deﬁnite matrix Σ ∈ R d × d , we consider sampling from a mixture of two Gaussians N ( a, Σ) and N ( − a, Σ) with equal weights.
Denote b := Σ − 1 a and Λ := Σ − 1 .
The potential is f ( x ) = 1 2 (cid:107) x − a (cid:107) 2 Λ − log(1 + exp( − 2 x (cid:62) b )) , and its gradient is ∇ f ( x ) =
Λ x − b + 2 b
(1 + exp( − 2 x (cid:62) b )) − 1 .
For each dimension i ∈ [ d ] , we set a [ i ] = √ i 2 d and set the covariance Σ = diag 1 ≤ i ≤ d ( id ) .
The potential is strongly convex if a (cid:62) Σ − 1 a < 1 , see e.g.,
Riou-Durand & Vogrinc (2022).
We set d = 10 in the experiment, and simply use the smallest and the largest eigenvalue of Λ to approximate the strong convexity constant m and the smoothness constant
L of the potential, which are ˆ m = 1 and ˆ L = 10 in this case.
Table 3 shows that the proposed method generates a larger effective sample size than the baseline.
4. B AYESIAN LOGISTIC REGRESSION
We also consider Bayesian logistic regression to evaluate the methods.
Given an observation ( z i , y i ) , where z i ∈ R d and y i ∈ { 0 , 1 } , the likelihood function is modeled as p ( y i |
z i , w ) =
[Equation]
Moreover, the prior on the model parameter w is assumed to follow a Gaussian distribution, p ( w ) =
N (0 , α − 1 I d ) , where α > 0 is a parameter.
The goal is to sample w ∈ R d from the posterior, p ( w |{ z i , y i } ni =1 ) = p ( w )Π ni =1 p ( y i |
z i , w ) , where n is the number of data points in a dataset.
The potential function f ( w ) can be written as [Equation] , where
[Equation]
[Table] We set α = 1 in the experiments.
We consider three datasets: Heart, Breast Cancer, and Diabetes binary classiﬁcation datasets, which are all publicly available online.
To approximate the strong convexity constant m and the smoothness constant L of the potential f ( w ) , we compute the smallest eigenvalue and the largest eigenvalue of the Hessian ∇ 2 f ( w ) at the maximizer of the posterior, and we use them as estimates of m and L respectively.
We apply Newton’s method to approximately ﬁnd the maximizer of the posterior.
The experimental results are reported on Table 4 in Appendix E.1 due to the space limit, which show that our method consistently outperforms the baseline.
4. S AMPLING FROM A hard DISTRIBUTION
We also consider sampling from a step-size-dependent distribution π ( x ) ∝ exp( − f h ( x )) , where the potential f h ( · ) is κ -smooth and 1 -strongly convex.
The distribution is considered in Lee et al.
(2021a) for showing a lower bound regarding certain Metropolized sampling methods using a constant integration time and a constant step size h of the leapfrog integrator.
More concretely, the potential is f h ( x ) :
= (cid:80) di
=1 f ( h ) i
( x i ) , where f ( h ) i ( x i ) =
[Equation]
In the experiment, we set κ = 50 and d = 10 .
The results are reported on Table 5 in Appendix E.2.
The scheme of the Chebyshev integration time is still better than the constant integration time for this task.
D ISCUSSION AND OUTLOOK
The Chebyshev integration time shows promising empirical results for sampling from a various of strongly log-concave distributions.
On the other hand, the theoretical guarantee of acceleration that we provide in this work is only for strongly convex quadratic potentials.
Therefore, a direction left open by our work is establishing some provable acceleration guarantees for general strongly log-concave distributions.
However, unlike quadratic potentials, the output (position, velocity) of a
HMC ﬂow does not have a closed-form solution in general, which makes the analysis much more challenging.
A starting point might be improving the analysis of Chen & Vempala (2019), where a contraction bound of two HMC chains under a small integration time η
= O ( 1 √ L ) is shown.
Since the scheme of the Chebyshev integration time requires a large integration time η = Θ (cid:16) 1 √ m (cid:17) at some iterations of HMC, a natural question is whether a variant of the result of Chen & Vempala (2019) can be extended to a large integration time η = Θ (cid:16) 1 √ m (cid:17) .
We state as an open question: can ideal HMC with a scheme of time-varying integration time achieve an accelerated rate O ( √ κ log 1 (cid:15) ) for general smooth strongly log-concave distributions? The topic of accelerating HMC with provable guarantees is underexplored, and we hope our work can facilitate the progress in the literature. After the preprint of this work was available on arXiv, Jiang (2022) proposes a randomized integration time with partial velocity refreshment and provably shows that ideal HMC with the proposed machinery has the accelerated rate for sampling from a Gaussian distribution.
Exploring any connections between the scheme of Jiang (2022) and ours can be an interesting direction.
A CKNOWLEDGMENTS
We thank the reviewers for constructive feedback, which helps improve the presentation of this paper. A A CONNECTION BETWEEN OPTIMIZATION AND SAMPLING To provide an intuition of why the technique of Chebyshev polynomials can help accelerate HMC for the case of the strongly convex quadratic potentials, we would like to describe the work of gradient descent with the Chebyshev step sizes Agarwal et al.
(2021) in more detail, because we are going to draw a connection between optimization and sampling to showcase the intuition.
Agarwal et al.
(2021) provably show that gradient descent with a scheme of step sizes based on the Chebyshev Polynomials has an accelerated rate for minimizing strongly convex quadratic functions compared to GD with a constant step size and their experiments show some promising results for minimizing smooth strongly convex functions beyond quadratics via the proposed scheme of step sizes.
More precisely, deﬁne f ( w ) =
12 w (cid:62)
Aw , where A ∈ R d × d is a positive deﬁnite matrix which has eigenvalues
L := λ 1 ≥ λ 2 ≥ · · · ≥ λ d =: m .
Agarwal et al.
(2021) consider applying gradient descent
[Equation]
to minimize f ( · ) , where η k is the step size of gradient descent at iteration k .
Let w be the unique global minimizer of f ( · ) .
It is easy to show that the dynamic of the distance evolves as
[Equation]
Hence, the size of the distance to w ∗ at iteration K + 1 is bounded by
[Equation]
This shows that the convergence rate of GD is governed by max j ∈ [ d ] | (cid:81)
Kk =1 (1 − η k λ j ) | .
By setting η k as the inverse of the Chebyshev root r ( K ) k or any permuted root r ( K ) σ ( k ) (see (8) for the deﬁnition), the polynomial (cid:81)
Kk =1 (1 − η k λ ) is actually the K -degree scale-and-shifted polynomial, i.e., (cid:81)
Kk =1 (1 − η k λ ) = (cid:81)
Kk = (cid:18)
[Equation]
−
[Equation]
[Equation]
k
[Equation]
(cid:19)
=
¯Φ k ( λ ) (see (7) for the deﬁnition).
It is well-known in the literature of optimization and numerical linear algebra that the K -degree scale-and-shifted polynomial satisﬁes
[Equation]
which is restated in Lemma 3 and its proof is replicated in Appendix B of our paper for the reader’s convenience.
Applying this result, one gets a simple proof of the accelerated linear rate of GD with the proposed scheme of step sizes for minimizing quadratic functions.
A nice blog article by
Pedregosa (2021) explains this in detail.
Now we are ready to highlight its connection with HMC.
In Lemma 1 of the paper, we restate a known result in HMC literature, where its proof is also replicated in Appendix B for the reader’s convenience.
The lemma indicates that the convergence rate of HMC is governed by max
[Equation]
) | .
By way of comparison to that of GD for minimizing quadratic functions, i.e., max j ∈ [ d ] | (cid:81)
Kk =1 (1 − η k λ j ) | , it appears that they share some similarity, which made us wonder if we could bound the former by the latter.
We show in Lemma 4 that cos( π 2 √ x ) ≤ 1 − x , which holds for all x ≥ 0 , and consequently,
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
The key lemma above implies that if we set the integration time as η ( K ) k
=
[Equation]
[Equation]
2 r ( K ) σ ( k ) , then we get acceleration of HMC.
B P ROOFS OF LEMMAS IN S ECTION
We restate the lemmas for the reader’s convenience.
Lemma 1.
(Vishnoi, 2021)
Let x 0 , y 0 ∈ R d .
Consider the following coupling: ( x t , v t ) =
HMC t ( x 0 , ξ ) and ( y t , u t ) =
HMC t ( y 0 , ξ ) for some ξ ∈ R d .
Then for all t ≥ 0 and for all j ∈ [ d ] , it holds that
[Equation]
Proof.
Given ( x t , v t ) :
= HMC t ( x 0 , ξ ) and ( y t , u t ) :
= HMC t ( y 0 , ξ ) , we have dv t dt − du t dt
= −∇ f
( x t ) + ∇ f ( y t ) =
2Λ( y t − x t ) .
Therefore, we have d 2 ( x t [ j ] − y t
[ j ]) dt= − 2 λ j ( x t
[ j ] − y t
[ j ]) , for all j ∈ [ d ] .
Because of the initial condition dx
[ j ] dt = dy
[ j ] dt = ξ
[ j ] , the differential equation implies that x t [ j ] − y t
[ j ] = cos (cid:0)(cid:112)
λ j t (cid:1) × ( x
[ j ] − y
[ j ]) .
It is noted that the result also follows directly from the explicit solution (4).
Lemma 2.
(Vishnoi, 2021) Let π ∝ exp( − f ) =
N (0 , Λ − 1 ) be the target distribution, where f ( x ) is deﬁned on (3).
Let ρ K be the distribution of x K generated by Algorithm 1 at the ﬁnal iteration K .
Then for any ρ 0 and any K ≥ 1 , we have
[Equation]
Proof.
Starting from x 0 ∼ ρ 0 , draw an initial point y 0 ∼ π such that ( x 0 , y 0 ) has the optimal
W 2 -coupling between ρ 0 and π .
Consider the following coupling at each iteration k : ( x k , v k ) =
HMC η ( K ) k
( x k − 1 , ξ k ) and ( y k , u k ) =
HMC η ( K ) k
( y k − 1 , ξ k ) where ξ k ∼ N (0 , I ) is an independent
Gaussian.
We collect { x k } K k =1 and { y k } K k =1 from Algorithm 1.
We know each y k ∼ π , since π is a stationary distribution of the HMC Markov chain.
Then by Lemma 1 we have
[Equation]
Taking the square root on both sides leads to the result.
Lemma 3.
(e [Equation] g., Section 2.3 in d’Aspremont et al.
For any positive integer K , we have
[Equation]
Proof.
Observe that the numerator of ¯Φ K ( [Equation] =
[Equation]
| Φ K ( h ( λ ))
| ≤ 1 , since h ( λ ) ∈ [ − 1 , 1] for λ ∈ [ m, L ] and that the Chebyshev polynomial satisﬁes | Φ K ( · ) | ≤ 1 when its argument is in [ − 1 , 1] by the deﬁnition.
It remains to bound the denominator, which is Φ K ( h (0))
=
cosh (cid:16) K arccosh (cid:16) L + m L − m (cid:17)(cid:17) .
Since
[Equation]
, where θ := √ L +
[Equation]
, we have
[Equation]
Combing the above inequalities, we obtain the desired result:
[Equation]
C P ROOF OF L EMMA
Lemma 4.
Denote | P
Cos
[Equation]
Suppose λ ∈ [ m, L ] .
Then, we have for any positive integer K ,
[Equation]
Proof.
We use the fact that the K -degree scaled-and-shifted Chebyshev Polynomial can be written as,
[Equation]
for any permutation σ ( · ) , since { r ( K ) σ ( k ) } are its roots and ¯Φ K (0)
= 1 .
So inequality (18) is equivalent to
[Equation]
To show (20), let us analyze the mapping ψ ( x ) :
= cos ( π 2 √ x ) 1 − x for x ≥ 0 , x (cid:54) = 1 , with ψ (1) =
π 4 by continuity, and show
that max x : x ≥ 0 | ψ ( x ) | ≤ 1 , as (20) would be immediate.
We have ψ (cid:48)
( x ) = [Equation]
Hence, ψ (cid:48)
( x ) =
when
[Equation]
Denote an extreme point of ψ ( x ) as ˆ x , which satisﬁes (21).
Then, using (21), we have | ψ (ˆ x ) | =
[Equation]
2 (cid:12)(cid:12)(cid:12)(cid:12) , where we used cos( π√ ˆ x ) =
π (1 − ˆ x ) √ 16ˆ x + π 2 (1 − ˆ x ) 2 or − π (1 − ˆ x ) √ 16ˆ x + π 2 (1 − ˆ x ) 2 .
The denominator (cid:112) 16ˆ x + π 2 (1 − ˆ x ) 2 has the smallest value at ˆ x = 0 , which means that the largest value of | ψ ( x ) | happens at x = 0 , which is 1 .
The proof is now completed.
D A COMPARISON OF THE TOTAL INTEGRATION TIME (J IANG , 2022)
Since the Chebyshev integration time are set to be some large values at some steps of HMC, it is natural to ask if the number of steps to get an (cid:15) 2-Wasserstein distance is a fair metric. In this section, we consider the total integration time (cid:80) Kk =1 η ( K ) k to get an (cid:15) distance as another metric for the convergence.
It is noted that the comparison between HMC with our integration time and HMC with the best constant integration time has been conducted by Jiang (2022), and our previous version did not have such a comparison.
Below, we reproduce the comparision of Jiang (2022).
Recall the number of iterations to get an (cid:15) 2-Wasserstein distance to the target distribution is
K = O (cid:0) √ κ log (cid:0) 1 (cid:15) (cid:1)(cid:1) of HMC with the Chebyshev integration time (Theorem 1 in the paper).
The average of the integration time is
[Equation]
where we recall that a permutation σ ( · ) does not affect the average.
Then, if K is even, we can rewrite the averaged integration time as
[Equation]
Otherwise, K is odd, and we can rewrite the averaged integration time as
[Equation]
We will show
[Equation]
for any k = { 1 , 2 , . . .
, (cid:98) K 2 (cid:99)} soon.
Given this, we can further upper-bound the averaged integration time as
[Equation]
when K is even; when K is odd, we can upper-bound the averaged integration time as
[Equation]
Using the deﬁnition of the Chebyshev root, we have
[Equation]
where the approximation is because ( (cid:98) K 2 (cid:99)− 12 )
[Equation]
K ≈ π 2 when K is large, and hence cos (cid:18) ( (cid:98) K
[Equation]
[Equation]
Similarly, we can approximate
[Equation]
≈
[Equation]
when K is large.
Also, we can approximate r ( K ) ( K +1) / 2 ≈ L + mwhen K is odd and large for the same reason.
Combining the above, the total integration time of HMC with the Chebyshev scheme can be approximated as number of iterations
[Equation]
average integration
[Equation]
When κ :=
Lm is large, the total integration time becomes
[Equation]
Now let us switch to analyzing HMC with the best constant integration time η = Θ (cid:16) 1 √ L (cid:17) (see e.g., (5), Vishnoi (2021)), which has the non-accelerated rate.
Speciﬁcally, it needs K = O (cid:0) κ log (cid:0)
[Equation]
(cid:1)(cid:1) iterations to converge to the target distribution.
Hence, the total integration time of HMC with the best constant integration time is number
[Equation]
By way of comparison ((22) vs. (23)), we see that the total integration time of HMC with the proposed scheme of Chebyshev integration time reduces by a factor √ κ , compared with HMC with the best constant integration time.
The remaining thing to show is the inequality
[Equation]
for
= {
, (cid:98)
(cid:99)} .
We have
[Equation]
Now let us deﬁne H ( k ) :
=
[Equation]

[Equation]
treat k as a continuous variable
[Equation]
The derivative of H ( k ) is
[Equation]
> 0 .
That is, H (cid:48)
( k ) is an increasing function of k when 1 ≤ k ≤ (cid:98) K 2 (cid:99) , which implies that the inequality (24).
Now we have completed the analysis.
E E XPERIMENTS E. B AYESIAN LOGISTIC REGRESSION
[Table]
Method
[Table]
E. S AMPLING FROM A hard DISTRIBUTION
[Table]