# Partitioned Variational Inference: A Framework for Probabilistic Federated
Learning

Matthew Ashman<EMAIL_ADDRESS>
University of Cambridge, Cambridge, UKThang D. Bui<EMAIL_ADDRESS>
University of Sydney, Sydney, AustraliaCuong V. Nguyen<EMAIL_ADDRESS>
KFSCIS, Florida International University, Florida, USAStratis Markou
<EMAIL_ADDRESS>
University of Cambridge, Cambridge, UKAdrian Weller<EMAIL_ADDRESS>
University of Cambridge, Cambridge, UKSiddharth Swaroop<EMAIL_ADDRESS>
University of Cambridge, Cambridge, UKRichard E. Turner<EMAIL_ADDRESS>
University of Cambridge, Cambridge, UK

###### Abstract

The proliferation of computing devices has brought about an opportunity to
deploy machine learning models on new problem domains using previously
inaccessible data. Traditional algorithms for training such models often
require data to be stored on a single machine with compute performed by a
single node, making them unsuitable for decentralised training on multiple
devices. This deficiency has motivated the development of federated learning
algorithms, which allow multiple data owners to train collaboratively and use
a shared model whilst keeping local data private. However, many of these
algorithms focus on obtaining point estimates of model parameters, rather than
probabilistic estimates capable of capturing model uncertainty, which is
essential in many applications. Variational inference (VI) has become the
method of choice for fitting many modern probabilistic models. In this paper
we introduce partitioned variational inference (PVI), a general framework for
performing VI in the federated setting. We develop new supporting theory for
PVI, demonstrating a number of properties that make it an attractive choice
for practitioners; use PVI to unify a wealth of fragmented, yet related
literature; and provide empirical results that showcase the effectiveness of
PVI in a variety of federated settings.

##  1 Introduction

As networks of connected devices, data centres and centralised servers
proliferate, there is a corresponding need for distributed machine learning
algorithms—algorithms that operate across multiple platforms, utilising the
computational capabilities and accounting for the privacy constraints of each.
The family of machine learning algorithms required for this setting differs
considerably from those required for traditional training pipelines, many of
which require data to be stored in a single machine with compute performed by
a single node. Federated learning attempts to bridge this gap, through
algorithms that allow multiple data owners—or clients—to collaboratively train
and use a shared model without sharing local data. Clients are then free to
employ whichever privacy preserving strategies they feel necessary, such as
differential privacy (Dwork et al., 2006; Jälkö et al., 2017). Federated
learning algorithms can be helpful in a wide variety of settings. For example,
a hospital may wish to develop a predictive model for a particular medical
diagnosis. However, the data collected by an individual hospital may be sparse
or biased towards a particular demographic. This has a detrimental effect on
properties of the trained model, such as accuracy, robustness and fairness.
Instead, it is desirable to collaborate with other hospitals to develop a
shared model that utilises the data collected by each. Federated learning
algorithms offer a route through which to achieve this whilst satisfying
privacy constraints. Federated learning is challenging in practice as modern
datasets are often inhomogeneously partitioned in terms of both the quantity
of data at each client and the distribution of data at each client, the
effects of which we describe in Section D. Consider the differences in data
collected by a teaching metropolis hospital and a rural community hospital,
for example. Further, communication channels between clients can be unreliable
and costly making structured communication protocols difficult to implement,
and those clients may also have different computational resources. Developing
efficient algorithms that are robust to such difficulties is an active area of
research that has received growing attention from the machine learning
community.

Of great importance in many real world applications is the ability to reason
in the presence of uncertainty. Probabilistic machine learning provides a
principled and elegant framework through which this can be achieved
(Ghahramani, 2015). However, performing inference is often infeasible in
practice and demands the use of approximations. A widely used approach is
variational inference (VI), in which an approximation to the posterior
distribution over model parameters is obtained as the solution to an
optimisation problem. Whilst widely applied to a range of popular machine
learning models—including Gaussian process models, latent topic models and
deep generative models—vanilla VI does not immediately extend to the federated
learning setting as it demands access to the entire dataset or unbiased
samples thereof.

The goal of this paper is to develop a framework, termed partitioned
variational inference (PVI), which generalises VI to the federated learning
setting. PVI unifies a wealth of existing schemes, providing practitioners
with the flexibility to develop unique methods whilst respecting the
constraints of a federated setting. We briefly summarise the contributions of
this paper, focusing on the suitability of PVI for federated learning and the
viewpoint of PVI as a unification of local and global VI methods.

###  1.1 Partitioned Variational Inference: Theoretical Contribution

The main theoretical contribution of this paper is to introduce PVI, described
in Section 2, a framework through which VI can be performed in the federated
setting. We generalise and derive new supporting theory (including PVI fixed-
point optimisation and mini-batch approximation), and demonstrate that PVI
satisfies a number of properties that make it an attractive choice for
practitioners.

###  1.2 Unification

Whilst the extension of VI to the federated learning setting motivates the
development of PVI, the framework can be viewed as a generalisation of pre-
existing local and global VI methods. From this viewpoint, in Section 4 PVI is
used to connect a large literature that has become fragmented with separated
strands of related, but mutually uncited work. More specifically we unify work
on: online VI (Ghahramani and Attias, 2000; Sato, 2001; Broderick et al.,
2013; Bui et al., 2017b; Nguyen et al., 2018); global VI (Sato, 2001; Hensman
et al., 2012; Hoffman et al., 2013; Salimans and Knowles, 2013; Sheth and
Khardon, 2016a; Sheth et al., 2015; Sheth and Khardon, 2016b); local VI
(Knowles and Minka, 2011; Wand, 2014; Khan and Lin, 2017); power EP and
related algorithms (Minka, 2001, 2004; Li et al., 2015; Hasenclever et al.,
2017; Gelman et al., 2014); and stochastic mini-batch variants of these
algorithms (Hoffman et al., 2013; Li et al., 2015; Khan and Lin, 2017).
Figures 2 and 3 and Table 1 present a summary of these relationships in the
context of PVI.

###  1.3 Partitioned Variational Inference: Empirical Contribution

In Section 5, we provide extensive experimentation showing that PVI is able to
recover the global VI solution for logistic regression models and Bayesian
neural networks in the presence of varying degrees and forms of inhomogeneous
data partitions. We outline the inadequacies of existing approaches for
performing approximate inference in the federated learning setting,
highlighting how these shortcomings are addressed by the PVI framework. The
experiments shed light on the settings in which different implementations of
PVI perform well, providing guidance to practitioners.

##  2 Partitioned Variational Inference

In this section, we introduce partitioned variational inference, a framework
that extends VI to the federated learning setting as well as encompassing many
other approaches to VI. We begin by framing PVI in terms of a series of
variational free-energy optimisation problems performed locally by individual
clients, proving several key properties of the algorithm that reveal the
relationship to global VI. We assume the presence of a central server which
coordinates communication between clients; however, it is important to note
that a central server is not a requirement provided each client is able to
communicate with all others. Skeleton code for the PVI algorithm is provided
in Algorithm 1, with different choices for step 1 yielding three variants of
PVI with different properties. In order to keep the development clear, we have
separated most of the discussion of related work into Section 4.

###  2.1 Setup

Consider the task of modelling a dataset $\mathbf{y}$ partitioned into $M$
local datasets $\\{\mathbf{y}_{m}\\}_{m=1}^{M}$, one for each of $M$ clients,
with a parametric probabilistic model defined by the prior
$p(\boldsymbol{\mathbf{\theta}}|\boldsymbol{\mathbf{\epsilon}})$ over
parameters $\boldsymbol{\mathbf{\theta}}$ and the likelihood function
$p(\mathbf{y}|\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})=\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})$.
For simplicity, we assume for the moment that the hyperparameters
$\boldsymbol{\mathbf{\epsilon}}$ are fixed and suppress them to lighten the
notation. We will discuss hyperparameter optimisation in Section A.

Exact Bayesian inference in this class of models is in general intractable, so
we resort to VI, in which we find the member of a family of tractable
distributions $\mathcal{Q}$ which minimising the KL-divergence between itself
and the true posterior distribution.

$q^{*}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,min}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\ {p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})}\right).$

This is equivalent to maximising the global (negative) free-energy
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$, which serves as a lower bound
to the log marginal likelihood $\log p(\mathbf{y})$:

$\displaystyle q^{*}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\underbrace{\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\log
p(\mathbf{y}|\boldsymbol{\mathbf{\theta}})\right]-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\
{p(\boldsymbol{\mathbf{\theta}})}\right)}_{\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}$
$\displaystyle=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\log
p(\mathbf{y})-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\ \|\
{p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})}\right).$

Maximisation of the global free-energy therefore jointly obtains an
approximation to the posterior, $q^{*}(\boldsymbol{\mathbf{\theta}})\approx
p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})$, and an approximation to the log
marginal likelihood,
$\mathcal{F}(q^{*}(\boldsymbol{\mathbf{\theta}}))\approx\log p(\mathbf{y})$,
two quantities of interest in probabilistic machine learning. See Wainwright
et al. (2008) and Zhang et al. (2018) for two excellent reviews of VI.
Throughout this paper, we refer to maximisation of the global free-energy as
global VI. Unfortunately, evaluating the global free-energy directly requires
access to the entire dataset, which violates the core constraint of federated
learning. Rather than considering a single, global distribution, PVI
decomposes the variational approximation in a manner that mimics the
decomposition of the true posterior,

$\displaystyle p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})$
$\displaystyle=\frac{1}{\mathcal{Z}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})$
(1)
$\displaystyle\approx\frac{1}{\mathcal{Z}_{q}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})=q(\boldsymbol{\mathbf{\theta}}),$

where $\mathcal{Z}_{q}$ is the normalising constant of the approximate
posterior and $\mathcal{Z}=p(\mathbf{y})$ is the normalising constant of the
true posterior. $\mathcal{Z}$ is also known as the marginal likelihood. This
decomposition is motivated by the fact that each client has access local data
$\mathbf{y}_{m}$, which shall be used to refine the approximate likelihood
$t_{m}(\boldsymbol{\mathbf{\theta}})$ such that it closely approximates the
effect the true likelihood term
$p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})$ has on the posterior. As we
typically restrict the approximate likelihoods to lie within the same
exponential family as the prior, computing the normalisation constant
$\mathcal{Z}_{q}$ is straightforward and is generally not required unless an
approximation to the log marginal likelihood is desired (see 2.2). In Section
2.4, we demonstrate that for one variant of PVI the product
$p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})$
is automatically normalised ($\mathcal{Z}_{q}=1$) regardless of the form of
the approximate likelihoods. Finally, we will show in Section 2.3 that PVI
will return an approximation to the log marginal likelihood
$\log\mathcal{Z}=\log p(\mathbf{y})$ in addition to the approximation of the
posterior.

###  2.2 PVI Skeleton

Algorithm 1 details the PVI algorithm. We describe each of the three steps in
detail below.

  1. 1

(_Server: client selection step._) At each iteration $i$, a central server
selects a set of approximate likelihoods
$\\{t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})\\}_{k\in b_{i}}$ to refine
according to a schedule $b_{i}\subseteq\\{1,\ldots M\\}$.

  2. 2

(_Client: update step._) Each approximate factor
$t_{k}(\boldsymbol{\mathbf{\theta}})$ is updated by communicating the most
recent approximate posterior $q^{(i-1)}(\boldsymbol{\mathbf{\theta}})$ to each
client $k\in b_{i}$, each of which then solves the local optimisation problem

$q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\underbrace{\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}}_{\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}}))}$
(2)

where the optimisation is over a tractable family $\mathcal{Q}$. We refer to
$\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}}))$ as the local
(negative) variational free-energy for client $k$. Given
$q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})$, the updated approximate
likelihood is found by division

$t_{k}^{(i)}(\boldsymbol{\mathbf{\theta}})\propto\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}}).$

The change in approximate likelihoods,
$\Delta^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\propto\frac{t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}$,
is communicated back to the server.

  3. 3

(_Server: update step._) Finally, the approximate posterior is updated as

