Large sample scaling analysis of the Zig-Zag algorithm for Bayesian inference

Piecewise deterministic Markov processes [PDMPs, 21] are a class of non-reversible Markov processes which have recently emerged as alternatives to traditional MCMC methods based on the Metropolis-Hastings algorithm. Practically implementable PDMP algorithms such as the Zig-Zag sampler [11], the Bouncy Particle Sampler [16] and others [18, 31, 33, 12, 7, 8, 29] have provided highly encouraging results for Bayesian inference as well as in Physics [26, for example]. As a result, there has been a surge of interest in the Monte Carlo community to better understand the theoretical properties of PDMPs [see e.g. 24, 22, 1, 9, 15, 14, 13]. This body of work demonstrates powerful convergence properties with impressive scalability as a function of the dimension.

However, the most exciting property of PDMPs, at least for statistical applications is that of principled sub-sampling [11], whereby randomly chosen subsets of data can be used to estimate likelihood values within the algorithm without biasing its invariant distribution. The authors only cover sub-sampling for the Zig-Zag algorithm, although the unbiasedness property holds for a rather general class of PDMP MCMC algorithms. This unbiasing property is in sharp contrast to most MCMC methods [32, 28, 5], although see [19]. In fact, [11] gives experiments to show that even for data sets of large size $n$, subsets of size $1$ can be used to estimate the likelihood, producing a computational speed-up of $O(n)$. This leads to the property of superefficiency whereby, after initialization, the entire running cost of the algorithm is dominated for large $n$ by the cost of evaluating the likelihood once. In [11], the authors provide two algorithms which utilise this sub-sampling idea: the straightforward ZZ-SS algorithm together with a control-variate variant ZZ-CV which will be formally introduced in Section 2.

In [11], the authors also point out that although principled sub-sampling gives these dramatic improvements in implementation time, the dynamics of ZZ-SS and ZZ-CV differ from that of canonical Zig-Zag: sample paths exhibit more diffusive behaviour to the extra velocity switches induced by the variation in the estimate of the likelihood. This might lead to slower convergence and inferior mixing of the underlying stochastic process (see Figure 2). For one-dimensional Zig-Zag, [9] studies the diffusive behaviour resulting from excess switches outside of a statistical context. In [2], the authors show that canonical Zig-Zag has superior mixing (in the sense of minimising asymptotic variance for estimating a typical $L_{2}$ function) than any alternative with excess switches. Since both ZZ-SS and ZZ-CV can be viewed as Zig-Zag algorithms with excess switches, this raises the question of whether the inferior mixing of sub-sampling Zig-Zag algorithms relative to their canonical parent might negate their computational advantage.

To date, no work has been done to study the extent of the algorithm mixing deterioration caused by these excess switches induced by sub-sampling. The goal of this work is to give the first rigorous theoretical results to understand this problem. We shall study the behaviour of the various Zig-Zag algorithms for large data sets, interplaying closely with large-sample Bayesian asymptotics, both in the correctly specified and misclassified cases.

Our contributions in this paper are as follows.

1. 1.

   For the Zig-Zag algorithm, we study the effect of sub-sampling on the dynamics of the underlying stochastic process as the data set size $n$ increases to infinity. In Theorem 3.3, we give weak convergence results for the trajectories of the process to smooth solutions of a system of ODEs. These ODEs can be explicitly characterized as presented in equation (3.4).

2. 2.

   We also study the effect of sub-sampling on the algorithm mixing in the stationary phase. We find that as $n$ becomes infinitely large, the trajectories of ZZ-SS converge to a reversible Ornstein-Uhlenbeck process as given in Theorem 4.2, mixing in $O(1)$ time. Whereas, both canonical Zig-Zag and ZZ-CV converge to a PDMP that mixes in time $O(n^{-1/2})$ as presented in Theorem 4.6 and 4.11.

3. 3.

   Our theoretical study leads to further guidance for implementing sub-sampling with Zig-Zag. In particular, our results highlight that for heavy-tailed models the vanilla sub-sampling outperforms ZZ-CV in the tails. In Section 5, we propose a mixed sub-sampling scheme which in addition to being exact, supersedes both ZZ-SS and ZZ-CV in performance.

4. 4.

   Finally, we outline how these methods can be generalised to more general classes of PDMPs.

