# Fearless Stochasticity in Expectation Propagation

Jonathan So
University of Cambridge
<EMAIL_ADDRESS>
&Richard E. Turner
University of Cambridge
<EMAIL_ADDRESS>

###### Abstract

Expectation propagation (EP) is a family of algorithms for performing
approximate inference in probabilistic models. The updates of EP involve the
evaluation of moments—expectations of certain functions—which can be estimated
from Monte Carlo (MC) samples. However, the updates are not robust to MC noise
when performed naively, and various prior works have attempted to address this
issue in different ways. In this work, we provide a novel perspective on the
moment-matching updates of EP; namely, that they perform natural-gradient-
based optimisation of a variational objective. We use this insight to motivate
two new EP variants, with updates that are particularly well-suited to MC
estimation; they remain stable and are most sample-efficient when estimated
with just a single sample. These new variants combine the benefits of their
predecessors and address key weaknesses. In particular, they are easier to
tune, offer an improved speed-accuracy trade-off, and do not rely on the use
of debiasing estimators. We demonstrate their efficacy on a variety of
probabilistic inference tasks.

##  1 Introduction

Expectation propagation (EP) [33, 36] is a family of algorithms that is
primarily used for performing approximate inference in probabilistic models
[7, 8, 9, 10, 13, 15, 17, 16, 18, 19, 20, 21, 25, 26, 27, 28, 31, 32, 34, 35,
37, 40, 43, 44, 46, 47, 48], although it can be used more generally for
approximating certain kinds of functions and their integrals [11].

EP involves the evaluation of moments—expectations of certain functions—under
distributions that are derived from the model of interest. EP is usually
applied to models for which these moments have convenient closed-form
expressions or can be accurately estimated using deterministic methods.
Moments can also be estimated using Monte Carlo (MC) samples, significantly
expanding the set of models EP can be applied to; however, the updates of EP
are not robust to MC noise when performed naively, and various prior works
have attempted to address this issue in different ways [14, 44, 47].

In this work we provide a novel perspective on the moment-matching updates of
EP; namely, that they perform natural-gradient-based optimization of a
variational objective (Section 3). We use this insight to motivate two new EP
variants, EP-$\eta$ (Section 3.2) and EP-$\mu$ (Section 3.3), with updates
that are particularly well-suited to MC estimation, remaining stable and being
most sample-efficient when estimated with just a single sample. These new
variants combine the benefits of their predecessors and address key
weaknesses. In particular, they are easier to tune, offer an improved speed-
accuracy trade-off, and do not rely on the use of debiasing estimators. We
demonstrate their efficacy on a variety of probabilistic inference tasks
(Section 4).

##  2 Background

In this section, we first introduce the problem setting. We then give an
overview of EP, followed by a discussion of issues related to sampled
estimation of EP updates.

Let $\mathcal{F}$ be the tractable, minimal exponential family of
distributions (see Appendix A), defined by the statistic function $s(.)$ with
respect to base measure $\nu(.)$. Let $\Omega$, $\mathcal{M}$, $A(.)$ and
$A^{*}(.)$ denote the natural domain, mean domain, log partition function, and
negative entropy function of $\mathcal{F}$, respectively.

Let $p_{0}$ be the member of $\mathcal{F}$ with natural parameter
$\theta_{0}$, so that
$p_{0}(z)=\exp\bm{(}\theta_{0}^{\top}s(z)-A(\theta_{0})\bm{)}$111We use e.g.
$p$ to refer to a distribution, and $p(.)$ or $p(z)$ for its density,
throughout., and assume that a distribution of interest $p^{*}$, the target
distribution, has a density of the form

$\displaystyle p^{*}(z)$ $\displaystyle\propto
p_{0}(z)\prod_{i=1}^{m}\exp\big(\ell_{i}(z)\big{missing}).$ (1)

In Bayesian inference settings, $p^{*}$ would be the posterior distribution
over parameters $z$ given some observed data $\mathcal{D}$, where $p_{0}$ may
correspond to a prior distribution, and $\\{\ell_{i}(z)\\}_{i}$ to log-
likelihood terms given some partition of $\mathcal{D}$.222Going forward, we
assume $i$ and $j$ to be taken from the index set $\\{1,\ldots,m\\}$ unless
otherwise specified. However, we consider the more general setting in which
the target distribution is simply a product of factors; note that we can
assume form (1) without loss of generality if we allow $p_{0}$ to be improper.
The inference problem typically amounts to evaluating quantities derived from
the normalised density $p^{*}(z)$, such as samples, summary statistics, or
expectations of given functions. In some special cases, these quantities can
be computed exactly, but this is not feasible in general, and approximations
must be employed.

###  2.1 Expectation propagation (EP)

Given a target distribution density of form (1), EP aims to find an
approximation $p\in\mathcal{F}$ such that

$\displaystyle p(z)$ $\displaystyle\propto
p_{0}(z)\prod_{i}\exp\big(\lambda_{i}^{\top}s(z)\big{missing})\approx
p^{*}(z).$ (2)

By assumption $\mathcal{F}$ is tractable, and so provided that
$(\theta_{0}+\sum_{i=1}^{m}\lambda_{i})\in\Omega$, $p$ is a tractable member
of $\mathcal{F}$. Each factor $\exp\bm{(}\lambda_{i}^{\top}s(z)\bm{)}$ is
known as a site potential, and can be roughly interpreted as a
$\mathcal{F}$-approximation to the $i$-th target factor,
$\exp\bm{(}\ell_{i}(z)\bm{)}$.333This interpretation is not precise, since it
is not necessarily the case that $\lambda_{i}\in\Omega$. $\lambda_{i}$ is
known as the $i$-th site parameter.

#### Variational problem

EP, and several of its variants [14, 21, 33], can be viewed as solving a
variational problem which we now introduce following the exposition of
Hasenclever et al. [14].

Let the $i$-th locally extended family, denoted $\mathcal{F}_{i}$, be the
exponential family defined by the statistic function
$s_{i}(z)=(s(z),\ell_{i}(z))$ with respect to base measure $\nu(.)$. Let
$\Theta_{i},\mathcal{M}_{i}$ and $A_{i}(.)$ denote the natural domain, mean
domain, and log partition function of $\mathcal{F}_{i}$, respectively. Note
that a member of $\mathcal{F}_{i}$ roughly corresponds to a distribution whose
density is the (normalised) product of a member of $\mathcal{F}$ with the
$i$-th target factor raised to a power. Unlike $\mathcal{F}$, we do not assume
$\mathcal{F}_{i}$ is minimal.

Fixed points of EP correspond to the solutions of the saddle-point problem

$\displaystyle\max_{\\{\theta_{i}\in\Omega\\}_{i}}\min_{\\{\lambda_{i}\\}_{i}}L(\theta_{1},\ldots,\theta_{m},\lambda_{1},\ldots,\lambda_{m}),\text{
where }$ $\displaystyle
L(\theta_{1},\ldots,\theta_{m},\lambda_{1},\ldots,\lambda_{m})=A\Big{(}\theta_{0}+\sum_{i}\lambda_{i}\Big{)}+\sum_{i}\beta_{i}\big{[}A_{i}\bm{(}(\theta_{i}-\beta_{i}^{-1}\lambda_{i},\beta_{i}^{-1})\bm{)}-A(\theta_{i})\big{]}.$
(3)

At a solution to (3), the EP approximation is given by (2). The
hyperparameters $\\{\beta_{i}\\}_{i}$ control the characteristics of the
approximation, and correspond to the power parameters of power EP.

#### EP updates

EP [33, 36], power EP [30] and double-loop EP [21], can all be viewed as
alternating between some number ($\geq 1$) of _inner updates_ to decrease $L$
with respect to $\\{\lambda_{i}\\}_{i}$, with an _outer update_ to increase
$L$ with respect to $\\{\theta_{i}\\}_{i}$ [14, 22]. EP is not typically
presented in this way, but by doing so we will be able to present a unified
algorithm that succinctly illustrates the relationship between the different
variants. The inner and outer updates are given by

Inner update:
$\displaystyle\lambda_{i}\leftarrow\lambda_{i}-\alpha\Big{(}\theta_{0}+\sum_{j}\lambda_{j}-\nabla
A^{*}\bm{(}\mathbb{E}_{p_{i}(z)}[s(z)]\bm{)}\Big{)},$ (4) Outer update:
$\displaystyle\theta_{i}\leftarrow\theta_{0}+\sum_{j}\lambda_{j},$ (5)

where $p_{i}\in\mathcal{F}_{i}$ denotes the $i$-th tilted distribution, with
density

$\displaystyle p_{i}(z)$ $\displaystyle\propto
p_{0}(z)\exp\Big((\theta_{i}-\beta^{-1}\lambda_{i})^{\top}s(z)+\beta^{-1}\ell_{i}(z)\Big{missing}).$
(6)

The hyperparameter $\alpha$ controls the level of damping, which is used to
aid convergence. An undamped inner update (with $\alpha=1$) can be seen as
performing _moment matching_ between $p$ and $p_{i}$. The inner updates can be
performed either serially (over $i$) or in parallel. In this work we assume
they are applied in parallel, but the ideas presented can easily be extended
to the serial case. Heskes and Zoeter [22] showed that (4) follows a decrease
direction in $L$ with respect to $\lambda_{i}$, and so is guaranteed to
decrease $L$ when $\alpha$ is small enough. See Appendix B for a derivation of
the inner update. When $\\{\lambda_{i}\\}_{i}$ are at a partial minimum of $L$
(for fixed $\\{\theta_{i}\\}_{i}$), the outer update performs an exact partial
maximisation of the primal form of the variational problem; see Hasenclever et
al. [14] for details.

#### Unified EP algorithm

The double-loop EP algorithm of Heskes and Zoeter [21] repeats the inner
update to convergence before performing each outer update, which ensures
convergence of the overall procedure. The usual presentation of EP combines
(4) and (5) into a single update in $\lambda_{i}$; however, Jylänki et al.
[27] observed that it can also be viewed as performing double-loop EP with
just a single inner step per outer update (which is not guaranteed to converge
in general). By taking this viewpoint, we are we are able to present EP, power
EP, and their double-loop counterparts as a single algorithm, presented in
Algorithm 1. Note that the dependence on
$\\{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\beta_{i}}\\}_{i}$
comes through the definition of $p_{i}(.)$.

Algorithm 1 EP ($\beta_{i}$=$1$,$n_{\text{inner}}$=$1$), power EP
($\beta_{i}\neq 1$,$n_{\text{inner}}$=$1$), and their double-loop variants
($n_{\text{inner}}$>$1$)

$\mathcal{F}$, $\theta_{0}$,
$\\{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\beta_{i}}\\}_{i}$,
$\\{\ell_{i}(z)\\}_{i}$, $\\{\lambda_{i}\\}_{i}$,
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}n_{\text{inner}}}$,
$\alpha$

while not converged do

$\theta_{i}\leftarrow\theta_{0}+\sum_{j}\lambda_{j}$ for $i=1$ to $m$

for $1$ to
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}n_{\text{inner}}}$
do

for $i=1$ to $m$ _in parallel_ do

$\lambda_{i}\leftarrow\lambda_{i}-\alpha\Big{(}\theta_{0}+\sum_{j}\lambda_{j}-\nabla
A^{*}\big{(}\mathbb{E}_{p_{i}(z)}[s(z)]\big{)}\Big{)}$

return $\\{\lambda_{i}\\}_{i}$

We can think of "exact" double-loop EP as corresponding to
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}n_{\text{inner}}}=\infty$,
with the inner loop exiting only once some convergence criteria has been
satisfied. In practice, a truncated inner loop is often used by setting
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}n_{\text{inner}}}$
to some small number [14, 27]. Upon convergence of Algorithm 1, the
approximation $p(z)\approx p^{*}(z)$ is given by (2). Going forwards we will
refer to the family of algorithms encompassed by Algorithm 1 simply as EP.

#### Stochastic moment estimation