$q^{(i)}(\boldsymbol{\mathbf{\theta}})\propto
q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\prod_{k\in
b_{i}}\Delta^{(i)}_{k}(\boldsymbol{\mathbf{\theta}}).$

Through modification of the update schedule executed by the server, three
variants of PVI can be implemented:

  1. 1.

sequential, in which a single approximate likelihood is refined at each
iteration;111The order in which the approximate likelihoods are updated can
make a difference in the convergence of PVI. This is discussed in more detail
in Section 2.4.

  2. 2.

synchronous, in which all approximate likelihoods are refined at each
iteration;

  3. 3.

asynchronous, in which approximate likelihoods are updated whenever clients
become available for performing local computation.

Algorithm 1 Partitioned Variational Inference

Input: data partition $\\{\mathbf{y}_{1},\ldots,\mathbf{y}_{M}\\}$, prior
$p(\boldsymbol{\mathbf{\theta}})$.

Initialise:  $\displaystyle t^{(0)}_{m}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle\vcentcolon=1\text{ for all }m=1,2,\ldots,M.$ $\displaystyle
q^{(0)}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle\vcentcolon=p(\boldsymbol{\mathbf{\theta}}).$

for $i=1,2,\ldots$ until convergence do

1 _Server: client selection step._

$b_{i}\vcentcolon=$ set of indices of the next approximate likelihoods to
refine.

Communicate $q^{(i-1)}(\boldsymbol{\mathbf{\theta}})$ to each client in
$b_{i}$.

2 _Client: update step._

for $k\in b_{i}$ do

Compute the new local approximate posterior:
$q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\vcentcolon=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}.$

Update the approximate likelihood:  $\displaystyle
t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle\propto\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$

Communicate
$\Delta^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\propto\frac{t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}$
to server.

end for

3 _Server: update step._ $q^{(i)}(\boldsymbol{\mathbf{\theta}})\propto
q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\prod_{k\in
b_{i}}\Delta^{(i)}_{k}(\boldsymbol{\mathbf{\theta}}).$

end for

###  2.3 PVI Properties

The three implementations, alongside their strengths and weaknesses, are
described in more detail in Section 2.4. Before doing so, we shall state three
properties of PVI, derived in Section B: 1) the local free-energy optimisation
is equivalent to a variational KL optimisation; 2) the sum of local free-
energies is equal to the global free-energy; and 3) any fixed point of the
algorithm is also a local optimum of global VI. These properties apply for
general $\mathcal{Q}$, and are not limited to the exponential
family.222However, we will only consider exponential family approximations in
the experiments in Section 5.

######  Property 2.1.

Maximising the local free-energy
$\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}}))$ is equivalent to the
KL optimisation

$\displaystyle
q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,min}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathrm{KL}\left(q(\boldsymbol{\mathbf{\theta}})~{}\|~{}\widehat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\right),$

where
$\widehat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\frac{1}{\widehat{\mathcal{Z}}^{(i)}_{k}}\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}{t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})=\frac{1}{\widehat{\mathcal{Z}}^{(i)}_{k}}p(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\prod_{m\notin
k}t^{(i-1)}_{m}(\boldsymbol{\mathbf{\theta}})$ is known as the tilted
distribution in the expectation propagation (EP, Minka, 2001) literature and
is intractable.

The proof is straightforward (see Section B.1). The tilted distribution can be
justified as a sensible target as it removes the approximate likelihood
$t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$ from the current approximate
posterior and replaces it with the true likelihood
$p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})$. In this way, the tilted
distribution comprises one true likelihood, $M-1$ approximate likelihoods and
the prior. The KL optimisation then incorporates the true likelihood’s effect,
in the context of the approximate likelihoods and the prior.

######  Property 2.2.

Let
$q(\boldsymbol{\mathbf{\theta}})=\frac{1}{\mathcal{Z}_{q}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})$,
$\mathcal{F}_{m}(q_{m}(\boldsymbol{\mathbf{\theta}}))=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q_{m}(\boldsymbol{\mathbf{\theta}})t_{m}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
be the local free-energy for client $m$ and
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
be the global free-energy. We have

$\displaystyle\sum_{m=1}^{M}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))+\log\mathcal{Z}_{q}=\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$

i.e. the sum of the local free-energies and normalisation constant is equal to
the global free-energy.

The proof is also straightforward (see Section B.2). The importance of this
result is that we can now perform efficient and principled hyperparameter
optimisation, which we describe in Section A. This additional utility of PVI
differentiates it from alternative schemes such as EP which, in general, do
not provide a lower bound on the log marginal likelihood and so cannot be used
for principled hyperparameter optimisation.

######  Property 2.3.

Let
$q^{*}(\boldsymbol{\mathbf{\theta}})=\frac{1}{\mathcal{Z}_{q^{*}}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t^{*}_{m}(\boldsymbol{\mathbf{\theta}})$
be a fixed point of Algorithm 1,
$\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q^{*}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})t^{*}_{m}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
be the local free-energy for client $m$, and
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
be the global free-energy. We have

$q^{*}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))\
\forall m\quad\implies\quad
q^{*}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$

i.e. a PVI fixed point $q^{*}(\boldsymbol{\mathbf{\theta}})$ is an optimum of
global VI.

This result is more complex to show, but can be derived by computing the
derivative and Hessian of the global free-energy and substituting into these
expressions the derivatives and Hessians of the local free-energies (see
Section B.3). The fact that the fixed point of PVI recovers a global VI
solution (both the optimal $q(\boldsymbol{\mathbf{\theta}})$ and the global
free-energy) is the main theoretical justification for employing PVI. However,
we do not believe that there is a Lyapunov function for PVI, indicating that
it may potentially oscillate or diverge. Although in practice we found that
damping prevented this behaviour, if convergence is required then, in theory,
the double-loop algorithm of Heskes and Zoeter (2002), later developed by
Hasenclever et al. (2017), can be employed which guarantees convergence for
power EP (Minka, 2004) which, as we shall discuss in Section 4, can recover
the solution of PVI. In practice, running the double-loop algorithm to
convergence is slow and the inner-loop is often truncated, typically to just a
single iteration. This removes any guarantee of convergence.333In fact,
running the inner-loop algorithm of the double-loop algorithm for a single
iteration recovers the original power EP algorithm. We do not believe this to
be widely known by the research community.

Figure 1: Steps of the PVI algorithm.

###  2.4 The Three Implementations of PVI

We now describe the three implementations of PVI described in Section 2.2,
their strengths and weaknesses, and the settings in which they are most
suitable.

####  2.4.1 Sequential PVI

In sequential PVI, a single approximate likelihood is refined at each
iteration. Sequential PVI enjoys the property that the update for the
posterior $q^{(i)}(\boldsymbol{\mathbf{\theta}})$ is automatically normalised
when we set
$t^{(i)}_{k}=\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$
and
$\Delta^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\frac{t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}$.
This allows us to be more flexible with our choice of approximate likelihoods
as we do not need to restrict them such that the normalisation constant can be
computed. The proof is straightforward:

$\displaystyle q^{(i)}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle=p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t^{(i)}_{m}(\boldsymbol{\mathbf{\theta}})=t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})p(\boldsymbol{\mathbf{\theta}})\prod_{m\neq
k}t_{m}^{(i-1)}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle=\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\underbrace{p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}^{(i-1)}(\boldsymbol{\mathbf{\theta}})}_{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}=q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}}).$

Since $q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})$ is normalised, so too is
$q^{(i)}(\boldsymbol{\mathbf{\theta}})$.

Damped updates are not required for sequential PVI, making it communication
efficient as damping often leads to a greater number of smaller updates to
attain convergence.444The total number of communications is taken to be the
number of times clients are communicated with, rather than the number of times
the server sends out the current approximate posterior. This offers an
important advantage when the cost of communication is large. However, as $K-1$
clients have to wait whilst a single approximate likelihood is updated at each
iteration, sequential PVI exhibits poor time efficiency. Further, whilst we
found sequential PVI to perform well in the homogeneous setting, we found that
in the presence of a large number of inhomogeneous data partitions sequential
PVI can struggle to converge without damping. We believe that inhomogeneous
settings generally require smaller updates to be made per communication round
so that all the approximate factors can come to an agreement, which damping
achieves.

The order in which approximate likelihoods are updated will affect the rate at
which sequential PVI converges. The simplest choice is to use a fixed, pre-
defined order, which is common when implementing sequential variants of EP or
belief propagation (BP) (Pearl, 1988). Although easy to implement, this is
unlikely to achieve the fastest convergence possible. Elidan et al. (2012)
propose an alternative scheme for sequential BP, in which the order in which
messages are updated depends on the magnitude of the previous update. They
demonstrate empirically that their residual BP algorithm achieves faster
convergence than the fixed-order baseline. A similar scheme for sequential PVI
could result in improvements in the rate of convergence. However, it is not
immediately clear which metric to use to measure the magnitude of the previous
update. Simple vector norms may not be applicable when many of the parameters
have little to no effect on the predictive distribution, such as in neural
network models.

Finally, we note that if each client is visited once then sequential PVI is
equivalent to online/continual learning. We discuss this relationship in more
detail in Section 4.3.

####  2.4.2 Synchronous PVI

In synchronous PVI, all approximate likelihoods are updated simultaneously at
each iteration. Unlike the sequential case, the aggregation of approximate
likelihoods is not automatically normalised. This typically means that we are
restricted to using approximate likelihoods that lie in the same exponential
family as the prior, so that the normalisation constant can be computed. A
second limitation of synchronous PVI is that the product aggregation of
approximate likelihoods may not be normalisable. To see this, consider the
updated approximate posterior $q^{(i)}(\boldsymbol{\mathbf{\theta}})$:

