The Zig-Zag Process and Super-Efficient Sampling for Bayesian Analysis of Big Data

Introduction

The importance of Markov chain Monte Carlo techniques in Bayesian inference shows no signs of diminishing. However, all commonly used methods are variants on the Metropolis-Hastings (MH) algorithm Metropolis et al. 1953; Hastings 1970 and rely on innovations which date back over 60 years. All MH algorithms simulate realisations from a discrete reversible ergodic Markov chain with invariant distribution $\pi$ which is (or is closely related to) the target distribution, i.e. the posterior distribution in a Bayesian context. The MH algorithm gives a beautifully simple though flexible recipe for constructing such Markov chains, requiring only local information about $\pi$ (typically pointwise evaluations of $\pi$ and, perhaps, its derivative at the current and proposed new locations) to complete each iteration.

However new complex modelling and data paradigms are seriously challenging these established methodologies. Firstly, the restriction of traditional MCMC to reversible Markov chains is a serious limitation. It is now well-understood both theoretically Hwang, Hwang-Ma and Sheu 1993; Chen and Hwang 2013; Rey-Bellet and Spiliopoulos 2015; Bierkens 2015; Duncan, Lelièvre and Pavliotis 2016 and heuristically Neal 1998 that non-reversible chains offer potentially massive advantages over reversible counterparts. The need to escape reversibility, and create momentum to aid mixing throughout the state space is certainly well-known, and motivates a number of modern MCMC methods, including the popular Hamiltonian MCMC Duane et al. 1987.

A second major obstacle to the application of MCMC for Bayesian inference is the need to process potentially massive data-sets. Since MH algorithms in their pure form require a likelihood evaluation – and thus processing the full data-set – at every iteration, it can be impractical to carry out large numbers of MH iterations. This has led to a range of alternatives that use sub-samples of the data at each iteration Welling and Teh 2011; Maclaurin and Adams 2014; Ma, Chen and Fox 2015; Quiroz, Villani and Kohn 2015, or that partition the data into shards, run MCMC on each shard, and then attempt to combine the information from these different MCMC runs Neiswanger, Wang and Xing 2013; Scott et al. 2016; Wang and Dunson 2013; Li, Srivastava and Dunson 2017. However most of these methods introduce some form of approximation error, so that the final sample will be drawn from some approximation to the posterior, and the quality of the approximation can be impossible to evaluate. As an exception the Firefly algorithm Maclaurin and Adams 2014 samples from the exact distribution of interest (but see the comment below).

This paper introduces the multi-dimensional Zig-Zag sampling algorithm (ZZ) and its variants. These methods overcome the restrictions of the lifted Markov chain approach of Turitsyn, Chertkov and Vucelja 2011 as they do not depend on the introduction of momentum generating quantities. They are also amenable to the use of sub-sampling ideas. The dynamics of the Zig-Zag process depends on the target distribution through the gradient of the logarithm of the target. For Bayesian applications this is a sum, and is easy to estimate unbiasedly using sub-sampling. Moreover, Zig-Zag with Sub-Sampling (ZZ-SS) retains the exactness of the required invariant distribution. Furthermore, if we also use control variate ideas to reduce the variance of our sub-sampling estimator of the gradient, the resulting Zig-Zag with Control Variates (ZZ-CV) algorithm has remarkable super-efficient scaling properties for large data sets.

We will call an algorithm super-efficient if it is able to generate independent samples from the target distribution at a higher efficiency than if we would draw independently from the target distribution at the cost of evaluating all data. The only situation we are aware of where we can implement super-efficient sampling is with simple conjugate models, where the likelihood function has a low-dimensional summary statistic which can be evaluated at cost $O(n)$, where $n$ is the number of observations, after which we can obtain independent samples from the posterior distribution at a cost of $O(1)$ by using the functional form of the posterior distribution. The ZZ-CV can replicate this computational efficiency: after a pre-computation of $O(n)$, we are able to obtain independent samples at a cost of $O(1)$. In this sense it contrasts with the Firefly algorithm Maclaurin and Adams 2014 which has an ESS per datum which decreases approximately as $1/n$ where $n$ is the size of the data, so that the gains of this algorithm do not increase with $n$; see (Bouchard-Côté, Vollmer and Doucet 2015, Section 4.6).

The use of PDMPs such as the Zig-Zag processes is an exciting and mostly unexplored area in MCMC. The first occurrence of a PDMP for sampling purposes is in the computational physics literature Peters and De With 2012, which in one dimension coincides with the Zig-Zag process. In Bouchard-Côté, Vollmer and Doucet 2015 this method is given the name Bouncy Particle Sampler. In multiple dimensions the Zig-Zag process and Bouncy Particle Sampler (BPS) are different processes: both are PDMPs which move along straight line segments, but the Zig-Zag process changes direction in only a single component at each switch, whereas the Bouncy Particle Sampler reflects the full direction vector in the level curves of the density function. As we will see in Section 2.4, this difference has a beneficial effect on the ergodic properties of the Zig-Zag process. The one-dimensional Zig-Zag process is analysed in detail in e.g. Fontbona, Guérin and Malrieu 2012; Monmarché 2014; Fontbona, Guérin and Malrieu 2016; Bierkens and Roberts 2017.

Since the first version of this paper was conceived already several other related theoretical and methodological papers have appeared. In particular we mention here results on exponential ergodicity of the BPS Deligiannidis, Bouchard-Côté and Doucet 2017 and ergodicity of the multi-dimensional Zig-Zag process Bierkens, Roberts and Zitt 2017. The Zig-Zag process has the advantage that it is ergodic under very mild conditions, which in particular means that we are not required to choose a refreshment rate. At the same time, the BPS seems more ‘natural’, in that it tries to minimise the bounce rate and the change in direction at bounces, and it may be more efficient for this reason. However it is a challenge to make a direct comparison in efficiency of the two methods since the efficiency depends both on the computational effort per unit of continuous time of the respective algorithms, as well as the mixing time of the underlying processes. Therefore we expect analysing the relative efficiency of PDMP based algorithms to be an important area of continued research for years to come.

