Complexity of zigzag sampling algorithm for strongly log-concave distributions

Let $x$ denote the state variable in $\mathbb{R}^{d}$ where $d$ is the dimension. The target distribution we want to sample from is denoted by

where $U(x)$ is the potential and $Z=\int_{\mathbb{R}^{d}}\exp(-U(x))\,\mathrm{d}x$ is the normalizing constant. Although the zigzag process can also be applied to sample non log-concave distributions, we will restrict our analysis to strongly log-concave distributions, namely, we make the following assumption throughout:

For any random variable $X$, we use $\rho(X)$ to denote its law. In this paper, we use chi-square divergence to measure the difference between two probability measures: for probability measures $\rho_{1},\rho_{2}$ that $\rho_{1}\ll\rho_{2}$, it is defined as

The zigzag sampling algorithm is based on a piecewise deterministic Markov process, called zigzag process. Besides the variable $x$, we augment the state space by an auxiliary velocity variable taking value in $\mathbb{R}^{d}$. A trajectory of the zigzag process, denoted by $(X_{t},V_{t})$, can be described as follows. Given some initial $(X_{0},V_{0})$, the position $X_{t}$ always evolves according to $\frac{\,\mathrm{d}}{\,\mathrm{d}t}X_{t}=V_{t}$, while the velocity $V_{t}$ is piecewise constant which only changes when bouncing or refreshing events occur at some random time following Poisson clocks. Bouncing events on the $j$-th direction occur with rate $(V_{t}^{(j)}\partial_{x_{j}}U(X_{t}))_{+}$, and at such an event the velocity $V_{t}$ changes by flipping its $j$-th component to $-V_{t}^{(j)}$. Refreshing events occur with rate $\lambda$ for some fixed $\lambda>0$, when the velocity $V_{t}$ is completely redrawn from the standard normal $\mathcal{N}(0,\mathrm{Id})$.

It has been established [andrieu2018hypocoercivity, bierkens2019ergodicity, lu2020explicit] that under Assumption 1, $\rho(X_{t},V_{t})$ converges to the invariant measure of the zigzag process, which is a product measure of the target measure in $x$ and the standard Gaussian in $v$:

Our analysis relies on the following more quantitative convergence result for zigzag process proved in [lu2020explicit], which also specifies the optimal choice of refreshing rate $\lambda$.¹¹1[lu2020explicit] shows exponential convergence for the backward equation. By duality the exponential convergence of the backward equation in $L^{2}(\bar{\mu})$ is equivalent to the exponential convergence of the forward equation in $\chi^{2}$ with the same rate. We would like to comment here that the choice of $\lambda=\sqrt{L}$ is completely technical since it optimizes the theoretical convergence rate (up to a universal constant) of the zigzag process established in [lu2020explicit]. The zigzag process is ergodic even if $\lambda=0$ and in practice the choice $\lambda=0$ is common.

The left-hand side of (2) controls desired divergence of $\rho(X)$ with respect to the target measure $\mu$, as we have

Moreover, we would take initial condition in the form of

which implies that $\chi^{2}(\bar{\mu}_{0}\,\|\,\bar{\mu})=\chi^{2}(\mu_{0}\,\|\,\mu)$. Therefore, we get

which suggests the total time $T$ needed to achieve control of chi-square divergence.

Of course, in practice, we cannot simulate the zigzag process directly, as simulating the Poisson process associated with the bouncing event would require integrating $(V_{t}^{(j)}\partial_{x_{j}}U(X_{t}))_{+}$ along the trajectory. To turn the zigzag process into an efficient and practical sampling algorithm, the Poisson process for the bouncing events are usually simulated using the Poisson thinning trick (see e.g., discussions in [bierkens2019zig]*Section 3). Under Assumption 1, we will use the following upper bound estimate for the rate:

This upper bound has the advantage of not involving evaluations of $U$ and its partial derivatives, which greatly reduces the computational cost, compared with using numerical quadrature for $d$ Poisson clocks. The price to pay is the increased frequency of potential bouncing events, which scales like $O(\sqrt{d})$ since the pessimistic bound for the partial derivative $|\partial_{x_{i}}U(x)|\leq|\nabla U(x)|\leq L|x|$ typically sacrifices a factor of $O(\sqrt{d})$ in the first inequality.

