# Coloured gravitational instantons, the strong CP problem and the companion
axion solution.

Zhe Chen<EMAIL_ADDRESS>Archil Kobakhidze
<EMAIL_ADDRESS>Sydney Consortium for Particle Physics and
Cosmology,
School of Physics, The University of Sydney, NSW 2006, Australia

###### Abstract

Quantum gravity introduces a new source of the combined parity (CP) violation
in gauge theories. We argue that this new CP violation gets bundled with the
strong CP violation through the coloured gravitational instantons.
Consequently, the standard axion solution to the strong CP problem is
compromised. Further, we argue that the ultimate solution to the strong CP
problem must involve at least one additional axion particle.

††preprint: CPPC-2021-03

##  I The strong CP problem

The Yang-Mills gauge theories exhibit a multiplicity of topologically distinct
vacua tHooft:1976rip ; tHooft:1976snw . Furthermore, quantum transitions
between vacua have non-zero probability due to instantons Belavin:1975fg  and,
thus, the ”true” vacuum, known as the $\theta$-vacuum, is given by a
superposition of those vacua. When applied to the gauge theory of strong
interactions, the Quantum Chromodynamics (QCD), this intricate vacuum
structure leads to the violation of the combined charge (C) and parity (P)
invariance, the CP invariance. The CP violation in QCD, however, is at odds
with the experimental limits on the neutron electric dipole moment, unless the
relevant CP parameter, $\theta_{\mathrm{QCD}}$, is fine-tuned to a tiny
value111Here $\theta_{\mathrm{QCD}}$ denotes the CP-violating parameter that
includes a contribution from the electroweak sector. Crewther:1979pi  (for a
review, see, e.g., Peccei:2006as ),

$\theta_{\mathrm{QCD}}\lesssim 10^{-10}~{}.$ (1)

The Peccei-Quinn mechanism Peccei:1977hh ; Peccei:1977ur  is arguably the most
elegant solution to the strong CP problem. It postulates gobal chiral symmetry
$U(1)_{\mathrm{PQ}}$, the Peccei-Quinn symmetry, which exhibits a mixed
$U(1)_{\mathrm{PQ}}$-QCD quantum anomaly, and the resulting violation of
$U(1)_{\mathrm{PQ}}$ charge conservation is due to the same QCD instantons. In
addition, the Peccei-Quinn symmetry is broken spontaneously, producing a light
pseudo-Goldstone boson, the axion Weinberg:1977ma ; Wilczek:1977pj ;
Kim:1979if ; Shifman:1979if ; Zhitnitsky:1980tq ; Dine:1981rt . In the absence
of any other source of $U(1)_{\mathrm{PQ}}$ symmetry breaking, the axion
automatically develops the vacuum expectation value that cancels out
$\theta_{\mathrm{QCD}}$, so that the vacuum state becomes CP-invariant. As a
bonus to the solution of the strong CP problem, the axion represents a viable
candidate for dark matter.