A continuous-time sequential Monte Carlo algorithm for scalable Bayesian inference with big data, the SCALE algorithm, is given in Pollock et al. 2016. Advantages that Zig-Zag has over SCALE is that it avoids the issue of controlling the stability of importance weights, and it is simpler to implement. Whereas the SCALE algorithm is well-adapted for the use of parallel architecture computing, and has particularly simple scaling properties for big data.

The Zig-Zag process

For a given $(\xi,\theta)\in E$, we may construct a trajectory of $(\Xi,\Theta)$ of the Zig-Zag process with initial condition $(\xi,\theta)$ as follows.

Let $(T^{0},\Xi^{0},\Theta^{0}):=(0,\xi,\theta)$.

Let $\xi^{k}(t):=\Xi^{k-1}+\Theta^{k-1}t$, $t\geq 0$

For $i=1,\dots,d$, let $\tau^{k}_{i}$ be distributed according to

Let $i_{0}:=\operatorname{argmin}_{i\in\{1,\dots,d\}}\tau^{k}_{i}$ and let $T^{k}:=T^{k-1}+\tau^{k}_{i_{0}}$.

The piecewise deterministic trajectories $(\Xi(t),\Theta(t))$ are now obtained as

Since the switching rates are continuous and hence bounded on compact sets, and $\Xi$ will travel a finite distance within any finite time interval, within any bounded time interval there will be finitely many switches almost surely.

The above procedure provides a mathematical construction of a Markov process as well as the outline of an algorithm which simulates this process. The only step in this procedure which presents a computational challenge is the simulation of the random times $(T_{i}^{k})$ and a significant part of this paper will consider obtaining these in a numerically efficient way.

Figure 1 displays trajectories of the Zig-Zag process for several examples of invariant distributions. The name of the process is derived by the zig-zag nature of paths that the process produces. Figure 1 shows an important difference in the output of the Zig-Zag process, as compared to a discrete-time MCMC algorithm: the output of is a continuous-time sample path. The bottom row of plots in Figure 1 also gives a comparison to a reversible MCMC algorithm, Metropolis Adjusted Langevin (Roberts and Tweedie 1996, MALA), and demonstrates an advantage of a non-reversible sampler: it can cope better with a heavy tailed target. This is most easily seen if we start the process out in the tail, as in the figure. The velocity component of the Zig-Zag process enables it to quickly return to the mode of the distribution, whereas the reversible algorithm behaves like a random walk in the tails, and takes much longer to return to the mode.

2 Invariant distribution

Suppose Assumption 2.1 holds. Let $\mu$ denote the probability distribution on $E$ such that $\mu$ has Radon-Nikodym derivative

where $Z=\int_{E}\exp(-\Psi)\ d\mu_{0}$. Then the Zig-Zag process $(\Xi,\Theta)$ with switching rates $(\lambda_{i})_{i=1}^{d}$ has invariant distribution $\mu$.

The proof is located in the Section 1 of the Supplementary Material. We see that under the invariant distribution of the Zig-Zag process, $\xi$ and $\theta$ are independent, with $\xi$ having density proportional to $\exp(-\Psi(\xi))$ and $\theta$ having a uniform distribution on the points in $\{-1,+1\}^{d}$.

The proof is located in Section 1 of the Supplementary Material.

3 Zig-Zag process for Bayesian inference

One application of the Zig-Zag process is as an alternative to MCMC for sampling from posterior distributions in Bayesian statistics. We show here that it is straightforward to derive a class of Zig-Zag processes that have a given posterior distribution as their invariant distribution. The dynamics of the Zig-Zag process only depend on knowing the posterior density up to a constant of proportionality.

We can write $\pi(\xi)$ in the form of the previous section,

will have the posterior density $\pi(\xi)$ as the marginal of its invariant distribution. We call the process with these rates the Canonical Zig-Zag process for the negative log density $\Psi$. As explained in Proposition 2.3, we can construct a family of Zig-Zag processes with the same invariant distribution by choosing any set of functions $\gamma_{i}(\xi,\theta)$, for $i=1,\ldots,d$, which take non-negative values and for which $\gamma_{i}(\xi,\theta)=\gamma_{i}(\xi,F_{i}[\theta])$, and setting

The intuition here is that $\lambda_{i}(\xi,\theta)$ is the rate at which we transition from $\theta$ to $F_{i}[\theta]$. The condition $\gamma_{i}(\xi,\theta)=\gamma_{i}(\xi,F_{i}[\theta])$ means that we increase by the same amount both the rate at which we will transition from $\theta$ to $F_{i}[\theta]$ and vice versa. As our invariant distribution places the same probability of being in a state with velocity $\theta$ as that of being in state $F_{i}[\theta]$, these two changes in rate cancel out in terms of their effect on the invariant distribution. Changing the rates in this way does impact the dynamics of the process, with larger $\gamma_{i}$ values corresponding to more frequent changes in the velocity of the Zig-Zag process, and we would expect the resulting process to mix more slowly.

for any initial condition $(\xi,\theta)\in E$. Sufficient conditions for ergodicity will be discussed in the following section. Taking $\gamma$ to be positive and bounded everywhere ensures ergodicity, as will be established in Theorem 2.10.

4 Ergodicity of the Zig-Zag process