EP is typically applied to models for which the updates either have closed-
form expressions, or can be accurately estimated using deterministic numerical
methods. However, as update (4) only depends on the target distribution
through expectations under the tilted distributions, this suggests that
updates could also be performed using sampled estimates of those expectations.
By using MC methods to estimate the tilted distribution moments, EP can be
applied in a black-box manner, dramatically expanding the space of models it
can be applied to. Instead of performing a single large sampling task—as would
be required by applying MC methods directly—EP can instead solve several
simpler ones, gaining significant computational advantages [14, 44, 47].
Unfortunately, when performed naively, update (4) is not robust to MC noise.
This is because the moment estimates are converted to the natural (site)
parameter space by mapping through $\nabla A^{*}(.)$, which is not linear in
general, and so MC noise in the estimates leads to biased updates of
$\lambda_{i}$.

##  3 Fearlessly stochastic EP algorithms

In this section we show an equivalence between the moment-matching updates of
EP and natural gradient descent (NGD). We use this view to motivate two new EP
variants which have several advantages when updates are estimated using
samples. We conclude with a review of related work.

###  3.1 Natural gradient view of EP

Several EP variants—EP, power EP, double loop EP, SNEP, and the new ones to
follow—differ only in how they perform the _inner_ optimisation in
$\eqref{eqn:epvariationalproblem}$ with respect to $\\{\lambda_{i}\\}_{i}$.
The optimisation can be viewed as one with respect to the $m+1$ distributions
$p,p_{1},\ldots,p_{m}$, which are jointly parameterised by
$\\{\lambda_{i}\\}_{i}$ [14]. Natural gradient descent (NGD) [1], which
involves preconditioning a gradient with the inverse Fisher information matrix
(FIM) of a distribution, is an effective method for optimising parameters of
probability distributions. It would therefore seems desirable to apply NGD to
the inner optimisation, yet doing so is not straightforward. While the FIM can
be naturally extended to the product of statistical manifolds that form the
solution space, computing and inverting it will not generally be tractable,
making (the natural extension of) NGD in this space infeasible.

Consider instead $\tilde{p}_{i}^{(t)}\in\mathcal{F}$ with natural parameter
$\eta_{i}^{(t)}(\lambda_{i})=\theta_{0}+\sum_{j}\lambda_{j}^{(t)}+\alpha^{-1}(\lambda_{i}-\lambda_{i}^{(t)})$.
The $t$ superscript on the site parameters indicates that they are fixed for
the current iteration, and so the distribution is fully parameterised by
$\lambda_{i}$. Note that when $\lambda_{j}=\lambda_{j}^{(t)}\ \forall\ j$, we
have $\tilde{p}_{i}^{(t)}=p$. Proposition 1, which we prove in Appendix C,
states that the moment-matching updates of EP can be viewed as performing NGD
in $L$ with respect to _mean parameters_ of $\tilde{p}_{i}^{(t)}$.

We will make use of the following properties: for an exponential family with
log partition function $A(.)$, the gradient $\nabla A(.)$ provides the map
from natural to mean parameters, and this mapping is one-to-one for minimal
families, with the inverse given by $\nabla A^{*}(.)$. See Appendix A for
details.

######  Proposition 1.

For $\alpha>0$, the moment-matching update of EP (4) is equivalent to
performing an NGD step in $L$ with respect to the mean parameters of
$\tilde{p}_{i}^{(t)}$ with step size $\alpha^{-1}$. That is, for
$\mu_{i}=\mathbb{E}_{\tilde{p}_{i}^{(t)}(z)}[s(z)]$, and
$\tilde{F}^{(t)}_{i}(\mu_{i})$ the FIM of $\tilde{p}_{i}^{(t)}$ with respect
to $\mu_{i}$, we have

$\displaystyle\mu_{i}-\alpha^{-1}\Big{[}\tilde{F}_{i}^{(t)}(\mu_{i})\Big{]}^{-1}\partialderivative{L}{\mu_{i}}$
$\displaystyle=\Big{(}\nabla
A\circ\eta_{i}^{(t)}\Big{)}\bigg{(}\lambda_{i}-\alpha\Big{(}\theta_{0}+\sum_{j}\lambda_{j}-\nabla
A^{*}\bm{(}\mathbb{E}_{p_{i}(z)}[s(z)]\bm{)}\Big{)}\bigg{)}.$ (7)

Note that the right hand side of (7) just maps update (4) to mean parameters
of $\tilde{p}_{i}^{(t)}$.444Note the apparent contradiction that by decreasing
$\alpha$ (increasing the damping) we actually _increase_ the NGD step size.
This is resolved by observing that the definition of $\tilde{p}^{(t)}_{i}$
also changes with $\alpha$. We discussed in Section 2 that noise in the moment
estimates results in bias in the update to $\lambda_{i}$ of EP, due to the
noise being passed through the nonlinear map $\nabla A^{*}(.)$. This alone
would not be a problem if the updates followed unbiased descent direction
estimates in some other fixed parameterisation. Let $\eta_{i}=\nabla
A^{*}(\mu_{i})$; then, from the definition of $L$ and $\mu_{i}$, we have

$\displaystyle\partialderivative{L}{\mu_{i}}$
$\displaystyle=\partialderivative{\eta_{i}}{\mu_{i}}\partialderivative{\lambda_{i}}{\eta_{i}}\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)},$
(8)

and so when $\mathbb{E}_{p_{i}(z)}[s(z)]$ is estimated with unbiased samples,
the NGD update (7) in $\mu_{i}$ is also unbiased. However, a sequence of
updates introduces bias. This is because the parameterisation changes from one
update to the next, with the mean parameters of $\tilde{p}_{i}^{(t)}$ mapped
to those of $\tilde{p}_{i}^{(t+1)}$ by

$\displaystyle\mu_{i}^{(t+1)}$ $\displaystyle=\Big{(}\nabla
A\circ\eta_{i}^{(t+1)}\circ\big{(}\eta_{i}^{(t)}\big{)}^{-1}\circ\nabla
A^{*}\Big{)}\big{(}\mu_{i}^{(t)}\big{)},$ (9)

where $\big{(}\eta_{i}^{(t)}\big{)}^{-1}(.)$ is the inverse of
$\eta_{i}^{(t)}(.)$, and $\eta_{i}^{(t+1)}(.)$ is defined using the site
parameters _after_ the update at time $t$. This map is not linear in general,
and so an unbiased, noisy update of $\mu_{i}$ at one time step, results in
bias when mapped to the next. This bias makes the EP updates unstable in the
presence of MC noise. This has led previous work to use methods, specific to
$\mathcal{F}$, for obtaining approximately unbiased natural parameter
estimates from samples. However, even with these debiasing methods, a
relatively large number of samples is still needed in practice [44, 47].

We now use the natural gradient interpretation of the EP updates to motivate
two new variants that are far more robust to MC noise; in particular, they are
stable, and most sample-efficient, when updates are estimated with just a
single sample. Furthermore, unlike methods that rely on deibasing estimators,
they do not require sample thinning (see Section 3.4). This largely eliminates
two hyperparameters for the practitioner. While practical considerations, such
as ensuring efficient use of parallel hardware, may justify using more than
one sample per update, this can in principle be tuned a priori, much like the
batch size hyperparameter of stochastic gradient descent.

###  3.2 EP-$\bm{\eta}$

We showed in the previous section that bias is introduced into the sequence of
NGD updates due to a nonlinear map from one parameterisation to the next. If
the map were linear, the sequence would remain unbiased. This can be achieved
by performing NGD with respect to $\eta_{i}$, the natural parameters of
$\tilde{p}_{i}^{(t)}$; then, the map from one parameterisation to the next (by
equating $\lambda_{i}$) is given by
$\eta_{i}^{(t+1)}\circ(\eta_{i}^{(t)})^{-1}$, which is linear. The NGD
direction in $L$ with respect to $\tilde{p}^{(t)}_{i}$ and $\eta_{i}$ is given
by

$\displaystyle-\Big{[}\tilde{F}_{i}(\eta_{i})\Big{]}^{-1}\partialderivative{L}{\eta_{i}}$
$\displaystyle=-\alpha\partialderivative{\eta_{i}}{\mu_{i}}^{\top}\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)};$
(10)

see Appendix D for a derivation. Note that $\alpha$ plays a similar role as
the NGD step size in this parameterisation, but we fix $\alpha=1$ here, and
introduce an explicit NGD step size $\epsilon$. The resulting update can be
expressed directly as an update in $\lambda_{i}$, by applying
$(\eta_{i}^{(t)})^{-1}$ to the updated $\eta_{i}$, giving

$\displaystyle\lambda_{i}$
$\displaystyle\leftarrow\lambda_{i}-\epsilon\partialderivative{\eta_{i}}{\mu_{i}}^{\top}\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)}.$
(11)

We can use automatic differentiation to efficiently compute (11), by
recognising the second term as a Jacobian-vector product (JVP) with respect to
$\eta_{i}(\mu_{i})=\nabla A^{*}(\mu_{i})$.555This can be computed using either
forward or reverse mode automatic differentiation, as
$\partialderivative{\eta_{i}}{\mu_{i}}$ is symmetric. We call the resulting
procedure—which is given in Algorithm 2—EP-$\eta$, to emphasise that we have
simply changed to a natural parameterisation for the NGD update of EP. The
unbiased updates of EP-$\eta$ allow it to be more sample-efficient than EP by
using a small number of samples per iteration. This is demonstrated using a
stochastic version of the clutter problem of Minka [33] in Figure 1a.

###  3.3 EP-$\bm{\mu}$

(a)  Expected decrease in $L$ after $100/n_{\text{samp}}$ steps. Performing
e.g. $100\times$ $1$-sample steps, or $10\times$ $10$-sample steps, achieves a
much larger decrease in $L$ than a single $100$-sample step.

(b)  Magnitude of the bias in $\lambda_{i}$ after a single parallel update,
averaged over all sites and dimensions. The bias of EP-$\mu$ shrinks far
faster as the step size decreases than that of EP. EP-$\eta$ is always
unbiased.

Figure 1: The effect of step size ($\alpha$ or $\epsilon$) and number of
samples ($n_{\text{samp}}$) on different EP variants, in the clutter problem
of Minka [33]. EP (naive) uses maximum likelihood estimation for the updates,
and EP (debiased) uses the estimator of Xu et al. [47]. Step size corresponds
to $\alpha$ for EP, and $\epsilon$ for EP-$\mu$ and EP-$\eta$. $1$-sample
updates are not possible for EP, hence those traces are not visible.

The bias introduced by the updates of EP cannot be mitigated by reducing
$\alpha$ (increasing the damping) without also proportionally sacrificing the
amount of progress made by each update. More concretely, let the bias in
dimension $d$ of $\mu_{i}^{(t+1)}$, after an update at time $t$, be defined as
$E[\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}]_{d}$, where $\bar{\mu}_{i}^{(t+1)}$
is the value of $\mu_{i}^{(t+1)}$ after a noise-free update, and expectation
is taken over the sampling distributions of all parallel updates at time $t$;
then, Proposition 2 below, which we prove in Appendix F, summarises the effect
of decreasing $\alpha$ on EP.

######  Proposition 2.

After update (4) is executed in parallel over $i$, as $\alpha\rightarrow
0^{+}$, both the expected decrease in $L$, and the bias
$E[\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}]_{d}$, are $O(\alpha)$ for all $d$.

Proposition 1 at the beginning of this section states that update (4) can be
viewed as performing an NGD step with respect to $\mu_{i}$ with a step size of
$\alpha^{-1}$. However, changing $\alpha$ also has the effect of changing the
definition of $\tilde{p}_{i}^{(t)}$. It is then natural to wonder what happens
if we fix $\alpha=1$ and introduce an explicit step size $\epsilon$ for NGD,
resulting in the update

$\displaystyle\mu_{i}$ $\displaystyle\leftarrow\mu_{i}-\epsilon\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)};$
(12)

see Appendix E for a derivation. It turns out that in doing so, when we
decrease $\epsilon$, the bias shrinks far faster than the expected decrease in
$L$.