In parametric Bayesian inference, given sufficient regularity in the model and the true data-generating distribution, it is well known that the posterior distribution converges to a degenerate distribution as the data set size $n\to\infty$ irrespective of whether the true data-generating distribution belongs to the assumed model or not [6, 25]. The point of concentration is the parameter value (if it exists and is unique) minimizing the Kullback-Liebler (KL) divergence between the true data-generating distribution and the assumed model. In well-specified settings, this is the same as the true parameter. Our results in the transient phase (Section 3) show that for large $n$, the Zig-Zag trajectories behave like deterministic smooth flows with a drift pushing them in the direction of decreasing KL-divergence. While the canonical Zig-Zag attains optimal speed ($\pm 1$) in the limit, the sub-sampling versions remain sub-optimal with a damping factor that depends on the model. In particular, we find that the vanilla sub-sampling starts with the optimal speed at infinity, suffering from an extreme slowdown as the process gets closer to the minimizer. On the other hand for sub-sampling with control variates, the drift remains positive at stationarity achieving optimal speed for model densities that are log-concave in the parameter.

Let $x_{0}$ be the unique parameter value minimizing the KL-divergence between the true data-generating distribution and the true parameter. In stationarity, the high concentration of the posterior mass around $x_{0}$ restricts the movement of the Zig-Zag process. Indeed, the canonical Zig-Zag remains close to $x_{0}$ once it hits $x_{0}$ (Theorem 4.1). In general, the smooth approximations from the transient phase are no longer informative since they only capture the large-scale movements of the Zig-Zag trajectories. Let $\hat{x}_{n}$ be a sequence of suitable estimators of $x_{0}$ such that the Bernstein von-Mises Theorem holds: in the rescaled coordinate $\xi(x)=n^{1/2}(x-\hat{x}_{n})$, the posterior converges to a Gaussian distribution with zero mean and variance-covariance matrix given by the inverse of the Fisher’s information [25]. Under regularity conditions, $\hat{x}_{n}$ can be taken to be the maximum likelihood estimators (MLE) for the (possibly misspecified) model. In the $\xi-$coordinate thus, the trajectories of the Zig-Zag process have a non-trivial limit. Although, the canonical Zig-Zag and control variate Zig-Zag mix at a rate $O(n^{-1/2})$ in contrast to $O(1)$ for the vanilla sub-sampling (Section 4).

In stationarity, the computational complexity of a sampling algorithm can be described as:

For Zig-Zag algorithms this further breaks down to

The canonical Zig-Zag incurs an $O(n)$ cost per proposed switch because the full data set needs to be processed to compute the switching rate at each proposed switching time. In contrast to this, the sub-sampling versions only cost $O(1)$ per proposed switch as only a small subset of the whole data set is utilized to estimate the switching rate at each proposed switching time. The time taken to obtain an essentially independent sample is $O(n^{-1/2})$ for canonical Zig-Zag and Zig-Zag with control variates whereas it is $O(1)$ for vanilla sub-sampling. The number of proposed switches depends on the computational bounds on the switching rates used in the implementation of the algorithm. In the rescaled coordinate, the effective switching rates are of the magnitude $O(n^{1/2})$ for canonical Zig-Zag and Zig-Zag with control variates and $O(n)$ for vanilla sub-sampling. Assuming the computational bounds are tight and an efficient implementation of the algorithm is possible, Table 1 summarizes the computational complexity of different versions of the Zig-Zag algorithm in the stationary phase.

The computational complexity of the canonical Zig-Zag and Zig-Zag with vanilla sub-sampling per independent sample is $O(n)$ provided tight computational bounds are available. In contrast, Zig-Zag with control variates achieves a computational speed-up of $O(n)$ when the sequence of reference points is chosen sufficiently close to the sequence of estimators $\hat{x}_{n}$.

The paper is constructed as follows. Section 2 details the Zig-Zag algorithm applied to the Bayesian setting with i.i.d observations along with a discussion of the sub-sampling method. Section 3 contains our main results on the transient phase of the algorithm. We study the stationary phase separately in Section 4. In Section 5, we discuss the results of the transient phase in more detail along with some illustrations on simple models. We conclude with a final discussion in Section 6. All the proofs can be found in the Appendix.