Ergodicity is directly related to the requirement that $(\Xi(t),\Theta(t))$ is irreducible, i.e. the state space is not reducible into components which are each invariant for the process $(\Xi(t),\Theta(t))$. For the one-dimensional Zig-Zag process, (exponential) ergodicity has already been established under mild conditions Bierkens and Roberts 2017. As we discuss below, irreducibility, and thus ergodicity, can be established for large classes of multi-dimensional target distributions, such as i.i.d. Gaussian distributions, and also if the switching rates $\lambda_{i}(\xi,\theta)$ are positive for all $i=1,\dots,d$, and $(\xi,\theta)\in E$.

Suppose $d=1$ and there exists $\xi_{0}>0$ such that

$\inf_{\xi\geq\xi_{0}}\lambda(\xi,+1)>\sup_{\xi\geq\xi_{0}}\lambda(\xi,-1)$, and

$\inf_{\xi\leq-\xi_{0}}\lambda(\xi,-1)>\sup_{\xi\leq-\xi_{0}}\lambda(\xi,+1)$.

(Bierkens and Roberts 2017, Theorem 5) Suppose Assumption 2.4 holds. Then there exists a function $f:E\rightarrow[1,\infty)$ which is norm-like such that the Zig-Zag process is $f$-exponentially ergodic, i.e. there exists a constant $\kappa>0$ and $0<\rho<1$ such that

As an example of fundamental importance, which will also be used in the proof of Theorem 2.10, consider a one-dimensional Gaussian distribution. For simplicity let $\pi(\xi)$ be centred, $\pi(\xi)=\frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left(-\frac{\xi^{2}}{2\sigma^{2}}\right)$ for some $\sigma>0$. According to (4) the switching rates take the form

As long as $\gamma$ is bounded from above, Assumption 2.4 is satisfied. In particular this holds if $\gamma$ is equal to a non-negative constant.

As long as $\gamma_{i}(\xi)=\gamma_{i}(\xi_{i})$, i.e. if $\gamma_{i}(\xi)$ only depends on the $i$-th coordinate of $\xi$, the switching rate of coordinate $i$ is independent of the other coordinates $\xi_{j}$, $j\neq i$. It follows that the switches of the $i$-th coordinate can be generated by a one-dimensional time inhomogeneous Poisson process, which is independent of the switches in the other coordinates. As a consequence the $d$-dimensional Zig-Zag process $(\Xi(t),\Theta(t))=(\Xi^{1}(t),\dots,\Xi_{d}(t),\Theta^{1}(t),\dots,\Theta^{d}(t))$ consists of a combination of $d$ independent Zig-Zag processes $(\Xi^{i}(t),\Theta^{i}(t))$, $i=1,\dots,d$.

Suppose $P(x,dy)$ is the transition kernel of a Markov chain on a state space $E$. We say that the Markov chain associated to $P$ is mixing if there exists a probability distribution $\pi$ on $E$ such that

For any continuous time Markov process with family of transition kernels $P^{t}(x,dy)$ we can consider the associated time-discretized process, which is a Markov chain with transition kernel $Q(x,dy):=P^{\delta}(x,dy)$ for a fixed $\delta>0$. The value of $\delta$ will be of no significance in our use of this construction.

This follows from the decomposition of the $d$-dimensional Zig-Zag process as $d$ one-dimensional Zig-Zag processes and Lemma 1.1 in the Supplementary material. ∎

Continuing Example 2.6, consider the simple case in which $\pi$ is of product form with each $\pi_{i}$ a centered Gaussian density function with variance $\sigma_{i}^{2}$. It follows from Proposition 2.8 and Example 2.6 that the multi-dimensional canonical Zig-Zag process (i.e. the Zig-Zag process with $\gamma_{i}\equiv 0$) is mixing. This is different from the Bouncy Particle Sampler Bouchard-Côté, Vollmer and Doucet 2015, which is not ergodic for an i.i.d. Gaussian without ‘refreshments’ of the momentum variable.

We now show that strict positivity of the rates ensures ergodicity.

Suppose $\lambda:E\rightarrow(0,\infty)^{d}$, in particular $\lambda_{i}(\xi,\theta)$ is positive for all $i=1,\dots,d$ and $(\xi,\theta)\in E$. Then there exists at most a single invariant measure for the Zig-Zag process with switching rate $\lambda$.

The proof of this result consists of a Girsanov change of measure with respect to a Zig-Zag process targeting an i.i.d. standard normal distribution, which we know to be irreducible. The irreducibility then carries over to the Zig-Zag process with the stated switching rates. A detailed proof can be found in the Supplementary material.

Based on numerous experiments, we conjecture that the canonical multi-dimensional Zig-Zag process is ergodic in general under only mild conditions. A detailed investigation of ergodicity will be the subject of a forthcoming paper Bierkens, Roberts and Zitt 2017.

Implementation

As mentioned earlier, the main computational challenge is an efficient simulation of the random times $T^{k}_{i}$ introduced in Section 2.1. We will focus on simulation by means of Poisson thinning.

Now for a given initial point $(\xi,\theta)\in E$, let $m_{i}(t):=\lambda_{i}(\xi+\theta t,\theta)$, for $i=1,\dots,d$, and suppose we have available continuous functions $M_{i}(t)$ such that $m_{i}(t)\leq M_{i}(t)$ for $i=1,\dots,d$ and $t\geq 0$. We call these $(M_{i})_{i=1}^{d}$ computational bounds for $(m_{i})_{i=1}^{d}$. We can use Proposition 3.1 to obtain the first switching times $(\widetilde{\tau}^{1}_{i})_{i=1}^{d}$ from a (theoretically infinite) collection of proposed switching times $(\tau^{1}_{i},\tau^{2}_{i},\dots)_{i=1}^{d}$ given the initial point $(\xi,\theta)$, and use the obtained skeleton point at time $\widetilde{\tau}^{1}:=\min_{i\in\{1,\dots,d\}}\widetilde{\tau}_{i}^{1}$ as a new initial point (which is allowed by the strong Markov property) with the component $i_{0}=\operatorname{argmin}_{i\in\{1,\dots,d\}}\widetilde{\tau}_{i}^{1}$ of $\theta$ switched.