Following the above discussions, the zigzag sampling algorithm is described in Algorithm 1, where Step LABEL:state:magrate uses the upper bound estimate in (5), while Steps LABEL:state:thinning1–LABEL:state:thinning2 correspond to the Poisson thinning step. Note that for each potential bouncing event, the algorithm requires one evaluation of $\partial_{x_{i}}U$ in Step LABEL:state:thinning1. In practice, typically accessing the partial derivatives of $U$ is the most time consuming step, therefore, in our complexity analysis, we focus on the number of access to partial derivatives.

We also need the following assumption for technical purposes, as will be discussed after stating the main results:

Let us explain the choice of $T$ in (9): For the zigzag sampling algorithm to reach the target $\varepsilon$ accuracy according to (4), the terminal time $T$ needs to be large enough. Meanwhile, the Assumption 2 guarantees that $T$ is not too large, as otherwise we cannot effectively control the number of bouncing events either due to a very large $V$ drawn from a velocity refreshing event or the trajectory reaching regions with large gradient. These motivate our previous Assumption 2 on the initial distribution $\mu_{0}$, as well as the restriction on $\varepsilon$ that it cannot be too small compared to $d$. We remark that the assumption on $\varepsilon$ is not prohibitive as we are interested in high dimensional cases and the error threshold is exponentially small in $d$.

The warm start condition (6) can be achieved if we start with a Gaussian distribution in $x$ and run Langevin Monte Carlo

where $h$ is the step size, and $\xi_{n}$ are i.i.d. $\mathcal{N}(0,\mathrm{Id})$ random variables. This leads to the following corollary:

Theorem 1 guarantees that the zigzag sampling algorithm (Algorithm 1) outputs a sample from a distribution with $\chi^{2}$-divergence at most $\varepsilon$ away from the target density for a computational complexity equivalent to $\widetilde{O}(d^{\frac{3}{2}}\kappa^{2})$ partial derivative evaluations (i.e., amounts to $\widetilde{O}(d^{\frac{1}{2}}\kappa^{2})$ gradient evaluations), in the regime $\max\{\kappa,\log\frac{1}{\varepsilon}\}\ll\frac{d}{\log d}$ with a warm-start condition. Corollary 1.4 establishes that the hybrid LMC-zigzag algorithm outputs a sample for a computational complexity $\widetilde{O}(d^{\frac{4}{5}}\kappa^{\frac{16}{5}})$ gradient evaluations. The initialization using LMC is added only for technical reasons as we currently do not have complexity guarantees otherwise with an explicit initial distribution, nor is it necessary for actual implementations. We would also like to comment that our goal is to obtain the best possible scaling in $d$, and the scaling in $\kappa$ might be possibly improved by a more careful analysis.

Our analysis is based on the quantitative convergence rate of the zigzag process established in [lu2020explicit], which is $O(\frac{m}{\sqrt{L}})$ for $m$-convex and $L$-smooth potentials. The rest of our proof is based on estimating $\sup|X_{t}|$ along a single trajectory of the zigzag process and subsequently turn this into an estimate on the number of potential bouncing events, and hence number of partial derivative evaluations. Our analysis utilizes the two important and desirable features of the zigzag sampling process:

• •

  The implementation of the zigzag process does not need time discretization, as the velocity in deterministic portion of the trajectory remains constant, which makes it possible to simulate the exact trajectories of the zigzag process while eliminating an important source of error. This is the reason that the complexity of the zigzag process only has logarithmic dependence on $\frac{1}{\varepsilon}$, without Metropolis acceptance/rejection.

• •

  Moreover, for each potential bouncing event of zigzag, only one evaluation of a partial derivative of the potential is required, which is $O(d)$ cheaper than a full gradient evaluation in computational cost for usual model of computation.

