Fearless Stochasticity in Expectation Propagation
Authors: Jonathan So
Richard E. Turner
Venue: NeurIPS 2024 spotlight (2024)
Keywords: expectation propagation
natural gradients
probabilistic methods
markov chain monte carlo
variational inference
Primary area: probabilistic_methods
TL;DR: Using natural gradients to improve handling of Monte Carlo noise in expectation propagation.

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.

Fearless Stochasticity in Expectation Propagation
Jonathan So University Cambridge j js2488@cam. ac.uk uk
Richard E. Turner University of Cambridge The Alan Turing Institute
Abstract
Expectation propagation (EP) is a family of algorithms for performing approxi- mate inference probabilistic models. The updates of EP involve the evaluation of moments srrrents-epeccattiion of certain functions wuhcc which can be estimated from Monte Carlo (MC) samples. However, the updates are not robust MC noise when performed naively, and various prior works have attempted to address this issue in different ways. this work, we provide a novel perspective on the noment-matching updates of EP; namely, that they perform natural-gradient-basse 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 sam- ple. These variants combine the benefits of their predecessors and address key weaknesses. In particular, they are easier to tune, offer an improved gpeed-accuracy trade-off, and do not rely the of debiasing estimators. We demonstrate their efficacy on a variety of probabilistic inference tasks.
Introduction
Expectation propagation (EP) [37, 40] is family of algorithms that is primarily used for performing
approximate inference probabilistic models [9, 10, 12, 13, 16, 18, 20, 19, 21, 22, 23, 24, 28, 29, 30, 32, 35, 36. 38, 39, 41, 44, 47, 48, 50, 52, 53], although it can be used more generally for approximating certain kinds of functions and their integrals [14].
EP involves the evaluation of cemments-expectatiiin derived from the model of interest. EP is usually applied 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 different ways [17, 48, 52).
of certain functions-under distributions that
In this work we provide a novel perspective on the moment moment-matching matching updates of EP: namely, that they perform natural- waalll-rrdddttbbaaee based optimization of a variational objective (Section 3). We use this insight to motivate two new EP variants, EP-r) (Section 3.2) and EP-p (Section 3.3), with updates that particularly well-suited MC estimation, remaining stable and being estimated with just single sample. These new variants combine the benefits of their predecessors particular, they are casier to tune, offer improved speed-accuracy
sample-efficient when
and address key weaknesses. trade-off, and do not rely on the use of debiasing estimators. We demonstrate their efficacy on 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 discussion of issues related sampled estimation of EP updates.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Let F be the tractable, minimal exponential family of distributions (see Appendix A), defined by the statistic function s(.) with respect to base measure v(.). Let Ω2 M, A(.) denote the natural domain, mean domain and log-partition function of F, respectively, and (*).) the convex dual of A(-).
Let po be the member of with natural parameter no. so that pz(2) exp(1) s(z) - A(mo)), and assume that a distribution of interest p", the target distribution, has density of the form
(²) (2(2)
Il exp (l,(z)).
(1)
In Bayesian inference settings, p' would be the posterior distribution over parameters given some observed data D, where Po may correspond to a prior distribution, and {f((z); to log-likelihood terms given some partition of 2 However, we consider the more general setting in which the target distribution is simply product factors. Note that we can assume form (1) without loss of generality if we allow Po be improper. The inference problem typically amounts computing 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 target distribution density of form (1), EP aims find an approximation p such that
p(2) x po(=) II, exp (A s(z)) = (z)
(2)
By assumption F is tractable, and provided that (m + +i Q p tractable member of F. Each factor exp(X, s(z)) is known as a sile potential, and can be roughly interpreted as a F-approximation to the i-th factor, exp(E,(2)). X is known as the i-th site parameter.
Variational problem EP, and several of its variants [17, 24, 37], can be viewed solving variational problem which we now introduce following the exposition Hasenclever et al. [17].
Let the i-th locally extended family, denoted F;, be the exponential family defined by the statistic function s2(2) (=))), (z()) with respect to base measure v1.). Let 94, M, and A,(.) denote the natural domain, mean domain, and log partition function of Fi, respectively. Note that member of Fi roughly corresponds distribution whose density the (normalised) product of a member of F with the i-th target factor raised to a power. Unlike F, we do not assume Fi is minimal.
Fixed points of EP correspond the solutions of the saddle-point problem
max min T(0,, -1.-.., Am), where BEN (A.):
L(O, AI. Ap..... Am) = A (m
Σ ) E,
-8,0000-''',, A(()).
(3)
At a solution to (3), the EP approximation is given by (2). The hyperparameters (3:); control the characteristics of the approximation, and correspond to the power parameters of power EP.
EP updates EP [37, 40]. power EP [34] and double-loop EP [24], can all be viewed as alternating between some number (2 1) of inner updates to decrease with respect to (Xi):. with an outer update to increase T. with respect to 0 [17, 25). EP is not typically presented in this way, but by doing so we will be able to present a unified algorithm that succinetly illustrates the relationship between the different variants. We show equivalence with the conventional presentation of EP in Appendix B. The inner and outer updates are given by
Inner update: As + A - a aam + E, A VA (B)(t)))())))
(4)
Outer update: 0 - 110 + Σ Ajr
(5)
We use e.g. to refer to distribution, and p(-) p(2) ) for its density, throughout.
GGoing forward, we assume and j to he taken from the index set {1,..., m) unless otherwise specified. This interpretation is not precise, since it is not necessarily the case that A E L.

---

where ps E F; denotes the i-th rilted distribution, with density
pc(z) x exp (10 A, ((z) + +,
(6)
The hyperparameter a controls the level of damping, which is used to aid convergence. An undamped inner update (with a 1) can be seen as performing moment matching between p and pt. The expectation in (4) is most often computed analytically, or estimated using deterministic numerical methods. It can also be estimated by sampling from pt. however, we will later show that the resulting stochasticity can lead to biased and unstable procedure. The inner updates can be performed either serially in parallel (over z). 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 [25] showed that (4) follows a decrease direction in with respect A and so is guaranteed decrease L when a small enough. See Appendix C for derivation of the inner update. When {A, h: are at partial minimum of L (for fixed 0), the outer update performs an exact partial maximisation of the primal form of the variational problem see Hasenclever et al. [17] for details. All of the EP variants considered in this paper differ only in their handling of the inner minimisation, and this is our focus.
Unified EP algorithm The double-loop EP algorithm of Heskes and Zoeter [24] 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 A (see Appendix B). however, Jylänki al. [30] observed that it can also be viewed as performing double-loop EP with just a single inner update per outer update (which is not guaranteed to converge in general). By taking this view, we are we are able to present EP, power EP, and their double-loop counterparts as single algorithm, presented in Algorithm 1. We do so primarily to illustrate how these variants are related to one another, and to the new variants of Section
Algorithm 1 EP (3,=1, Minner 1),, power EP
and their double-loop variants (Tlinner>>
Require: F, no {8.}.. (+)())), Duji Dinner, a while not converged do
Stochastic estimation of this expectation can lead to biased and unstable updates.
0 4 170 + ンi
for 1 to Ninner do
for i = to m in parallel do 2 x anno + Σラ VA (E)(())((2))))
return {h/
Note that dependence on (3.}% comes through the definition of p.(.). We can think of "exact" double-loop EP as corresponding to Timner oc, with the inner loop exiting once some convergence criterion has been satisfied. In practice, a truncated inner loop is often used by setting Tenner to some small number [17, 30]. Upon convergence of Algorithm 1, the approximation p(2) p²(2) is given by (2). Going forwards we will refer to the family of algorithms encompassed by Algorithm 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 used in black-box manner, dramatically expanding the set of models it can be applied to. Instead of performing single large sampling task as would be required by applying MC methods directly- EP can instead solve several simpler ones, gaining significant computational advantages [17, 48, 52]. Unfortunately, when performed naively, update (4) is not robust to MC noise. This is because the moment estimates are converted the natural (site) parameter space by mapping through V (.),., which is not linear in general, and so MC noise in the estimates leads to biased updates of X;-
3 Fearlessly stochastic EP algorithms
In this section we show an equivalence between the moment 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 review of related work.
3

---

3.1 Natural gradient view of EP
Several EP variants EP [40, 371, power EP, double loop EP [34], stochastic natural gradient EP [17], and the new ones to follow-dife only in how they perform the inner optimisation in (3) with respect to {. The optimisation can be viewed as one with respect to the m + distributions P.P1......P Pm. which are jointly parameterised by {Ar}e [17]. Natural gradient descent (NGD) [1], which involves preconditioning a gradient with the inverse Fisher information matrix (FIM) of a distribution, is effective method for optimising parameters of probability distributions. It would therefore seem 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 is the solution space, computing and inverting it will not generally be tractable, making (the natural extension of) NGD in this space infeasible.
Consider instead p;(0) E F with natural parameter v); (A.) Noo + Σ, + aA; 1,,,00+0~'000 The & superseript on the site parameters indicates that they are fixed for the current iteration, and so the distribution is fully parameterised by A₁ Note that when 'ラ (4) A j. we have D. = p. Proposition 1, which we prove in Appendix D, states that the noment-matching updates of EP can be viewed as performing NGD in L with respect to mean parameters of ()())
- 1))
We will make use of the following properties: for exponential family with log partition function A()), the gradient VA(.) provides the map from natural to mean parameters, and this mapping is one-to-one for minimal families, with the inverse given by VA (.). See Appendix A for details.
Proposition 1. For a > O, the moment- moment-matching update of EP (4) is equivalent to performing an NGD step in L with respect to the mean parameters of (t) with step size a=! ! That is, for H4 = -ep,0 [z)], and (14)) the FIM of pe with respect to μas we have
(14)
M4 a-
[ - (1))
aL 8p,
(VAOT))) ) ( -a-a ΣS DA-VA' (E(()(()))())) (7)
Note that the right hand side of (7) just maps update (4) mean parameters of p. (1) We discussed in Section 2 that noise in the moment estimates results in bias in the update to X of EP, due to the
noise being passed through the nonlinear map (.). This alone would not be a problem the
updates followed unbiased descent direction estimates in some other fixed parameterisation. Let T = VA* (Hx), then, from the definition of L and pi, we have
aL On: ai
(VA(To + Σ, -M)-Ep.ce - (p)))))())))
(8)
0p: ap: On:
and so when E [=)==]((=)]] is estimated with unbiased samples, the NGD update (7) in μt 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 (() mapped those of p, (I+1) by
p (++1)
(VAoT) (++1), o (1²")) VA") (n("),
where (1))))))())) is the inverse of 1 (.), and ny (.) is defined using the site parameters after the update at time L. This map is not linear in general, and so an unbiased, noisy update of i at one time step, results in bias when mapped to the next. This bias can make the EP updates unstable in the presence of MC noise, leading previous work to use methods, specific to F, for obtaining approximately unbiased natural parameter estimates from samples. However, even with these debiasing methods, relatively large number of samples is still needed in practice [48, 52]. Here we have assumed parallel application of (4), but the bias at site is induced whenever the parameters for other sites change from one update to the next, which occurs in both parallel and serial settings.
We now use the natural gradient interpretation of the EP updates 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.
*Note the apparent contradiction that by decreasing (increasing the damping) actually increase the NGD step size. This is resolved by observing that the definition of g! also changes with a.

---

Progress per 100 samples
Bias induced by single update
10
10
Expocted docrense
10-2
Biaos
10-1
10-3
10
10-4
10
10-0
10-5
10
10
10-2
10
0.0
0.2
0.4
0.6
0.8
1.0
Step size
Stop size
EP (naive)
(debiased)
EP-7
EF-
Teamp
Teamp
Thamp
The effect of step size (a or c) and number of MC samples (Msamp) on different EP variants in a stochastic version of the clutter problem of Minka [37]. EP (naive) uses maximum likelihood estimation for the updates, and EP (debiased) uses the estimator of Xu et al. [52]. Step size corresponds to a for EP, and for EP-p and EP-1 Only EP-r) and EP- can perform 1-sample updates, hence the other traces are not visible. The left panel shows the expected decrease I. after s0p steps. Performing e.g. 100x 1-sample steps, 10x 10-sample steps, achieves much larger decrease in L than a single 100-sample step. The right panel shows the magnitude of the bias A after a single parallel update, averaged over all sites and dimensions. The bias of EP-p shrinks far faster as the step size decreases than that of EP. EP-r is always unbiased and so is not visible.
Figure 1:
3.2 EP-r)
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 affine, the sequence would remain unbiased. This can be achieved by performing NGD) with respect 7/s the natural parameters of p. (f), Then, the map from one parameterisation to the next (by equating X₂) is given by m (++1) (n) (0) which linear. The NGD direction L with respect (((t) and n)s is given by
-[Fr(r]]
OL On Onp
(VA(T + E,,,) - -p(tt)((()))
(10)
Dp,
See Appendix E for a derivation. Note that a plays a similar role as the NGD step size in this parameterisation, but it also affects (t), and to obtain more faithful NGD interpretation we fix and introduce an explicit NGD step size to decouple from the effect on p0(0) The resulting
update can expressed directly as an update in A; by applying (1)( -1 the updated 7/s, giving
On:
A X, 0j4,
(VA(TO E, -Epp)Ep.
(11)
We can use automatic differentiation to efficiently compute (11), by recognising the second term as a Jacobian-vector product (JVP) with respect to T((44) = V (w:)) summary, by performing NGD with respect natural parameters of (,) instead of mean parameters, the resulting sequence of updates is unbiased in A₃ We call the resulting procedure whichh is given in Algorithm 2 2-EP-1 to emphasise that we have simply changed the NGD parameterisation of EP. The unbiased updates of EP-r allow it to be more sample-efficient than EP by using smaller number of samples per iteration. The left panel of Figure demonstrates this effect in stochastic version of the clutter problem [37].
3.3
EP-
The bias introduced by the updates of EP cannot be mitigated by reducing (increasing the damping) without also proportionally sacrificing the amount of progress made by each update. More concretely, let the bias in dimension d of μ4 (t+1), after an update at time f, be defined as E[u, (++1) P. (++1)) la- where p. (++1) is the value of μ. (t++1) after a noise-free update, and expectation taken over the sampling distributions of all parallel updates at time f. Then, Proposition 2 below, which we prove in Appendix G, summarises the effect of decreasing a on EP.
TThis be computed using either forward reverse mode automatic differentiation. on is symmetric. Bps
5

---

Proposition 2. After update (4) is executed in parallel over i, as - 0*, both the expected decrease in L. and the bias Elp, (++1) F."" (e(++))] I]a, are O(a) for all d.
Proposition at the beginning of this section states that update (4) can be viewed as performing an NGD step with respeet to p4 with a step size of a a= However, changing a also has the effect of changing the definition of p, (9) It is then natural to wonder what happens if we fix 1 and introduce an explicit step size for NGD, e, resulting in the update
M4 - P4 - ( A + -
(12)
See Appendix F for derivation. It turns out that in doing so, when we decrease the bias shrinks far faster than the expected decrease in L.
Proposition 3. After update (12) is executed in parallel over i, as 0+, the expected decrease in L is O(e), and the bias E [u, (1+1) - [i:(-11) is ((e²), for all d.
Proposition 3, which we prove in Appendix G, 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 A by applying (1)(0) I VA", giving
A < VA (11 (1-1)) ( + b,, + c & (p(())(())))
170
j Az.
(13)
which has the simple interpretation of performing EP, but with damping of the moments instead of the site (natural) parameters. In summary, if we perform NGD with respect to the same mean parameterisation as EP, but treat the NGD step size as a free parameter, e, we can obtain updates that are approximately unbiased while still making progress. We call this variant EP-p, again to indicate that we are simply performing the NGD update of EP in the mean parameter space. The resulting procedure is also given by Algorithm 2.
Algorithm 2 EP-1 and EP-u (differences with Algorithm I are highlighted in green)
Require: F, no. {Be k. (=)),. Di Minners while not converged do 0 < 1/0 + Σラ>;
for 1 to Vinnner do for i to m in parallel do update A, using (1I) for EP-1, or (13) for EP-μ
return {
The computational cost of EP-u is lower than that of EP-7 because it does not require JVPs through (.) The drawback is that the updates of EP-u do still retain some bias, however, we find that the bias of EP-u typically has negligible impact on its performance relative to EP-7). This is evident in the clutter problem of Minka [37], as demonstrated in the left panel of Figure 1, as well as in the larger scale experiments of Section 4. The right panel of Figure illustrates how quickly the bias in A shrinks as the step size of EP- is decreased. Note that the results of this subsection are stated in terms of bias in μú. but it is straightforward to show that equivalent results also hold for X; using Taylor series arguments.
3.4 Related work
The link between natural gradients and moment matching in an exponential family is well known [6], but the connection with EP shown here which relies on identification of the distribution p and parameterisation n7. (") - is new, the best of our knowledge. Bui et al. [11] showed that the updates of (power) EP are equivalent to performing NGD in a local variational free energy objective when taking the limit of the power parameter as B, - oc, and Wilkinson al. [51] showed that this extends to the (global) variational free energy objective for models with a certain structure. Our result different in that it relates to NGD of the EP variational objective, and applies for all values of 3,
'EP too is performing NGD in mean parameters, but using a "step size" that also affects the distribution p.. "See Appendix I for detailed discussion of the computational costs of EP-1) and EP-p.

---

EP was combined with Markov chain Monte Carlo (MCMC) moment estimates by Xu et al. [52] in a method called sampling via moment sharing (SMS), and later by Vehtari al. [48] in the context of hierarchical models. In both works, when F was multivariate normal (MVN) -arguably the most useful case-the authors used an estimator for the updates that is unbiased when P. is also MVN. 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 known number of independent samples. Samples drawn using MCMC methods generally autocorrelated, necessitating the use of sample thinning before the estimators can be applied. This adds another hyperparameter to the procedure- the thinning ratio- 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), 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. [17], which is closely related to our work, also optimises the EP variational objective using a form of NGD. The SNEP updates are unbiased in the presence of MC noise, allowing them to be performed with little as one sample, without relying on F-specific debiasing estimators, or the sample thinning that would entail. In SNEP, NGD is performed with respect to mean parameters of the site potentials, which are treated as bona fide distributions in F. In contrast, we showed that EP can already be viewed as performing NGD, but with respect to the distributions ("), and our new variants, EP-1 and EP-14, able to gain the same advantages as SNEP, but using the same distributions for NGD as EP. Hasenclever et al. [17] showed that in some settings, SNEP can obtain accurate point estimates fairly rapidly. However, we find that it typically converges far slower than both EP and our new variants, consistent with findings in Vehtari al. [48]. We argue that this is because the site potentials (when considered as distributions) bear little resemblance with the distributions that are ultimately being optimised, and so their geometry is largely irrelevant for the optimisation of L. In contrast, the geometry of p, (t) is closely related to that of L, as we show in Appendix H.
We note that Xu et al. [52] and Hasenclever al. [17] 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-1 and EP-H 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 D. All of the models in these experiments followed the same general structure, consisting of a minimal exponential family prior over 2, Po, and partition of the data (m where each block D, depends on both and an additional vector of local latent variables wA that also depend on 2. That is, the joint density has the form
Po(z) II, p(u; 2)p(D;: Ws 2),
(14)
where po F. This structure is shown graphically in Appendix J. To perform approximate inference of z, we first define (=(z) as the log likelihood of D; given with w; marginalised out. That is,
(=)) = log / p(D,,un) =))d;
(15)
Then, given Po(2) and (4(=)), we define p²(z) as (1), and proceed to find an approximation p(z) R p" (z) using the methods described in earlier sections. Note that sampling from the tilted distribution (z( requires jointly sampling over and w; and taking the marginal. EP is 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 [48]. In our experiments we used NUTS [27] to perform the sampling, consistent with prior work [48, 52]. In each experiment we compared EP-1 and EP-u with
7

---

Hierarchical logistic regression (MVN prior)
Hierarchical logistic regression (NIW prior)
10
Divergence (nata)
10'
10"
10"
10-1
10
10?
10
Cosmii radiation model
Noural response model
10a
Divengence (sats)
10'
10
10
10
Steps
Steps
EP
EP-1)
EP-4
SNEP
Figure 2: Pareto frontiers showing the number of NUTS steps (r-axis) against the KL divergence
from p to an estimate of the optimum (y-axis). Each point on the plot marks the lowest average KI. divergence attained by any hyperparameter setting by that step count. Error bars mark the full range of values for the marked hyperparameter setting across random seeds.
EP, the manner used by Xu al. [52] and Vehtari al. [48]. When JF was MVN we also compared to SNEP [17]. For the models with inveree Wishart (NIW) F, we were unable to find an initialisation for SNEP that would allow to perform meaningful comparison, We discuss this point further in Appendix J.
To evaluate the performance of the different variants, monitored KL divergence 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-7 and EP-a were performed with ssmmp = 1 and Pinner 1. The results of our evaluation are summarised by the Pareto frontiers in Figure 2.
estimate
We see that EP- and EP- can typically reach given level of accuracy (distance from an estimate of the optimum) faster than EP, often significantly We stress that the frontiers for EP-n and EP-u traced out with namp and Tiinner each fixed to 1, with no sample thinning, varying only the step size hyperparameter with higher accuracy/cost regions of the frontier corresponding lower values of In contrast, to trace out the frontier of EP we must jointly vary a, Rsamp. and the thinning ratio. See Appendix K for examples of this effect. The performances of EP-1 and EP- are very similar. We found SNEP to be significantly slower than the other methods, consistent with findings by Vehtari al. [48]. Pareto frontiers with respect wall-clock time can be found in Appendix L. Code for these experiments can be found at https: //gi /ggtttb....//./// com/ htt//rmmbrieeee laaeeemmlgffareeeeeeeeeeeee now provide a brief overview of individual experiments, with further details given in Appendix
ep. We
Hierarchical logistic regression with MVN prior Hierarchical logistic regression (HLR) used to perform binary classification in number of groups when the extent to which data should be pooled across groups unknown [15]. We performed approximate Bayesian inference in HLR model using MVN prior over the group-level parameters. We applied this model to synthetic data, generated using the same procedure as Vehtari al. [48], which was designed to be challenging for EP. this experiment E R$, w; IRd V i, and each of the upper-left panel of Figure we see that EP-7) and EP-p typically reach a given level of accuracy faster than EP and SNEP.
16 groups had n 20 observations. In the

---

Hierarchical logistic regression with NIW prior We also performed approximate Bayesian inference in a similar HLR model, but using normal-inverse Wishart (NIW) prior over group-level parameters, allowing for correlation of regression coefficients within groups. We applied this model to political survey data, where each of m = 50 groups corresponded to responses from particular US state, with 7 predictor variables corresponding to characteristics of given survey respondent, so that Wi E IR? V i. There were n = 97 survey respondents per state. The data, taken from the 2018 Cooperative Congressional Election Study, related to support for allowing employers to decline coverage of abortions in insurance plans [33, 43). The upper-right panel of Figure 2 shows that EP-1) and E consistently reach a given level of aceuracy significantly faster than EP.
Cosmic radiation model A hierarchical Bayesian model was used by Vehtari et al. [48] to capture the nonlinear relationship between diffuse galactic far ultraviolet radiation and 100-um infrared emission in various sectors of the observable universe, using data obtained from the Galaxy Evolution Explorer telescope. In this model each w, E R 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 E RR18, which parameterised the section-level parameter densities. The prior, Po E F, was MVN. The specifics of the nonlinear regression model are quite involved and we refer the reader to Vehtari et al. [48] for details. We were unable to obtain the dataset, and so we generated synthetic data using parameters that were tuned by hand to try and match the qualitative properties of the original data - see Appendix N for examples. We used a reduced number of m = 36 sites and n 200 observations per site to allow us to perform comprehensive hyperparameter search. The lower-left panel of Figure 2 shows once again that EP-r) and EP-u consistently reach a given level of accuracy significantly faster than EP and SNEP.
Neural response model A common task in neuroscience is model the firing rates of neurons under various conditions. We performed inference in hierarchical Bayesian neural response model, using recordings of V1 complex cells in an anesthesised adult cat [4]. 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 the 10 neurons in each trial as Poisson, with latent (log) intensities being MVN, with mean and covariance (collectively = E TR drawn from a NIW prior po E F. Inference of z amounted to inferring the means, and inter-neuron covariances, 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; E R match or beat the speed of EP for reaching any given level of accuracy. For high levels of accuracy the performances of the methods appear to converge with one another, but note that the asymptote is not exactly zero because the optimal solution was estimated by running EP with a large (but finite) number of samples.
The lower-right panel of Figure 2 shows that EP-r) and EP- either
5 Limitations
In this work we have 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 I, along with strategies for reducing computational overheads. We note, however, that wall-clock time results for evaluation experiments (in Appendix L) are in line with the results of Section 4, indicating that the drawing of MC samples was indeed the dominant cost.
When EP and its variants-includding those introduced here-are 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-1) and EP-μ 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 (for all EP variants) is still greater than that for 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.
Our focus has been on developing improved methods for performing EP when the updates must be estimated with noise, with the motivation for this kind of setting having been discussed by prior

---

NUTS EP-1 kernel
"Oracle" EP-r kernel
Pairwise posterior marginals
Divergence (uale)
10a
10"
bogs
10a
10
10'
10a
10a
10
10
log 02
A4 MCMC
Time (+)
Independent samples
EP-1
Figure 3: Comparison BP-1 with conjugate-compptation variational inference (CVI) on hierarchi- cal logistic regression model. The two leftmost plots show forward (solid) and reverse (dashed) KL divergences between the approximation of each method and a MVN distribution estimated directly from MCMC samples. The left panel shows comparison with respect to wall-clock time, when NUTS is used as the underlying sampling kernel for EP-7). The left-middle panel shows a similar comparison, but with respect to the number of samples drawn, and using "oracle" sampling kernel for EP-r- Hyperparameters were tuned for each method and shaded regions show the range of trajectories across 5 random seeds. The right and right-middle panels show pairwise marginals of the various MVN approximations overlaid on contours the true posterior. Coloured dots and ellipses correspond means and 2-standard-deviate contours, respectively. See Section for discussion of these results, and Appendix M for further details.
work [52, 17, 48). EP and its variants are agmostic to the choice sampling kernel, and as such we have not discussed issues relating to the performance of the underlying samplers. In the left panel of Figure however, we highlight the relative inefficiency of using EP-r with NUTS as the underlying sampler when compared to
variational inference (CVI), efficient variational inference method [31]. While EP-r is able to obtain much more faithful posterior approximations than CVI, it is significantly slower when measured in wall-clock time, despite the gains made over its predecessors. However, this due to the relative cost of drawing independent samples with NUTS, the sample efficiency of EP-q is roughly equivalent that of CVI when an (illustrative) "oracle" sampling kernel is used, as demonstrated in the left-middle panel of Figure 3. Further discussion and details of the experiments behind Figure 3 can be found in Appendix M.
Future work
In Section 5, we demonstrated that the efficiency EP variants is constrained by that of of the underlying samplers. There are several ways in which their performance could be improved by using knowledge of p to guide sampling from pt. One such improvement would to use p to set MCMC hyperparameters directly, obviating the need to adapt them during warm-up phases, e.g. when p E F MVN, its precision matrix could used directly as the mass matrix in Hamiltonian Monte Carlo schemes. Another potential improvement is that samples from could be used to re- re-initialise MCMC chains at each iteration, providing approximate independence between updates for small additional cost. It would also be worthwhile considering how our methods on modern compute hardware to better take advantage of the inherently parallel nature of EP, e.g. practical performance could be improved by using sampling kernels that themselves designed for efficient parallel execution [26]. Performance may be further improved in distributed settings by using asynchronous extensions of EP in order to minimise communication-related overheads and delays [17, 52]. We leave investigation of these ideas to future work.
be efficiently implemented
7 Conclusion
In this work, we used 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 casier to tune.
10

---

8 Acknowledgements
We would like thank the anonymous reviewers for providing valuable feedback, and Jihao Andreas Lin for sharing his plotting expertise. Jonathan So is supported by the University of Cambridge Harding Distinguished Postgraduate Scholars Programme, Richard E. Turner supported by Google, Amazon, ARM, Improbable, EPSRC grant EP/T005386/1, and the EPSRC Probabilistic AI Hub (ProbAI, EP/Y028783/1).
References
[1] S.-i. Amari. Natural gradient works efficiently in learning. Neural Computation, 1998.
[2] R. Anil, V. Gupta, T. Koren, K. Regan, and Singer. Scalable second order optimization for deep learning, arXiv:2002.090: 090018, 2021.
[3] C. M. Bishop. Pattern Recognition and Machine Learning. Springer, 2007.
[4] T. Blanche. Multi-neuron recordings in primary visual cortex., 2009.
[5] M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer, F. Llinares- F. Pedregosa, and
J.PP. Vert. Efficient and modular implicit differentiation. arXiv:2105. /5/03, 2021.
[6] A. Bouchard-Côté, Campbell, G. Pleiss, and Surjanovie. MCMC-driven learning. arXiv:2402.09598, 2024.
[7] J. Bradbury, R. Frostig. P. Hawkins, J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke,
J. VanderPlas, S. Wanderman- and Q. Zhang. JAX: composable transformations of Python+NummPy programs, 2018.
[8] S. Brooks, A. Gelman, G. Jones, and X.-L. Meng. Handbook of Markov Chain Monte Carlo. CRC Press, 2011.
[9] T. Bui, D. Hernandez-oobato, J. Hernandez-Lobtto Y. Li, and R. Turner. Deep Gaussian processes for regression using approximate expectation propagation. International Conference on Machine Learming, 2016.
[10] T. D. Bui, J. Yan, and R. E. Turner. A unifying framework for Gaussian process pseudo-poin approximations using power expectation propagation. Journal of Machine Learning Research, 2017.
[11] T. D. Bui, C. V. Nguyen, S. Swaroop, and R. E. Turner. Partitioned variational inference: A unified framework encompassing federated and continual learning- 18X.v 2018.
[12] W. Chu and Z. Ghahramani. Gaussian processes for ordinal regression. Journal of Machine Learning Research, 2005.
[13] J. P. Cunningham, K. V. Shenoy, and M. Sahani. Fast Gaussian process methods for point process intensity estimation. International Conference on Machine Learning, 2008
[14] J. P. Cunningham, P. Hennig, and S. Lacoste-Julien. Gaussian probabilities and expectation propagation. //1.::///1...... 2013.
[15] A. Gelman, J. Carlin, H. Stern, D. Dunson, A. Vehtari, and D. Rubin. Bayesian Data Analysis, Third Edition. Chapman & Hall / CRC Press, 2013.
[16] J. Gonzalez, M. Osborne, and N. Lawrence. Glasses: Relieving the myopia of Bayesian optimisation. International Conference on Artificial Intelligence and Statistics, 2016.
[17] I. Hasenclever, Webb, T. Lienart, S. Vollmer, Lakshhinarayynaan Blundell, and Y. W. Teh. Distributed Bayesian learning with stochastic natural gradient expectation propagation and the posterior server. Journal of Machine Learning Research, 2017.
[18] Hennig. Expectation propagation on the maximum correlated normal variables. Technical report, University of Cambridge, 2009.

---

[19] P. Hennig and C. J. Schuler. Entropy search for cetemmattio--ffiiiennee global optimization. Journal of Machine Learning Research, 2012.
[20] Hennig, D. Stern, and T. Graepel. Coherent inference on optimal play in game trees. International Conference on Artificial Intelligence and Statistics, 2010.
[21] R. Herbrich, T. Minka, and T. Graepel. TrueSkillMM; A Bayesian skill rating system. Advances in Neural Information Processing Systems, 2006.
[22] D. Hernandez-Lobate and J. M. Hernandez-Lobail Scalable Gaussian process classification via expectation propagation. International Conference on Artificial Intelligence and Statistics, 2016.
[23] J. M. Hernández- Lobato and R. P. Adams. Probabilistic backpropagation for scalable learning of Bayesian neural networks. International Conference on Machine Learning. 2015.
[24] T. Heskes and O. Zoeter. Expectation propagation for approximate inference in dynamic Bayesian networks. Uncertainty in Artificial Intelligence, 2002.
[25] T. Heskes and O. Zoeter. Extended version of "Expectation propagation for approximate inference in dynamic Bayesian networks". Technical report, University of Nijmegen, 2003.
[26] M. Hoffman, A. Radul, and P. Sountsov. An adaptive-MCMC scheme for setting trajectory lengths in Hamiltonian Monte Carlo. International Conference Artificial Intelligence and Statistics, 2021.
[27] 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.
[28] A. F. R. independent component analysis. Neural Computation, 2002.
O. Winther, and L. K. Hansen. Mean-fiell approaches to
[29] W. Jitkrittum, A. Gretton, N. M. O. Heess, S.M.A 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.
[30] P. Jyläänki, J. Vanhatalo, and A. Vehtari. Robust Gaussian process regression with Student-t likelihood. Journal of Machine Learning Research, 2011.
[31] M. Khan and W. Lin. Conjugate-commputtionn variational inference: : Converting variational in- ference in non-conjugate models inferences in conjugate models. In International Conference on Artificial Intelligence and Statistics, 2017.
[32] Y. Li, M. Hemánneez-obatto, and R. E. Turner. Stochastic expectation propagation. Advances in Neural Information Processing Systems, 2015.
[33] Lorez-Martin, J. Phillips, and Gelman. Multilevel regression and poststratification case studies, 2022.
[34] T. Minka. Power EP. Technical report, Microsoft, 2004.
[35] T. Minka and J. Lafferty. Expectation propagation for the generative aspect model. Uncertainty in Artificial Intelligence, 2002.
[36] T. Minka, R. Cleven, and Y. Zaykov. TrueSkill 2: An improved Bayesian skill rating system. Technical report, Microsoft, 2018
[37] T. P. Minka. A family of algorithms for approximate Bayesian inference. PhD thesis, Mas- sachusetts Institute of Technology, 2001.
[38] A. Naish-guzman and S. Holden. The generalized FTTC approximation. Advances in Neural Information Processing Systems, 2007.
[39] U. Nodelman, Koller, and C. Shelton. Expectation propagation for continuous time Bayesian networks. Uncertainty in Artificial Intelligence, 2005.
12

---

[40] M. Opper and O. Winther. Gaussian processes for classification: Mean-fiell algorithms. Neural Computation, 2000.
[41] M. Opper, U. Paquet, and O. Winther. Improving on expectation propagation. Advances in Neural Information Processing Systems, 2008
[42] D. Phan, N. Pradhan, and M. Jankowiak. Composable effects for flexible and accelerated probabilistic programming in numpyro, 2019.
[43] B. Schaffner, S. Ansolabehere, and S. Luks. CCES common content 2018, 2019.
[44] M. Seeger and H. Nickisch. Fast convergent algorithms for expectation propagation approximate Bayesian inference. International Conference on Artificial Intelligence and Statistics, 2011.
[45] So. Estimating the normal-inverse- Wishart distribution. arXiv:2405. 16088, 2024.
[46] So and R. E. Turner. Optimising distributions with natural gradient surrogates. International Conference Artificial Intelligence and Statistics, 2024.
[47] R. E. Turner and M. Sahani. Probabilistic amplitude and frequency demodulation. Advances in Neural Information Processing Systems, 2011.
[48] A. Vehtari, A. Gelman, T. Sivula, P. Jylänki, D. Tran, S. Sahai, P. Blomstedt, J. P. Cunningham,
D. Schiminovich, and P. Robert. Expectation propagation as way of life: A framework for Bayesian inference on partitioned data. Journal of Machine Learning Research, 2020.
[49] M. Wainwright and M. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 2008
[50] M. Welling and J. J. Lim. A distributed message passing algorithm for sensor localization. International Conference Artificial Neural Networks, 2007.
[51] W. J. Wilkinson, S. Särkkä, and A. Solin. Bayes-Newton methods for approximate Bayesian inference with PSD guarantees. Journal of Machine Learning Research, 2023.
[52] 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.
[53] O. Zoeter and T. Heskes. Gaussian quadrature based expectation propagation. Proceedings of the Tenth International Workshop on Artificial Intelligence and Statistics, 2005.
13

---

A Exponential families of distributions
In this appendix, we give very brief overview of exponential families of distributions. This is largely a condensed summary of relevant material from Wainwright and Jordan [49], and we refer the reader to the original work for a far more detailed treatment.
The exponential family of distributions F, defined by the d-dimensional, vector-valued sufficient statistic function s(.), and base measure v(.), has density function
p²=) = exp(1) T s(z) - A(1)),
(16)
taken with respect to v(.), where A(1) = log fex((1) + [xxxpp]](==)xddddd s(=))du(z) is known as the log-partition
function. 17 9 are known as the natural parameters, and are used to index specific member of F. is the set of normalisable natural parameters F, given by
{ E Rd / eexp(1 (oo0nnTs()))d//// s(z))du(=) 00 )
(17)
and is known as the natural domain of F. P is a convex set, and when it is open, the family F is said be regular we shall only consider regular families in this work. When the components of s(.) are linearly independent, F is said to be minimal. For minimal families, each distribution in the family is associated with a unique natural parameter vector 7). An exponential family that not minimal said to be overcomplete, in which case, each distribution is associated with an entire affine subspace of Ω.
The expected sufficient statistics of distribution with density p(.) with respect to base measure vv.). are given by [=)]] Let M be defined as the set of expected sufficient statistics that can be attained by any density p(-) with respect to base measure v.)) that is,
3p(.): ff p(3))(2)dr(()= =
M 4 EERRR
(18)
All vectors in the interior of .., MM, realisable by a member of F, and μ M provides an alternative parameterisation of F, in which are known as the mean parameters, and M is known as the mean domain..
For minimal families, A(.) is strictly convex on Ω. The convex dual of A(.) is defined as
(m) = sup rj T ( 7EQ
(19)
and on MM, (1) is equal to the negative entropy of the member of F with mean parameter ja.
The mean parameters of the member of F with natural parameter 1) can be obtained by μ VA(m).. For minimal families, this mapping is one-to-one, and the reverse map is given by 1 = VA*())) For this reason, V-(-) and (.) are sometimes referred to as the forward mapping and backward mapping respectively. In this work, we say that a family F is traciable 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 17. is given by, Fn) V A2)(. Furthermore, in minimal families, the FIM with respect to mean parameters y is given by F(x) = (1) Using the forward and backward mappings, we also have
0p
F(g)
(20)
On
Ōn F(R) Fp'
(21)
and from the inverse function theorem,
3n
F(n)
(22)
F(R)
(23)
Om
14

---

B Conventional presentation of EP updates
In this appendix give more conventional presentation of the EP updates, and show equivalence with the atypical presentation of Section 2. We will assume here that 3; = 1 A i, but similar reasoning be used show equivalence in the power EP case, where 83 * 1. Note that the update of power EP is often presented as having scale proportional to St with a separate damping parameter. We subsume both into a, with a corresponding the default damping suggested by Minka [34].
Given a target distribution with the form
(=) exp (no s())) IIT exp (l((z))
(24)
EP attempts to find an approximation p E F such that
p(=) x exp (no s(=)) II, exp exp (x]) p² (z).
(25)
The i-th site parameter X; can loosely be interpreted as the natural parameters of a F-approximation to the i-th target factor, exp((,(=)) EP is also often used to approximate the normalising constant of (x)) we have omitted details here, but see c.g. Bishop [3].
In order to update A₃ EP performs a KL--divereence minimisation between p(z) and local approxi- mation the target distribution, consisting of just one "real" factor and (m - 1) approximate ones. Specifically,
A f- arg min KL p(=;
(26)
A
where
(=()) x exp (c(()) II,
exp (A] s()))
(27)
jtt
and we have used the notation p(2; X,) make the dependence on A explicit. It can be easily shown that the solution to (26) is found by moment matching. That is, the KL-divergence minimised when
[p)=))))((<=] (=))))((=))..
(28)
For a minimal exponential family F. there is unique member of FF with mean parameters (p)())((=))}, and its natural parameters are given by V (B))(s))(()))). By equating this with the parameters of p(2), we have
== (F)()s))))))
Σ
(29)
jzi
Update (29) is the "standard" update of EP, and can be performed either serially or in parallel (over i). Often damping factor is used, giving
A a)>, + (VA' (EA(((() 7o -
(30)
In Section 2 we presented EP and several variants as performing some number of "inner" updates, to decrease with respect to X; with "outer" update with respect to 6. These updates are given by
Inner update: X; + A a n10 + Σ, A3 VA' (Bp((t))))s(())))
(31)
Outer update: 0+ 110 + Az.
(32)
Immediately after outer update we have that 0 170 E, Aj. and so for 3, we have
(z(2) x exp ((0 - X1) s(2) + +x(2))
= exp (no + E )1) 5(2) + f((2)
X (2()).
(33)
We stated in Section 2 that (non-double-loop) can be viewed as performing single inner update outer update, and so by rolling (31) and (32) into single update for X; we have
A f A - a a(10
Σ X VA (Fp,,c2))(())))
100 +
(1 a) VA' (Bnncs)(s(2)) -no -E
(34)
which is identical to (30).
15

---

C EP inner update derivation
Begin by taking gradients of (3), with respect to A, giving
A Am) VA
nu +, B, B),
(35)
At a minimum of L with respect to A, both sides of (35) must be equal to zero. We can use this to define fixed-point condition with respect to Xs and its updated value X!
sزر 2/) VA To X + ΣStt
20== - A, 3, 8,, 1,
=Ep.(=)(s(2))
(36)
where p; E F;;; denotes the i-th tilted distribution, defined by (6). Applying VA*(.), the inverse of VA(.), to both sides of (36), and rearranging terms, recover the moment-matching update of EP
A - VA (B) no
(37)
Often damping parameter is used aid convergence, giving the more general update
A (1 a) (VA (E,. (Ep,())(s(2)))) 170 E}ps>a) N))
\ + a VA (Ep n10 ,,,,)
(38)
where the level of damping is given by (1 a). It was shown by Heskes and Zoeter [25] that (4) follows a decrease direction in with respect Ao. and so guaranteed to decrease when is small enough.
D Natural gradient view of EP
+a (9 Let D. (6) be the member of with natural parameter 77, (A,) = 7o0 + -,,,0000'0 Proposition states that that the moment-matching updates of EP can be viewed performing NGD in L with respect to the mean parameters, μt,, of p; (6), 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 (t))) The left hand side of (7) is simply a NGD update in pę-
Proposition 1. For a > O, the moment- an NGD step in L with respect to the mean parameters of D with step size x That is, for μ4 = (s [z)),, and (1a)) the FIM of p. with respect to μss we have
moment-matching matching update of EP (4) is equivalent to performing
#4 a- (E
0L
(VAon)")
a Σ, 23 - VA VA' (E)(((c))((3)))) (7)
0p:
Proof. Taking the left hand side of (7), we have
#4 - a ((((1 (Hi)) -1 QL,
-1 OA AL Inn Oii
p4 a -((((()) (ui)
Om, OΛ: ØL
ph a
Dns On: Od:
aL
= 14 0x;
14 (VA Jo E,,, E (x)=))((()))
= #4 (m - (s)(s))(())))
[s()))
(39)
where the penultimate equality follows from the definition μ4: = R0 (2) [ [s(=)], and the fact that 2, = = v before the update. All that remains is to show that the right hand side of (7) is equal
16

---

to (x))))((22).
(VAom) {1 - aa + +,, A VA (B()e(ccsss())))
VA a (Ax-1()) + (Ep(s
(40)
E
EP-n update direction derivation
Using the relation = -m² 17,² and therefore A; = (0)(10))((11)
we have
-1 aL
= [((m))]
022 OL
Om:
Omi 0d:
a Ope (((P))
8p4
-1
ØT.
Om
Om,
04
( an: F: (p())
dmp
0L.
api
0p:
a app On: 0ps 0d,
On, 0pg Om:
OL
Om
(V + E, A1) - [Ex(t)l((>))
(41)
0pt
EP-p update derivation
a = 1, we have
N4 - -(Fi(us))
0L Opt;
p4 (F(lue)
1 On: Od, ÔL
Tyn Ōij O%:
O4:
Om Od,
(V + E, - -Ep(t((c))s(()))
p4 8pm!
8u, On
14 V. + E, A1) - (x(c))((22))
(42)
G Bias of EP and EP-μ
Let the bias in dimension d of μu (++1), after an update at time I, be defined as E[ - نū, (1+1) a where p. (++1) is the value μ (t-1) after a noise-free update, and expectation 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 a and on EP and EP-u, respectively.
Proposition 2. After update (4) executed in parallel over i, as -+ 0+, both the expected decrease in L, and the bias Elu; (t+1). f2,11 (t+1) Jd. are O for all d.
Proof. Let n (1) and uel (0) be the pre-update natural and mean parameters of pr (9 respectively. Furthermore, let
14 ((0) -(( - -a- (:(00 (mo) (14)) Ob2 ØL
- -1
Dp, 0x, + +
-Ep((e)((()) + Ei.
(43)
be the mean parameters of ((()) after an inner update (4) but before being mapped to mean parameters of 5, (1+1) through map (9). & some zero-mean noise, e.g. due to MC variation. The
17

---

second equality follows from the results of Appendix D. Similarly, let n = V (14 be the corresponding natural parameters. Now apply map (9) convert to mean parameters of p, (++1),
Ms
(++1)
V
17 (++1)
(mh²)) VA V-A-) (4)
V.
(11 11: (1+1) (1/² " (4())))))
VA VA (0) + E, V (u - v())))))
(44)
where the last equality follows from the definition of m/ (.), and by observing that
2 1g+11) 25 = a(1i 7//)
= =a A" (u)
10)) (v)()))
(45)
for all j. Let (1) be the mean parameters of pj (1) after a noise-free update, but before mapping
mean parameters of po: (1+11, so that 14(0) = 11. (1) + E.- Substituting this in, we have
(++1) VA V 14
E, (V (0) +5, - vA
(46)
Taking a Tayloor expansion around a we have
p4 (t-1) = =
(1)
VA' + - + vx) + 0(a2). (47)
Now subtract the noise-free update zi, (++1), and take expectations with respect to Ej for I,..., m.
B p[p((tt) 11: (1-1)
= A VA (u(0)
Σ, (n + +2, (nfor) + +((a))
(48)
the first term on the right hand side does not have zero expectation in general general-VA*G is not generally affine-and so the bias is (aa)) as 0+.
To that the expected change in L also O(a), take the Taylor expansion of L along the update direction-obaainn by subtracting A from the right hand side of (4)) function of the step size around a O,
-a no Σ jæi A, E,
(A(E
T OL ay,
+O(a),
(49)
which is clearly ((x) as a + 0+.
Proposition 3. After update (12) is executed in parallel over i, 0+, the expected decrease in O()), and the bias E [ps (1+1) (.((-)) is O(e), for all
Proof. Let (1) and (0) be the pre-update natural and mean parameters of 5, respectively. Furthermore, let
= - VA(mo + E, 21) - (p.cc))s(x)) &.,
1 (t)
< (1) Ep((s)]s(3]]- -s)))
(6)
(50)
be the mean parameters pú after an EP-p update (12) before being mapped to mean parameters of (() through (9). & is some zero-mean noise, e.g. due MC variation. The second equality
18

---

follows from the definition 14 ( = E (2) [2]]], (=)(())) and the fact that py = before an update. Now apply map (9) convert to mean parameters of p. (t+!),
((+1)_ pi
VA o 171
((+1)
(1z (1( - VA-)
(44 (4400
VA
(1+1)
(m²²²) (1) VA'" (o("))
VA VA 1 (1) + +,, (AA" (us* VA" (u5
(51)
where the last equality is obtained from the definition of n7 (1) for 1, and using
2 x(t11 2 7/3 πl
= A" (m) (09) VA (u²))
(52)
for j. Substituting in the definition of 10, we have
Mi ((-1) VA V (x Σ,
VA
V (1 (1-c))p )) < ([))())))=)]]+++ ( ( -va
(53)
Taking a Taylor expansion around 0, we have
p4 ((+1) = (1 + (22A(VA'
Σ, a o
VA" (11 (1-c)u + (Bpp)())(((()) (B + {3 (,) +0(e))
= = +e <<V<AAVAA'
19
(2 E o
(1 c) + (c (s {1 }/1 +0((2)
10 14 + [22 A V (m4 A' [ E, (
(x)())(((2)) + & " + +
(54)
Finally, subtract the noise free j (1+1), and take expectations (over &, for j = 1,...,,,,, m) to obtain the bias
E
(1+1)
14: C
[(+1)
E V
(55)
By assumption E[6,] 0 j. hence the first order term disappears, and the bias is O
To see that the expected decrease L is also Oc)), note that we are simply following the gradient in L with respect to μ. multiplied by the inverse Fisher and scaled by Neither the reparameterisation (from A to (4;) or the Fisher depend on e, and so the change in L must be (()) as + 0+ by simple Taylor expansion arguments.
H Implicit geometries of EP variants
In this paper we have shown that the updates of EP, EP-1, and EP-p can be interpreted as performing NGD of L with respect to the distributions {p, (1) The SNEP method Hasenclever et [17], in contrast, performs NGD of L with respect to the site potentials, treating them as bona fide distributions in for the purposes of NGD that is, NGD is performed with respect to the distributions with densities given by exp(A s(=) - A(x:) for i l..... In Section 3.4, we argued that the statistical manifolds (p, (0) ) are far more relevant for the optimisation of T. than those of the site potential pseudo-distributions In this appendix we shall justify this claim.
19

---

The loss function L, given by (3), has a unique optimum with respect to Det when Vx, L(0, >...... Am) i, which implies that
A +,, - A, 8 Ne, 0-, +
E pp(x) Fptt)
(56)
The minimisation of I. with respect to {de}e can therefore be viewed joint optimisation in the m + 1 distributions p. Pl. .... Pin. with the minimum found when their expected F-statistics match. Intuitively, it seems that optimisation of L with respect to X should be performed with consideration for the geometry of p and p:. By instead performing NGD with respect to the distributions {p, (+) J. in the case of EP, EP and EP- or with respect to the site potential psueco-distributionn.. in the case of SNEP, these methods can be interpreted as using surrogate distributions for NGD [46]. It possible to show that the Fisher information matrix (FIM) of p. (t) with respect to X; is equal to that of p. up to scalar constant, and furthermore, it can also be seen as close approximation to that of ps. This motivates the use of NGD with respect to p, (*) for the optimisation of L. In contrast, the FIMs of the site-potential pseudo-distributiions used by SNEP, can bear little relation to those of p or pt- In particular, the difference will be greatest when the full approximation no + Σjdy is significantly different than X; and we believe this may be why Hasenclever et al. [17] found that increasing the number of sites reduced the convergence speed of SNEP.
We can make these arguments more concrete by considering the curvature of L. Taking second derivatives of L with respect to A we have
V2 L(,, AI...., Am) = V2 A(To + E, + A, 87 (.
(57)
Let 2, (1) Xg at time t for j = L,... m. then
V2 A 11= L(O, A1 ......, Am) = V2 A(1o E, 100h + (A; + A:((0 311 1)
0p: + 8 3 A: 8, ".. )):
(58)
where TI and μ4 are the natural and mean parameters of 7t respectively. The Hessian of the second term will not generally be available in closed form, or easily invertible. We can however find convenient approximation. Note that after an outer update, at time t, we have 6 = no + Σ,,, (9. Also we have that e exp(X, (A, s(z) = ex)() at least near convergence, where c some unknown scalar constant. Combining these, we have
V3, A, - 3, Ae,, = V2, log f exp (10-8~>x) ax=)+++ + (x(=)) du(z)
Vi log
exp
(no + Σ, , 11t) - -8- +-12 T ((=))
V (7o + Σ, 110) + - - -N))
A A no + E, 2/00 + - 1001)
B 2 aps
(59)
Combining this with (58), we have
0p
V2
A L(0, X1 Am) (1 + 1)
(60)
and if we use this curvature approximation to perform a quasi-Newton step in L with respect to A and step size ē, we have
A + A - 2(1 + 1111+1111111 .m: VA, L(0, A1,,.,,,,,,
= 2 + 1 1 Om:
(A
- -
(61)
0μh
20

---

By letting = E(1 + 3, we have equivalence with the update of EP-7, given by (11). We therefore have that EP-1 is equivalent performing a quasi-Newton step in L. Furthermore, the approximate curvature becomes exact at the optimum if (=2) is an affine function of s(2), e.g. if (x)( is an (unnormalised) member of F.
In contrast, let us consider the implicit curvature approximation used by SNEP. Let be mean parameters of the member of F with natural parameter A, To perform the inner minimisation of the variational objective, SNEP performs NGD in L with respect to %. The implicit curvature matrix used by the SNEP update is then the Fisher with respect %Es which is given by dA: Now consider the actual curvature of L with respect to Y
(7).
8%
a% V3, L(0, 4(m) = 0%
aL.
Am)
V2 7: A())),
(62)
0%
k=
where we have used Ae(%)) = (n)( to denote the backwards map. Near the optimum we have 0L 8x, 0, and so the curvature will be approximately equal to the first term approach the optimum. We have already shown that
0pi (1- (1- V + Σ,,>i A1). Om:
V L(0, 4 A1 A.. ~ (1-)
(63)
and so the curvature with respect to 7/s is
Odi V2 A + E, )1) ad:
V2 4(:) ~ (1 - 9,
0%
07; v2
=(1 - -11-3,7 3,
A(no + ,,,,,) (V²
(64)
0%
Compare this with the implicit curvature used by SNEP of aAs and it is clear that the two are proportional only when To + +, = A₁ This suggests that if there are many sites (m is large), or the prior is very informative, the curvature matrix implicity used by SNEP is likely to be significantly different from the true curvature, even near the optimum.
I Computational cost analysis
Let Gsamp be the cost of drawing a single sample from of the tilted distributions. Let Cwwd be the cost of converting from natural to mean parameters in the base family F (the forward mapping). Similarly, let Obed be the cost of converting from mean to natural parameters in 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 nsmpp samples from each of the m sites. It then performs m backward mappings, to map from moments back updated site parameters. The total cost is then masamp MCbwd-
EP-7 EP-m draws smp 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 (A"(.). The primals of this JVP are the same for each site, so the linearisation only needs to be performed once per update, costing Ry 20owd, but have tangents. The overall cost of EP-m iteration is therefore s mmsaampSammp + Oww + (2 + m. Obwd-
EP-μ EP-p draws Nsamp samples from each of the 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-u iteration is therefore mMNssmppSSmmp + Crwd + M0bwd-
SNEP SNEP draws Rsamp samples from each of the m sites, and performs m backward mappings per iteration. The total cost is then MMsamppSump Moowd. which is the same EP.
In this work we largely assume that Gamp is the dominant cost. This often the case when the sampling is performed using MCMC. For example, in the NUTS sampler with default numpyro settings - as used in our evaluation each sample can involve evaluating up to 1024 gradients of the tilted distribution log density.
21

---

The other costs, Crwd and Cowd. depend on F. In the case of a MVN family with dense covariance, the foward and backward mappings are O(d)). for the dimensionality of 2. When d is large, Cwdd and Obed become more significant, in which case the balance will shift in favour of using more samples per update (and taking larger steps). When d very large, however, it may be necessary to choose a diagonal covariance base family instead, in which case the mappings can be performed in ((d) time. Also note that when d increases, Csamp will also typically increase at a rate faster than d. even with optimal tuning [8]. In our evaluation used just single sample per update for EP-1) and EP- 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 Cwwd and Cowd 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-learnigg optimisation scheme. We did not make use of such optimisations in our experiments.
J Evaluation 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 large number of samples (Nsamp 10 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 a--a in Figure 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 L.
Hyperparameter search spaces The step size for EP, a, was drawn log-uniformly in the range (10 4, 1), and >smp was drawn fog-uniformly between [d + 2.5, 10000.5) and then rounded to the nearest integer, where d is the dimensionality of 2. Note that the debiasing estimator used by Xu et al. [52] for MVN families is not defined for 7samp I d + 2. The thinning ratio was drawn uniformly from {1, 3, 4). The step size for EP-1 and EP-J4, C, was drawn log- uniformll in the range (10-5, 10 2). and nsamp was fixed to 1. The step size for SNEP was drawn log-uniformly in the range (10-5, 10 10-2), with rksamp and Ninner (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 Rinner > 1 inner steps per outer update. We did not find this necessary in any of our experiments, and fixed Tianer for EP, EP-1 and EP-p. For SNEP we included Nimer in the search space (as described above) to ensure 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 single core.
Software Implementations were written in JAX [7], with NUTS [27] used the underlying sampler. We used the numpyro [42] 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 [45]. with JAXopt [5] used to perform implicit differentiation through the iterative solve.
MCMC hyperparameter adaptation We performed regular warm-up phases in order to adapt the samplers to the constantly evolving tilted distributions, consistent with prior work [52, 48). The frequency and duration of these warm-up phases was also included in the hyperparameter search as follows. We drew the duration (number of samples) of each warm-up phase log- ungguniformly in the
22

---

range [99.5, 1000.5) and then rounded to the nearest integer. We then drew the sampling to warm-up ratio log- log-uniformly in the range (1, This ratio then determined how frequently the warm-up was performed, which we rounded to the nearest positive integer number of updates. Note that the Pareto frontier plots in Figures and include the computation time spent during warm-up phases.
Site parameter initialisation For EP, BP-1 and EP- we initialised the site parameters to A; Y 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 JF was MVN we used A = (2m) '7o A i, consistent with Vehtari et al. [48]. Unfortunately this initialisation strategy for SNEP is not always valid when F is NIW. This is because the site potentials in SNEP 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 F; namely, the political survey hierarchical logistic regression and neural response model experiments.
2
w
D
II
Figure 4: Directed graphical model for the experiments of Section 4.
We set 3; V in all experiments, corresponding to the original (not power) EP of Minka [37]. All of the models in these experiments followed the same general structure, given by
po II, p(w; 2)p(P;; Wi, =),
(65)
where is member of tractable, minimal exponential family F. This is shown graphically in Figure 4. We now provide details about individual experiments.
J.1 Hierarchical logistic regression
In the hierarchical logistic regression (HLR) experiments, there m groups (logistic regression problems) each with n covariate/response pairs, so that D; = {(Eig> Bi.j)) Ti, Rd and Viyy E m, and j 1 to n. Each group had its unobserved vector of regression
(0, 1), for coefficients E Rd, with density p(w; parameterised by global parameters z.
Synthetic data with MVN prior In the synthetic experiment, we had m = 16 groups, with d = 4 and n 20. F was an MVN family, and z (M1, log of ..., M4. log of) E g corresponded to the means and log- variances for the independent normal density p(w; 2) = [[.,,. k3.a3)) The prior had density P2(2) = -O, (iag(((,2) 2, 4, 2, .)). where diag(.) constructs a diagonal matrix from its vector- valued argument. The dataset was generated using the same procedure as Vehtari al. [48].
Political survey data with NIW prior In the political survey data experiment, F was a normal- inverse-Wishar (NTW) family of distributions, and (μ, vech(E)) was used to parameterise the group-level coefficient density p(w; 2) NWw; | μ, E). Po was a NIW prior on μ and Σ, 1, and V (following Wikipedia notation). We used log log-Choleso Cholesky parameterisation of Σ for the purposes of sampling. The dataset consisted of binary responses the statement " "Allow employers to decline coverage of abortions 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 predictor variables from the dataset, corresponding to characteristies of a given survey participant. The predictors were binary variables conveying: age (3 groups), ethicity (2 groups), education (3
with parameters Ho 0,
'Note that in this section use to denote the mean of normal or multivariate normal distribution. This should not be confused with the usage for exponential family mean parameters the main text.
23

---

groups) and gender. We also included a state-level intercept to capture variation in the base response level between states, so that This setup is based on one used by Lopez-Martin al. [33]. The data is available at https: //tgt/gbt../ttt/// hwwp://ooeoeaaattt///eeeeaaeeeee
J.2 Cosmic radiation model
In this experiment, there were nonlinear regression problems, each corresponding model of the relationship between diffuse galactic far ultraviolet radiation (FUV) and 100-um infrared (i100) emission in a particular sector of the observable universe.
The nonlinear regression model for sector had parameters w, E IR. The m regression problems were related through common hyperparameters E R which parameterised the section-level densities P(D,, W/ 2) where D, is the observed data for sector i. FF was the family of MVN distributions, and the prior
E IR 18 had density po(z) (O, 102). The specifies of p(D,, w; 2) are quite involved and we refer the reader to Vehtari et al. [48] for further details.
Vehtari et al. [48] applied this model to data obtained from the Galaxy Evolution Explorer telescope. We were unable to obtain the dataset, and we generated synthetic data using hyperparameters that were tuned by hand to try and match the qualitative properties of the original data set (see Appendix for examples). We used a reduced number of m = 36 sites and n 200 observations per site to reduce computation, allowing us to perform comprehensive hyperparameter search.
Using the conventional site parameter initialisation of X = 0 V resulted in most runs EP failing during early iterations. We found that this was resolved by initialising with the method used by Vehtari et [48] for SNEP, that is, A (2m) No i, and so we used this initialisation for all methods. We note, however, that the performances of EP-17 and EP-p were largely unaffected by this change.
J.3 Neural response model
In this experiment we performed inference in a hierarchical Bayesian neural response model, using recordings of VI complex cells in anaesthetised adult cat. 10 neurons in specific area of cat V1 were simultaneously recorded under the presentation of 18 different visual stimuli, each repeated 8 times, for 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:// //rcnss//////////////ee org/data- sets/vc/pvc- [4].
2 = (1,, vech(E)) E R was used to parameterise N (logTyi μ4, E). the density for latent log firing
rates in trial for j I..... 144. po FF was NIW with parameters oo 1, 1/ 12, X = 2.5, and
= 1.25/ (again following Wikipedia notation). The observed spike count for neuron in trial j. JjA N, modelled Poisson(C,,A exp(rr,x)). for j 1,..., 144, k 1, ... 10. We again used a log-Cholesky parameterisation of for sampling.
We grouped trials together into
batches of
18 trials, so that
w; = (log 118005-1)+1;..... log r180:) R150 was the concatenation log firing rates for all 10 neu- rons across 18 trials, with D, ((.).1...... x3,10)) the observed spike counts, for j 180(; - 1) to 180i.
K Hyperparameter sensitivity
In this appendix, we examine the effect of varying hyperparameters of EP and EP-1 on the synthetic hierarchical logistic regression experiment of Section 4. The results for EP-p are similar to those of EP-1) and we do not consider it separately here.
EP In the left panel of Figure 5 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. The middle panel, similarly, shows the effect of varying the step size, and the right panel shows the effect of varying the thinning ratio Together these plots illustrate the difficulty of tuning HP in stochastic settings - tracing out the frontier a three-dimensional problem. For example, to achieve the best accuracy within compute budget of 10" steps, the practitioner would need to set
24

---

400 < samp I 500, 0.1 < S 0.3, and T = 2, and any deviation from this would seemingly result in suboptimal accuracy.
Sample count
Step size
Thinning ratio
10
Divergence (nats)
100
Paamp
E
(0, 400)
0..0.1) 0.1)
(0.1. 0.3]
(500, 108]
(0.3,
10c
10
10
10c
10
10
10
10
10
Steps
Steps
Steps
Figure 5: The effect of varying EP hyperparameters. Partial Pareto frontiers show the number of NUTS steps (I-axis) against the KL. divergence from p estimate of the optimum (y-axis).
EP-n In the left panel of Figure 6 we show the effect of varying the number of samples used for estimating updates in EP-m - the frontier traced out with 7samp I. We note however that the difference between Tesamp 1 and Reamp 10 is relatively small, and so it may be sensible to choose Bsamp" 1 to make efficient use parallel hardware, or minimise per-iteration overheads (see Appendix 1). The middle panel, similarly, shows the effect of varying the step size, with more accurate regions of the frontier corresponding
smaller step size. This also suggests that practice may make sense to set relatively large at first and then gradually decay it to improve the accuracy. Finally, right panel shows the effect of varying Tinners the number of inner steps performed per outer update. YAimeer corresponds to using EP-7) in "double-loop" mode. We did not find it necessary to increase Ninnrr above one to obtain convergence our experiments, but Figure (right) demonstrates that doing would have relatively small impact on performance. Together these plots illustrate that hyperparameter tuning for EP-r) is relatively straightforward. The frontier can be largely be traced out by varying e, and is relatively insensitive to Rsamp and Tiiner-
Sample count
Step size
Inner step count
10a
Divergence (nats)
10
C
Pinncr
(3e- le
10
10
100
(3e 4. 1]-
100
10a
10"
10a
10a
107
10"
10a
107
10
Steps
Steps
Steps
Figure 6: The effect of varying EP-r7 hyperparameters. Partial Pareto frontiers show the number of NUTS steps (r-axis) against the KL divergence from p to estimate of the optimum (y-axis).
L Pareto frontiers in wall-clock time
In Figure 7, 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 Figure 2. These results are in broad agreement with those of Figure 2, which suggests that the sampling cost does indeed dominate the computational overheads of EP-7 and EP-p in these experiments. The wall-clock time results for EP-1 and EP-μ would likely be improved by using more than one sample per update, by making more efficient use of hardware resources and minimising per- -eerittatioon 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 we discuss approaches for minimising per-iteration overheads. These would likely improve wall-clock time performance for all methods, but should disproportionately favour EP-1) and EP-ya, due their frequent updates and larger per-iteration overheads compared to EP and SNEP.
25

---

Hierarchical logistic regression (MVN prior)
Hierarchical logistic regression (NIW prior)
10'
Divergence (natts
10a
10"
10a
10a
10a
10a
10a
Cosmic radiation model
Neural response model
10a
Divergence (natas
10a
10"
109
10-1
10a
10a
10a
10a
Time (s)
Time (8)
EP
EP-1)
EP-4
SNEP
Figure 7: Pareto frontiers showing the number of seconds elapsed (I-axis) against the KL divergence from p to an estimate of the optimum (y-axis). Fach point on the plot marks the lowest KI. 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.
M
Comparison with conjugate-computatie variational inference (CVI)
Our focus in this work has been on developing improved methods for performing EP when the updates must be estimated with noise. A major motivation for this setting is that it potentially enables EP to be used as a so-called black-box inference method, significantly expanding the set of models can be applied EP has several computational advantages over direct MCMC approaches, the prevailing dominant class of black-box mehods, as discussed by earlier works [52. 17, 48]. There another popular class of black-box inference methods however; namely, those of variational inference (VI)." In order to understand the trade-offs involved in using EP and its variants over VI, we performed experiments explore the relative strengths and weaknesses of EP-7 and conjugate-computation variational inference (CVI) [31], an efficient VI method.
Experimental setup We used EP-1 and CVI to perform approximate Bayesian inference using the same hierarchical logistic regression model (with MVN prior) and synthetic dataset as Section 4. For our evaluation metries, we used both forward and reverse KL divergences between the current (MVN) approximation, and a MVN "target" distribution. The target distribution was a MVN distribution fitted using 500,000 samples from the posterior, obtained by running NUTS for million samples and then discarding the first half as a warm-up. For CVI, we used structured MVN approximation that had the same conditional independence structure as the true posterior that is, modelled all pairwise dependencies within within each w,, and between and each w, but did not model direct dependence between Wi, wy for i /f j. We used 500 random hyperparameter settings for each method, using the search spaces described below. Unless stated otherwise (below), all other details were as described in Appendix
Hy perepaaareeeer search spaces For EP-1 used the same hyperparameter search space as described in Section For CVI, the step size was drawn log-uniformly in the range (10-5, 5 1). and the number of samples used estimate the update drawn log- Jog-uniformly in the range [.5, 100.5) and then rounded to the nearest integer. The site parameters for CVI were initialised to be the
"Expectation propagation is also variational inference method. but we follow convention here by using variational inference to refer specifically variational optimisation of the free energy / evidence lower bound.
26

---

parameters of a zero mean MVN distribution with scaled identity covariance matrix, with the scale drawn log-unine segg--nifomly from (10 5, 10%).
Wall-clock time performance comparison We compared the wall-clock time performance of EP-1 and CVI, with the former using NUTS as the underlying sampler. For both methods, we chose the hyperparameters that gave the lowest reverse divergence between the approximation and the MCMC target distribution after 100 seconds. The results of this experiment are shown in the left panel of Figure 3. We can see that EP-1 converges least one order of magnitude slower than CVI on this problem, when measured in wall-clock time. This difference is largely due to the relative cost of drawing samples for the two methods. While CVI also uses samples to estimate its updates, it uses samples from the current MVN approximation, which are cheap to generate using standard methods. EP-r on the other hand (as with other EP variants), requires samples from the tilted distributions to estimate the updates, and in this experiment we used NUTS to draw those samples. Each NUTS sample can involve many evaluations of the energy function gradient, and furthermore, because MCMC samples are generally autocorrelated, the cost of drawing approximately independent samples can be even higher. We used NUTS to be consistent with prior work, and because our focus in this work has primarily been the relative performance of different EP variants, we did not make particular efforts to improve the performance of the underlying samplers, beyond tuning their hyperparameters in fairly standard manner. However, all of the EP variants considered in this paper, including EP-m are agnostic to the choice of sampling kernel used, and we believe there is scope to significantly improve the efficiency of the underlying samplers, for reasons we briefly discuss in Section 6.
Sample efficiency In order to decouple the performance of EP-1 from that of the underlying sampler, we repeated the experiment using an "oracle" sampling kernel for EP-7. This oracle kernel was simply NUTS but with thinning ratio of 100, so that each sample was an approximately independent sample from the tilted distribution. We then compared EP-1 and CVI using the same metrics as the previous experiment, but with "time" measured in samples. We chose the hyperparameter setting for each method that achieved the lowest reverse KL divergence after I, 1,500 samples. When measured on this basis we see that the convergence speeds of CVI and EP-1) are roughly equivalent, demonstrating that the sample efficiency of the two methods are similar. This should not be too surprising, as CVI is (mathematically, but not algorithmically) equivalent to the limiting case of power EP (as B, - o0) [11, 51].
Approximation quality Figure demonstrates that EP-1 is able to achieve a more faithful ap- proximation of the target distribution when measured by either the forward or reverse KL. At first may seem surprising that EP is able to obtain a lower reverse KI. divergence than CVI, but there are two reasons for this apparent anomaly. First, with EP, we have MVN approximation for only, with the approximate posterior over local variables (w; for 1, 1,.... m) represented only implicitly with samples. CVI on the other hand necessarily optimises joint MVN over all variables, and the z marginal of the optimal VI posterior over all variables is not the same as the optimal VI posterior over z. The second reason is that our target distribution is essentially a mathed MVN approximation of the true posterior, which would naturally tend to favour the forward KL divergence that is (approximately) targeted by EP. Nevertheless, qualitative assessment of the pairwise marginals, as seen in the two rightmost panels of Figure 3, shows that the true posterior is more faithfully approximated by EP-1 than CVI, with the latter significantly underestimating uncertainty in the inter-group variability parameters, (log ). The full set of pairwise posterior marginals can be seen in Figure 8. Note that the CVT approximation appears more constrained than might be expected based on the marginal plots alone. This is due to the zero-avoiding nature of the reverse KL divergence non- non--Gussianity in one variable can constrain the marginal distributions of others, particularly when there is strong dependence. The local variables, {w:}, can have both significant non-Gaussian posterior structure, and strong dependence on {log J
N Cosmic radiation data
Vehtari et al. [48] used a hierarchical Bayesian model to capture the nonlinear relationship between diffuse galactic far ultraviolet radiation (FUV) and 100-um 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 that
27

---

log 01
log
log
log 74
log
Logas
EP-7
CVI
MCMC
Figure 8: Contours of pairwise posterior marginals for the hierarchical logistic regression experiment of Section 4. overlaid with the MVN approximations of EP-7, CVI, and one obtained by fitting directly to MCMC samples.
were tuned by hand try and match the qualitative properties of the original data Example data generated using these hyperparameters is shown in Figure 9, For comparison, see Figure 9 of Vehtari al. [48].
28

---

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

---

1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The novel perspective on EP is given in Section 3, with proof in Appendix D. The new EP variants also introduced in Section 3, and we demonstrate the stated advantages in the results of Section 4, and Appendices and L. The abstract and introduction also make clear that proposed methods primarily of use when the updates of EP are to be estimated with MC samples.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in Section 5, and computational efficiency is also discussed separately in Appendix I.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
30

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Proposition is proven in Appendix D, and Propositions 2 and 3 are proven in Appendix G.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide pseudocode of the proposed EP variants in Algorithm 2. Details of individual experiments are given in Appendices J and M. Code for our experiments is available from the URL given in Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
31

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide pseudocode of the proposed EP variants in Algorithm 2. Details of individual experiments, including links datasets, are given in Appendix J. Synthetic data was generated from models that are described either in Appendix J, or in other referenced works. Code for experiments is available from the URL given in Section 4.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /ttg//udde/////// www.SSuimmiiiooooololooooo
/ /nips.
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental settings are summarised in Section 4, with full details, including hyperparameter search spaces, given in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The main results are summarised by the Pareto frontiers Figures and The figures display enror bars, corresponding to the minimum and maximum value observed for chosen hyperparameter settings over random seeds, which were used to seed the MCMC samplers, state in Section 4.
32

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2- 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative етгог rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix gives details of the hardware used to run the experiments, as well the number of runs carried out. This, combined with the wall-clock time results in L, give an estimate of the compute resources needed to recreate our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification: We have read the NeurIPS Code of Ethics and believe our work confirms with it in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
33

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper proposes ways improve the practicality and efficiency of existing probabilistic inference methods, and we cannot foresee direct societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not believe there is reasonable risk of misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, see Appendix J.
Guidelines:
34

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. condattaees has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: We have made our experiment code public, but do not consider it one of the paper's main contributions.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourced experiments or research involving human subjects were conducted as part of this work.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
35

---

Answer:
[NA]
Justification:
No research involving human subjects was conducted as part of this work.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
36


Reviews:
[{"review_id": "DUV0is0hLv", "reviewer": "Reviewer_6XAG", "summary": "This paper considers new inference algorithms for EP. By framing the moment-matching equations of EP as a natural gradient update of a variational objective, they propose two new algorithms that are better suited for reducing/removing the bias introduced when sampling is required.", "questions": "1)\tWhat is c_samp?\n\n2)\tIn Eqn 24 should there be a $beta_i$ scaling term?\n\n3)\tIn fig 1 a) why do most of the curves have a `u’ shape? Additionally why do the dashed orange/green curves not?\n\n4)\tDoes sequential update (instead of in parallel) affect the bias?\n\n5)\tWhy do you fix alpha to 1 and introduce a new step size eps? This looks like the only difference between Eqn 10 and the natural gradient part of 11. \n\n6)\tThe link between natural gradient steps  and EP updates has been established previously (in Heskes and also in Hasenclever). Why is the view taken in paper a novel perspective compared to these previous works?\n\n7)\tHow does this work relate to `Bayes-Newton Methods for Approximate Bayesian Inference with PSD Guarantees’  Wilkinson et al, 2021 ? Which consider EP style algorithms with a natural gradient framework.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I like the attempt at generalising and encapsulating the EP literature, and I find the `trick’ of mitigating the sampling inducing bias through a change of parameterisation is clever.", "weaknesses": "1)\tIn general I find the paper quite hard to follow and of course this is not helped by the form of the standard EP equations. For example the main point of the paper is to handle the bias introduced by using sampling to estimate certain quantities within the EP equations. However the source of this sampling is only explicitly mentioned on line 132 in text. It would be much clearer if this was explicitly in Alg 1, and would also make it more convincing that Alg. 1 actually encapsulated many EP style algorithms.\n\n2)\tFollowing above some terms are not defined until later on in the paper. For example the convex conjugate in eqn 4, and on line 100 only defined one page later on line 123. \n\n3)\tThe paper is very descriptive but does not explain results. For example, the experiments only describe the set up, but all results and any conclusions are pushed into the appendix. This makes it hard to actually assess the contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for taking the time to read our paper and for providing thoughtful feedback. The reviewer’s concerns relate to the paper presentation. We will take specific steps to address these, detailed below. In light of these changes, would the reviewer be willing to reconsider their position on the paper?\n1. We agree the presentation of EP updates is somewhat atypical. We felt it necessary to provide them in the context of the saddle-point problem, as this perspective is crucial for the later sections. Space constraints meant that providing both this and the conventional view in the main paper was difficult. We think this would be best addressed by referring to a short appendix after acknowledging the atypical presentation on line 69. The appendix would provide the conventional view – motivated as KL-divergence-minimising projections into the approximating family – and show that the resulting updates are equivalent to ours. We will also introduce the sampling source earlier by adding the following at the end of line 72: _“The expectation in (4) is most often computed analytically, or estimated using deterministic numerical methods. In principle it can also be estimated by sampling from $p_i$, however, we will later show that the resulting stochasticity can introduce bias into the procedure.”_. Note that this is also discussed starting on line 91. Regarding Algorithm 1, we will highlight the expectation in the update and add the comment _“stochastic estimation of this expectation can lead to biased updates”_.\n2. $A^*(.)$ is introduced on line 38, but we will add the following text after that definition: _“$A(.)$ and $A^*(.)$ are convex conjugates of one another.”_.  We will also add reminders throughout to help the reader.\n3. As the results in Figure 2 were similar we summarised them jointly at the beginning of Section 4 (lines 251-259) before giving an overview of the individual experiments. We will expand on this and give a brief summary of the results for each experiment at the end of their descriptions. If space permits we will also move some of the hyperparameter sensitivity results from Appendix I into this section.\n\nIn response to the questions:\n1. $c_\\text{samp}$ is the cost of drawing a sample from one of the tilted distributions. It is defined at the beginning of Appendix G (line 884) and not used outside of that section.\n2. Yes, we will fix this.\n3. If the step size is too large the expected progress can decrease. This can simply be due to overshooting the minimum along the step direction. Alternatively, as a larger step size also results in higher variance, it can lead to some steps that are far worse than the expected step location, or even outside of the valid domain. If any steps in our sample went out of the valid domain we could not compute the average, and so the line can end abruptly (above a certain step size).\n4. Bias results from other site parameters changing between updates, which happens in both sequential and parallel settings. We will explain this on line 140 as other readers are likely to have the same question.\n5. $\\alpha$ plays a similar role to a NGD step size in (10), but it also affects the distribution $\\tilde{p}_i$ and hence the map from $\\lambda_i$ to $\\eta_i$. Note that (10) is an update direction in $\\eta_i$ whereas (11) is an update for $\\lambda_i$. If we did not fix $\\alpha=1$, the resulting coefficient in (11) would be $\\alpha^2$. To obtain a more faithful NGD interpretation we fix the distribution and vary the step size, which requires introducing $\\epsilon$. We will elaborate on this on line 155.\n6. We do not believe any prior work has made a direct connection between the updates of EP and NGD, except for the limiting case in which power EP coincides with VI (see answer to 7 below). Hasenclever et al. (2017) derived _new_ updates (different from the standard ones) which performed NGD of the variational objective with respect to a different distribution and parameterisation; in contrast, we show that the standard EP updates can already be viewed as performing NGD. We discuss this from line 210, but will amend the wording to make the distinction clearer. We believe the reviewer may be referring to the extended version of Heskes and Zoeter (2002) in which the authors propose two methods for finding saddle-points of the objective. One performs joint gradient ascent/descent in the parameters, and the other is what could be called “standard” double-loop EP. With respect to the first, the authors say it _“can be interpreted as a kind of natural gradient descent in γ and ascent in δ”_. However, the method referred to simply follows the standard gradient. We asked one of the authors about the meaning behind the statement through private correspondence, to which they replied _“I wouldn't dare to claim that there is a direct connection to Amari's natural gradient, perhaps more to his idea of information geometry and em algorithms.”_. In any case, the statement is about the gradient-based method and not the “standard” updates. We are not aware of any other work by Heskes suggesting a connection between EP and NGD.\n\n7. Wilkinson et al. (2021) present a unifying framework encompassing several algorithms, viewing them as performing online Newton optimisations with localised updates. This framework elegantly illustrates connections between EP, VI and posterior linearisation. Their view is complementary to ours, and does not make any claims about a connection between EP and NGD in the general case. The authors do however show that the updates of power EP coincide with those of natural gradient VI when the variational limit of the power parameter is taken ($\\beta_i \\rightarrow \\infty$, in our notation); this point was also made by Bui et al. (2018). Our connection is more general, and applies for any value of $\\beta_i$. This result is clearly relevant however, and we will reference it after Proposition 1 on line 127.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dZVNEgrxB1", "reviewer": "Reviewer_YMg2", "summary": "This work addresses the sensitivity of expectation propagation (EP) to the randomness of Monte Carlo (MC) estimates involved in its update steps. It tackles this issue by recasting the moment matching step in EP as natural gradient descent (NGD) in the mean space of an exponential family distribution. The author identifies that the instability of EP to MC noise is due to the nonlinearity of the mirror map defined by the log partition function of the exponential family. By cleverly moving the NGD from the mean space to the natural parameter space, this issue is bypassed.\nThe author also studies the influence of the stepsize on the accumulation of MC error, finding that decreasing the stepsize of the NGD updates helps to reduce the bias.", "questions": "- How do you estimate the KL divergence in the experiments? Maybe I'm missing something, but the target posterior is only known up to a constant, so unbiased KL est is not feasible?\n\n- I'm curious to see the relative performance of this work to NGVI [1] (specified in the Weakness section).", "limitations": "The limitation is well discussed.", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This is a well-written work. It starts with a succinct overview of the expectation propagation (EP) algorithm and summarizes several variants of EP updates in a unified manner (in Algorithm 1). This review also makes the contribution of this work very clear. It provides a sufficient discussion on why EP is sensitive to stochasticity and clearly motivates the design of the methodology. I also appreciate the self-contained review of the exponential family, natural gradient descent, and detailed derivation of EP updates provided in the appendix. Although these materials are standard, they help to make the work more accessible to a general audience.\n\n- The feasibility of performing unbiased NGD updates in the natural space of $\\tilde p_i$ (Prop1) is a smart observation, and this simple modification makes the inner loop significantly more robust to the stochasticity in MC est. This simple modification (from mean space to natural parameter space) makes the inner loop significantly more robust to the stochasticity in Monte Carlo estimation. I'm indeed impressed that EP can work with single-sample estimation in the inner loop.\n\n- I also appreciate the rigorous investigation of the effect of $\\alpha$ and $\\epsilon$ to the bias of estimated mean parmaeter.\n\n- The empirical performance of proposed methods are superior; it outperforms the standard EP by a significant margin (as shown in Fig2).", "weaknesses": "The technical side of this work is very strong in my perspective, and I don't observe many weakness on this regard. \nHowever, I think the novelty of this NGD perspective is overclaimed. To my knowledge, the relationship between moment matching in the exponential family and NGD is very well known (e.g., [1][2][3]). Once it is identified that $\\tidle p_i$ is in the exponential family, the derivation of NGD (in either mean space Prop1 or natural parameter space) becomes quite straightforward. I hope the authors can elaborate on this point, and include discussion on some relavant literature arolund prop 1.\n\nI believe the real contribution of this work lies in the rigorous understanding of the influence of the stepsize on the accumulation of Monte Carlo (MC) error and the identification of the nonlinearity of the mirror map that leads to this error accumulation, which is already impactful to me.\n\n\n\n\n- [1]: Conjugate-Computation Variational Inference : Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models, 2018\n- [2]: MCMC-driven learning, 2024\n- [3]: Distributed Bayesian Learning with Stochastic Natural Gradient Expectation Propagation and the Posterior Server", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for taking the time to read our paper and for providing thoughtful feedback. With regards to the novelty of our NGD interpretation, we agree that the link between NGD and exponential family moment matching is known, and our intention is not to claim otherwise; we will change the wording around Proposition 1 to make this clearer, with reference to the relevant literature.  We do believe however that the connection of this result with the updates of EP is novel; without the identification of an exponential family distribution $\\tilde{p}_i$, parameterised by $\\lambda_i$ such that the natural parameters are given by the specific affine function $\\eta_i^{(t)}(\\lambda_i)$, the connection with NGD does not hold. While this is perhaps a straightforward consequence of the link between NGD and moment matching, we do not believe this connection has been made before and it was a necessary insight for the contributions of the later sections. Indeed the only connection we are aware of having been made between the updates of (power) EP and NGD is in the limiting case (as $\\beta_i \\rightarrow \\infty$) for which the updates of power EP and natural gradient variational inference (NGVI) coincide (Bui et al. 2018, Wilkinson et al. 2021).\n\nIn response to the questions:\n1. The metric displayed in the plots shows the KL divergence from the current approximation ($p$) to (an estimate of) a converged EP solution. The converged solution was found by running EP for a large number of iterations, with the moments estimated using a very large number of samples; we discuss this on lines 247 and 923. We felt this was the most meaningful metric for our experiments, given that we were benchmarking the optimisation performance of different EP variants. Another choice we considered was to show the KL divergence between $p$ and a moment-matched Gaussian, for which the moments are obtained by running long MCMC chains. However, we found this metric to be problematic, as it often does not monotonically decrease during optimisation; for example, converged EP solutions tend to have lower variance than the true posterior (see, for example, Vehtari et al. 2020, Cunningham et al. 2011), and so we often saw a pattern of the KL decreasing as the approximate posterior variance shrinks towards the true posterior variance, and then rising as the approximation variance continues to decrease, before eventually settling at a higher KL than the earlier transient. In contrast, the KL divergence to a converged EP solution was typically monotonic decreasing for stable hyperparameter settings, and had the additional benefit of the attainable lower bound being (approximately) zero.\n\n2. The main focus of our work was to provide more effective methods for optimising the EP variational problem in the presence of Monte Carlo noise, with the case for using EP in such settings having been been made by prior work (Xu et al. 2014, Hasenclever et al. 2017, Vehtari et al. 2020). That being said, we agree it would be interesting to compare NGVI with the different EP variants on various performance metrics as a function of computation cost / time, to gain understanding of the tradeoffs involved. Due to time constraints we have not been able to run these experiments for the rebuttal, but we will do so and add the results to the final paper. As an aside, we note that Bui et al. (2017) studied the effect of the power parameter of power EP ($\\beta_i$ in our notation), including the limiting case of VI ($\\beta_i \\rightarrow \\infty$); they found that intermediate values of the power parameter were consistently better across tasks and performance metrics when compared to either the EP or VI extremes. As our methods apply to the general power EP objective, they too can be used to find such intermediate approximations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iGRNtwfzd3", "reviewer": "Reviewer_1uyT", "summary": "The authors introduce two methods for Expectation Propagation [EP] (which itself can be understood as a method for approximate Bayesian inference) that are robust to the stochastic noise introduced by using approximate expectations in the inner loop of EP. To do so the authors re-interpret the moment-matching update as a particular Natural Gradient Descent (NGD) step. While stochasticity remains due to the expectation, by moving to natural parameters MCMC estimates of the expectation do not suffer from bias-introducing non-linearities. By significantly reducing bias (to zero in the case of EP-$\\eta$) the authors show that stable and computationally efficient EP is possible with single-sample moment-matching updates. In experiments the authors show that the method is easier to tune than alternatives and exhibits promising empirical performance.", "questions": "- math typos: line 804\n- NIW is undefined where first introduced\n- i guess in sec. 2.1 you should emphasize more that the domain of $\\theta_i$ is $\\Omega$, since $\\theta_i \\notin \\Theta_i$ as one might expect? similarly i guess in practice given the form of the outer update $\\theta_i = \\theta_j$ for all $i, j > 0$? maybe this should be emphasized? can you please clarify?\n- i understand that this is a paper about EP, but wouldn't it be valuable to benchmark against variational inference in at least one case? whether the authors like it or not, this remains a relevant comparison in my opinion.", "limitations": "To my mind the major limitation is the lack of comparison to alternative methods for approximate inference like variational inference.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- To my mind the major strength of this submission is that it offers a well-reasoned technical improvement to a class of algorithms (EP) that are arguably underexplored in the literature (probably in large part due to the dominance of variational inference).\n- Moreover, empirical results support the effectiveness of the proposed algorithms. Since the authors work hard to reduce dependence on tricky hyperparameter choices, the reader has reason to be confident in the basic validity of the presented empirical results.\n- It is conceivable that the (novel, to the best of my knowledge) NGD interpretation offered by the authors could help motivate yet other EP variants.", "weaknesses": "To my mind the major weakness of this submission is that the exposition is rather dense and at times hard to follow. This may be somewhat inherent in the technical nature of the topic, but I believe the authors could do a better job of guiding the reader. In particular a lot of space is devoted to somewhat extraneous modeling details (\"...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...\"), and this space in the main text could be re-purposed to offer more details on the algorithm, discuss the implications of the experimental results in more detail, or give more color on the \"unified\" EP presentation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for taking the time to read our paper and for providing thoughtful feedback. We agree with the reviewer that the exposition can be dense at times; as the reviewer correctly mentioned, this was to some extent unavoidable due to the nature of the ideas being discussed, but there were several places that we would have liked to expand on points or offer additional explanation and were unable to do so due to space constraints. With the additional page we will look to add additional signposting throughout, including the following specific changes in line with the reviewers suggestions:\n1. We will add sentences to add some intuition and summarise the idea behind EP-$\\eta$ on line 158, before _“We call the resulting procedure…”_, and likewise for EP-$\\mu$, on line 184, before _“We call this variant…”_.\n2. We will expand on the summary of experimental results (lines 251-259) and discuss their implications. We will also add a brief summary of the results for each individual experiment at the end of their descriptions (lines 270, 279, 293 and 304). If space permits we will also look to move some of the hyperparameter sensitivity results from Appendix I into this section.\n3. We will add a paragraph, beginning on line 79, to inform the reader that we are about to introduce a unified EP algorithm, and explain that the reason for doing so is primarily to show how several variants, including the new ones to be introduced later, are related to one another. We will also add an explanation of Algorithm 1, before line 86.\n\nIn response to the reviewer’s questions:\n1. Yes, we will fix this.\n2. Agreed, we will define it at first usage.\n3. Yes, agreed and we will emphasise both points. It is indeed the case that after an outer update $\\theta_i = \\theta_j$ for all $i, j$. We could therefore obtain an equivalent procedure using just a single variable $\\theta$ in the outer optimisation, but we kept the current presentation to be consistent with prior work. We believe this presentation may be a legacy of earlier works on double-loop EP where additional model structure was assumed, in which cases the distributions involved in the optimisation can be over different subsets of variables at each site.\n4. The main focus of our work was to provide more effective methods for optimising the EP variational problem in the presence of Monte Carlo noise, with the case for using EP in such settings having been been made by prior work (Xu et al. 2014, Hasenclever et al. 2017, Vehtari et al. 2020). That being said, we agree it would be interesting to compare VI with the different EP variants on various performance metrics as a function of computation cost / time, to gain understanding of the tradeoffs involved. Due to time constraints we have not been able to run these experiments for the rebuttal, but we will do so and add the results to the final paper. More specifically, we will compare with natural gradient VI (Khan and Lin, 2017) as suggested by another reviewer.\nAs an aside, we note that Bui et al. (2017) studied the effect of the power parameter of power EP ($\\beta_i$ in our notation), including the limiting case of VI ($\\beta_i \\rightarrow \\infty$); they found that intermediate values of the power parameter were consistently better across tasks and performance metrics when compared to either the EP or VI extremes. As our methods apply to the general power EP objective, they too can be used to find such intermediate approximations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all of the reviewers for taking the time to review our paper. All reviewers offered valuable feedback, which we have taken on board. Two reviewers highlighted weaknesses related to the paper presentation; we have taken specific steps to address the points raised, which are detailed in the rebuttals below. One reviewer questioned the novelty of our NGD interpretation of EP; we believe we have addressed this below, and will add wording to the paper to make the extent of our contribution clearer. Two reviewers also asked for variational inference to be included as an additional baseline, which we have committed to doing for the camera-ready paper.

Meta-review: This paper presents a new view of the moment-matching updates of expectation propagation (EP) as  natural gradient descent (NGD) of a variational objective. This view allows the authors to propose two new methods that address the deficiencies of previous approaches, in particular that of not being robust to Monte Carlo noise when estimating the corresponding expectations. 
 
Overall, the reviews acknowledge the novelty of the approach (which has been further clarified in the rebuttal) and have praised the paper for the “clever” trick of reducing the bias of the MC expectation via change of parameterization and for being technically very strong. Indeed, one of the reviewers mentions “I'm indeed impressed that EP can work with single-sample estimation in the inner loop”. The main issues raised by the reviewers concern (1) the clarity of the presentation of the EP algorithm, (2) description of the results and (3) relation to previous work (in particular wrt to the work of Hasenclever et al, 2017). I believe the authors have addressed these issues in their rebuttal successfully with some of the reviewers updating their score accordingly. 

With this, I recommend acceptance.
Decision: Accept (spotlight)