The strong Markov property of the Zig-Zag process simplifies the computational procedure further: we can draw for each component $i=1,\dots,d$ the first proposed switching time $\tau_{i}:=\tau_{i}^{1}$, determine $i_{0}:=\operatorname{argmin}_{i\in\{1,\dots,d\}}\tau_{i}$ and decide whether the appropriate component of $\theta$ is switched at this time with probability $m_{i_{0}}(\tau)/M_{i_{0}}(\tau)$, where $\tau:=\tau_{i_{0}}$. Then since $\tau$ is a stopping time for the Markov process, we can use the obtained point of the Zig-Zag process at time $\tau$ as new starting point, regardless of whether we switch a component of $\theta$ at the obtained skeleton point. A full computational procedure for simulating the Zig-Zag process is given by Algorithm 1.

We now come to the important issue of obtaining computational bounds for the Zig-Zag Process, i.e. useful upper bounds for the switching rates $(m_{i})$. If we can compute the inverse function $G_{i}(y):=\inf\{t\geq 0:H_{i}(t)\geq y\}$ of $H_{i}:t\mapsto\int_{0}^{t}M_{i}(s)\ ds$, we can simulate $\tau_{1},\dots,\tau_{d}$ using the CDF inversion technique, i.e. by drawing i.i.d. uniform random variables $U_{1},\dots,U_{d}$ and setting $\tau_{i}:=G_{i}(-\log U_{i})$, $i=1,\dots d$.

The computational bounds are directly related to the algorithmic efficiency of Zig-Zag Sampling. From Algorithm 1, it is clear that for every simulated time $\tau$ a single component of $\lambda$ needs to be evaluated, which corresponds by (4) to the evaluation of a single component of the gradient of the negative log density $\Psi$. The magnitude of the computational bounds, $(M_{i})$, will determine how far the Zig-Zag process will have moved in the state space before a new evaluation of a component of $\lambda$ is required, and we will pay close attention to the scaling of $M_{i}$ with respect to the number of available observations in a Bayesian inference setting.

2 Example: globally bounded log density gradient

Algorithm 1 may be used with $M_{i}\equiv c_{i}$ for $i=1,\dots,d$ at every iteration.

This situation arises with heavy-tailed distributions. E.g. if $\pi$ is Cauchy, then $\Psi(\xi)=\log(1+\xi^{2})$, and consequently $\lambda(\xi,\theta)=\left(\frac{2\theta\xi}{1+\xi^{2}}\right)^{+}\leq 1$.

3 Example: negative log density with dominated Hessian

Applying this inequality we obtain for $i=1,\dots,d$,

Hence the general Zig-Zag Algorithm may be applied taking

with $a_{i}$ and $b_{i}$ as specified above. A complete procedure for Zig-Zag Sampling for a log density with dominated Hessian is provided in Algorithm 2.

It is also possibly to apply inequality (6) in such a way as to obtain the estimate

This requires us to compute $Q\theta$ whenever $\theta$ changes (a computation of $O(d)$).

Big data Bayesian inference by means of error-free sub-sampling

Throughout this section we assume the derivatives of $\Psi$ admit the representation

where $\Psi^{j}(\xi)=-\log\pi_{0}(\xi)-n\log f(x^{j}|\xi)$, and we could choose $E_{i}^{j}(\xi)=\partial_{i}\Psi^{j}(\xi)$. It is crucial that every $E_{i}^{j}$ is a factor $O(n)$ cheaper to evaluate than the full derivative $\partial_{i}\Psi(\xi)$.

We will describe two successive improvements over the basic Zig-Zag Sampling (ZZ) algorithm specifically tailored to the situation in which (7) is satisfied. The first improvement consists of a sub-sampling approach where we need calculate only one of the $E_{i}^{j}s$ at each simulated time, rather than sum of all $n$ of them. This sub-sampling approach (referred to as Zig-Zag with Sub-Sampling, ZZ-SS) comes at the cost of an increased computational bound. Our second improvement is to use control variates to reduce this bound, resulting in the Zig-Zag with Control Variates (ZZ-CV) algorithm.

Let $(\xi(t))_{t\geq 0}$ denote a linear trajectory originating in $(\xi,\theta)\in E$, i.e. $\xi(t)=\xi+\theta t$. Define a collection of switching rates along the trajectory $(\xi(t))$ by

Algorithm 3 generates a skeleton of a Zig-Zag process with switching rates given by

and invariant distribution $\mu$ given by (3).

Conditional on $\tau$, the probability that component $i_{0}$ of $\theta$ is switched at time $\tau$ is seen to be

By Proposition 3.1 we thus have an effective switching rate $\lambda_{i}$ for switching the $i$-th component of $\theta$ given by (10). Finally we verify that the switching rates $(\lambda_{i})$ given by (10) satisfy (2). Indeed,

By Theorem 2.2, the Zig-Zag process has the stated invariant distribution. ∎