$\displaystyle q^{(i)}(\boldsymbol{\mathbf{\theta}})$ $\displaystyle\propto
p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}^{(i)}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle\propto
p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}\frac{q^{(i)}_{m}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{m}(\boldsymbol{\mathbf{\theta}})$
$\displaystyle=\left(\frac{1}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\right)^{M}\underbrace{p(\boldsymbol{\mathbf{\theta}})\left(\prod_{m=1}^{M}t^{(i-1)}_{m}(\boldsymbol{\mathbf{\theta}})\right)}_{\propto
q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\left(\prod_{m=1}^{M}q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\right)$
$\displaystyle\propto\frac{\prod_{m=1}^{M}q^{(i)}_{m}(\boldsymbol{\mathbf{\theta}})}{\left[q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\right]^{K-1}}.$

This bears close resemblance to the update scheme of the Bayesian committee
machine (BCM, Tresp, 2000). There is no guarantee that the update is
normalisable, which results in an improper distribution for
$q^{(i)}(\boldsymbol{\mathbf{\theta}})$. To overcome this, we use damping of
the approximate likelihoods:
$t_{k}^{(i)}(\boldsymbol{\mathbf{\theta}})\propto\left(\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\right)^{\rho}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$
for some $\rho\in(0,1]$. We found choosing $\rho\propto\frac{1}{M}$ results in
stable convergence.

The use of damping makes synchronous PVI less communication efficient than
sequential PVI, particularly in the case of a large number of clients when
more damping is necessary. Yet, synchronous PVI can exhibit better time
efficiency as all approximate likelihoods are updated in parallel.555Although
not as time efficient as one might hope due to the need for damping. In fact,
for homogeneous data partitions, choosing $\rho\propto\frac{1}{M}$ means both
sequential and synchronous PVI effectively update a single approximate
likelihood at each iteration, thus having similar time efficiencies. We found
synchronous PVI to be particularly effective in the presence of inhomogeneous
data partitions. In general, the PVI approximate posterior will be biased
towards the most recently visited dataset. However, as synchronous PVI updates
the posterior by considering the entire dataset, it performs well even when
partitions are inhomogeneous. In contrast, sequential PVI can perform poorly
in the presence of inhomogeneous partitions, and can even fail to converge
when $M$ is large unless damping is used.

####  2.4.3 Asynchronous PVI

Rather than waiting for all clients to finish their local free-energy
optimisation, in asynchronous PVI an updated approximate likelihood is
integrated into the approximate posterior as soon as it becomes available. The
corresponding client then receives the updated posterior from the server and
resumes local free-energy optimisation. This is particularly useful when
communicating over an unreliable channel, or when the update-frequencies are
highly inhomogeneous due to inhomogeneous hardware or dataset sizes. As the
approximate posterior is changed whilst clients perform local free-energy
optimisation, the returned approximate likelihood may be stale given the
update context. This also means that the updated approximate posterior may be
improper, necessitating the use of damping.666We found a similar damping
factor to that used in synchronous PVI was required for asynchronous PVI.

Out of all three PVI implementations, the communication efficiency of
asynchronous PVI can be expected to be the worst.777To see why asynchronous
PVI can be less communication efficient than synchronous PVI, consider the
case in which a single client has significantly less data than the others. The
small data client will finish its local optimisation faster than the others
and, assuming clients communicate as soon as their local optimisation
converges, communicate its update approximate likelihood back to the server.
In synchronous PVI, the small client must wait until the other large data
clients have finished their optimisation before proceeding with its next local
optimisation. However, in asynchronous PVI the small client can perform its
next local optimisation immediately—this optimisation is redundant as it will
have received no information from any of the large data clients. Thus,
unnecessary communications can be performed. Yet, asynchronous PVI can be
expected to be the most time-efficient.

####  2.4.4 Summary

Out of the three PVI implementations, sequential PVI offers the most
theoretically appealing properties and is the most communication efficient
and, as no damping is required, it is also often time efficient. However, due
to the sequential visiting of data partitions and the accumulation of
approximation errors, sequential PVI can struggle to converge when data
partitions are inhomogeneous. Further, the ordering in which partitions are
visited is likely to effect the rate of convergence, yet it is unclear which
ordering to use. Synchronous PVI is in theory the most time efficient;
however, it requires careful tuning of the damping factor to ensure this
efficiency is realised without encountering improper distributions.
Asynchronous PVI is the least communication efficient and requires damping,
yet can be the most-time efficient and is the only implementation which can be
deployed without modification in settings where communication between clients
is unreliable.

##  3 Client Update Step: Approaches for Optimising the Local Free-energies

Having laid out the general framework for PVI, what remains to be decided is
the method used for optimising the local free-energies at each client. We
consider three choices: analytic updates, off-the-shelf optimisation methods
and fixed-point iterations, as well as discussing how stochastic
approximations can be combined with these approaches.

###  3.1 Analytic Local Free-energy Updates

Each local free-energy
$\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}}))$ is equivalent in form
to a global free-energy with an effective prior
$p_{\mathrm{eff}}(\boldsymbol{\mathbf{\theta}})\propto
q^{(i-1)}(\boldsymbol{\mathbf{\theta}})/t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$.
Thus, in conjugate exponential family models the KL optimisations will be
available in closed-form, for example in GP regression, and can be substituted
back into the local variational free-energies to yield locally-collapsed
bounds that are useful for hyperparameter optimisation (Bui et al.,
2017a).888This leads to a potential advantage of PVI over global VI in terms
of speed of convergence: collapsed bounds can be leveraged on local datasets
within each client where an application to the entire dataset would be
computationally intractable.

###  3.2 Off-The-Shelf Optimisers for Local Free-Energy Optimisation

If analytic updates are not tractable, the local free-energy optimisations can
be carried out using standard optimisers. If the partitions are small enough,
non-stochastic optimisers such as BFGS be leveraged.999Again, this offers a
potential computational advantage of PVI over global VI: non-stochastic
optimisers can be deployed in the large data setting by partitioning the data
across multiple clients. Of course, if a stochastic approximation like Monte
Carlo VI is employed for each local optimisation, stochastic optimisers such
as RMSProp (Tieleman and Hinton, 2012) or Adam (Kingma and Ba, 2014a) might be
more appropriate choices. In all cases, since the local free-energy is
equivalent in form to a global free-energy with an effective prior
$p_{\mathrm{eff}}(\boldsymbol{\mathbf{\theta}})\propto
q^{(i-1)}(\boldsymbol{\mathbf{\theta}})/t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$,
PVI can be implemented via trivial modification to existing code for global
VI. Whilst other local VI approaches, such as such as variational message
passing (Winn et al., 2005; Winn and Minka, 2009; Knowles and Minka, 2011),
can be deployed in the federated learning setting, they do not share the same
flexibility and ease of implementation of PVI, instead requiring bespoke and
closed-form updates for different likelihoods and cavity distributions.

For gradient-based updates, the training dynamics of global VI can be
reproduced by synchronous PVI if a single local update is performed
synchronously by all clients:

######  Property 3.1.

Terminating the local optimisation step in synchronous PVI after a single
gradient-based update results in identical dynamics as global VI for
$q(\boldsymbol{\mathbf{\theta}})$, given by the following equation, regardless
of the partition of the data:

$\boldsymbol{\mathbf{\eta}}^{(i)}_{q}=\boldsymbol{\mathbf{\eta}}^{(i-1)}_{q}+\rho\left.\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\mathcal{F}(\boldsymbol{\mathbf{\eta}}_{q})\right|_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{(i-1)}_{q}}$

where
$\mathcal{F}(\boldsymbol{\mathbf{\eta}}_{q})=\sum_{n=1}^{N}\mathbb{E}_{q}\left[\log
p(y_{n}|\boldsymbol{\mathbf{\theta}})\right]-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\ {p(\boldsymbol{\mathbf{\theta}})}\right)$ is the global free-energy.

See Section B.5 for the proof. This demonstrates a further equivalence between
PVI and global VI: employing a single iteration of gradient based updates for
synchronous PVI is the same as employing gradient based updates for global
VI.101010It should be noted that global (batch) VI retains only a single set
of global natural parameters, and thus is more memory efficient than PVI which
retains $M$ sets of local natural parameters. The equivalence provides us with
an alternative method of reproducing the dynamics of global VI in the
federated learning setting.111111The alternative method which recovers global
VI with biased mini-batches is to communicate $N=\sum_{k=1}^{K}N_{k}$ and
$p(\boldsymbol{\mathbf{\theta}})$ to each client, and have each client
sequentially update the approximate posterior by performing a single/several
optimisation step(s) the modified free-energy
$\mathcal{F}_{k}(q(\boldsymbol{\mathbf{\theta}}))=\frac{N_{k}}{N}\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\ {p(\boldsymbol{\mathbf{\theta}})}\right)$. This method, however, requires
clients to release their number of datapoints $N_{k}$. The fact that PVI
supports multiple optimisation steps per communication is key to achieving
communication efficiency. In most settings performing more than a single
optimisation step improves communication efficiency, although in highly
inhomogeneous settings performing fewer optimisation steps might be
advantageous as this has an effect comparable to damping.

It is important to stress that if the local optimisation step is run to
convergence then the partition of the data does influence the dynamics for
$q(\boldsymbol{\mathbf{\theta}})$, and the dynamics of global VI are not
recovered—if they were, PVI would offer no improvement in communication
efficiency over this federated implementation of global VI.

###  3.3 Local Free-Energy Fixed Point Updates, Natural Gradient Methods, and
Mirror Descent

An alternative to using off-the-shelf optimisers is to derive fixed-point
update equations by zeroing the gradients of the local free-energy. These
fixed-point updates have elegant properties for approximate posterior
distributions that are in the exponential family.

######  Property 3.2.

Let the prior and approximate likelihood factors be in the un-normalised
exponential family
$t_{m}(\boldsymbol{\mathbf{\theta}})=t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{m})=\exp(\boldsymbol{\mathbf{\eta}}_{m}^{\intercal}T(\boldsymbol{\mathbf{\theta}}))$
so that the variational distribution is in the normalised exponential family
$q(\boldsymbol{\mathbf{\theta}})=\exp(\boldsymbol{\mathbf{\eta}}_{q}^{\intercal}T(\boldsymbol{\mathbf{\theta}})-A(\boldsymbol{\mathbf{\eta}}_{q}))$.
Let $q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})$ be the stationary point of the
local free-energy such that
$\left.\frac{\mathrm{d}\mathcal{F}_{k}^{(i)}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\right|_{q=q^{(i)}_{k}}=\boldsymbol{\mathbf{0}}$,
then, noting that the update step
$t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$
implies
$\boldsymbol{\mathbf{\eta}}^{(i)}_{k}=\boldsymbol{\mathbf{\eta}}^{(i)}_{q_{k}}-\boldsymbol{\mathbf{\eta}}^{(i-1)}_{q}+\boldsymbol{\mathbf{\eta}}^{(i-1)}_{k}$,

$\displaystyle\boldsymbol{\mathbf{\eta}}^{(i)}_{k}$
$\displaystyle=\left.\mathbb{C}^{-1}\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\mathbb{E}_{q}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]\right|_{q=q^{(i)}_{k}}$

where
$\mathbb{C}\vcentcolon=\frac{\mathrm{d}^{2}A(\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}=\mathrm{cov}_{q(\boldsymbol{\mathbf{\theta}})}[T(\boldsymbol{\mathbf{\theta}})T^{\intercal}(\boldsymbol{\mathbf{\theta}})]$
is the Fisher Information. Moreover, the Fisher Information can be written as
$\mathbb{C}=\frac{\mathrm{d}\boldsymbol{\mathbf{\mu}}_{q}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}$
where
$\boldsymbol{\mathbf{\mu}}_{q}=\mathbb{E}_{q}\left[T(\boldsymbol{\mathbf{\theta}})\right]$
is the mean parameter of $q(\boldsymbol{\mathbf{\theta}})$. Hence,

$\displaystyle\boldsymbol{\mathbf{\eta}}^{(i)}_{k}$
$\displaystyle=\left.\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\mu}}_{q}}\mathbb{E}_{q}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]\right|_{q=q^{(i)}_{k}}.$
(3)

For some approximate posterior distributions
$q(\boldsymbol{\mathbf{\theta}})$, taking derivatives of the average log-
likelihood with respect to the mean parameters is analytic (e.g.
Gaussian121212However, in the stochastic case where Monte Carlo samples are
employed to estimate these quantities, it can reduce variance to explicitly
estimate the Fisher information, even if the derivative with respect to the
mean parameters is analytic, see Salimans and Knowles (2013).) and for some it
is not (e.g. gamma).

These conditions, derived in Section B.4, can be used as fixed point
equations. That is, they can be iterated possibly with damping $\rho$,

$\displaystyle\boldsymbol{\mathbf{\eta}}^{(i)}_{k}\leftarrow(1-\rho)\boldsymbol{\mathbf{\eta}}^{(i)}_{k}+\left.\rho\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\mu}}_{q}}\mathbb{E}_{q}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]\right|_{q=q^{(i)}_{k}}.$
(4)

These iterations, which form an inner-loop in PVI, are themselves not
guaranteed to converge (there is no Lyapunov function in general and so, for
example, the local free-energy will not reduce at every step).

The fixed point updates are the natural gradients of the local free-energy and
the damped versions are natural gradient ascent (Sato, 2001; Hoffman et al.,
2013; Khan and Nielsen, 2018; Khan and Rue, 2021). The natural gradients could
also be used in other optimisation schemes (Hensman et al., 2012; Salimbeni
and Hensman, 2018). The damped updates are also equivalent to performing
mirror-descent (Raskutti and Mukherjee, 2015; Khan and Lin, 2017), a general
form of proximal algorithm (Parikh and Boyd, 2014) that can be interpreted as
trust-region methods. For more details about the relationship between these
methods, see Section B.8.

For these types of updates, there is an equivalent relationship between
synchronous PVI and global VI as 3.1:

######  Property 3.3.

Terminating the local optimisation step in synchronous PVI after a single
iteration of the damped fixed-point update (4) results in identical dynamics
as global VI for $q(\boldsymbol{\mathbf{\theta}})$ given by the following
equation, regardless of the partitioning of the data

