Piecewise Deterministic Markov Processes for Scalable Monte Carlo on Restricted Domains

Introduction

Markov chain Monte Carlo (MCMC) methods have been central to the wide-spread use of Bayesian methods. However their applicability to some modern applications has been limited due to their high computational cost, particularly in big-data, high-dimensional settings. This has led to interest in new MCMC methods, particularly non-reversible methods which can mix better than standard reversible MCMC , and variants of MCMC that require accessing only small subsets of the data at each iteration .

One of the main technical challenges associated with likelihood-based inference for big data is the fact that likelihood calculation is computationally expensive (typically $O(N)$ for data sets of size $N$). MCMC methods built from piecewise deterministic Markov processes (PDMPs) offer considerable promise for reducing this $O(N)$ burden, due to their ability to use sub-sampling techniques, whilst still being guaranteed to target the true posterior distribution . Furthermore, factor graph decompositions of the target distribution can be leveraged to perform sparse updates of the variables .

PDMPs explore the state space according to constant velocity dynamics, but where the velocity changes at random event times. The rate of these event times, and the change in velocity at each event, are chosen so that the position of the resulting process has the posterior distribution as its invariant distribution. We will refer to this family of sampling methods as Piecewise Deterministic Monte Carlo methods (PDMC).

Piecewise Deterministic Monte Carlo on Restricted Domains

Here we present the general PDMC algorithm in a restricted domain. Specific implementations of PDMC algorithms can be derived as continuous-time limits of familiar discrete-time MCMC algorithms , and these derivations convey much of the intuition behind why the algorithms have the correct stationary distribution. Our presentation of these methods is different, and more general. We first define a simple class of PDMPs and show how these can be simulated. We then give simple recipes for how to choose the dynamics of the PDMP so that it will have the correct stationary distribution.

To ensure that $X_{t}$ remains confined within $\mathcal{O}$ the velocity of the process is updated whenever $X_{t}$ hits $\partial\mathcal{O}$ so that the process moves back into $\mathcal{O}$. We shall refer to such updates as reflections even though they need not be specular reflections.

The resulting stochastic process is a Piecewise Deterministic Markov Process (PDMP, ). For it to be useful as the basis of a Piecewise Deterministic Monte Carlo (PDMC) algorithm we need to (i) be able to easily simulate this process; and (ii) have simple recipes for choosing the intensities, $(\lambda_{i})_{i=1}^{N}$, and transition kernels, $(Q_{i})_{i=1}^{N}$, such that the resulting process has $\pi(x)$ as its marginal stationary distribution. We will tackle each of these problems in turn.

The key challenge in simulating our PDMP is simulating the event times. The intensity of events is a function of the state of the process. But as the dynamics between event times are deterministic, we can easily represent the intensity for the next event as a deterministic function of time. Suppose that the PDMP is driven by a single inhomogeneous Poisson process with intensity function

We can simulate the first event time directly if we have an explicit expression for the inverse function of the monotonically increasing function

In this case the time until the next event is obtained by (i) simulating a realization, $y$ say, of an exponential random variable with rate $1$; and (ii) setting the time until the next event as the value $\tau$ that solves $\int_{0}^{\tau}\widetilde{\lambda}(s;X_{t},V_{t})\,ds=y$.

Inverting (1) is often not practical. In such cases simulation can be carried out via thinning . This requires finding a tractable upper bound on the rate, $\overline{\lambda}(u)\geq\widetilde{\lambda}(u;X_{t},V_{t})$ for all $u>0$. Such an upper bound will typically take the form of a piecewise linear function or a step function. Note that the upper bound $\overline{\lambda}$ is only required to be valid along the trajectory $u\mapsto(X_{t}+uV_{t},V_{t})$ in $\mathcal{O}\times\mathcal{V}$. Therefore the upper bound may depend on the starting point $(X_{t},V_{t})$ of the line segment we are currently simulating. We then propose potential events by simulating events from an inhomogenous Poisson process with rate $\overline{\lambda}(u)$, and accept an event at time $u$ with probability $\widetilde{\lambda}(u;X_{t},V_{t})/\overline{\lambda}(u)$. The time of the first accepted event will be the time until the next event in our PDMP.

To handle boundary reflections, at every given time $t$, we also keep track of the next reflection event in the absence of a switching event, i.e. we compute