The important advantage of using Zig-Zag in combination with sub-sampling is that at every iteration of the algorithm we only have to evaluate a single component of $E^{j}_{i}$, which reduces algorithmic complexity by a factor $O(n)$. However this may come at a cost. Firstly, the computational bounds $(M_{i})$ may have to be increased which in turn will increase the algorithmic complexity of simulating the Zig-Zag sampler. Also, the dynamics of the Zig-Zag process will change, because the actual switching rates of the process are increased. This increases the diffusivity of the continuous time Markov process, and affects the mixing properties in a negative way.

2 Zig-Zag with Sub-Sampling (ZZ-SS) for globally bounded log density gradient

A straightforward application of sub-sampling is possible if we have (8) with $\nabla\Psi^{j}$ globally bounded, i.e. there exist positive constants $(c_{i})$ such that

so that (9) is satisfied. The corresponding version of Algorithm 3 will be called Zig-Zag with Sub-Sampling (ZZ-SS).

3 Zig-Zag with Control Variates (ZZ-CV)

The reason for defining $E_{i}^{j}(\xi)$ in this manner is to try and reduce its variability as we vary $j$. By the Lipschitz condition we have $E_{i}^{j}(\xi)\leq|\partial_{i}\Psi(\xi^{\star})|+C_{i}\|\xi-\xi^{\star}\|_{p}$, and thus the variability of the $E_{i}^{j}(\xi)$s will be small if 1) the reference point $\xi^{\star}$ is close to the mode of the posterior and 2) $\xi$ is close to $\xi^{\star}$. Under standard asymptotics we expect a draw from the posterior for $\xi$ to be $O_{p}(n^{-1/2})$ from the posterior mode. Thus if we have a procedure for finding a reference point $\xi^{\star}$ which is within $O(n^{-1/2})$ of the posterior mode then this would ensure $\|\xi-\xi^{\star}\|_{2}$ is $O_{p}(n^{-1/2})$ if $\xi$ is drawn from the posterior. For such a choice of $\xi^{\star}$ we would have $\partial_{i}\Psi(\xi^{\star})$ of $O_{p}(n^{1/2})$.

Using the Lipschitz condition, we can now obtain computational bounds of $(m_{i})$ for a trajectory $\xi(t):=\xi+\theta t$ originating in $(\xi,\theta)$. Define

where $a_{i}:=\left(\theta_{i}\partial_{i}\Psi(\xi^{\star})\right)^{+}+C_{i}\|\xi-\xi^{\star}\|_{p}$ and $b_{i}:=C_{i}d^{1/p}$. Then (9) is satisfied. Indeed, using Lipschitz continuity of $y\mapsto(y)^{+}$,

Implementing this scheme requires some pre-processing of the data. First we need a way of choosing a suitable reference point $\xi^{\star}$ to find a value close to the mode using an approximate or exact numerical optimization routine. The complexity of this operation will be $O(n)$. Once we have found such a reference point we have an one-off $O(n)$ cost of calculating $\partial_{i}\Psi(\xi^{\star})$ for each $i=1,\ldots,d$. However, once we have paid this upfront computational cost, the resulting Zig-Zag sampler can be super-efficient. This is discussed in more detail in Section 5, and demonstrated empirically in Section 6. The version of Algorithm 3 resulting from this choice of $E^{j}_{i}$ and $M_{i}$ will be called Zig-Zag with Control Variates (ZZ-CV).

When choosing $p\geq 1$, there will be a trade-off between the magnitude of $C_{i}$ and of $\|\xi-\xi^{\star}\|_{p}$, which may influence the scaling of Zig-Zag sampling with dimension. We will see in Section 6.3 that for i.i.d. Gaussian components, the choice $p=\infty$ is optimal. When the situation is less clear, choosing the Euclidean norm ($p=2$) is a reasonable choice.

Scaling analysis

where $x^{j}$ are i.i.d. drawn from $f(x^{j}\mid\xi_{0})$. Let $\widehat{\xi}$ denote the maximum likelihood estimator (MLE) for $\xi$ based on data $x^{1},\ldots,x^{n}$. Introduce the coordinate transformation

As $n\rightarrow\infty$ the posterior distribution in terms of $\phi$ will converge to a multivariate Gaussian distribution with mean 0 and covariance matrix given by the inverse of the expected information $i(\theta_{0})$; see e.g. Johnson 1970.

First let us obtain a Taylor expansion of the switching rate for $\xi$ close to $\widehat{\xi}$. We have

The first term vanishes by the definition of the MLE. Expressed in terms of $\phi$, the switching rates are

With respect to the coordinate $\phi$, the canonical Zig-Zag process has constant speed $\sqrt{n}$ in each coordinate, and by the above computation, a switching rate of $O(\sqrt{n})$. After a rescaling of the time parameter by a factor $\sqrt{n}$, the process in the $\phi$-coordinate becomes a Zig-Zag process with unit speed in every direction and switching rates

If we let $n\rightarrow\infty$, the switching rates converge almost surely to those of a Zig-Zag process with switching rates

where $i(\theta_{0})$ denotes the expected information. These switching rates correspond to the limiting Gaussian distribution with covariance matrix $(i(\theta_{0}))^{-1}$.

In this limiting Zig-Zag process, all dependence on $n$ has vanished. Starting from equilibrium, we require a time interval of $O(1)$ (in the rescaled time) to obtain an essentially independent sample. In the original time scale this corresponds to a time interval of $O(n^{-1/2})$. As long as the computational bound in the Zig-Zag algorithm is $O(n^{1/2})$, this can be achieved using $O(1)$ proposed switches. The computational cost for every proposed switch is $O(n)$, because the full data $(x^{i})_{i=1}^{n}$ needs to be processed in the computation of the true switching rate at the proposed switching time.