######  Proposition 3.

After update (12) is executed in parallel over $i$, as $\epsilon\rightarrow
0^{+}$, the expected decrease in $L$ is $O(\epsilon)$, and the bias
$E[\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}]_{d}$ is $O(\epsilon^{2})$, for all
$d$.

Proposition 3, which we prove in Appendix F, tells us that by using update
(12), we can reduce the bias to arbitrarily small levels while still making
progress in decreasing $L$. Update (12) can also be expressed directly as an
update in $\lambda_{i}$ by applying $(\eta^{(t)})^{-1}\circ\nabla A^{*}$,
giving

$\displaystyle\lambda_{i}$ $\displaystyle\leftarrow\nabla
A^{*}\Big{(}(1-\epsilon)\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}+\epsilon\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)}-\theta_{0}-\sum_{j\neq
i}\lambda_{j},$ (13)

which has the simple interpretation of performing EP, but with damping of the
moments instead of the site parameters. We call this variant EP-$\mu$, again
to indicate that we are simply performing the NGD update of EP in the mean
parameter space.666EP too is performing NGD in mean parameters, but using a
step size that also affects the distribution $\tilde{p}_{i}$. The resulting
procedure is also given by Algorithm 2.

Algorithm 2 EP-$\eta$ and EP-$\mu$ (differences with Algorithm 1 are
highlighted in  green)

$\mathcal{F}$, $\theta_{0}$, $\\{\beta_{i}\\}_{i}$, $\\{\ell_{i}(z)\\}_{i}$,
$\\{\lambda_{i}\\}_{i}$, $n_{\text{inner}}$, $\epsilon$

while not converged do

$\theta_{i}\leftarrow\theta_{0}+\sum_{j}\lambda_{j}$ for $i=1$ to $m$

for $1$ to $n_{\text{inner}}$ do

for $i=1$ to $m$ _in parallel_ do

update $\lambda_{i}$ using (11) for EP-$\eta$, or (13) for EP-$\mu$

return $\\{\lambda_{i}\\}_{i}$

The computational cost of EP-$\mu$ is lower than that of EP-$\eta$ because it
does not require JVPs through $\nabla A^{*}(.)$.777See Appendix G for a
detailed discussion of the computational costs of EP-$\eta$ and EP-$\mu$. The
drawback is that the updates of EP-$\mu$ do still retain some bias; however,
we find that the bias of EP-$\mu$ typically has negligible impact on its
performance relative to EP-$\eta$. This is evident in the clutter problem of
Minka [33], as shown in Figure 1a, as well as in the larger scale experiments
of Section 4. Figure 1b illustrates how quickly the bias in $\lambda_{i}$
shrinks as the step size of EP-$\mu$ is decreased; note that the results of
this subsection are stated in terms of $\mu_{i}$, but it is straightforward to
show that equivalent results also hold for $\lambda_{i}$ using Taylor series
arguments.

###  3.4 Related work

EP was combined with MC moment estimates by Xu et al. [47] in a method called
sampling via moment sharing (SMS), and later by Vehtari et al. [44] in the
context of hierarchical models. In both works, when $\mathcal{F}$ was
multivariate normal—arguably the most useful case—the authors used an
estimator for the updates that is unbiased when $\tilde{p}_{i}$ is also
multivariate normal. Although this estimator is only approximately unbiased in
general, it can help to stabilise the updates significantly. Even so, it is
necessary to use a relatively large number of samples for the updates, with
several hundred or more being typical. Using many samples per update is
inefficient, as the update direction can change from one iteration to the
next. Furthermore, such estimators typically rely on a known number of
independent samples; samples drawn using Markov chain Monte Carlo (MCMC)
methods are generally autocorrelated, necessitating the use of sample thinning
before the estimators can be applied. This adds another hyperparameter to the
procedure—the thinning ratio—and it is also inefficient due to the discarding
of samples. The practitioner is required to choose the number of samples, the
thinning ratio, and the amount of damping, all of which affect the accuracy,
stability, and computational efficiency of the procedure. There is no easy way
to choose these a priori, forcing the practitioner to either set them
conservatively (favouring accuracy and stability), or to find appropriate
settings by trial and error, both of which are likely to expend time and
computation unnecessarily.

The stochastic natural gradient EP (SNEP) method of Hasenclever et al. [14]
optimises the same variational objective as EP, but performs stochastic
natural gradient descent with respect to a mean parameterisation (under
$\mathcal{F}$) of the site potentials in order to perform the inner
minimisation. Its natural gradient updates are unbiased in the presence of MC
noise, allowing them to be performed with as little as one sample, without
relying on $\mathcal{F}$-specific debiasing estimators, or the sample thinning
that would entail. The authors found that in some settings, accurate point
estimates can be obtained fairly rapidly with SNEP. However, we find that SNEP
typically converges far slower than EP, consistent with findings in Vehtari et
al. [44]. In SNEP the site potentials are treated as bona fide distributions
for the purposes of optimisation, and can be seen as playing the role of
surrogate distributions for NGD [42]. The results of this section showed that
EP, EP-$\eta$ and EP-$\mu$ are closely related to SNEP; they too can be viewed
as performing NGD of $L$ with respect to surrogate distributions
($\tilde{p}_{i}$), but with distributions that are more closely matched with
those being optimised. This allows them to retain the same benefits as SNEP,
but with faster convergence.

We note that Xu et al. [47] and Hasenclever et al. [14] also proposed methods
for performing updates in an asynchronous fashion, significantly reducing the
frequency and cost of communication between nodes in distributed settings. In
principle, these methods could be combined with those presented in this paper,
but we do not consider them further here.

##  4 Evaluation

In this section we demonstrate the efficacy of EP-$\eta$ and EP-$\mu$ on a
variety of probabilistic inference tasks. In each experiment, the task was to
perform approximate Bayesian inference of unobserved parameters $z$, given
some observed data $\mathcal{D}$. All of the models in these experiments
followed the same general structure, consisting of a minimal exponential
family prior over $z$, and a partition of the data
$\\{\mathcal{D}_{i}\\}_{i=1}^{m}$, where each block $\mathcal{D}_{i}$ depends
on both $z$ and an additional vector of _local_ latent variables $w_{i}$ that
also depend on $z$; that is, the joint density has structure

$\displaystyle p_{0}(z)\prod_{i=1}^{m}p(w_{i}\mid z)p(\mathcal{D}_{i}\mid
w_{i}),$ (14)

where $p_{0}\in\mathcal{F}$. This is shown graphically in Appendix H. To
perform approximate inference of $z$, we first define $\ell_{i}(z)$ as the log
likelihood of $\mathcal{D}_{i}$ given $z$, with $w_{i}$ marginalised out; that
is,

$\displaystyle\ell_{i}(z)$ $\displaystyle=\log\int p(\mathcal{D}_{i},w_{i}\mid
z)\mathrm{d}w_{i}.$ (15)

Then, given $p_{0}(z)$ and $\\{\ell_{i}(z)\\}_{i}$, we define $p^{*}(z)$ as
(1), and proceed to find an approximation $p(z)\approx p^{*}(z)$ using the
methods described in earlier sections. Note that sampling from the tilted
distribution $p_{i}(z)$ requires jointly sampling over $z$ and $w_{i}$ and
taking the marginal. EP is a particularly appealing framework for performing
inference in this setting, as the dimensionality of the sampled distributions
is constant with respect to $m$, mitigating the curse of dimensionality
experienced by conventional MCMC approaches [44]. In our experiments we used
NUTS [24] to perform the sampling, consistent with prior work [44, 47]. In
each experiment we compared EP-$\eta$ and EP-$\mu$ with EP, in the manner used
by Xu et al. [47], Vehtari et al. [44]. Where $\mathcal{F}$ was MVN, we also
compared to SNEP [14]. For the models with NIW $\mathcal{F}$, we were unable
to find an initialisation for SNEP that would allow us to perform a meaningful
comparison; we discuss this more in Appendix H.

To evaluate the performance of the different variants, we monitored KL
divergence to an estimate of the optimum, obtained by running EP with a large
number of samples. We used 500 different hyperparameter settings for each
variant, chosen using random search, and repeated each run 5 times using
different random seeds for NUTS. All runs of EP-$\eta$ and EP-$\mu$ were
performed with $n_{\text{samp}}=1$ and $n_{\text{inner}}=1$. The results of
these experiments are summarised in Figure 2

(a)  Hierarchical logistic regression: MVN prior

(b)  Hierarchical logistic regression: NIW prior

(c)  Cosmic radiation model

(d)  Neural response model

Figure 2: Pareto frontier showing the number of NUTS steps ($x$-axis) against
the KL divergence from $p$ to an estimate of the optimum ($y$-axis). Each
point on the plot marks the lowest average KL divergence attained by _any_
hyperparameter setting by that step count. Error bars mark the full range of
values for the marked hyperparameter setting across 5 random seeds.

We see that EP-$\eta$ and EP-$\mu$ can typically reach a given level of
accuracy (distance from optimum) faster than EP, often significantly so. We
stress that the frontiers for EP-$\eta$ and EP-$\mu$ are traced out with
$n_{\text{samp}}$ and $n_{\text{inner}}$ each fixed to $1$, with no sample
thinning, varying only the step size hyperparameter $\epsilon$, with higher
accuracy/cost regions of the frontier corresponding to lower values of
$\epsilon$. In contrast, to trace out the frontier of EP we must jointly vary
$\alpha$, $n_{\text{samp}}$, and the thinning ratio; see Appendix I for
examples of this effect. The performances of EP-$\eta$ and EP-$\mu$ are very
similar. We found SNEP to be significantly slower than the other methods,
consistent with findings by Vehtari et al. [44]. Pareto frontiers with respect
to wall-clock time can be found in Appendix J. We now provide a brief overview
of individual experiments, with further details given in Appendix H.

###  4.1 Hierarchical logistic regression

Hierarchical logistic regression is used to perform binary classification in a
number of groups when the extent to which data should be pooled across groups
is unknown [12]. Each of the $m$ groups has its own covariates and binary
response variables, collectively denoted by $\mathcal{D}_{i}$, as well as
regression coefficients $w_{i}$. The groups are related through prior
parameters $z$, which parameterise the group-level coefficient prior densities
$p(w_{i}\mid z)$. We performed two experiments in this setting.

#### Synthetic data with MVN prior

In this experiment, the group-level prior density $p(w_{i}\mid z)$ was a
product of independent normal densities, with $p_{0}\in\mathcal{F}$ being a
multivariate normal (MVN) prior over the means and log-variances. We used the
same synthetic data generation as Vehtari et al. [44], which was designed to
be challenging for EP. In this experiment $z\in\mathbb{R}^{8}$,
$w_{i}\in\mathbb{R}^{4}\ \forall\ i$, $m=16$, and each of the $m$ regression
tasks had $n=20$ observations. Results are shown in Figure 2a.

#### Political survey data with NIW prior

In this experiment, each of $m=50$ nested regression tasks corresponded to
political survey responses from a particular US state, with $7$ predictor
variables corresponding to characteristics of a given survey respondent, so
that $w_{i}\in\mathbb{R}^{7}$ for $i=1,\ldots,m$. The state-level prior
density $p(w_{i}\mid z)$ was that of a MVN distribution, with
$p_{0}\in\mathcal{F}$ being a normal-inverse-Wishart (NIW) prior on the mean
and covariance matrix. This structure allows for correlation of regression
coefficients across states. The data, taken from the 2018 Cooperative
Congressional Election Study, related to support for allowing employers to
decline coverage of abortions in insurance plans [29, 39]. We truncated the
data slightly so that each state contained $n=97$ survey responses. Results
are shown in Figure 2b.

###  4.2 Cosmic radiation model

