#  Stochastic Gradient Piecewise Deterministic Monte Carlo Samplers

Paul Fearnhead
School of Mathematical Sciences,
Lancaster University
<EMAIL_ADDRESS>
Sebastiano Grazzi
Department of Statistics,
University of Warwick
<EMAIL_ADDRESS>
Chris Nemeth
School of Mathematical Sciences,
Lancaster University
<EMAIL_ADDRESS>
Gareth Roberts
Department of Statistics,
University of Warwick
<EMAIL_ADDRESS>
###### Abstract

Recent work has suggested using Monte Carlo methods based on piecewise
deterministic Markov processes (PDMPs) to sample from target distributions of
interest. PDMPs are non-reversible continuous-time processes endowed with
momentum, and hence can mix better than standard reversible MCMC samplers.
Furthermore, they can incorporate exact sub-sampling schemes which only
require access to a single (randomly selected) data point at each iteration,
yet without introducing bias to the algorithm’s stationary distribution.
However, the range of models for which PDMPs can be used, particularly with
sub-sampling, is limited. We propose approximate simulation of PDMPs with sub-
sampling for scalable sampling from posterior distributions. The approximation
takes the form of an Euler approximation to the true PDMP dynamics, and
involves using an estimate of the gradient of the log-posterior based on a
data sub-sample. We thus call this class of algorithms stochastic-gradient
PDMPs. Importantly, the trajectories of stochastic-gradient PDMPs are
continuous and can leverage recent ideas for sampling from measures with
continuous and atomic components. We show these methods are easy to implement,
present results on their approximation error and demonstrate numerically that
this class of algorithms has similar efficiency to, but is more robust than,
stochastic gradient Langevin dynamics.

Keywords: Bouncy Particle Sampler; Control-variates; MCMC; Stochastic gradient
Langevin dynamics; Sub-sampling; Zig-Zag Sampler

##  1 Introduction

Whilst Markov chain Monte Carlo (MCMC) has been the workhorse of Bayesian
statistics for the past thirty years, it is known to scale poorly for large
datasets, since each iteration of MCMC requires the evaluation of the log-
posterior density which typically scales linearly with data size. As a result,
approximate MCMC methods that use only a subsample of data at each iteration
have become popular. The first such method, and arguably the most widely used,
is the stochastic gradient Langevin dynamics (SGLD) algorithm of Welling and
Teh (2011). The idea of SGLD is to approximately simulate a Langevin diffusion
that has the posterior as its stationary distribution. The method involves two
approximations: firstly it simulates an Euler discretisation of the Langevin
diffusion; and secondly it approximates the gradient of the log-posterior,
i.e. the drift of the diffusion, based on a subsample of the data. The SGLD
algorithm has been applied to applications such as topic models Baker et al.
(2018), Bayesian neural networks Gawlikowski et al. (2023) and probabilistic
matrix factorisation Şimşekli et al. (2017); and has also been extended to
more general dynamics - e.g. Chen et al. (2014); Ma et al. (2015). See Nemeth
and Fearnhead (2021) for a review.

Recently there has been interest in developing efficient continuous time MCMC
algorithms known as Piecewise Deterministic Markov processes (PDMP) Vanetti et
al. (2017); Davis (1984). Such methods include the Bouncy Particle Sampler
Bouchard-Côté et al. (2018) and the Zig-Zag algorithm Bierkens et al. (2019)
amongst others Fearnhead et al. (2018). Due to their non-reversibility, these
methods often mix better than traditional reversible MCMC algorithms Diaconis
et al. (2000); Bierkens (2016). Furthermore, Bierkens et al. (2019) show how
these methods can be implemented whilst only using a small subsample of data
at each iteration, and yet still target the true posterior distribution. PDMPs
(without subsampling) have been also developed and successfully applied in
computational physics for example for the simulation of hard sphere models
Peters and de With (2012); Bernard et al. (2009); Krauth (2021).

However PDMP samplers can be challenging to implement, particularly when using
subsamples, as they require bounding the gradient of the log-density. There
has been some work on automating the simulation of PDMPs using numerical
methods Corbella et al. (2022); Pagani et al. (2024), but these methods
largely require exact calculation of the gradient of the log-posterior at each
iteration and are incompatible with subsampling ideas.

In this paper, we investigate the approximate simulation of PDMPs with
subsampling as a means of achieving scalable MCMC. The idea is to discretise
time into intervals of length $\varepsilon$, for each interval we sample a
single data point, and then simulate the (approximate) dynamics of the PDMP
using only the information from this data point. Loosely speaking, this can be
viewed as simulating an Euler approximation to the dynamics of the PDMP with
subsampling. This simple idea was suggested, but not investigated, by Bertazzi
et al. (2022) and Bertazzi et al. (2023). We call these methods stochastic-
gradient PDMPs, or SG-PDMPs. We show that they give results that are
competitive with SGLD, but have the advantage of being more robust to large
discretisation sizes. They simulate continuous trajectories, and we implement
a version that leverages this property to easily sample from a trans-
dimensional posterior distribution for models that incorporate variable
selection.

The idea of approximately simulating a PDMP algorithm with subsampling has
been previously suggested by Pakman et al. (2017). Their algorithm, called the
stochastic bouncy particle sampler (SBPS) uses Poisson thinning to simulate
events of the Bouncy Particle Sampler with subsampling. The idea is that, if
we can upper bound the actual event rate, and can simulate events with this
upper bound rate, then we can thin (or remove) events with an appropriate
probability to obtain events simulated with the required rate. For most target
distributions we cannot calculate an upper bound for the rate, so Pakman et
al. (2017) suggests estimating this upper bound based on information on the
rate from sub-samples of data at proposed event times. In practice one of the
main differences between our proposal and SBPS is that, for each discrete
time-step we only need to sample a single data point. By comparison the
stochastic bouncy particle sampler uses sub-sample batch sizes $n$ that are of
the order of, say 10% of, the full-data size $N$ (a smaller sub-sample batch
size would compromise its performance). Empirical results suggest our ability
to use a sub-sample of one can lead to a substantial increase in efficiency.
For example Figure 1 shows the trace of the first coordinate and the
autocorrelation function of two versions of the SBPS algorithm and of our SG-
BPS algorithm, for a Bayesian logistic regression problem (see Appendix B for
details). In Appendix B, we present a more detailed comparison of our SG-PDMPS
and the SBPS, for a logistic regression model and a linear regression model.
In all cases we found SBPS mixes more slowly for a fixed computational cost.
Also SBPS sometimes does not converge if started in the tail of the posterior.
Our ability to use, at every iteration, a mini-batch size of 1 data point is a
key advantage of our method and, in most contexts, we expect worse performance
(relative to CPU cost) for larger mini-batches (see Appendix D for numerical
simulations varying the batch-sizes of SG-PDMPs).

Figure 1: Top panels: traces of SBPS, SBPS with pre-conditioning (PSBPS)
(left) and our Euler approximation of the bouncy particle sampler, SG-BPS
(right) with step-size $h=10^{-4}$. Bottom panel: auto-correlation function of
the first coordinate for the three algorithms. All algorithms were implemented
with the same CPU cost, and output thinned to give 10,000 samples.

The paper is structured as follows. Section 2 reviews PDMPs with subsampling,
while in Section 3 we describe in detail SG-PDMPs and present the main
theoretical results which shows the order of error for this class of
algorithms. Section 4 highlights the main benefit of using SG-PDMPs on an
illustrative linear regression example. In Section 5, our algorithms are
applied to a logistic regression model, with artificial data, and for a
Bayesian neural network model with a variety of real datasets. Finally,
Section 6 outlines promising research directions stemming from this work.

##  2 PDMP samplers

Throughout we will consider the problem of sampling from a target density on
$\mathbb{R}^{d}$ of the form $\pi(x)\propto\exp(-U(x))$, and assume that

$U(x)=\sum_{j=1}^{N}U_{j}(x)$ (1)

for some large value $N$ and differentiable functions
$U_{j}:\mathbb{R}^{d}\to\mathbb{R},\,j=1,2,\ldots,N$. This is a common problem
in, for example, Bayesian inference where $\pi(x)$ is the posterior density
and the factors, $U_{j}(x)$ for $j=1,\ldots,N$, corresponds to either the log
prior or the log-likelihood of conditionally independent observations.

Like other MCMC algorithms, a PDMP sampler simulates a Markov process that is
designed to have $\pi(x)$ as its stationary distribution. However, whilst
standard MCMC algorithms simulate discrete-time, and generally reversible,
Markov process, a PDMP sampler simulates a continuous-time, non-reversible
process.

The dynamics of the PDMP are defined in the product space of position and
velocity $E=\mathbb{R}^{d}\times\mathcal{V}$ where
$\mathcal{V}\subset\mathbb{R}^{d}$. Evolution of the process is Markovian on
$E$ and given (in our cases) by piecewise constant velocity dynamics
interspersed by a sequence of random event times at which the velocity
component of the PDMP changes.

We will denote the state of the PDMP by $z=(x,v)\in E$, with $x$ denoting its
position and $v$, its velocity. Let $\phi_{s}(z)$ denote the change in state
due to the deterministic dynamics over a time-period of length $s$. That is
$z_{t+s}=\phi_{s}(z_{t})$, where