Although theoretically we may choose a new velocity pointing outwards and have an immediate second jump, we will for algorithmic purposes assume that the probability measure $Q_{b}(x,u,\cdot)$ for $(x,u)\in\partial\mathcal{\mathcal{}}{O}\times\mathcal{V}$ is concentrated on those directions $v$ for which $(v\cdot n(x))\leq 0$, where $n(x)$ is the outward normal at $x\in\partial\mathcal{O}$.

For a PDMP driven by $N$ inhomogeneous Poisson processes with intensities $(\lambda_{i})_{i=1}^{N}$ the previous steps lead to the following algorithm for simulating the next event of our PDMP. This algorithm can be iterated to simulate the PDMP for a chosen number of events or a pre-specified time-interval.

Initialize: Set $t$ to the current time and $(X_{t},V_{t})$ to the current position and velocity.

Determine bound: For each $i\in 1,\ldots,N$, find a convenient function $\overline{\lambda}_{i}$ satisfying $\overline{\lambda}_{i}(u)\geq\widetilde{\lambda_{i}}(u;X_{t},V_{t})$ for all $u\geq 0$, depending on the initial point $(X_{t},V_{t})$ from which we are departing.

Propose event: For $i=1,\ldots,N$ simulate the first event times $\tau_{i}^{\prime}$ of a Poisson process with rate function $\overline{\lambda}_{i}$. Compute the next boundary reflection time $\tau_{b}$.

Let $i_{\min}=\argmin_{j=1,\ldots,N}\tau^{\prime}_{j}$ and $\tau^{\prime}=\tau^{\prime}_{i_{min}}$.

If $\tau_{b}<\tau^{\prime}$ then set $\tau=\tau_{b}$; set $X_{t+\tau}=X_{t}+\tau V_{t}$; sample a new velocity $V_{t+\tau}\sim Q_{b}(X_{t+\tau},V_{t},\cdot)$.

accept the event at time $\tau=\tau^{\prime}$.

Upon acceptance: set $X_{t+\tau}=X_{t}+\tau V_{t}$; sample a new velocity $V_{t+\tau}\sim Q_{i_{\min}}(X_{t+\tau},V_{t},\cdot)$.

Upon rejection: set $X_{t+\tau}=X_{t}+\tau V_{t}$ and set $V_{t+\tau}=V_{t}$.

Update: Record the time $t+\tau^{\prime}$ and state $(X_{t+\tau^{\prime}},V_{t+\tau^{\prime}})$.

2 Output of PDMC algorithms

The output of these algorithms will be a sequence of event times $t_{1},t_{2},t_{3},\ldots,t_{K}$ and associated states $(X_{1},V_{1})$, $(X_{2},V_{2}),\ldots,(X_{K},V_{K})$. To obtain the value of the process at times $t\in[t_{k},t_{k+1})$, we can linearly interpolate the continuous path of the process between event times, i.e. $X_{t}=X_{t_{k}}+V_{k}(t-t_{k})$. Time integrals $\int_{0}^{t}f(X_{s})\,ds$ of a function $f$ of the process $X_{t}$ can often be computed analytically from the output of the above algorithm. If not they can be approximated by numerically integrating the one dimensional integral along the piecewise linear trajectory of the PDMP. Alternatively we can sample the PDMP at a set of evenly spaced time points along the trajectory and use this collection as an approximate sample from our target distribution.

Under the assumption that the resulting PDMP is ergodic (for sufficient conditions see e.g. ) and that the marginal density on $\mathcal{O}$ of the stationary distribution of $(X_{t},V_{t})$ is equal to $\pi$, we have the following version of the law of large numbers for the PDMP $(X_{t},V_{t})_{t\geq 0}$: For all $f\in L^{2}(\pi)$ we have that, with probability one,

It is this formula which allows us to use PDMPs for Monte Carlo purposes.

3 Choosing the intensity and transition kernels

Assume, as most existing PDMC methods do , that the target density, $\pi(x):\mathcal{O}\rightarrow(0,\infty)$ is differentiable. Under this condition we can provide criteria on the switching intensities $(\lambda_{i})$ and transition kernels $Q_{i}$ and $Q_{b}$ which must hold for a given probability distribution to be a stationary distribution of $Z_{t}$. We shall consider stationary distributions for which $x$ and $v$ are independent, i.e. distributions of the form $\pi(x)dx\otimes\rho(dv)$ on $E$. Furthermore we assume that $\pi(x)\propto\exp(-U(x))$ where $U$ is continuously differentiable.