In this section, we review how the Zig-Zag process can be designed to sample from a Bayesian posterior distribution in the setting where the data are assumed to be generated independently from a fixed probability model. For details, we refer to [11].

This section establishes fluid limits for the Zig-Zag process in Bayesian inference. We consider the identifiable situation where, as more data becomes available, the posterior distribution contracts to a point mass. In this case, the gradient of the log-likelihood grows linearly in $n$ and, as a result, the switching intensities become proportionally large (see Proposition 3.1). This results in a limiting process that switches immediately to a ‘favourable’ direction and asymptotically the Zig-Zag process behaves as the solution of an ordinary differential equation as we will detail below. Throughout this section, we assume the following about the sequence of reference points defined in (2.9).

For each $n$, let $(Z^{n}_{t})_{t\geq 0}=(X^{n}_{t},V^{n}_{t})_{t\geq 0}$ be a Zig-Zag process targeting $\Pi^{(n)}$ with switching rates $(\lambda^{n}_{i})_{i=1}^{d}$ defined as in (2.6). For each $n$ and $i=1,\dots,d$, define the sets

Let $H^{n}=\cup_{i=1}^{d}H^{n}_{i}$. Since the switching rates are continuous, $(H^{n})^{c}$ is open. Define for each $n$,

where we set $\tau_{n}=\infty$ if $H^{n}=\emptyset$. For some cemetary state $\Delta$, we consider the sequence of stopped processes $\tilde{Z}^{n}_{t}$ defined as,

The stopped process is again a PDMP that is sent to the cemetery state as soon as the position process enters $H^{n}$ and stays there. The corresponding semigroup can be appropriately modified [20]. However, we introduce the cemetery state only for the sake of rigour; it can be ignored for all practical purposes.

The following proposition is a consequence of the strong law of large numbers for $U-$statistics and applies to any choice of switching rates as defined in Section 1. We remind the reader that we omit the subscripts “can”, “ss”, and “cv” here.

The proof is located in the Appendix A. Proposition 3.1 indicates that the switching rates are $O(n)$. This means that on average the process flips direction every $O(n^{-1})$ time. But since the position component has unit speed in each coordinate, the process only travels $O(n^{-1})$ distance between each direction flip. As $n$ goes to infinity, the number of switches per unit time goes to infinity but the distance travelled between each switch becomes infinitesimally small. Consequently, on the usual time scale, the position component $(X^{n}_{t})_{t\geq 0}$ appears to traverse a smooth path and is feasibly approximated by the solution to an ODE.

For each $n$, define the sequence of drifts $b^{n}:\mathbb{R}^{d}\to[-1,1]^{d}$ such that,

for all $i=1,\dots,d$. Then, $b^{n}$ is well-defined. For each fixed $\omega\in\Omega$, let $\lambda$ be as in the Proposition 3.1. Define, for each $i=2,\dots,d$,

and let $H=\cup_{i=1}^{d}H_{i}$. Note that the sets $H_{i}$ may vary with $\omega$ e.g. in ZZ-CV where the sum of switching rates depends on the limiting reference point. For each $\omega\in\Omega$, define the asymptotic drift $b(\omega,\cdot):\mathbb{R}^{d}\setminus H(\omega)\to[-1,1]^{d}$ as,

The following proposition follows easily from the last one.

The proof is located in the Appendix A. Extend the solution $(X_{t})_{t\leq\tau}$ to $(\tilde{X}_{t})_{t\geq 0}$ as,

We now state the main result of this section.

The proof is located in the Appendix B.

The asymptotic drift $b$ has the following form (see Appendix A):

where $\lambda$ is the asymptotic switching rate from Proposition 3.1 and $S_{i}(x;y)=\nabla_{x,i}\log f(y;x)$. In the large $n$ limit, the function $b$ entirely characterizes the dynamics of the Zig-Zag trajectories. To understand the effect of different sub-sampling schemes, it is sufficient to understand how the denominator $\lambda_{i}(x,-\boldsymbol{1}_{d})+\lambda_{i}(x,\boldsymbol{1}_{d})$ behaves in each situation. For the canonical Zig-Zag,

for all $x$ such that the expectation is non-zero. This corresponds to the optimal speed achievable by the Zig-Zag process i.e. $\pm 1$. Now suppose the subsample size $m$ is fixed. Then,