$\phi_{s}(z_{t})=\phi_{s}((x_{t},v_{t}))=(x_{t}+sv_{t},v_{t}).$

The dynamics of the PDMP are then specified through the rate at which events
occur and how the velocity changes at each event. As a PDMP is Markovian, the
instantaneous rate, $\lambda(z)$ of an event only depends on the current
state. The change of velocity at an event will be defined through a Markov
kernel that only acts on the velocity.

There are many choices of event rate and Markov kernel that will have $\pi(x)$
as the $x$-marginal of the stationary distribution of the resulting PDMP. We
are interested in PDMP samplers that use sub-sampling ideas, for which the
rate and transition kernel can each be written as a sum of terms, each of
which depends on just a single factor $U_{j}(x)$. We call these S-PDMPs.

###  2.1 PDMP samplers with subsampling

The idea of the S-PDMP is that we define the dynamics in terms of dynamics
associated with each factor $U_{j}(x)$. For each factor, $j=1,\ldots,N$,
introduce a reflection $F^{j}_{x}(\cdot)$, such that if
$v^{\prime}=F^{j}_{x}(v)$ then $v=F^{j}_{x}(v^{\prime})$. The rate of events
in the S-PDMP is then of the form

$\lambda(z)=\frac{1}{N}\sum_{j=1}^{N}\lambda^{j}(z)$

where $\lambda^{j}\colon E\to\mathbb{R}^{+}$ satisfies

$\lambda^{j}(x,F_{x}^{j}(v))-\lambda^{j}(x,v)=N\left(v\cdot\nabla
U_{j}(x)\right)$

and the Markov kernel for the transition at an event allows for $N$ possible
transitions. If the state immediately before the event is $z=(x,v)$, then the
state after the event is $z^{\prime}=(x,v^{\prime})$ with

$\Pr(v^{\prime}=F^{j}_{x}(v))=\frac{\lambda^{j}(z)}{N\lambda(z)},\mbox{ for
$j=1,\ldots,N.$}$

It can be shown that for these choices, the S-PDMP targets $\pi(x)$ Bierkens
et al. (2019); Fearnhead et al. (2018); Chevallier et al. (2024).

The $N\nabla U_{j}(x)$ term that appears in the rate $\lambda^{j}(z)$ can be
viewed as an unbiased estimator of $\nabla U(x)$. Just as control variates can
be used to reduce the variance of such an estimator, we can transform each
factor $U_{j}(x)$ to $\hat{U}_{j}(x)$ so that $N\nabla\hat{U}_{j}(x)$ is a
better approximation of $\nabla U(x)$. This is done by defining
$\hat{U}_{j}(x)=U_{j}(x)-x(U_{j}(\hat{x})-\sum_{j=1}^{N}U_{j}(\hat{x}))$, for
some centering value $\hat{x}$. This gives

$\nabla\hat{U}_{j}(x)=N\left\\{\nabla U_{j}(x)-\nabla
U_{j}(\hat{x})\right\\}+\sum_{j=1}^{N}\nabla U_{j}(\hat{x}).$ (2)

We will use such a set of transformed factors in the following, with $\hat{x}$
assumed to be an estimate of the mode of $\pi$, obtained through an initial
run of stochastic gradient descent. However, the following arguments will
apply to other choices. As we will show, and as is consistent with SGLD Baker
et al. (2018), choices of factors that give lower variance estimators of
$\nabla U(x)$ will lead to better algorithms.

####  2.1.1 Bouncy Particle Sampler with subsampling

The original BPS (Bouchard-Côté et al., 2018) takes values in
$\mathbb{R}^{d}\times\mathbb{R}^{d}$ and targets a density proportional to
$\pi(x)p(v)$, where $p(v)$ is a density for the velocity that is independent
of position and is symmetric. There are two common choices for $p(v)$, one is
a uniform density over a $d$-dimensional hypersphere, and the other is a
standard $d$-dimensional Gaussian distribution. More generally, we could
consider any distribution defined in terms of an arbitrary distribution for
the speed $||v||$, together with an independent uniform distribution for the
direction of $v$, i.e. $v/||v||$. The dynamics of the BPS are unaffected by
this choice, other than the initialisation, which should involve drawing
$v_{0}$ from $p(v)$, and at the refresh events (see below).

When using subsampling, the reflection event associated with factor $j$ is

$F^{j}_{x}(v)=v-\frac{v\cdot\nabla\hat{U}_{j}(x)}{\|\nabla\hat{U}_{j}(x)\|^{2}}\nabla\hat{U}_{j}(x),$

and the rates are

$\lambda^{j}(z)=\max(v\cdot\nabla\hat{U}_{j}(x),0).$ (3)

Furthermore, to ensure the process is irreducible, with constant rate
$\lambda_{\mathrm{ref}}>0$, the velocity component is refreshed with an
independent draw $v\sim p(v)$.

####  2.1.2 Zig-Zag sampler with subsampling

The Zig-Zag sampler with subsampling (Bierkens et al., 2019) has velocity
$v\in\\{-1,1\\}^{d}$. It involves $d$ possible events, each of which flips one
component of the velocity.

As there are $d$ possible events, an additional superscript, $i$, is
introduced for each type of flip. For $i=1,\ldots,d$ we define

$\mathrm{R}_{i}(v):=(v_{1},v_{2},\dots,v_{i-1},-v_{i},v_{i+1},\dots,v_{d}),$

i.e. only the $i$th component is flipped. This transition will be the same for
all factors, $j$, that is $F_{x}^{i,j}(v)=\mathrm{R}_{i}(v)$. The rate
associated with this transition is

$\lambda^{i,j}(z)=\max(v\partial_{x_{i}}\hat{U}_{j}(x),0),$ (4)

where, by analogy to eq. (2),

$\partial_{x_{i}}\hat{U}_{j}(x)=N\left\\{\partial_{x_{i}}U_{j}(x)-\partial_{x_{i}}U_{j}(\hat{x})\right\\}+\sum_{j=1}^{N}\partial_{x_{i}}U_{j}(\hat{x}).$

###  2.2 Simulating S-PDMPs

S-PDMPs are continuous-time Markov processes that have $\pi(x)$ as the
$x$-marginal of their stationary distribution. For reasons of practicality, we
require a method for simulating these processes. The challenge in simulating
an S-PDMP lies in simulating the event-times, as the other dynamics are
simple.

Given current state $z=(x,v)$, the rate of the next event just depends on the
time until the next event, which we define as
$\lambda_{z}(s):=\lambda((x+vs,v))$, which uses the fact that, if $z_{t}=z$
and there has been no further event, then at time $t+s$ the state is
$z_{t+s}=(x+vs,v)$. Thus the time until the next event will be the time of the
first event in an in-homogeneous Poisson process (IPP) of rate
$\lambda_{z}(s)$. We will denote the time until the next event as
$\tau\sim\text{IPP}(s\to\lambda_{z}(s))$. If this rate is constant, $\tau$ is
distributed as an exponential random variable with rate $\lambda(z)$, i.e.
$\tau\sim\text{Exp}(\lambda(z))$.

There exists a range of techniques for simulating from an in-homogeneous
Poisson process. The most common general methods are based on the idea of
thinning Lewis and Shedler (1979): we upper bound the rate
$\lambda(\phi_{s}(z))$ by a piecewise linear function of $s$; we simulate
possible events with this upper bound rate; and we accept these as actual
events with probability equal to the true rate divided by the bound. The
challenge with efficiently sampling an S-PDMP then comes from finding good
upper bounds for the rates.

Bierkens et al. (2019) show that if we obtain an upper bound,
$\lambda^{+}_{z}(s)$, that bounds the rate corresponding to each factor of the
target distribution, which in other words bounds $\lambda^{j}(\phi_{s}(z))$
for $j=1\ldots,N$, then we can simulate the S-PDMP exactly whilst accessing
only a single factor at each iteration. This works by (i) simulating the next
possible event, at a rate $\lambda^{+}_{z}(s)$; (ii) sample a factor uniformly
at random; (iii) if the possible event is at time $\tau$, and we have sampled
the $j$th factor, we accept the event with probability
$\lambda^{j}(\phi_{\tau}(z))/\lambda^{+}_{z}(\tau)$, and change the state
according to the kernel $Q_{j}(\phi_{\tau}(z),\cdot)$. Importantly, given the
upper bound, the only step that depends on the target is step (iii) and this
involves only a single factor.

In settings where S-PDMPs can be applied, they are shown to have an overall
complexity of $\mathcal{O}(1)$, relative to the sample size $N$. This is
theoretically motivated and numerically illustrated for logistic regression
models in Bierkens et al. (2019) and Bierkens et al. (2020). However, the set
of models where we can find appropriate upper bounds is currently limited and
motivates our interest in approximate stochastic gradient PDMP samplers.

##  3 Stochastic gradient PDMP samplers

By analogy with stochastic gradient Langevin algorithms, where one
approximately simulates a Langevin diffusion with $\pi(x)$ as its stationary
distribution, we introduce stochastic gradient PDMP algorithms that
approximately simulates an S-PDMP algorithm, and called SG-PDMP algorithms.
The idea of simulating approximations to PDMPs is introduced in Bertazzi et
al. (2022, 2023).