A sufficient condition for (2) is that each $Q_{i}$ is reversible with respect to $\rho$, i.e. for every $i=1,\ldots,N$ and $x\in\mathcal{O}$, we have that $Q_{i}(x,v,du)\rho(dv)=Q_{i}(x,u,dv)\rho(du)$.

Moreover, we shall require the following condition which relates the probability flow with the switching intensities $\lambda_{i}$:

Finally, the boundary transition kernel should satisfy

where for $x\in\partial\mathcal{O}$, we denote by $n(x)$ the outward unit normal of $\partial\mathcal{O}$.

4 Example: The Bouncy Particle Sampler

Current PDMC algorithms differ in terms of how the $Q_{i}$ and $\lambda_{i}$ are chosen such that the above equation holds for some simple distribution for the velocity. Here we discuss how the Bouncy Particle Sampler (BPS), introduced in and explored in , is an example of the framework introduced here. In the supplementary material, Section 1.1, the Zig-Zag sampler is described as a second example. In the following example $\delta_{x}$ denotes the Dirac-measure centered in $x$.

where $P_{y}:z\mapsto\frac{z\cdot y}{\|y\|^{2}}y$ denotes orthogonal projection along the one dimensional subspace spanned by $y$.

As noted in this algorithm suffers from reducibility issues. These can be overcome by refreshing the velocity by drawing a new velocity independently from $\rho(dv)$. In the simplest case the refreshment times come from an independent Poisson process with constant rate $\lambda_{\text{ref}}$. This also fits in the framework above by choosing $\widetilde{\lambda}=\lambda_{\text{BPS}}+\lambda_{\text{ref}}$ and

As boundary transition kernel it is natural to choose

for $s\in\partial\mathcal{O}$, so that the process $X_{t}$ reflects specularly at the boundary (i.e. angle of incidence equals angle of reflection of process with respect to the boundary normal). It is straightforward to check that condition (2) holds at the boundary and that (5) is satisfied.

As a generalization of the BPS, one can consider a preconditioned version, which is obtained by introducing a constant positive definite symmetric matrix $M$ to rescale the velocity process. The choice of $M$ plays a very similar role to the mass matrix in HMC, and careful tuning can give rise to dramatic increases in performance .

Subsampling

When using PDMC to sample from a posterior, we can use sub-samples of data at each iteration of the algorithm, as described in , which reduces the computational complexity of the algorithm from $O(N)$ to $O(1)$, where $N$ is the size of the data, without affecting the theoretical validity of the algorithm. In the following we will assume that we can write the posterior as $\pi(x)\propto\prod_{i=1}^{N}f(y_{i};x),$ for some function $f$. For example this would be the likelihood for a single IID data point times the $1/N$th power of the prior.

The idea of using sub-sampling, within say the Bouncy Particle Sampler (BPS), is that at each iteration of our PDMC algorithm we can replace $\nabla U(x)$ by an unbiased estimator in step (3). We need to use the same estimate both when calculating the actual event rate in the accept/reject step and, if we accept, when simulating the new velocity. The only further alteration we need to the algorithm is to choose an upper bound $\overline{\lambda}$ that holds for all realizations of $\widehat{\nabla U}$. A more comprehensive explanation of this argument can be found in in the context of the Zig-Zag sampler, and in for the bouncy particle sampler.

We first present a way for estimating $\nabla U$ unbiasedly using control variates . For any $x,\hat{x}\in\mathcal{O}$ we note that $\nabla U(x)=\nabla U(\hat{x})+\left[\nabla U(x)-\nabla U(\hat{x})\right]$. We can then introduce the estimator $\widehat{\nabla U}(x)$ of $\nabla U(x)$ by

where $I$ is drawn uniformly from $\{1,\ldots,N\}$.

Note that this gain in computational efficiency does not come for free, as it follows from Jensen’s inequality that the overall rate of events will be higher. This makes mixing of the PDMC process slower. It is also immediate that the bound, $\overline{\lambda}$, we will have to use will be higher. However show that if our estimator of $\widehat{\nabla U}(x)$ has sufficiently small variance, then we can still gain substantially in terms of efficiency. In particular they give an example where the CPU cost effective sample size does not grow with $N$ – by comparison all standard MCMC algorithms would have a cost that is at least linear in $N$.