A hierarchical Bayesian model was used by Vehtari et al. [44] to capture the
nonlinear relationship between diffuse galactic far ultraviolet radiation and
100-µm infrared emission in various sectors of the observable universe, using
data obtained from the Galaxy Evolution Explorer telescope. In this model each
$w_{i}\in\mathbb{R}^{9}$ parameterised a nonlinear regression model using data
obtained from one of $m$ sections of the observable universe. The $m$
regression problems were related through hyperparameters
$z\in\mathbb{R}^{18}$, which parameterised the section-level coefficient prior
densities $p(w_{i}\mid z)$ for $i=1,\ldots,m$. $p_{0}\in\mathcal{F}$ was MVN.
The specifics of the nonlinear regression model are quite involved and we
refer the reader to Vehtari et al. [44] for details.

Vehtari et al. [44] used EP to perform inference in this model. We were unable
to obtain the dataset, and so we generated synthetic data using
hyperparameters parameters that were tuned by hand to try and match the
qualitative properties of the original data; see Appendix K for examples. We
used a reduced number of $m=36$ sites and $n=200$ observations per site to
allow us to perform a comprehensive hyperparameter search. Results are shown
in Figure 2c.

###  4.3 Neural response model

A common task in neuroscience is to model the firing rates of neurons under
various conditions. We performed inference in a hierarchical Bayesian neural
response model, using recordings of V1 complex cells in an anesthesised adult
cat [3]. In this dataset, 10 neurons in a specific area of cat V1 were
simultaneously recorded under the presentation of 18 different visual stimuli,
each repeated 8 times, for a total of 144 trials. We modelled the observed
spike counts of the 10 neurons in each trial as Poisson, with latent (log)
intensities being MVN, with mean and covariance (collectively
$z\in\mathbb{R}^{65}$) drawn from a NIW prior $p_{0}\in\mathcal{F}$. Inference
of $z$ amounted to inferring the means, and inter-neuron correlations, of
latent log firing rates across stimuli. We grouped the data into $m=8$ batches
of $n=18$ trials each, so that the latent log firing rates for all trials in
batch $i$ were jointly captured by $w_{i}\in\mathbb{R}^{180}$. Results are
shown in Figure 2d.

##  5 Conclusion

In this work, we used a novel interpretation of the moment-matching updates of
EP to motivate two new EP variants that are far more robust to MC noise. We
demonstrated that these variants can offer an improved speed-accuracy trade-
off compared to their predecessors, and are easier to tune.

#### Limitations

We largely assumed that the drawing of MC samples is the dominant cost. In
practice, other costs become relevant, which may shift the balance in favour
of using more samples to estimate updates. We discuss this in Appendix G,
along with strategies for reducing computational overheads. We note, however,
that wall-clock time results for our experiments (in Appendix J) are in line
with the results of Section 4. When EP and its variants—including those
introduced here—are combined with MC methods, the underlying samplers
typically have hyperparameters of their own, which are often adapted using so-
called warmup phases. While EP-$\eta$ and EP-$\mu$ significantly reduce the
complexity of tuning hyperparameters specific to EP, they do not help with
those of the underlying samplers; this means that the complexity of
hyperparameter tuning is still greater than that of direct MC methods. We used
comprehensive hyperparameter searches in our experiments in order to perform a
meaningful comparison with baseline methods, necessarily limiting the scale of
problems we could tackle.

#### Future work

It would be worthwhile considering how our algorithms can be efficiently
implemented for modern parallel compute hardware; e.g., practical performance
could be improved by using sampling kernels that are themselves designed for
efficient parallel execution [23]. Performance may also be improved in
distributed settings by using asynchronous extensions to minimise
communication overheads [14, 47]. The need to tune sampling hyperparameters
could be obviated by using information from $p$ to set them directly; e.g.,
when $p\in\mathcal{F}$ is MVN, its precision matrix could be used directly as
the mass matrix hyperparameter in Hamiltonian Monte Carlo schemes. We would
also like to apply our methods to large-scale problems. We leave investigation
of these ideas to future work.