We conclude that the computational complexity of the Zig-Zag (ZZ) algorithm per independent sample is $O(n)$, provided that the computational bound is $O(n^{1/2})$. This is the best we can expect for any standard Monte Carlo algorithm (where we will have a $O(1)$ number of iterations, but each iteration is $O(n)$ in computational cost).

To compare, if the computational bound is $O(n^{\alpha})$ for some $\alpha>1/2$, then we require $O(n^{\alpha-1/2})$ proposed switches before we have simulated a total time interval of length $O(n^{-1/2})$, so that, with a complexity of $O(n)$ per proposed switching time, the Zig-Zag algorithm has total computational complexity $O(n^{\alpha+1/2})$. So, for example, with global bounds we have that the computational bound is $O(n)$ (as each term in the log density is $O(1)$), and hence ZZ will have total computational complexity of $O(n^{3/2})$.

Consider Algorithm 2 in the one-dimensional case, with the second derivative of $\Psi$ bounded from above by $Q>0$. We have $Q=O(n)$ as $\Psi^{\prime\prime}$ is the sum of $n$ terms of $O(1)$. The value of $b$ is kept fixed at the value $b=Q=O(n)$. Next $a$ is given initially as

and increased by $b\tau$ until a switch happens and $a$ is reset to $\theta\Psi^{\prime}(\xi)$. Because of the initial value for $a$, switches will occur at rate $O(n^{1/2})$ so that $\tau$ will be $O(n^{-1/2})$, and the value of $a$ will remain $O(n^{1/2})$. Hence the magnitude of the computational bound $M(t)=(a+bt)^{+}$ is $O(n^{1/2})$.

2 Scaling of Zig-Zag with Control Variates (ZZ-CV)

Now we will study the limiting behaviour as $n\rightarrow\infty$ of ZZ-CV introduced in Section 4.3. In determining the computational bounds we take $p=2$ for simplicity, e.g. in (12). Also for simplicity assume that $\xi\mapsto\partial_{\xi_{i}}\log f(x^{j}\mid\xi)$ has Lipschitz constant $k_{i}$ (independent of $j=1,\dots,n$) and write $C_{i}=nk_{i}$, so that (12) is satisfied. In practice there may be a logarithmic increase with $n$ in the Lipschitz constants $k_{i}$ as we have to take a global bound in $n$. For the present discussion we ignore such logarithmic factors. We assume reference points $\xi^{\star}$ for growing $n$ are determined in such a way that $\|\xi^{\star}-\widehat{\xi}\|_{2}$ is $O(n^{-1/2})$. For definiteness, suppose there exists a $d$-dimensional random variable $Z$ such that $n^{1/2}(\xi^{\star}-\widehat{\xi})\rightarrow Z$ in distribution, with the randomness in $Z$ independent of $(x^{j})_{j=1}^{\infty}$.

We can look at ZZ-CV with respect to the scaled coordinate $\phi$ as $n\rightarrow\infty$. Denote the reference point for the rescaled parameter as $\phi^{\star}:=\sqrt{n}(\xi^{\star}-\widehat{\xi})$.

The essential quantities to consider are the switching rate estimators $E_{i}^{j}$. We estimate

We find that $|E_{i}^{j}(\xi)|=O(n^{1/2})$ under the stationary distribution.

By slowing down the Zig-Zag process in $\phi$ space by $\sqrt{n}$, the continuous time process generated by ZZ-CV will approach a limiting Zig-Zag process with a certain switching rate of $O(1)$. In general this switching rate will depend on the way that $\xi^{\star}$ is obtained. To simplify the exposition, in the following computation we assume $\xi^{\star}=\widehat{\xi}$. Rescaling by $n^{-1/2}$, and developing a Taylor approximation around $\widehat{\xi}$,

By Theorem 4.1, the rescaled effective switching rate for ZZ-CV is given by

Just as with ZZ, the rescaled Zig-Zag process underlying ZZ-CV converges to a limiting Zig-Zag process with switching rate $\widetilde{\lambda}_{i}(\phi,\theta)$. Since the computational bounds of ZZ-CV are $O(n^{1/2})$, a completely analogous reasoning to the one for ZZ algorithm above (Section 5.1) leads to the conclusion that $O(1)$ proposed switches are required to obtain an independent sample. However, in contrast with the ZZ-algorithm, the ZZ-CV algorithm is designed in such a way that the computational cost per proposed switch is $O(1)$.

We conclude that the computational complexity of the ZZ-CV algorithm is $O(1)$ per independent sample. This provides a factor $n$ increase in efficiency over standard MCMC algorithms, resulting in an asymptotically unbiased algorithm for which the computational cost of obtaining an independent sample does not depend on the size of the data.

3 Remarks

The arguments above assume we are at stationarity – and how quickly the two algorithms converge is not immediately clear. Note however that for sub-sampling Zig-Zag it is possible to choose the reference point $\xi^{\star}$ as starting point, thus avoiding much of the issues about convergence.

In some sense, the good computational scaling of ZZ-CV is leveraging the asymptotic normality of the posterior, but in such a way that ZZ-CV always samples from the true posterior. Thus when the posterior is close to Gaussian it will be quick; when it is far from Gaussian it may well be slower but will still be “correct”. This is fundamentally different from other algorithms (Neiswanger, Wang and Xing 2013; Scott et al. 2016; Bardenet, Doucet and Holmes 2015, e.g.) that utilise the asymptotic normality in terms of justifying their approximation to the posterior. Such algorithms are accurate if the posterior is close to Gaussian, but may be inaccurate otherwise, and it is often impossible to quantify the size of the approximation in practice.

Examples and experiments

There are essentially two different ways of using the Zig-Zag skeleton points which we obtain by using e.g. Algorithms 1, 2, or 3.