Following Bertazzi et al. (2022), our approximation is based on discretising
time into intervals of size $\varepsilon$, and for each interval choosing a
factor at random and then simulating the dynamics of the underlying S-PDMP
based on the rate and dynamics for that factor. As a starting point, and to
further simplify the simulation of any event in the time interval, we fix the
event rate based on the state at the start of the time interval, and simulate
at most one event. As we show below, using these two simplifications does not
impact the order of the approximation.

The resulting general algorithm is given in Algorithm 1. We assume our true
PDMP has $K$ types of events, and introduce a rate and transition for each
type and each factor. For example, with the Zig-Zag sampler, we have $d$ event
types, one for each possible component of the velocity to flip. See eq. (4)
for the rates associated with each pair of event type and factor, with event
$i$ having the same transition regardless of the factor, and with
$v^{\prime}=R_{i}(v)$ for the new velocity. For the Bouncy Particle Sampler,
we have two types of events. The first is a reflection, with rate for factor
$j$ given by (3), and transition given by $v^{\prime}=F_{x}^{j}(v)$. The
second type of event is a refresh event. This is the same for all factors and
has constant rate $\lambda_{\mathrm{ref}}$ with transition $v$ drawn from the
stationary distribution for the velocity.

The algorithm then loops over time intervals of length $\varepsilon$,
simulates a factor $J$, and then an event time for each type of event. These
occur with a constant rate defined as the rate for each event $i$, and factor
$J$, for the current state of the process. We then calculate the time, $\tau$,
and event type $i^{*}$ which occurs first. If that event occurs within the
interval, we simulate the exact continuous-time dynamics of the PDMP over the
time interval with an event of type $i^{*}$ for factor $J$ at time $\tau$. If
the first event does not occur within the time interval, we just simulate the
PDMP dynamics over the time interval with no events.

Algorithm 1 Stochastic gradient PDMP sampler

0: $(x,v)\in E$, step size $\varepsilon>0$, time horizon $T>0$.

$t=0$,

Rates $\lambda^{i,j}(z)$ for event of type $i$ associated with potential
$U_{j}$, for $i=1,\ldots,K$ and $j=1,\ldots,N$.

Transition kernel $Q^{i,j}(\cdot;z)$ associated with event of type $i$ and
potential $U_{j}$.

while $t<T$ do

$J\sim\mathrm{Unif}(1,2,\dots,N)$

$\tau_{i}\sim\mathrm{Exp}(\lambda^{i,J}(z)),\mbox{ for }i=1,\ldots,K$

$\tau=\min_{i=1,\ldots,K}\,\tau_{i},\,i^{*}=\operatorname*{argmin}\tau_{i}$

if $\tau<\varepsilon$ then

$x\leftarrow x+v\tau$

$v\sim Q^{i^{*},J}(\cdot;(x,v))$ {Simulate new velocity}

$x\leftarrow z+v(\varepsilon-\tau)$

else

$x\leftarrow x+v\varepsilon$

end if

$t\leftarrow t+\varepsilon$

Save $(x,t)$

end while

Using results in Bertazzi et al. (2022) we can show that the resulting SG-PDMP
sampler can give an $O(\varepsilon)$ approximation to the distribution of the
true PDMP over any fixed time interval $t$:

######  Proposition 3.1.

Let $\overline{\mathcal{P}}_{t}(z,\cdot)$ and ${\mathcal{P}}_{t}(z,\cdot)$ be
the transition kernels for the stochastic gradient PDMP process of Algorithm 1
and for the true underlying PDMP process, respectively. Assume the PDMP
processes have bounded velocities, so $\|v\|<C_{0}$ for some $C_{0}$. Assume
that for any state $z=(x,v)$ and any $i=1,\ldots,K$ and $j=1,\ldots,N$ the
function $\lambda^{i,j}(x+vt,v)$ has a continuous derivative with respect to
$t$. Then there exists constants $C(z,T)$, that depend on the initial state,
$z$ and time interval $T$, and $\varepsilon_{0}>0$ such that for all
$\varepsilon<\varepsilon_{0}$

$||{\mathcal{P}}_{T}(z,\cdot)-\overline{\mathcal{P}}_{T}(z,\cdot)||_{TV}\leq
C(z,T)\varepsilon.$

See Appendix A for the proof.

While this result shows that the SG-PDMP algorithm is an $O(\varepsilon)$
approximation of the true S-PDMP, a more informal analysis gives insight into
the approximation error. Consider the probability of simulating an event in
the next interval of length $\varepsilon$ given a current state $z=(x,v)$. Let
$\lambda^{j}(z)=\sum_{i=1}^{K}\lambda^{i,j}(z)$, then the probability for the
true S-PDMP of no event in an interval of length $\varepsilon$ is
$\exp\\{-\sum_{j=1}^{N}(1/N)\int_{0}^{\varepsilon}\lambda^{j}(x+vt,v)\mbox{d}t\\}$,
whereas for the SG-PDMP it is
$(1/N)\sum_{j=1}^{N}\exp\\{-\lambda^{j}(z)\varepsilon\\}$. We can decompose
the difference into two terms

$\displaystyle\exp\left\\{-\sum_{j=1}^{N}\int_{0}^{\varepsilon}\frac{\lambda^{j}(x+vt,v)}{N}\mbox{d}t\right\\}-\exp\left\\{-\sum_{j=1}^{N}\frac{\lambda^{j}(z)\varepsilon}{N}\right\\}$
$\displaystyle+$
$\displaystyle\exp\left\\{-\sum_{j=1}^{N}\frac{\lambda^{j}(z)\varepsilon}{N}\right\\}-\frac{1}{N}\sum_{j=1}^{N}\exp\left\\{-\lambda^{j}(z)\varepsilon\right\\}.$

The first difference is due to the use of constant event rates over the
interval. Assuming the rates are Lipschitz continuous, the change in event
rate over an interval of length $\varepsilon$ is $O(\varepsilon)$, so the
difference in the integral of the rates is $O(\varepsilon^{2})$. We could
reduce this error by using a better approximation to the event rate over the
interval Bertazzi et al. (2023). The second difference is due to the
stochastic gradient approximation, that is we consider just one factor for
each time interval. There are two points to make. The first is that by
Jensen’s inequality, this difference is always negative. This means that the
stochastic gradient approximation reduces the probability of an event. As the
PDMPs only introduce events when moving into areas of lower probability
density, the impact of this is that it samples from a heavier-tailed
approximation to the target. Second, we can use a Taylor expansion to get the
highest order, in $\varepsilon$, term of the approximation. This is the
$O(\varepsilon^{2})$ term, as the first order terms cancel, whose coefficient
is

$-\left(\frac{1}{N}\sum_{j=1}^{N}\lambda^{j}(z)^{2}-\left(\sum_{j=1}^{N}\frac{\lambda^{j}(z)}{N}\right)^{2}\right),$

which is (minus) the variance of the $\lambda^{j}(z)$s. This means that the
approximation error depends on the variability of the estimator of the
gradient of $U(x)$ across the different factors. Furthermore, this error is
$O(\varepsilon^{2})$ which means that using a better approximation for how the
rates vary over the time interval would not improve the rate in $\varepsilon$
of the approximation we introduce. Proposition 3.1 provides an order of
approximation in terms of the transition kernels of our algorithms. For
results on the approximation of the invariant measure, we refer to Section 6
in Huggins and Zou (2017): under appropriate mixing conditions, the order of
the error (in terms of step-size) of the approximation in the transition
density will imply a similar order of error in the invariant measure.

One disadvantage of Algorithm 1 for coarse discretisations, i.e. when
$\varepsilon$ is large, is that it does not allow more than one event in the
interval. Thus a simple improvement is, if there is an event at time
$\tau<\varepsilon$ within the interval, to then iterate the simulation of
events for the time interval $[\tau,\varepsilon]$. See Appendix C for details
of the resulting algorithm for stochastic gradient versions of the Zig-Zag
sampler and the Bouncy Particle Sampler. These are the algorithms we evaluate
in the empirical results in Sections 4 and 5.

###  3.1 Extension with sticky components for variable selection

The methods developed here can be naturally extended to the PDMPs with
boundary conditions recently developed in Chevallier et al. (2023); Bierkens
et al. (2023a); Chevallier et al. (2024); Bierkens et al. (2023b) for sampling
from target densities which are only piecewise smooth and for reference
measures which are mixture of continuous and atomic components. The only
difference these methods have is with the dynamics between events. For
example, if there is a boundary then the PDMP may reflect off the boundary if
it hits it.

For sampling from Bayesian posteriors for models with variable selection,
where the target distribution is for the co-efficients of the model, it is
common to have a prior that has positive probability on co-efficients being
zero. The sticky PDMP of Bierkens et al. (2023a) can deal with this by setting
the velocity of components of $x$ to zero if that component of $x$ is zero,
and then re-introducing the non-zero velocity at a certain rate. This is also
easily incorporated into the SG-PDMP algorithm by appropriately changing the
dynamics of the process. See Appendix E for further details and the
corresponding SG-PDMP algorithms.

Figure 2:  Left top panel: error in standard deviation estimation of each
coordinate (5) $\mathcal{E}^{(h,d)}$ as a function of the stepsize $h$
($x$-axis in log-scale). Left bottom panel: $\mathcal{E}^{(h,d)}$ as a
function of $d$. Right panels: trace plots of the first coordinate.

##  4 Illustrative example: Linear regression model