$\displaystyle\boldsymbol{\mathbf{\eta}}^{(i)}_{q}$
$\displaystyle=(1-\rho)\boldsymbol{\mathbf{\eta}}^{(i-1)}_{q}+\rho\left(\boldsymbol{\mathbf{\eta}}_{0}+\left.\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\mu}}_{q}}\mathbb{E}_{q}\left[\log
p(\mathbf{y}|\boldsymbol{\mathbf{\theta}})\right]\right|_{q=q^{(i-1)}}\right)$
$\displaystyle=(1-\rho)\boldsymbol{\mathbf{\eta}}^{(i-1)}_{q}+\rho\left(\boldsymbol{\mathbf{\eta}}_{0}+\sum_{n=1}^{N}\left.\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\mu}}_{q}}\mathbb{E}_{q}\left[\log
p(y_{n}|\boldsymbol{\mathbf{\theta}})\right]\right|_{q=q^{(i-1)}}\right).$

See Section B.6 for the proof. This provides a method through which damped
fixed-point updates of global VI can be performed whilst conforming to the
constraints of federated learning. As in the case of gradient-based updates,
PVI achieves improved communication efficiency relative to the federated
implementation of global VI by iterating the fixed-point update equation on
each client until convergence, before communicating the result back to the
server.

###  3.4 Mini-Batch Approximations and Hyperparameter Learning

In the case where each partition includes a large number of data points,
stochastic approximations can lead to faster convergence of the local
optimisation procedure. There are two different ways of splitting up the local
partitions:  1. fixed sub-partitions of data;  2. randomly sampled mini-
batches of data.  Section B.7 compares and contrasts these two difference
approaches using damped fixed-point updates.

Finally, many probabilistic models rely on an appropriate choice of the
hyperparameters $\boldsymbol{\mathbf{\epsilon}}$ to achieve good predictive
performance on a task. The PVI framework offers a principled method for both
performing hyperparameter optimisation through maximisation of the approximate
marginal likelihood and performing approximate inference over hyperparameters
if a distributional estimate is desired. See Section A for details.

##  4 Unification of Previous Work

The preceding sections have introduced PVI as a framework extending VI to the
federated learning setting. However, if we depart from the federated learning
setting and consider the data partitioning scheme to be a choice made by the
practitioner, then PVI can more generally be interpreted as a unification of
existing approaches to VI. These methods include global VI (Section 4.1),
local VI (Section 4.2), online VI (Section 4.3) and a number of methods based
on power EP (Sections 4.4, 4.5, 4.6 and 4.7). A schematic showing the
relationships between these methods at a high level is shown in Figure 2. The
literature has been organised into Figure 3 and Table 1.

Figure 2: Variational inference schemes encompassed by the PVI framework.

###  4.1 Global VI Fixed-Point Methods

There has been a long history of applying the fixed-point updates for global
VI (PVI where $M=1$). Sato (2001) derived them for conjugate exponential
family models, showing they recover the closed form updates for
$q(\boldsymbol{\mathbf{\theta}})$, and noting that damped fixed-point updates
are equivalent to natural gradient ascent with unit step size ($\rho=1$).
Sato’s insight was subsequently built upon by several authors. Honkela et al.
(2010) considered non-conjugate models, employed a Gaussian variational
distribution and used natural gradients to update the variational
distribution’s mean. Hensman et al. (2012) and Hoffman et al. (2013) applied
the insight to conjugate models when optimising collapsed variational free-
energies and deriving stochastic natural gradient descent, respectively.
Salimans and Knowles (2013) apply the fixed-points to non-conjugate models
where the expectations over $q(\boldsymbol{\mathbf{\theta}})$ are intractable
and use Monte Carlo to approximate them, but they explicitly calculate the
Fisher information matrix. Sheth and Khardon (2016a) and Sheth et al. (2015)
treat non-conjugate models with Gaussian latent variables, employ the
cancellation of the Fisher information, and analyse convergence properties.
Sheth and Khardon (2016b) further extend this to two level-models through
Monte Carlo essentially applying the Fisher information cancellation to
Salimans and Knowles (2013), but they were unaware of this prior work.

###  4.2 Local VI Fixed-Point Methods

Whereas global VI seeks a single approximation to the full posterior
distribution, local VI methods employ an approximate posterior distribution
which decomposes as a product of factors, with each factor being iteratively
refined. Typically, each factor involves only a subset of datapoints (PVI
where $M>1$) and may also include only a subset of variables over which the
posterior is defined. Similar to global VI, there has been a long history of
applying fixed-point updates in the local VI setting. Knowles and Minka (2011)
derive them for non-conjugate variational message passing (VMP), but
explicitly calculate the Fisher information matrix (except in a case where
$q(\boldsymbol{\mathbf{\theta}})$ was univariate Gaussian case where they do
employ the cancellation). Wand (2014) simplified VMP by applying the Fisher
information cancellation to the case where $q(\boldsymbol{\mathbf{\theta}})$
is multivariate Gaussian. The work is closely related to Salimans and Knowles
(2013) and Sheth and Khardon (2016b), since although these papers use fixed-
point updates for global VI, they show that these decompose over data points
and thereby derive mini-batch updates that closely resemble fixed-point per-
datapoint local VI. This is a result of 3.3.

Khan and Lin (2017) also extend VMP to employ Monte Carlo approximation and
the Fisher information cancellation. They were unaware of Wand (2014), but
extend this work by treating a wider range of approximate posteriors and
models, stochastic updates, and a principled approach to damping. Their
approach starts from a global perspective, and uses the independence of
likelihood terms to split up the posterior. Any partition of data can be used
for their method (including $M=1$), and natural-gradient updates must be used
on any non-conjugate terms. If the mini-batches are selected randomly at each
iteration, then a set of natural parameters for each datapoint must be stored.
However, if the mini-batches are kept fixed across iterations then only a
single set of natural parameters for each mini-batch needs to be maintained.
In doing so, their method becomes exactly equivalent to PVI using fixed-point
updates and terminating the local optimisation procedure after a single
iteration.

###  4.3 Online, Streaming, Incremental and Continual VI as a Single Pass of
PVI

If PVI makes a single sequential pass through the data, the approximate
likelihoods do not need to be explicitly computed or stored as data partitions
are not revisited. In this case PVI reduces to initialising the approximate
posterior to be the prior,
$q^{(0)}(\boldsymbol{\mathbf{\theta}})=p(\boldsymbol{\mathbf{\theta}})$ and
then proceeds by optimising a sequence of local free-energies

$\displaystyle
q^{(i)}(\boldsymbol{\mathbf{\theta}})\vcentcolon=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}.$

These have the form of standard variational inference with the prior replaced
by the previous variational distribution
$q^{(i-1)}(\boldsymbol{\mathbf{\theta}})$. This idea—combining the likelihood
from a new batch of data with the previous approximate posterior and
projecting back to a new approximate posterior—underpins online variational
inference (Ghahramani and Attias, 2000; Sato, 2001), streaming variational
inference (Broderick et al., 2013; Bui et al., 2017b), and variational
continual learning (Nguyen et al., 2018; Swaroop et al., 2019; Loo et al.,
2021). Early work on online VI used conjugate models and analytic updates
(Ghahramani and Attias, 2000; Sato, 2001; Broderick et al., 2013; Bui et al.,
2017b). This was followed by off-the-shelf optimisation approaches for non-
conjugate models (Bui et al., 2017b) and further extended to leverage Monte
Carlo approximations of the local-free energy (Nguyen et al., 2018). Recently
Zeno et al. (2018) use the variational continual learning framework of Nguyen
et al. (2018), but employ fixed-point updates instead. As these methods can be
interpreted as finding the functional approximation
$t^{(i)}(\boldsymbol{\mathbf{\theta}})=\frac{q^{(i)}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}$
involving just a single datapoint or a group of datapoints, they too can be
interpreted as local VI methods.

###  4.4 Power EP as a Local VI Fixed-Point Method

There is also an important relationship between PVI methods employing fixed-
point updates and power EP (PEP) (Minka, 2004, 2005). 4.1 below states that
the local VI fixed-point equations are recovered from the PEP algorithm as
$\alpha\rightarrow 0$.

######  Property 4.1.

The damped fixed-point equations are precisely those returned by the PEP
algorithm, shown in Algorithm 2, in the limit as $\alpha\rightarrow 0$.

Algorithm 2 One step of the PEP algorithm at the $i$-th iteration, for the
$k$-th data partition

Compute the tilted distribution:
$\hat{p}^{(i)}_{\alpha}(\boldsymbol{\mathbf{\theta}})=q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\left(\frac{p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{t_{k}^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\right)^{\alpha}$.

Moment match:
$q_{\alpha}(\boldsymbol{\mathbf{\theta}})=\mathrm{proj}(\hat{p}^{(i)}_{\alpha}(\boldsymbol{\mathbf{\theta}}))\;\text{such
that}\;\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[T(\boldsymbol{\mathbf{\theta}})\right]=\mathbb{E}_{\hat{p}^{(i)}_{\alpha}(\boldsymbol{\mathbf{\theta}})}\left[T(\boldsymbol{\mathbf{\theta}})\right]$.

Update the posterior distribution with damping $\rho$:
$q^{(i)}(\boldsymbol{\mathbf{\theta}})=\left(q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\right)^{1-\rho/\alpha}\left(q_{\alpha}(\boldsymbol{\mathbf{\theta}})\right)^{\rho/\alpha}$.

Update the approximate likelihood:
$t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\frac{q^{(i)}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$.

Although we suspect Knowles and Minka (2011) knew of this relationship, and it
is well known that PEP has the same fixed-points as VI in this case, it does
not appear to be widely known that variationally limited PEP yields exactly
the same algorithm as fixed-point local VI. See Section B.9 for the proof.

###  4.5 Alpha-Divergence EP as a Local VI Method With Off-The-Shelf
Optimisation

PVI is intimately related to alpha-divergence EP. If PVI’s KL divergence is
replaced by an alpha divergence
$\mathrm{D}_{\alpha}[p(\boldsymbol{\mathbf{\theta}})||q(\boldsymbol{\mathbf{\theta}})]=\frac{1}{\alpha(1-\alpha)}\int\left[\alpha
p(\boldsymbol{\mathbf{\theta}})+(1-\alpha)q(\boldsymbol{\mathbf{\theta}})-p(\boldsymbol{\mathbf{\theta}})^{\alpha}q(\boldsymbol{\mathbf{\theta}})^{1-\alpha}\right]\mathrm{d}\boldsymbol{\mathbf{\theta}}$,
we recover the alpha-divergence formulation of the PEP algorithm (Minka, 2004,
2005), which encompasses the current case as $\alpha\rightarrow 0$ and EP when
$\alpha\rightarrow 1$ (Minka, 2001). The updates using this formulation are
shown in Algorithm 3.

Algorithm 3 One step of the PEP algorithm, as in Algorithm 2, but with alpha
divergence minimisation

Compute the tilted distribution:
$\hat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\frac{p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{t_{k}^{(i-1)}(\boldsymbol{\mathbf{\theta}})}$.

Find the posterior distribution:
$q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\vcentcolon=\operatorname*{arg\,min}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathrm{D}_{\alpha}[\hat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})||q(\boldsymbol{\mathbf{\theta}})]$.

Update the approximate likelihood:
$t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})$.

The alpha divergence is typically very difficult to compute when more than one
non-Gaussian likelihood is included in a data partition $\mathbf{y}_{m}$,
meaning that for general alpha it would be appropriate to set $M=N$. The
variational KL is the exception as it decomposes over data points.

###  4.6 Stochastic Power EP as a Stochastic Global VI Fixed-Point Method

The stochastic PEP (SPEP) algorithm (Li et al., 2015) reduces the memory
overhead of PEP by maintaining a single likelihood approximation that
approximates the average effect a likelihood has on the posterior
$q(\boldsymbol{\mathbf{\theta}})=p(\boldsymbol{\mathbf{\theta}})t(\boldsymbol{\mathbf{\theta}})^{M}$.
Taking the variational limit of this algorithm, $\alpha\rightarrow 0$, we
recover global VI ($M=1$) with damped simplified fixed-point updates that
employ a stochastic (mini-batch) approximation (Hoffman et al., 2013).