## References

  * Amari [1998] S.-i. Amari.  Natural Gradient Works Efficiently in Learning.  _Neural Computation_ , 1998.
  * Anil et al. [2021] R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer.  Scalable second order optimization for deep learning.  _arXiv:2002.09018_ , 2021.
  * Blanche [2009] T. Blanche.  Multi-neuron recordings in primary visual cortex., 2009.
  * Blondel et al. [2021] M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer, F. Llinares-López, F. Pedregosa, and J.-P. Vert.  Efficient and modular implicit differentiation.  _arXiv:2105.15183_ , 2021.
  * Bradbury et al. [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.
  * Brooks et al. [2011] S. Brooks, A. Gelman, G. Jones, and X.-L. Meng.  _Handbook of Markov Chain Monte Carlo_.  Chapman and Hall/CRC, 2011.
  * Bui et al. [2016] T. Bui, D. Hernandez-Lobato, J. Hernandez-Lobato, Y. Li, and R. Turner.  Deep gaussian processes for regression using approximate expectation propagation.  _International Conference on Machine Learning_ , 2016.
  * Bui et al. [2017] T. D. Bui, J. Yan, and R. E. Turner.  A unifying framework for gaussian process pseudo-point approximations using power expectation propagation.  _Journal of Machine Learning Research_ , 2017.
  * Chu and Ghahramani [2005] W. Chu and Z. Ghahramani.  Gaussian processes for ordinal regression.  _Journal of Machine Learning Research_ , 2005.
  * Cunningham et al. [2008] J. P. Cunningham, K. V. Shenoy, and M. Sahani.  Fast gaussian process methods for point process intensity estimation.  _International Conference on Machine Learning_ , 2008.
  * Cunningham et al. [2013] J. P. Cunningham, P. Hennig, and S. Lacoste-Julien.  Gaussian probabilities and expectation propagation.  _arXiv:1111.6832_ , 2013.
  * Gelman et al. [2013] A. Gelman, J. Carlin, H. Stern, D. Dunson, A. Vehtari, and D. Rubin.  _Bayesian Data Analysis, Third Edition_.  Chapman and Hall/CRC, 2013.
  * Gonzalez et al. [2016] J. Gonzalez, M. Osborne, and N. Lawrence.  Glasses: Relieving the myopia of bayesian optimisation.  _International Conference on Artificial Intelligence and Statistics_ , 2016.
  * Hasenclever et al. [2017] L. Hasenclever, S. Webb, T. Lienart, S. Vollmer, B. Lakshminarayanan, C. Blundell, and Y. W. Teh.  Distributed bayesian learning with stochastic natural gradient expectation propagation and the posterior server.  _Journal of Machine Learning Research_ , 2017.
  * Hennig [2009] P. Hennig.  Expectation propagation on the maximum of correlated normal variables.  Technical report, University of Cambridge, 2009.
  * Hennig and Schuler [2012] P. Hennig and C. J. Schuler.  Entropy search for information-efficient global optimization.  _Journal of Machine Learning Research_ , 2012.
  * Hennig et al. [2010] P. Hennig, D. Stern, and T. Graepel.  Coherent inference on optimal play in game trees.  _International Conference on Artificial Intelligence and Statistics_ , 2010.
  * Herbrich et al. [2006] R. Herbrich, T. Minka, and T. Graepel.  Trueskill™: A bayesian skill rating system.  _Advances in Neural Information Processing Systems_ , 2006.
  * Hernandez-Lobato and Hernandez-Lobato [2016] D. Hernandez-Lobato and J. M. Hernandez-Lobato.  Scalable gaussian process classification via expectation propagation.  _International Conference on Artificial Intelligence and Statistics_ , 2016.
  * Hernández-Lobato and Adams [2015] J. M. Hernández-Lobato and R. P. Adams.  Probabilistic backpropagation for scalable learning of bayesian neural networks.  _International Conference on Machine Learning_ , 2015.
  * Heskes and Zoeter [2002] T. Heskes and O. Zoeter.  Expectation propagation for approximate inference in dynamic bayesian networks.  _Uncertainty in Artificial Intelligence_ , 2002.
  * Heskes and Zoeter [2003] T. Heskes and O. Zoeter.  Extended version of “expectation propagation for approximate inference in dynamic bayesian networks.  Technical report, University of Nijmegen, 2003.
  * Hoffman et al. [2021] M. Hoffman, A. Radul, and P. Sountsov.  An adaptive-mcmc scheme for setting trajectory lengths in hamiltonian monte carlo.  _International Conference on Artificial Intelligence and Statistics_ , 2021.
  * Hoffman and Gelman [2014] M. D. Hoffman and A. Gelman.  The no-u-turn sampler: Adaptively setting path lengths in hamiltonian monte carlo.  _Journal of Machine Learning Research_ , 2014.
  * Højen-Sørensen et al. [2002] P. A. d. F. R. Højen-Sørensen, O. Winther, and L. K. Hansen.  Mean-field approaches to independent component analysis.  _Neural Computation_ , 2002.
  * Jitkrittum et al. [2015] W. Jitkrittum, A. Gretton, N. M. O. Heess, S. M. A. Eslami, B. Lakshminarayanan, D. Sejdinovic, and Z. Szabó.  Kernel-based just-in-time learning for passing expectation propagation messages.  _Uncertainty in Artificial Intelligence_ , 2015.
  * Jylänki et al. [2011] P. Jylänki, J. Vanhatalo, and A. Vehtari.  Robust gaussian process regression with a student-$t$ likelihood.  _Journal of Machine Learning Research_ , 2011.
  * Li et al. [2015] Y. Li, J. M. Hernández-Lobato, and R. E. Turner.  Stochastic expectation propagation.  _Advances in Neural Information Processing Systems_ , 2015.
  * Lopez-Martin et al. [2022] J. Lopez-Martin, J. Phillips, and A. Gelman.  Multilevel regression and poststratification case studies, 2022.
  * Minka [2004] T. Minka.  Power ep.  Technical report, Microsoft, 2004.
  * Minka and Lafferty [2002] T. Minka and J. Lafferty.  Expectation-propagation for the generative aspect model.  _Uncertainty in Artificial Intelligence_ , 2002.
  * Minka et al. [2018] T. Minka, R. Cleven, and Y. Zaykov.  Trueskill 2: An improved bayesian skill rating system.  Technical report, Microsoft, 2018.
  * Minka [2001] T. P. Minka.  _A family of algorithms for approximate Bayesian inference_.  PhD thesis, Massachusetts Institute of Technology, 2001.
  * Naish-guzman and Holden [2007] A. Naish-guzman and S. Holden.  The generalized fitc approximation.  _Advances in Neural Information Processing Systems_ , 2007.
  * Nodelman et al. [2005] U. Nodelman, D. Koller, and C. Shelton.  Expectation propagation for continuous time bayesian networks.  _Uncertainty in Artificial Intelligence_ , 2005.
  * Opper and Winther [2000] M. Opper and O. Winther.  Gaussian processes for classification: Mean-field algorithms.  _Neural Computation_ , 2000.
  * Opper et al. [2008] M. Opper, U. Paquet, and O. Winther.  Improving on expectation propagation.  _Advances in Neural Information Processing Systems_ , 2008.
  * Phan et al. [2019] D. Phan, N. Pradhan, and M. Jankowiak.  Composable effects for flexible and accelerated probabilistic programming in numpyro, 2019.
  * Schaffner et al. [2019] B. Schaffner, S. Ansolabehere, and S. Luks.  CCES common content 2018, 2019.
  * Seeger and Nickisch [2011] M. Seeger and H. Nickisch.  Fast convergent algorithms for expectation propagation approximate bayesian inference.  _International Conference on Artificial Intelligence and Statistics_ , 2011.
  * So [2024] J. So.  Estimating the normal-inverse-wishart distribution.  _arXiv:2405.16088_ , 2024.
  * So and Turner [2024] J. So and R. E. Turner.  Optimising distributions with natural gradient surrogates.  _International Conference on Artificial Intelligence and Statistics_ , 2024.
  * Turner and Sahani [2011] R. E. Turner and M. Sahani.  Probabilistic amplitude and frequency demodulation.  _Advances in Neural Information Processing Systems_ , 2011.
  * Vehtari et al. [2020] A. Vehtari, A. Gelman, T. Sivula, P. Jylänki, D. Tran, S. Sahai, P. Blomstedt, J. P. Cunningham, D. Schiminovich, and C. P. Robert.  Expectation propagation as a way of life: A framework for bayesian inference on partitioned data.  _Journal of Machine Learning Research_ , 2020.
  * Wainwright and Jordan [2008] M. Wainwright and M. Jordan.  Graphical models, exponential families, and variational inference.  _Foundations and Trends in Machine Learning_ , 2008.
  * Welling and Lim [2007] M. Welling and J. J. Lim.  A distributed message passing algorithm for sensor localization.  _International Conference on Artificial Neural Networks_ , 2007.
  * Xu et al. [2014] M. Xu, B. Lakshminarayanan, Y. W. Teh, J. Zhu, and B. Zhang.  Distributed bayesian posterior sampling via moment sharing.  _Advances in Neural Information Processing Systems_ , 2014.
  * Zoeter and Heskes [2005] O. Zoeter and T. Heskes.  Gaussian quadrature based expectation propagation.  _Proceedings of the Tenth International Workshop on Artificial Intelligence and Statistics_ , 2005.

##  Appendix A Exponential families of distributions

In this appendix, we give a very brief overview of exponential families of
distributions. This is a highly condensed summary of relevant material from
Wainwright and Jordan [45], and we refer the reader to the original work for a
far more detailed treatment.

The exponential family of distributions $\mathcal{F}$, defined by the
($d$-dimensional, vector-valued) sufficient statistic function $s(.)$ and base
measure $\nu(.)$, has density function

$\displaystyle p(z)$
$\displaystyle=\exp\bm{(}\theta^{\top}s(z)-A(\theta)\bm{)},$ (16)

taken with respect to $\nu(.)$, where
$A(\theta)=\log\int\exp\bm{(}\theta^{\top}s(z)\bm{)}\mathrm{d}\nu(z)$ is known
as the _log partition_ function. $\theta\in\Omega$ are known as the _natural
parameters_ , and are used to index a specific member of $\mathcal{F}$.
$\Omega$ is the set of normalisable natural parameters of $\mathcal{F}$, given
by

$\displaystyle\Omega=\left\\{\theta\in\mathbb{R}^{d}\ \Bigg{\lvert}\
\int\exp\bm{(}\theta^{\top}s(z)\bm{)}\mathrm{d}\nu(z)<\infty\right\\},$ (17)

and is known as the _natural domain_ of $\mathcal{F}$. $\Omega$ is a convex
set, and when it is open, the family $\mathcal{F}$ is said to be _regular_ ;
we shall only consider regular families in this work. When the components of
$s(.)$ are linearly independent, $\mathcal{F}$ is said to be _minimal_. For
minimal families, each distribution in the family is associated with a unique
natural parameter vector $\theta$. An exponential family that is not minimal
is said to be _overcomplete_ , in which case, each distribution is associated
with an entire affine subspace of $\Omega$.

The expected sufficient statistics of a distribution with density $p(.)$ with
respect to base measure $\nu(.)$, are given by $\mathbb{E}_{p(z)}[s(z)]$. The
_mean domain_ , $\mathcal{M}$, is defined as the set of expected sufficient
statistics that can be attained by _any_ density $p(.)$ with respect to base
measure $\nu(.)$; that is,

$\displaystyle\mathcal{M}=\left\\{\mu\in\mathbb{R}^{d}\ \Bigg{\lvert}\ \exists
p(.):\int p(z)s(z)\mathrm{d}\nu(z)=\mu\right\\}.$ (18)

All vectors in the interior of $\mathcal{M}$, $\mathcal{M}^{\circ}$, are
realisable by a member of $\mathcal{F}$, and so $\mu\in\mathcal{M}^{\circ}$
provides an alternative parameterisation of $\mathcal{F}$, in which $\mu$ are
known as the _mean parameters_.

For minimal families, $A(.)$ is strictly convex on $\Omega$. The convex dual
of $A(.)$ is defined as

$\displaystyle A^{*}(\mu)=\max_{\theta\in\Theta}\theta^{\top}\mu-A(\theta),$
(19)

and on $\mathcal{M}^{\circ}$, $A^{*}(\mu)$ is equal to the negative entropy of
the member of $\mathcal{F}$ with mean parameter $\mu$, hence $A^{*}(.)$ is
also referred to as the _negative entropy function_.

The mean parameters of the member of $\mathcal{F}$ with natural parameter
$\theta$ can be obtained by $\mu=\nabla A(\theta)$. For minimal families, this
mapping is one-to-one, and the reverse map is given by $\theta=\nabla
A^{*}(\theta)$. For this reason, $\nabla A(.)$ and $\nabla A^{*}(.)$ are
sometimes referred to as the _forward mapping_ and _backward mapping_
respectively. In this work, we say that a family $\mathcal{F}$ is _tractable_
if both the forward and backward mappings can be evaluated efficiently.

The Fisher information matrix (FIM) of an exponential family distribution,
with respect to natural parameters $\theta$, is given by,
$F(\theta)=\nabla^{2}A(\theta)$. Furthermore, in minimal families, the FIM
with respect to mean parameters $\mu$ is given by
$F(\mu)=\nabla^{2}A^{*}(\mu)$. Using the forward and backward mappings, we
also have

$\displaystyle F(\theta)$ $\displaystyle=\partialderivative{\mu}{\theta}$ (20)
$\displaystyle F(\mu)$ $\displaystyle=\partialderivative{\theta}{\mu},$ (21)

and from the inverse function theorem,

$\displaystyle F(\theta)^{-1}$ $\displaystyle=\partialderivative{\theta}{\mu}$
(22) $\displaystyle F(\mu)^{-1}$
$\displaystyle=\partialderivative{\mu}{\theta}.$ (23)

##  Appendix B EP inner update derivation

Begin by taking gradients of (3), with respect to $\lambda_{i}$, giving

$\displaystyle\nabla_{\lambda_{i}}L(\eta_{1},\ldots,\eta_{m},\lambda_{1},\ldots,\lambda_{m})$
$\displaystyle=\nabla
A\left(\eta_{0}+\sum_{j}\lambda_{j}\right)-\nabla_{\lambda_{i}}A_{i}\bm{(}(\eta_{i}-\beta_{i}^{-1}\lambda_{i},\beta_{i}^{-1})\bm{)}.$
(24)

At a minimum of $L$ with respect to $\lambda_{i}$, both sides of (24) must be
equal to zero. We can use this to define a fixed-point condition with respect
to $\lambda_{i}$, and its updated value $\lambda_{i}^{\prime}$

$\displaystyle\nabla A\left(\eta_{0}+\lambda_{i}^{\prime}+\sum_{j\neq
i}\lambda_{j}\right)$
$\displaystyle=\nabla_{\lambda_{i}}A_{i}\bm{(}(\eta_{i}-\beta_{i}^{-1}\lambda_{i},\beta_{i}^{-1})\bm{)}$
$\displaystyle=\mathbb{E}_{p_{i}(z)}[s(z)]$ (25)

where $p_{i}\in\mathcal{F}_{i}$ denotes the $i$-th tilted distribution,
defined by (6). Applying $\nabla A^{*}(.)$, the inverse of $\nabla A(.)$, to
both sides of (25), and rearranging terms, we recover the moment-matching
update of EP

$\displaystyle\lambda_{i}$ $\displaystyle\leftarrow\nabla
A^{*}\big{(}\mathbb{E}_{p_{i}(z)}[s(z)]\big{)}-\eta_{0}-\sum_{j\neq
i}\lambda_{j}.$ (26)

Often a damping parameter is used to aid convergence, giving the more general
update

$\displaystyle\lambda_{i}$
$\displaystyle\leftarrow(1-\alpha)\lambda_{i}+\alpha\left(\nabla
A^{*}\big{(}\mathbb{E}_{p_{i}(z)}[s(z)]\big{)}-\eta_{0}-\sum_{j\neq
i}\lambda_{j}\right)$ $\displaystyle\leftarrow\lambda_{i}+\alpha\left(\nabla
A^{*}\big{(}\mathbb{E}_{p_{i}(z)}[s(z)]\big{)}-\eta_{0}-\sum_{j}\lambda_{j}\right),$
(27)

where the level of damping is given by $(1-\alpha)$. It was shown by Heskes
and Zoeter [22] that (4) follows a decrease direction in $L$ with respect to
$\lambda_{i}$, and so is guaranteed to decrease $L$ when $\alpha$ is small
enough.

##  Appendix C Natural gradient view of EP

Let $\tilde{p}_{i}^{(t)}$ be the member of $\mathcal{F}$ with natural
parameter
$\eta_{i}^{(t)}(\lambda_{i})=\theta_{0}+\sum_{j}\lambda_{j}^{(t)}+\alpha^{-1}(\lambda_{i}-\lambda_{i}^{(t)})$.
Proposition 1 states that that the moment-matching updates of EP can be viewed
as performing NGD in $L$ with respect to the mean parameters, $\mu_{i}$, of
$\tilde{p}_{i}^{(t)}$. We restate the proposition below, and give a proof.
Note that the right hand side of (7) maps the EP update (4) onto mean
parameters of $\tilde{p}_{i}^{(t)}$. The left hand side of (7) is simply a NGD
update in $\mu_{i}$.

See 1

###### Proof.

Taking the left hand side of (7), we have

$\displaystyle\mu_{i}-\alpha^{-1}\big{[}\tilde{F}_{i}^{(t)}(\mu_{i})\big{]}^{-1}\partialderivative{\lambda_{i}}{\mu_{i}}$
$\displaystyle=\mu_{i}-\alpha^{-1}\big{[}\tilde{F}_{i}^{(t)}(\mu_{i})\big{]}^{-1}\partialderivative{\lambda_{i}}{\mu_{i}}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=\mu_{i}-\alpha^{-1}\bigg{(}\partialderivative{\eta_{i}}{\mu_{i}}\bigg{)}^{-1}\partialderivative{\eta_{i}}{\mu_{i}}\partialderivative{\lambda_{i}}{\eta_{i}}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=\mu_{i}-\partialderivative{L}{\lambda_{i}}$
$\displaystyle=\mu_{i}-\bigg{(}\nabla
A\Big{(}\theta_{0}+\sum_{j}\lambda_{j}\Big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\bigg{)}$
$\displaystyle=\nabla
A\Big{(}\theta_{0}+\sum_{j}\lambda_{j}\Big{)}-\bigg{(}\nabla
A\Big{(}\theta_{0}+\sum_{j}\lambda_{j}\Big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\bigg{)}$
$\displaystyle=\mathbb{E}_{p_{i}(z)}[s(z)],$ (28)

where the penultimate equality follows from the definition
$\mu_{i}=\mathbb{E}_{\tilde{p}_{i}^{(t)}(z)}[s(z)]$, and the fact that
$\lambda_{j}^{(t)}=\lambda_{j}\ \forall\ j$ before the update. All that
remains is to show that the right hand side of (7) is equal to
$\mathbb{E}_{p_{i}(z)}[s(z)]$;

$\displaystyle\big{(}\nabla
A\circ\eta_{i}^{(t)}\big{)}\bigg{(}\lambda_{i}-\alpha\Big{(}\theta_{0}+\sum_{j}\lambda_{j}-\nabla
A^{*}\bm{(}\mathbb{E}_{p_{i}(z)}[s(z)]\bm{)})\Big{)}\bigg{)}$ $\displaystyle=\
$ $\displaystyle\nabla
A\bigg{(}\alpha^{-1}(\lambda_{i}-\lambda_{i}^{(t)})+\nabla
A^{*}(\mathbb{E}_{p_{i}(z)}[s(z)]\bm{)})\bigg{)}$ $\displaystyle=\ $
$\displaystyle\mathbb{E}_{p_{i}(z)}[s(z)]$ (29)

∎

##  Appendix D EP-$\bm{\eta}$ update direction derivation

Using the relation $\eta_{i}=\eta_{i}^{(t)}(\lambda_{i})$, and therefore
$\lambda_{i}=\big{(}\eta_{i}^{(t)}\big{)}^{-1}(\eta_{i})$, we have

$\displaystyle-\Big{[}\tilde{F}_{i}(\eta_{i})\Big{]}^{-1}\partialderivative{L}{\eta_{i}}$
$\displaystyle=-\Big{[}\tilde{F}_{i}(\eta_{i})\Big{]}^{-1}\partialderivative{\lambda_{i}}{\eta_{i}}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=-\alpha\Big{(}\partialderivative{\mu_{i}}{\eta_{i}}\tilde{F}_{i}(\mu_{i})\partialderivative{\mu_{i}}{\eta_{i}}^{\top}\Big{)}^{-1}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=-\alpha\Big{(}\partialderivative{\eta_{i}}{\mu_{i}}^{\top}\tilde{F}_{i}(\mu_{i})^{-1}\partialderivative{\eta_{i}}{\mu_{i}}\Big{)}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=-\alpha\Big{(}\partialderivative{\eta_{i}}{\mu_{i}}^{\top}\partialderivative{\mu_{i}}{\eta_{i}}\partialderivative{\eta_{i}}{\mu_{i}}\Big{)}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=-\alpha\partialderivative{\eta_{i}}{\mu_{i}}^{\top}\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)}.$
(30)