We illustrate the benefits of SG-PDMP samplers on a simple linear regression
model $y=Ax+\varepsilon$, for a response variable $y\in\mathbb{R}^{N}$,
covariates $A\in\mathbb{R}^{N\times d}$, parameters $x\in\mathbb{R}^{d}$ and
$\varepsilon\sim\mathcal{N}(0,Nc\mathrm{I}_{N})$ (the variance of the noise is
re-scaled by a factor to aid comparison across different values of $N$, as the
posteriors will have similar variance regardless of $N$). For this problem, we
compare sample averages computed with the trajectories of each sampler against
the true expectation of selected functionals of the true posterior.

We set $A_{:,1}=\mathbf{1}$ to account for the intercept and simulate the data
as $A_{i,j}\sim\mathcal{N}(0,\Sigma)$ with
$\Sigma_{i,j}=\Sigma_{j,i}\sim\text{Unif}(0.4,0.8)^{|i-j|},i=1,2,\dots,d;\,j=2,3,\dots,d$
and set $x_{1}=0$. We simulate
$x_{i}\overset{\mathrm{iid}}{\sim}\mathcal{N}(0,1),i=2,3,\dots,d$ with prior
$\mathcal{N}(0,100\times\mathrm{I}_{d})$ for $x$.

We consider two experimental settings. First, we simulate $N=10^{6}$
covariates with $d=5$ features and run SGLD, SG-ZZ and SG-BPS for different
step-sizes and $T=5\times 10^{6}$ iterations, initialising each sampler at the
ordinary least squares estimate. Figure 2 (left top panel) shows

$h\mapsto\mathcal{E}^{(d,h)}=\frac{1}{d}\sum_{i=1}^{d}\left(\frac{\hat{\sigma}_{i}^{(h)}-\sigma_{i}}{\sigma_{i}}\right)^{2}$
(5)

for each sampler, where $\hat{\sigma}^{(h)}_{i}$ is the sample standard
deviation of the $i$th component estimated with the trace. The $x$-axis is
displayed on a log-scale. We also consider the setting where we initialise the
algorithms in the tails of the distribution and plot in Figure 2 (left panels)
the first few iteration of the trace of the first coordinate (whose true value
was set to 0) and including the SG-SZZ sampler.

Second, we set $N=10^{5}$ and simulate multiple datasets with varying
$d=10,\dots,10^{2}$. For each dataset, we run SGLD, SG-ZZ and SG-BPS. We fix
the step size equal to $h=5\times 10^{-7}$ and the number of iterations
$T=2\times 10^{6}$, such that all samplers have a comparable error for the
dataset with the smallest number of dimensions. Figure 2 (bottom left panel)
shows $d\mapsto\mathcal{E}^{(d,h)}$ for each sampler. In all simulations, each
sampler uses a stochastic gradient with control variates and 1 data point at
each iteration (2). The number of gradient evaluations and the running time is
comparable among the different algorithms.

This simple tractable example provides useful insight into the advantages of
SG-PDMP over SGLD. Firstly we see greater stability of SG-PDMPS for large
discretisation steps. This is in contrast with SGLD which diverges to infinity
if the discretisation step is too large. (In these simulations, the maximum
step size that can be taken for SGLD before diverging is 0.004.) Secondly, the
error in SG-PDMPs grows much slower compared to SGLD as the dimensionality of
the parameter increases: in these simulations, SGLD was unstable for $d>50$.
This offers opportunities for SG-PDMPs especially for high dimensional
problems. Finally, the deterministic dynamics of SG-PDMPs allows these
algorithms to deal easily with the variable dimensional posterior in the case
where we use a prior on each component on $x$ that includes a point-mass at
zero. This is achieved using the stochastic gradient version of the sticky
Zig-Zag sampler of Section 3.1. The trace of the sampler is shown in the
bottom-right panel, and we see it successfully enforces sparsity of the first
parameter, whose true value is equal to 0.

##  5 Numerical experiments

In this section, we assess the performance of SG-PDMPs and SGLD on intractable
posterior distributions. We consider logistic regression and a Bayesian neural
network model.

We analyse the performance of each sampler by computing various metrics to
assess bias and numerical accuracy. In the logistic regression example, we
simulate a large dataset so that the posterior distribution is well-
approximated by a Gaussian distribution formed from a Laplace approximation.
Then, for each sampler, we compute the sum of squared errors as in eq. (5)
between the estimated standard deviations and the standard deviations of the
Laplace approximation.

The samplers considered here are approximate and asymptotically biased, which
means that standard MCMC diagnostics such as Effective Sample Size (ESS) are
not appropriate (Nemeth and Fearnhead, 2021). A popular performance metric for
stochastic gradient Monte Carlo algorithms is the Kernel Stein Discrepancy
(Gorham and Mackey, 2017); see Appendix F.1. This metric can detect both poor
mixing and bias from the target distribution.

For all the problems considered, predictive performance of each sampler was
assessed by splitting the dataset $\mathcal{D}$ into a training set
$\mathcal{T}$ and its complementary test set $\mathcal{T}^{c}$ so that
$\mathcal{T}\sqcup\mathcal{T}^{c}=\mathcal{D}$. We set
$|\mathcal{T}|=0.9|\mathcal{D}|$. Each datum $\mathcal{D}_{i}=(X_{i},y_{i})$
consists of a pair of covariates together with a dependent variable. Then, for
each point $x$ of the output of each algorithm, its predictive accuracy is
given by

$\frac{1}{|\mathcal{T}^{c}|}\sum_{(X_{i},y_{i})\in\mathcal{T}^{c}}\ell(X_{i},y_{i},x)$
(6)

where $\ell$ is a non-negative loss function between $y_{i}$ and the predicted
outcome given the covariates $X_{i}$ and parameter $x$. Specific loss
functions will be specified in context.

Each algorithm was implemented $M$ times, each time with a different random
permutation of training and test datasets. Average loss was computed across
these $M$ realisations. In all examples, control-variate stochastic gradient
estimates were employed (Equation (2)) and samplers were initialised at the
same control variate point, with no burn-in period. The control variate was
computed using the stochastic optimization algorithm ADAM (Kingma and Ba,
2014), with $10^{6}$ iterations and a subsample size equal to 1% of the
dataset and all the other parameters suggested therein.

###  5.1 Logistic regression

A Bayesian logistic regression problem was considered with standard Gaussian
prior with variance $10$ for each component. We set the dimension of the
parameter $p=10$ and we simulate $N=10^{5}$ covariates
$X_{i,j}\sim\mathcal{N}(0,\Sigma)$ where
$\Sigma_{i,j}=\Sigma_{j,i}\sim\text{Unif}(-\rho,\rho)^{|i-j|},\,\rho=0.4$ for
$i=1,2,\dots,N$ and $j=1,2,\dots,p$. $y$ was simulated with true parameters:
$x^{\star}_{i}\sim\mathcal{N}(0,1)$.

SGLD was implemented with minibatches of sizes 1, 10 and 100 data points and
compared against SG-PDMP samplers SG-ZZ, SG-BPS and SG-SZZ. All algorithms
were run for $T=10^{6}$ iterations. The running time, and number of gradient
evaluations, for SGLD with minibatches of sizes 10 and 100 were significantly
higher compared to the SG-PDMP algorithms. Each simulation was implemented for
several values of $h=10^{-6},\dots,10^{-3}$. Figure 3 shows
$h\to\mathcal{E}^{(h,d)}$ \- eq. (5) between the estimated standard deviation
and the standard deviation of the Laplace approximation. In Appendix F, we
give the Kernel Stein discrepancy and the loss function (6). For all
experiments, SG-PDMPs performs similarly to SGLD for small step-sizes and out-
perform SGLD with 1, 10, 100 mini-batch sizes for larger step-sizes (the
maximum step-size for SGLD before the algorithms diverge is $10^{-4}$).

Figure 3: Error between standard deviation estimation of each coordinate and
the one relative to the Laplace approximation as in (5) for the logistic
regression as a function of the stepsize $h$ ($x$-axis in log-scale). Dashed
lines corresponds to SGLD algorithms, solid lines to SG-PDMPs.

A second experiment considering an over-parameterised regime was also
considered. In this setting we took $p=10^{2}$ parameters and $N=10^{2}$
covariates, with true parameter being 0 with probability 0.5. SGLD, SG-ZZ, SG-
BPS, SG-SZZ were all implemented for $T=10^{7}$ and $h=10^{-4}$. SG-SZZ
utilised a spike-and-slab prior with spike weight equal to $w=0.5$. Table 1
shows the mean squared error between true parameters and sample mean and
median.

| SGLD | SG-ZZ | SG-BPS | SG-SZZ
---|---|---|---|---
mean | 3.08155 | 3.15546 | 2.93989 | 1.42291
median | 3.02602 | 3.12318 | 2.90635 | 1.18081
Table 1: Mean squared error relative to the sample mean and sample median for
over-parameterised logistic regression.

###  5.2 Bayesian neural networks

Bayesian approaches to neural networks (BNNs) provides a powerful calculus to
quantify uncertainty and reduce the risk of overfitting by incorporating
techniques such as dropout (Gal and Ghahramani, 2016) and imposing sparsity-
inducing priors (Polson and Ročková, 2018).