We would also remark that we quantify the error of distribution in terms of $\chi^{2}$-divergence, which provides stronger guarantee than total variation, KL divergence or $2$-Wasserstein distance. While $\chi^{2}$-divergence is relatively convenient for obtaining convergence rates of continuous processes based on Poincarè inequality [cao2019explicit, lu2020explicit], it does not seem easy to use for analyzing discretization error of SDEs. The work [vempala2019rapid] made assumptions of Poincarè inequality for the discrete invariant measure, which is difficult to verify. We are fortunate to avoid such problem for zigzag sampler, thanks to the fact that zigzag does not need time discretization. After the first version of this work appears online, [erdogdu2020convergence] established convergence of LMC in $\chi^{2}$- and Rényi divergence, using the exponential convergence of continuous time overdamped Langevin dynamics in Rényi divergence [cao2019exponential, vempala2019rapid].

Here we focus on results on non-asymptotic analysis of sampling algorithms, which has been a focused research area in recent years. Many sampling algorithms have been analyzed including algorithms based on overdamped Langevin dynamics [dalalyan2017theoretical, durmus2019high, durmus2019analysis, vempala2019rapid, li2019stochastic, ding2020randomlangevin], underdamped Langevin dynamics [cheng2018underdamped, dalalyan2020sampling, ma2019there, shen2019randomized, ding2020randomunderdamped, monmarche2020high], Hamiltonian Monte Carlo [mangoubi2017rapid, lee2018algorithmic, chen2019optimal, mangoubi2018dimensionally, bou2020coupling], or high order Langevin dynamics [mou2019high], among others. These methods involve discretization of ODEs or SDEs, which yields an error that scales polynomially with step size. Thus the complexity of these algorithms has polynomial dependence on $\varepsilon^{-1}$, where $\varepsilon$ is the desired accuracy threshold.

Metropolized variants of sampling algorithms, including Metropolized HMC and Metropolis Adjusted Langevin Algorithm (MALA), have also been studied in [dwivedi2018log, chen2020fast, lee2020logsmooth], the complexities of which have only logarithmic dependence on $\varepsilon^{-1}$, similar to the zigzag sampling process analyzed here. In [dwivedi2018log] the complexity upper bound for MALA is established as $\widetilde{O}(\kappa d+\kappa^{\frac{3}{2}}d^{\frac{1}{2}})$ under warm start condition, and $\widetilde{O}(\kappa d^{2}+\kappa^{\frac{3}{2}}d^{\frac{3}{2}})$ with a feasible start. In [chen2020fast] the complexity upper bound for MALA is improved to $\widetilde{O}(\kappa d+\kappa^{\frac{3}{2}}d^{\frac{1}{2}})$ with feasible start (where $\mu_{0}=\mathcal{N}(0,\frac{1}{L}\mathrm{Id})$). The work [chen2020fast] also established bounds for Metropolized HMC, which is $\widetilde{O}(\kappa d^{\frac{11}{12}})$ with warm start (which is in fact more stringent than our Assumption 2) in the regime $\kappa=O(d^{\frac{2}{3}})$, and $\widetilde{O}(\kappa^{\frac{3}{4}}d+\kappa^{\frac{7}{4}}d^{\frac{1}{2}})$ with feasible start if the target potential function has a bounded Hessian. The complexity upper bound has been improved in [lee2020logsmooth] to $\widetilde{O}(\kappa d)$ for both Metropolized HMC and MALA with a feasible start, based on a refined analysis using concentration of gradient norm. In comparison, our result for zigzag relies on a warm start (which is achievable by LMC), while the complexity upper bound has better dependence in $d$. The issue of feasible start will be further discussed in Section 3.

Regarding asymptotic analysis for the convergence of zigzag process, the ergodicity was first established in [bierkens2019ergodicity]. Exponential convergence of the zigzag process is established in [fontbona2016long, bierkens2017piecewise] using a Lyapunov function argument. A central limit theorem of the zigzag process is established in [bierkens2017limit], and a large deviation principle is established for the empirical measure in [bierkens2019large]. The spectrum of the zigzag process has been studied in [bierkens2019spectral, guillin2020low]. A dimension independent exponential convergence rate for the zigzag process is established in [andrieu2018hypocoercivity], using the hypocoercivity framework developed in [dolbeault2015hypocoercivity]. Finally, a more quantitative convergence estimate was established in [lu2020explicit], for which our analysis of the sampling algorithm is based on.