##  Appendix E EP-$\bm{\mu}$ update derivation

For $\alpha=1$, we have

$\displaystyle\mu_{i}-\epsilon\Big{[}\tilde{F}_{i}(\mu_{i})\Big{]}^{-1}\partialderivative{L}{\mu_{i}}$
$\displaystyle=\mu_{i}-\epsilon\Big{[}\tilde{F}_{i}(\mu_{i})\Big{]}^{-1}\partialderivative{\eta_{i}}{\mu_{i}}\partialderivative{\lambda_{i}}{\eta_{i}}\partialderivative{L}{\lambda_{i}}$
$\displaystyle=\mu_{i}-\epsilon\Big{(}\partialderivative{\eta_{i}}{\mu_{i}}\Big{)}^{-1}\partialderivative{\eta_{i}}{\mu_{i}}\partialderivative{\lambda_{i}}{\eta_{i}}\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)}$
$\displaystyle=\mu_{i}-\epsilon\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]\Big{)}.$
(31)

##  Appendix F Bias of EP and EP-$\bm{\mu}$

Let the bias in dimension $d$ of $\mu_{i}^{(t+1)}$, after an update at time
$t$, be defined as $E[\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}]_{d}$, where
$\bar{\mu}_{i}^{(t+1)}$ is the value of $\mu_{i}^{(t+1)}$ after a noise-free
update, and expectation is taken over the sampling distributions of all
parallel updates at time $t$. Propositions 2 and 3, which we restate and prove
below, summarise the effect of decreasing $\alpha$ on EP and EP-$\mu$.

See 2

###### Proof.

Let $\eta_{i}^{(t)}$ and $\mu_{i}^{(t)}$ be the pre-update natural and mean
parameters of $\tilde{p}_{i}{(t)}$ respectively. Furthermore, let

$\displaystyle\mu_{i}^{(t)^{\prime}}$
$\displaystyle=\mu_{i}^{(t)}-\alpha^{-1}\big{[}\tilde{F}_{i}^{(t)}\big{]}^{-1}(\mu_{i})\partialderivative{\lambda_{i}}{\mu_{i}}\partialderivative{L}{\lambda_{i}}+\xi_{i}$
$\displaystyle=\mathbb{E}_{p_{i}(z)}[s(z)]+\xi_{i},$ (32)

be the mean parameters of $\tilde{p}_{i}^{(t)}$ after an EP inner update (4)
but _before_ being mapped to mean parameters of $\tilde{p}_{i}^{(t+1)}$
through map (9). $\xi_{i}$ is some zero-mean noise, e.g. due to MC variation.
The second equality follows from the results of Appendix C. Similarly, let
$\eta_{i}^{(t)^{\prime}}=\nabla A^{*}(\mu_{i}^{(t)^{\prime}})$ be the
corresponding natural parameters. Now apply map (9) to convert to mean
parameters of $\tilde{p}_{i}^{(t+1)}$,

$\displaystyle\mu_{i}^{(t+1)}$ $\displaystyle=\Big{(}\nabla
A\circ\eta_{i}^{(t+1)}\circ\big{(}\eta_{i}^{(t)}\big{)}^{-1}\circ\nabla
A^{*}\Big{)}\big{(}\mu_{i}^{(t)^{\prime}}\big{)}$ $\displaystyle=\nabla
A\bigg{(}\Big{(}\eta_{i}^{(t+1)}\circ\big{(}\eta_{i}^{(t)}\big{)}^{-1}\Big{)}\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)^{\prime}}\big{)}\Big{)}\bigg{)}$
$\displaystyle=\nabla A\bigg{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}+\alpha\sum_{j}\Big{(}\nabla
A^{*}\big{(}\mu_{j}^{(t)^{\prime}}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)}\Big{)}\bigg{)},$ (33)

where the last equality follows from the definition of $\eta_{i}^{(t)}(.)$,
and by observing that

$\displaystyle\lambda_{j}^{(t+1)}-\lambda_{j}^{(t)}$
$\displaystyle=\alpha(\eta_{j}^{\prime}-\eta_{j})$
$\displaystyle=\alpha\Big{(}A^{*}\big{(}\mu_{j}^{(t)^{\prime}}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)}\Big{)},$ (34)

for all $j$. Let $\bar{\mu}_{i}^{(t)^{\prime}}=\mathbb{E}_{p_{i}(z)}[s(z)]$ be
the mean parameters of $\tilde{p}_{i}^{(t)}$ after a _noise-free_ update (but
before mapping to those of $\tilde{p}_{i}^{(t+1)}$), so that
$\mu_{i}^{(t)^{\prime}}=\bar{\mu}_{i}^{(t)^{\prime}}+\xi_{i}$. Substituting
this in, we have

$\displaystyle\mu_{i}^{(t+1)}$ $\displaystyle=\nabla A\bigg{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}+\alpha\sum_{j}\Big{(}\nabla
A^{*}\big{(}\bar{\mu}_{j}^{(t)^{\prime}}+\xi_{j}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)}\Big{)}\bigg{)}.$ (35)

Taking a Taylor expansion around $\alpha=0$, we have

$\displaystyle\mu_{i}^{(t+1)}$
$\displaystyle=\mu_{i}^{(t)}+\alpha\left[\nabla^{2}A\bigg{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\bigg{)}\right]\left(\sum_{j}\nabla
A^{*}\big{(}\bar{\mu}_{j}^{(t)^{\prime}}+\xi_{j}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)}\right)+O(\alpha^{2}).$ (36)

Now subtract the noise-free update $\bar{\mu}_{i}^{(t+1)}$, and take
expectations with respect to $\xi_{j}$ for $j=1,\ldots,m$,

$\displaystyle\mathbb{E}\Big{[}\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}\Big{]}$
$\displaystyle=\mathbb{E}\Bigg{[}\alpha\left[\nabla^{2}A\bigg{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\bigg{)}\right]\bigg{(}$
$\displaystyle\qquad\sum_{j}\nabla
A^{*}\big{(}\bar{\mu}_{j}^{(t)^{\prime}}+\xi_{j}\big{)}-\nabla
A^{*}\big{(}\bar{\mu}_{j}^{(t)^{\prime}}\big{)}\bigg{)}+O(\alpha^{2})\Bigg{]}.$
(37)

the first term on the right hand side does not have zero expectation in
general—$\nabla A^{*}(.)$ is not linear—and so the bias is $O(\alpha)$ as
$\alpha\rightarrow 0^{+}$.

To see that the expected change in $L$ is also $O(\alpha)$, take the Taylor
expansion of $L$ along the update direction—obtained by subtracting
$\lambda_{i}$ from the right hand side of (4))—as a function of the step size
around $\alpha=0$,

$\displaystyle-\alpha\Big{(}\theta_{0}+\sum_{j\neq i}\lambda_{j}-\nabla
A\big{(}E_{p_{i}(z)}[s(z)]+\xi_{i}\big{)}\Big{)}^{\top}\partialderivative{L}{\lambda_{i}}+O(\alpha^{2}),$
(38)

which is clearly $O(\alpha)$ as $\alpha\rightarrow^{+}$. ∎

See 3

###### Proof.

Let $\eta_{i}^{(t)}$ and $\mu_{i}^{(t)}$ be the pre-update natural and mean
parameters of $\tilde{p}_{i}{(t)}$ respectively. Furthermore, let

$\displaystyle\mu_{i}^{(t)^{\prime}}$
$\displaystyle=\mu_{i}^{(t)}-\epsilon\Big{(}\nabla
A\big{(}\theta_{0}+\sum_{j}\lambda_{j}\big{)}-\mathbb{E}_{p_{i}(z)}[s(z)]-\xi_{i}\Big{)}$
$\displaystyle=\mu_{i}^{(t)}-\epsilon\Big{(}\mu_{i}^{(t)}-\mathbb{E}_{p_{i}(z)}[s(z)]-\xi_{i}\Big{)},$
(39)

be the mean parameters of $\tilde{p}_{i}^{(t)}$ after an EP-$\mu$ update (12)
but _before_ being mapped to mean parameters of $\tilde{p}_{i}^{(t)}$ through
(9). $\xi_{i}$ is some zero-mean noise, e.g. due to MC variation. The second
equality follows from the definition
$\mu_{i}^{(t)}=\mathbb{E}_{\tilde{p}_{i}^{(t)}(z)}[s(z)]$, and the fact that
$\tilde{p}_{i}^{(t)}=p$ before an update. Now apply map (9) to convert to mean
parameters of $\tilde{p}_{i}^{(t+1)}$,

$\displaystyle\mu_{i}^{(t+1)}$ $\displaystyle=\Big{(}\nabla
A\circ\eta_{i}^{(t+1)}\circ\big{(}\eta_{i}^{(t)}\big{)}^{-1}\circ\nabla
A^{*}\Big{)}\big{(}\mu_{i}^{(t)^{\prime}}\big{)}$ $\displaystyle=\nabla
A\bigg{(}\Big{(}\eta_{i}^{(t+1)}\circ\big{(}\eta_{i}^{(t)}\big{)}^{-1}\Big{)}\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)^{\prime}}\big{)}\Big{)}\bigg{)}$
$\displaystyle=\nabla A\bigg{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}+\sum_{j}\Big{(}\nabla
A^{*}\big{(}\mu_{j}^{(t)^{\prime}}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)}\Big{)}\bigg{)},$ (40)

where the last equality is obtained from the definition of $\eta_{i}^{(t)}(.)$
for $\alpha=1$, and using

$\displaystyle\lambda_{j}^{(t+1)}-\lambda_{j}^{(t)}$
$\displaystyle=\eta_{j}^{\prime}-\eta_{j}$
$\displaystyle=A^{*}\big{(}\mu_{j}^{(t)^{\prime}}\big{)}-\nabla
A^{*}\big{(}\mu_{j}^{(t)}\big{)},$ (41)

for all $j$. Substituting in the definition of $\mu_{j}^{(t)^{\prime}}$, we
have

$\displaystyle\mu_{i}^{(t+1)}=\nabla A\Bigg{(}$ $\displaystyle\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}+\sum_{j}\bigg{[}$ $\displaystyle\qquad\nabla
A^{*}\Big{(}(1-\epsilon)\mu_{j}^{(t)}+\epsilon\big{(}\mathbb{E}_{p_{j}(z)}[s(z)]+\xi_{j}\big{)}\Big{)}-\nabla
A^{*}\Big{(}\mu_{j}^{(t)}\Big{)}\bigg{]}\Bigg{)}.$ (42)

Taking Taylor expansions around $\epsilon=0$, we have