Three datasets were considered from the _UCI machine learning repository_
111https://archive.ics.uci.edu/, varying in dimension and data size labelled
as boston, concrete, kin8mn. A two-layer Bayesian neural network was utilised:

$y=a_{2}(W_{2}(a_{1}(W_{1}x+b_{1})+b_{2})+N(0,1)$

where $a_{1}(x)=x$ and $a_{2}(x)=\max(0,x)$, for unknown parameters
$W_{1}\in\mathbb{R}^{50\times p},W_{2}\in\mathbb{R}^{1\times
50},b_{1}\in\mathbb{R}^{50}$ and $b_{2}\in\mathbb{R}$ for $p$ covariates in
each dataset. An independent prior $\mathcal{N}(0,10)$ was chosen for each
component of $x=(W_{1},W_{2},b_{1},b_{2})$. With this setting, the number of
dimensions of the parameter $x$ ranged from 501 to 751.

SG-ZZ, SG-SZZ, SG-BPS and SGLD were implemented for $N=10^{6}$ iterations on
$M=3$ random permutations of the training and test datasets, varying
$h=10^{-6},\dots,10^{-3}$. The loss function was chosen to be the mean squared
error. Figure 4 shows the loss function trace plot for each sampler in the
first permutation of training and test set, for the dataset ‘boston’ which has
$N=507$ data points and $p=13$ covariates. For this experiment, SGLD was
unstable for $h=10^{-3}$. In Table 2 we display the average loss of the trace
which was computed by averaging the $M=3$ random permutations of training and
test datasets. The results for the other three datasets considered here are
qualitatively similar and can be found in Appendix G.

Figure 4: Trace of the loss function for each sampler for different step-sizes $h$ relative to the first permutation of the training and test set of the dataset ‘boston’. | h
---|---
Sampler | $10^{-3}$ | $10^{-4}$ | $10^{-5}$ | $10^{-6}$
SGLD | _NaN_ | _0.4940_ | _0.4586_ | 0.4526
SG-ZZ | 0.4720 | 0.4572 | 0.4518 | 0.4555
SG-BPS | 0.4692 | 0.4533 | 0.4543 | 0.4574
SG-SZZ | 0.4704 | 0.4586 | 0.4564 | _0.4643_
Table 2: Average loss for each sampler for different values of step size $h$
relative to the dataset ‘boston’. In bold the best performance and italic the
worst performance given a step size $h$.

##  6 Discussion

We have presented SG-PDMPs as a competitive and robust alternative to the
popular SGLD algorithm for approximate Bayesian inference. In particular, we
highlighted that SG-PDMPs are generally more stable than SGLD for large step
sizes, and can take advantage of the continuous-time sample paths through e.g.
the sticky dynamics for regression with model choice. The better robustness
that SG-PDMPs have is due to the fixed velocity dynamics, and is reminiscent
of that of relativistic Hamiltonian dynamics Lu et al. (2017), and the
stochastic gradient Barker dynamics of Mauri and Zanella (2024).

There are several natural extensions to this paper which could be explored as
future work. It is known that suitable pre-conditioning can improve mixing of
PDMPs. An adaptive preconditioning mass matrix for PDMPs has been studied in
Bertazzi and Bierkens (2022) and Pakman et al. (2017). As noted in Livingstone
and Zanella (2022), the convergence of adaptive algorithms can be drastically
improved if algorithms are robust to the choice of step size, as SG-PDMPs
appear to be. Second, SGLD is the Monte Carlo analogue of the popular
stochastic optimisation algorithm _Stochastic Gradient Decent_ (SGD).
Similarly, Piecewise Deterministic Markov processes which converge to global
minima appeared for example in Monmarché (2014). It is therefore natural to
develop a stochastic optimization method based on SG-PDMP dynamics. Finally,
it would be interesting to develop higher-order approximation schemes. This
will be non-trivial, as, in Section 3, we showed that higher order
approximations of the Poisson rates, as developed in Bertazzi et al. (2023),
cannot be directly adopted in this context since the error given by taking the
stochastic gradient dominates the overall error of the algorithm.

## References

  * Baker et al. (2018) Baker, J., Fearnhead, P., Fox, E. and Nemeth, C. (2018). Large-scale stochastic sampling from the probability simplex. Advances in Neural Information Processing Systems 31.
  * Bernard et al. (2009) Bernard, E. P., Krauth, W. and Wilson, D. B. (2009). Event-chain Monte Carlo algorithms for hard-sphere systems. Phys. Rev. E 80, 056704.
  * Bertazzi and Bierkens (2022) Bertazzi, A. and Bierkens, J. (2022). Adaptive schemes for piecewise deterministic Monte Carlo algorithms. Bernoulli 28(4), 2404–2430.
  * Bertazzi et al. (2022) Bertazzi, A., Bierkens, J. and Dobson, P. (2022). Approximations of piecewise deterministic Markov processes and their convergence properties. Stochastic Processes and their Applications 154, 91–153.
  * Bertazzi et al. (2023) Bertazzi, A., Dobson, P. and Monmarché, P. (2023). Piecewise deterministic sampling with splitting schemes. ArXiv:2301.02537.
  * Bierkens (2016) Bierkens, J. (2016). Non-reversible Metropolis-Hastings. Statistics and Computing 26(6), 1213–1228.
  * Bierkens et al. (2019) Bierkens, J., Fearnhead, P. and Roberts, G. (2019). The Zig-Zag process and super-efficient sampling for Bayesian analysis of big data. Ann. Statist. 47(3), 1288–1320.
  * Bierkens et al. (2020) Bierkens, J., Grazzi, S., Kamatani, K. and Roberts, G. (2020). The Boomerang Sampler. In: International Conference on Machine Learning, PMLR, 908–918.
  * Bierkens et al. (2023a) Bierkens, J., Grazzi, S., Meulen, F. v. d. and Schauer, M. (2023a). Sticky PDMP samplers for sparse and local inference problems. Statistics and Computing 33(1), 8.
  * Bierkens et al. (2023b) Bierkens, J., Grazzi, S., Roberts, G. and Schauer, M. (2023b). Methods and applications of PDMP samplers with boundary conditions. ArXiv:2303.08023.
  * Bouchard-Côté et al. (2018) Bouchard-Côté, A., Vollmer, S. J. and Doucet, A. (2018). The bouncy particle sampler: A nonreversible rejection-free Markov chain Monte Carlo method. Journal of the American Statistical Association 113(522), 855–867.
  * Chen et al. (2014) Chen, T., Fox, E. and Guestrin, C. (2014). Stochastic gradient Hamiltonian Monte Carlo. In: International Conference on Machine Learning, PMLR, 1683–1691.
  * Chevallier et al. (2023) Chevallier, A., Fearnhead, P. and Sutton, M. (2023). Reversible jump PDMP samplers for variable selection. Journal of the American Statistical Association 118, 2915–2927.
  * Chevallier et al. (2024) Chevallier, A., Power, S., Wang, A. Q. and Fearnhead, P. (2024). PDMP Monte Carlo methods for piecewise-smooth densities. Advances in Applied Probability, to appear .
  * Corbella et al. (2022) Corbella, A., Spencer, S. E. and Roberts, G. O. (2022). Automatic zig-zag sampling in practice. Statistics and Computing 32(6), 1–16.
  * Davis (1984) Davis, M. H. (1984). Piecewise-deterministic markov processes: A general class of non-diffusion stochastic models. Journal of the Royal Statistical Society: Series B (Methodological) 46(3), 353–376.
  * Diaconis et al. (2000) Diaconis, P., Holmes, S. and Neal, R. M. (2000). Analysis of a nonreversible Markov chain sampler. Annals of Applied Probability 10, 726–752.
  * Fearnhead et al. (2018) Fearnhead, P., Bierkens, J., Pollock, M. and Roberts, G. O. (2018). Piecewise Deterministic Markov Processes for Continuous-Time Monte Carlo. Statistical Science 33(3), 386–412.
  * Gal and Ghahramani (2016) Gal, Y. and Ghahramani, Z. (2016). Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In: International Conference on Machine Learning, PMLR, 1050–1059.
  * Gawlikowski et al. (2023) Gawlikowski, J., Tassi, C. R. N., Ali, M., Lee, J., Humt, M., Feng, J., Kruspe, A., Triebel, R., Jung, P., Roscher, R., Shahzad, M., Yang, W., Bamler, R. and Zhu, X. X. (2023). A survey of uncertainty in deep neural networks. Artificial Intelligence Review 56, S1513–S1589.
  * Gorham and Mackey (2017) Gorham, J. and Mackey, L. (2017). Measuring sample quality with kernels. In: International Conference on Machine Learning, PMLR, 1292–1301.
  * Huggins and Zou (2017) Huggins, J. and Zou, J. (2017). Quantifying the accuracy of approximate diffusions and markov chains. In: Artificial Intelligence and Statistics, PMLR, 382–391.
  * Kingma and Ba (2014) Kingma, D. P. and Ba, J. (2014). Adam: A method for stochastic optimization. ArXiv:1412.6980.
  * Krauth (2021) Krauth, W. (2021). Event-chain monte carlo: Foundations, applications, and prospects. Frontiers in Physics 9.
  * Lewis and Shedler (1979) Lewis, P. W. and Shedler, G. S. (1979). Simulation of nonhomogeneous Poisson processes by thinning. Naval Research Logistics Quarterly 26(3), 403–413.
  * Livingstone and Zanella (2022) Livingstone, S. and Zanella, G. (2022). The Barker proposal: combining robustness and efficiency in gradient-based MCMC. Journal of the Royal Statistical Society Series B: Statistical Methodology 84(2), 496–523.
  * Lu et al. (2017) Lu, X., Perrone, V., Hasenclever, L., Teh, Y. W. and Vollmer, S. (2017). Relativistic Monte Carlo . In: Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (eds. A. Singh and J. Zhu), volume 54 of Proceedings of Machine Learning Research, PMLR, 1236–1245.
  * Ma et al. (2015) Ma, Y.-A., Chen, T. and Fox, E. (2015). A complete recipe for stochastic gradient MCMC. Advances in Neural Information Processing Systems 28.
  * Mauri and Zanella (2024) Mauri, L. and Zanella, G. (2024). Robust approximate sampling via stochastic gradient barker dynamics. In: International Conference on Artificial Intelligence and Statistics, PMLR, 2107–2115.
  * Monmarché (2014) Monmarché, P. (2014). Piecewise deterministic simulated annealing. ArXiv:1410.1656.
  * Nemeth and Fearnhead (2021) Nemeth, C. and Fearnhead, P. (2021). Stochastic gradient Markov chain Monte Carlo. Journal of the American Statistical Association 116(533), 433–450.
  * Pagani et al. (2024) Pagani, F., Chevallier, A., Power, S., House, T. and Cotter, S. (2024). NuZZ: numerical Zig-Zag sampling for general models. Statistics and Computing 34, 61.
  * Pakman et al. (2017) Pakman, A., Gilboa, D., Carlson, D. and Paninski, L. (2017). Stochastic bouncy particle sampler. In: International Conference on Machine Learning, PMLR, 2741–2750.
  * Peters and de With (2012) Peters, E. A. J. F. and de With, G. (2012). Rejection-free monte carlo sampling for general potentials. Physical Review E 85(2).
  * Polson and Ročková (2018) Polson, N. G. and Ročková, V. (2018). Posterior concentration for sparse deep learning. Advances in Neural Information Processing Systems 31.
  * Şimşekli et al. (2017) Şimşekli, U., Durmus, A., Badeau, R., Richard, G., Moulines, E. and Cemgil, A. T. (2017). Parallelized stochastic gradient Markov chain Monte Carlo algorithms for non-negative matrix factorization. In: 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), IEEE, 2242–2246.
  * Vanetti et al. (2017) Vanetti, P., Bouchard-Côté, A., Deligiannidis, G. and Doucet, A. (2017). Piecewise-Deterministic Markov Chain Monte Carlo. ArXiv:1707.05296.
  * Welling and Teh (2011) Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient Langevin dynamics. In: Proceedings of the 28th International Conference on Machine Learning (ICML-11), 681–688.