To obtain such a low-variance estimator requires a good choice of $\hat{x}$, so that with high probability $x$ will be closer to $\hat{x}$. This involves a preprocessing step to find a value $\hat{x}$ close to the posterior mode, a preprocessing step to then calculate $\nabla U(\hat{x})$ is also needed.

We now illustrate how to find an upper bound on the event rate. Following , if we assume $L$ is a uniform (in space and $i$) upper bound on the largest eigenvalue of the Hessian of $U^{i}$, and if $\|v\|=1$:

Thus the upper bound on the intensity is of the form $\bar{\lambda}(\tau)=a+b\cdot\tau$ with $a,b\geq 0$. In this case the first arrival time can be simulated as follows

An alternative and complementary approach to improve the efficiency of this subsampling procedure is to use an estimator of the gradient (6) where $I$ is drawn according to a distribution dependent on the observations .

Software and Numerical Experiments

A open-source Julia package PDMP.jl has been developed to provide efficient implementations of various recently developed piecewise deterministic Monte Carlo methods for sampling in (possibly restricted) continuous spaces. A variety of algorithms are implemented including the Zig-Zag sampler and the Bouncy Particle Sampler with full and local refreshment along with control variate based sub-sampling for these methods. The package has been specifically designed with extensibility in mind, permitting rapid implementation of new PDMP based methods. The library along with code and documentation is available at github.com/alan-turing-institute/PDSampler.jl.

We use Bayesian binary logistic regression as a testbed for our newly proposed methodology and perform a simulation study. The data $y_{i}\in\{-1,1\}$ is modelled by

For simplicity we use a flat prior over the space of parameters values consistent with our constraints. By Bayes’ rule the posterior $\pi$ satisfies

where $\mathcal{O}$ is the space of parameter values consistent with our constraints. We implement the BPS with subsampling. As explained in the introduction, subsampling is a key benefit of using piecewise deterministic sampling methods; see Section 3. We use reflection at the boundary i.e. $Q_{b}(s,v,du)=\delta_{(I-2P_{n(s)})v}(du)$ for $s\in\partial\mathcal{O}$. We can bound the switching intensity by a linear function of time, even when we use the subsampling estimator for the switching rate. See the supplementary material, Section 2, for details on the application of subsampling in this example. We use $n=10,000$ and $p=20$ and generate artificial data based on $\xi$ and $x^{\star}$ whose components are a realization i.i.d. of uniformly distributed random variables satisfying the imposed constraints.

We compare the performance of BPS to standard MALA and HMC schemes, in terms of effective sample size (ESS) per epoch of data evaluation. For each scheme we obtain the distribution of ESS based on $10$ independent realisations of each chain. In Figure 1(a) we plot for each scheme, the distribution of ESS per epoch with respect to the function $f_{1}(x)=\frac{1}{p}(x_{1}+\ldots+x_{p})$. Similarly, In Figure 1(b) we plot the ESS per epoch for each chain with respect to the function $f_{2}(x)=\log\pi(x)$. The performance of MALA and HMC appears commensurate and the BPS demonstrates a clear advantage over both in terms of ESS per epoch.

The HMC and MALA schemes were tuned by minimising the ESS with respect to the step-size, calculated from exploratory runs. For HMC we use $5$ leap-frog steps. We find that we must tune both HMC and MALA to have a small step size due to proposals being rejected at the boundary. The ESS is estimated based on asymptotic variance using the batch means method; see for details.

For specific types of constraints more efficient implementations of HMC and MALA are possible, either by introducing an appropriate transformation of the restricted state space, or by reflecting the posterior distribution along the constraint boundaries. Moreover, we note that there exists a version of HMC which can sample from truncated Gaussian distributions . However, to our knowledge there is no efficient HMC or MALA scheme able to handle generally restricted domains.

The Bouncy Particle Sampler for this model was implemented using PDSampler.jl while the corresponding HMC and MALA samplers implemented with Klara.jl. The code for this numerical experiment along with results are carefully presented in github.com/tlienart/ConstrainedPDMP/.

Discussion

