Fearless Stochasticity in Expectation Propagation

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:

1. 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…”.
2. 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.
3. 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.

In response to the reviewer’s questions:

1. Yes, we will fix this.
2. Agreed, we will define it at first usage.
3. 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.
4. 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. 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.

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).

In response to the questions:

1. 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.

2. 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.

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?

1. 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”.
2. $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.
3. 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.

In response to the questions:

1. $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.

2. Yes, we will fix this.

3. 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).

4. 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.

5. $\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.

6. 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.

7. 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.