######  Proposition 4.1.

The mini-batch fixed-point equations are precisely those returned by the SPEP
algorithm, shown in Algorithm 4 in the limit that $\alpha\rightarrow 0$.

Algorithm 4 One step of the SPEP algorithm at the $i$-th iteration, for the
$k$-th data partition

Compute the tilted distribution:
$\hat{p}^{(i)}_{k,\alpha}(\boldsymbol{\mathbf{\theta}})=q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\left(\frac{p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{t^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\right)^{\alpha}$.

Moment match:
$q^{(i)}_{k,\alpha}(\boldsymbol{\mathbf{\theta}})=\mathrm{proj}(\hat{p}^{(i)}_{k,\alpha}(\boldsymbol{\mathbf{\theta}}))\;\text{such
that}\;\mathbb{E}_{q^{(i)}_{k,\alpha}(\boldsymbol{\mathbf{\theta}})}\left[T(\boldsymbol{\mathbf{\theta}})\right]=\mathbb{E}_{\hat{p}^{(i)}_{\alpha}(\boldsymbol{\mathbf{\theta}})}\left[T(\boldsymbol{\mathbf{\theta}})\right]$.

Update the posterior distribution with damping $\rho$:
$q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\left(q^{(i-1)}(\boldsymbol{\mathbf{\theta}})\right)^{1-M\rho/\alpha}\left(q^{(i)}_{k,\alpha}(\boldsymbol{\mathbf{\theta}})\right)^{M\rho/\alpha}.$

Update the approximate likelihood:
$t^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})=\left(\frac{q^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})}{p(\boldsymbol{\mathbf{\theta}})}\right)^{1/M}$.

In this way the relationship between EP and SEP is the same as the
relationship between fixed-point PVI and fixed-point mini-batch global VI (see
Section B.7 where the two approaches differ by removing either an average
natural parameter or a specific one). Similarly, if we altered PVI to maintain
a single average likelihood approximation, as SEP does, we would recover mini-
batch global VI.

###  4.7 Distributed (Power) EP Methods

The convergent distributed PEP approach of Hasenclever et al. (2017) recovers
a version of PVI as $\alpha\rightarrow 0$ with convergence guarantees. The PVI
approach can be used in similar spirit to Gelman et al. (2014); Xu et al.
(2014); Hasenclever et al. (2017) who use EP to split up datasets into small
parts that are amenable to MCMC. Here PVI could be used to split up datasets
so that they are amenable for optimisation.

Table 1: Variational inference schemes encompassed by the PVI framework. (See
next page.) Selected past work has been organised into four categories: global
VI (PVI with $M=1$), local VI (PVI with $M>1$), online VI and PEP variants.
The citation to the work is provided along with the granularity of the method.
The optimisation used from the PVI perspective on this work is noted.
Abbreviations used here are: conjugate gradient (CG) and Monte Carlo (MC). The
model class that the scheme encompasses is noted (conjugate versus non-
conjugate) along with the specific models that the scheme was tested on. Model
abbreviations are: non-linear state-space model (NSSM), non-linear factor
analysis (NFA), latent Dirichlet allocation (LDA), Poisson mixed model (PMM),
heteroscedastic linear regression (HLR), sparse Gaussian processes (SGPs),
graphical model (GM), logistic regression (LR), beta-binomial (BB), stochastic
volatility model (SV), probit regression (PR), multinomial regression (MR),
neural network (NN), Bayesian neural network (BNN), gamma factor model (GFM),
Poisson gamma matrix factorisation (PGMF), mixture of Gaussians (MoG),
Gaussian latent variable (GLV), switching linear dynamical system (SLDS). If
the scheme proposed by the method has a name, this is noted in the final
column. Abbreviations of the inference scheme are: variational inference (VI),
stochastic variational inference (SVI), automatic differentiation VI (ADVI),
incremental VI (IVI), variational message passing (VMP), non-conjugate
variational message passing (NC-VMP), simplified NC-VMP (SNC-VMP), conjugate-
computation VI (CVI), power EP (PEP), alpha-divergence PEP (ADPEP), convergent
power EP (CPEP), stochastic natural gradient EP (SNEP), stochastic power EP
(SPEP), variational continual learning (VCL), Bayesian gradient descent (BGD).
Figure 3: The local VI framework unifies prior work. The granularity of the
approximation and the optimisation method employed are two fundamental
algorithmic dimensions that are shown as axes. Fixed-point updates are
identical to natural gradient ascent with unit step size. The models
encompassed by each paper are indicated by the color. See Table 1 for more
information.

##  5 Experiments

In this section we investigate the performance of sequential and synchronous
PVI. Our aim is to demonstrate that PVI can recover the solution of global VI
with greater communication efficiency in both simple and complex models, and
in settings reflective of those in which PVI is likely to be deployed.
Throughout, we make comparisons with global VI, VCL, streaming VB and BCM. We
briefly describe these methods below. Global VI is implemented in the
federated learning setting using synchronous PVI and performing a single local
optimisation step using stochastic gradient descent. As discussed in Section
3.2, this recovers the exact dynamics of global VI. A more communication
efficient but biased implementation of global VI is to communicate the current
approximate posterior to a single client and perform a single stochastic
update using a biased approximation of the global free-energy
$\tilde{\mathcal{F}}_{k}(q(\boldsymbol{\mathbf{\theta}}))=\frac{N_{k}}{N}\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\ {p(\boldsymbol{\mathbf{\theta}})}\right)$. We provide results for this
implementation in Section E.

  1. Bayesian committee machine:

the Bayesian committee machine (BCM) (Tresp, 2000) performs VI for each
dataset independently of the others, giving $M$ local approximate posteriors
$\\{q_{m}(\boldsymbol{\mathbf{\theta}})\\}_{m=1}^{M}$. The local approximate
posteriors are then aggregated. We consider two BCM strategies:

BCM (same): $\displaystyle
p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})\propto\frac{\prod_{m=1}^{M}p(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{p(\boldsymbol{\mathbf{\theta}})^{M-1}}\approx\frac{\prod_{m=1}^{M}q_{m}(\boldsymbol{\mathbf{\theta}})}{p(\boldsymbol{\mathbf{\theta}})^{M-1}}$
BCM (split): $\displaystyle
p(\boldsymbol{\mathbf{\theta}}|\mathbf{y})\propto\prod_{m=1}^{M}p(\boldsymbol{\mathbf{\theta}})^{N_{m}/N}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})\approx\prod_{m=1}^{M}q_{m}(\boldsymbol{\mathbf{\theta}}).$

As alluded to in Section 2.4, BCM (same) is equivalent to a single round of
synchronous PVI with no damping. Whilst the true global posterior decomposes
as a normalised product of local posteriors, there is no guarantee that the
normalised product of approximations to the local posteriors will form a good
approximation to the global posterior due to the interaction between
approximation errors. This is particularly so in the case of non-identifiable
models such as BNNs, for which each local mean-field Gaussian approximation
can be drawn to a different mode, and when they are combined the disagreement
between the local approximations leads to a poor posterior approximation.

  2. Variational continual learning:

a continual learning method, variational continual learning (VCL) (Nguyen et
al., 2018) visits each client in sequential order, performing VI with the
prior equal to the previous approximate posterior,

$q^{(i)}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})}\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\log
p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})\right]-\mathrm{KL}\left({q(\boldsymbol{\mathbf{\theta}})}\
\|\ {q^{(i-1)}(\boldsymbol{\mathbf{\theta}})}\right).$ (5)

This is equivalent to performing sequential PVI for a single round of client
updates. Whilst exact implementation of the online Bayesian update rule is
insensitive to the order in which clients are visited, this is not true for
VCL. In general, we expect the final approximate posterior learnt by VCL to be
biased towards performing well on the most recently visited dataset due to the
accumulation of approximation errors.

  3. Streaming variational Bayes:

analogous to the relationship between assumed density filtering (ADF) and EP,
streaming variational Bayes (VB) performs sequential PVI without replacement,
minimising the local free-energy (5) rather than (2).

In each experiment we either evaluate the performance of a logistic regression
model or Bayesian neural network, which we define below.

  1. Logistic regression model:

we use a logistic regression model defined by the likelihood

$p(y|\boldsymbol{\mathbf{\theta}},\mathbf{x})=\sigma\left(\tilde{\mathbf{x}}^{\top}\boldsymbol{\mathbf{\theta}}\right)^{y}\left(1-\sigma\left(\tilde{\mathbf{x}}^{\top}\boldsymbol{\mathbf{\theta}}\right)\right)^{1-y}$

where $\sigma(z)=\frac{1}{1+e^{-z}}$ is the sigmoid function and
$\tilde{\mathbf{x}}=\left[1,\ \mathbf{x}^{\top}\right]^{\top}$. In all
experiments we use a standard normal prior on the parameters and a mean-field
Gaussian variational approximation. The posterior predictive distribution is
approximated using the probit approximation (Spiegelhalter and Lauritzen,
1990):

$\displaystyle
p(y_{*}=1|\mathbf{x}_{*},\mathbf{y},\mathbf{X})\approx\sigma\left(\frac{\boldsymbol{\mathbf{\mu}}^{\top}_{q}\tilde{\mathbf{x}}_{*}}{\sqrt{1+\pi\tilde{\mathbf{x}}_{*}^{\top}\text{diag}(\boldsymbol{\mathbf{\sigma}}^{2}_{q})\tilde{\mathbf{x}}_{*}}}\right)$

where $\boldsymbol{\mathbf{\mu}}_{q}$ and $\boldsymbol{\mathbf{\sigma}}_{q}$
and the mean and standard deviation of $q(\boldsymbol{\mathbf{\theta}})$.

  2. Bayesian neural network:

we use a fully connected neural network with ReLU activation functions. In all
experiments we use a standard normal prior on the parameters and a mean-field
Gaussian variational approximation. The posterior predictive distribution is
approximated using the Monte Carlo estimate

$p(y_{*}=1|\mathbf{x}_{*},\mathbf{y},\mathbf{X})\approx\frac{1}{S}\sum_{s=1}^{S}p(y_{*}=1|\boldsymbol{\mathbf{\theta}}^{(s)},\mathbf{x}_{*})\quad\text{where}\
\boldsymbol{\mathbf{\theta}}^{(s)}\sim q(\boldsymbol{\mathbf{\theta}}).$

###  5.1 UCI Classification

The purpose of this experiment is to evaluate the performance of PVI on
relatively simple datasets under both homogeneous and inhomogeneous data
distributions. We consider three binary classification tasks available on the
UCI machine learning repository (Dua and Graff, 2017): adult, with the task
being to predict whether income exceeds $50,000/year based on census data;
bank, with the task being to predict if a client will subscribe to a term
deposit; and credit, with the task being the predict whether individuals
default on payments. All three tasks contain sensitive financial information,
thus are datasets for which the use of federated learning is suitable. For
each of the federated learning schemes, we perform approximate inference of
the parameters of a logistic regression model. We follow the dataset
distribution scheme of Sharma et al. (2019) to split each dataset into $M=10$
partitions—five small and five large partitions. For each task, we consider a
homogeneous split (A) and highly inhomogeneous split (B). A detailed
description of the datasets and corresponding partitions is provided in
Section E.1.

(a) Adult

(b) Bank

(c) Credit

Figure 4: A comparison between the predictive performance of the federated
learning schemes on the test datasets of each of the UCI classification tasks.
The mean $\pm$ standard deviation averaged over five random initialisations is
shown. See text for details.