We can also estimate posterior quantiles by using the quantiles of the sample $\Xi_{1},\ldots,\Xi_{m}$, as with standard MCMC output. An issue with this approach is that we have to decide on the number, $m$, of samples we wish to use. Whilst the more samples we use the greater the accuracy of our approximation to $\pi(f)$, this comes at an increased computational and storage cost. The trade-off in choosing an appropriate value for $m$ is equivalent to the choice of how much to thin output from a standard MCMC algorithm.

It is important that one does not make the mistake of using the switching points of the Zig-Zag process as samples, as these points are not distributed according to $\pi$. In particular, the switching points are biased towards the tails of the target distribution.

An alternative approach is intrinsically related to the continuous time and piecewise linear nature of the Zig-Zag trajectories. This approach consists of continuous time integration of the Zig-Zag process. By the continuous time ergodic theorem, for $f$ as above, $\pi(f)$ can be estimated as

Since the output of the Zig-Zag algorithms consists of a finite number of skeleton points $(T^{i},\Xi^{i},\Theta^{i})_{i=0}^{k}$, we can express this as

2 Beating one ESS per epoch

We use the term epoch as a unit of computational cost, corresponding to the number of iterations required to evaluate the complete gradient of $\log\pi$. This means that for the basic Zig-Zag algorithm (without sub-sampling), an epoch consists of exactly one iteration, and for the sub-sampled variants of the Zig-Zag algorithm, an epoch consists of $n$ iterations. The CPU running times per epoch of the various algorithms we consider are equal up to a constant factor. To assess the scaling of various algorithms, we use ESS per epoch. The notion of ESS is discussed in the supplementary material (Bierkens, Fearnhead and Roberts 2017, Section 2). Consider any classical MCMC algorithm based upon the Metropolis-Hastings acceptance rule. Since every iteration requires an evaluation of the full density function to compute the acceptance probability, we have that the ESS per epoch for such an algorithm is bounded from above by one. Similar observations apply to all other known MCMC algorithms capable of sampling asymptotically from the exact target distribution.

There do exist several conceptual innovations based on the idea of sub-sampling, which have some theoretical potential to overcome the fundamental limitation of one ESS per epoch sketched above.

The Pseudo-Marginal Method (PMM, Andrieu and Roberts 2009) is based upon using a positive unbiased estimator for a possibly unnormalized density. Obtaining an unbiased estimator of a product is much more difficult than obtaining one for a sum. Furthermore, it has been shown to be impossible to construct an estimator that is guaranteed to be positive without other information about the product, such as a bound on the terms in the product (Jacob and Thiery 2015). Therefore the PMM does not apply in a straightforward way to vanilla MCMC in Bayesian inference.

In the supplementary material (Bierkens, Fearnhead and Roberts 2017, Section 3) we analyse the scaling of Stochastic Gradient Langevin Dynamics (SGLD, Welling and Teh 2011) in an analogous fashion to the analysis of ZZ and ZZ-CV in Section 5. From this analysis we conclude that it is in general not possible to implement SGLD in such a way that the ESSpE has a larger order of magnitude than $O(1)$. We compare SGLD to Zig-Zag in experiments of Sections 6.3 and 6.5.

3 Mean of a Gaussian distribution

Consider the illustrative problem of estimating the mean of a Gaussian distribution. This problem has the advantage that it allows for an analytical solution which can be compared with the numerical solutions obtained by Zig-Zag Sampling and other methods. Conditional on a one-dimensional parameter $\xi$, the data is assumed to be i.i.d. from $N(\xi,\sigma^{2})$. Furthermore a $N(0,1/\rho^{2})$ prior on $\xi$ is specified. Data are generated from the true distribution $N(\xi_{0},\sigma^{2})$ for some fixed $\xi_{0}$. For a detailed description of the experiment and computational bounds, see Section 4 of the supplementary material.

Results for this experiment are displayed in Figure 2. The MSE for the second moment using SGLD does not decrease beyond a fixed value, indicating the presence of bias in SGLD. This bias does not appear in the different versions of Zig-Zag sampling, agreeing with the theoretical result that ergodic averages over Zig-Zag trajectories are consistent. Furthermore we see a significant relative increase in efficiency for ZZ-(so)CV over basic ZZ when the number of observations is increased, agreeing with the scaling results of Section 5. A poor choice of reference point (as in ZZ-soCV) is seen to have only a small effect on the efficiency.

4 Logistic regression

Combined with a flat prior distribution, this induces a posterior distribution $\xi$ given observations of $(x^{j},y^{j})$ for $j=1,\dots,n$; see the supplementary material for implementational details (Bierkens, Fearnhead and Roberts 2017, Section 5).

The results of this experiment are shown in Figure 3. In both the plots of ESS per epoch (see (a) and (c)), the best linear fit for ZZ-CV has slope approximately 0.95, which is in close agreement with the scaling analysis of Section 5. The other algorithms have roughly a horizontal slope, corresponding to a linear scaling with the size of the data. We conclude that, among the algorithms tested, ZZ-CV is the only algorithm for which the ESS per CPU second is approximately constant as a function of the size of the data (see Figure 3, (b) and (d)). Furthermore ZZ-CV obtains an ESSpE which is roughly linearly increasing with the number of observations $n$ (see Figure 3,(a) and (c)). whereas the other versions of the Zig-Zag algorithms, and MALA, have an ESSpE which is approximately constant with respect to $n$. These statements apply regardless of the dimensionality of the problem.

5 A non-identifiable logistic regression example with unbounded Hessian