$\displaystyle\mu_{i}^{(t+1)}$
$\displaystyle=\mu_{i}^{(t)}+\epsilon\left[\nabla^{2}A\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\Big{)}\right]\sum_{j}\partialderivative{\epsilon}\Big{\\{}$
$\displaystyle\qquad\nabla
A^{*}\Big{(}(1-\epsilon)\mu_{j}^{(t)}+\epsilon\big{(}\mathbb{E}_{p_{j}(z)}[s(z)]+\xi_{j}\big{)}\Big{)}\Big{\\}}\Big{\rvert}_{0}+O(\epsilon^{2})$
$\displaystyle=\mu_{i}^{(t)}+\epsilon\left[\nabla^{2}A\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\Big{)}\right]\left[\nabla^{2}A^{*}\big{(}\mu_{i}^{(t)}\big{)}\right]\sum_{j}\partialderivative{\epsilon}\Big{\\{}$
$\displaystyle\qquad(1-\epsilon)\mu_{j}^{(t)}+\epsilon\big{(}\mathbb{E}_{p_{j}(z)}[s(z)]+\xi_{j}\big{)}\Big{\\}}\Big{\rvert}_{0}+O(\epsilon^{2})$
$\displaystyle=\mu_{i}^{(t)}+\epsilon\left[\nabla^{2}A\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\Big{)}\right]\left[\nabla^{2}A^{*}\big{(}\mu_{i}^{(t)}\big{)}\right]\sum_{j}\Big{(}$
$\displaystyle\qquad\mathbb{E}_{p_{j}(z)}[s(z)]+\xi_{j}-\mu_{j}^{(t)}\Big{)}+O(\epsilon^{2}).$
(43)

Finally, subtract the noise free $\bar{\mu}_{i}^{(t+1)}$, and take
expectations (over $\xi_{j}$ for $j=1,\ldots,m$) to obtain the bias

$\displaystyle\mathbb{E}\Big{[}\mu_{i}^{(t+1)}-\bar{\mu}_{i}^{(t+1)}\Big{]}$
$\displaystyle=\mathbb{E}\bigg{[}\epsilon\left[\nabla^{2}A\Big{(}\nabla
A^{*}\big{(}\mu_{i}^{(t)}\big{)}\Big{)}\right]\left[\nabla^{2}A^{*}\big{(}\mu_{i}^{(t)}\big{)}\right]\sum_{j}\xi_{j}+O(\epsilon^{2})\bigg{]}.$
(44)

By assumption $E[\xi_{j}]=\mathbf{0}\ \forall\ j$, hence the first order term
disappears, and the bias is $O(\epsilon^{2})$.

To see that the expected decrease in $L$ is also $O(\epsilon)$, note that we
are simply following the gradient in $L$ with respect to $\mu$, multiplied by
the inverse Fisher and scaled by $\epsilon$. Neither the reparameterisation
(from $\lambda_{i}$ to $\mu_{i}$) or the Fisher depend on $\epsilon$, and so
the change in $L$ must be $O(\epsilon)$ as $\epsilon\rightarrow 0^{+}$ by
simple Taylor expansion arguments. ∎

##  Appendix G Computational cost analysis

Let $c_{\text{samp}}$ be the cost of drawing a single sample from one of the
tilted distributions. Let $c_{\text{fwd}}$ be the cost of converting from
natural to mean parameters in the base family $\mathcal{F}$ (the forward
mapping). Similarly, let $c_{\text{bwd}}$ be the cost of converting from mean
to natural parameters in $\mathcal{F}$ (the backward mapping). Recall that $m$
is the number of sites. We now state the computational complexity of a round
of parallel updates in each of the EP variants evaluated in this paper.

#### EP

EP draws $n_{\text{samp}}$ samples from each of the $m$ sites. It then
performs $m$ backward mappings, to map from moments back to updated site
parameters. The total cost is then
$O(mn_{\text{samp}}c_{\text{samp}}+mc_{\text{bwd}})$

#### EP-$\bm{\eta}$

EP-$\eta$ draws $n_{\text{samp}}$ samples from each of the $m$ sites. Each
update also involves one forward mapping per iteration. In addition, each
iteration also requires $m$ JVPs through the backward mapping $\nabla
A^{*}(.)$. The primals of this JVP are the same for each site, so the
linearisation only needs to be performed once per update, costing
$O(2c_{\text{bwd}})$, but we have $m$ cotangents. The overall cost of an
EP-$\eta$ iteration is therefore
$O(mn_{\text{samp}}c_{\text{samp}}+c_{\text{fwd}}+(2+m)c_{\text{bwd}})$.

#### EP-$\bm{\mu}$

EP-$\mu$ draws $n_{\text{samp}}$ samples from each of the $m$ sites. Each
update involves one forward mapping, and $m$ backward mappings per iteration,
but notably does _not_ require any JVPs. The overall cost of an EP-$\mu$
iteration is therefore
$O(mn_{\text{samp}}c_{\text{samp}}+c_{\text{fwd}}+mc_{\text{bwd}})$.

#### SNEP

SNEP draws $n_{\text{samp}}$ samples from each of the $m$ sites, and performs
$m$ backward mappings per iteration. The total cost is then
$O(mn_{\text{samp}}c_{\text{samp}}+mc_{\text{bwd}})$, which is the same as EP.

In this work we largely assume that $c_{\text{samp}}$ is the dominant cost.
For example, in the NUTS sampler with default numpyro settings, (as used in
our experiments), each sample can involve evaluating up to 1024 gradients of
the tilted distribution log density.

The other costs, $c_{\text{fwd}}$ and $c_{\text{bwd}}$, depend on
$\mathcal{F}$. In the case of a MVN family with dense covariance, the foward
and backward mappings are $O(d^{3})$, for $d$ the dimensionality of $z$. When
$d$ is large, $c_{\text{fwd}}$ and $c_{\text{bwd}}$ become more significant,
in which case the balance will shift in favour of using more samples per
update (and taking larger steps). When $d$ is very large, however, it may be
necessary to choose a diagonal covariance base family instead, in which case
the mappings can be performed in $O(d)$ time. Also note that when $d$
increases, $c_{\text{samp}}$ will also typically increase at a rate faster
than $d$, even with optimal tuning [6]. In the main experiments, we used just
a single sample per update for EP-$\eta$ and EP-$\mu$, and they significantly
outperformed the baselines when measure in both NUTS steps _and_ wall-clock-
time.

We also note that the per-iteration cost of $c_{\text{fwd}}$ and
$c_{\text{bwd}}$ could be reduced by exploiting the results of previous
iterations. For example, in the case of a MVN with dense covariance, the cost
of the forward and backward mappings are dominated by matrix inversions. This
cost could be significantly reduced by using the result from the previous
iteration to warm-start a Newton-style iterative inversion routine. This
approach was used by Anil et al. [2] to reduce the cost of computing inverse
matrix roots as part of a wider deep learning optimisation scheme. We did not
make use of such optimisations in our experiments.

##  Appendix H Experiment details

In this appendix, we give details about the experiments presented in Section
4. To evaluate the performance of the different variants, we monitored KL
divergence to an estimate of the optimum, obtained by running EP to
convergence with a large number of samples ($n_{\text{samp}}=10^{5}$ for the
final iterations). We used 500 different hyperparameter settings for each
variant, chosen using random search, and repeated each run using 5 different
random seeds, which were used to seed the MCMC samplers. Hyperparameter
settings that failed in any of the 5 runs were discarded. We used this setup
for all experiments in Section 4.

The $x$-axis in Figure 2 of Section 4 corresponds to the number of NUTS steps,
or more specifically, the number of NUTS candidate evaluations. This number is
hardware and implementation agnostic, and roughly corresponds to the total
computational cost incurred up to that point. Wall-clock time results are
given in Appendix J.

#### Hyperparameter search spaces

For EP, the step size ($\alpha$) was drawn log-uniformly in the range
$(10^{-4},1)$. $n_{\text{samp}}$ was drawn log-uniformly between
$[d+2.5,10000.5)$ and then rounded to the nearest integer, where $d$ is the
dimensionality of $z$; note that the debiasing estimator used by Xu et al.
[47] for MVN families is not defined for $n_{\text{samp}}\leq d+2$. The
thinning ratio was drawn uniformly from $\\{1,2,3,4\\}$. For EP-$\eta$ and
EP-$\mu$, the step size ($\epsilon$) was drawn log-uniformly in the range
$(10^{-5},10^{-2})$, and $n_{\text{samp}}$ was fixed to $1$. For SNEP, the
step size was drawn log-uniformly in the range $(10^{-5},10^{-2})$.
$n_{\text{samp}}$ and $n_{\text{inner}}$ were (independently) drawn log-
uniformly in the range $[.5,10.5)$ and then rounded to the nearest integer.

#### Double-loops

All the methods considered can be used in a double-loop manner by taking
$n_{\text{inner}}>1$ inner steps per outer update. We did not find this
necessary in any of our experiments, and so we fixed $n_{\text{inner}}=1$ for
EP, EP-$\eta$ and EP-$\mu$. For SNEP we included $n_{\text{inner}}$ in the
search space (as described above) to make sure that setting it to $1$ was not
negatively impacting its performance.

#### Hardware

All experiments were executed on 76-core Dell PowerEdge C6520 servers, with
256GiB RAM, and dual Intel Xeon Platinum 8368Q (Ice Lake) 2.60GHz processors.
Each individual run was assigned to a single core.

#### Software

Implementations were written in JAX [5], with NUTS [24] used as the underlying
sampler. We used the numpyro [38] implementation of NUTS with default
settings. For experiments with a NIW base family we performed mean-to-natural
parameter conversions using the method of So [41], with JAXopt [4] used to
perform implicit differentiation through the iterative solve.

#### Sampling hyperparameter adaptation

We performed regular warmup phases in order to adapt the samplers to the
constantly evolving tilted distributions, consistent with prior work [47, 44].
The frequency and duration of these warmup phases was also included in the
hyperparameter search as follows. We drew the duration (number of samples) of
each warmup phase log-uniformly in the range $[99.5,1000.5)$ and then rounded
to the nearest integer. We then drew the sampling-to-warmup ratio log-
uniformly in the range $(1,4)$. This ratio then determined how frequently the
warmup was performed, which we rounded to the nearest positive integer number
of updates. Note that the Pareto frontier plots in Figures 2 and 6 include the
computation/time spent during warmup phases.

#### Site parameter initialisation

For EP, EP-$\eta$ and EP-$\mu$, we initialised the site parameters to
$\lambda_{i}=\mathbf{0}\ \forall\ i$ for all models, with the exception of the
cosmic radiation model for reasons we discuss later. SNEP is not compatible
with improper site potentials, and so for models where $\mathcal{F}$ was MVN
we used $\lambda_{i}=(2m)^{-1}\theta_{0}\ \forall\ i$, consistent with Vehtari
et al. [44]. Unfortunately this initialisation strategy for SNEP is not always
valid when $\mathcal{F}$ is NIW. This is because the site potentials in SNEP
are restricted to being proper distributions, and this constraint was violated
when using this initialisation in our experiments. We tried various other
initialisation strategies, but none produced satisfactory results or allowed
us to make a meaningful comparison, hence we omitted SNEP from the experiments
with NIW $\mathcal{F}$; namely, the political survey hierarchical logistic
regression and neural response model experiments.

$\mathcal{D}$$w$$z$$i=1,\ldots,m$ Figure 3: Directed graphical model for the
experiments of Section 4.

We set $\beta_{i}=1\ \forall\ i$ in all experiments, corresponding to the
original (not power) EP of Minka [33]. All of the models in these experiments
followed the same general structure, given by

$\displaystyle p_{0}(z)\prod_{i=1}^{m}p(w_{i}\mid z)p(\mathcal{D}_{i}\mid
w_{i}),$ (45)

where $p_{0}$ is a member of a tractable, minimal exponential family
$\mathcal{F}$. This is shown graphically in Figure 3. We now provide details
about individual experiments.

###  H.1 Hierarchical logistic regression