Figures 4(a), LABEL:, 4(b), LABEL: and 4(c) compare the predictive performance
of the approximate posteriors learnt by each of the federated learning schemes
on each of the three UCI binary classification tasks. A damping factor of
$\rho=0.2$ is used for synchronous PVI. For both splits, we observe that both
sequential and synchronous PVI converge towards the predictive performance
obtained by the ‘optimal’ approximate posterior learnt by global VI—this is in
accordance with the theory presented in Section 2. For the homogeneous split
A, convergence is achieved after a single round of communications.131313This
corresponds to a single communication for synchronous PVI and $K$
communications for sequential PVI. Convergence takes noticeably longer for the
inhomogeneous split B, but is nonetheless achieved after only a few rounds of
communications. Importantly, the results demonstrate that PVI is able to
recover the predictive performance obtained using global VI with far greater
communication efficiency. Despite being relatively simple tasks, the plots
reveal the comparatively poor performance of both BCM methods and VCL relative
to PVI in the presence of inhomogeneity. Although the posterior distribution
is identifiable, the client-level approximate factors are interdependent and
so a single round of communications—as with BCM methods—is unable to identify
the globally optimal approximate posterior. Similarly, the poor performance of
VCL arises due to the interdependence of approximate factors and accumulation
of approximation errors, the latter resulting in an approximate posterior
biased towards performing well on the distribution of the dataset visited
last. This explains the relatively high variance in final test set performance
of VCL. The inadequacies of BCM and VCL demonstrate the importance of
revisiting clients to model interdependencies between approximate factors.
Yet, this is not sufficient—we also need the deletion step. Without the
deletion step of PVI, streaming VB ‘over-counts’ the observations resulting in
over-sharpening of the likelihood approximation and overfitting, as
illustrated by the predictive performance plot on the credit dataset.141414The
same degradation in predictive performance does not occur for the larger adult
and bank datasets. Because the datasets are so large and the prior
distribution relatively weak, the likelihood is much sharper than the prior
and the posterior distribution resembles a delta on the maximum likelihood
estimate. Since the maximum likelihood estimate is independent of the number
of data points, both streaming VB and PVI converge towards the same solution.

###  5.2 Mortality Prediction

In this experiment, our aim is to investigate the performance of PVI as the
complexity of both the dataset and model is increased. We consider the
mortality prediction task of the MIMIC-III dataset (Johnson et al., 2016;
Harutyunyan et al., 2019). The presence of sensitive medical data motivates
the use of federated learning. The dataset consists of recordings of 17
clinical variables recorded over the first 48 hours of each ICU stay, of which
there are 42,276, corresponding to 33,798 unique patients. Each of the 17
recordings is divided into seven subsequences—full time series; first 10% of
time; first 25% of time, first 50% of time, last 50% of time, last 25% of
time; and last 10% of time—from which six statistical features are extracted.
In total, each ICU stay is represented by a $17\times 7\times 6=714$
dimensional feature vector.

We construct $M=5$ inhomogeneous partitions of the data by performing k-means
clustering on the inputs, the purpose being to simulate the inhomogeneity of
patients that one might observe at different hospital—see Section E.2 for
details. We compare the performance of the federated learning methods in
performing approximate inference in a simple logistic regression model and
more complex two-layered BNN with 50 hidden units per layer. A damping factor
of $\rho=0.4$ is used for synchronous PVI. Figures 5(a), LABEL: and 5(b) plot
the predictive performance of the methods.

(a) Logistic regression model.

(b) Two-layered BNN with 100 hidden units.

Figure 5: A comparison between the predictive performance of the federated
learning schemes on the mortality prediction task. The mean $\pm$ standard
deviation averaged over five random initialisations is shown. See text for
details.

Both PVI methods converge towards the same predictive performance achieved by
global VI—doing so in far fewer communications—consolidating the observations
made in Section 5.1 and demonstrating the effectiveness of PVI in more complex
models. Convergence is noticeably smoother for the simpler logistic regression
model than for the BNN, a consequence of the approximate posterior of the BNN
being more complex.

###  5.3 Federated MNIST

In this experiment, we evaluate the performance of PVI on a complex multi-
class classification problem and investigate how performance is affected in
the presence of severe inhomogeneity. The MNIST multi-class classification
task (LeCun et al., 1998) consists of 60,000 $28\times 28$ greyscale images of
handwritten digits between 0 and 9. We consider performing federated learning
on two distinct partitions of the dataset: a homogeneous split, in which the
dataset is partitioned equally into $M=10$ groups of 6,000 datapoints each;
and the inhomogeneous split of McMahan et al. (2017), in which the dataset is
partitioned into $M=100$ groups of 600 datapoints, with each group consisting
of examples of just two digits. The inhomogeneous split is highly non-IID and
was introduced with the intention of testing federated learning algorithms to
their limits.

We investigate the performance of the federated learning methods in performing
approximate inference in a single-layered BNN with 200 hidden units. For
synchronous PVI a damping factor of $\rho=0.2$ is used for the homogeneous
split and $\rho=0.025$ is used for the inhomogeneous split. In both cases, we
found this to be the smallest damping factor necessary for stable convergence.
For sequential PVI, we used a damping factor of $\rho=0.3$ for the
inhomogeneous split which we found to result in faster convergence than the
use of no damping.

####  5.3.1 Homogeneous Partition

Figure 6 compares the performance of federated learning schemes on the
homogeneous split of the MNIST dataset.

Figure 6: Plots of the predictive performance of the federated learning
schemes on the homogeneous split of the MNIST classification task. The split
is constructed by randomly partitioning the 60,000 training images into $M=10$
partitions. The mean $\pm$ standard deviation computed using five random
initialisations is shown. See text for details.

As for the simpler binary classification tasks, the predictive performance of
both PVI methods converges towards the same asymptote as global VI with far
greater communication efficiency. Both BCM methods perform poorly, which again
is due to the non-identifiability of the approximate posterior. VCL achieves
reasonable predictive performance; yet, it is evident that by revisiting
clients and modifying their contribution to the approximate posterior—as in
sequential PVI—predictive performance can be improved. Without the deletion
step, streaming VB overfits the training dataset which results in a
degradation in test set performance.

A pathology of mean-field variational approximations to BNNs is the over-
pruning effect (Trippe and Turner, 2018). Figure 6 demonstrates this in full
effect, showing that as the number of pruned parameters increases the
predictive performance worsens.151515A parameter, $\theta_{i}$, is marked as
pruned if $\textrm{KL}(q(\theta_{i})\|p(\theta_{i}))<0.1$. To mitigate this
effect, we consider applying early stopping to the local free-energy
optimisations by checking for convergence of the expected log-likelihood term
$\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\log
p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})\right]$. Rather than initialise
at the context distribution, it is important to initialise the approximate
posteriors at a ‘tight’ distribution for all updates to ensure early stopping
can be applied.161616If initialised at the context distribution as before, the
local free-energies do not get the opportunity to stop early. Figure 7
compares the performance of the federated learning schemes with early stopping
based on the expected log-likelihood term on the homogeneous split.

Figure 7: Plots of the predictive performance of the federated learning
schemes on the homogeneous split of the MNIST classification task. The solid
lines indicate the use of early stopping, whereas the dashed lines indicate
the free-energies are optimised until convergence. The mean $\pm$ standard
deviation computed using five random initialisations is shown. See text for
details.

Observe that the use of early stopping is able to identify the best predictive
performance for global VI with good precision, indicating its effectiveness.
For the PVI methods and streaming VB, the approximate posterior overfits to
the training dataset. Unlike streaming VB without early stopping, however,
this does not correspond to a worsening in test set predictive performance:
all three methods achieve a final test set negative log-likelihood better than
that achieved without the use of early stopping, with streaming PVI and VB
outperforming global VI.

####  5.3.2 Inhomogeneous Partition

Figure 8 compares the performance of the federated learning schemes on the
inhomogeneous split of the MNIST dataset. Note that the performance metrics
are plotted on the log-scale, as the variation between the performance of each
method is significantly greater than in any of the previous experiments.

Figure 8: Plots of the predictive performance of the federated learning
schemes on the inhomogeneous split of the MNIST classification task. The split
is constructed by partitioning the 60,000 training images into $M=100$
partitions, each with at most two digits. The mean $\pm$ standard deviation
computed using five random initialisations is shown. See text for details.

The communication efficiency of all methods is considerably worse than in the
homogeneous case. For global VI, this is because 100 communications are
required to simulate a single iteration of the exact VI dynamics (one for each
client), rather than 10 communications in the homogeneous case. The initial
communication efficiency of both PVI methods is far better than that of global
VI; however, they achieve a worse final predictive performance with a test set
error of 3%. This is a consequence of 1) the inhomogeneity of client data
distributions, making it difficult for the approximate likelihoods to agree
upon aggregation, and 2) the large number of clients, which means that
approximate likelihoods are updated infrequently in the case of sequential PVI
and a crippling large degree of damping is used in the synchronous PVI case.
The interaction of these two factors hinders the effectiveness of both
methods. Nonetheless, we emphasise that the degree of inhomogeneity present in
this experiment is significant, and likely to be greater than that encountered
in many practical situations. We expect extreme inhomogeneity to demand
relatively small updates are made at each client before communication, which
makes it impossible to develop communication efficient algorithms.

##  6 Conclusion

This paper provided a general framework for performing variational inference
in the federated learning setting, partitioned variational inference (PVI).

We provided theoretical and empirical results showing that PVI is able to
recover the global VI solution whilst respecting the data privacy constraints
of federated learning, demonstrating its efficacy in the presence of varying
degrees and forms of inhomogeneity across data partitions. Departing from the
federated learning setting, we showed that the PVI framework flexibly subsumes
many existing VI methods as special cases and sheds light on the connections
between a wealth of techniques.

One of the key contributions of this work is the connection of deterministic
local message passing methods with global optimisation-based schemes. Each of
these different branches of approximate inference has been arguably developed
independently of each other, for example, existing probabilistic programming
toolkits tend to work primarily with one approach but not both (see e.g. Tran
et al., 2017; Minka et al., 2018). This paper suggests these methods are
inter-related and practitioners could benefit from a unified framework, i.e.
there are ways to expand the existing probabilistic programming packages to
gracefully handle both approaches. Additionally, the PVI framework could be
used to automatically choose a granularity level and an optimisation scheme
that potentially offer a better inference method for the task or the model at
hand. It is, however, unclear how flexible variational approximations such as
mixtures of exponential family distributions (see e.g. Sudderth et al., 2010)
or normalising flows (Rezende and Mohamed, 2015) can be efficiently and
tractably accommodated in the PVI framework. We leave these directions as
future work.

The experiments in Section 5 demonstrated that PVI is well-suited to learning
with decentralised data. Deployment of PVI in this setting is practical as its
implementation only requires a straightforward modification of existing global
VI implementations. We have also explored how this algorithm allows data
parallelism—each local client stores a complete copy of the model—and
communication efficient, uncertainty-aware updates between clients. A
potential future extension of the proposed approach is to use model
parallelism. That is, in addition to decentralising the data and computation
across the clients, the model itself is partitioned. As commonly done in many
deep learning training algorithms, model parallelism could be achieved by
assigning the parameters (and computation) of different layers of the network
to different devices. Alternatively, for spatio-temporal datasets modelled
using spatio-temporally distributed parameters, each client could support part
of the full model relevant for its local region of space. Another potential
research direction is PVI with coordinator-free, peer-to-peer only
communication between clients. This could be achieved by a client passing each
update to several randomly selected other clients, who then apply the changes,
rather than to a central parameter server.

Whilst our implementation of PVI performed either sequential or synchronous
local optimisation until convergence, the relationship between PVI and global
VI—namely, that they share the same fixed-points—holds regardless of how
optimisation is performed. In Section 3 we demonstrated how the exact dynamics
of global VI can be recovered by terminating the local optimisation procedure
of synchronous PVI after a single update step is performed. The space in-
between performing a single update and running to convergence is largely
unexplored, and could result in algorithms that improve on the communication
efficiency of our implementation of PVI. Other directions we are yet to
explore include stochastic implementations of synchronous PVI, the use of
alpha-divergences to define the local free-energy, and a dynamic scheduling of
clients.

## Acknowledgements