##  Appendix A Proof of Proposition 3.1

The proof follows by application of Theorem 4.17 of Bertazzi et al. (2022).
Within Bertazzi et al. (2022), the authors discuss approximations of the Zig-
Zag sampler with sub-sampling – see their Example 5.7. However they comment
that such an approximation is different from the algorithms they consider and
just sketch how their proofs could be extended to such an algorithm. Here we
consider a different approach, by showing that we can re-formulate our SG-PDMP
algorithm as an specific case of their Algorithm 3, and then directly apply a
result for that algorithm.

Let the current state be $z$. The distribution of the time to the next,
$\tau$, that is simulated in Algorithm 1 satisfies

$\Pr(\tau>t)=\sum_{j=1}^{N}\frac{1}{N}\exp\left\\{-t\sum_{i=1}^{K}\lambda^{i,j}(z)\right\\}.$

This follows by averaging $\Pr(\tau>t|J=j)$ over the possible values of $j$.
Given an event at time $\tau$, type of event $i^{*}$ and factor $J$ has
probability mass function

$\Pr(i^{*}=i,J=j|\tau)=\frac{\lambda^{i,j}(z)\exp\left\\{-t\sum_{k=1}^{K}\lambda^{k,j}(z)\right\\}}{\sum_{l=1}^{N}\left(\sum_{k=1}^{K}\lambda^{k,l}(z)\right)\exp\left\\{-t\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}}.$
(7)

This follows as, by Bayes theorem, this is proportional to the density of
choosing factor $j$, having an event at time $\tau$ and the event being of
type $i$ – which is the term in the numerator. The denominator is then just
the normalising constant of the probability mass function.

In Algorithm 3 of Bertazzi et al. (2022) they simulate the time $\tau$ such
that

$\Pr(\tau>t)=\exp\left\\{-\int_{0}^{t}\bar{\lambda}(z,s)\mbox{d}s\right\\},$
(8)

for some time-inhomogeneous rate $\bar{\lambda}(z,s)$, that is the sum of $m$
event specific rates
$\bar{\lambda}(z,s)=\sum_{k=1}^{m}\bar{\lambda}_{k}(z,s)$. The the type of
event, $k$ say, is simulated with probability
$\bar{\lambda}_{k}(z,\tau)/\bar{\lambda}(z,\tau)$.

To relate the two algorithms, we first let $m=KN$. We then slightly adapt the
notation of Bertazzi et al. (2022) and subscript the event rates by the pair
$(i,j)$, rather than a single index. Define the event specific time-
inhomogeneous rates by, for $i=1,\ldots,K$ and $j=1,\ldots,N$,

$\bar{\lambda}_{i,j}(z,s)=\lambda_{i,j}(z)\frac{\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,j}(z)\right\\}}{\sum_{l=1}^{N}\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}}.$
(9)

Then

$\bar{\lambda}(z,s)=\sum_{j=1}^{N}\sum_{i=1}^{K}\lambda^{i,j}(z)\frac{\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,j}(z)\right\\}}{\sum_{l=1}^{N}\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}},$

and by noting that

$\frac{\mbox{d}}{\mbox{d}s}\log\left\\{\sum_{l=1}^{N}\frac{1}{N}\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}\right\\}=-\bar{\lambda}(z,s),$
$\int_{0}^{t}\bar{\lambda}(z,s)\mbox{d}s=\left[-\log\left\\{\frac{1}{N}\sum_{l=1}^{N}\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}\right\\}\right]_{0}^{t}=-\log(1)+\log\left\\{\frac{1}{N}\sum_{l=1}^{N}\exp\left\\{-t\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}\right\\}.$

Substituting into eq. (8), we see that the distribution of $\tau$ in Algorithm
3.1 of Bertazzi et al. (2022) is the same of the distribution of $\tau$
simulated by Algorithm 1. Conditional on $\tau$, the probability of choosing
event $i,j$ in Algorithm of Bertazzi et al. (2022) is

$\frac{\bar{\lambda}_{i,j}(z,\tau)}{\bar{\lambda}(z,\tau)}=\frac{\lambda^{i,j}(z)\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,j}(z)\right\\}}{\sum_{l=1}^{N}\left(\sum_{k=1}^{K}\lambda^{k,l}(z)\right)\exp\left\\{-s\sum_{k=1}^{K}\lambda^{k,l}(z)\right\\}},$

which is the same as the probability for Algorithm 1 – see (7).

Thus Algorithm 1 is equivalent to Algorithm 3 of Bertazzi et al. (2022), with
the specific choice of rates, $\bar{\lambda}_{i,j}(z,s)$ given in (9).
Proposition 3.1 follows immediately from Theorem 4.17 of Bertazzi et al.
(2022) if we can show their Assumption 4.14 holds. As they comment (see their
Note 4.16), this will hold if the state of the true and approximate PDMP has
bounded norm for a finite time horizon, and if their Assumption 4.6 holds. The
former is true as we are considering PDMPs with bounded velocity.

Their Assumption 4.6 requires (A) an $\bar{M}(z)$ such that for $0\leq
s\leq\varepsilon\leq\varepsilon_{0}$, and all $i,j$

$|\bar{\lambda}_{i,j}(z,s)-\lambda_{i,j}((x+sv,v))|\leq\varepsilon\bar{M}(z),$

where $z=(x,v)$; and (B) that for any future time time $n\varepsilon<T$ for
positive integer $n$, that

$\mbox{E}_{z}[\bar{M}(\bar{Z}_{n\varepsilon})]\leq M(n\varepsilon,z)<\infty,$

where $\bar{Z}_{t}$ is that state of the approximate PDMP at time $t$ and
expectation is with respect to this state assuming $\bar{Z}_{0}=z$.

Part (A) follows (i) as the two rates are identical at $s=0$, i.e.
$\bar{\lambda}_{i,j}(z,0)=\lambda^{i,j}((x,v))$ and (ii) both rates have
bounded derivative with respect to $s$ (this is simple to show for
$\bar{\lambda}_{i,j}(z,s)$ from its definition, and is by assumption for
$\lambda_{i,j}((x+sv,v))$). The constant $\bar{M}(z)$ can be defined by to be
the sum of the maximum of each of these derivatives for $0\leq
s\leq\varepsilon_{0}$. Part (B) will then follow immediately as the region of
possible values of $\bar{Z}_{n\varepsilon}$ within a finite time interval is
compact.

##  Appendix B Comparison with Stochastic Bouncy Particle Sampler