In these experiments, there were $m$ logistic regression problems, each with
$n$ covariate/response pairs, so that
$\mathcal{D}_{i}=\\{(x_{i,j},y_{i,j})\\}$, $x_{i,j}\in\mathbb{R}^{d}$, and
$y_{i,j}\in\\{0,1\\}$. Each regression problem had its own (unobserved) vector
of regression coefficients $w_{i}\in\mathbb{R}^{d}$, with prior density
$p(w_{i}\mid z)$ parameterised by global parameters $z$.

#### Synthetic data with MVN prior

In the synthetic experiment, we had $m=16$ regression problems with $d=4$ and
$n=20$. $\mathcal{F}$ was the family of MVN distributions, and
$z=(\mu_{1},\log\sigma_{1}^{2},...,\mu_{4},\log\sigma_{4}^{2})\in\mathbb{R}^{8}$
corresponded to the means and log-variances for the independent normal prior
with density $p(w_{i}\mid
z)=\prod_{j=1}^{4}\mathcal{N}(w_{i,j}\mid\mu_{j},\sigma_{j}^{2})$.888Note that
in this section we use $\mu$ to denote the mean of a normal or multivariate
normal distribution. This should not to be confused with the usage for
exponential family mean parameters in the main text. The prior on $z$ had
density
$p_{0}(z)=\mathcal{N}(\mathbf{0},\text{diag}\bm{(}(4,2,4,2,\ldots)\bm{)})$,
where $\text{diag}(.)$ constructs a diagonal matrix from its vector-valued
argument. The dataset was generated using the same procedure as Vehtari et al.
[44].

#### Political survey data with NIW prior

In the political survey data experiment, $\mathcal{F}$ was the family of
normal-inverse-Wishart (NIW) distributions, and $z=(\mu,\text{vech}(\Sigma))$
was used to parameterise the group-level coefficient prior with density
$p(w_{i}\mid z)=\mathcal{N}(w_{i}\mid\mu,\Sigma)$. $p_{0}$ was a NIW prior on
$\mu$ and $\Sigma$, with parameters $\mu_{0}=\mathbf{0}$, $\nu=9$,
$\lambda=1$, and $\Psi=I$ (following Wikipedia notation). We used a log-
Cholesky parameterisation of $\Sigma$ for the purposes of sampling. The
dataset consisted of binary responses to the statement "Allow employers to
decline coverage of abortions in insurance plans (Support / Oppose)". We
constructed $m=50$ regression problems, corresponding to the $50$ US states,
and truncated the data so that there were exactly $n=97$ responses for each
state. We used 6 predictor variables from the dataset, corresponding to
characteristics of a given survey participant. The predictors were binary
variables conveying: age (3 groups), ethicity (2 groups), education (3 groups)
and gender. We also included a state-level intercept to capture variation in
the base response level between states, so that $d=7$. This setup is based on
one used by Lopez-Martin et al. [29]. The data is available at
https://github.com/JuanLopezMartin/MRPCaseStudy.

###  H.2 Cosmic radiation model

In this experiment, there were $m$ nonlinear regression problems, each
corresponding to a model of the relationship between diffuse galactic far
ultraviolet radiation (FUV) and 100-µm infrared (i100) emission in a
particular section of the observable universe.

Each nonlinear regression model was parameterised by parameters
$w_{i}\in\mathbb{R}^{9}$. The $m$ regression problems were related through
common hyperparameters $z\in\mathbb{R}^{18}$, which parameterised the section-
level densities $p(\mathcal{D}_{i},w_{i}\mid z)$ where $\mathcal{D}_{i}$ is
the observed data for sector $i$. $\mathcal{F}$ was the family of MVN
distributions, and the prior on $z\in\mathbb{R}^{18}$ had density
$p_{0}(z)=\mathcal{N}(\mathbf{0},10I)$. The specifics of
$p(\mathcal{D}_{i},w_{i}\mid z)$ are quite involved and we refer the reader to
Vehtari et al. [44] for further details.

Vehtari et al. [44] applied this model to data obtained from the Galaxy
Evolution Explorer telescope. We were unable to obtain the dataset, and so we
generated synthetic data using hyperparameters parameters that were tuned by
hand to try and match the qualitative properties of the original data set (see
Appendix K for examples). We used a reduced number of $m=36$ sites and $n=200$
observations per site to reduce computation, allowing us to perform a
comprehensive hyperparameter search.

Using the conventional site parameter initialisation of
$\lambda_{i}=\mathbf{0}\ \forall i$ resulted in most runs of EP failing during
early iterations. We found that this was resolved by initialising with the
method used by Vehtari et al. [44] for SNEP, that is,
$\lambda_{i}=(2m)^{-1}\theta_{0}\ \forall i$, and so we used this
initialisation for all methods. We note, however, that the performances of
EP-$\eta$ and EP-$\mu$ were largely unaffected by this change.

###  H.3 Neural response model

In this experiment we performed inference in a hierarchical Bayesian neural
response model, using recordings of V1 complex cells in an anaesthetised adult
cat. 10 neurons in a specific area of cat V1 were simultaneously recorded
under the presentation of 18 different visual stimuli, each repeated 8 times,
for a total of 144 trials. Neural data were recorded by Tim Blanche in the
laboratory of Nicholas Swindale, University of British Columbia, and
downloaded from the NSF-funded CRCNS Data Sharing website
https://crcns.org/data-sets/vc/pvc-3 [3].

In this experiment $\mathcal{F}$ was the family of NIW distributions.
$z=\bm{(}\mu,\text{vech}(\Sigma)\bm{)}\in\mathbb{R}^{65}$ was used to
parameterise $\mathcal{N}(\log r_{j};\mu,\Sigma)$, the prior distribution
density for latent log firing rates in trial $j$, for $j=1,\ldots,144$.
$p_{0}$ was NIW with parameters $\mu_{0}=\mathbf{1}$, $\nu=12$, $\lambda=2.5$,
and $\Psi=1.25I$ (again following Wikipedia notation). The observed spike
count for neuron $k$ in trial $j$, $x_{j,k}\in\mathbb{N}$, was modelled as
$\text{Poisson}\bm{(}c_{j,k},\exp(r_{j,k})\bm{)}$, for $j=1,\ldots,144$,
$k=1,\ldots,10$. We again used a log-Cholesky parameterisation of $\Sigma$ for
sampling.

We grouped trials together into $m=8$ batches of $n=18$ trials, so that
$w_{i}=(\log r_{(i-1)180+1},\ldots,\log r_{(i-1)180+180})\in\mathbb{R}^{180}$
was the concatenation of (log) firing rates for all $10$ neurons across $18$
trials, with
$\mathcal{D}_{i}=\\{(x_{j,1},\ldots,x_{j,10})\\}_{j=(i-1)180+1}^{(i-1)180+180}$
the observed spike counts.

##  Appendix I Hyperparameter sensitivity

In this appendix, we examine the effect of varying hyperparameters of EP and
EP-$\eta$ on the synthetic hierarchical logistic regression experiment of
Section 4.1. The results for EP-$\mu$ are similar to those of EP-$\eta$ and so
we do not consider it separately here.

#### EP

In Figure 4a we show the effect of varying the number of samples used for
estimating updates in EP; more accurate regions of the frontier generally
require more samples. Figure 4b similarly shows the effect of varying the step
size; more accurate regions of the frontier generally correspond to a smaller
step size. Finally, Figure 4c shows the effect of varying the thinning ratio
$\tau$; more accurate regions of the frontier correspond to more thinning.
Together these plots illustrate the difficulty of tuning EP in stochastic
settings. Tracing out the frontier is a three-dimensional problem. To achieve
the best accuracy within the compute budget used for this experiment, the
practitioner would need to set $400<n_{\text{samp}}\leq 500$, $0.1<\alpha\leq
0.3$, and $\tau=2$; and any deviation from this would seemingly result in
suboptimal accuracy.

(a)  Number of samples

(b)  Step size

(c)  Thinning ratio

Figure 4: The effect of varying EP hyperparameters; partial Pareto frontiers
showing the number of NUTS steps ($x$-axis) against the KL divergence from $p$
to an estimate of the optimum ($y$-axis).

#### EP-$\bm{\eta}$

In Figure 5a we show the effect of varying the number of samples used for
estimating updates in EP-$\eta$; the frontier is traced out with
$n_{\text{samp}}=1$. We note however that the difference between
$n_{\text{samp}}=1$ and $n_{\text{samp}}=10$ is relatively small, and so it
may be sensible to choose $n_{\text{samp}}>1$ to make efficient use of
parallel hardware, or to minimise per-iteration overheads (see Appendix G).
Figure 5b similarly shows the effect of varying the step size; with more
accurate regions of the frontier corresponding to a smaller step size. This
also suggests that in practice it may make sense to set $\epsilon$ relatively
large at first and then gradually decay it to improve the accuracy. Finally,
Figure 5c shows the effect of varying $n_{\text{inner}}$, the number of inner
steps performed per outer update. $n_{\text{inner}}>1$ corresponds to using
EP-$\eta$ in "double-loop" mode. We did not find it necessary to increase
$n_{\text{inner}}$ above one to obtain convergence in our experiments, but
Figure 5c demonstrates that doing so would have a relatively small impact on
performance. Together these plots illustrate that hyperparameter tuning for
EP-$\eta$ is relatively straightforward. The frontier can be largely be traced
out by varying $\epsilon$, and is relatively insensitive to $n_{\text{samp}}$
and $n_{\text{inner}}$.

(a)  Number of samples

(b)  Step size

(c)  Number of inner steps

Figure 5: The effect of varying EP-$\eta$ hyperparameters; partial Pareto
frontiers showing the number of NUTS steps ($x$-axis) against the KL
divergence from $p$ to an estimate of the optimum ($y$-axis).

##  Appendix J Wall-clock time results

In Figure 6, we show Pareto frontiers, with the y-axis showing the lowest KL
divergence achieved by any hyperparameter setting, and the x-axis shows the
cumulative number of seconds elapsed; that is, the wall-clock time equivalent
of Figure 2. These results are in broad agreement with those of Figure 2,
which suggests the sampling cost does indeed dominate the computational
overheads of EP-$\eta$ and EP-$\mu$ in these experiments. The wall-clock time
results for EP-$\eta$ and EP-$\mu$ would likely be improved by using more than
one sample per update, by making more efficient use of hardware resources and
minimising per-iteration overheads.

We note that these times are necessarily implementation and hardware
dependent. We did not make particular efforts to optimise for wall-clock time.
In Appendix G we discuss approaches for minimising per-iteration overheads;
these would likely improve wall-clock time performance for all methods, but
should disproportionately favour EP-$\eta$ and EP-$\mu$, due to their frequent
updates and larger per-iteration overheads compared to EP and SNEP.

(a)  Hierarchical logistic regression: MVN prior

(b)  Hierarchical logistic regression: NIW prior

(c)  Cosmic radiation model

(d)  Neural response model

Figure 6: Pareto frontier showing the number of seconds elapsed ($x$-axis)
against the KL divergence from $p$ to an estimate of the optimum ($y$-axis).
Each point on the plot marks the lowest KL divergence attained by _any_
hyperparameter setting by that time. Error bars mark the full range of values
for the marked hyperparameter setting across 5 random seeds.

##  Appendix K Cosmic radiation data

Vehtari et al. [44] used a hierarchical Bayesian model to capture the
nonlinear relationship between diffuse galactic far ultraviolet radiation
(FUV) and 100-µm infrared emission (i100) in various sectors of the observable
universe, using data obtained from the Galaxy Evolution Explorer telescope. We
were unable to obtain the dataset, and so we generated synthetic data using
hyperparameters parameters that were tuned by hand to try and match the
qualitative properties of the original data set. Example data generated using
these hyperparameters is shown in Figure 7. For comparison, see Figure 9 of
Vehtari et al. [44].

Figure 7: Synthetic data, generated using the cosmic radiation model of
Vehtari et al. [44]. Each plot shows galactic far ultraviolet radiation (FUV)
versus infrared radiation (i100) for a single sector of the observable
universe.