We would like to thank Jonathan So for his insights on the properties of the
double-loop algorithm for EP. Matthew Ashman is supported by the George and
Lilian Schiff Foundation. Adrian Weller acknowledges support from a Turing AI
Fellowship under grant EP/V025379/1, The Alan Turing Institute, and the
Leverhulme Trust via CFI. Siddharth Swaroop is supported by an EPSRC DTP
studentship and Microsoft Research EMEA PhD Award. Richard E. Turner is
supported by Google, Amazon, ARM, Improbable and EPSRC grant EP/T005386/1.

## References

  * Archambeau and Ermis (2015) Cedric Archambeau and Beyza Ermis.  Incremental variational inference for latent Dirichlet allocation.  _arXiv preprint arXiv:1507.05016_ , 2015.
  * Beal (2003) Matthew James Beal.  _Variational algorithms for approximate Bayesian inference_.  PhD thesis, UCL, 2003.
  * Broderick et al. (2013) Tamara Broderick, Nicholas Boyd, Andre Wibisono, Ashia C. Wilson, and Michael I. Jordan.  Streaming variational Bayes.  In _Advances in Neural Information Processing Systems_ , 2013.
  * Bui et al. (2017a) Thang D. Bui, Cuong V. Nguyen, and Richard E. Turner.  Streaming sparse Gaussian process approximations.  In _Advances in Neural Information Processing Systems_ , 2017a.
  * Bui et al. (2017b) Thang D. Bui, Josiah Yan, and Richard E. Turner.  A unifying framework for Gaussian process pseudo-point approximations using power expectation propagation.  _Journal of Machine Learning Research_ , 18(104):1–72, 2017b.
  * Dua and Graff (2017) Dheeru Dua and Casey Graff.  UCI machine learning repository, 2017.  URL http://archive.ics.uci.edu/ml.
  * Dwork et al. (2006) Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith.  Calibrating noise to sensitivity in private data analysis.  In _Theory of cryptography conference_ , pages 265–284. Springer, 2006.
  * Elidan et al. (2012) Gal Elidan, Ian McGraw, and Daphne Koller.  Residual belief propagation: Informed scheduling for asynchronous message passing.  _arXiv preprint arXiv:1206.6837_ , 2012.
  * Gelman et al. (2014) Andrew Gelman, Aki Vehtari, Pasi Jylänki, Christian Robert, Nicolas Chopin, and John P Cunningham.  Expectation propagation as a way of life.  _arXiv preprint arXiv:1412.4869_ , 2014.
  * Ghahramani (2015) Zoubin Ghahramani.  Probabilistic machine learning and artificial intelligence.  _Nature_ , 521(7553):452–459, 2015.
  * Ghahramani and Attias (2000) Zoubin Ghahramani and H. Attias.  Online variational Bayesian learning.  In _NIPS Workshop on Online Learning_ , 2000.
  * Harutyunyan et al. (2019) Hrayr Harutyunyan, Hrant Khachatrian, David C Kale, Greg Ver Steeg, and Aram Galstyan.  Multitask learning and benchmarking with clinical time series data.  _Scientific data_ , 6(1):1–18, 2019.
  * Hasenclever et al. (2017) Leonard Hasenclever, Stefan Webb, Thibaut Lienart, Sebastian Vollmer, Balaji Lakshminarayanan, Charles Blundell, and Yee Whye Teh.  Distributed Bayesian learning with stochastic natural gradient expectation propagation and the posterior server.  _Journal of Machine Learning Research_ , 18(106):1–37, 2017.  URL http://jmlr.org/papers/v18/16-478.html.
  * Hensman et al. (2012) James Hensman, Magnus Rattray, and Neil D. Lawrence.  Fast variational inference in the conjugate exponential family.  In _Advances in Neural Information Processing Systems_ , pages 2888–2896, 2012.
  * Hensman et al. (2013) James Hensman, Nicolo Fusi, and Neil D Lawrence.  Gaussian processes for big data.  In _Uncertainty in Artificial Intelligence_ , page 282, 2013.
  * Heskes and Zoeter (2002) Tom Heskes and Onno Zoeter.  Expectation propagation for approximate inference in dynamic Bayesian networks.  In _Proceedings of the Eighteenth Conference on Uncertainty in Artificial Intelligence_ , UAI’02, page 216–223, 2002.  ISBN 1558608974.
  * Hinton and Van Camp (1993) Geoffrey E Hinton and Drew Van Camp.  Keeping the neural networks simple by minimizing the description length of the weights.  In _Conference on Computational Learning Theory_ , pages 5–13, 1993\.
  * Hoffman et al. (2013) Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley.  Stochastic variational inference.  _Journal of Machine Learning Research_ , 14(1):1303–1347, 2013.
  * Honkela et al. (2010) Antti Honkela, Tapani Raiko, Mikael Kuusela, Matti Tornio, and Juha Karhunen.  Approximate Riemannian conjugate gradient learning for fixed-form variational Bayes.  _Journal of Machine Learning Research_ , 11(Nov):3235–3268, 2010.
  * Jälkö et al. (2017) Joonas Jälkö, Onur Dikmen, and Antti Honkela.  Differentially private variational inference for non-conjugate models.  In _Conference on Uncertainty in Artificial Intelligence_ , 2017.
  * Johnson et al. (2016) Alistair EW Johnson, Tom J Pollard, Lu Shen, H Lehman Li-Wei, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark.  Mimic-iii, a freely accessible critical care database.  _Scientific data_ , 3(1):1–9, 2016.
  * Khan et al. (2015) Mohammad E Khan, Pierre Baqué, François Fleuret, and Pascal Fua.  Kullback-Leibler proximal variational inference.  In _Advances in Neural Information Processing Systems_ , pages 3402–3410, 2015.
  * Khan and Lin (2017) Mohammad Emtiyaz Khan and Wu Lin.  Conjugate-computation variational inference : Converting variational inference in non-conjugate models to inferences in conjugate models.  In _International Conference on Artificial Intelligence and Statistics_ , 2017.
  * Khan and Nielsen (2018) Mohammad Emtiyaz Khan and Didrik Nielsen.  Fast yet simple natural-gradient descent for variational inference in complex models.  In _2018 International Symposium on Information Theory and Its Applications (ISITA)_ , pages 31–35. IEEE, 2018.
  * Khan and Rue (2021) Mohammad Emtiyaz Khan and Håvard Rue.  The Bayesian learning rule, 2021.
  * Khan et al. (2016) Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama.  Faster stochastic variational inference using proximal-gradient methods with general divergence functions.  In _Conference on Uncertainty in Artificial Intelligence_ , pages 319–328, 2016.
  * Kingma and Ba (2014a) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2014a.
  * Kingma and Ba (2014b) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014b.
  * Knowles and Minka (2011) David A. Knowles and Tom Minka.  Non-conjugate variational message passing for multinomial and binary regression.  In _Advances in Neural Information Processing Systems_ , pages 1701–1709, 2011.
  * Kucukelbir et al. (2017) Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei.  Automatic differentiation variational inference.  _Journal of Machine Learning Research_ , 18(1):430–474, 2017.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Li et al. (2015) Yingzhen Li, José Miguel Hernández-Lobato, and Richard E. Turner.  Stochastic expectation propagation.  In _Advances in Neural Information Processing Systems_ , pages 2323–2331, 2015.
  * Loo et al. (2021) Noel Loo, Siddharth Swaroop, and Richard E Turner.  Generalized variational continual learning.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=_IM-AfFhna9.
  * McMahan et al. (2017) Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pages 1273–1282. PMLR, 2017.
  * Minka et al. (2018) T. Minka, J.M. Winn, J.P. Guiver, Y. Zaykov, D. Fabian, and J. Bronskill.  /Infer.NET 0.3, 2018.  Microsoft Research Cambridge. http://dotnet.github.io/infer.
  * Minka (2001) Thomas P Minka.  Expectation propagation for approximate Bayesian inference.  In _Conference on Uncertainty in Artificial Intelligence_ , pages 362–369, 2001.
  * Minka (2004) Thomas P. Minka.  Power EP.  Technical report, Microsoft Research Cambridge, 2004.
  * Minka (2005) Thomas P. Minka.  Divergence measures and message passing.  Technical report, Microsoft Research Cambridge, 2005.
  * Nguyen et al. (2018) Cuong V. Nguyen, Yingzhen Li, Thang D. Bui, and Richard E. Turner.  Variational continual learning.  In _International Conference on Learning Representations_ , 2018.
  * Parikh and Boyd (2014) Neal Parikh and Stephen Boyd.  Proximal algorithms.  _Foundations and Trends® in Optimization_ , 1(3):127–239, 2014.
  * Pearl (1988) Judea Pearl.  _Probabilistic reasoning in intelligent systems: networks of plausible inference_.  Morgan kaufmann, 1988.
  * Raskutti and Mukherjee (2015) Garvesh Raskutti and Sayan Mukherjee.  The information geometry of mirror descent.  _IEEE Transactions on Information Theory_ , 61(3):1451–1457, March 2015.
  * Rezende and Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , pages 1530–1538, 2015.
  * Salimans and Knowles (2013) Tim Salimans and David A Knowles.  Fixed-form variational posterior approximation through stochastic linear regression.  _Bayesian Analysis_ , 8(4):837–882, 2013.
  * Salimbeni and Hensman (2018) Hugh Salimbeni and James Hensman.  Natural gradients in practice: Non-conjugate variational inference in Gaussian process models.  In _International Conference on Artificial Intelligence and Statistics_ , 2018.
  * Sato (2001) Masa-Aki Sato.  Online model selection based on the variational Bayes.  _Neural Computation_ , 2001.
  * Sharma et al. (2019) Mrinank Sharma, Michael Hutchinson, Siddharth Swaroop, Antti Honkela, and Richard E Turner.  Differentially private federated variational inference.  _arXiv preprint arXiv:1911.10563_ , 2019.
  * Sheth and Khardon (2016a) Rishit Sheth and Roni Khardon.  A fixed-point operator for inference in variational Bayesian latent Gaussian models.  In Arthur Gretton and Christian C. Robert, editors, _International Conference on Artificial Intelligence and Statistics_ , volume 51 of _Proceedings of Machine Learning Research_ , pages 761–769, Cadiz, Spain, 09–11 May 2016a. PMLR.
  * Sheth and Khardon (2016b) Rishit Sheth and Roni Khardon.  Monte carlo structured SVI for non-conjugate models.  _arXiv preprint arXiv:1612.03957_ , 2016b.
  * Sheth et al. (2015) Rishit Sheth, Yuyang Wang, and Roni Khardon.  Sparse variational inference for generalized GP models.  In Francis Bach and David Blei, editors, _International Conference on Machine Learning_ , volume 37, pages 1302–1311, 2015.
  * Spiegelhalter and Lauritzen (1990) David J Spiegelhalter and Steffen L Lauritzen.  Sequential updating of conditional probabilities on directed graphical structures.  _Networks_ , 20(5):579–605, 1990.
  * Sudderth et al. (2010) Erik B Sudderth, Alexander T Ihler, Michael Isard, William T Freeman, and Alan S Willsky.  Nonparametric belief propagation.  _Communications of the ACM_ , 53(10):95–103, 2010\.
  * Swaroop et al. (2019) Siddharth Swaroop, Cuong V. Nguyen, Thang D. Bui, and Richard E. Turner.  Improving and understanding variational continual learning.  _arXiv preprint arXiv:1905.02099_ , 2019.
  * Theis and Hoffman (2015) Lucas Theis and Matthew D Hoffman.  A trust-region method for stochastic variational inference with applications to streaming data.  In _International Conference on Machine Learning_ , 2015.
  * Tieleman and Hinton (2012) Tijmen Tieleman and Geoffrey E Hinton.  Lecture 6.5—RmsProp: Divide the gradient by a running average of its recent magnitude.  COURSERA: Neural Networks for Machine Learning, 2012.
  * Tran et al. (2017) Dustin Tran, Matthew D. Hoffman, Rif A. Saurous, Eugene Brevdo, Kevin Murphy, and David M. Blei.  Deep probabilistic programming.  In _International Conference on Learning Representations_ , 2017.
  * Tresp (2000) Volker Tresp.  A Bayesian committee machine.  _Neural computation_ , 12(11):2719–2741, 2000\.
  * Trippe and Turner (2018) Brian Trippe and Richard Turner.  Overpruning in variational Bayesian neural networks.  _arXiv preprint arXiv:1801.06230_ , 2018.
  * Wainwright et al. (2008) Martin J Wainwright, Michael I Jordan, et al.  Graphical models, exponential families, and variational inference.  _Foundations and Trends® in Machine Learning_ , 1(1–2):1–305, 2008.
  * Wand (2014) Matt P. Wand.  Fully simplified multivariate Normal updates in non-conjugate variational message passing.  _Journal of Machine Learning Research_ , 15:1351–1369, 2014\.
  * Winn and Minka (2009) John Winn and Tom Minka.  Probabilistic programming with Infer.NET.  Machine Learning Summer School, 2009.
  * Winn et al. (2005) John Winn, Christopher M. Bishop, and Tommi Jaakkola.  Variational message passing.  _Journal of Machine Learning Research_ , 6:661–694, 2005\.
  * Xu et al. (2014) Minjie Xu, Balaji Lakshminarayanan, Yee Whye Teh, Jun Zhu, and Bo Zhang.  Distributed Bayesian posterior sampling via moment sharing.  _Advances in Neural Information Processing Systems_ , 27, 2014.
  * Zeno et al. (2018) Chen Zeno, Itay Golan, Elad Hoffer, and Daniel Soudry.  Bayesian gradient descent: Online variational Bayes learning with increased robustness to catastrophic forgetting and weight pruning, 2018.
  * Zhang et al. (2018) Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt.  Advances in variational inference.  _IEEE transactions on pattern analysis and machine intelligence_ , 41(8):2008–2026, 2018.