In this section, we compare SG-PDMPs with the Stochastic Bouncy Particle
Sampler (SBPS), and its adaptive preconditioned version (pSBPS) as presented
in Pakman et al. (2017). The SBPS and pSBPS algorithms attempt to tackle a
similar underlying problem as our paper: how to approximately implement a PDMP
sampler with sub-sampling. The challenge lies in how to efficiently sample
events of the PDMP with sub-sampling. As we shall explain below, our approach
is simpler and computationally more efficient than SBPS and pSBP.

There are a number of substantive and practical differences between this
approach and our algorithms. Our approach discretises time, and simulates
events with a constant rate within each time-interval. This constant rate is
simple to calculate, and involves accessing a single data point and
calculating the current event rate associated with the process. By comparison,
SBPS involves fitting a linear model to the rate based on the observed rates
for mini-batches of data for rejected events since the last event. The SBPS
approach needs to account for the randomness across the choice of mini-batch,
and deal with the challenge of fitting a (probably incorrect) linear model to
few (initially just one) data points. Together this means that SBPS uses much
larger mini-batch sizes: the authors considered mini-batch sizes of
approximately $n=0.1N$ compared to $n=1$ for our Euler scheme (a lower mini-
batch can affect negatively the performance of SBPS as the thinning rate grows
linearly with $N/n$ – see Equation 10 in Pakman et al. (2017)). The
uncertainty in fitting the linear model means one has to take conservative
upper bounds, which can also lead to a resulting loss of efficiency.

We consider two Baysian problems: a Bayesian linear regression and a Bayesian
logistic regression. In both examples, we set an improper uniform prior. In
both cases, we set the dimension to be $d=10$ and we simulate $N=10^{5}$
covariates $X_{i,j}\sim\mathcal{N}(0,\mathrm{I}_{d}))$. $y$ was simulated from
the model with the true parameters $x^{\star}_{i}\sim\mathcal{N}(-5,\,100)$
(for the linear regression, we re-scaled the noise with $\sqrt{N}$). For both
problems, we perform analogous experiments and we set the sub-sample batch
size for the SBPS to be equal to 10%$N$ as in Pakman et al. (2017) and all the
other parameters suggested therein. Note that the need for SBPS to scale
minibatch size with $N$ means that the computational advantages of our
approach will grow (linearly in $N$) for larger data sets.

For both problems, we compare SBPS and pSBPS with SG-PDMPs in two numerical
experiments. To make the comparison fair, for all implementations we have
thinned the continuous-time output so that, for the same CPU cost, each
algorithm outputs the same number of samples. In the first experiment we
initialise all algorithms far away from the bulk of the posterior measure, by
setting $x_{0}=\mathcal{N}(0,100^{2})$.

With this initialisation, we observe that, in the logistic regression problem,
SBPS (respectively pSBPS) fails to approximate the Bouncy Particle Sampler
with sub-sampling: it estimates an upper bound which fails to bound the real
rate more than $50\%$ (respectively $22\%$). This undesirable behaviour
translates in an algorithm which does not converge to the posterior density.
In contrast, SG-PDMPs converge to the bulk of the posterior measure, for
different values of step-size $h$. In the linear regression example, SBPS and
pSBPS converge to the bulk of the distribution, but the rate of convergence is
significantly slower compared to SG-PDMPs.

Figure 7 and Figure 10 display the trace of the first coordinate of SBPS, SG-
BPS, SG-ZZ with step-size equal to $h=10^{-3}$ and $h=10^{-4}$ for the
logistic regression and linear regression examples.

In the second experiment, we initialise all algorithms close to the posterior
mode. In this setting, SBPS behaves better: for example, for the logistic
regression, SBPS (respectively pSBPS) fails to upper bound the real rate $3\%$
($4\%$) of the time. However, when comparing SBPS to SG-PDMPs, we notice that
SBPS is substantially worse at mixing for both of the problems considered. We
show in Figure 5 and in Figure 8 the traces of the first coordinate fo SBPS,
SG-BPS and SG-ZZ (with $h=10^{-3},\,10^{-4}$) and the auto-correlation
functions for the logistic and linear regression problem. Furthermore, for
each algorithm we compute

$\varepsilon(t_{1},t_{2})=\frac{1}{d}\sum_{i=1}^{d}\left(\frac{\hat{\sigma}_{i}^{(t_{1},t_{2})}-\sigma_{i}}{\sigma_{i}}\right)^{2}$

where $\hat{\sigma}_{i}^{(t_{1},t_{2})},\,t_{1}\leq t_{2}$ is the empirical
standard deviation of the $ith$ coordinate computed with the output with
burnin equal to $t_{1}$ and total number of sample points equal to $t_{2}$ and
$\sigma_{i}$ is the empirical standard deviation computed with a long-run
($10^{8}$ iterations) of the SG-BPS with small step-size $h=10^{-5}$. Figure 7
and Figure 10 show for each algorithm
$\varepsilon(\lfloor\frac{t}{2}\rfloor,t)$ for different values of $t$
respectively for the logistic and linear regression problem.

Figure 5:  From the top-left to the bottom-right. Traces and auto-correlation
function of the first coordinate of SBPS, SG-BPS, SG-ZZ (with step-sizes
$10^{-3},10^{-4}$) relative to the Bayesian logistic regression model. All
algorithms where initialised near the mode of the posterior (second
experiment).

Figure 6:  Traces of the first coordinate of SBPS, SG-BPS, SG-ZZ (with step-
sizes $10^{-3},10^{-4}$) relative to the Bayesian logistic regression model.
All algorithms where initialised far away from the bulk of the posterior
density (first experiment).

Figure 7:  $\varepsilon(t\div 2,t)$ for $t=1,2,\dots,6000$ relative to the
Bayesian logistic regression model. All algorithms where initialised near the
mode of the posterior (second experiment).

Figure 8:  From the top-left to the bottom-right. Traces and auto-correlation
function of the first coordinate of SBPS, SG-BPS, SG-ZZ (with step-sizes
$10^{-3},10^{-4}$) relative to the Bayesian linear regression model. All
algorithms where initialised near the mode of the posterior (second
experiment).

Figure 9:  Traces of the first coordinate of SBPS, SG-BPS, SG-ZZ (with step-
sizes $10^{-3},10^{-4}$) relative to the Bayesian linear regression model. All
algorithms where initialised far away from the bulk of the posterior density
(first experiment).

Figure 10:  $\varepsilon(t\div 2,t)$ for $t=1,2,\dots,6000$ relative to the
Bayesian linear regression model. All algorithms where initialised near the
mode of the posterior (second experiment).

##  Appendix C Stochastic Gradient ZZ and BPS

In Algorithms 2 and 3 we describe the specific implementation of stochastic
gradient versions of the Zig-Zag sampler and the Bouncy Particle Sampler.

Algorithm 2 Stochastic gradient Zig-Zag (SG-ZZ) sampler

0: $(x,v)\in E$, step size $\varepsilon>0$, time horizon $T>0$.

$t=0$, $\delta_{t}=\varepsilon$

while $t<T$ do

$J\sim\mathrm{Unif}(1,2,\dots,N)$

$\tau_{i}\sim\mathrm{Exp}(\lambda^{i,J}(z)),\mbox{ for }i=1,\ldots,d$

$\tau=\min_{i=1,\ldots,d}\,\tau_{i},\,i^{*}=\operatorname*{argmin}\tau_{i}$

if $\tau<\delta_{t}$ then

$\delta_{t}\leftarrow\delta_{t}-\tau,\,x\leftarrow x+v\tau,\,t\leftarrow
t+\tau$

$v\leftarrow\mathrm{R}_{i^{*}}(v)$ {Flip $i^{*}$ velocity component}

else

$x\leftarrow x+v\delta_{t},\,t\leftarrow t+\delta_{t}$

$\delta_{t}=\varepsilon$

Save $(x,t)$

end if

end while

Algorithm 3 Stochastic gradient Bouncy Particle Sampler (SG-BPS)

0: $(x,v)\in E$, step size $\varepsilon>0$, time horizon $T>0$.

$t=0$, $\delta_{t}=\varepsilon$

while $t<T$ do

$J\sim\mathrm{Unif}(1,2,\dots,d)$

$\tau\sim\mathrm{Exp}(\lambda^{J}(z))$

if $\tau<\delta_{t}$ then

$v^{(\mathrm{new})}=F^{J}_{x}(v)$ {Velocity reflection}

$\delta_{t}\leftarrow\delta_{t}-\tau,\,x\leftarrow x+v\tau,\,t\leftarrow
t+\tau$

$v=v^{\mathrm{(new)}}$

else

$x\leftarrow x+v\delta_{t},\,t\leftarrow t+\delta_{t}$

$\delta_{t}=\varepsilon$

Save $(x,t)$

end if

end while

##  Appendix D Comparison of SG-PDMPs with varying mini-batch size

PDMP samplers with subsampling and SG-PDMPs, as presented in Section 2-3, can
be extended by using, at every iteration, a batch-size of $n<N$ random data
points instead of a single one. We view the ability to use a mini-batch size
of 1 to be a key advantage of our method and, in most contexts, we would
expect worse performance (relative to CPU cost) for larger mini-batches. In
this section we support our claim with numerical experiments. A summary of our
experiments is shown in Figure 11 where we run the SG-Zig-Zag algorithm with
different mini-batch sizes for a linear regression problem with $10^{6}$
observations. As can be seen, for a fixed CPU cost using a mini-batch size of
1 is the best – this is most strikingly shown in the right-hand plot.