In Figure 4 we compare trace plots for the Zig-Zag algorithms (ZZ, ZZ-CV) to trace plots for Stochastic Gradient Langevin Dynamics (SGLD) and the Consensus Algorithm Scott et al. 2016. SGLD and Consensus are seen to be strongly biased, whereas ZZ and ZZ-CV target the correct distribution. However this comes at a cost: ZZ-CV loses much of its efficiency in this situation (due to the combination of lack of posterior contraction and unbounded Hessian); in particular it is not super-efficient. The use of multiple reference points may alleviate this problem, see also the discussion in Section 7.

Discussion

We have introduced the multi-dimensional Zig-Zag process and shown that it can be used as an alternative to standard MCMC algorithms. The advantages of the Zig-Zag process are that it is a non-reversible process, and thus has the potential to mix better than standard reversible MCMC algorithms, and that we can use sub-sampling ideas when simulating the process and still be guaranteed to sample from the true target distribution of interest. We have shown that it is possible to implement sub-sampling with control-variates in a way that we can have super-efficient sampling from a posterior: the cost per effective sample size is sub-linear in the number of data points. We believe the latter aspect will be particularly useful for applications where the computational cost of calculating the likelihood for a single data point is high.

As such, the Zig-Zag process holds substantial promise. However, being a completely new method, there are still substantial challenges in implementation which will need to be overcome for Zig-Zag to reach the levels of popularity of standard discrete-time MCMC. The key challenges to implementing the Zig-Zag efficiently are

to simulate from the relevant time-inhomogeneous Poisson process; and

in order to realise the advantages of Zig-Zag for large datasets, reasonable centering points need to be found before commencing the MCMC algorithm itself.

For the first of these challenges, we have shown how this can be achieved through bounding the rate of the Poisson process, but the overall efficiency of the simulation algorithm then depends on how tight these bounds are. In Subsection 3.1 we describe efficient ways to carry this out. Moreover, as pointed out by a reviewer, there is a substantial literature on simulating stochastic processes that involve simulating such time-inhomogeneous Poisson processes Gibson and Bruck 2000; Anderson 2007. Ideas from this literature could be leveraged both to extend the class of models for which we can simulate the Zig-Zag process, and also to make implementation of simulation algorithms more efficient.

The second challenge applies when using the ZZ-CV algorithm to obtain super-efficiency for big data as discussed in Subsection 4.3. Although in our experience finding appropriate centering points is rarely a serious problem, it is difficult to give a prescriptive recipe for this step.

On the face of it, these challenges may limit the practical applicability of Zig-Zag, at least in the short term. With that in mind, we have released an R/Rcpp package for logistic regression, as well as the code which reproduces the experiments of Section 6 Bierkens 2017.

In addition, while Zig-Zag is an exact approximate simulation method, there are various short-cuts to speed it up at the expense of the introduction of an approximation. For instance, there are already ideas of approximately simulating the continuous-time dynamics, through approximate bounds on the Poisson rate Pakman et al. 2016. These ideas can lead to efficient simulation of the Zig-Zag process for a wide class of models, albeit with the loss of exactness. Understanding the errors introduced by such an approach is an open area.

The most exciting aspect of the Zig-Zag process is the super-efficiency we observe when using sub-sampling with control variates. Already this idea has been adapted and shown to apply to other recent continuous-time MCMC algorithms Fearnhead et al. 2018; Pakman et al. 2016. We have shown in Subsection 6.5 that Zig-Zag can be applied effectively within highly non-Gaussian examples where rival approximate methods such as SGLD and the Consensus Algorithm are seriously biased. So there is no intrinsic reason to expect Zig-Zag to rely on the target distribution being close to Gaussian, although posterior contraction and the ability to find tight Poisson process rate bounds play important roles as we saw in our examples. There is much to learn about how the efficiency of Zig-Zag depends on the statistical properties of the posterior distribution. However, unlike its approximate competitors, Zig-Zag will still remain an exact approximate method whatever the structure of the target distribution.

In truly ‘big data’ settings, in principle we still need to process all the data once, although a suitable reference point can be determined using a subset of the data, we do need to evaluate the full gradient of the log density once at this reference point, and this computation is $O(n)$. This operation however is much easier to parallelize than MCMC is, and after this approximately independent samples can be obtained at a cost of $O(1)$ each. Thus if we wish to obtain $k$ approximately independent samples, the computational efficiency of ZZ-CV is $O(k+n)$ while the complexity of traditional MCMC algorithms is $O(kn)$. This is confirmed by the experiment in Section 6.4.

The idea for control variates we present in this paper is just one, possibly the simplest, implementation of this idea. There are natural extensions to deal with e.g. multi-modal posteriors or situations where we do not have posterior concentration for all parameters. The simplest of these involve using multiple reference points and monitoring the computational bound we get within the CV-ZZ algorithm and switching to a different algorithm when we stray so far from a reference point that this bound becomes too large. More sophisticated approaches include using the ideas from Dubey et al. 2016, where we introduce a reference point for each data point and update the reference points for data within the subsample at each iteration of the algorithm. This would lead to the estimate of the gradient that we center our control variate estimator around to depend on the recent history of the Zig-Zag process, and thus could be accurate even if we explore multiple modes or the tails of the target distribution.

The authors are grateful for helpful comments from referees, the editor and the associate editor which have improved the paper. Furthermore the authors acknowledge Matthew Moores (University of Warwick) for helpful advice on implementing the Zig-Zag algorithms as an R package using Rcpp. All authors acknowledge the support of EPSRC under the ilike grant: EP/K014463/1.

Supplementary Material

Supplement: Supplement to “The Zig-Zag Process and Super-Efficient Sampling for Bayesian Analysis of Big Data” (doi: COMPLETED BY THE TYPESETTER; .pdf). Mathematics of the Zig-Zag process, scaling of SGLD, details on the experiments including how to obtain computational bounds.

References