##  A Hyperparameter Learning

Many probabilistic models depend on a set of hyperparameters
$\boldsymbol{\mathbf{\epsilon}}$ and it is often necessary to learn suitable
settings from data to achieve good performance on a task. One method is to
optimise the variational free-energy thereby approximating maximum likelihood
learning. The gradient of the global variational free-energy decomposes into a
set of local computations, as shown in Section C,

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\epsilon}}}\mathcal{F}(\boldsymbol{\mathbf{\epsilon}},q(\boldsymbol{\mathbf{\theta}}))=$
$\displaystyle\sum_{m=1}^{M}\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\epsilon}}}\log
p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})\right]+\mathbb{E}_{q(\boldsymbol{\mathbf{\theta}})}\left[\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\epsilon}}}\log
p(\boldsymbol{\mathbf{\theta}}|\boldsymbol{\mathbf{\epsilon}})\right].$

This expression holds for general $q(\boldsymbol{\mathbf{\theta}})$ and is
valid both for coordinate ascent (iterating between updating
$\boldsymbol{\mathbf{\epsilon}}$ with $q(\boldsymbol{\mathbf{\theta}})$ fixed
and updating $q(\boldsymbol{\mathbf{\theta}})$ with
$\boldsymbol{\mathbf{\epsilon}}$ fixed) and for optimising the collapsed bound
(where the approximate posterior optimises the global free-energy
$q(\boldsymbol{\mathbf{\theta}})=q^{*}(\boldsymbol{\mathbf{\theta}})$ and
therefore depends implicitly on $\boldsymbol{\mathbf{\epsilon}}$). The latter
does not require us to compute $q(\boldsymbol{\mathbf{\theta}})$ using PVI,
but rather query the gradients of the log-likelihood for each data partition
with respect to the hyperparameters. Notice that this expression is amenable
to stochastic approximation which leads to optimization schemes that use only
local information at each step. When combined with different choices for the
optimization of the local free-energies w.r.t.
$q(\boldsymbol{\mathbf{\theta}})$, this leads to a wealth of possible
hyperparameter optimization schemes.

In cases where a distributional estimate for the hyperparameters is necessary,
e.g. in continual learning, the PVI framework above can be extended to handle
the hyperparameters. In particular, the approximate posterior in (1) can be
modified as follows,

$\displaystyle q(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})$
$\displaystyle=\frac{1}{\mathcal{Z}_{q}}p(\boldsymbol{\mathbf{\epsilon}})p(\boldsymbol{\mathbf{\theta}}|\boldsymbol{\mathbf{\epsilon}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})\approx\frac{1}{\mathcal{Z}}p(\boldsymbol{\mathbf{\epsilon}})p(\boldsymbol{\mathbf{\theta}}|\boldsymbol{\mathbf{\epsilon}})\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})=p(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}}|\mathbf{y}),$

where the approximate likelihood factor
$t_{m}(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})$ now
involves both the model parameters and the hyperparameters. Similar to (2),
the approximate posterior above leads to the following local variational free-
energy,

$\displaystyle\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}}))=\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\mathrm{d}\boldsymbol{\mathbf{\epsilon}}q(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})\log\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})}{q(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}},\boldsymbol{\mathbf{\epsilon}})}.$

Note that this approach retains all favourable properties of PVI such as local
computation and flexibility in choosing optimization strategies and stochastic
approximations.

##  B Proofs

###  B.1 Relating Local KL Minimisation to Local Free-Energy Minimisation:
Proof of 2.1

Substituting the tilted distribution
$\widehat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})$ into the KL divergence
yields,

$\displaystyle\mathrm{KL}\left(q(\boldsymbol{\mathbf{\theta}})\|\widehat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\right)$
$\displaystyle=\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}})\log\frac{q(\boldsymbol{\mathbf{\theta}})\widehat{\mathcal{Z}}^{(i)}_{k}t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}$
$\displaystyle=\log\widehat{\mathcal{Z}}^{(i)}_{k}-\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}})\log\frac{q^{(i-1)}(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{k}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})t^{(i-1)}_{k}(\boldsymbol{\mathbf{\theta}})}.$

Hence
$\mathcal{F}^{(i)}_{k}(q(\boldsymbol{\mathbf{\theta}}))=\log\widehat{\mathcal{Z}}^{(i)}_{k}-\mathrm{KL}\left(q(\boldsymbol{\mathbf{\theta}})\|\widehat{p}^{(i)}_{k}(\boldsymbol{\mathbf{\theta}})\right)$.

###  B.2 The Sum of Local Free-Energies is Equal to the Global Free-Energy

Let
$q(\boldsymbol{\mathbf{\theta}})=\frac{1}{\mathcal{Z}_{q}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})$.
We have

$\displaystyle\sum_{m=1}^{M}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))+\log\mathcal{Z}_{q}$
$\displaystyle=\sum_{m=1}^{M}\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{q(\boldsymbol{\mathbf{\theta}})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})t_{m}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}+\log\mathcal{Z}_{q}$
$\displaystyle=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}+\log\mathcal{Z}_{q}$
$\displaystyle=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{\frac{1}{\mathcal{Z}_{q}}p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}t_{m}(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
$\displaystyle=\int
q(\boldsymbol{\mathbf{\theta}})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m=1}^{M}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}})}\mathrm{d}\boldsymbol{\mathbf{\theta}}$
$\displaystyle=\mathcal{F}(q(\boldsymbol{\mathbf{\theta}})).$

###  B.3 Maximisation of Local Free-Energies Implies Maximisation of Global
Free-Energy: Proof of 2.3

Let $\boldsymbol{\mathbf{\eta}}_{q}$ and $\boldsymbol{\mathbf{\eta}}^{*}_{q}$
be the variational parameters of $q(\boldsymbol{\mathbf{\theta}})$ and
$q^{*}(\boldsymbol{\mathbf{\theta}})$ respectively. We can write
$q(\boldsymbol{\mathbf{\theta}})=q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})$
and
$q^{*}(\boldsymbol{\mathbf{\theta}})=q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})=\frac{1}{\mathcal{Z}_{q^{*}}}p(\boldsymbol{\mathbf{\theta}})\prod_{m}t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})$.
First, we show that at convergence, the derivative of the global free-energies
equals the sum of the derivatives of the local free-energies. The derivative
of the local free-energy $\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))$
w.r.t. $\boldsymbol{\mathbf{\eta}}_{q}$ is:

$\displaystyle\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle\quad-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}}.$

Thus, at convergence when
$\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}$,

$\displaystyle\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}.$

Summing both sides over all $m$,

$\displaystyle\sum_{m}\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{\prod_{m}t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle\
-\cancelto{0}{\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\mathcal{Z}_{q^{*}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}_{q}^{*}}}.$

Now consider the derivative of the global free-energy
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$:

$\displaystyle\frac{\mathrm{d}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle=\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}}.$

Hence,

$\displaystyle\frac{\mathrm{d}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=\sum_{m}\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}.$

For all $m$, since
$q^{*}(\boldsymbol{\mathbf{\theta}})=\operatorname*{arg\,max}_{q(\boldsymbol{\mathbf{\theta}})\in\mathcal{Q}}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))$,
we have
$\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\big{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=0$,
which implies:

$\displaystyle\frac{\mathrm{d}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=\sum_{m}\frac{\mathrm{d}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=0.$

Thus, $q^{*}(\boldsymbol{\mathbf{\theta}})$ is an extremum of
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$.

Now we show that $q^{*}(\boldsymbol{\mathbf{\theta}})$ is a maximum of
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$ by considering the Hessian
$\frac{\mathrm{d}^{2}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}$
of the global free-energy at convergence. Similar to the derivative case, we
now show that the Hessian of the global free-energies equals the sum of the
Hessians of the local free-energies. The Hessian of the local free-energy
$\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))$ w.r.t.
$\boldsymbol{\mathbf{\eta}}_{q}$ is:

$\displaystyle\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\biggl{(}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle\hskip
225.0pt-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}}\biggr{)}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d^{2}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle\
-\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}\log
q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}.$

The final term can be written as

$\displaystyle\left(\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}\log
q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\right)^{\intercal}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\frac{\mathrm{d}\log
q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}^{2}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}=\boldsymbol{\mathbf{0}}.$

Thus

$\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d^{2}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}.$

At convergence when
$\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}$,

$\displaystyle\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$

Summing both sides over all $m$,

$\displaystyle\sum_{m}\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{\prod_{m}t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle\
-\cancelto{0}{\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\mathcal{Z}_{q^{*}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}_{q}^{*}}}.$
$\displaystyle\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}+\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}^{2}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad-\cancelto{0}{\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}}.$

At convergence when
$\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}$,

$\displaystyle\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}.$

Summing both sides over all $m$,

$\displaystyle\sum_{m}\frac{\mathrm{d}^{2}\mathcal{F}_{m}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{\prod_{m}t_{m}(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}^{*}_{q})}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}^{*}_{q}}$
$\displaystyle\
-\cancelto{0}{\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\mathcal{Z}_{q^{*}}\biggr{|}_{\boldsymbol{\mathbf{\eta}}_{q}=\boldsymbol{\mathbf{\eta}}_{q}^{*}}}.$

Now consider the Hessian of the global free-energy
$\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))$:

$\displaystyle\frac{\mathrm{d}^{2}\mathcal{F}(q(\boldsymbol{\mathbf{\theta}}))}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}$
$\displaystyle=\frac{\mathrm{d}^{2}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}$
$\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}}\left(\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}q(\boldsymbol{\mathbf{\theta}})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})\frac{\mathrm{d}\log
q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}}\right)$
$\displaystyle=\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}^{2}q(\boldsymbol{\mathbf{\theta}})}{\mathrm{d}\boldsymbol{\mathbf{\eta}}_{q}\mathrm{d}\boldsymbol{\mathbf{\eta}}^{\intercal}_{q}}\log\frac{p(\boldsymbol{\mathbf{\theta}})\prod_{m}p(\mathbf{y}_{m}|\boldsymbol{\mathbf{\theta}})}{q(\boldsymbol{\mathbf{\theta}};\boldsymbol{\mathbf{\eta}}_{q})}-\cancelto{0}{\int\mathrm{d}\boldsymbol{\mathbf{\theta}}\frac{\mathrm{d}\log