This work provides a framework for describing a general class of PDMC methods which are ergodic with respect to a given target probability distribution. Open questions remain on how the choice of intensity function, velocity transition kernel as well as other parameters of the system influence the overall performance of the scheme. The problem of understanding the true computational cost of such PDMC schemes is more subtle than for classical discrete time MCMC schemes: often one needs to find a balance between fast mixing of the continuous time Markov process and having a switching rate that is relatively cheap to simulate. For example, when using subsampling the mixing of the Markov process is slower than without subsampling, but the computational cost per simulated switch is significantly smaller. Further investigation is required to understand this delicate balance.

Acknowledgements

All authors thank the Alan Turing Institute and Lloyds registry foundation for support. S.J.V. gratefully acknowledges funding through EPSRC EP/N000188/1. J.B., P.F. and G.R. gratefully acknowledge EPSRC ilike grant EP/K014463/1. A.B.D. acknowledges grant EP/L020564/1 and A.D. acknowledges grant EP/K000276/1. We kindly acknowledge comments from the editor and anonymous reviewers which have significantly improved the exposition in this paper.

References

Stationary distribution for PDMPs on restricted domains

From [8, Section 5] the process $Z_{t}$ will have infinitesimal generator given by the closure of the operator

where $\mathcal{D}(\mathcal{L})$ is the set of functions which are continuously differentiable with respect to $x$ on $\mathcal{O}$, which is decaying to infinity as $\lVert x\rVert\rightarrow\infty$ and such that

for all $(x,v)\in\partial\mathcal{O}\times\mathcal{V}$. Based on this identification of the infinitesimal generator we can now provide a formal proof that the conditions of Proposition 1 of the paper are sufficient to ensure invariance of $\pi\otimes\rho$.

Without loss of generality we take $\pi(x)=\exp(-U(x))$, i.e. the proportionality factor in $\pi(x)\propto\exp(-U(x))$ is assumed to be 1. We shall only provide a formal proof of this result, by demonstrating that

so that $\mathcal{L}$ is infinitesimally invariant. A rigorous proof would require establishing that $\mathcal{D}(\mathcal{L})$ as defined above is a core for the extended generator. This is a technical result which we defer for future work.

where the boundary term arises from integration by parts with respect to $x$. Considering the boundary integral, by applying (4) (in the paper) which is assumed to hold on $\partial\mathcal{O}$ and (S2) (above) we obtain

so that the boundary term evaluates to zero.

so that $\pi(x)\,dx\otimes\rho(dv)$ is infinitesimally invariant with respect to $Z_{t}.$ ∎

Another possible behaviour at the boundary is to generate the new reflected direction independently of the angle of incidence. This will also preserve the invariant distribution provided that $\rho$ is isotropic.

Consider the process $Z_{t}$ as in the previous proposition, such that conditions (2) and (3) (of the paper) hold and the distribution $\rho$ has mean zero. Then $\pi(x)\,dx\otimes\rho(dv)$ will be an invariant distribution for the process $Z_{t}$ if $Q_{b}(x,v,du)$ is independent of $v$ for all $x\in\partial\mathcal{O}$.

Let $f\in\mathcal{D}(\mathcal{L})$, so that $f$ satisfies (S2). By the assumptions on $Q_{b}$ in Proposition 1 (of the paper), this implies that $f(x,v)=f(x)$ for all $x\in\partial\mathcal{O}$. Following the proof of Proposition 1 above, the boundary integral term becomes

which is zero if $\rho$ has mean zero, as required. ∎

The Zig-Zag sampler

The Zig-Zag sampler can be recovered by choosing $N=d$ and picking as velocity space $\mathcal{V}=\{-1,+1\}^{d}$ equipped with discrete uniform distribution $\rho$, defining switching rates $\lambda_{i}(x,v)=\max(v_{i}\partial_{x_{i}}U(x),0)$. The corresponding switching kernels over new directions are given by

where $F_{i}:\mathcal{V}\rightarrow\mathcal{V}$ denotes the operation of flipping the $i$-th component, i.e. $(F_{i}v)(i)=-v(i)$, and $(F_{i}v)(j)=v(j)$ for $j\neq i$.

Derivation of dominating intensity for logistic regression example

A valid choice of $L$ can be derived as follows: Notice that $\left(\log f(z)\right)^{\prime}=f(-z)$ and $f^{{}^{\prime}}(z)=f(z)\left(1-f(z)\right)$ so that we obtain

So Equation (7) (of the paper) holds with

This is a linear upper bound on the intensity which can be used to sample according to (8) (of the paper) and then used for thinning as introduced in Section 1 of the paper.