It has been known for quite some time that gravitational instantons222For a
comprehensive review of gravitational instantons, see Ref. Eguchi:1980jx . may
support new CP-violating topological terms in the standard model Deser:1980kc
; Holman:1992ah ; Arunasalam:2018eaz 333Destabilising role of gravitational
instantons is also discussed in Ref. Dvali:2005an  within the dual description
of the strong CP problem. For example, electrically charged Eguchi-Hanson (EH)
gravitational instantons Eguchi:1978xp ; Eguchi:1978gw  induce the respective
electromagnetic and gravitational $\theta$-terms in the effective Lagrangian
of the standard model Holman:1992ah ; Arunasalam:2018eaz . The impact of this
new CP-violation on the standard axion solution to the strong CP problem has
been examined in Ref. Holman:1992ah  and found that the electromagnetic
$\theta_{\mathrm{EM}}$ contribution to the effective axion potential is small
enough (at least within the standard model) to affect the original
solution444It has been also pointed out, that quantum gravity effects may
explicitly break the global Peccei-Quinn symmetry and thus potentially
jeopardise the axion solution to the strong CP problem Kamionkowski:1992mf .
Such breaking is typically accounted for within the effective field theory by
high-dimensional operators with _a priory_ unknown strengths. Some may argue
that the operators are actually sufficiently suppressed to alter the Peccei-
Quinn solution Kallosh:1995hi ; Dvali:2013cpa , or there could be additional
theoretical input to avoid the ”axion quality” problem (for a recent
discussion see, e.g. Cox:2019rro  .

However, there are known non-Abelian counterparts of the electromagnetic
gravitational instantons BoutalebJoutei:1979rw ; Chakrabarti:1987kz , in
particular, the QCD coloured Eguchi-Hanson (CEH) instantons. In this paper, we
describe these instantons and compute their contribution to the effective
axion potential. We find that CEH instantons bundle the new gravity-induced
source of CP-violation with the CP-violation in the sector of strong
interactions in such a way that compromises the original axion solution to the
strong CP problem. Further, we propose an _à la_ Peccei-Quinn remedy to this
re-emerged strong CP problem by introducing spontaneously broken anomalous
$U(1)^{\prime}_{\mathrm{PQ}}$ in addition to the standard
$U(1)_{\mathrm{PQ}}$. Consequently, in our scenario, the Peccei-Quinn axion is
necessarily accompanied by a companion axion. The phenomenology of this
companion-axion solution to the strong CP problem are discussed in
Chen:2021hfq ; Chen:2021wcf .

The paper is organised as follows. In the next section, we describe
asymptotically locally flat CEH instantons. In section III, we compute the CEH
contribution to the effective axion potential and demonstrate that the
standard Peccei-Quinn mechanism fails to solve the strong CP problem.
Following this, in section IV, we propose a remedy to this problem by
introducing a companion axion in addition to the standard Peccei-Quinn axion.
The last section V we reserve for conclusions.

##  II Coloured gravitational instantons

In this section, we recall the salient features of the combined Yang-Mills -
gravitational instantons. First, we point out that since we are interested in
processes described by (asymptotically) flat spacetime S-matrix, we restrict
ourselves to asymptotically Euclidean (AE) or asymptotically locally Euclidean
(ALE) manifolds. The AE and ALE vacuum manifolds are known to be Ricci flat,
$R=0$, and have non-negative Einstein-Hilbert action, $S_{\mathrm{EH}}\geq 0$,
according to the positive action theorem Schon:1979uj ; Witten:1981mf .
Furthermore, while for AE manifolds $S_{\mathrm{EH}}=0$ implies that they are
Riemann-flat (no gravity), ALE manifolds with $S_{EH}=0$ support non-trivial
(anti)self-dual Riemann curvature. Hence, it is reasonable to think that ALE
vacuum manifolds support topologically non-trivial vacuum structure of a
combined Yang-Mills-gravity theory in addition to the Yang-Mills instanton
vacuum structure in flat spacetime.

The combined Yang-Mills - gravitational instantons BoutalebJoutei:1979rw ;
Chakrabarti:1987kz  can be defined on multi-centred ALE spacetimes
Gibbons:1978tef  and their topological properties are dramatically different
from their flat spacetime counterparts. This is largely related to the fact
that the infinity of ALE spacetime is topologically equivalent to
$S^{3}/Z_{N}$ ($N=1$ is the usual flat spacetime). Hence, $SU(2)$ instantons
in ALE spacetimes describe transitions between topologically distinct vacua
that are classified homotopy of the mappings $S^{3}\to S^{3}$ modulo $Z_{N}$.
In what follows we describe the simplest instanton on the EH ALE space
($N=2$), where the $SU(2)$ Yang-Mills configuration is defined through the EH
spin-connection and hence its topological index is entirely determined by the
topology of the base space. It turns out that this so-called ”standard
embedding” is the most relevant instanton for the strong CP problem and its
axion solution.

The metric for the pure gravitational EH instanton is given by:

$\displaystyle ds^{2}$ $\displaystyle=$
$\displaystyle\frac{1}{1-\frac{a^{4}}{r^{4}}}dr^{2}+\frac{r^{2}}{4}\left[d\theta^{2}+\sin^{2}\theta
d\phi^{2}\right.$ (2) $\displaystyle+$
$\displaystyle\left.\left(1-\frac{a^{4}}{r^{4}}\right)\left(d\psi+\cos\theta
d\phi\right)^{2}\right]~{},$

where $r$ is the 4D radial coordinate and $(\theta,\phi,\psi)$ are angular
coordinates; $a$ is an arbitrary integration constant interpreted as the
instanton size. This instanton is known to correspond to the two-centre
instanton of Ref. Gibbons:1978tef  (see the explicit coordinate
transformations that relate two metrics in Ref. Prasad:1979kg )

The metric (2) is evidently singular at $r=a$. However, this singularity can
be removed by performing a $\mathbb{Z}_{2}$ identification of the coordinates.
This can be seen as follows. Letting $u=r\sqrt{1-a^{4}/r^{4}}$, it can be
shown that near $r=a$ (or $u=0$), the metric can be rewritten in terms of the
Euler angles on $S^{3}$ as Eguchi:1978gw :

$\displaystyle
ds^{2}\simeq\frac{1}{4}\mathop{du^{2}}+\frac{1}{4}u^{2}(\mathop{d\psi}+\cos\theta\mathop{d\phi})^{2}+\frac{a^{2}}{4}(\mathop{d\theta^{2}+\sin^{2}\theta\mathop{d\phi^{2}}})$
(3)

Here, it is evident that at fixed $\theta$ and $\phi$, the metric becomes the
usual metric of a plane with radial coordinate $u$ and angular coordinate
$\psi$. Therefore, to remove the apparent singularity at $u=0$, one must
restrict the domain of $\psi$ to $[0,2\pi)$. With this modification, it is
clear that the topology of this space near the horizon, $r=a$ is that of
$S^{2}\times\mathbb{R}^{2}$ where the sphere is parametrised by
$(\theta,\phi)$ and the plane by $(u,\psi)$. As $r\to\infty$, the metric
asymptotically approaches that of flat spacetime but with the restriction in
the domain of $\psi$, we see that the boundary at infinity is in fact
$S^{3}/\mathbb{Z}_{2}=\mathbb{R}\mathrm{P}^{3}$.

The EH spacetime (2) is Ricci flat, i.e., $R=0$. One can easily verify that
the Gibbons-Hawking boundary term is also vanishing and hence the total action
is zero Eguchi:1978xp ; Eguchi:1978gw . This may create a false impression
that the EH instanton could mediate unsuppressed anomalous processes with
chiral charge violation. However, there are no normalisable fermion zero-modes
in the background of the pure gravitational EH instanton, and hence such
instantons have no relation to the chiral anomaly.

Interestingly, the metric (2) supports a self-dual $SU(2)$ gauge field
configurations, which are relevant to QCD once embedded into $SU(3)$. Among
them, the focus of our interest here is the so-called ”standard embedding”
configuration, where $SU(2)$-valued gauge potential $A_{\mu}^{a}$ ($a=1,2,3$
are $SU(2)$ gauge indices) is expressed through the EH spin-connection
$\omega_{\mu}^{AB}$ ($A,B=0,1,2,3$ are tangent space indices) as:

$\displaystyle A^{a}_{\mu}=\frac{1}{2}\eta_{AB}^{a}\omega_{\mu}^{AB}~{},$
$\displaystyle\omega^{01}_{\theta}=\omega^{23}_{\theta}=\omega^{02}_{\phi}=\omega^{31}_{\phi}=\frac{1}{2}\sqrt{1-\frac{a^{4}}{r^{4}}}~{},$
$\displaystyle\omega^{03}_{\psi}=\omega^{12}_{\psi}=\frac{1}{2}\left(1+\frac{a^{4}}{r^{4}}\right)~{},$
(4)

where $\eta^{a}_{AB}$ is the self-dual ’tHooft symbol:
$\eta^{a}_{A0}=\delta_{aA}$ and $\eta^{a}_{AB}=\epsilon_{aAB}$ for
$A,B=1,2,3$. The above EH spin-connection is self-dual with respect to tangent
space indices:
$\omega_{\mu}^{AB}=-\frac{1}{2}\epsilon_{ABCD}\omega_{\mu}^{CD}$.
Correspondingly, the Rieman tensor and the $SU(2)$ field strength are self-
dual as well:
$R_{\mu\nu\rho\sigma}=\tilde{R}_{\mu\nu\rho\sigma}=\frac{\sqrt{g}}{2}\epsilon_{\mu\nu\alpha\beta}g^{\alpha\gamma}g^{\beta\delta}R_{\gamma\delta\rho\sigma}$
and
$F^{a}_{\mu\nu}=\tilde{F}^{a}_{\mu\nu}=\frac{\sqrt{g}}{2}\epsilon_{\mu\nu\alpha\beta}g^{\alpha\gamma}g^{\beta\delta}F^{a}_{\gamma\delta}$,
respectively ($\epsilon_{\mu\nu\rho\sigma}$ denotes the Levi-Chivita tensor
with $\epsilon_{0123}=1$). Since the pure gravitational part of the action is
zero for EH metric, the action for CEH instanton (4) is defined entirely by
the Yang-Mills part of the action:

$\displaystyle S_{\rm CEH}$ $\displaystyle=\frac{1}{4g^{2}}\int
d^{4}x\sqrt{g}F^{a\mu\nu}F^{a}_{\mu\nu}=\frac{4\pi^{2}}{g^{2}}\times 3~{}.$
(5)

Note the $4\pi^{2}$ factor in Eq. (5) vs the standard $8\pi^{2}$ which appears
in the action for pure Yang-Mills instantons. This is due $Z_{2}$
identification of antipodal points in the EH space, which becomes half of the
(asymptotic) Euclidean space. We also note that, unlike the non-gravitational
instantons, the final result in (5) originates from the two boundary
contributions evaluated at $r=a$ and at $r=\infty$.

The factor 3 in Eq. (5) can be associated with the gravitational Pontryagin
index which, due to the relations in (4), is twice the $SU(2)$ gauge
Pontryagin index:

$\displaystyle p$ $\displaystyle=$ $\displaystyle\frac{1}{16\pi^{2}}\int
d^{4}x\sqrt{g}R_{\mu\nu\rho\sigma}\tilde{R}^{\mu\nu\rho\sigma}$ (6)
$\displaystyle=$ $\displaystyle\frac{2}{32\pi^{2}}\int
d^{4}x\sqrt{g}F_{\mu\nu}^{a}\tilde{F}^{a\mu\nu}=3~{}.$

The above Pontryagin index in turn is related to the topological Hirzebruch
signature of EH spacetime, $\tau_{EH}=p/3=1$.

The vaccum-to-vaccum transitions mediated by CEH and BPST Belavin:1975fg
instantons are accounted by adding to the action $S[\Phi]$ two topological
terms:

$\displaystyle S_{eff}[\Phi]$ $\displaystyle=$ $\displaystyle
S[\Phi]+\frac{\theta_{QCD}}{32\pi^{2}}\int
d^{4}x\sqrt{g}F_{\mu\nu}^{a}\tilde{F}^{a\mu\nu}$ (7) $\displaystyle+$
$\displaystyle\frac{\theta_{EH}}{48\pi^{2}}\int
d^{4}x\sqrt{g}R_{\mu\nu\rho\sigma}\tilde{R}^{\mu\nu\rho\sigma}$

The normalisation of the topological terms is fixed such that the BPST-
instanton vacua are parametrised by $\theta_{QCD}$ and CEH-instanton vacua by
$\theta_{CEH}=3\theta_{QCD}/2+\theta_{EH}$.

In the presence of fermions, the transition amplitudes get modified. This is
captured by another important topological characteristic of the CEH instantons
the index of the Dirac operator. Following the generalised index theorem
Atiyah:1975jf , this index has been explicitly computed in Bianchi:1996zj  for
a generic representation $R$ of $SU(2)$ of dimension $d_{R}$. Adopting these
calculations for the ”standard embedding” EH instantons, we obtain:

$\displaystyle\nu^{(R)}_{1/2}=\left\\{\begin{tabular}[]{ll}$\frac{1}{4}d_{R}(d_{R}^{2}-2)$~{},&for
$d_{R}=2,4,...$\\\ \\\ $\frac{1}{4}d_{R}(d_{R}^{2}-1)$~{},&for
$d_{R}=1,3,...$\\\ \end{tabular}\right.$ (11)

We observe that for gauge-singlet fermions $d_{R}=1$ and $\nu^{(R)}_{1/2}=0$,
hence pure gravitational EH instantons do not mediate chiral-symmetry breaking
transitions; on the other hand, for the fermions (relevant for the Standard
Model quarks) in the fundamental representations $d_{R}=2$ the spin-1/2
fermion index is defined through the Hirzebruch signature,
$\nu^{(R)}_{1/2}=\tau_{EH}=1$. This is the same as the index in the background
of flat spacetime BPST instanton, although the Pontryagin indices of CEH and
BPST instantons differ. Consequently, the number of fermionic zero-modes is
the same in the backgrounds of CEH and BPST instantons and, hence, these
instantons generate in the effective Lagrangian the chiral symmetry breaking
terms of the form
$\propto\mathrm{e}^{-\frac{12\pi^{2}}{g^{2}}}\mathrm{det}\left(\psi_{L}\bar{\psi}_{R}\right)$
and
$\propto\mathrm{e}^{-\frac{8\pi^{2}}{g^{2}}}\mathrm{det}\left(\psi_{L}\bar{\psi}_{R}\right)$,
respectively.

The calculation of transition amplitudes requires the knowledge of the
explicit form of CEH-instanton zero modes. From topological considerations,
there are 12 bosonic zero modes Bianchi:1996zj  for $SU(2)$ CEH instanton, but
they are not known explicitly. We will present more detailed analysis of zero-
mode solutions elsewhere forthcoming2 , while here we estimate transition
amplitudes based on the following educated guess. Since both amplitudes
generated by the ordinary QCD and CEH instantons have the same dimension, the
Jacobian of collective coordinates of the CEH instanton must contain the same
dimensional factor factor $\rho^{-5}d\rho$. The norm of each of 12 zero modes
of $SU(2)$555We expect further 4 gauge zero modes from $SU(2)$ embedding into
$SU(3)$ Bernard:1979qt ; Vainshtein:1981wh . contributes a factor
$(8\pi^{2}/g^{2})^{1/2}$. Finally, the renormalisation will introduce a factor
$(\rho\Lambda)^{3b/2}$, where $b$ is the one-loop beta function coefficient
and $\Lambda$ is the regularisation scale. Other numerical factors are
expected to be the same order of magnitude as in the standard QCD
calculations. With this knowledge, we are ready now to compute the instanton-
induced axion potential.

##  III The axion potential

In the context of the standard model, the results from the previous section
imply that we have two independent $\theta$ parameters,
$\theta_{\mathrm{QCD}}$ and $\theta_{\mathrm{CEH}}$ (7), that parametrise
topologically non-trivial vacuum structure666We ignore here possible
electromagnetic and weak $\theta$ parameters, since their contributions are
negligible for the problem we are interested in.. It is important to note that
even if one tunes $\theta_{QCD}$ to be zero, the relevant parameter that must
satisfy the experimental constraint (1) becomes $\theta_{CEH}=\theta_{EH}$.
Therefore, any viable axion solution must ensure the cancellation of both CP-
violating contributions to the desired accuracy (see also Ref. Dvali:2005an
for a related discussion).

Let us consider a generic invisible axion model with the anomalous
$U(1)_{\mathrm{PQ}}$ Peccei-Quinn symmetry being spontaneously broken at a
high energy scale $f_{a}$. The leading contribution to the axion potential
comes from the low-energy theory large-scale QCD and CEH instantons, where the
relevant fermion degrees of freedom are $N_{f}=3$ light quarks
$q_{i}~{}(i=u,d,s)$. The $|\nu_{\mathrm{QCD}}|=1$ pure QCD and
$|\nu_{\mathrm{CEH}}|=1$ CEH instantons of sizes in the range
$[\rho,\rho+d\rho]$ induce 6-fermion effective operators (’t Hooft vertices):

$\displaystyle\Delta\mathcal{L}$
$\displaystyle=-d\left(\frac{2\pi}{\alpha_{s}(\Lambda)}\right)^{6}\mathrm{exp}\left[-\frac{2\pi}{\alpha_{s}(\Lambda)}+i\xi\frac{a}{f_{a}}+i\theta_{\mathrm{QCD}}\right]$
$\displaystyle\times\frac{d\rho}{\rho^{5}}(\Lambda\rho)^{b}\mathrm{det}\left(q_{iL}\bar{q}_{jR}\right)$
$\displaystyle-\bar{d}\left(\frac{2\pi}{\alpha_{s}(\Lambda)}\right)^{8}\mathrm{exp}\left[-\frac{3\pi}{\alpha_{s}(\Lambda)}+i\bar{\xi}\frac{a}{f_{a}}+i\theta_{\mathrm{CEH}}\right]$
$\displaystyle\times\frac{d\rho}{\rho^{5}}(\Lambda\rho)^{3b/2}\mathrm{det}\left(q_{iL}\bar{q}_{jR}\right)+\mathrm{h.c.}$
(12)

Here, $d\simeq\frac{\mathrm{e}^{-3.7+0.26N_{f}}}{\pi^{2}}\approx 0.017$ is the
renormalisation factor computed in the minimal subtraction scheme (we expect
$\bar{d}\simeq d$), $\alpha_{s}(\Lambda)$ is the QCD running coupling
evaluated at scale $\Lambda$ and $b=11-\frac{2}{3}N_{f}$ is the one-loop
$\beta$-function that governs evaluation of $\alpha_{s}$ at one-loop:
$\alpha^{-1}_{s}(\rho)=\alpha^{-1}_{s}(\Lambda)-\frac{b}{2\pi}\mathrm{ln}\left(\Lambda\rho\right)$;
$\xi=\sum_{q}\left[2Q_{\rm PQ}C_{2}(R)\right]$ and
$\bar{\xi}=\sum_{q}\left[3Q_{\rm PQ}C_{2}(R)+Q_{\rm PQ}d_{R}/4\right]$ is the
mixed $U(1)_{\mathrm{PQ}}$-QCD and the combined $U(1)_{\mathrm{PQ}}$-QCD-
gravity anomaly coefficients, respectively, that depend on the Peccei-Quinn
charge assignment of coloured fermions in a specific model. We note that, as
per standard approximation, we have ignored in (12) contributions from
instantons with higher topological charge (see, e.g. Rennecke:2020zgb ), as
well as contributions that steams from the overlap of QCD and CEH instantons.
All these contributions will be subdominant in the standard dilute instanton
gas approximation, where the instanton size is much smaller that average
inter-instanton distance.

To calculate axion potential we tie up the fermion (quasi)zero-modes via their
masses, i.e. $\left(q_{iL}\bar{q}_{iR}\right)\to m_{i}\rho$ in (12)
Vainshtein:1981wh , and integrate over instanton size $\rho$ [note,
$m_{i}\rho<<1$ in the dilute instanton approximation]. We obtain:

$\displaystyle V(a)$
$\displaystyle=-2K_{\mathrm{QCD}}\cos\left(\xi\frac{a}{f_{a}}+\theta_{\mathrm{QCD}}\right)$
$\displaystyle-2K_{\mathrm{CEH}}\cos\left(\bar{\xi}\frac{a}{f_{a}}+\theta_{\mathrm{CEH}}\right)~{},$
(13)

where,

$\displaystyle
K_{\mathrm{QCD}}=\frac{d}{8}\left(\frac{2\pi}{\alpha_{s}(\Lambda)}\right)^{6}\mathrm{e}^{-\frac{2\pi}{\alpha_{s}(\Lambda)}}m_{u}m_{d}m_{s}\Lambda$
$\displaystyle=-\frac{m_{u}m_{d}m_{s}\langle
0|\bar{q}q|0\rangle}{m_{u}m_{d}+m_{u}m_{s}+m_{d}m_{s}}\approx\frac{m_{\pi}^{2}f_{\pi}^{2}m_{u}m_{d}}{(m_{u}+m_{d})^{2}}~{},$
(14)
$\displaystyle\frac{K_{\mathrm{CEH}}}{K_{\mathrm{QCD}}}\approx\frac{16}{25}\left(\frac{2\pi}{\alpha_{s}(\Lambda)}\right)^{2}\mathrm{e}^{-\frac{\pi}{\alpha_{s}(\Lambda)}}\approx
0.06-0.4~{}.$ (15)

In Eq. (14) we have equated instanton calculation of $K_{\mathrm{QCD}}$ with
the same coefficient computed in chiral perturbation theory. The numerical
values in Eq. (15) are obtained by using the value of the strong coupling
constant calculated in the minimal substruction scheme,
$\alpha_{s}(\Lambda=1GeV)=0.4-0.6$ Deur:2016tte . We observe that CEH
contribution originates from large-size instantons and thus is not
exponentially suppressed, unlike the contribution from the electrically
charged instantons estimated in Ref. Holman:1992ah . This has important
ramifications for CP violation in QCD. It is easy to see that the axion ground
state described by the minimisation of the axion potential (13) is not CP-
invariant, unless two independent sources of CP violation are fine tuned,
$\xi\theta_{\mathrm{CEH}}=\bar{\xi}\theta_{\mathrm{QCD}}$. Consequently, the
neutron EDM bound (1) is not satisfied.

##  IV Companion axion solution to the strong CP problem

We saw in the previous section that the standard Peccei-Quinn mechanism with a
single axion fails to resolve the strong CP problem. Following the footsteps
of Peccei and Quinn, it is not difficult to find a remedy to this problem,
however. Let us assume that the theory posses a new spontaneously broken
anomalous symmetry $U(1)^{\prime}_{PQ}$ in addition to the standard
$U(1)_{PQ}$777We note in passing that in Ref. Dvali:2013cpa , it was suggested
that neutrino condensation and the associated ”$\eta$-meson” can resolve the
potential problem with gravitational anomalies.. in the low energy limit the
theory then would contain a new axion field $a^{\prime}$, the companion field
of the standard Peccei-Quinn axion $a$. The two-axion potential then reads:

$\displaystyle V(a,a^{\prime})$
$\displaystyle=-2K_{\mathrm{QCD}}\cos\left(\xi\frac{a}{f_{a}}+\xi^{\prime}\frac{a^{\prime}}{f_{a^{\prime}}}+\theta_{\mathrm{QCD}}\right)$
$\displaystyle-2K_{\mathrm{CEH}}\cos\left(\bar{\xi}\frac{a}{f_{a}}+\bar{\xi}^{\prime}\frac{a^{\prime}}{f_{a^{\prime}}}+\theta_{\mathrm{CEH}}\right)~{},$
(16)

where $\xi^{\prime}$ and $\bar{\xi}^{\prime}$ are the anomaly coefficient
associated with $U(1)^{\prime}_{PQ}$ and $f_{a^{\prime}}$ is the scale of
$U(1)^{\prime}_{PQ}$ breaking. The quick inspection of the above potential
shows that its minimum is achieved for CP-conserving vacuum state, that is,
the vacuum expectation values of the axion fields cancel out both
$\theta$-terms.

Note that two axion states are mixed through the interactions in Eq. (16) [the
mixing with neutral $\pi$ and $\eta$ mesons are neglected here] and hence do
not represent mass eigenstates of physical axions. We readily compute the
masses of physical axions,

$\displaystyle m^{2}_{\pm}$
$\displaystyle=K_{QCD}\left(\frac{\xi^{2}}{f^{2}_{a}}+\frac{\xi^{\prime
2}}{f^{2}_{a^{\prime}}}\right)+K_{CEH}\left(\frac{\bar{\xi}^{2}}{f^{2}_{a}}+\frac{\bar{\xi^{\prime}}^{2}}{f^{2}_{a^{\prime}}}\right)$
$\displaystyle\pm\left[\bigg{(}K_{QCD}\left(\frac{\xi^{2}}{f^{2}_{a}}-\frac{\xi^{\prime
2}}{f^{2}_{a^{\prime}}}\right)+K_{CEH}\left(\frac{\bar{\xi}^{2}}{f^{2}_{a}}-\frac{\bar{\xi^{\prime}}^{2}}{f^{2}_{a^{\prime}}}\right)\bigg{)}^{2}\right.$
$\displaystyle+16\left.\left(K_{QCD}\frac{\xi\xi^{\prime}}{f_{a}f_{a^{\prime}}}+K_{CEH}\frac{\bar{\xi}\bar{\xi^{\prime}}}{f_{a}f_{a^{\prime}}}\right)^{2}\right]^{1/2}~{},$
(17)

and the mixing angle between two states:

$\displaystyle\tan
2\alpha=\frac{2K_{QCD}\frac{\xi\xi^{\prime}}{f_{a}f_{a^{\prime}}}+2K_{CEH}\frac{\bar{\xi}\bar{\xi^{\prime}}}{f_{a}f_{a^{\prime}}}}{K_{QCD}\left(\frac{\xi^{2}}{f^{2}_{a}}-\frac{{\xi^{\prime}}^{2}}{f^{2}_{a^{\prime}}}\right)+K_{CEH}\left(\frac{\bar{\xi}^{2}}{f^{2}_{a}}-\frac{\bar{\xi^{\prime}}^{2}}{f^{2}_{a^{\prime}}}\right)}~{}.$
(18)

There is a wide range of options for masses and mixing, including hierarchical
masses with small or large mixing of the axion states. The couplings to the
standard model fields can be worked out straightforwardly. Some of the related
phenomenology is discussed in Refs. Chen:2021hfq ; Chen:2021wcf .

##  V Conclusion

In this paper, we have argued that the standard Peccei-Quinn solution to the
strong CP problem is compromised due to the new source of CP violation in
gravity-QCD sector. It is well-known quantum gravity effects via gravitational
instantons support new CP-violating topological terms in the effective
lagrangian of the standard model Deser:1980kc ; Holman:1992ah ;
Arunasalam:2018eaz . Namely, coloured gravitational instantons considered in
this paper together with the flat spacetime QCD instantons support both QCD as
well as gravitational $\theta$-terms (see, eq. (7)). Furthermore, CEH
”standard embedding” instantons mediate vacuum-to-vacuum transitions that
violate chiral and Peccei-Quinn global symmetries explicitly due to the
combined QCD and gravitational anomalies, and hence result in the dependence
of physical observables both on QCD and gravitational $\theta$-terms via
$\theta_{\rm CEH}$. The standard single axion model is incapable of removing
both the standard $\theta_{\rm QCD}$ and the additional $\theta_{\rm CEH}$
from the physical observables (e.g., from the neutron electric dipole moment),
and hence the standard Peccei-Quinn mechanism is compromised.

In this regard, it is worth to compare the contribution of CEH instantons to
the axion potential with other known quantum gravity induced contributions
Holman:1992ah ; Giddings:1987cg . In Ref. Holman:1992ah , the contribution of
electromagnetic Eguchi-Hanson instantons were estimated and it has found that
the dominant contribution comes from small-scale instantons of size $\sim
1/f_{a}$. The key reasons for that are (i) the running electromagnetic
coupling becomes larger at small scales, thus reducing the WKB-suppression
factor for instanton induced amplitudes as well as (ii) the fact that the
fermion zero-modes must be tied up through the Higgs boson loop in order to
have non-vanishing contribution. Contrary to that, the dominant contribution
to the axion potential comes from large-scale coloured Eguchi-Hanson
instantons (similar to the standard QCD instantons) where the fermion zero
modes are tied up by their masses. Simple comparison of WKB-factors reveal
that electromagnetic Eguchi-Hanson instanton contribution is negligible (see,
eq. (10) in Ref. Holman:1992ah ), while coloured Eguchi-Hanson instanton
contribution is comparable with the pure QCD instanton contribution (see, eq.
(15)).

Another known quantum gravity contribution to the axion potential is related
to the gravitational wormhole solutions Giddings:1987cg . Wormholes also break
the Peccei-Quinn symmetry explicitly and hence introduce additional terms in
the axion potential. However, the strength of this breaking is significant
only for small-size wormholes, $L\sim 1/\sqrt{M_{P}f_{a}}$, and is generally
negligible unless unless $f_{a}$ is too large, $f_{a}\gtrsim 10^{16}$ GeV
Alonso:2017avz . We must stress again, that the CEH instantons in contrast
dominate at low energies, and their effect cannot be offset by constraining
high-energy parameters.

With these conclusions, we then have argued that the genuine solution to the
strong CP problem _à la_ Peccei-Quinn must involve an additional companion
axion field. Such two-axion models have rich phenomenological and cosmological
implications that were partially studied in Chen:2021hfq ; Chen:2021wcf .

###### Acknowledgements.

This work of AK was partially supported by the Australian Research Council
through the Discovery Project grant DP210101636 and by Shota Rustaveli
National Science Foundation of Georgia (SRNSFG) through the grant DI-18-335.

## References

  * (1) G. ’t Hooft, “Symmetry Breaking Through Bell-Jackiw Anomalies,” Phys. Rev. Lett. 37, 8-11 (1976)
  * (2) G. ’t Hooft, “Computation of the Quantum Effects Due to a Four-Dimensional Pseudoparticle,” Phys. Rev. D 14, 3432-3450 (1976) [erratum: Phys. Rev. D 18, 2199 (1978)]
  * (3) A. A. Belavin, A. M. Polyakov, A. S. Schwartz and Y. S. Tyupkin, “Pseudoparticle Solutions of the Yang-Mills Equations,” Phys. Lett. B 59, 85-87 (1975)
  * (4) R. J. Crewther, P. Di Vecchia, G. Veneziano and E. Witten, “Chiral Estimate of the Electric Dipole Moment of the Neutron in Quantum Chromodynamics,” Phys. Lett. B 88, 123 (1979) [erratum: Phys. Lett. B 91, 487 (1980)]
  * (5) R. D. Peccei, “The Strong CP problem and axions,” Lect. Notes Phys. 741, 3-17 (2008) [arXiv:hep-ph/0607268 [hep-ph]].
  * (6) R. D. Peccei and H. R. Quinn, “CP Conservation in the Presence of Instantons,” Phys. Rev. Lett. 38, 1440-1443 (1977)
  * (7) R. D. Peccei and H. R. Quinn, “Constraints Imposed by CP Conservation in the Presence of Instantons,” Phys. Rev. D 16, 1791-1797 (1977)
  * (8) S. Weinberg, “A New Light Boson?,” Phys. Rev. Lett. 40, 223-226 (1978)
  * (9) F. Wilczek, “Problem of Strong $P$ and $T$ Invariance in the Presence of Instantons,” Phys. Rev. Lett. 40, 279-282 (1978)
  * (10) J. E. Kim, “Weak Interaction Singlet and Strong CP Invariance,” Phys. Rev. Lett. 43, 103 (1979)
  * (11) M. A. Shifman, A. I. Vainshtein and V. I. Zakharov, “Can Confinement Ensure Natural CP Invariance of Strong Interactions?,” Nucl. Phys. B 166, 493-506 (1980)
  * (12) A. R. Zhitnitsky, “On Possible Suppression of the Axion Hadron Interactions. (In Russian),” Sov. J. Nucl. Phys. 31, 260 (1980)
  * (13) M. Dine, W. Fischler and M. Srednicki, “A Simple Solution to the Strong CP Problem with a Harmless Axion,” Phys. Lett. B 104, 199-202 (1981)
  * (14) T. Eguchi, P. B. Gilkey and A. J. Hanson, “Gravitation, Gauge Theories and Differential Geometry,” Phys. Rept. 66, 213 (1980).
  * (15) S. Deser, M. J. Duff and C. J. Isham, “Gravitationally induced CP effects,” Phys. Lett. B 93, 419-423 (1980)
  * (16) R. Holman, T. W. Kephart and S. J. Rey, “Semiclassical gravity and invisible axions,” Phys. Rev. Lett. 71, 320-323 (1993) [arXiv:hep-ph/9207208 [hep-ph]].
  * (17) S. Arunasalam and A. Kobakhidze, “Charged gravitational instantons: extra CP violation and charge quantisation in the Standard Model,” Eur. Phys. J. C 79, no.1, 49 (2019) [arXiv:1808.01796 [hep-th]].
  * (18) G. Dvali, “Three-form gauging of axion symmetries and gravity,” [arXiv:hep-th/0507215 [hep-th]].
  * (19) T. Eguchi and A. J. Hanson, “Asymptotically Flat Selfdual Solutions to Euclidean Gravity,” Phys. Lett. B 74, 249-251 (1978)
  * (20) T. Eguchi and A. J. Hanson, “Selfdual Solutions to Euclidean Gravity,” Annals Phys. 120, 82 (1979)
  * (21) M. Kamionkowski and J. March-Russell, “Planck scale physics and the Peccei-Quinn mechanism,” Phys. Lett. B 282, 137-141 (1992) [arXiv:hep-th/9202003 [hep-th]].
  * (22) R. Kallosh, A. D. Linde, D. A. Linde and L. Susskind, “Gravity and global symmetries,” Phys. Rev. D 52, 912-935 (1995) [arXiv:hep-th/9502069 [hep-th]].
  * (23) G. Dvali, S. Folkerts and A. Franca, “How neutrino protects the axion,” Phys. Rev. D 89, no.10, 105025 (2014) [arXiv:1312.7273 [hep-th]].
  * (24) P. Cox, T. Gherghetta and M. D. Nguyen, “A Holographic Perspective on the Axion Quality Problem,” JHEP 01, 188 (2020) [arXiv:1911.09385 [hep-ph]].
  * (25) H. Boutaleb-Joutei, A. Chakrabarti and A. Comtet, “Gauge Field Configurations in Curved Space-times. 3. Selfdual SU(2) Fields in Eguchi-hanson Space,” Phys. Rev. D 21, 979-983 (1980)
  * (26) A. Chakrabarti, “Classical Solutions of Yang-Mills Fields. (Selected Topics),” Fortsch. Phys. 35, 1-64 (1987)
  * (27) Z. Chen, A. Kobakhidze, C. A. J. O’Hare, Z. S. C. Picker and G. Pierobon, “Phenomenology of the companion-axion model: photon couplings,” [arXiv:2109.12920 [hep-ph]].
  * (28) Z. Chen, A. Kobakhidze, C. A. J. O’Hare, Z. S. C. Picker and G. Pierobon, “Cosmology of the companion-axion model: dark matter, gravitational waves, and primordial black holes,” [arXiv:2110.11014 [hep-ph]].
  * (29) R. M. Schon and S. T. Yau, “Proof of the positive action conjecture in quantum relativity,” Phys. Rev. Lett. 42, 547-548 (1979)
  * (30) E. Witten, “A Simple Proof of the Positive Energy Theorem,” Commun. Math. Phys. 80, 381 (1981)
  * (31) G. W. Gibbons and S. W. Hawking, “Gravitational Multi - Instantons,” Phys. Lett. B 78, 430 (1978)
  * (32) M. K. Prasad, “Equivalence of Eguchi-Hanson metric to two-center Gibbons-Hawking metric,” Phys. Lett. B 83, 310-310 (1979)
  * (33) M. F. Atiyah, V. K. Patodi and I. M. Singer, “Spectral asymmetry and Riemannian Geometry 1,” Math. Proc. Cambridge Phil. Soc. 77, 43 (1975). doi:10.1017/S0305004100049410; ibid. 78, 405 (1976). ibid. 79, 71 (1976).
  * (34) M. Bianchi, F. Fucito, G. Rossi and M. Martellini, Nucl. Phys. B 473, 367-404 (1996) doi:10.1016/0550-3213(96)00240-4 [arXiv:hep-th/9601162 [hep-th]].
  * (35) Z. Chen and A. Kobakhidze, “Gauge zero-modes on ALE spaces,” _in progress_.
  * (36) C. W. Bernard, “Gauge Zero Modes, Instanton Determinants, and QCD Calculations,” Phys. Rev. D 19, 3013 (1979)
  * (37) A. I. Vainshtein, V. I. Zakharov, V. A. Novikov and M. A. Shifman, “ABC’s of Instantons,” Sov. Phys. Usp. 25, 195 (1982)
  * (38) F. Rennecke, “Higher topological charge and the QCD vacuum,” Phys. Rev. Res. 2, no.3, 033359 (2020) [arXiv:2003.13876 [hep-th]].
  * (39) A. Deur, S. J. Brodsky and G. F. de Teramond, Nucl. Phys. 90, 1 (2016) doi:10.1016/j.ppnp.2016.04.003 [arXiv:1604.08082 [hep-ph]].
  * (40) S. B. Giddings and A. Strominger, “Axion Induced Topology Change in Quantum Gravity and String Theory,” Nucl. Phys. B 306, 890-907 (1988)
  * (41) R. Alonso and A. Urbano, “Wormholes and masses for Goldstone bosons,” JHEP 02, 136 (2019) [arXiv:1706.07415 [hep-ph]].


# Piecewise Deterministic Markov Processes for Bayesian Neural Networks

Ethan Goan Queensland University of Technolgy  Dimitri Perrin  Queensland
University of Technolgy  Kerrie Mengersen  Queensland University of Technolgy
Clinton Fookes  Queensland University of Technolgy

###### Abstract

Inference on modern Bayesian Neural Networks (BNNs) often relies on a
variational inference treatment, imposing violated assumptions of independence
and the form of the posterior. Traditional MCMC approaches avoid these
assumptions at the cost of increased computation due to its incompatibility to
subsampling of the likelihood. New Piecewise Deterministic Markov Process
(PDMP) samplers permit subsampling, though introduce a model-specific
inhomogenous Poisson Process (IPPs) which is difficult to sample from. This
work introduces a new generic and adaptive thinning scheme for approximate
sampling from these IPPs, and demonstrates how this approach can accelerate
the application of PDMPs for inference in BNNs. Experimentation illustrates
how inference with these methods is computationally feasible, can improve
predictive accuracy, MCMC mixing performance, and provide informative
uncertainty measurements when compared against other approximate inference
schemes.

##  1 Introduction

Since Hamiltonian Monte Carlo (HMC) was first developed for Bayesian inference
Neal [2012], sampling methods have seen relatively little application to
Bayesian Neural Networks (BNNs). Flexibility, inference diagnostics and
asymptotic guarantees of HMC comes at the cost of computational complexity as
each data point needs to be used to compute the entire likelihood, and to
perform Metropolis Hastings corrections. As models and data sets have grown,
this expense has not been offset by the considerable performance increase in
computational hardware. A recent study found that the fitting of a HMC model
for ResNet20 required a computational cost equivalent to 60 million SGD epochs
to obtain only 240 samples from three chains Izmailov et al. [2021].

(a)

(b)

(c)

Figure 1: Example of correlations between the parameters in the first layer of
a BNN for a simple regression task. Plot (a) samples of predictive posterior
from proposed method, (b) correlation between all parameters on the first
layer, (c) kernel density estimate for a single parameter.

To circumnavigate the computational expense, much research has explored the
application of approximate inference through the lens of Variational Inference
(VI) Jordan et al. [1999], Wainwright and Jordan [2008], Blei et al. [2017] or
through exploiting properties of SGD Mandt et al. [2017]. VI replaces the true
target distribution with an approximate distribution that can be easily
manipulated, typically using a mean-field approach where independence between
parameters is assumed. These methods are attractive due to their reduced
computational complexity and their amenability to stochastic optimisation.
However, the suitability of these methods relies heavily on the expressiveness
of the approximate posterior to accurately model the true distribution. Given
the known correlations and frequent multi-modal structure amongst parameters
within BNNs Barber and Bishop [1998], MacKay [1995], a mean-field
approximation can be unsuitable for accurate inference. Figure 1 illustrates
these properties for a simple BNN. Stochastic gradient MCMC methods such as
Stochastic Gradient Langevin Dynamics (SGLD) aim to address this issue but
requires prohibitively small and decreasing learning rates to target the
posterior that limits their applicability Nagapetyan et al. [2017].

This work explores a new set of “exact” inference methods based on Piecewise
Deterministic Markov Processs (PDMPs) Davis [1984] to perform Bayesian
inference. PDMP methods can maintain the true posterior as its invariant
distribution during inference whilst permitting sub-sampling of the likelihood
at each update. This property is attractive for BNNs which typically are of
large dimension in terms of parameters and data sets. Furthermore, previous
research has highlighted PDMP methods for favourable performance in terms of
mixing and sampling efficiency Bouchard-Côté et al. [2018], Bierkens et al.
[2019], Wu and Robert [2017], Bierkens et al. [2020]. The dynamics of these
samplers are simple to simulate, though simulating the times to update these
dynamics is controlled by an Inhomogeneous Poisson Process (IPP) which can be
difficult to sample. This work explores an adaptive procedure to approximate
sampling of these event times to allow for approximate inference within the
context of BNNs. The contributions of this paper are the following,

  * •

Propose a novel adaptive thinning method for approximate sampling from IPPs
events

  * •

Develop a GPU-accelerated package for applying these methods to general models

  * •

Evaluate the performance of these methods applied to computer vision tasks
using BNNs;

  * •

Evaluate the suitability of PDMP samplers for BNNs and investigate how they
can improve predictive accuracy, calibration and posterior exploration when
compared against SGLD.

MCMC methods have often been seen as computationally prohibitive for models
with many parameters or where modern large data sets are used. It is hoped
that this work will demonstrate that approximate inference using MCMC
approaches for BNNs can be practically feasible, offer insightful results, and
to show how we can leverage exact methods for approximate inference to more
accurately target posterior distributions in BNNs.

##  2 Preliminaries

Following the description from Fearnhead et al. [2018], PDMP are defined by
three key components: piecewise deterministic dynamics, an event rate, and the
transition kernel. For inference, the goal is to design these three key
components such that we can use the properties of a PDMP to sample from the
posterior distributions of our parameters $\mathbf{\omega}$. We represent the
deterministic dynamics as $\Psi(\mathbf{\omega},\mathbf{v},t)$, where
$\mathbf{v}$ is an auxiliary velocity variable to guide posterior exploration
with known distribution $\Phi(\mathbf{v})$ and $t$ represents time. At random
events, these dynamics are updated in accordance to a specified transition
kernel. Upon an update event, the piecewise deterministic dynamics of the
system update according to the kernel, and the state $\mathbf{\omega}$ at the
time of the update event serves as the starting position for the next segment
such that they are all connected.

An IPP with rate function $\lambda(\mathbf{\omega}(t),\mathbf{v}(t))$ governs
the update times for the dynamics. All rate functions in this work rely upon
the negative joint log probability of the model,

$U(\omega)=-\log\Big{(}p(\omega)p(\mathcal{D}|\omega)\Big{)},$ (1)

where $p(\omega)$ is a prior or reference measure and $p(\mathcal{D}|\omega)$
is our likelihood, If these three components are suitably defined, these
processes can sample from a given posterior distribution. For derivations on
how to design these components to target a posterior distribution, the reader
can refer to Fearnhead et al. [2018], Vanetti et al. [2017], Davis [1993]. We
now introduce the samplers used within this work.

###  2.1 Bouncy Particle Sampler

The dynamics of the Bouncy Particle Sampler (BPS) Bouchard-Côté et al. [2018]
are given by
$\Psi(\mathbf{\omega},\mathbf{v},t)=\mathbf{\omega}^{i}+\mathbf{v}^{i}t$,
where the superscripts indicate a deterministic segment. The velocity remains
constant within these segments and the parameter space is explored linearly.
The velocity is updated at event times given by
$\tau\sim\text{IPP}(\lambda(\mathbf{\omega}(t)),\mathbf{v})$, where,

$\lambda(\mathbf{\omega}(t),\mathbf{v})=\max\\{0,\nabla
U(\mathbf{\omega})\cdot\mathbf{v}^{i}\\}.$ (2)

Once an event time is sampled, the state of our variable “bounces” according
to a lossless inelastic Newtonian collision,

$\mathbf{v}^{i+1}=\mathbf{v}^{i}-2\dfrac{\nabla
U(\mathbf{\omega}^{i+1})\cdot\mathbf{v}^{i}}{\left\lVert\nabla
U(\mathbf{\omega}^{i+1})\right\rVert^{2}}\nabla U(\mathbf{\omega}^{i+1})$ (3)

where $\mathbf{\omega}^{i+1}$ represents the end of the previous segment at
time $\tau$, and serves as the starting position for the following segment.
The BPS provides linear dynamics that are simple to simulate, though relies
only on local gradient information, which can lead to inefficient exploration
for BNNs. Preconditioning can allow us to address this.

###  2.2 Preconditioned BPS

To accelerate posterior exploration in directions of interest, we can
precondition the gradients to include more information about the structure of
our posterior space. Introduction of a preconditioning matrix $A$ results in
new dynamics of
$\Psi(\mathbf{\omega},\mathbf{v},t)=\mathbf{\omega}^{i}+A\mathbf{v}^{i}t$, and
a new event rate,
$\lambda(\mathbf{\omega}(t),\mathbf{v})=\max\\{0,\mathbf{v}\cdot A\nabla
U(\mathbf{\omega}+\mathbf{v}t)\\}$. Upon events, the velocity is updated
according to,

$\mathbf{v}^{i+1}=\mathbf{v}^{i}-2\dfrac{A\nabla
U(\mathbf{\omega}^{i+1})\cdot\mathbf{v}^{i}}{\left\lVert A\nabla
U(\mathbf{\omega}^{i+1})\right\rVert^{2}}A\nabla U(\mathbf{\omega}^{i+1}).$
(4)

With careful choice of $A$, exploration along certain axis can be
appropriately scaled. Pakman et al. [2017] propose a preconditioner similar to
Li et al. [2015], though our preliminary experimentation found inconsistent
results when applied to BNNs. Instead, we opt to build on the approach of
Bertazzi and Bierkens [2020], where we use variance information of our samples
to precondition our dynamics. We choose the preconditioner such that
$A=\text{diag}\big{(}\Sigma^{\frac{1}{2}}\big{)}$, where $\Sigma$ is the
estimated covariance in our sample found during a warm-up period. As such, we
refer to this sampler as the $\sigma$BPS .

###  2.3 Boomerang Sampler

The Boomerang Sampler Bierkens et al. [2020] introduces non-linear dynamics
for both parameter and velocity terms, and the inclusion of a Gaussian
reference measure for the parameters and velocity
$\mathcal{N}(\mathbf{\omega}_{\star},\Sigma_{\star})\otimes\mathcal{N}(0,\Sigma_{\star})$.
The first term in this reference measure will appear in the target
distribution similar to a prior in the joint probability over parameters, and
the second as the known distribution for the velocity component. The
parameters $\mathbf{\omega}_{\star}$ and $\Sigma_{\star}$ can be specified as
traditional prior, or can be specified in an empirical approach where they are
learnt from the data. Within this work, we will set $\mathbf{\omega}_{\star}$
to the MAP estimate. In the original paper, $\Sigma_{\star}$ is set to the
inverse of the Hessian, however, this can be computationally prohibitive for
BNNs. Instead, we approximate the diagonal of the Hessian matrix as,

$\Sigma_{\star}=\gamma\Big{[}\sum_{i=0}^{N}\frac{\partial}{\partial\mathbf{\omega}^{2}}-\log
p(\mathbf{y}_{i}|\mathbf{x}_{i},\mathbf{\omega})\Big{]}^{-1}$ (5)

where $N$ is the number of mini-batches present,
$\mathbf{x}_{i},\mathbf{y}_{i}$ represents data from the $i$’th mini-batch,
and $\gamma$ is a hyperparameter to adjust the scale as needed.

Unlike the BPS samplers, the velocity does not remain constant between events.
The dynamics of the Boomerang sampler for $\mathbf{\omega}$ and $\mathbf{v}$
within events are given by
$\Psi(\mathbf{\omega},\mathbf{v},t)_{\mathbf{\omega}}=\mathbf{\omega}_{\star}+(\mathbf{\omega}^{i}-\mathbf{\omega}_{\star})\cos(t)+\mathbf{v}^{i}\sin(t)$,
$\Psi(\mathbf{\omega},\mathbf{v},t)_{\mathbf{v}}=-(\mathbf{\omega}^{i}-\mathbf{\omega}_{\star})\sin(t)+\mathbf{v}^{i}\cos(t)$,
where the subscripts denote the parameter and velocity trajectory within the
deterministic segment. The event rate is the same as the BPS, and the starting
velocity for the next segment is updated upon events as,

$\mathbf{v}=\mathbf{v}-2\dfrac{\nabla
U(\mathbf{\omega})\cdot\mathbf{v}}{\left\lVert\Sigma_{\star}^{\frac{1}{2}}\nabla
U(\mathbf{\omega})\right\rVert^{2}}\Sigma_{\star}\nabla U(\mathbf{\omega}).$
(6)

###  2.4 Velocity Refreshment

All of the samplers introduced fail to target the posterior explicitly when
using the above dynamics alone. Introduction of a refreshment step rectifies
this, which is governed by a homogeneous PP
$\tau_{ref}\sim\lambda(\lambda_{ref})$. When $\tau_{ref}<\tau$, the velocity
is instead randomly sampled from the known reference distribution
$\Phi(\mathbf{v})$, and $\tau_{ref}$ is used for the update event time. For
BPS samplers in this work, we use a refreshment distribution of the form
$\mathcal{N}(0,\sigma^{2})$, where for the BPS $\sigma$ is a hyper-parameter
to be set and for for the $\sigma$BPS is used as the standard deviation learnt
during the warmup stage, and the Boomerang sampler requires
$\Phi(\mathbf{v})=\mathcal{N}(0,\Sigma_{\star})$. A summary of PDMP algorithms
for inference is described in Algorithm 1.

Result: Samples from posterior distribution

while  _Sampling_ do

// Simulate event time

// event times in this work simulated with Algorithm 2

$\tau\sim\text{PP}(\lambda(\mathbf{\omega},\mathbf{v}))$;

// Simulate time of refresh event

$\tau_{\text{ref}}\sim\text{PP}(\lambda_{\text{ref}})$;

$\tau^{i}=\text{min}(\tau,\tau_{\text{ref}})$;

// find end of current piecewise-deterministic segment, which will form start
of next segment

$\omega^{i+1}=\Psi(\mathbf{\omega},\mathbf{v},\tau^{i})_{\mathbf{\omega}}$;

if  _$\tau^{i}=\tau$_ then

// update according kernel

$\mathbf{v}^{i+1}=R(\mathbf{\omega}^{i+1},\mathbf{v}^{i})$;

else

// refresh velocity

$\mathbf{v}^{i+1}\sim\Phi(\mathbf{v})$;

end if

end while

Algorithm 1 Application of PDMP samplers for Inference

###  2.5 Problems with the event rate

With the deterministic dynamics illustrated in these samplers, the main
challenge in implementation of these methods is due to the sampling of the
event times. Analytic sampling from $\text{IPP}\big{(}\lambda(t)\big{)}$
requires being able to invert the integral of the event rate w.r.t. time,

$\Lambda(t)=\int_{0}^{\tau}\lambda(t)dt=\int_{0}^{\tau}\max\\{0,\mathbf{v}\cdot
A\nabla U(\mathbf{\omega}(\mathbf{v},t)\\}dt,$ (7)

where $A=\mathbf{I}$ for the BPS and Boomerang samplers. Inverting the above
integral is feasible only for simple models. A general case for sampling from
IPPs is available through thinning Lewis and Shedler [1979]. This requires
introducing an additional rate function that we can sample from $\mu(t)$ that
is also a strict upper bound on the event rate function of interest such that
$\mu(t)\geq\lambda(t)$ for all $t\geq 0$.

Figure 2: Example of the progression of the proposed envelope scheme used for
thinning. The blue line represents the true event rate, orange section depicts
the active regions for which we sample a new proposal time, and the red
section depicts previous segments in the envelope. Starting from the left, an
initial segment is found through interpolation between time points $t_{0}$ and
$t_{init}$. In the next segment, active regions of the envelope are found by
interpolating between the two prior points, which extends to create a new
segment to propose times. This process continues until a proposed time is
accepted from thinning.

The efficiency of any thinning scheme relies on the tightness of the upper
bound; the greater the difference between the upper bound and the true rate,
the more likely a proposed time will be rejected when sampling. Pakman et al.
[2017] propose a Bayesian linear regression method to generate an upper bound
suitable for thinning, though requires the calculation of variance within
gradients to formulate a suitable upper bound. They calculate this variance
empirically, which requires computing the gradient for each data point
individually within a mini-batch. This computation prohibits use for BNNs
where automatic differentiation software is used. Furthermore, the selection
of a suitable prior distribution for this is non-trivial for neural networks.
In the next section, we address this issue by instead introducing an
interpolation-based scheme for creating efficient and adaptive approximate
upper bounds that avoids excessive gradient computations and the need to
specify a prior on the approximate bound.

##  3 Adaptive Bounds for samplers

###  3.1 Sampling from IPPs with Linear Event Rates

Our goal is to create a piecewise-linear envelope $h(t)$ that will serve as an
approximate upper bound of our true event rate, where each segment in $h(t)$
is represented by $a_{i}t+b_{i}$. This envelope will serve as the event rate
for a proposal IPP that will be suitable for use with the thinning method of
Lewis and Shedler [1979]. Acceptance of an event time $t$ is given by,

$U\leq\frac{\lambda(t)}{h(t)},$ (8)

where $U\sim\text{Uniform}[0,1]$. We begin by building on the work of Klein
and Roberts [1984] to demonstrate how to sample times from an IPP with a
piecewise-linear event rate which we can use with thinning.

Within our proposal IPP with rate $h(t)$, we wish to generate the next event
time $t_{i}$ given the previous event $t_{i-1}$. The probability of events
occuring within the range of $[t_{i-1},t_{i}]$ is given by Devroye [2006],

$F(x)=1-\text{exp}\\{-(\Lambda(t_{i})-\Lambda(t_{i-1}))\\}.$ (9)

$t_{i}$ is distributed as,

$t_{i-1}+\Lambda^{-1}(U),$ (10)

where $U\sim\text{Uniform[0,1]}$. For linear segments, the solution to this
system can be written as Klein and Roberts [1984],

$\displaystyle t_{i}=$ $\displaystyle-b_{i}/a_{i}$
$\displaystyle+\sqrt{b_{i}^{2}+a_{i}^{2}t_{i-1}^{2}+2a_{i}b_{i}t_{i-1}-2a_{i}log(1-U)}\big{)}/a_{i}.$
(11)

This provides a framework for sampling from IPPs with a linear event rate. We
now describe how we create a piecewise-linear envelope for a proposal process
that can be used for thinning.

###  3.2 Piecewise Interpolation for Event Thinning

We begin by introducing a modified event rate for which we will form our
envelope,

$\hat{\lambda}(\mathbf{\omega}(t),\mathbf{v})=\max\\{0,\alpha\nabla
U(\mathbf{\omega})\cdot\mathbf{v}^{i}\\},$ (12)

where $\alpha\geq 1$ is a positive scaling factor to control the tightness of
the approximate bound on the rate. The use of $\hat{\lambda}$ for creating our
envelope is valid, since for values of $\alpha\geq 1$,
$\hat{\lambda}(t)\geq\lambda(t)$. The scaling factor included in this event
rate is designed to provide flexibility to end users with respect to
computational time and bias that will be introduced during inference. The
closer $\alpha$ is to one, the lower the probability for rejection of proposed
event times, but the greater the probability that the generated event rate
will not be a strict upper bound.

Our goal is to create a piecewise-linear upper bound suitable for proposing
event times using Equation 3.1. To achieve this we have two growing sets, one
for proposed event times $T=\\{t_{0},...,t_{n}\\}$ and the value of the
adjusted event rates at these times
$L=\\{\hat{\lambda}(t_{0}),\dots,\hat{\lambda}(t_{n})\\}$ for which we can
create a set of functions,

$h(t)=a_{i}t+b_{i},\hskip 14.22636ptt\geq t_{i}.$ (13)

The values for $a_{i}$ and $b_{i}$ are found by interpolating between the
points $(t_{i-1},\hat{\lambda}(t_{i-1}))$ and $(t_{i},\hat{\lambda}(t_{i}))$.

At the beginning of every deterministic PDMP segment, the sets $T$ and $L$
will be empty. To initialise the sets and create our first linear segment, we
evaluate the event rate at two points, $t_{0}=0$ and $t=t_{init}$, where
$t_{init}>t_{0}$. To evaluate the values for $a_{0}$ and $b_{0}$, we
interpolate between these two segments. Once the values for the first linear
segment are found, $t_{0}$ and $\hat{\lambda}(t_{0})$ are appended to their
corresponding sets, and $t_{init}$ and $\hat{\lambda}(t_{init})$ are
discarded. With this initial linear segment, we can propose a time $t_{i}$
through Equation 3.1. This proposed time is either accepted or rejected from
Equation 8.

If the proposed time is accepted, then the dynamics of the PDMP sampler are
updated at the given event time and the sets $T$ and $L$ are cleared, ready to
be re-initialised for the new dynamics. If the time is rejected, the proposed
time $t_{i}$ and envelope evaluation $\hat{\lambda}(t_{i})$ are appended to
their respective sets, and a new linear segment is calculated to interpolate
between this rejected proposal and the previous elements in the sets $T$ and
$L$. The rejected proposal time will serve as the new starting point
$(t_{i-1})$ for the new linear segment to propose the next time using Equation
3.1. This will continue until the proposed event time is accepted. This
process depicted visually in Figure 2 and summarised in Algorithm 2.

Within this work, we limit ourselves to models where the envelope provided by
$h(t)$ will only be an approximate upper bound, meaning bias will likely be
introduced during inference. Diagnosis and correction of this can be
identified through the acceptance ratio $\lambda(t)/h(t)$; if this value is
greater than one, the condition of $h(t)$ being a local upper bound is
violated. The amount of potential bias introduced can be mitigated by
increasing the scaling factor $\alpha$ in Equation 12 at the expense of
increasing computation load. We also introduce an additional rejection
condition such that if $\lambda(t)/h(t)>R$, then we will reject the the
proposed time. For this work, $R=2$. These properties are investigated in
Supp. Material A. In the following sections, we evaluate the proposed event
thinning scheme for BNNs to identify the suitability of different samplers for
inference in these challenging models, and how they can outperform other
stochastic approximation methods in terms of calibration, posterior
exploration, sampling efficiency and predictive performance.

Result: Proposed PDMP Event Time $\tau$

Initialize $T,L$;

Evaluate $(0,\lambda(0)),(t_{init},\lambda(t_{init}))$;

$i=1$;

Compute $a_{i},b_{i}$;

$T_{0}\leftarrow 0,L_{0}\leftarrow\lambda(t)$;

Discard $t_{init},\lambda(t_{init})$;

while  _not accepted_ do

// propose event time with $t_{i-1},a_{i}$ and $b_{i}$

$t_{i}\sim PP\Big{(}h(t)$);

$u\sim\text{Uniform}[0,1]$;

$r=\lambda(t_{i})/h(t)$;

// if valid proposal

if  _$u\leq r$ and $r<R$_ then

// sample is accepted

$\tau=t_{i}$;

accepted = True;

end if

else

// increment counter

$i+=1$;

$T_{i}\leftarrow t_{i},L_{i}\leftarrow\lambda(t_{i})$;

// update linear segment

$a_{i},b_{i}$ = update($L,T$);

end if

end while

Algorithm 2 Sampling event rate using proposed adaptive thinning method.

##  4 Related Work

The samplers used within this work require the use of an additional reference
process to provide velocity refreshments. The Generalised BPS Wu and Robert
[2017] is an updated variant of the BPS algorithm that incorporates a
stochastic update of the velocity which alleviates the need for a refreshment
process. Simulations have shown comparable performance to the BPS for simple
models and how it can reduce the need for fine-tuning the reference parameter
$\tau_{ref}$.

Another prominent sampler is the Zig-Zag Process (ZZP) Bierkens et al. [2019],
where at events the dynamics of a single parameter are updated. For the one-
dimensional case, this sampler represents the same process as the BPS. This
sampler has shown favourable results in terms of mixing performance and can
achieve ergodicity for certain models where the BPS cannot. A key
characteristic of this method is that each parameter is assigned an individual
event rate, making implementation for high-dimensional BNN models challenging.

Another class of algorithms designed for subsampling are discrete stochastic
MCMC methods Wenzel et al. [2020], Chen et al. [2014], Ma et al. [2015],
Welling and Teh [2011], Li et al. [2015]. These models have shown favourable
performance, with a recent variant achieving comparable predictive accuracy on
the ImageNet data set Heek and Kalchbrenner [2019]. Compared to algorithms
related to PDMPs, it has been shown that high variance related to naive
subsampling limits these methods to provide only an approximation to the
posterior Betancourt [2015]. The bias that is introduced due to subsampling
can be controlled by reducing the step-size for these methods at the expense
of mixing performance and posterior explorationNagapetyan et al. [2017],
Brosse et al. [2018], Teh et al. [2016]. We investigate the effect of this
property for SGLD and compare performance with PDMP samplers in the following
section.

##  5 Experiments

We now validate the performance of PDMPs using the proposed event sampling
method on a number of synthetic and real-world data sets for regression and
classification. To analyse performance for predictive tasks, the predictive
posterior needs to be evaluated. In this work, we discretise samples from the
trajectory to allow for Monte Carlo integration,

$p(y^{*}|x^{*},\mathcal{D})=\int\pi(\mathbf{\omega})p(y^{*}|\mathbf{\omega},x^{*})d\mathbf{\omega}\\\
\approx\dfrac{1}{N}\sum_{i=1}^{N}p(y^{*}|\mathbf{\omega}_{i},x^{*})\hskip
14.22636pt\mathbf{\omega}_{i}\sim\pi(\mathbf{\omega}),$ (14)

where parameter samples of $\mathbf{\omega}^{(i)}$ are taken from the values
encountered at event times. Experimentation is first conducted on synthetic
data sets to allow us to easily visualise predictive performance and
uncertainty in our models, followed by more difficult classification tasks
with Bayesian Convolutional Neural Networks (CNNs) on real data sets. For all
experimentation, we set our scaling factor from Equation 12 to $\alpha=1.0$ to
promote computational efficiency. To enable these experiments, we deliver a
Python package titled Tensorflow PDMP (TPDMP). This package utilises the
Tensorflow Probability library Dillon et al. [2017], allowing for hardware
acceleration and graph construction of all our models to accelerate
computation. We deliver kernels to implement the BPS, $\sigma$BPS , and
Boomerang sampler with our proposed event thinning scheme. Code is available
at https://github.com/egstatsml/tpdmp.git.

(a) BPS

(b) $\sigma$BPS

(c) Boomerang

(d) SGLD

(e) SGLD-ND

Figure 3: Examples of the different PDMP samplers using the proposed event
thinning procedure on synthetic regression task compared against SGLD with
decaying learning rate and constant learning rate (SGLD-ND).

###  5.1 Regression and Binary Classification with Synthetic Data

To visualise predictive performance and uncertainty estimation, regression and
binary classification tasks are formed on synthetic data sets. Description of
the networks used for these tasks is described in Supp. Material E. Before
sampling, a MAP estimate was first found using stochastic optimisation, and
was used to initialise each sampler. 2,000 samples were generated using each
sampling method, with each sampler initialised from the same MAP estimate. The
$\sigma$BPS requires an additional warmup period to identify suitable values
for the preconditioner. We achieve this by performing 1,000 initial samples
using the BPS, and standard deviation parameters used for the preconditioner
are estimated from these samples using the Welford algorithm Welford [1962].
These preconditioner values are then fixed throughout the sampling process.
The PDMP methods are compared against SGLD which starts with a learning rate
that decays to zero as required Welling and Teh [2011], Nagapetyan et al.
[2017], and with no decay of the learning rate as is commonly done in practice
(SGLD-ND). Examples of the predictive posterior distribution for regression
and binary classification are shown in Figures 3 and 4 respectively, with full
analysis in Supp. Material B. All PDMP models are fit with the proposed
adaptive event thinning procedure.

Results from these experiments affirm that inference from the PDMP models is
suitable for predictive reasoning, with low variance seen within the range of
observed data and greater variance as distance from observed samples
increases. We similarly see an increase in uncertainty along the decision
boundary, which is a desireable property. This is in contrast to SGLD, which
is able to provide low variance predictions for in-distribution data, but
offers less predictive uncertainty whilst extrapolating, for even in the case
for larger non-decreasing learning rates. This highlights the known
limitations of SGLD, that with a decaying learning rate it can fail to explore
the posterior, and with a larger non-decreasing learning rate will converge to
dynamics offered by traditional SGD Brosse et al. [2018], Nagapetyan et al.
[2017].

These tests indicate promising performance in terms of predictive accuracy and
uncertainty estimates. To further demonstrate classification performance, we
move to larger and more complicated models for performing regression and
classification on real-world data sets.

(a) BPS

(b) SGLD

Figure 4: Examples of the predictive mean and variance for synthetic
classification task. Left column illustrates results using the BPS and the
right using SGLD. We see increased uncertainty for the BPS outside the range
of observed data, whilst SGLD shows greater certainty.

####  5.1.1 UCI-Datasets

We further evaluate the performance of the PDMP samplers enabled by the
proposed event sampling scheme on datasets from the UCI repository Newman et
al. [1998]. In Table 1, we show performance metrics on the Boston houses
dataset, with the Naval, Energy, Yacht, and Concrete datasets evaluated in
Supp. Material E.2. Each model is fit with 1,000 samples. For these
experiments, we further include the naive Stochastic Gradient Hamiltonian
Monte Carlo (SGHMC) Chen et al. [2014]. Predictive performance of these models
is measured with Root Mean Squared Error (RMSE) and Negative Log-Likelihood
(NLL). Sampling efficiency is evaluated with Effective Sample Size (ESS)
Robert et al. [1999]. Due to the high dimension of our models, we perform PCA
on returned samples and project them onto the first principal component to
report ESS on the direction of greatest variance within samples.

From these results, we see SGLD and SGHMC provide a slight improvement in
terms of RMSE, though we see that the Boomerang Sampler considerably
outperforms both of these methods in terms of sample efficiency. This result
follows from the previous sections where we see that SGLD frequently converges
to the SGD solution space, whilst the PDMP samplers can explore the posterior
space. Additional results in Supp. Material E.2 further validate these
results, and show how predictive performance can be improved with PDMP
samplers.

Table 1: Summary of predictive performance using PDMP samplers with the
proposed event time sampling methods on the Boston Houses dataset. Negative
log-likelihood (NLL) and Root Mean Squared Error (RMSE) are reported.
Effective sample size (ESS) is measured over the first principal component of
samples. Results are shown over 5 independent runs with standard deviations
reported.

Inference | NLL $\downarrow$ | RMSE $\downarrow$ | ESS $\uparrow$
---|---|---|---
BPS | 0.96 $\pm$ 0.00 |  2.76 $\pm$ 0.11 |  2.71 $\pm$ 0.02
$\sigma$BPS | 0.96 $\pm$ 0.01 |  2.73 $\pm$ 0.12 |  2.70 $\pm$ 0.01
Boomerang | 0.96 $\pm$ 0.01 |  2.77 $\pm$ 0.13 | 901.07 $\pm$ 105.03
SGLD | 0.96 $\pm$ 0.00 | 2.68 $\pm$ 0.01 |  2.89 $\pm$ 0.00
SGHMC | 0.96 $\pm$ 0.00 |  2.69 $\pm$ 0.05 |  2.71 $\pm$ 0.00

###  5.2 Multi-Class Classification

We now evaluate the performance of the proposed sampling procedures on the
popular MNIST LeCun et al. [1998], Fashion MNIST Xiao et al. [2017], SVHN
Netzer et al. [2011], CIFAR-10 and CIFAR-100 Krizhevsky and Hinton [2009] data
sets using CNNs. For MNIST and Fashion-MNIST, the LeNet5 architecture was used
whilst for SVHN, CIFAR-10, and CIFAR-100 the modified ResNet20 architecture
from Wenzel et al. [2020] was used. For the Boomerang Sampler, we follow the
original Bierkens et al. [2020] and implement a flat prior, and for other
samplers a normal prior is used.

Table 2: Summary of predictive performance using PDMP samplers with the
proposed event time sampling methods. Negative log-likelihood (NLL) is
reported, along with calibration measured using the expected calibration error
(ECE) Guo et al. [2017]. Effective sample size (ESS) is measured over the
first principal component of samples. Mean and standard deviation in results
presented over 5 independent runs.

Inference | ACC $\uparrow$ | NLL $\downarrow$ | ECE $\downarrow$ | ESS $\uparrow$
---|---|---|---|---
MNIST
BPS | 0.99 $\pm$ 0.00 |  0.08 $\pm$ 0.04 |  5.22 $\pm$ 3.44 |  2.88 $\pm$ 0.07
$\sigma$BPS | 0.99 $\pm$ 0.00 | 0.02 $\pm$ 0.00 |  0.48 $\pm$ 0.22 |  2.72 $\pm$ 0.02
Boomerang | 0.99 $\pm$ 0.00 |  0.03 $\pm$ 0.00 |  0.85 $\pm$ 0.37 | 175.26 $\pm$ 29.38
SGLD | 0.99 $\pm$ 0.00 | 0.02 $\pm$ 0.00 |  0.19 $\pm$ 0.02 |  22.74 $\pm$ 0.18
SGLD-ND | 0.99 $\pm$ 0.00 | 0.02 $\pm$ 0.00 | 0.15 $\pm$ 0.01 |  3.00 $\pm$ 0.00
SGHMC | 0.99 $\pm$ 0.00 | 0.02 $\pm$ 0.00 |  0.17 $\pm$ 0.02 |  2.71 $\pm$ 0.00
Fashon-MNIST
BPS | 0.91 $\pm$ 0.00 | 0.25 $\pm$ 0.00 | 1.00 $\pm$ 0.38 |  2.73 $\pm$ 0.04
$\sigma$BPS | 0.91 $\pm$ 0.00 |  0.30 $\pm$ 0.01 |  3.94 $\pm$ 0.30 |  2.74 $\pm$ 0.02
Boomerang | 0.91 $\pm$ 0.00 |  0.28 $\pm$ 0.01 |  3.82 $\pm$ 2.24 | 143.05 $\pm$ 31.32
SGLD |  0.90 $\pm$ 0.01 |  0.30 $\pm$ 0.01 |  3.09 $\pm$ 1.49 |  102.52 $\pm$ 84.14
SGLD-ND | 0.91 $\pm$ 0.00 |  0.33 $\pm$ 0.00 |  4.71 $\pm$ 0.01 |  2.97 $\pm$ 0.00
SGHMC | 0.91 $\pm$ 0.00 |  0.34 $\pm$ 0.01 |  4.88 $\pm$ 0.13 |  2.71 $\pm$ 0.00
SVHN
BPS |  0.95 $\pm$ 0.00 |  0.18 $\pm$ 0.00 | 0.61 $\pm$ 0.09 |  2.71 $\pm$ 0.02
$\sigma$BPS |  0.95 $\pm$ 0.00 |  0.19 $\pm$ 0.01 |  0.63 $\pm$ 0.18 |  2.75 $\pm$ 0.02
Boomerang |  0.95 $\pm$ 0.00 |  0.21 $\pm$ 0.01 |  3.58 $\pm$ 0.82 | 172.02 $\pm$ 40.73
SGLD | 0.96 $\pm$ 0.00 | 0.16 $\pm$ 0.00 |  1.33 $\pm$ 0.03 |  21.84 $\pm$ 0.25
SGLD-ND | 0.96 $\pm$ 0.00 | 0.16 $\pm$ 0.00 |  1.34 $\pm$ 0.00 |  2.97 $\pm$ 0.00
SGHMC |  0.95 $\pm$ 0.00 |  0.19 $\pm$ 0.00 |  0.83 $\pm$ 0.27 |  2.71 $\pm$ 0.00
CIFAR-10
BPS |  0.80 $\pm$ 0.00 |  0.67 $\pm$ 0.00 |  4.91 $\pm$ 0.65 |  2.70 $\pm$ 0.02
$\sigma$BPS |  0.79 $\pm$ 0.01 |  0.72 $\pm$ 0.01 |  8.55 $\pm$ 0.64 |  2.73 $\pm$ 0.02
Boomerang | 0.81 $\pm$ 0.00 | 0.59 $\pm$ 0.02 | 4.17 $\pm$ 1.55 | 200.00 $\pm$ 0.00
SGLD | 0.81 $\pm$ 0.00 |  1.01 $\pm$ 0.00 |  15.32 $\pm$ 0.04 |  19.78 $\pm$ 0.01
SGLD-ND | 0.81 $\pm$ 0.00 |  0.99 $\pm$ 0.00 |  15.15 $\pm$ 0.01 |  2.85 $\pm$ 0.00
SGHMC |  0.80 $\pm$ 0.00 |  0.96 $\pm$ 0.00 |  14.64 $\pm$ 0.15 |  2.71 $\pm$ 0.00
CIFAR-100
BPS |  0.63 $\pm$ 0.00 |  1.43 $\pm$ 0.01 |  10.16 $\pm$ 0.21 |  2.72 $\pm$ 0.02
$\sigma$BPS |  0.63 $\pm$ 0.00 | 1.42 $\pm$ 0.00 | 8.70 $\pm$ 0.00 |  2.80 $\pm$ 0.00
Boomerang |  0.63 $\pm$ 0.01 | 1.42 $\pm$ 0.08 |  12.67 $\pm$ 3.26 | 54.96 $\pm$ 18.20
SGLD | 0.64 $\pm$ 0.00 | 1.42 $\pm$ 0.00 |  11.24 $\pm$ 0.32 |  26.79 $\pm$ 0.73
SGLD-ND | 0.64 $\pm$ 0.00 |  1.45 $\pm$ 0.00 |  12.22 $\pm$ 0.05 |  2.69 $\pm$ 0.00
SGHMC |  0.63 $\pm$ 0.00 |  1.46 $\pm$ 0.00 |  12.63 $\pm$ 0.12 |  2.71 $\pm$ 0.00

Similar to the experiments on regression, a MAP estimate is found and used to
initialise each sampler. 2,000 samples for each model are then generated,
though a thinning factor of 10 is used to reduce the number of returned
samples used for prediction to 200. For these models, we measure predictive
performance and calibration through the Accuracy, NLL, and Expected
Calibration Error (ECE) Guo et al. [2017], and similarly measure sampling
efficiency using the ESS with on samples after performing PCA. A full
description of the models used, and experiment parameters is shown in Supp.
Material E.3. Table 2 summarises the results of these experiments.

(a) ACF first principal component

(b) Trace first principal component

Figure 5: Example of ACF and trace plots for the first principal component of
the samples from network fit on SVHN dataset.

(a) Image Class 0

(b) Boomerang Sampler

(c) SGLD

Figure 6: Examples from predictive posterior for difficult-to-classify samples
from SVHN. Top row shows the original image and the bottom row shows the
predictive distribution for the Boomerang sampler and SGLD. The mean for each
class is represented by the dot, and the 95% credible intervals shown with the
error bars.

These results highlight favourable performance for PDMP samplers. The BPS and
$\sigma$BPS samplers frequently provide more calibrated predictions in terms
of ECE, whilst maintaining comparable predictive accuracy to SGLD methods.
Most importantly, we note the Boomerang sampler consistently outperforms other
samplers in terms of effective sample size, whilst also promoting competitive
predictive accuracy and calibration. This highlights the potential for these
samplers for probabilistic inference within neural networks.

With measures of predictive performance and ESS within our models, we wish to
further investigate the mixing properties of the samplers presented within to
identify how well the posterior space is being explored. ESS only gives a
measure to approximate the number of independent samples within our MCMC
chain, though we are also interested in how well the support for the posterior
is being explored. Given the large number of parameters seen within a BNN, it
is infeasible to evaluate the coordinate trace and autocorrelation plots for
individual parameters as is typically done for MCMC models. Instead, we again
perform PCA to reduce the dimension of our data and investigate the trace
plots of the first principal component as illustrated in Figure. 5. From these
figures, we can identify strong correlation between samples from the BPS,
$\sigma$BPS , SGHMC, and SGLD-ND solutions. SGLD offers reduced correlation in
samples, however as seen in the trace plot, samples fail to explore the
posterior and instead converge to a steady state, whilst the Boomerang sampler
provides considerably reduced correlation and more favourable mixing, however
we note in Supp. Material C that the Boomerang sampler can exhibit a mode
seeking behaviour due to the nature of the reference measure. Convergence of
the SGLD samples can be attributed to the reduction in learning rate required
to target the posterior. We verify this result in Supp. Material C, where we
provide further analysis into results from all networks and remaining
principal components. The effect of this convergence in terms of predictive
uncertainty is illustrated within Figure 6, where the PDMP sampler is able to
provide more meaningful uncertainty estimates for difficult-to-classify
samples, and the SGLD predictive results converge to that similar of a point
estimate. Additional examples of the predictive distributions is shown in
Supp. Material H.

Probabilistic methods have shown favourable performance in terms of Out of
Distribution (OOD) detection Grathwohl et al. [2019], Maddox et al. [2019].
Given the point-estimate-like nature of the results returned by SGLD, we wish
to compare with results from the Boomerang sampler to see if both can offer
similar performance for OOD data. We see in Figure 7 that the Boomerang
sampler offers greater entropy for OOD data, indicating a desireable increase
in aleatoric uncertainty. Additional analysis is provided in Supp. Material G.

(a) In-distribution

(b) OOD

Figure 7: Entropy histograms comparing SGLD and Boomerang sampler fit on the
CIFAR-10 dataset. OOD data represented by SVHN. We see the predictive entropy
from the Boomerang sampler increases as desired for OOD data, whilst SGLD
remains overly confident for erroneous samples.

##  6 Discussion and Limitations

Whilst the PDMP methods have shown favourable performance in terms of
predictive accuracy, calibration and uncertainty in BNNs, there are certain
challenges with fitting them. The PDMP samplers used within this work are
designed to target the joint distribution,

$p(\mathbf{\omega},\mathbf{v})=\pi(\mathbf{\omega})\Phi(\mathbf{v})$ (15)

where $\pi(\mathbf{\omega})$ is the target posterior and $\Phi(\mathbf{v})$ is
the distribution of the auxiliary velocity components which must be set by
users in the form of the refreshment distribution. For the BPS and $\sigma$BPS
samplers, it has been shown that with a reference distribution may be a
Gaussian or restricted to the unit hypersphere Bouchard-Côté et al. [2018],
DURMUS et al. [2020]. For the Boomerang sampler, the velocity distribution is
designed with respect to a reference measure to ensure invariance to the
target distribution, such that
$\Phi(\mathbf{v})=\mathcal{N}(0,\Sigma_{\star})$, where $\Sigma_{\star}$ is
the same factor used to precondition the dynamics. The choice in distribution
used for the velocity component has an explicit effect on the mixing
capabilities of the models when applied to BNNs. We demonstrate this effect in
Supp. Material D.1. We find that a velocity distribution with too much
variance can cause effects similar to that of divergences seen in HMC and
NUTS. Furthermore, we see that with variance set too low, the samplers can
fail to explore the posterior sufficiently to provide the desired meaningful
uncertainty estimates. A similar effect can be seen for the choice of
refreshment rate, which we investigate in Supp. Material D.2. We highlight
these limitations as areas for future research to enable robust application of
PDMP methods for BNNs.

The Boomerang sampler as implemented within this work and the original paper
is probabilistic, though is not purely Bayesian. This is due to the reference
measure being identified through the data itself. A strictly Bayesian approach
can be recovered by setting the reference measure and associated
preconditioner matrix from a prior distribution, though we would lose some
favourable sampling performance offered by this sampler. We can view the
approach implemented within similar to an empirical Bayes, where we are
gleaning information about the reference measure for the Boomerang sampler
from the data itself. Given the difficulty of specifying a meaningful and
informative prior, and the success seen when using empirical priors for BNNs
Krishnan et al. [2020], we believe the use of such an approach for the
Boomerang sampler is justified.

##  7 Conclusion

Within this work, we demonstrate how PDMPs can be used for BNNs. We provide a
flexible piecewise linear bound to enable sampling of event times within these
frameworks that permits inference in BNNs. A GPU-accelerated software package
is offered to increase the availability of PDMPs for a wide array of models.
Experimentation on BNNs for regression and classification indicates comparable
or improved predictive performance and calibration, though were able to
consistently improve sampling efficiency and uncertainty estimation when
compared against existing stochastic inference methods.

##  8 Corrigendum

The published version of this paper contained a missing term in Equation 3.1
that has is corrected. Furthermore an error in the software for this paper was
identified where the scaling term to correct for mini-batch sizes for the
unbiased gradient estimates was not being called. Experiments have been rerun
to with software corrected and additional information provided in the
supplementary material highlighting how samplers can be explore the posterior
space for BNNs.

## References

  * Barber and Bishop [1998] David Barber and Christopher M Bishop.  Ensemble learning in bayesian neural networks.  _Nato ASI Series F Computer and Systems Sciences_ , 168:215–238, 1998.
  * Bertazzi and Bierkens [2020] Andrea Bertazzi and Joris Bierkens.  Adaptive schemes for piecewise deterministic monte carlo algorithms.  _arXiv preprint arXiv:2012.13924_ , 2020.
  * Betancourt [2015] Michael Betancourt.  The fundamental incompatibility of scalable hamiltonian monte carlo and naive data subsampling.  In _International Conference on Machine Learning_ , ICML’15, pages 533–540. JMLR.org, 2015.
  * Bierkens et al. [2019] Joris Bierkens, Paul Fearnhead, Gareth Roberts, et al.  The zig-zag process and super-efficient sampling for bayesian analysis of big data.  _The Annals of Statistics_ , 47(3):1288–1320, 2019.
  * Bierkens et al. [2020] Joris Bierkens, Sebastiano Grazzi, Kengo Kamatani, and Gareth Roberts.  The boomerang sampler.  _arXiv preprint arXiv:2006.13777_ , 2020.
  * Blei et al. [2017] David M Blei, Alp Kucukelbir, and Jon D McAuliffe.  Variational inference: A review for statisticians.  _Journal of the American statistical Association_ , 112(518):859–877, 2017.
  * Bouchard-Côté et al. [2018] Alexandre Bouchard-Côté, Sebastian J Vollmer, and Arnaud Doucet.  The bouncy particle sampler: A nonreversible rejection-free markov chain monte carlo method.  _Journal of the American Statistical Association_ , 113(522):855–867, 2018.
  * Brosse et al. [2018] Nicolas Brosse, Alain Durmus, and Eric Moulines.  The promises and pitfalls of stochastic gradient langevin dynamics.  In _Advances in Neural Information Processing Systems_ , pages 8268–8278, 2018.
  * Chen et al. [2014] Tianqi Chen, Emily Fox, and Carlos Guestrin.  Stochastic gradient hamiltonian monte carlo.  In _International conference on machine learning_ , pages 1683–1691, 2014.
  * Davis [1984] Mark HA Davis.  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, 1984.
  * Davis [1993] Mark HA Davis.  _Markov models & optimization_, volume 49.  CRC Press, 1993.
  * Devroye [2006] Luc Devroye.  Nonuniform random variate generation.  _Handbooks in operations research and management science_ , 13:83–121, 2006.
  * Dillon et al. [2017] Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous.  Tensorflow distributions.  _arXiv preprint arXiv:1711.10604_ , abs/1711.10604, 2017.
  * DURMUS et al. [2020] ALAIN DURMUS, ARNAUD GUILLIN, and PIERRE MONMARCHÉ.  Geometric ergodicity of the bouncy particle sampler.  _The Annals of Applied Probability_ , 30(5):2069–2098, 2020.
  * Fearnhead et al. [2018] Paul Fearnhead, Joris Bierkens, Murray Pollock, Gareth O Roberts, et al.  Piecewise deterministic markov processes for continuous-time monte carlo.  _Statistical Science_ , 33(3):386–412, 2018.
  * Grathwohl et al. [2019] Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky.  Your classifier is secretly an energy based model and you should treat it like one.  In _International Conference on Learning Representations_ , 2019.
  * Guo et al. [2017] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In _International conference on machine learning_ , pages 1321–1330. PMLR, 2017.
  * Heek and Kalchbrenner [2019] Jonathan Heek and Nal Kalchbrenner.  Bayesian inference for large scale image classification.  _arXiv preprint arXiv:1908.03491_ , 2019.
  * Izmailov et al. [2021] Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Wilson.  What are bayesian neural network posteriors really like?  _arXiv preprint arXiv:2104.14421_ , 2021.
  * Jordan et al. [1999] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul.  An introduction to variational methods for graphical models.  _Machine learning_ , 37(2):183–233, 1999.
  * Klein and Roberts [1984] Robert W. Klein and Stephen D. Roberts.  A time-varying poisson arrival process generator.  _Simulation_ , 43:193–195, 1984.
  * Krishnan et al. [2020] Ranganath Krishnan, Mahesh Subedar, and Omesh Tickoo.  Specifying weight priors in bayesian deep neural networks with empirical bayes.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 4477–4484, 2020.  URL https://ojs.aaai.org/index.php/AAAI/article/view/5875.
  * Krizhevsky and Hinton [2009] Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  2009\.
  * LeCun et al. [1998] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Lewis and Shedler [1979] PA W Lewis and Gerald S Shedler.  Simulation of nonhomogeneous poisson processes by thinning.  _Naval research logistics quarterly_ , 26(3):403–413, 1979.
  * Li et al. [2015] Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin.  Preconditioned stochastic gradient langevin dynamics for deep neural networks.  _arXiv preprint arXiv:1512.07666_ , 2015.
  * Ma et al. [2015] Yi-An Ma, Tianqi Chen, and Emily Fox.  A complete recipe for stochastic gradient mcmc.  In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett, editors, _Advances in Neural Information Processing Systems 28_ , pages 2917–2925. Curran Associates, Inc., 2015.  URL http://papers.nips.cc/paper/5891-a-complete-recipe-for-stochastic-gradient-mcmc.pdf.
  * MacKay [1995] David JC MacKay.  Probable networks and plausible predictions–a review of practical bayesian methods for supervised neural networks.  _Network: computation in neural systems_ , 6(3):469–505, 1995.
  * Maddox et al. [2019] Wesley J. Maddox, Pavel Izmailov, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson.  A simple baseline for bayesian uncertainty in deep learning.  In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett, editors, _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pages 13132–13143, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/118921efba23fc329e6560b27861f0c2-Abstract.html.
  * Mandt et al. [2017] Stephan Mandt, Matthew D Hoffman, and David M Blei.  Stochastic gradient descent as approximate bayesian inference.  _Journal of Machine Learning Research_ , 18:1–35, 2017\.
  * Nagapetyan et al. [2017] Tigran Nagapetyan, Andrew B Duncan, Leonard Hasenclever, Sebastian J Vollmer, Lukasz Szpruch, and Konstantinos Zygalakis.  The true cost of stochastic gradient langevin dynamics.  _arXiv preprint arXiv:1706.02692_ , 2017.
  * Neal [2012] Radford M Neal.  _Bayesian learning for neural networks_ , volume 118.  Springer Science & Business Media, 2012.
  * Netzer et al. [2011] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  In _NIPS workshop on deep learning and unsupervised feature learning_ , volume 2011, page 5, 2011.
  * Newman et al. [1998] D.J. Newman, S. Hettich, C.L. Blake, and C.J. Merz.  Uci repository of machine learning databases, 1998.  URL http://www.ics.uci.edu/~mlearn/MLRepository.html.
  * Pakman et al. [2017] Ari Pakman, Dar Gilboa, David Carlson, and Liam Paninski.  Stochastic bouncy particle sampler.  In _International Conference on Machine Learning_ , pages 2741–2750, 2017.
  * Robert et al. [1999] Christian P Robert, George Casella, and George Casella.  _Monte Carlo statistical methods_ , volume 2.  Springer, 1999.
  * Teh et al. [2016] Yee Whye Teh, Alexandre H Thiery, and Sebastian J Vollmer.  Consistency and fluctuations for stochastic gradient langevin dynamics.  _Journal of Machine Learning Research_ , 17, 2016.
  * Vanetti et al. [2017] Paul Vanetti, Alexandre Bouchard-Côté, George Deligiannidis, and Arnaud Doucet.  Piecewise-deterministic markov chain monte carlo.  _arXiv preprint arXiv:1707.05296_ , 2017.
  * Wainwright and Jordan [2008] Martin J Wainwright and Michael Irwin Jordan.  _Graphical models, exponential families, and variational inference_ , volume 1.  Now Publishers Inc, 2008.
  * Welford [1962] BP Welford.  Note on a method for calculating corrected sums of squares and products.  _Technometrics_ , 4(3):419–420, 1962.
  * Welling and Teh [2011] Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th international conference on machine learning (ICML-11)_ , pages 681–688, 2011.
  * Wenzel et al. [2020] Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin.  How good is the bayes posterior in deep neural networks really?  In _International Conference of Machine Learning_ , 2020.
  * Wu and Robert [2017] Changye Wu and Christian P Robert.  Generalized bouncy particle sampler.  _arXiv preprint arXiv:1706.04781_ , art. arXiv:1706.04781, Jun 2017\.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  _arXiv preprint arXiv:1708.07747_ , 2017.