Figure 11: Left panel: error in the standard deviation estimation of each
coordinate $\mathcal{E}^{(h,d)}$ (y-axis) against the step-size divided by the
mini-batch size $\frac{h}{n}$ (x-axis, on a log-scale). We run all algorithms
for the same time horizon so that $\frac{h}{n}$ is a proxy for the inverse of
the running time of each algorithm. Middle panel: traces of one coordinate
relative of SG-ZZ with step-size h = $10^{-3}$ and mini-batch size equal to 1
(blue), 10 (red), $10^{2}$ (green). All algorithms were initialised out of
stationarity. The star symbols indicate the location of each Markov chain
after $25\times 10^{3}$ gradient evaluations. Right panel: traces normalised
by the mini-batch size.

##  Appendix E Sticky Zig-Zag sampler

In what follows, we present the stochastic gradient sticky Zig-Zag (SG-SZZ)
sampler which approximates the sticky Zig-Zag (Bierkens et al. (2023a)). In
this case, the target distribution is assumed to be of the form

$\pi(\mathrm{d}x)\propto\exp(-U(x))\prod_{i=1}^{d}(\mathrm{d}x_{i}+\frac{1}{\kappa_{i}}\delta_{0}(\mathrm{d}x_{i})$
(10)

where $U(x)=\sum_{j=1}^{N}U_{j}(x)$. With simple algebra, it is not difficult
to see that the taregt measure given by a smooth log-likelihood
$\ell(x)=\sum_{i=1}^{N}\ell_{i}(x)$ and spike and slab prior

$\pi_{0}(\mathrm{d}x)=\otimes_{i=1}^{d}(1-w_{i})\pi_{i}(x_{i})\mathrm{d}x_{i}+w_{i}\delta_{0}(\mathrm{d}x_{i})$
(11)

is of the form of equation (10), with

$U_{i}(x)=C-\ell_{i}(x)-\frac{1}{N}\sum_{j=1}^{d}\log(\pi_{j}(x_{j})),\qquad\kappa_{i}=\frac{1-w_{i}}{w}\pi_{i}(0),$

for some constant $C$ indepent from $x$. In algorithm 4, we present the SG-SZZ
as a minor modification of the SG-ZZ, enabling to approximate measures which
are of the form of (10).

Algorithm 4 Stochastic gradient sticky Zig-Zag (SG-SZZ) samplers

0: $(x,v,s)\in E$, step size $\varepsilon>0$, time horizon $T>0$, sticky
parameter $\kappa\in(0,\infty)^{d}$.

$t=0$, $\delta_{t}=\varepsilon$

$v^{c}=v$ {Copy of velocity}

$A=\\{1,2,\dots,d\\}$, $A^{c}=\emptyset$

while $t<T$ do

$J\sim\mathrm{Unif}(1,2,\dots,d)$

$\tau=\min_{i\in A}\\{\tau_{i}\sim\mathrm{Exp}(\lambda^{i,J}(z))\\}$

$\tau^{\star}=\min_{i\in A}\\{t_{i}^{\star}\colon x_{i}+v_{i}t_{i}=0\\}$

$\tau^{\circ}=\min_{i\in A^{c}}\\{t_{i}^{\circ}\sim\mathrm{Exp}(k_{i})\\}$

if $\tau<\min(\delta_{t},\tau^{\star},\tau^{\circ})$ then

$x=x+v\tau,\,t=t+\tau,\,\delta_{t}\leftarrow\delta_{t}-\tau$

$v\leftarrow\mathrm{R}_{j}(v)$ where $j:=\mathrm{argmin}_{i\in
A}\\{\tau_{i}\\}$

else if $\tau^{\star}<\min(\delta_{t},\tau^{\circ})$ then

$x=x+v\tau^{\star},\,t=t+\tau^{\star},\,\delta_{t}\leftarrow\delta_{t}-\tau^{\star}$

$A^{c}\leftarrow j^{\star}$ where $j^{\star}:=\mathrm{argmin}_{i\in
A}\\{\tau_{i}^{\star}\\}$

$v_{i}=0$ {Stick }

else if $\tau^{\circ}<\delta_{t}$ then

$x=x+v\tau^{\circ},\,t=t+\tau^{\circ},\,\delta_{t}\leftarrow\delta_{t}-\tau^{\circ}$

$A\leftarrow j^{\circ}$ where $j^{\circ}:=\mathrm{argmin}_{i\in
A^{c}}\\{\tau_{i}^{\circ}\\}$

$v_{i}=v_{i}^{c}$ {Unstick}

else

$x=x+v\delta_{t},\,t=t+\delta_{t}$

$\delta_{t}=\varepsilon$

Save $(t,x)$

end if

end while

##  Appendix F Logistic regression

###  F.1 Stein discrepancy kernel

The Stein discrepancy kernel evaluated on $K$ sample points with dimenision
$d$ is defined as

$\sum_{k=1}^{d}\sqrt{\sum_{i,j=1}^{K}\frac{\kappa_{k}(x^{(i)},x^{(j)})}{K^{2}}}$

where

$\kappa_{k}(y,z)=\partial_{x_{k}}U(x)\partial_{y_{k}}U(y)\kappa(x,y)+\partial_{x_{k}}U(x)\partial_{y_{k}}\kappa(x,y)+\partial_{y_{k}}U(y)\partial_{x_{k}}\kappa(x,y)+\partial_{x_{k}}\partial_{y_{k}}\kappa(x,y).$

Following Nemeth and Fearnhead (2021) Section 4, we choose
$\kappa(x,y)=(c^{2}+\|x-y\|^{2})^{\beta},\,\beta\in(-1,0),\,c>0$. Gorham and
Mackey (2017) shows that this metric is able to detect both poor mixing and
bias from the target distribution.

###  F.2 Further simulations

We compute the Stein discrepancy kernel relative to the traces obtained with
SGLD with 1, 10, 100 minibatch sample points and SG-ZZ and SG-BPS. We run each
algorithm varying $h$ and we plot in Figure 12, top panel the results. Figure
12, bottom panel, displays the loss function of the samplers above and
additionaly SG-SZZ with spike-and-slab prior as in Equation (11) with
$w_{i}=0.5,\,i=1,2,\dots,d$. The loss function has been averaged over $M=10$
random permutation of train and test datasets, as described in Section 4.1 of
the main manuscript. For this model, we set the loss function of Equation (6)
equal to

$\ell(X_{i},y_{i},x)=y_{i}\log\left(\frac{1}{1+\exp(-\langle
x,X_{i}\rangle)}\right)+(1-y_{i})\log\left(1-\frac{1}{1+\exp(-\langle
x,X_{i}\rangle)}\right).$

Figure 12: Stein discrepancy kernel (left panel) and average loss function on
$M=10$ random permutation of train and test sets (right panel) for the
logistic regression as a function of the stepsize $h$ ($x$-axis in log-scale).
Dashed lines corresponds to SGLD algorithms, solid lines to SG-PDMPs. For
SGLD, the Stein Discrepancy kernel was finite only for $h\leq 10^{-5}$.

##  Appendix G Bayesian neural networks

In this section, we show the trace of the loss function (Figure 13) and report
the average loss on the test sets relative to the traces obtained by running
SGLD, SG-ZZ, SG-BPS, SG-SZZ on $M=3$ different permuation of train and test
sets for the other 2 datasets from _UCI machine learning repository_
222https://archive.ics.uci.edu/ labeled as ‘concrete’ and ‘kin8nm’ (Table 3)
which have respectively $N=1031,N=45731$ and $p=8,p=9$ covariates.

Figure 13: Trace of the loss function for each sampler for different step-
sizes $h$ relative to the first permutation of the training and test set of
the dataset ‘concrete’ (left panel) and ‘kin8nm’ (right panel). SGLD was
unstable for $h=10^{-3}$ for the dataset ‘concrete’ and for
$h=10^{-3},10^{-4}$ for the dataset ‘kin8nm’.

Concrete | h
---|---
sampler | $10^{-3}$ | $10^{-4}$ | $10^{-5}$ | $10^{-6}$
SGLD | _NaN_ | _0.5219_ | 0.4625 | 0.4513
SG-ZZ | 0.5103 | 0.4803 | 0.4526 | 0.4739
SG-BPS | 0.4681 | 0.4606 | 0.4759 | _0.4885_
SG-SZZ | 0.5167 | 0.5017 | _0.4932_ | 0.4877

Kin8mn | h
---|---
sampler | $10^{-3}$ | $10^{-4}$ | $10^{-5}$ | $10^{-6}$
SGLD | _NaN_ | _NaN_ | _0.4480_ | 0.4143
SG-ZZ | 0.4536 | 0.4465 | 0.4224 | 0.4250
SG-BPS | 0.4257 | 0.4272 | 0.4230 | 0.4460
SG-SZZ | 0.4520 | 0.4375 | 0.4233 | 0.4506

Table 3: Average loss for each sampler for different values of step size $h$
relative to the dataset ‘concrete’ (left table) and ‘kin8mn’ (right table). In
bold the best performance and italic the worst performance given a step size
$h$.