where $Y_{1},\dots,Y_{m}\overset{iid}{\sim}P$. Similarly, for Zig-Zag with control variates the denominator is given, when the expectation is non-zero, by

where $Y_{1},\dots,Y_{m}\overset{iid}{\sim}P$ and $X^{*}$ is as in Assumption 3.1.

The sets $H^{n}_{i}$ take different forms for different switching rates. For canonical Zig-Zag, each $H^{n}_{i}$ is the hypersurface $H^{n}_{i}=\{x\in\mathbb{R}^{d}:\sum_{j=1}^{n}s_{i}^{j}(x)=0\}$. For ZZ-SS, each $H^{n}_{i}=\emptyset$ since the sum is always positive. For ZZ-CV with reference point $x^{*}_{n}$, the sets $H^{n}_{i}$ can be either as small as the empty sets in the ZZ-SS case e.g. if the reference point is not close to any of the conditional modes. Or they can be as large as the hypersurfaces in the canonical case e.g. when the model density is Gaussian in which case the ZZ-CV switching rates are equivalent to the canonical Zig-Zag irrespective of the reference point $x^{*}_{n}$. However, note that since the canonical rates are the smallest, for any $\lambda^{n}_{i}$ of the form (2.6),

for all $(x,v)\in E$. And so, for any choice of the switching rates,

When $X^{n}_{t}$ hits $H^{n}_{i}$ for some $i$, the total switching rate, and hence the momentum, in the $i-$th coordinate, becomes zero. However, other coordinates may still carry momentum and affect the overall dynamics of the process. This leads to interesting behaviour: the process either jumps back out of $H^{n}$ with a certain velocity or remains stuck to $H^{n}_{i}$ for a random amount of time (see Figure 3). It is not straightforward to extend the proof of Theorem 3.3 as the process behaves $(d-1)-$dimensional on these sets. Consequently, the stability analysis of the sets $H^{n}_{i}$ is beyond the scope of this paper.

As noted in Section 1, the fluid limits only capture deviations in the Zig-Zag trajectories that are big enough on the usual space scale. Suppose Assumptions 2.3 and 2.4 hold. As $n$ becomes large the posterior concentrates around $\hat{x}_{n}$. Consequently, the posterior mass inside the interval $(\hat{x}_{n}-\epsilon,\hat{x}_{n}+\epsilon)$ goes to $1$ for any $\epsilon>0$. This restricts the movement of the Zig-Zag process to regions close to $\hat{x}_{n}$ (see Figure 4). Once the Zig-Zag process targeting $\Pi^{(n)}$ hits $\hat{x}_{n}$, it is hence reasonable to expect that the process stays within $\epsilon-$distance of $\hat{x}_{n}$ with high probability. The following result illustrates this in one dimension under the additional assumption of unimodality of the posterior. The proof is located in the Appendix C.

A sufficient condition for the assumption in Theorem 4.1 to hold is when $f(y;x)$ is log-concave as a function of $x$ for any $y$ (where $f(y;x)$ is the data generating density for a given parameter $x$), and $\hat{x}_{n}$ satisfying $\sum_{j}s^{j}(\hat{x}_{n})=0$ exists. Since the sequence $\hat{x}_{n}$ converges to $x_{0}$ in probability, Theorem 4.1 implies that once it gets close to $x_{0}$, the canonical Zig-Zag stays close to $x_{0}$ as $n\to\infty$. To obtain a non-trivial limit in stationarity the process needs to be analyzed on its natural scale which depends on the rate of the posterior contraction. Consider the reparameterization,

The posterior distribution $\Pi^{(n)}$ in terms of the $\xi-$parameter converges to a multivariate Gaussian distribution. Under this reparameterization, it is natural to expect that the Zig-Zag process converges to a non-trivial stochastic process targeting the asymptotic Gaussian distribution.

In the $\xi$ coordinate, the Zig-Zag process has a speed $n^{1/2}$. However, the switching rates for ZZ-SS and ZZ-CV in terms of $\xi$ may differ in magnitudes and hence require separate analysis. In what follows we assume that the process starts from stationarity i.e. $(X^{n}_{0},V^{n}_{0})\sim\Pi^{(n)}\times\text{Uniform}(\{-1,1\}^{d})$ for all $n$.