# Bayesian Learning for Neural Networks:
an algorithmic survey

Martin Magris and Alexandros Iosifidis
Department of Electrical and Computer Engineering, Aarhus University, Denmark
<EMAIL_ADDRESS>

###### Abstract

The last decade witnessed a growing interest in Bayesian learning. Yet, the
technicality of the topic and the multitude of ingredients involved therein,
besides the complexity of turning theory into practical implementations, limit
the use of the Bayesian learning paradigm, preventing its widespread adoption
across different fields and applications. This self-contained survey engages
and introduces readers to the principles and algorithms of Bayesian Learning
for Neural Networks. It provides an introduction to the topic from an
accessible, practical-algorithmic perspective. Upon providing a general
introduction to Bayesian Neural Networks, we discuss and present both standard
and recent approaches for Bayesian inference, with an emphasis on solutions
relying on Variational Inference and the use of Natural gradients. We also
discuss the use of manifold optimization as a state-of-the-art approach to
Bayesian learning. We examine the characteristic properties of all the
discussed methods, and provide pseudo-codes for their implementation, paying
attention to practical aspects, such as the computation of the gradients.

##  1 Introduction

Machine Learning (ML)techniques have been proven to be successful in many
prediction and classification tasks across natural language processing (Young
et al., 2018), computer vision (Krizhevsky et al., 2012), time-series
(Längkvist et al., 2014) and finance applications (Dixon et al., 2020), among
the several others. The widespread of ML methods in diverse domains is found
due to their ability to scale and adapt to data, and their flexibility in
addressing a variety of problems while retaining high predictive ability.
Recently, Bayesian methods have gained considerable interest in ML as an
attractive alternative to the classical methods providing point estimations
for their inputs. Despite the numerous advantages that traditional ML methods
offer, they are, broadly speaking, prone to overfitting dimming their
generalization capabilities and performance on unseen data. Furthermore, an
implicit consequence of the classical point estimation and modeling setup is
that it delivers models that are generally incapable of addressing
uncertainties. This inability is twofold as it includes both the estimation
and prediction aspects. Indeed, as opposed to the typical practice of
statistical modeling and, e.g., econometrics methods, ML methods do not
directly tackle aspects related to the significance and uncertainties
associated with the estimated parameters. At the same time, predictions
correspond to simple point estimates without reference to the confidence
levels that such estimates have. Whereas some models have been developed to,
e.g., provide confidence intervals over the forecasts (e.g. Gal and
Ghahramani, 2016), it has been observed that such models are generally
overconfident. To estimate uncertainties implicitly embedded in ML models,
Bayesian inference provides an immediate remedy and stands out as the main
approach.

Bayesian methods have gained considerable interest as an attractive
alternative to point estimation, especially for their ability to address
uncertainty via posterior distribution, generalize while reducing overfitting
(Hoeting et al., 1999), and for enabling sequential learning (Freitas et al.,
2000) while retaining prior and past knowledge. Although Bayesian principles
have been proposed in ML decades ago (e.g. Mackay, 1992, 1995; Lampinen and
Vehtari, 2001), it has been only recently that fast and feasible methods
boosted a growing use of Bayesian methods in complex models, such as deep
neural networks (Osawa et al., 2019; Khan et al., 2018a; Khan and Nielsen,
2018).

The most challenging task in following the Bayesian paradigm is the
computation of the posterior. In the typical ML setting characterized by a
high number of parameters and a considerable size of data, traditional
sampling methods are prohibitive, and alternative estimation approaches such
as Variational Inference (VI)have been shown to be suitable and successful
(Saul et al., 1996; Wainwright and Jordan, 2008; Hoffman et al., 2013; Blei et
al., 2017). Furthermore, recent research advocates the use of natural
gradients for boosting the optimum search and the training (Wierstra et al.,
2014), enabling fast and accurate Bayesian learning algorithms that are
scalable and versatile.

Recent years witnessed enormous growth in the interest related to Bayesian ML
methodologies and several contributions in the field. This survey aims at
summarizing the major methodologies nowadays available, presenting them from
an algorithmic, empirically-oriented, perspective. With this rationale, this
paper aims at providing the reader with the basic tools and concepts to
understand the theory behind Bayesian Deep Learning (DL)and walk through the
implementation of the several Bayesian estimation methodologies available. We
should note that the focus of this paper is on purely Bayesian methods. Indeed
there are a number of network architectures that can resemble a Bayesian
framework by e.g. creating a distribution for the outputs, e.g. Deep Ensembles
(Osband et al., 2018), Batch Ensembles (Wen et al., 2020), Layer Ensembles
(Oleksiienko and Iosifidis, 2022), or Variational Neural Networks (Oleksiienko
et al., 2022). These solutions, based on particular network designs, are
however not implicitly Bayesian and out of scope in our context. Other surveys
and tutorials do exist on the general topic (e.g. Jospin et al., 2022;
Heckerman, 2008, along with several lecture notes available online), yet the
focus of this paper is on algorithms and mainly devoted to Variational
Inference (VI)methods. In fact, despite the wide number of VI and non-VI
methods published in the last decade, a comprehensive survey embracing and
discussing all of them (or perhaps the major ones) is missing, and non-experts
will easily find themselves lost in their pursuit to comprehend and different
notions and processing steps in different methodologies. By filling this gap
we aim at promoting applications and research in this area.

###  1.1 The Bayesian paradigm

The Bayesian paradigm in statistics is often opposed to the pure frequentist
paradigm, a major area of distinction being in hypothesis testing (Etz et al.,
2018). The Bayesian paradigm is based on two simple ideas. The first is that
probability is a measure of belief in the occurrence of events, rather than
just some limit in the frequency of occurrence when the number of samples goes
toward infinity. The second is that prior beliefs influence posterior beliefs
(Jospin et al., 2022). The above two are summarized in the Bayes theorem which
we now review. Let $\mathcal{D}$ denote the data and
$p\left(\mathcal{D}|{\bm{\theta}}\right)$ the likelihood of the data based on
a postulated model with ${\bm{\theta}}\in\Theta$ a $k$-dimensional vector of
model parameters. Let $p\left(\bm{\theta}\right)$ be the prior distribution on
${\bm{\theta}}$. The goal of Bayesian inference is the estimation of the
posterior distribution

$p\left({\bm{\theta}}|\mathcal{D}\right)=\frac{p\left(\mathcal{D},{\bm{\theta}}\right)}{p\left(\mathcal{D}\right)}=\frac{p\left(\bm{\theta}\right)p\left(\mathcal{D}|{\bm{\theta}}\right)}{p\left(\mathcal{D}\right)},$
(1)

where $p\left(\mathcal{D}\right)$ is referred to as evidence or marginal
likelihood, since
$p\left(\mathcal{D}\right)=\int_{\Theta}p\left(\bm{\theta}\right)p\left(\mathcal{D}|{\bm{\theta}}\right)d{\bm{\theta}}$.
$p\left(\mathcal{D}\right)$ acts as a normalization constant for retrieving an
actual probability distribution for $p\left({\bm{\theta}}|\mathcal{D}\right)$.
In this light, as opposed to the frequentist approach, it becomes clear that
the unknown parameter ${\bm{\theta}}$ is treated as a random variable. The
prior probability $p\left(\bm{\theta}\right)$, which intuitively expresses in
probabilistic terms any knowledge about the parameter before the data has been
collected, is updated in the posterior probability
$p\left({\bm{\theta}}|\mathcal{D}\right)$, mixturing prior knowledge and
evidence supported by the data through the model’s likelihood. Bayesian
inference is generally difficult due to the fact that the marginal likelihood
is often intractable and of unknown form. Indeed, only for a limited class of
models, where the prior is so-called conjugate to the likelihood, the
calculation of the posterior is analytically tractable. Standard examples are
Normal likelihoods and prior (resulting in Normal posteriors) or Poisson
likelihoods with Gamma priors (resulting in Negative Binomial posteriors). Yet
already for the simple linear regression example, Bayesian derivation is
rather tedious, and already for the logistic regression closed-form solutions
are not generally available. It is clear that in complex models, such as deep
neural networks typically used in ML applications, Bayesian inference can be
tackled neither analytically nor numerically (consider that the integral in
the marginal likelihood is multivariate, over as many dimensions as the number
of parameters).

Monte-Carlo (MC)methods for sampling the posterior are certainly a possibility
that has been early explored and adopted. While it still remains a valid and
appropriate method for performing Bayesian inference in retractable settings,
especially in high-dimensional applications, the MC approach is challenging
and may become infeasible, mainly because of the need for an implicit high-
dimensional sampling scheme, which is generally time-consuming and
computationally demanding. As an alternative approach, VI gained much
attention in recent years. VI turns the integration Bayesian problem in eq.
(1) into an optimization problem. The idea behind VI is that of targeting an
approximate form of the posterior distribution, perhaps chosen within a family
of well-behaved distributions, and finding the corresponding parameter that
optimizes a specific objective, i.e. that is optimal under some criterion.

In the following subsection, we review the standard non-Bayesian approach for
neural network parameter estimation (Section 1.2.1), we introduce Bayesian
Neural Networks (BNNs) (Section 1.4), and we provide some motivation for their
use, also recalling some literature about their applications (Section 1.3).
After providing the reader an introduction to standard and Bayesian neural
networks, we introduce VI in Section 1.5, we describe the standard framework
used in Bayesian learning, and we discuss how the standard Stochastic Gradient
Descent (SGD)approach can be used for solving the optimization problem therein
(Section 1.5.1).

###  1.2 Standard and Bayesian Neural Networks

A Bayesian Neural Network (BNN)is an Artificial Neural Network (ANN)trained
with Bayesian Inference (Jospin et al., 2022). In the following, we provide a
quick overview of ANN s and their typical estimation based on Backpropagation
(Section 1.2.1). We then describe what a Bayesian Neural Network (BNN)is
(Section 1.4), provide motivations on why to use a BNN, over a standard ANN
(Section 1.3), and lastly introduce Variational Inference (VI)(Section 1.5).

####  1.2.1 Artificial Neural Networks and classical Backpropagation

Neurons are elementary building blocks which can be thought of as processing
units that, when combined, constitute a neural network. Each neuron processes
the information presented to its input by applying a transformation to it.
When affine neurons are used, the transformation corresponds to computing the
weighted sum of the inputs to the neuron (received from the neurons that are
connected to it or corresponding to the inputs to the neural network) and
generates a value, which is further introduced to a (usually nonlinear)
activation function to produce the neuron’s output (input to other neurons or
the neural network output). In order to account for the need of a shift to the
value needed to produce an activation response, a bias is also added as an
input to the activation function, which is commonly included in the weighted
sum by augmenting the input to the neuron with an additional input with a
constant value of 1, which associated with the corresponding bias term. While
activation functions squeezing their outputs to a pre-determined range of
values, like the sigmoid (with outputs in $[0,1]$) or the tanh (with outputs
in $[-1,1]$) functions, have been widely used in the past, piece-wise linear
functions, like the Rectified Linear Unit (ReLU) or the parametric ReLU
functions (He et al., 2015), are nowadays widely adopted in building neural
networks. A common characteristic of activation functions used in neural
networks is that they are differentiable with respect to their parameters over
the range of their inputs. The transformation performed by an affine neuron is
illustrated in Figure 1.

Figure 1: Representation of the operations within the $j$-th neuron at layer
is $l$. Connections between this neuron and neurons in layer $l-1$ are
represented by lines corresponding to weights $\theta^{l}_{\cdot j}$. The
inputs to the neuron $o^{l-1}_{\cdot}$ interact with the weights
$\theta^{l}_{\cdot j}$, computing the weighted sum $a^{l}_{j}$. The so-called
activation function $g(\cdot)$ is applied to $a^{l}_{j}$ leading to the output
$o^{l}_{j}$, which is sent to nodes at layer $l+1$.

Whenever the information flow between neurons has no feedback (i.e., neurons
do not process information resulting from their outputs), in the sense that
information flows from the input through the neurons producing the output of
the network, the network is referred to as feedforward. Neurons are arranged
in layers, and a network formed by neurons in one layer is called single layer
network. When more than one layer are forming a neural network, layers are
generally called hidden layers since they stand between the input and the
output, i.e., the “tangible” information which consists of the input samples
and their classification targets/outputs. A feedforward neural network
receiving as input a $d$-dimensional vector and producing a $3$-dimensional
output is shown in Figure 2.

Figure 2: A feedforward network with multiple layers. Dots represent neurons
across different layers (colors). The $d$-dimensional input vector
$\bm{x}_{i}=[x_{i}^{1},\dots,x_{i}^{d}]^{T}$ is sequentially parsed to the
output, from left to right, following the connections represented in grey
which correspond to the weights of the network’s layers.

The most relevant feature of a neural network is its capacity of learning.
This corresponds to the ability to improve its outputs (performance in
classification) by tuning the parameters (weights and biases) of its neurons.
Learning algorithms of neural networks use a set of training data to
iteratively update the parameters of a neural network such that some error
measure is decreased or some performance measure is increased (see e.g.
Goodfellow et al., 2016). The data $\mathcal{D}$ consists of vectors
$\mathcal{D}_{i}=\left\\{\bm{y}_{i},\bm{x}_{i}\right\\}$, with $\bm{x}_{i}$
representing an input and $\bm{y}_{i}$ the corresponding target for
$i=1,...,N$. Let $\hat{\bm{y}}_{i}$ denote the output of the network
corresponding to the sample $\bm{x}_{i}$, that is
$\hat{\bm{y}}_{i}=\text{NN}_{\bm{\theta}}\left(\bm{x}_{i}\right)$, with
$\text{NN}_{\bm{\theta}}\left(\bm{x}_{i}\right)$ denoting a Neural Network
parametrized over ${\bm{\theta}}$ and evaluated at $\bm{x}_{i}$. An error
function $E\left(\mathcal{D},{\bm{\theta}}\right)$ is defined at a particular
parameter ${\bm{\theta}}$, which is used to guide the learning process.
Several error functions have been used to this end, the most widely adopted
ones being the mean-squared error (suitable for regression problems) and the
cross-entropy (suitable for classification problems). The gradient of the
error between the network’s outputs $\hat{\bm{y}}_{i}$ and the targets
$\bm{y}_{i}$ over the entire data set (full-batch) or a subset of the data
(mini-batch) is commonly used to update the network parameter values through
an iterative optimization process, commonly a variant of the Backpropagation
algorithm (Rumelhart et al., 1986). Widely used iterative optimization methods
are the Stochastic Gradient Descent (SGD)(Robbins and Monro, 1951), Root Mean
Squared Propagation (RMSProp)(Tieleman and Hinton, 2012) and Adaptive Moment
Estimation (ADAM)(Kingma and Ba, 2014).

While feedforward neural networks with affine neurons have been briefly
described above, a large variety of neural networks have been proposed and
used for modeling different input-output data relationships. Such networks
follow the main principles as those described above (i.e., they are formed by
layers of neurons, which perform transformations followed by differentiable
activation functions), but they are realized by using different types of
neurons and/or transformations. Examples include the Radial Basis Function
(RBF) networks (Broomhead and Lowe, 1988) which replace affine transformation
with distance-based transformation, Convolutional Neural Networks (Homma et
al., 1987) which receive a tensor input and use neurons performing
convolution, Recurrent Neural Networks (e.g., Long-Short Term Memory (LSTM)
(Hochreiter and Schmidhuber, 1997) and Gated Recurrent Unit (GRU) (Cho et al.,
2014) networks) which model sequences of their inputs by using recurrent
units, and specialized types of neural networks, such as the Temporal-
Augmented Bilinear Layer (TABL) network (Tran et al., 2019) based on bilinear
mapping, and the Neural Bag-of-Features network (Passalis et al., 2020)
extending the classical Bag-of-Features model with a differentiable processing
suitable to be used in combination with other types of neural network layers.

###  1.3 Motivation for adopting Bayesian Neural Networks

Bayesian Neural Networks are interesting tools under three perspectives: (i)
theoretical, (ii) methodological, and (iii) practical. In the following, we
shall briefly discuss what we mean by the above three interconnected
perspectives.

From a theoretical perspective, BNN s allow for differentiating and
quantifying two different sources of uncertainty, namely epistemic uncertainty
and aleatoric uncertainty (see e.g. Der Kiureghian and Ditlevsen, 2009, from a
ML perspective). Epistemic uncertainty is the one referring to the lack of
knowledge and it is captured by $p\left({\bm{\theta}}|\mathcal{D}\right)$. In
the light of the Bayes theorem, epistemic uncertainty can be reduced with the
use of additional data, so that the lack of knowledge is addressed as more
data are collected. After the data is collected, this results in the update of
the prior belief (before the experiment is conducted) to the posterior. Thus,
the Bayesian perspective allows mixing expert knowledge with experimental
evidence. This is quite relevant in small-sample applications where the amount
of collected data is inappropriate for classical statistical tools and results
to apply (e.g. inference based on asymptotic theory), yet it nevertheless
allows the update of the a-priori belief on the parameters,
$p\left(\bm{\theta}\right)$, into the posterior. On the other hand, the
likelihood term captures the aleatoric uncertainty, that is the intrinsic
uncertainty naturally embedded in the data, i.e. $p\left(y|\theta\right)$, in
the Bayesian framework is clearly distinguished and separated from the
aleatoric one.

Methodologically, the ability of Bayesian methods to learn from small data and
eventually converge to, e.g., non-Bayesian maximum likelihood estimates, or in
general to agree with alternative frequentist methods, when the amount of
collected data overwhelms the role of the prior in the likelihood-prior
mixture, enables the possibility of seeing Bayesian methods as generalizations
of standard non-Bayesian approaches. Within the Bayesian methods family,
certain research areas such as PAC-Bayes (Alquier, 2021), Empirical Bayes
(Casella, 1985) and Approximate Bayes Computations (Csilléry et al., 2010)
deal with such connections very tightly. In this regard there are many
examples in the statistics literature; we focus on the ML perspective. For
instance regularization, ensemble and meta-learning, Monte-Carlo dropout,
etc., can be all understood as Bayesian methods, and, e.g., Variational Bayes
can be seen as standard linear regression (Salimans and Knowles, 2013). More
in general, many ML methods can be seen as approximate Bayesian methods, whose
approximate nature makes them simpler and of practical use. Furthermore, as
the learned posterior can be reused and re-updated once new data become
available, Bayesian learning methods are well-suited for online learning
(Opper and Winther, 1999). In this regard, also the explicit use of the prior
in Bayesian formulations is aligned with the No-Free-Lunch Theorem (Wolpert,
1996) whose philosophical interpretation, among the others, is that any
supervised algorithm implicitly embeds and encodes some form of prior,
establishing a tight connection with Bayesian theory (Serafino, 2013; Guedj
and Pujol, 2021).

From a practical perspective, the Bayesian approach implicitly allows for
dealing with uncertainties, both in the estimated parameters and in the
predictions. For a practitioner, this is by far the most relevant aspect in
shifting from a standard ANN approach to BNN s. Thus, with little surprise,
Bayesian methods have been well-received in high-risk application domains
where quantifying uncertainties is of high importance. Examples can be found
across different fields, such as industrial applications (Vehtari and
Lampinen, 1999), medical applications (e.g. Chakraborty and Ghosh, 2012; Kwon
et al., 2020; Lisboa et al., 2003), finance (e.g. Jang and Lee, 2017; Sariev
and Germano, 2020; Magris et al., 2022a, b), fraud detection (e.g. Viaene et
al., 2005), engineering (e.g. Cai et al., 2018; Du et al., 2020; Goh et al.,
2005), and genetics (e.g. Ma and Wang, 1999; Liang and Kelemen, 2004;
Waldmann, 2018).

As widely recognized, the estimation of BNN is not a simple task due to the
generally non-conjugacy between the prior and the likelihood and the non-
trivial computation of the integral involved in the marginal likelihood. For
this reason, application of BNN s is relatively infrequent and their use is
not widespread across the different domains. As of now, applying Bayesian
principles in a plug-and-play fashion is challenging for the general
practitioner. On top of that, several estimation approaches have been
developed, and navigating through them can indeed be confusing. In this
survey, we collect and present parameter estimation and inference methods for
Bayesian DL at an accessible level to promote the use of the Bayesian
framework.

###  1.4 Bayesian Neural Networks

From the description in Section 1.2.1, it can be seen that the goal of
approximating a function relating the input to the output in classical ANN s
is treated under an entirely deterministic perspective. Switching towards a
Bayesian perspective in mathematical terms is rather straightforward. In place
of estimating the parameter vector ${\bm{\theta}}$, BNN s target the
estimation of the posterior distribution
$p\left({\bm{\theta}}|\mathcal{D}_{x},\mathcal{D}_{y}\right)$, that is:

$p\left({\bm{\theta}}|\mathcal{D}_{x},\mathcal{D}_{y}\right)=\frac{p\left(\mathcal{D}_{y}|\mathcal{D}_{x},{\bm{\theta}}\right)p\left({\bm{\theta}}\right)}{\int_{\Theta}p\left(\mathcal{D}_{y}|\mathcal{D}_{x},{\bm{\theta}}^{\prime}\right)p\left({\bm{\theta}}^{\prime}\right)d{\bm{\theta}}^{\prime}}\text{,}$
(2)

which stands as a simple application of the Bayes theorem. Here we assume, as
it is usually the case, that the data $\mathcal{D}$ is composed of an input
set $\mathcal{D}_{x}$ and the corresponding set of outputs $\mathcal{D}_{y}$.
In general, $\mathcal{D}_{x}$ is a matrix of regressors and $\mathcal{D}_{y}$
is either the vector or matrix (depending on whether the nature of the output
is univariate or not) of the variables that the networks aim at modeling based
on $\mathcal{D}_{x}$. Alternatively but analogously, $\mathcal{D}$ can be
thought as the collection of all input-output pairs
$\mathcal{D}=\left\\{\bm{y}_{i},\bm{x}_{i}\right\\}_{i=1}^{N}$, where $N$
denotes the sample size, and $\bm{x}_{i}$ and $\bm{y}_{i}$ are the input and
output vectors of observations for the $i$-th sample, respectively. Using this
notation, $\mathcal{D}_{x}=\left\\{\bm{x}_{i}\right\\}_{i=1}^{N}$ and
$\mathcal{D}_{y}=\left\\{\bm{y}_{i}\right\\}_{i=1}^{N}$.

While eq. (2) provides a theoretical prescription for obtaining the posterior
distribution, in practice solving for the form of the posterior distribution
and retrieving its parameters is a very challenging task. The estimation of a
BNN with MC techniques and VI is discussed in the remainder of the review,
here we continue the discussion towards different aspects.

Equation (2) involves all the ingredients required for performing Bayesian
inference on ML models, and specifically neural networks. In the first place,
eq. (2) involves a likelihood function for the data $\mathcal{D}_{y}$
conditional on the observed sample $\mathcal{D}_{x}$ and the parameter vector
${\bm{\theta}}$. The forward pass parses the input into predictions via some
parameter values, such as outputs (conditional on the data and the parameters)
follow a prescribed likelihood function. Intuitive examples are the Gaussian
likelihood (for regression) and the Binomial one (for classification). An
underlying neural network is implicit in the likelihood term
$p\left(\mathcal{D}_{y}|\mathcal{D}_{x},{\bm{\theta}}\right)$, that links the
inputs to the outputs. In other words, as is the case for ANN s, the first
step in designing a BNN is that of identifying a suitable neural network
architecture (e.g., how many layers and of which kind and size) followed by a
reasonable assumption for the likelihood function.

A major difference between ANN s and BNN s is that the latter requires the
introduction of the prior distribution over the model parameters. After all, a
prior must be in place for Bayesian inference to be performed, thus priors are
required in the BNN setup. This means that the practitioner needs to decide on
the parametric form of the prior over the parameters.

######  Example 1.

Consider a BNN to model the variables
$\mathcal{D}_{y}=\left\\{y_{i}\right\\}_{i=1}^{N}$ where
$y_{i}\in\left\\{0,1\right\\}$, based on the matrix of covariates $D_{x}$. The
likelihood is of a certain form and parametrized over a neural network whose
weights are denoted by ${\bm{\theta}}$, i.e.
$\text{NN}_{\bm{\theta}}\left(\cdot\right)$.

We can approach the above problem as a 2-class classification with
$y_{i}\in\left[0,1\right]$, and derive the likelihood from the Bernoulli
distribution

$p\left(D_{y}|\mathcal{D}_{x},{\bm{\theta}}\right)=\prod_{i=1}^{N}\hat{p}_{i}^{y_{i}}\left(1-\hat{p}_{i}\right)^{1-y_{i}}\text{,}$
(3)

where $\hat{p}_{i}=\text{NN}_{\bm{\theta}}\left(\bm{x}_{i}\right)$ denotes the
output of the network for the $i$-th sample, that is the probability that
sample $i$ belongs to class 1. The prior (on the network parameters) can be a
diagonal Gaussian
$p\left({\bm{\theta}}\right)=\mathcal{N}\left({\bm{\theta}}|0,\tau I\right)$,
where $\tau>0$ is a scalar and $I$ the identity matrix.

We can also approach the above problem as a regression to
$\bm{y}_{i}\in\mathbb{R}^{d}$ and derive the likelihood from the Multivariate
Normal distribution

$\displaystyle p\left(D_{y}|D_{x},{\bm{\theta}}\right)=$
$\displaystyle\left(2\pi\right)^{-Nk/2}|\text{det}\left(\Sigma\right)|^{-N/2}\exp\left(-\frac{1}{2}\sum_{i=1}^{N}\left(\bm{y}_{i}-\hat{\bm{y}}_{i}\right)^{\top}\Sigma^{-1}\left(\bm{y}_{i}-\hat{\bm{y}}_{i}\right)\right)$
(4)

where $\hat{\bm{y}}_{i}=\text{NN}_{\bm{\theta}}\left(\bm{x}_{i}\right)$.
Assuming that the covariance matrix $\Sigma^{-1}$ is known, the prior on
${\bm{\theta}}$ could be as well a diagonal Gaussian. If $\Sigma$ is unknown,
the prior could be the product of the above Gaussian prior with, e.g., an
Inverse Wishart prior distribution on $\Sigma$. In this case, the goal of the
Bayesian inference is that of estimating the joint posterior of
$\left({\bm{\theta}},\Sigma\right)$.

The inference goal is the posterior distribution. (i) If the problem has a
form for which the posterior can be solved analytically, we find
$p\left({\bm{\theta}}|\mathcal{D}_{x}\mathcal{D}_{y}\right)$ to be of a known
parametric form and identify the parameters characterizing it. (ii) In
general, we may proceed via MC sampling, in which case the estimation leads to
a sample, of arbitrary size, approximating the true posterior. The true
posterior remains unknown in its exact form, yet MC enables sampling from it
and thus estimating an approximate representation. (iii) Alternatively,
following VI, one sets a certain chosen parametric form for the posterior and
optimizes its parameters for a certain objective function (see Section 1.5).
While the actual posterior remains unknown, in VI one seeks an approximation
that is optimal in some sense of optimization of a certain objective on the
provided data.

Figure 3: A BNN with multiple layers. Connections correspond to random
variables, and outputs here correspond to a tri-variate distribution, whose
marginals are represented in the rightmost boxes.

Figure 3 provides an analogous representation of Figure 2, now for a BNN.
Opposed to traditional ANN s, weights in BNN s are stochastic and represented
with distributions. A probability distribution over the weights is learned by
updating the prior with the evidence supported by the data. Even though Figure
3 might give the opposite impression, the posterior over the weights is in
general a truly multivariate distribution where independence among its
dimensions, in general, does not hold.

While the above clarifies that the estimation goal is a distribution whose,
e.g., variance can be indicative of the level of confidence in the estimated
parameters, the uncertainty associated with the outputs and the generation of
the model outputs themselves remains unaddressed. The predictive distribution
is defined as

$p\left(\bm{y}_{i}|\bm{x}_{i},\mathcal{D}\right)=\int_{\Theta}p\left(\bm{y}_{i}|\bm{x}_{i},{\bm{\theta}}\right)p\left({\bm{\theta}}|\mathcal{D}\right)d{\bm{\theta}}.$
(5)

As the posterior (eq. (2)) is solved also the predictive distribution can be
recovered. Yet, in practice, it is indirectly sampled. Indeed, an intuitive
MC-related approach for approximating the predictive distribution is that of
sampling $N_{s}$ values from the posterior to create $N_{s}$ realizations of
the neural network, each based on a different parameter sample, which are used
to provide predictions. This results in a collection of predictions that
approximate the actual predictive distribution. In this way, it is relatively
simple to recover (approximations of) the predictive distribution from which,
e.g., confidence intervals can be constructed. A way to reduce the sample
forecast to single values conveying relevant information is by, e.g., using
common (sampling) moment estimators. One may evaluate

$\hat{\bm{y}}_{i}=\frac{1}{N_{s}}\sum_{j=1}^{N_{s}}\text{NN}_{{\bm{\theta}}_{j}}\left(\bm{x}_{i}\right)$
(6)

to approximate the posterior mean through model averaging (across the
different realizations ${\bm{\theta}}_{j},\>j=1,\dots,N_{s}$ and thus
different outputs) or compute

$\hat{\Sigma}_{\bm{y}_{i}}=\frac{1}{N_{s}-1}\sum_{j=1}^{N_{s}}\bm{\varepsilon}_{j,i}\bm{\varepsilon}_{j,i}^{\top}\text{,}$
(7)

with

$\quad\bm{\varepsilon}_{j,i}=\text{NN}_{{\bm{\theta}}_{j}}\left(\bm{x}_{i}\right)-\hat{\bm{y}}_{i}\text{,}$
(8)

to approximate the covariance matrix, which is indicative of the uncertainty
associated with the prediction. $N_{s}$ corresponds to the number of samples
generated from the posterior and used to generate the prediction of the
network $\text{NN}_{{\bm{\theta}}_{j}}(\cdot)$ receiving as input
$\bm{x}_{i}$. In classification, one may analogously approximate predictive
densities for the joint probability of the different classes and average such
probabilities to summarize the average probabilities of each class and
implicitly the uncertainties associated with a certain class decision, which
is typically determined by the predicted class of maximum probability:

$\hat{y}_{i}=\operatorname*{arg\,max}_{c\,\in\,C}\hat{p}_{i,c}$ (9)

with $C$ being the total number of classes and $\hat{p}_{i,c}$ the predicted
probability of class $c$ for the sample $i$.

###  1.5 Variational Inference (VI)

Let $\mathcal{D}$ denote the data and
$p\left(\mathcal{D}|{\bm{\theta}}\right)$ the likelihood of the data based on
a postulated model with ${\bm{\theta}}\in\Theta$ a $k$-dimensional vector of
model parameters. Let $p\left(\bm{\theta}\right)$ be the prior distribution on
${\bm{\theta}}$. The goal of Bayesian inference is the posterior distribution

$p\left({\bm{\theta}}|\mathcal{D}\right)=\frac{p\left(\mathcal{D},{\bm{\theta}}\right)}{p\left(\mathcal{D}\right)}=\frac{p\left(\mathcal{D}|{\bm{\theta}}\right)p\left(\bm{\theta}\right)}{p\left(\mathcal{D}\right)},$
(10)

where $p\left(\mathcal{D}\right)$ is referred to as evidence or marginal
likelihood, since
$p\left(\mathcal{D}\right)=\int_{\Theta}\left(\mathcal{D}|{\bm{\theta}}\right)p\left(\bm{\theta}\right)d{\bm{\theta}}$.
$p\left({\bm{\theta}}\right)$ acts as a normalization constant for retrieving
an actual probability distribution for
$p\left({\bm{\theta}}|\mathcal{D}\right)$. Bayesian inference is generally
difficult due to the fact that the evidence is often intractable and of
unknown form. In high-dimensional applications, Monte-Carlo methods for
sampling the posterior turn challenging and infeasible, and Variational
Inference (VI)is an attractive alternative.

VI consists in an approximate method where the posterior distribution is
approximated by the so-called variational distribution. The variational
distribution is a probability density $q\left(\bm{\theta}\right)$, belonging
to some tractable class of distributions $\mathcal{Q}$ such as, e.g., the
Exponential family. VI thus turns the Bayesian inference problem in eq. (10)
into that of finding the best approximation
$q^{\star}\left(\bm{\theta}\right)\in\mathcal{Q}$ to
$p\left({\bm{\theta}}|\mathcal{D}\right)$ by minimizing the Kullback-Leibler
(KL)divergence from $q\left(\bm{\theta}\right)$ to
$p\left({\bm{\theta}}|\mathcal{D}\right)$,

$q^{\star}=\operatorname*{arg\,min}_{q\in\mathcal{Q}}\>\text{KL}\left(q||p\left({\bm{\theta}}|\mathcal{D}\right)\right)=\operatorname*{arg\,min}_{q\in\mathcal{Q}}\int
q\left(\bm{\theta}\right)\log\frac{q\left(\bm{\theta}\right)}{p\left({\bm{\theta}}|\mathcal{D}\right)}d{\bm{\theta}}\text{.}$
(11)

By simple manipulations of the KL divergence definition, it can be shown that

$\text{KL}\left(q||p\left({\bm{\theta}}|\mathcal{D}\right)\right)=-\int
q\left(\bm{\theta}\right)\log\frac{p\left(\mathcal{D}|{\bm{\theta}}\right)p\left(\bm{\theta}\right)}{q\left(\bm{\theta}\right)}d{\bm{\theta}}+\log
p\left(\mathcal{D}\right)\text{.}$ (12)

Since $\log p\left(\mathcal{D}\right)$ is a constant not depending on the
model parameters, the KL minimization problem is equivalent to the
maximization problem of the so-called Lower Bound (LB)on $\log
p\left(\mathcal{D}\right)$,

$\mathcal{L}\left(q\right)\coloneqq\int
q\left(\bm{\theta}\right)\log\frac{p\left(\mathcal{D}|{\bm{\theta}}\right)p\left(\bm{\theta}\right)}{q\left(\bm{\theta}\right)}d{\bm{\theta}}=\mathbb{E}_{q}\left[\log\frac{p\left(\mathcal{D}|{\bm{\theta}}\right)p\left(\bm{\theta}\right)}{q\left(\bm{\theta}\right)}\right]\text{.}$
(13)

For any random vector ${\bm{\theta}}$ and a function
$g\left(\bm{\theta}\right)$ we denote by
$\mathbb{E}_{f}[g\left(\bm{\theta}\right)]$ the expectation of
$g\left(\bm{\theta}\right)$ where ${\bm{\theta}}$ follows a probability
distribution with density $f$, i.e.
$\mathbb{E}_{f}[g\left(\bm{\theta}\right)]=\mathbb{E}_{{\bm{\theta}}\sim
f}[g\left(\bm{\theta}\right)]$. To make explicit the dependence of the LB on
some vector of parameters ${\bm{\zeta}}$ parametrizing the variational
posterior we write
$\mathcal{L}\left({\bm{\zeta}}\right)=\mathcal{L}\left(q_{\bm{\zeta}}\right)=\mathbb{E}_{q_{\bm{\zeta}}}\left[\log
p\left(\bm{\theta}\right)-\log
q_{\bm{\zeta}}\left(\bm{\theta}\right)+p\left(\mathcal{D}|{\bm{\theta}}\right)\right]$.
We operate within the Fixed-Form Variational Inference (FFVI)framework, where
the parametric form of the variational posterior is set. I.e. FFVI seeks at
finding the best $q\equiv q_{\bm{\zeta}}$ in the class $\mathcal{Q}$ of
distributions indexed by a vector parameter ${\bm{\zeta}}$ that minimizes the
LB $\mathcal{L}\left({\bm{\zeta}}\right)$. In this context, ${\bm{\zeta}}$ is
called variational parameter. A common choice for $\mathcal{Q}$ is the
Exponential family and ${\bm{\zeta}}$ is the corresponding natural parameter.

####  1.5.1 Estimation with Stochastic Gradient Descent (SGD)

A straightforward approach to maximize $\mathcal{L}\left({\bm{\zeta}}\right)$
is that of using a gradient-based method such as Stochastic Gradient Descent
(SGD), Adaptive Moment Estimation (ADAM)(Kingma and Ba, 2014), or Root Mean
Squared Propagation (RMSProp)(Tieleman and Hinton, 2012). The form of the
basic SGD update is

${\bm{\zeta}}_{t+1}={\bm{\zeta}}_{t}+\beta_{t}\left.\left[\hat{\nabla}_{\bm{\zeta}}\mathcal{L}\left({\bm{\zeta}}\right)\right]\right|_{{\bm{\zeta}}={\bm{\zeta}}_{t}}\text{,}$
(14)

where $t$ denotes the iteration, $\beta_{t}$ a (possibly adaptive) step size,
and $\hat{\nabla}_{\bm{\zeta}}\mathcal{L}\left({\bm{\zeta}}\right)$ a
stochastic estimate of
$\nabla_{\bm{\zeta}}\mathcal{L}\left({\bm{\zeta}}\right)$. The derivative,
considered with respect to ${\bm{\zeta}}$, is evaluated at
${\bm{\zeta}}={\bm{\zeta}}_{t}$.

Under a pure Gaussian variational assumption, it is instinctive to optimize
the LB for the mean vector ${\bm{\zeta}}_{1}={\bm{\mu}}$ and variance-
covariance matrix ${\bm{\zeta}}_{2}=\Sigma$. In the wider FFVI setting with
$\mathcal{Q}$ being the Exponential family, the LB is often optimized in terms
of the natural parameter ${\bm{\lambda}}$ (Wainwright and Jordan, 2008). The
application of the SGD update based on the standard gradient is problematic
because it ignores the information geometry of the distribution
$q_{\bm{\zeta}}$ (Amari, 1998), as it implicitly relies on the Euclidean
distance to capture the dissimilarity between two distributions in terms of
the euclidean norm $||{\bm{\zeta}}_{t}-{\bm{\zeta}}||^{2}$, which can be a
quite poor and misleading measure of dissimilarity (Khan and Nielsen, 2018).
By replacing the Euclidean norm with the KL divergence, the SGD update results
in the following natural gradient update:

$\displaystyle{\bm{\lambda}}_{t+1}={\bm{\lambda}}_{t}+\beta_{t}\left[\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}\right)\right]\text{.}$
(15)

The natural gradient update results in better step directions toward the
optimum when optimizing the distribution parameter. The natural gradient of
$\mathcal{L}\left({\bm{\lambda}}\right)$ is obtained by rescaling the gradient
$\nabla_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}\right)$ by the inverse
of the Fisher Information Matrix (FIM)$\mathcal{I}_{\bm{\lambda}}$,

$\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}\right)=\mathcal{I}^{-1}_{\bm{\lambda}}\nabla\mathcal{L}_{\bm{\lambda}}\left({\bm{\lambda}}\right),$
(16)

where subscript in $\mathcal{I}^{-1}_{\bm{\lambda}}$ remarks that the FIM is
expressed in terms of the natural parameter ${\bm{\lambda}}$. By replacing in
the above $\nabla_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}\right)$ with a
stochastic estimate
$\hat{\nabla}_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}\right)$ one
obtains a stochastic natural gradient update.

######  Example 2 (Problem statement for Variational Inference (VI)).

Consider a BNN to model the targets $\bm{y}_{i}$, based on the covariates
$\bm{x}_{i}$. The likelihood, of a certain form, is parametrized over a neural
network, whose weights are denoted by ${\bm{\theta}}$. The prior could be a
Gaussian distribution with, e.g., zero-mean, diagonal
$p\left(\bm{\theta}\right)=\mathcal{N}\left({\bm{\theta}}|\bm{0},I/\tau\right)$
or not
$p\left(\bm{\theta}\right)=\mathcal{N}\left({\bm{\theta}}|\bm{0},\Sigma_{0}\right)$.
$\mathcal{Q}$ is the set of multivariate Gaussian distributions, specified
e.g. in terms of the natural parameter ${\bm{\lambda}}$.

The objective is that of finding the corresponding variational parameter such
that the LB $\mathbb{E}_{q_{\bm{\lambda}}}\left[\log
p\left(\bm{\theta}\right)-\log
q_{\bm{\lambda}}\left(\bm{\theta}\right)+p\left(\mathcal{D}|{\bm{\theta}}\right)\right]$
is maximized. The update of the variational parameter ${\bm{\lambda}}$ follows
a gradient-based method, with natural gradients. The training terminates after
the LB $\mathcal{L}\left({\bm{\lambda}}\right)$ does not improve for a certain
number of iterations: the terminal ${\bm{\lambda}}$ provides the natural
parameter of the variational posterior approximation minimizing the KL
divergence to the true posterior $p\left({\bm{\theta}}|\mathcal{D}\right)$.

##  2 Sampling methods

###  2.1 Monte Carlo Markov Chain (MCMC)

Monte Carlo Markov Chain (MCMC)is a set of methods for sampling from a
probability distribution. MCMC s have numerous applications, and especially in
Bayesian statistics are a fundamental tool. The foundation of MCMC methods are
Markov Chains, stochastic models describing a sequence of events in which the
probability of each event depends only on the state of the previous one
(Gagniuc, 2017). By constructing a Markov Chain that has the desired
distribution as its stationary distribution, towards which the sequence
eventually converges, one can obtain samples from it, i.e. one can sample any
generic probability distribution, including e.g. a complex, perhaps multi-
modal, Bayesian posterior. Early samples may be autocorrelated and not
representative of the target distribution so that MCMC methods generally
require a burnout period before attaining the so-called stationary
distribution. In fact, while the construction of a Markov Chain converging to
the desired distribution is relatively simple, determining the number of steps
to achieve such convergence with an acceptable error is much more challenging
and strongly dependent on the initial setup and starting values. With burnout,
the large collection of samples is practically subsampled by discarding an
initial fraction of draws (e.g. 20%) to obtain a collection of approximately
independent samples from the desired distribution. An accessible introduction
to Markov Chains can be found in (Gagniuc, 2017), for a dedicated monograph on
MCMC methods oriented toward Bayesian statistics and applications see, e.g.,
(Gamerman and Lopes, 2006).

Within the class of MCMC methods some popular ones are not effective in large
Bayesian problems such as BNN s. For example, the plain Gibbs sampler (Geman
and Geman, 1984) despite its simplicity and desirable properties (Casella and
George, 1992), it suffers from residual autocorrelation between successive
samples and becomes increasingly difficult as the dimensionality increases in
multivariate distributions (e.g. Lynch, 2007, Ch. 4). We review the most
widespread MC approaches in the context of performing Bayesian learning for
Neural networks.

###  2.2 Metropolis-Hastings (MH)

The Metropolis-Hastings (MH)algorithm (Metropolis et al., 1953; Hastings,
1970) is particularly helpful in Bayesian inference as it allows drawing
samples from any probability distribution $p$, given that a function $f$
proportional up to a constant to $p$ can be computed. This is particularly
convenient as it allows to sample a Bayesian posterior by only evaluating
$f\left({\bm{\theta}}\right)=p\left({\bm{\theta}}|y\right)p\left(\bm{\theta}\right)$,
completely excluding the normalization factor from the computations. The
values of the Markov Chain are sampled iteratively with each value depending
solely on the preceding one: at each iteration, based on the current value,
the algorithm picks a candidate value ${\bm{\theta}}$ (proposed value) which
is either accepted or rejected randomly, with a probability that depends on
the current and earlier values. Upon acceptance, the proposed value is used
for the next iteration, otherwise is discarded and the current value is used
in the next iteration. As the algorithm proceeds and more sample values are
generated, the sampled-value distribution more and more closely approximates
the target distribution $p$.

A key ingredient in MH is the proposal density determining the drawing of the
proposed value at each iteration. This is formalized by an arbitrary
probability density $g\left({\bm{\theta}}^{\star}|\cdot\right)$, upon which
depends the probability of drawing ${\bm{\theta}}^{\star}$ given the previous
value ${\bm{\theta}}$. $g$ is usually assumed symmetric and a common choice is
provided by a Gaussian distribution centered on ${\bm{\theta}}$. Algorithm 1
summarizes the above steps.

Algorithm 1 Metropolis-Hastings

Initialize ${\bm{\theta}}_{0}$

for $t=0$ to $N_{\text{max}}$ do

Simulate $x^{\star}\sim g\left({\bm{\theta}}^{\star}|{\bm{\theta}}_{t}\right)$

Calculate
$r=\frac{f\left({\bm{\theta}}^{\star}\right)q\left({\bm{\theta}}_{t}|{\bm{\theta}}^{\star}\right)}{f\left({\bm{\theta}}_{t}\right)q\left({\bm{\theta}}^{\star}|{\bm{\theta}}_{t}\right)}$

Simulate $\alpha\sim\mathcal{U}\left[0,1\right]$

if $\alpha<r$ then

${\bm{\theta}}_{t+1}={\bm{\theta}}^{\star}$

else

${\bm{\theta}}_{t+1}={\bm{\theta}}_{t}$

end if

end for

return $\left\\{{\bm{\theta}}_{t}\right\\}_{t=B}^{N_{\text{max}}}$

The acceptance ratio $\alpha$ is representative of the likelihood of the
proposed sample ${\bm{\theta}}^{\star}$ over the current one
${\bm{\theta}}_{t}$ according to $p$. Indeed,
$\alpha=f\left({\bm{\theta}}^{\star}\right)/f\left({\bm{\theta}}_{t}\right)=p\left({\bm{\theta}}^{\star}\right)/p\left({\bm{\theta}}_{t}\right)$
as $f$ is proportional to $p$. A proposed sample value ${\bm{\theta}}^{\star}$
that is more probable than ${\bm{\theta}}_{t}$ ($\alpha>1$), is always
accepted, otherwise, it may be rejected with probability $\alpha$. The
algorithm thus moves around the sample space tending to stay in regions where
$p$ is of high density and, seldomly, in regions of low density. The final
collection of samples follows the distribution $p$. As the Markov chain
eventually converges to the target distribution $p$, initial samples may be
quite incompatible with $p$ especially if the algorithm is initialized at a
low-density region. Thus, it is customary to discard a number $B$ of samples
and retain only the subsample $\left\\{{\bm{\theta}}_{t}\right\\}_{t=B}^{N}$.
Note that by construction successive samples of the Markov chain are
correlated. Even though the chain eventually converges to $p$ nearby samples
are correlated causing a reduction of the effective sample size (e.g. for
$\mathbb{E}_{{\bm{\theta}}\sim p_{\bm{\theta}}}\left[{\bm{\theta}}\right]$ the
central limit theorem applies but e.g. the limiting variance is inflated by
the non-zero autocorrelation in the chain).

An important feature of the MH algorithm is that it is applicable to high
dimensions as it does not suffer from the course of dimensionality problem
causing an increasing rejection rate as the number of dimensions increases.
This makes MH suitable for large Bayesian inference problems such as training
BNN s.

###  2.3 Hamiltonian Monte Carlo (HMC)

Hamiltonian Monte Carlo (HMC)generates efficient transitions by using the
derivatives of the density function being sampled by using approximate
Hamiltonian dynamics, later corrected for performing an MH-like acceptance
step (Neal et al., 2011).

HMC augments the target probability density $p\left(\bm{\theta}\right)$ by
introducing an auxiliary momentum variable $\rho$ and generating draws from

$p\left(\rho,{\bm{\theta}}\right)=p\left(\rho|{\bm{\theta}}\right)p\left(\bm{\theta}\right)\text{.}$
(17)

Typically the auxiliary density is taken as a multivariate Gaussian
distribution, independent of ${\bm{\theta}}$:

$\rho\sim\mathcal{N}\left(\bm{0},\Sigma\right)\text{.}$ (18)

$\Sigma$ can be conveniently set to the identity matrix, restricted to a
diagonal matrix, or estimated from warm-up draws. The Hamiltonian is defined
upon the joint density $p\left(\rho,{\bm{\theta}}\right)$:

$\displaystyle H\left(\rho,{\bm{\theta}}\right)=-\log
p\left(\rho,{\bm{\theta}}\right)$ $\displaystyle=-\log
p\left(\rho|{\bm{\theta}}\right)-\log p\left(\bm{\theta}\right)$ (19)
$\displaystyle=T\left(\rho|{\bm{\theta}}\right)+V\left(\bm{\theta}\right)\text{.}$
(20)

The term $T\left(\rho|{\bm{\theta}}\right)=-\log
p\left(\rho|{\bm{\theta}}\right)$ is usually called kinetic energy and
$V\left(\bm{\theta}\right)=-\log p\left(\bm{\theta}\right)$ is called
potential energy. To generate transitions to a new state, first, a value for
the momentum is drawn independently from the current ${\bm{\theta}}$, then,
Hamilton’s equations are adopted to describe the evolution of the joint system
$\left(\rho,{\bm{\theta}}\right)$, i.e.:

$\displaystyle\frac{d{\bm{\theta}}}{dt}$ $\displaystyle=+\frac{\partial
H}{\partial\rho}=+\frac{\partial T}{\partial\rho}\text{,}$ (21)
$\displaystyle\frac{d\rho}{dt}$ $\displaystyle=-\frac{\partial
H}{\partial{\bm{\theta}}}=-\frac{\partial
T}{\partial{\bm{\theta}}}-\frac{\partial V}{\partial{\bm{\theta}}}\text{.}$
(22)

By having the momentum density being independent of the target density
$p\left(\rho|{\bm{\theta}}\right)=p\left(\rho\right)$, $\partial
T/\partial{\bm{\theta}}=\bm{0}$, the transitions are governed by the
derivatives

$\displaystyle\frac{d{\bm{\theta}}}{dt}$ $\displaystyle=\frac{\partial
H}{\partial\rho}\text{,}$ (23) $\displaystyle\frac{d\rho}{dt}$
$\displaystyle=-\frac{\partial V}{\partial{\bm{\theta}}}\text{.}$ (24)

Note that $-\partial V/\partial{\bm{\theta}}$ is simply the gradient of the
negative loglikelihood, which can be computed using automatic differentiation.
The main difficulty is the simulation of the Hamiltonian dynamics, for which
there is a variety of approaches (see e.g. Leimkuhler and Reich, 2005; Berry
et al., 2015; Hoffman et al., 2014). Yet, to solve the above system of
differential equations, a leapfrog integrator is generally used due to its
simplicity and volume-preservation and reversibility properties (Neal et al.,
2011). The leapfrog integrator is a numerically stable integration algorithm
specific for Hamiltonian systems. It discretizes time using a step size
$\varepsilon$ and alternates half-step momentum updates and full-step
parameter updates:

$\displaystyle\rho$ $\displaystyle=\rho-\frac{\varepsilon}{2}\frac{\partial
V}{\partial{\bm{\theta}}}\text{,}$ (25) $\displaystyle{\bm{\theta}}$
$\displaystyle=-{\bm{\theta}}+\varepsilon\Sigma^{-1}\rho\text{,}$ (26)
$\displaystyle\rho$ $\displaystyle=\rho-\frac{\varepsilon}{2}\frac{\partial
V}{\partial{\bm{\theta}}}\text{.}$ (27)

By repeating the above steps $L$ times, a total of $L\varepsilon$ time is
simulated and the resulting state is
$\left(\rho^{\star},{\bm{\theta}}^{\star}\right)$. Note that both $L$ and
$\varepsilon$ are hyperparameters, and their tuning is often difficult in
practice. In this regard, see the Generalized HMC approach of (Horowitz, 1991)
and developments aimed at resolving the tuning of the leapfrog iterator
(Fichtner et al., 2020; Hoffman and Sountsov, 2022),

Instead of generating a random momentum vector right away and sampling a new
state $\left(\rho^{\star},{\bm{\theta}}^{\star}\right)$, to account for
numerical errors in the leapfrog integrator (an analysis in this regard is
found in (Leimkuhler and Reich, 2005)), a Metropolis-hasting step is used. The
probability of accepting the proposal
$\left(\rho^{\star},{\bm{\theta}}^{\star}\right)$ by transitioning from
$\left(\rho,{\bm{\theta}}\right)$ is

$\min\left(1,e^{-H\left(\rho,{\bm{\theta}}\right)+H\left(\rho^{\star},{\bm{\theta}}^{\star}\right)}\right)\text{.}$
(28)

If the proposal $\left(\rho^{\star},{\bm{\theta}}^{\star}\right)$ is accepted,
the leapfrog integrator is initialized with a new momentum draw and
${\bm{\theta}}^{\star}$, otherwise the same $\left(\rho,{\bm{\theta}}\right)$
parameters are returned to start the next iteration. The HMC procedure is
summarized in Algorithm 2. Besides the difficulty of calibrating the
hyperparameters $L$ and $\varepsilon$, HMC suffers from multimodality yet the
Hamiltonian boosts the local exploration for unimodal targets.

Algorithm 2 Hamiltonian Monte Carlo

Initialize ${\bm{\theta}}_{0}$, $\rho_{0}$, set $\varepsilon$

for $t=0$ to $N_{\text{max}}$ do

Simulate $\rho\sim\mathcal{N}\left(0,\Sigma\right)$

$\left(x^{(0)},p^{(0)}\right)=\left(x_{i-1},p\right)$

for $j=1$ to $L$ do

$\rho^{(j-\frac{1}{2}\varepsilon)}=\rho^{(j-1)}-\frac{\varepsilon}{2}\left.\nabla
V\left(\bm{\theta}\right)\right|_{{\bm{\theta}}={\bm{\theta}}_{j-1}}$

${\bm{\theta}}^{(j)}={\bm{\theta}}^{(j-1)}+\varepsilon\Sigma^{-1}\rho^{(j-\frac{1}{2})}$

$\rho^{(j)}=\rho^{(j-\frac{1}{2})}-\frac{\varepsilon}{2}\left.\nabla
V\left(\bm{\theta}\right)\right|_{{\bm{\theta}}={\bm{\theta}}_{j}}$

end for

$\left(x^{\star},p^{\star}\right)=\left(x^{(L)},p^{(L)}\right)$

Simulate $\alpha\sim\mathcal{U}\left[0,1\right]$

if
$\alpha<\min\left(1,\exp\left(H\left(\rho,{\bm{\theta}}\right)-H\left(\rho^{\star},{\bm{\theta}}^{\star}\right)\right)\right)$
then

$\left(x_{t},p_{t}\right)=\left(x^{\star},p^{\star}\right)$

else

$\left(x_{t},p_{t}\right)=\left(x_{t-1},p_{t-1}\right)$

end if

end for

return $\left\\{{\bm{\theta}}_{t},\rho_{t}\right\\}_{t=B}^{N\text{max}}$

##  3 Monte Carlo Dropout (MCD)

Monte Carlo Dropout (MCD)stands as an indirect method to perform Bayesian
inference. Dropout has been earlier proposed as a regularization method for
avoiding overfitting and improving neural networks’ predictive performance
(Srivastava et al., 2014). This is achieved by applying a multiplicative
Bernoulli noise on the neurons constituting the layers of the network. This
corresponds to randomly switching off some neurons at each training step. The
dropout rate sets the probability $p_{i}$ of a neuron $i$ being switched off.
Though Bernoulli noise is the most common choice, note that other types of
noise can be as well adopted (e.g. Shen et al., 2018). Neurons are randomly
switched off only in the training phase, and the very same network
configuration in terms of the activated and disabled neurons is used during
backpropagation for computing gradients for weights’ calibration. On the other
hand, all the neurons are left activated for predictions. Though it is
intuitive that the above procedure implicitly connects to model averaging
across different randomly pruned architectures obtained from a certain DL
network, the exact connection between MC dropout and Bayesian inference
follows a quite elaborated theory.

(Gal and Ghahramani, 2016) shows that a neural network of arbitrary depth and
non-linearity with dropout applied before every single layer is mathematically
equivalent to an approximation to the probabilistic deep Gaussian Process
(GP)model (Damianou and Lawrence, 2013), and (Jakkala, 2021) for a recent
survey. That is, the dropout objective minimizes the KL divergence between a
certain approximate variational model and the deep GP. A treatment limited to
multi-layer perceptron networks is provided in (Gal and Ghahramani, 2015).

With $\hat{\bm{y}}$ being the output of a Neural Network with $L$ layers whose
loss function is $E$, for each layer $i=1,\dots,L$ let $W_{i}$ denote the
corresponding weight matrix of dimension $K_{i}\times K_{i-1}$, and
$\bm{b}_{i}$ the bias vector of dimension $K_{i}$. Be $\bm{y}_{n}$ the target
for the input $\bm{x}_{n}$ for $n=1,\dots,N$ and denote the input and output
sets respectively with $\mathcal{D}_{x}$ and $\mathcal{D}_{y}$. A typical
optimization objective includes a regularization term weighted by some decay
parameter $\lambda$, that is

$\mathcal{L}_{\text{dropout}}=\frac{1}{N}\sum_{n=1}^{N}E\left(\bm{y}_{n},\hat{\bm{y}}_{n}\right)+\lambda\sum_{i=1}^{L}\left(||{W_{i}}||^{2}_{2}+||{b_{i}}||^{2}_{2}\right)\text{.}$
(29)

Now consider a deep Gaussian process for modeling distributions over functions
corresponding to different network architectures. Assume its covariance is of
the form

$K\left(\bm{x},\bm{y}\right)=\int
p\left(\bm{w}\right)p\left(b\right)\sigma\left(\bm{w}^{\top}\bm{x}+b\right)\sigma\left(\bm{w}^{\top}\bm{y}+b\right)\>\text{d}\bm{w}\>\text{d}b\text{,}$
(30)

where $\sigma\left(\cdot\right)$ is an element-wise non-linearity, and
$p\left(\bm{w}\right)$, $p\left(b\right)$ distributions. Now let $W_{i}$ be a
random matrix of size $K_{i}\times K_{i-1}$ for each layer $i$, be
$\bm{\omega}=\left\\{W_{i}\right\\}_{i=1}^{L}$. The predictive distribution of
the deep GP model can be expressed as

$p\left(\bm{y}_{n}|\bm{x}_{n},\mathcal{D}_{x},\mathcal{D}_{y}\right)=\int
p\left(\bm{y}_{n}|\bm{x}_{n},\bm{\omega}\right)p\left(\bm{\omega}|\mathcal{D}_{x},\mathcal{D}_{y}\right)\>\text{d}\bm{\omega}\text{,}$
(31)

where $p\left(\bm{\omega}||\mathcal{D}_{x},\mathcal{D}_{y}\right)$ is the
posterior distribution. $p\left(\bm{y}_{n}|\bm{x}_{n},\bm{\omega}\right)$ is
determined by the likelihood, while $\hat{\bm{y}}_{n}$ is a function of
$\bm{x}_{n}$ and $\bm{\omega}$:

$\displaystyle p\left(\bm{y}_{n}|\bm{x}_{n},\bm{\omega}\right)$
$\displaystyle=\mathcal{N}\left(\bm{y}_{n};\hat{\bm{y}}_{n},I/\tau\right)\text{,}$
(32)
$\displaystyle\hat{\bm{y}}\equiv\hat{\bm{y}}\left(\bm{x}_{n},\bm{\omega}\right)$
$\displaystyle=\sqrt{\frac{1}{K_{L}}}W_{L}\sigma\left(\dots\sqrt{\frac{1}{K_{1}}}W_{2}\sigma\left(W_{1}\bm{x}_{n}+\bm{m}_{1}\right)\dots\right)\text{.}$
(33)

$\bm{m}_{i}$ are vectors of size $K_{i}$ for each GP layer. For the
intractable posterior
$p\left(\bm{\omega}|\mathcal{D}_{x},\mathcal{D}_{y}\right)$, (Gal and
Ghahramani, 2016) uses the variational approximation
$q\left(\bm{\omega}\right)$ defined as

$\displaystyle\bm{\omega}$ $\displaystyle=\left\\{W_{i}\right\\}_{i=1}^{L},$
(34) $\displaystyle W_{i}$
$\displaystyle=M_{i}\,\text{diag}\left(\left[\bm{z}_{i,j}\right]_{j=1}^{K_{i}}\right),$
(35) $\displaystyle z_{i,j}$
$\displaystyle\sim\text{Bernoulli}\left(p_{i}\right),$ (36)
$\displaystyle\text{\quad\, for }i=1,\dots,L\text{, }j=1,\dots,K_{i-1},$

where the collection of probabilities $p_{i}$ and matrices $M_{i}$,
$i=1,\dots,L$ constitute the variational parameter. Thus, $q$ stands as a
distribution over (non-random) matrices whose columns are randomly set to
zero, and $z_{i,j}=0$ implies that the unit $j$ in layer $i-1$ is dropped as
an input to layer $i$. For minimizing the KL divergence form $q$ to
$p\left(\bm{\omega}|\mathcal{D}_{x},\mathcal{D}_{y}\right)$, the objective
corresponds to

$-\int q\left(\bm{\omega}\right)\log
p\left(\mathcal{D}_{y}|\mathcal{D}_{x},\bm{\omega}\right)+KL\left(q\left(\bm{\omega}\right)||p\left(\bm{\omega}\right)\right)\text{,}$
(37)

By use of Monte Carlo integration and some further approximations (see Gal and
Ghahramani, 2016, for details), the objective reads

$\mathcal{L}\propto\frac{1}{\tau N}\sum_{n=1}^{N}-\log
p\left(\bm{y}_{n}|\bm{x}_{n},\hat{\bm{\omega}}_{n}\right)+\sum_{i=1}^{L}\left(\frac{p_{i}l^{2}}{2\tau
N}||{M_{i}}||_{2}^{2}+\frac{l^{2}}{2\tau N}||{\bm{m}_{i}}||^{2}_{2}\right)$
(38)

which, up to the constant $\frac{1}{\tau N}$, is a feasible and unbiased MC
estimator of eq. (37) where $\hat{\bm{\omega}}$ denotes a single MC draw from
the posterior $\hat{\bm{\omega}}_{n}\sim q\left(\bm{\omega}\right)$. By taking
$E\left(\bm{y}_{n},\hat{\bm{y}}_{n}\right)=-\log
p\left(\bm{y}_{n}|\bm{x}_{n},\hat{\bm{\omega}}_{n}\right)/\tau$ eq. (38) and
eq. (29) are equivalent for an appropriate choice of the hyperparameters
$\tau$ and $l$. This shows that the minimization of the loss in eq. (29) with
dropout is equivalent to minimizing the KL divergence from $q$ to
$p\left(\bm{\omega}|\mathcal{D}_{x},\mathcal{D}_{y}\right)$, thus performing
VI on the deep Gaussian process.

With an SGD approach, one can maximize the above LB and estimate the
variational parameters from which one can simply obtain samples from the
predictive distribution $q\left(\bm{y}^{\star}|\bm{x}^{\star}\right)$, and
approximate its mean by the naive MC estimator:

$\mathbb{E}_{q\left(\bm{y}^{\star}|\bm{x}^{\star}\right)}\left(\bm{y}^{\star}\right)\approx\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\hat{\bm{y}}^{\star}\left(\bm{x}^{\star},\bm{\omega}^{s}=\left\\{W^{s}_{1},\dots,W^{s}_{L}\right\\}\right)\text{.}$
(39)

$\bm{x}^{\star}$ denotes a new observation, not in $\mathcal{D}_{x}$, for
which the corresponding prediction is $\hat{\bm{y}}^{\star}$. I.e., the
predictive mean is obtained by performing $N_{s}$ forward passes through the
network with Bernoulli realizations
$\left\\{\bm{z}^{s}_{1},\dots,\bm{z}^{s}_{L}\right\\}_{s=1}^{N_{s}}$ with
$\bm{z}^{s}_{i}=[\bm{z}^{s}_{i,j}]_{j=1}^{K_{i}}$ for $s=1,\dots,N_{s}$,
giving $\left\\{W^{s}_{1},\dots,W^{s}_{L}\right\\}_{s=1}^{N_{s}}$. Such
average predictions are generally referred to as MC dropout estimates.
Similarly, by simple moment-matching, one can estimate the predictive variance
and higher-order statistics synthesizing the properties of
$q\left(\bm{y}^{\star}|\bm{x}^{\star}\right)$.

The predictive distribution is in general a multi-modal distribution resulting
from superposing bi-modal distribution on each weight matrix column. This
constitutes a drawback of MCD, as well the implicit VI on a GP. Furthermore,
the VI approximation in eqs. (34)-(36) may be adequate or not. It is clear
that even though MCD is a possibility for VI in deep-learning models, it is
constrained by the very specific form in eq. (34) of the variational posterior
that implicitly corresponds to performing VI on a deep GP. Furthermore, there
is evidence that MCD does not fully capture uncertainty associated with model
predictions (Chan et al., 2020), and there are issues related to the use of
improper priors and singularity of the approximate posterior. The latter ones
are addressed and explored in (Hron et al., 2018), suggesting the use of the
so-called Quasi-KL divergence as a remedy. Clearly, high dropout rates drive
the convergence rate slow, expand the network training time, and can cause
important training data to be missed or given little relative importance.
However, compared to the traditional approach for neural networks, applying
dropout places no additional effort and is often of faster training than other
VI methods. Furthermore, if a network has been trained with dropout, only by
including an additional form of regularization acting as a prior turns the ANN
into a BNN, without requiring re-estimation (Jospin et al., 2022).

##  4 Bayes-By-Backprop (BBB)

A common approach for estimating the variational posterior over the networks’
weights is the Bayes-By-Backprop (BBB)method of Blundell et al. (2015),
perhaps a breakthrough in probabilistic deep learning as a practical solution
for Bayesian inference.

The key argument in (Blundell et al., 2015) is the use of the local
reparametrization trick under which the derivative of an expectation can be
expressed as the expectation of a derivative. It introduces a random variable
$\bm{\varepsilon}$ having a probability density given by
$q\left(\bm{\varepsilon}\right)$ and a deterministic transform
$t\left({\bm{\theta}},\bm{\varepsilon}\right)$ such that
$\bm{w}=t\left({\bm{\theta}},\bm{\varepsilon}\right)$. The main idea is that
the random variable $\bm{\varepsilon}$ is a source of noise that does not
depend on the variational distribution, and the weights $\bm{w}$ are sampled
indirectly as a deterministic transformation of $\bm{\varepsilon}$, leading to
a training algorithm that is analogous to that used in training regular
networks. Indeed, by writing $\bm{w}$ as
$\bm{w}=t\left({\bm{\theta}},\bm{\varepsilon}\right)$, in place of evaluating

$\frac{\partial}{\partial{\bm{\theta}}}\mathbb{E}_{q\left(\bm{w}|{\bm{\theta}}\right)}\left[f\left(\bm{w},{\bm{\theta}}\right)\right]=\frac{\partial}{\partial{\bm{\theta}}}\int
q\left({\bm{\theta}}|\bm{w}\right)f\left(\bm{w},{\bm{\theta}}\right)\text{d}\bm{w}\text{,}$
(40)

which can be complex and rather tedious, under the assumption
$q\left(\bm{\varepsilon}\right)d\bm{\varepsilon}=q\left(\bm{w}|{\bm{\theta}}\right)\text{d}\bm{w}$,
Blundell et al. (2015) prove that

$\displaystyle\frac{\partial}{\partial{\bm{\theta}}}\mathbb{E}_{q\left(\bm{w}|{\bm{\theta}}\right)}\left[f\left(\bm{w},{\bm{\theta}}\right)\right]$
$\displaystyle=\mathbb{E}_{q\left(\bm{\varepsilon}\right)}\left[\frac{\partial}{\partial{\bm{\theta}}}f\left(t\left({\bm{\theta}},\bm{\varepsilon}\right),{\bm{\theta}}\right)\right]$
(41)
$\displaystyle=\mathbb{E}_{q\left(\bm{\varepsilon}\right)}\left[\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial\bm{w}}\frac{\partial\bm{w}}{\partial{\bm{\theta}}}+\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial{\bm{\theta}}}\right]\text{.}$
(42)

With $f\left(\bm{w},{\bm{\theta}}\right)=\log
q\left(\bm{w}|{\bm{\theta}}\right)-\log
p\left(\bm{w}\right)p\left(y|\bm{w}\right)$, the right side of eq. (41) and
eq. (42) provide an alternative approach for the estimation of the gradients
of the cost function with respect to the model parameters.

In fact, upon sampling $\bm{\varepsilon}$ and obtaining $\bm{w}$, $\log
q\left(\bm{w}|{\bm{\theta}}\right)-\log
p\left(\bm{w}\right)p\left(y|\bm{w}\right)$ is a stochastic approximation of
the VI objective
$\text{KL}\left[q\left(\bm{w}|{\bm{\theta}}\right)||p\left(\bm{w}|\mathcal{D}\right)\right]=\mathbb{E}_{q\left(\bm{w}|{\bm{\theta}}\right)}\left[\log
q\left(\bm{w}|{\bm{\theta}}\right)-\log
p\left(\bm{w}\right)p\left(\mathcal{D}|\bm{w}\right)\right]$ to be minimized.

The sampled value $\bm{\varepsilon}\sim q\left(\bm{\varepsilon}\right)$,
resampled at each iteration, is independent of the variational parameters,
while $\bm{w}$ is not directly sampled but here it is a deterministic function
of $\bm{\varepsilon}$. Given $\bm{\varepsilon}$, all the quantities in the
square bracket of eq. (42) are non-stochastic, enabling the use of
backpropagation. A single draw for $\bm{\varepsilon}$ approximates the right
side of eq. (41), and suffices for providing an unbiased stochastic gradient
estimation of the relevant gradient on the left side. Equation (41) makes
explicit the possibility of using automatic differentiation to compute the
gradient of $f$ with respect to the parameter ${\bm{\theta}}$. By using a
single sampled draw ${\bm{\theta}}$ for approximating the expectation on the
right side of eq. (41), the only parameter in the loss is ${\bm{\theta}}$ and
the use of backpropagation for evaluating the gradients is straightforward.
Equation (42) instead employs backpropagation in the “usual” sense, involving
gradients of the cost with respect to the network parameters $\bm{w}$, further
rescaled by $\partial\bm{w}/\partial{\bm{\theta}}$ and shifted by $\partial
f\left(\bm{w},{\bm{\theta}}\right)/\partial{\bm{\theta}}$. Equation (42)
concerns the usual backpropagation computations in terms of the network’s
weights, the specific form of the partial derivative with respect to
${\bm{\theta}}$ that the choice of $t$ implies, while the last term depends on
the chosen form of the variational posterior only ($\bm{w}$ is here not seen
as a function of ${\bm{\theta}}$, as the form of eq. (42) results from
applying the multi-variable chain rule). This results in a general framework
for learning the posterior distribution over the network’s weights. The
following algorithm 3 summarizes the BBB approach.

Algorithm 3 Bayes-By-Backprop

repeat

Simulate $\bm{\varepsilon}\sim q\left(\bm{\varepsilon}\right)$

$\bm{w}=t\left({\bm{\theta}},\bm{\varepsilon}\right)$

$f\left(\bm{w},{\bm{\theta}}\right)=\log
q\left(\bm{w}|{\bm{\theta}}\right)-\log
p\left(\bm{w}\right)p\left(\mathcal{D}|\bm{w}\right)$

$\nabla_{\bm{\theta}}f=\text{backprop}_{\bm{\theta}}\left(f\right)$

${\bm{\theta}}={\bm{\theta}}-\beta\nabla_{\bm{\theta}}f$

until stopping criterion is met

Algorithm 3 is initialized by preliminary setting the initial values of the
variational parameter ${\bm{\theta}}$ and of course by specifying the form of
the prior and the posterior along with the form of the likelihood involving
the outputs of the forward pass obtained from the specified underlying network
structure. The update is very similar to the one employed in standard non-
Bayesian settings, where standard optimizers such as ADAM are applicable. It
is the applicability of standard optimization algorithms and the use of
classic backpropagation that constitute the major breakthrough element in BBB,
making it a feasible approach for Bayesian learning.

To make the description more explicit, and aligned with the following
sections, we present the case where the variational posterior is a diagonal
Gaussian with mean ${\bm{\mu}}$ and covariance matrix $\sigma^{2}I$. In this
case, the transform $t$ takes the simple and convenient form

$\bm{w}=t\left({\bm{\theta}},\bm{w}\right)={\bm{\mu}}+\sigma\bm{\varepsilon}\text{.}$
(43)

As $\sigma$ is required to be always non-negative, (Blundell et al., 2015)
adopts the reparametrization $\sigma=\log\left(1+\exp\left(\rho\right)\right)$
and the variational posterior parameter
${\bm{\theta}}=\left({\bm{\mu}},\rho\right)$. In this case, Algorithm 4
summarizes the BBB approach.

Algorithm 4 Bayes-By-Backprop for a diagonal Gaussian variational posterior

repeat

Simulate $\bm{\varepsilon}\sim\mathcal{N}\left(\bm{0},I\right)$

$\bm{w}={\bm{\mu}}+\log\left(1+\exp\left(\rho\right)\right)\odot\bm{\varepsilon}$

$f\left(\bm{w},{\bm{\theta}}\right)=\log
q\left(\bm{w}|{\bm{\theta}}\right)-\log
p\left(\bm{w}\right)p\left(\mathcal{D}|\bm{w}\right)$

$\nabla_{\bm{\mu}}f=\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial\bm{w}}+\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial{\bm{\mu}}}$

$\nabla_{\rho}f=\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial\bm{w}}\frac{\bm{\varepsilon}}{1+\exp\left(-\rho\right)}+\frac{\partial
f\left(\bm{w},{\bm{\theta}}\right)}{\partial\rho}$

${\bm{\mu}}={\bm{\mu}}-\beta\nabla_{\mu}f$

$\rho=\rho-\beta\nabla_{\mu}f$

until stopping criterion is met

As for Algorithm 4, one may backpropagate the gradients of $f$ w.r.t.
${\bm{\mu}}$ and $\rho$ directly. Alternatively, as for Algorithm 3, one may
use backpropagation for computing the gradients $\partial
f\left(\bm{w},{\bm{\theta}}\right)/\partial\bm{w}$, which are furthermore
shared across the updates for ${\bm{\mu}}$ and $\rho$, or, if preferred, adopt
a general automatic differentiation setup, if, e.g., the form of the
variational likelihood does not allow for a simple analytic form of the
gradient.

##  5 Exponential family and Natural gradients

Assume $q_{\bm{\lambda}}\left(\bm{\theta}\right)$ belongs to an exponential
family distribution. Its probability density function is parametrized as

$q_{\bm{\lambda}}\left(\bm{\theta}\right)=h\left(\bm{\theta}\right)\exp\left(\phi\left(\bm{\theta}\right)^{\top}{\bm{\lambda}}-A\left({\bm{\lambda}}\right)\right)\text{,}$
(44)

where ${\bm{\lambda}}\in\Omega$ is the natural parameter,
$\phi\left(\bm{\theta}\right)$ the sufficient statistic.
$A\left({\bm{\lambda}}\right)=\log\int
h\left(\bm{\theta}\right)\exp(\phi\left(\bm{\theta}\right)^{\top}{\bm{\lambda}})d\nu$
is the log-partition function, determined upon the measure $\nu$, $\phi$ and
the function $h$. The natural parameter space is defined as
$\Omega=\\{{\bm{\lambda}}\in\mathbb{R}^{d}:A\left({\bm{\lambda}}\right)<+\infty\\}$.
When $\Omega$ is a non-empty open set, the exponential family is referred to
as regular. Furthermore, if there are no linear constraints among the
components of ${\bm{\lambda}}$ and $\phi\left(\bm{\theta}\right)$, the
exponential family in eq. (44) is said of minimal representation. Non-minimal
families can always be reduced to minimal families through a suitable
transformation and reparametrization, leading to a unique parameter vector
${\bm{\lambda}}$ associated with each distribution (Wainwright and Jordan,
2008). The mean (or expectation) parameter $\bm{m}\in\mathcal{M}$ is defined
as a function of ${\bm{\lambda}}$,
$\bm{m}\left({\bm{\lambda}}\right)=\mathbb{E}_{q_{\bm{\lambda}}}\left[\phi\left(\bm{\theta}\right)\right]=\nabla_{\bm{\lambda}}A\left({\bm{\lambda}}\right)$.
Moreover, for the Fisher Information Matrix
$\mathcal{I}_{\bm{\lambda}}=-\mathbb{E}_{q_{\bm{\lambda}}}\left[\nabla_{\bm{\lambda}}^{2}\log
q_{\bm{\lambda}}\left(\bm{\theta}\right)\right]$ it holds that
$\mathcal{I}_{\bm{\lambda}}=\nabla^{2}_{\bm{\lambda}}A\left({\bm{\lambda}}\right)=\nabla_{\bm{\lambda}}\bm{m}\text{.}$
Under minimal representation, $A\left({\bm{\lambda}}\right)$ is convex, thus
the mapping $\nabla_{\bm{\lambda}}A=\bm{m}:\Omega\rightarrow\mathcal{M}$ is
one-to-one, and $\mathcal{I}_{\bm{\lambda}}$ is positive definite and
invertible (Nielsen and Garcia, 2009). $\mathcal{M}$ denotes the set of
realizable mean parameters. Therefore, under minimal representation we can
express ${\bm{\lambda}}$ in terms of $\bm{m}$ and thus
$\mathcal{L}\left({\bm{\lambda}}\right)$ in terms of
$\mathcal{L}\left(\bm{m}\right)$ and vice versa (Khan and Nielsen, 2018).

######  Example 3 (The Gaussian distribution as an exponential-family member).

The multivariate Gaussian distribution
$\mathcal{N}\left({\bm{\mu}},\Sigma\right)$ with $k$-dimensional mean vector
${\bm{\mu}}$ and covariance matrix $\Sigma$ can be seen as a member of the
exponential family (eq. (44)). Its density reads

$q_{\bm{\lambda}}\left(\bm{\theta}\right)=\left(2\pi\right)^{k/2}\exp\\{\phi\left(\bm{\theta}\right)^{\top}{\bm{\lambda}}-\frac{1}{2}{\bm{\mu}}^{\top}\Sigma^{-1}{\bm{\mu}}-\frac{1}{2}\log|\Sigma|\\}\text{,}$
(45)

where

$\phi\left(\bm{\theta}\right)=\begin{bmatrix}\theta\\\
\theta\theta^{\top}\end{bmatrix}\text{,}\quad{\bm{\lambda}}=\begin{bmatrix}{\bm{\lambda}}_{1}\\\
{\bm{\lambda}}_{2}\end{bmatrix}=\begin{bmatrix}\Sigma^{-1}{\bm{\mu}}\\\
-\frac{1}{2}\Sigma^{-1}\end{bmatrix}\text{,}\quad\bm{m}=\begin{bmatrix}\bm{m}_{1}\\\
\bm{m}_{2}\end{bmatrix}=\begin{bmatrix}{\bm{\mu}}\\\
\Sigma+{\bm{\mu}}{\bm{\mu}}^{\top}\end{bmatrix}\text{,}$ (46)

and
$A\left({\bm{\lambda}}\right)=-\frac{1}{4}{\bm{\lambda}}_{1}^{\top}{\bm{\lambda}}_{2}^{-1}{\bm{\lambda}}_{1}-\frac{1}{2}\log\left(-2{\bm{\lambda}}_{2}\right)$.
On the other hand,
${\bm{\zeta}}=\left[{\bm{\zeta}}_{1}^{\top},{\bm{\zeta}}_{2}^{\top}\right]^{\top}$
with ${\bm{\zeta}}_{1}={\bm{\mu}}=\bm{m}_{1}$ and
${\bm{\zeta}}_{2}=\Sigma=\bm{m}_{2}-{\bm{\mu}}{\bm{\mu}}^{\top}$, constitutes
the common parametrization of the multivariate Gaussian distribution in terms
of its mean and variance-covariance matrix.

By applying the chain rule,
$\nabla_{\bm{\lambda}}\mathcal{L}=\nabla_{\bm{\lambda}}\bm{m}\nabla_{\bm{m}}\mathcal{L}=\nabla_{\bm{\lambda}}\left(\nabla_{\bm{\lambda}}A\right)\mathcal{L}=\nabla^{2}_{\bm{\lambda}}A\left({\bm{\lambda}}\right)\mathcal{L}=\mathcal{I}_{\bm{\lambda}}\nabla_{\bm{m}}\mathcal{L}$,
from which

$\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}=\mathcal{I}^{-1}_{\bm{\lambda}}\nabla_{\bm{\lambda}}\mathcal{L}=\nabla_{\bm{m}}\mathcal{L}\text{,}$
(47)

The quantity $\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}$ is referred to as the
natural gradient of $\mathcal{L}$ with respect to $\lambda$ and it is obtained
by pre-multiplying the Euclidean gradient by the inverse of the FIM
(parametrized in terms of ${\bm{\lambda}}$). In general, $\mathcal{L}$ can be
a generic function whose derivative with respect to a parameter
${\bm{\lambda}}$ (not necessarily the natural parameter) exists. The standard
reference for natural gradients computation is the seminal work of Amari
(1998). Within a SGD context, the application of simple Euclidean gradients is
problematic as it ignores the information geometry of the distribution
$q_{\bm{\lambda}}$. Euclidean gradients implicitly rely on the Euclidean norm
to capture the dissimilarity between two distributions which can be a quite
poor dissimilarity measure (Khan and Nielsen, 2018). In fact, the SGD update
can be obtained by writing

${\bm{\lambda}}_{t+1}=\operatorname*{arg\,min}_{{\bm{\lambda}}}{\bm{\lambda}}^{\top}\left[\nabla_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}_{t}\right)\right]-\frac{1}{2\beta}||{\bm{\lambda}}-{\bm{\lambda}}_{t}||^{2}$
(48)

and setting to zero its derivative. Although the above implies that
${\bm{\lambda}}$ moves in the direction of the gradient, it remains close to
the previous ${\bm{\lambda}}_{t}$ in terms of Euclidean distance. As
${\bm{\lambda}}$ is a parameter of a distribution, the adoption of the
Euclidean measure is misleading. An Exponential family distribution induces a
Riemannian manifold with a metric defined by the FIM (Khan and Nielsen, 2018).
By replacing the Euclidean metric with the Riemannian one,

${\bm{\lambda}}_{t+1}=\operatorname*{arg\,min}_{{\bm{\lambda}}}{\bm{\lambda}}^{\top}\left[\nabla_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}_{t}\right)\right]-\frac{1}{2\beta}\left({\bm{\lambda}}-{\bm{\lambda}}_{t}\right)^{\top}\mathcal{I}_{\bm{\lambda}}\left({\bm{\lambda}}-{\bm{\lambda}}_{t}\right)$
(49)

the resulting update is indeed expressed in terms of the natural parameter:

${\bm{\lambda}}_{t+1}={\bm{\lambda}}_{t}+\beta\mathcal{I}^{-1}_{\bm{\lambda}}\nabla_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}_{t}\right)\text{,}$
(50)

generally referred to as natural gradient update. More in general, one could
replace the Euclidean distance with a proximity function such as the Bregman
divergence and obtain richer classes of SGD-like updates, like mirror descent
(which can be interpreted as natural gradient descent), see e.g. (Nielsen,
2020). A very interesting point on the limitations of plain gradient search is
made in (Wierstra et al., 2014) concerning the impossibility of locating, even
in a one-dimensional case, a quadratic optimum. The example provided therein
involves the Gaussian distribution, pivotal in VI. For an one-dimensional
Gaussian distribution with mean $\mu$ and standard deviation $\sigma$, the
gradient of $\mathcal{L}$ with respect to the parameters $\mu$ and $\sigma$
lead to the following SGD updates:

$\displaystyle\mu$
$\displaystyle=\mu+\beta\nabla_{\mu}\mathcal{L}=\mu+\beta\frac{z-\mu}{\sigma^{2}}$
(51) $\displaystyle\sigma$
$\displaystyle=\sigma+\beta\nabla_{\sigma}\mathcal{L}=\frac{\left(z-\mu\right)^{2}-\sigma^{2}}{\sigma^{3}}\text{.}$
(52)

For the updates to converge and the optimum to be precisely located, $\sigma$
must decrease (i.e. the distribution shrinks around $\mu$). The fact that
$\sigma$ appears in the denominator of both the updates is problematic: as it
decreases, the variance of the updates increases as
$\Delta_{\mu}\propto\frac{1}{\sigma}$ and
$\Delta_{\sigma}\propto\frac{1}{\sigma}$. The updates become increasingly
unstable, and a large overshooting update makes the search start all over
again rather than converging. Increased population size and small learning
rates cannot avoid the problem. The choice of the starting value is
problematic too: starting with $\sigma>>1$ makes the updated minuscule,
conversely $\sigma<<1$ makes them huge and unstable. (Wierstra et al., 2014)
discusses how the use of natural gradients fixes this issue, that e.g. may
arise with BBVI.

Algorithm 5 summarizes the generic scheme upon the implementation of a natural
gradient update. In Algorithm 5, ${\bm{\zeta}}$ denotes a generic variational
parameter, e.g. the natural parameter or not, while methods for evaluating
$\nabla_{\theta}\mathcal{L}$, $\mathcal{I}$, and efficiently computing its
inverse $\mathcal{I}^{-1}$ are discussed in the following sections.

Algorithm 5 Stochastic Gradient Descent with natural gradients

repeat

Compute: $\nabla_{\bm{\zeta}}\mathcal{L}$, $\mathcal{I}$

${\bm{\zeta}}={\bm{\zeta}}+\beta\mathcal{I}^{-1}\nabla_{\bm{\zeta}}\mathcal{L}$

until stopping condition is met

##  6 Black-Box methods

A major issue in VI is that it heavily relies upon model-specific
computations, on which a generalized, ready-to-use, and plug-and-play
optimizer is difficult to design. Black-box methods aim at providing solutions
that can be immediately applied to a wide class of models with little effort.
In the first instance, the ubiquitous use of model’s gradients that
traditional ML and VI approaches rely upon struggles with this principle. As
(Ranganath et al., 2014) describes, for a specific class of models, where
conditional distributions have a convenient form and a suitable variational
family exists, VI optimization can be carried out using closed-form coordinate
ascent methods (Ghahramani and Beal, 2000). In general, there is no close-form
solution resulting in model-specific algorithms (Jaakkola and Jordan, 1997;
Blei and Lafferty, 2007; Braun and McAuliffe, 2010) or generic algorithms that
involve model-specific computations (Knowles and Minka, 2011; Paisley et al.,
2012). As a consequence model assumptions and model-specific functional forms
play a central role, making VI practical. The general idea of Black-box VI is
that of rewriting the gradient of the LB objective as the expectation of an
easy-to-compute function of the latent and observed variables. The expectation
is taken with respect to the variational distribution, and the gradient is
estimated by using stochastic samples from it in a MC fashion. Such stochastic
gradients are used to update the variational parameters following an SGD
optimization approach. Within this framework, the end-user is required to
develop functions only for evaluating the model log-likelihood, while the
remaining calculations are easily implemented in libraries of general use
applicable to several classes of models. Black-box VI falls within stochastic
optimization where the optimization objective is the maximization of the LB
using noisy, unbiased, estimates of its gradient. As such, variance reduction
methods have a major impact on stability and convergence, among them control
variates are the most effective and of immediate implementation.

###  6.1 Black-Box Variational Inference (BBVI)

Black-Box Variational Inference (BBVI)optimizes the LB with stochastic
optimization, through an unbiased estimator of its gradients obtained from
samples from the variational posterior (Ranganath et al., 2014). By using the
LB definition and the log-derivative trick on the gradient of the LB with
respect to the variational parameter, $\nabla_{\bm{\zeta}}\mathcal{L}$ can be
expressed as

$\nabla_{\bm{\zeta}}\mathcal{L}=\mathbb{E}_{q}\left[\nabla_{\bm{\zeta}}\log
q\left({\bm{\theta}}|{\bm{\zeta}}\right)\left(\log
p\left(\mathcal{D},{\bm{\theta}}\right)-\log
q\left({\bm{\theta}}|{\bm{\zeta}}\right)\right)\right]\text{,}$ (53)

where ${\bm{\zeta}}$ denotes the parameter of the variational distribution
$q_{\bm{\zeta}}$. The above expression rewrites the gradient as an expectation
of a quantity that does not involve the model’s gradients but only those of
$\log q\left(w|{\bm{\zeta}}\right)$. A naive noisy unbiased estimate of the
gradient of the LB is immediate to obtain with $N_{s}$ samples obtained from
the variational distribution,

$\nabla_{\bm{\zeta}}\mathcal{L}=\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\nabla_{\bm{\zeta}}\log
q\left({\bm{\theta}}_{s}|{\bm{\zeta}}\right)\left[\log
p\left(\mathcal{D},{\bm{\theta}}_{s}\right)-\log
q\left({\bm{\theta}}_{s}|{\bm{\zeta}}\right)\right]\text{,}$ (54)

where ${\bm{\theta}}_{s}\sim q\left({\bm{\theta}}|{\bm{\zeta}}\right)$. The
above MC estimator enables the immediate and feasible computation of the LB
gradients as, given a sample ${\bm{\theta}}_{s}$, $\log
q\left({\bm{\theta}}_{s}|{\bm{\theta}}\right)$ is a quantity that solely
depends on the form of the variational posterior and can be of simple form. On
the other hand, $\log p\left(\mathcal{D},{\bm{\theta}}\right)-\log
q\left({\bm{\theta}}|{\bm{\theta}}\right)$ is immediate to compute as it only
requires evaluating the logarithm of the joint
$p\left(\mathcal{D},{\bm{\theta}}_{s}\right)$ and the density of the
variational distribution in ${\bm{\theta}}_{s}$. This process is summarized in
Algorithm 6. If sensible, one may assume that $\log
p\left(\mathcal{D},{\bm{\theta}}\right)=\log
p\left(\mathcal{D}|{\bm{\theta}}\right)p\left({\bm{\theta}}\right)$ but this
is not explicitly required as of (Ranganath et al., 2014): there are no
assumptions on the form of the model; the approach only requires the gradient
of the variational likelihood with respect to the variational parameters to be
feasible to compute.

Algorithm 6 Black-Box Variational Inference

repeat

Simulate ${\bm{\theta}}_{s}\sim q_{\bm{\zeta}}$, for $s=1,\dots,N_{s}$

$\nabla_{\bm{\zeta}}\mathcal{L}=\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\nabla_{\bm{\zeta}}\log
q\left({\bm{\theta}}_{s}|{\bm{\zeta}}\right)\left[\log
p\left(\mathcal{D}|{\bm{\theta}}_{s}\right)p\left({\bm{\theta}}_{s}\right)-\log
q\left({\bm{\theta}}_{s}|{\bm{\zeta}}\right)\right]$

${\bm{\zeta}}={\bm{\zeta}}+\beta\nabla_{\bm{\zeta}}\mathcal{L}$

until stopping criterion is met

In (Ranganath et al., 2014), the authors employ an adaptive learning rate
satisfying the Robbins Monroe conditions $\sum_{t}\beta_{t}=\infty$ and
$\sum_{t}\beta^{2}_{t}<\infty$, and for controlling the variance of the
stochastic gradient estimator adopt Rao-blackewllization (Rao, 1945;
Blackwell, 1947; Robert and Roberts, 2021) and use the of control variates
(e.g. Lemieux, 2014; Robert et al., 1999, Ch. 3) within Algorithm 6.

###  6.2 Natural-Gradient Black-Box Variational Inference (NG-BBVI)

We shall review the approach of Trusheim et al. (2018) boosting BBVI with
natural gradients, referred to as Natural-Gradient Black-Box Variational
Inference (NG-BBVI). The FIM corresponds to the outer product of the score
function with itself (see Section 5) and is furthermore equal to the second
derivative of the KL divergence to the approximate posterior
$q\left(x|{\bm{\zeta}}\right)$:

$\displaystyle F\left({\bm{\zeta}}\right)$
$\displaystyle=\left.\frac{d^{2}\text{KL}\left[q_{\bm{\zeta}}\left(\bm{\theta}\right)||q_{\hat{{\bm{\zeta}}}}\left(\bm{\theta}\right)\right]}{\left(d{\bm{\zeta}}\right)^{2}}\right|_{\hat{{\bm{\zeta}}}={\bm{\zeta}}}=\mathbb{E}_{q_{\bm{\zeta}}}\left[\nabla_{\bm{\zeta}}\log
q_{\bm{\zeta}}\left(\bm{\theta}\right)\nabla_{\bm{\zeta}}\log
q_{\bm{\zeta}}\left(\bm{\theta}\right)^{\top}\right]\text{.}$ (55)

For the practical implementation, (Trusheim et al., 2018) uses a mean-field
restriction on the variational model, i.e. the joint is factorized into the
product of $K$ independent terms, where each term is in general a multivariate
distribution:

$q_{\bm{\zeta}}\left(\bm{\theta}\right)=\prod_{k=1}^{K}q_{{\bm{\zeta}}_{k}}\left({\bm{\theta}}_{k}\right)\text{.}$
(56)

The above restriction is also suggested by Ranganath et al. (2014) in order to
allow for Rao-Blackwellization (Robert and Roberts, 2021) as a tool to be used
in conjunction with control variates (e.g. Lemieux, 2014, Ch. 3) for reducing
the variance of the stochastic gradient estimator. Under the above assumption,
the FIM simplifies to:

$\mathcal{I}_{\bm{\zeta}}=\begin{cases}\mathbb{E}_{q_{i}\left({\bm{\theta}}|{\bm{\zeta}}\right)}\left[\nabla_{{\bm{\zeta}}_{i}}\log
q_{{\bm{\zeta}}_{i}}\left({\bm{\theta}}_{i}\right)\nabla_{{\bm{\zeta}}_{i}}\log
q_{{\bm{\zeta}}_{i}}\left({\bm{\theta}}_{i}\right)\nabla_{{\bm{\zeta}}_{i}}^{\top}\right]\text{,}&i=j,\\\
0\text{,}&i\neq j,\end{cases}$ (57)

which significantly simplifies the general form
$q_{\bm{\zeta}}\left(\bm{\theta}\right)$ while implicitly enabling Rao-
Blackwellization with the variable-wise local expectations and thus reducing
the variance of the FIM, estimated via a Monte-Carlo approach. In fact,
besides a few variational models it is difficult to compute the above
expectations analytically so Trusheim et al. (2018) adopts the following naive
MC estimator:

$\hat{\mathcal{I}}_{\bm{\zeta}}=\begin{cases}\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\left[\nabla_{{\bm{\zeta}}_{i}}\log
q_{{\bm{\zeta}}_{i}}\left({\bm{\theta}}_{i}^{\left(s\right)}\right)\nabla_{{\bm{\zeta}}_{i}}\log
q_{{\bm{\zeta}}_{i}}\left({\bm{\theta}}_{i}^{\left(s\right)}\right)^{\top}\right]\text{,}&i=j\\\
0\text{,}&i\neq j\end{cases}$ (58)

with ${\bm{\theta}}_{i}^{\left(s\right)}\sim
q_{{\bm{\zeta}}_{i}}\left({\bm{\theta}}_{i}\right)$ denoting a sample from the
$i$-th factor of the posterior mean-field approximation. Note that the above
does not introduce additional computations as the score of the samples
${\bm{\theta}}_{i}^{\left(s\right)}$ is anyway required in the computation of
the LB gradient. Furthermore, instead of using a plain SGD-like update,
(Trusheim et al., 2018) adopts an ADAM-like version, boosted with natural
gradient computations. Algorithm 7 summarizes the NG-BBVI approach.

The NG-BBVI implementation is slightly more complex than the original BBVI,
see Algorithm 7. The MC computation involves both the black-box estimation
stochastic gradients and the optimal control variate coefficient $a^{\star}$.
Thus the posterior samples are split into two subsets. The first one $X$ aimed
at estimating $a^{\star}$, and the second one $Y$ at implementing the MC
estimators, independently from $X$, and with the control variate correction
term $a^{\star}$ earlier computed. The computation of the FIM follows
immediately from eq. (55), and the computation of
$\nabla_{{\bm{\zeta}}_{k}}\mathcal{L}$ is analogous to BBVI. The last four
lines of Algorithm 7 correspond to the implementation of the ADAM update,
operators are intended to be applied element-wise, $\beta_{1}$, $\beta_{2}$
(exponential decay rates) are typical ADAM hyperparameters, $\varepsilon>0$ is
a small offset preventing divisions by zero.

Algorithm 7 Natural-Gradient Black-Box Variational Inference

repeat

for $s=1$ to $N_{s}$ do

Simulate ${\bm{\theta}}\left[s\right]\sim q_{\bm{\zeta}}$

end for

$X\leftarrow\text{subset of }{\bm{\theta}}$,
$Y\leftarrow{\bm{\theta}}\setminus X$, $N_{X}=\text{numel}\left(X\right)$,
$N_{Y}=\text{numel}\left(Y\right)$

for $k=1$ to $K$ do

$a^{\star}=0$

for $s=1$ to $N_{X}$ do

$h_{X}^{\left(s\right)}=\nabla_{{\bm{\zeta}}_{k}}\log
q_{{\bm{\zeta}}_{k}}(X_{k}^{\left(s\right)})$

$f_{X}^{\left(s\right)}=h_{X}^{\left(s\right)}\left(\nabla_{{\bm{\zeta}}_{k}}\log
p_{k}(X_{k}^{\left(s\right)},\mathcal{D})-\log
q_{{\bm{\zeta}}_{k}}(X_{k}^{\left(s\right)})-a^{\star}\right)$

end for

$a^{\star}=\text{Cov}\left(h_{X},f_{X}\right)/\text{Var}\left(h_{X}\right)$

for $s=1$ to $N_{Y}$ do

$h_{Y}^{\left(s\right)}=\nabla_{{\bm{\zeta}}_{k}}\log
q_{{\bm{\zeta}}_{k}}(Y_{k}^{\left(s\right)})$

$f_{Y}^{\left(s\right)}=h_{Y}^{\left(s\right)}\left(\nabla_{{\bm{\zeta}}_{k}}\log
p_{k}(Y_{k}^{\left(s\right)},\mathcal{D})-\log
q_{{\bm{\zeta}}_{k}}(Y_{k}^{\left(s\right)})-a^{\star}\right)$

end for

$\hat{\mathcal{I}}_{{\bm{\zeta}}_{k}}=\frac{1}{N_{Y}}\sum_{s=1}^{N_{Y}}h_{Y}^{\left(s\right)}{h_{Y}^{\left(s\right)\top}}$

$\hat{\nabla}_{{\bm{\zeta}}_{k}}\mathcal{L}\leftarrow\hat{\mathcal{I}}^{-1}_{{\bm{\zeta}}_{k}}\left(\frac{1}{N_{Y}}\sum_{s=1}^{N_{Y}}f_{Y}^{\left(s\right)}\right)$

end for

$\bm{m}=\beta_{1}\bm{m}+\left(1-\beta_{2}\right)\hat{\nabla}_{{\bm{\zeta}}}\mathcal{L}$

$\bm{v}=\beta_{2}\bm{v}+\left(1-\beta_{1}\right)||{\hat{\nabla}_{{\bm{\zeta}}}\mathcal{L}}||^{2}$

$\hat{\bm{m}}=\frac{\bm{m}}{1-\beta_{2}^{t}}$,
$\hat{\bm{v}}=\frac{\bm{v}}{1-\beta_{2}^{t}}$

${\bm{\zeta}}={\bm{\zeta}}+\beta\frac{\hat{\bm{m}}}{\sqrt{\hat{\bm{v}}}+\varepsilon}$

until $||{d{\bm{\zeta}}}||<\text{threshold}$

(Trusheim et al., 2018) differs from BBVI by the use of natural gradients (and
the adoption of the ADAM-like update, though applicable to BBVI as well). On
the other hand, the use of control variates and Rao-Blackwellization for
variance reduction is found in both BBVI and NG-BBVI. As the natural gradient
approach is preferable for the reasons discussed in Section 5, NG-BBVI is
favored over BBVI.

The use of the black-box framework for computing the gradients of the LB along
the MC estimator for the FIM renders NG-BBVI of general applicability and not
constrained to a certain form of the variational posterior. Yet the MC-
computations of the FIM are implicitly approximate, whereas for certain
distributions the FIM computation can be carried out analytically and in an
exact form. NG-BBVI furthermore requires the inversion of the FIM, which is a
computational bottleneck. The following VON Khan and Nielsen (2018), VADAM
(Khan et al., 2018a) and VOGN (Khan et al., 2018a; Osawa et al., 2019) methods
indeed fix this issue: assuming a variational posterior within the exponential
distribution family, natural gradients are enabled without the direct
computation of the FIM and its inverse.

##  7 Natural gradient methods for exponential-family variational
distributions

In the following subsections, we review methods based on Natural gradients and
Exponential-family variational approximations. The following techniques are
built on natural parameter updates in the natural parameter space and rely on
simplified but exact FIM computations based on the natural/expectation
parameter duality (eq. (47)).

###  7.1 Exact gradient computations for the exponential family

The computation of the FIM required in the natural gradient computation is, in
general, not trivial. In a generic perspective, not bound to a specific
variational form, the sampling approach for the FIM estimation of (Trusheim et
al., 2018) is feasible. Yet for certain distributions, namely for those in the
Exponential family class, natural gradients can be computed in an exact form
with an analytical solution which furthermore does not involve the computation
of the FIM.

The theoretic foundation of such a viable approach is provided in (Khan and
Nielsen, 2018) and traces back to eq. (47). For an Exponential family of
minimal representation, the natural gradient with respect to the natural
parameter ${\bm{\lambda}}$ is equal to the gradient with respect to the
expectation parameter $\bm{m}$. This is a powerful result that allows the
computation of the natural gradient as an Euclidean gradient, avoiding the
computation of the FIM and its inversion.

This section presents some baseline methods using the above duality for the
natural gradient computation. Differently from BBB, BBVI, and NG-BBVI the
following approaches explicitly deal with variational distributions members of
the Exponential family with a focus on updating their natural parameter:

$\displaystyle{\bm{\lambda}}_{t+1}$
$\displaystyle={\bm{\lambda}}_{t}+\beta\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}\left({\bm{\lambda}}_{t}\right)$
(59)
$\displaystyle={\bm{\lambda}}_{t}+\beta\mathcal{I}^{-1}_{{\bm{\lambda}}_{t}}\nabla_{{\bm{\lambda}}_{t}}\mathcal{L}\left({\bm{\lambda}}_{t}\right)={\bm{\lambda}}_{t}+\beta\nabla_{\bm{m}}\mathcal{L}\left(\bm{m}_{t}\right).$
(60)

From the above updates on natural parameters, update rules for alternative and
perhaps a more usual parametrization can often be obtained, see e.g. Section
7.2.

###  7.2 Natural-Gradient Variational Inference (NGVI)

Natural-Gradient Variational Inference (NGVI)constitutes a baseline
methodology for natural gradient computation under a Gaussian variational
distribution, upon which several other approaches have been developed.

The natural gradients in the natural parameter space can be computed under the
expectation parametrization as Euclidean gradients. (Khan and Nielsen, 2018)
shows that such gradients are of simple form and correspond to

$\displaystyle\tilde{\nabla}_{{\bm{\lambda}}_{1}}\mathcal{L}=\nabla_{\bm{m}_{1}}\mathcal{L}$
$\displaystyle=\nabla_{\bm{\mu}}\mathcal{L}-2\left[\nabla_{\Sigma}\mathcal{L}\right]{\bm{\mu}},$
(61)
$\displaystyle\tilde{\nabla}_{{\bm{\lambda}}_{2}}\mathcal{L}=\nabla_{\bm{m}_{2}}\mathcal{L}$
$\displaystyle=\nabla_{\Sigma}\mathcal{L}.$ (62)

By using the definition of natural gradients in terms of ${\bm{\mu}}$ and
$\Sigma$, the update in eq. (60) for the natural parameter
${\bm{\lambda}}_{1}={\bm{\mu}}$, ${\bm{\lambda}}_{2}=-\frac{1}{2}\Sigma^{-1}$
rewrites as

$\displaystyle\Sigma^{-1}_{t+1}$
$\displaystyle=\Sigma^{-1}_{t}\nabla_{\bm{\mu}}\mathcal{L}-2\beta\left[\nabla_{\Sigma}\mathcal{L}\right]{\bm{\mu}}$
(63) $\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle=\Sigma_{t+1}\left[\Sigma^{-1}_{t}{\bm{\mu}}+\beta\left(\nabla_{\bm{\mu}}\mathcal{L}-2\left[\nabla_{\Sigma}\mathcal{L}\right]{\bm{\mu}}\right)\right]$
(64)
$\displaystyle=\Sigma_{t+1}\left[\left(\Sigma^{-1}_{t}-2\beta\left[\nabla_{\Sigma}\mathcal{L}\right]{\bm{\mu}}\right)+\beta\nabla_{\bm{\mu}}\mathcal{L}\right]$
(65)
$\displaystyle=\Sigma_{t+1}\left[\Sigma^{-1}_{t+1}{\bm{\mu}}_{t}+\beta\nabla_{\bm{\mu}}\mathcal{L}\right]$
(66)
$\displaystyle={\bm{\mu}}_{t}+\beta\Sigma_{t+1}\left[\nabla_{\bm{\mu}}\mathcal{L}\right].$
(67)

The above two constitute the NGVI update rules for updating the mean
${\bm{\mu}}$ and covariance matrix $\Sigma$ of the variational posterior with
a natural gradient update, that however does not involve the computation of
the FIM as it relies on Euclidean gradients.

For a diagonal covariance matrix
$\Sigma=\text{diag}\left({\sigma^{2}}\right)$, the corresponding NGVI updates
read

$\displaystyle\sigma^{-2}_{t+1}$
$\displaystyle=\sigma^{-2}_{t}-2\beta\left[\nabla_{\sigma^{2}}\mathcal{L}_{t}\right],$
(68) $\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}+\beta\sigma^{2}_{t+1}\odot\left[\nabla_{\bm{\mu}}\mathcal{L}_{t}\right].$
(69)

With respect to the NGVI update two points are important to stress out. First,
at each iteration, the update for ${\bm{\mu}}$ implicitly requires
$\Sigma_{t+1}$. This means that the update for ${\bm{\mu}}$ follows that for
$\Sigma^{-1}$ and that ${\bm{\mu}}$ readily uses the one-step ahead updated
information on $\Sigma$. Though it may appear counter-intuitive, (Lyu and
Tsang, 2021; Magris et al., 2022b) show that this update is not optimal (in
the terms therein discussed), while an update of the form
${\bm{\mu}}_{t+1}={\bm{\mu}}_{t}+\beta\Sigma_{t}\left[\nabla_{\bm{\mu}}\mathcal{L}\right]$
would be. Also, note that the update for ${\bm{\mu}}$ involves $\Sigma_{t+1}$
and not $\Sigma^{-1}_{t+1}$, meaning that in the NGVI an online inversion of
$\Sigma^{-1}$ is implicitly required at each iteration. Clearly, for the
diagonal case, this is trivial and effortless to obtain. Second, in the full-
covariance case, there is no guarantee that the updates guarantee $\Sigma$ to
be a positive-definite covariance matrix. This issue is tackled in Section 8.
For the diagonal case, the constraint on $\Sigma$ results in guaranteeing the
positivity of the entries in the diagonal. This can be achieved via a proper
reparametrization, e.g. BBVI updates $\rho$ where
$\sigma=\log\left(1+\exp\left(\rho\right)\right)$, or (e.g. Tan, 2021) updates
the Cholesky factor. Alternatively, the learning rate can be adapted to
guarantee that the step size does not drive the updates $\sigma^{-2}$ negative
(e.g., Khan and Nielsen, 2018; Magris et al., 2022c).

###  7.3 Variational Online Newton (VON)

A computational burden in NGVI is that the gradients of the LB are still
required: Variational Online Newton (VON)develops on NGVI but does not require
the gradients of the variational objective. Furthermore, it only involves the
gradient and Hessian of the model log-likelihood which can be computed with
usual backpropagation.

Khan et al. (2018b) express the lower bound as

$\mathcal{L}=\mathbb{E}_{q_{\bm{\lambda}}}\left[-Nf\left(\bm{\theta}\right)+\log
p\left(\bm{\theta}\right)-\log q\left(\bm{\theta}\right)\right]\text{,}$ (70)

where $N$ is the sample size, and
$f\left(\bm{\theta}\right)=-\frac{1}{N}\sum_{i=1}^{N}\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)$ is negative log-likelihood of the
model, i.e. standard MLE objective, where $\mathcal{D}_{i}$ denotes a data
example, i.e. $\mathcal{D}_{i}=\left(\bm{y}_{i},\bm{x}_{i}\right)$. VON uses
the theoretical results of (Opper and Archambeau, 2009; Rezende et al., 2014)
to express the gradients of the LB objective in terms of gradient and Hessian
of $f\left(\bm{\theta}\right)$. By linearity of the expectation, the gradients
of the $\mathcal{L}$ consist of the sum of the gradients of three expectation
terms, in particular:

$\displaystyle\nabla_{\bm{\mu}}\mathbb{E}_{q}\left[f\left(\bm{\theta}\right)\right]$
$\displaystyle=\mathbb{E}_{q}\left[\nabla_{\bm{\theta}}f\left(\bm{\theta}\right)\right]=\mathbb{E}_{q}\left[\bm{g}\left(\bm{\theta}\right)\right],$
(71)
$\displaystyle\nabla_{\Sigma}\mathbb{E}_{1}\left[f\left(\bm{\theta}\right)\right]$
$\displaystyle=\frac{1}{2}\mathbb{E}_{q}\left[\nabla^{2}_{{\bm{\theta}}{\bm{\theta}}}f\left(\bm{\theta}\right)\right]=\frac{1}{2}\mathbb{E}_{q}\left[H\left(\bm{\theta}\right)\right],$
(72)

where $\bm{g}=\nabla_{\bm{\theta}}f\left(\bm{\theta}\right)$ and
$H\left(\bm{\theta}\right)=\nabla^{2}_{{\bm{\theta}}{\bm{\theta}}}f\left(\bm{\theta}\right)$
denote the gradient and Hessian of the MLE objective, respectively. With these
relations the gradients of the LB objective write

$\displaystyle\nabla_{\bm{\mu}}\mathcal{L}$
$\displaystyle=\nabla_{\bm{\mu}}\mathbb{E}_{q}\left[-Nf\left(\bm{\theta}\right)+\log
p\left(\bm{\theta}\right)-\log q\left(\bm{\theta}\right)\right]$ (73)
$\displaystyle=-\left(\mathbb{E}_{q}\left[N\nabla_{\bm{\theta}}f\left(\bm{\theta}\right)\right]+0+{\bm{\lambda}}{\bm{\mu}}\right)$
(74)
$\displaystyle=-\left(\mathbb{E}_{q}\left[N\bm{g}\left(\bm{\theta}\right)\right]+{\bm{\lambda}}{\bm{\mu}}\right)$
(75)

and

$\displaystyle\nabla_{\Sigma}\mathcal{L}$
$\displaystyle=\frac{1}{2}\mathbb{E}_{q}\left[-N\nabla^{2}_{{\bm{\theta}}{\bm{\theta}}}f\left(\bm{\theta}\right)\right]+0-\frac{1}{2}{\bm{\lambda}}I+\frac{1}{2}\Sigma^{-1}$
$\displaystyle=\frac{1}{2}\mathbb{E}_{q}\left[-NH\left(\bm{\theta}\right)\right]-\frac{1}{2}{\bm{\lambda}}I+\frac{1}{2}\Sigma^{-1}.$

By using these gradients in the NGVI update, one obtains

$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}-\beta\Sigma_{t+1}\left[\mathbb{E}_{q}\left[N\bm{g}\left(\bm{\theta}\right)\right]+{\bm{\lambda}}{\bm{\mu}}\right],$
(76) $\displaystyle\Sigma^{-1}_{t+1}$
$\displaystyle=\left(1-\beta\right)\Sigma^{-1}_{t}+\beta\left(\mathbb{E}_{q}\left[NH\left(\bm{\theta}\right)\right]+{\bm{\lambda}}I\right),$
(77)

where the expectations can be again evaluated via MC sampling. By using a
single draw
${\bm{\theta}}_{t}\sim\mathcal{N}\left({\bm{\theta}}|{\bm{\mu}}_{t},\Sigma_{t}\right)$,
the feasible update reads

$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}-\beta\Sigma_{t+1}\left[N\bm{g}\left({\bm{\theta}}_{t}\right)+{\bm{\lambda}}{\bm{\mu}}\right],$
(78) $\displaystyle\Sigma^{-1}_{t+1}$
$\displaystyle=\left(1-\beta\right)\Sigma^{-1}_{t}+\beta\left(NH\left({\bm{\theta}}_{t}\right)+{\bm{\lambda}}I\right).$
(79)

To obtain a form for the update that resembles Newton’s method where the
scaling matrix is estimated online, (Khan et al., 2018b) defines
$S_{t}=\left(\Sigma^{-1}_{t}-{\bm{\lambda}}I\right)/N$ and conversely
$\Sigma_{t}=\left(N\left(S_{t}+{\bm{\lambda}}I/N\right)\right)^{-1}$, and
write the final form of the VON update

$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}-\beta\left(S_{t+1}+\frac{{\bm{\lambda}}}{N}I\right)^{-1}\left(g\left({\bm{\theta}}_{t}\right)+\frac{{\bm{\lambda}}}{N}{\bm{\mu}}_{t}\right),$
(80) $\displaystyle S_{t+1}$ $\displaystyle=\left(1-\beta\right)S_{t}+\beta
H\left({\bm{\theta}}_{t}\right).$ (81)

Similarly, for a diagonal covariance matrix (thus under a mean-field
assumption), with
$\sigma^{2}_{t}=\left[N\left(\bm{s}_{t}+\frac{{\bm{\lambda}}}{N}\right)\right]^{-1}=\left[N\left(\bm{s}_{t}+\tilde{{\bm{\lambda}}}\right)\right]^{-1}$
and
${\bm{\theta}}_{t}\sim\mathcal{N}\left({\bm{\theta}}|{\bm{\mu}}_{t},\text{diag}\left({\sigma^{2}_{t}}\right)\right)$

$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}-\beta\left(\bm{g}\left({\bm{\theta}}_{t}\right)+\frac{{\bm{\lambda}}}{N}{\bm{\mu}}_{t}\right)/\left(s_{t+1}+\frac{{\bm{\lambda}}}{N}\right)\text{,}$
(82) $\displaystyle\bm{s}_{t+1}$
$\displaystyle=\left(1-\beta\right)\bm{s}_{t}+\beta\,\text{diag}\left({H\left({\bm{\theta}}_{t}\right)}\right)\text{,}$
(83)

where the division is intended to be element-wise. Algorithm 8 summarizes the
main elements of VON implementation.

Algorithm 8 Variational Online Newton

repeat

$\sigma=\text{diag}(1/(N(s+\tilde{{\bm{\lambda}}})^{\frac{1}{2}}))$

Simulate ${\bm{\theta}}\sim q_{{\bm{\mu}},\sigma}$

Randomly sample a data example $\mathcal{D}_{i}$

$\hat{\bm{g}}=\nabla_{\bm{\theta}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)\right]$

$\hat{H}=\nabla^{2}_{{\bm{\theta}}{\bm{\theta}}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)\right]$

$\bm{s}=\left(1-\beta\right)\bm{s}+\beta\,\text{diag}\left(H\right)$

${\bm{\mu}}={\bm{\mu}}_{t}-\beta(\hat{\bm{g}}+\tilde{{\bm{\lambda}}}{\bm{\mu}})/(\bm{s}+\tilde{{\bm{\lambda}}})$

until stopping condition is met

In a mini-batch setting for estimating the stochastic gradient, with
$\mathcal{M}$ denoting a mini-batch containing $M$ samples, the stochastic
estimates

$\displaystyle\hat{g}\left({\bm{\theta}}_{t}\right)$
$\displaystyle=\frac{1}{M}\sum_{i\in\mathcal{M}}\nabla_{\bm{\theta}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}_{t}\right)\right]\text{,}$ (84)
$\displaystyle\hat{H}\left({\bm{\theta}}_{t}\right)$
$\displaystyle=\frac{1}{M}\sum_{i\in\mathcal{M}}\nabla^{2}_{{\bm{\theta}}{\bm{\theta}}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}_{t}\right)\right]\text{,}$ (85)

enable the practical implementation of the VON update by replacing $\bm{g}$
and $H$. To make this statement clear, think of $f\left(\bm{\theta}\right)$ as
the typical negative log-likelihood of a sample (as it is an average across
samples), then $\bm{g}$ is the typical gradient for a sample in
${\bm{\theta}}$ and, analogously, $H$ is interpreted as the typical (average)
value of the Hessian evaluated in ${\bm{\theta}}$, resulting when using a
single data point. Stochastic gradient estimation estimates $\bm{g}$ by using
a single observation $\mathcal{D}_{i}$ picked at random as an unbiased
estimate of the actual gradient of
$f\left(\bm{\theta}\right)=-\frac{1}{N}\sum_{i=1}^{N}\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)$,
$\bm{g}=\frac{1}{N}\sum_{i=1}^{n}\nabla_{\bm{\theta}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)\right]$, which would require the
parsing of the entire sample. Analogously, one constructs a stochastic
estimate of the Hessian with one or $M$ observations (the higher $M$ the lower
the variance of the estimator, which is in any case unbiased).

###  7.4 Variational ADAM (VADAM)

The principle of Variational ADAM (VADAM)is that of augmenting the natural
gradient update by incorporating a momentum factor, i.e.,

${\bm{\lambda}}_{t+1}={\bm{\lambda}}_{t}+\beta\tilde{\nabla}_{\bm{\lambda}}\mathcal{L}+\gamma\left({\bm{\lambda}}_{t}-{\bm{\lambda}}_{t-1}\right)$
(86)

which slightly extends the form of the update in eq. (59).

Under a Gaussian variational $q$, (Khan et al., 2018b) expresses the momentum
update as a VON update with momentum and recovers a variational version of an
RMSProp update, to obtain the following updates

$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}-\bar{\beta}_{t}\left[\frac{1}{\sqrt{\bm{s}_{t+1}}+\tilde{{\bm{\lambda}}}}\right]\left(\bm{g}\left(\bm{\theta}\right)+\tilde{{\bm{\lambda}}}{\bm{\mu}}_{t}\right)+\frac{\beta_{t}}{1-\beta_{t}}\left[\frac{\bm{s}_{t}+\tilde{{\bm{\lambda}}}}{\bm{s}_{t+1}+\tilde{{\bm{\lambda}}}}\right]\left({\bm{\mu}}_{t}-{\bm{\mu}}_{t-1}\right),$
(87) $\displaystyle\bm{s}_{t+1}$
$\displaystyle=\left(1-\bar{\beta}_{t}\right)\bm{s}_{t}+\bar{\beta}_{t}\left(\bm{g}\left(\bm{\theta}\right)\right)^{2}\text{,}$
(88)

where $\bar{\beta}_{t}=\beta\frac{1-\gamma_{1}}{1-\gamma_{1}^{t}}$,
$\bar{\gamma}_{t}=\gamma_{1}\left(1-\gamma_{1}^{t-1}\right)\left(1-\gamma_{1}^{t}\right)$
and $\beta$,$\gamma_{1}$ are learning rates. Note that in the above updates
the Hessian is estimated as a squared gradient: details are provided in
Section 7.5. These updates can be implemented and used in their actual form,
yet they correspond to an ADAM-like update. Indeed the above update has the
same form of an adaptive version of Polyak’s heavy ball method. (Wilson et
al., 2017) establishes a relation between the form of eq. (87) and the ADAM
update, and in particular that the ADAM update can be written as an adaptive
version of the Polyak’s heavy ball method. Upon introducing the typical bias
correction terms of ADAM, (Khan et al., 2018b) expresses eq. (87) as an ADAM
update. With respect to a true ADAM update, the model weights are
stochastically sampled from the posterior, resulting in a Variational version
of ADAM (VADAM). For the full derivation, which is quite elaborate and
extensive, refer to (Khan and Nielsen, 2018). Algorithm 9 summarizes the VADAM
approach for a Gaussian variational posterior with diagonal covariance.

Algorithm 9 Variational ADAM

repeat

$\sigma=\text{diag}(1/(N(s+\tilde{{\bm{\lambda}}})^{\frac{1}{2}})$

Simulate ${\bm{\theta}}\sim q_{{\bm{\mu}},\sigma}\left({\bm{\theta}}\right)$

Randomly sample a data example $\mathcal{D}_{i}$

$\hat{g}=\nabla_{\bm{\theta}}\left[-\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}\right)\right]$

$\bm{m}=\gamma_{1}\,\bm{m}+\left(1-\gamma_{1}\right)\left(\hat{\bm{g}}+\tilde{{\bm{\lambda}}}{\bm{\mu}}\right)$,

$\bm{s}=\gamma_{2}\,\bm{s}+\left(1-\gamma_{2}\right)\left(\hat{\bm{g}}\odot\hat{\bm{g}}\right)$

$\bm{m}=\bm{m}/\left(1-\gamma_{1}^{t}\right)$,

$\bm{s}=\bm{s}/\left(1-\gamma_{2}^{t}\right)$

${\bm{\mu}}={\bm{\mu}}-\beta\,\bm{m}/\left(\sqrt{\bm{s}}+\tilde{{\bm{\lambda}}}\right)$

until stopping condition is met

###  7.5 Variational Online Gauss-Newton (VOGN)

In the diagonal VON update, the Hessian drives the update for the scaling
vector $\bm{s}$ which determines the covariance matrix
$\text{diag}\left({\sigma^{2}}\right)$. The Hessian can be negative, a
situation that could turn $\sigma^{2}$ negative, which is meaningless. Instead
of indirectly tackling the issue by using a constrained optimization approach
(which could be difficult to implement), such as a controlled adaptive
learning rate, or model reparametrization, (Khan et al., 2018b) proposes the
use of the Generalized Gauss-Newton approximation for the Hessian:

$\nabla_{{\bm{\theta}}_{j}{\bm{\theta}}_{h}}^{2}f\left(\bm{\theta}\right)\approx\frac{1}{M}\sum_{i\in\mathcal{M}}\left[\nabla{\bm{\theta}}_{j}f_{i}\left(\bm{\theta}\right)\right]^{2}:=\hat{h}_{j}\left(\bm{\theta}\right).$
(89)

This enables a minor but important difference with respect to VON: with an
initial positive value for $\sigma^{2}$, the above approximation will remain
positive leading to valid covariance updates. This provides an algorithmic
advantage over VON as constraints on $\sigma^{2}$ are implicitly satisfied.
The above implementation of the Hessian estimation, within VON, consists in
the Variational Online Gauss-Newton (VOGN)approach (Khan et al., 2018b; Osawa
et al., 2019). The implementation of the above approximation is not immediate
as it requires per-sample gradients. The approximation averages squared
gradients evaluated on a sample-per-sample basis, as opposed to batch-gradient
computation which directly computes the sum of the gradients over mini-batches
(Osawa et al., 2019), which can be seen by comparing eq. (89) with eq. (90)

The gradient-magnitude approximation that makes use of the mini-batch squared
gradient as an approximation for the Hessian,

$\nabla_{{\bm{\theta}}_{j}{\bm{\theta}}_{h}}^{2}f\left(\bm{\theta}\right)\approx\left[\frac{1}{M}\sum_{i\in\mathcal{M}}\nabla{\bm{\theta}}_{j}f_{i}\left(\bm{\theta}\right)\right]^{2}=\left[\hat{\bm{g}}_{j}\left(\bm{\theta}\right)\right]^{2}\text{,}$
(90)

introduces a bias in the Hessian estimation. In fact, increasing the mini-
batch size is not advisable as it introduces more bias. Based on the above
approximation, (Khan et al., 2018b) advances an RMSProp version of the VON
update.

The practical implementation of VOGN is extensively discussed in (Osawa et
al., 2019), where the efficient implementation of the per-sample gradient
computation for certain network layers is discussed: the additional
computations needed to access individual gradients bring the run-time within
2-5 times of that of ADAM. Algorithm 10 summarizes the implementation of the
VOGN optimizer.

Algorithm 10 Variational Online Gauss-Newton

repeat

Randomly sample a mini-batch $\mathcal{M}$ of size $M$

for $s=1,\dots,N_{s}$ do

Simulate ${\bm{\theta}}_{s}={\bm{\mu}}+\varepsilon\sigma$, where
$\sigma=(1/(N(\bm{s}+\tilde{{\bm{\lambda}}})))^{\frac{1}{2}}$,
$\bm{\varepsilon}\sim\mathcal{N}\left(0,I\right)$

for each $i$ in $\mathcal{M}$ do

$\bm{g}_{i}^{\left(s\right)}=\nabla_{{\bm{\theta}}}\log
p\left(\mathcal{D}_{i}|{\bm{\theta}}_{s}\right)$

end for

$\hat{\bm{g}}_{s}=\frac{1}{M}\sum_{i\in\mathcal{M}}\bm{g}_{i}^{\left(s\right)}$

$\hat{\bm{h}}_{s}=\frac{1}{M}\sum_{i\in\mathcal{M}}(\bm{g}_{i}^{\left(s\right)})^{2}$

end for

$\hat{\bm{g}}=\frac{1}{N_{s}}\sum_{s}\hat{\bm{g}}_{s}$

$\hat{\bm{h}}=\frac{1}{N_{s}}\sum_{s}\hat{\bm{h}}_{s}$

$\bm{m}=\beta_{1}\bm{m}+(\hat{\bm{g}}+\tilde{{\bm{\lambda}}}{\bm{\mu}})$,

$\bm{s}=(1-\beta_{2})\bm{m}+\beta_{2}\hat{\bm{h}}$

${\bm{\mu}}={\bm{\mu}}-\beta\,\bm{m}(\bm{s}+\tilde{{\bm{\lambda}}})$

until stopping condition is met

The form of Algorithm 10 slightly differs from that of VOGN/ADAM. The sampling
of the random weights is analogous to that of Algorithm 9 and Algorithm 8, yet
here posterior samples are built over standard-normal random numbers rather
than directly sampling from the multivariate diagonal posterior by the use of
the reparametrization trick. Note the index $i$ referring to the individual
samples in the mini-batch $\mathcal{M}$. While VOGN uses a single sample for
evaluating the stochastic gradients, here $N_{s}$ draws are averaged to reduce
the approximation variance. In particular, the nested for loop computes the
single-observation gradient used for the Hessian approximation, each computed
in the sampled weight vector ${\bm{\theta}}_{s}$. Draw-specific gradients and
Hessian $\hat{\bm{g}}_{s}$ and $\hat{\bm{h}}_{s}$ are thus averaged across
samples (leading to $\hat{\bm{g}}$ and $\hat{\bm{h}}$) and used in the
implementation of the ADAM-like update based on momentum (thus the
hyperparameters $\beta_{1}$, $\beta_{2}$). The pseudo-code in (Osawa et al.,
2019) involves an additional tempering parameter and data-augmentation factor
along with details for the VOGN parallel implementation, to which we refer for
further insights.

(Osawa et al., 2019) furthermore discusses practical implementation aspects
typical in ML such as batch normalization, data augmentation, momentum, and
distributed computing. The feasibility of the VOGN update for large-scale
experiments with big-data sizes and deep network architectures on standard
datasets promotes VOGN as a state-of-the-art method for Bayesian DL. As a
remark, among its limitations, note that VOGN applies to Gaussian variational
posteriors with a diagonal covariance matrix only.

###  7.6 Quasi Black-Box Variational Inference (QBVI)

The BBVI framework of (Ranganath et al., 2014) can benefit from the use of the
natural gradients. In fact, in (Trusheim et al., 2018) natural gradients are
estimated via MC sampling. On the other hand eq. (60) provides an exact
framework for computing natural gradients without relying on sampling methods,
applicable for the wide class of variational posteriors within the Exponential
family, yet model-specific derivations, i.e. the computation of the gradients
and Hessian, are involved. The Quasi Black-Box Variational Inference
(QBVI)approach (Magris et al., 2022c) merges the BBVI setting with the exact
natural gradient computation. QBVI uses eq. (60) to turn the computation of
the natural gradients into Euclidean gradients of the LB, which are computed
by the use of the score estimation, resembling the BBVI framework. On a
general level, the QBVI update estimates the gradient of the LB with respect
to the natural parameters as

$\displaystyle\tilde{\nabla}\mathcal{L}_{\bm{\lambda}}\left({\bm{\lambda}}\right)$
$\displaystyle=\tilde{\nabla}_{\bm{\lambda}}\mathbb{E}_{q_{\bm{\lambda}}}\left[\log\frac{p_{\bm{\eta}}\left(\bm{\theta}\right)}{q_{\bm{\lambda}}\left(\bm{\theta}\right)}\right]+\tilde{\nabla}_{\bm{\lambda}}\mathbb{E}_{q_{\bm{\lambda}}}\left[\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]$ (91)
$\displaystyle=\bm{\eta}-{\bm{\lambda}}+\tilde{\nabla}_{\bm{\lambda}}\mathbb{E}_{q_{\bm{\lambda}}}\left[\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\text{,}$ (92)

which along with a plain SGD lead to the update rule

$\displaystyle{\bm{\lambda}}_{t+1}=\left(1-\beta\right){\bm{\lambda}}_{t}+\beta\left(\bm{\eta}+\tilde{\nabla}_{\bm{\lambda}}\mathbb{E}_{q_{\bm{\lambda}}}\left[\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\right)\text{.}$ (93)

Here the exact computation of the natural gradient is carried out in terms of
eq. (60), so that the QBVI update for a generic variational distribution and
prior (both within the exponential family) reads, for the natural parameters,
as:

$\displaystyle{\bm{\lambda}}_{t+1}$
$\displaystyle=\left(1-\beta\right){\bm{\lambda}}_{t}+\beta\left(\eta+\mathbb{E}_{q_{\bm{\lambda}}}\left[\nabla_{\bm{m}}\left[\log
q_{\bm{\lambda}}\left(\bm{\theta}\right)\right]\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\right)\text{.}$ (94)

Similarly to (Khan and Nielsen, 2018), eq. (92) uses the properties of the
Exponential family distribution for the prior $p$, with natural parameter
$\bm{\eta}$, and $q$ to simplify the first term on the right-side of eq. (91).
This results in the natural-parameter difference $\bm{\eta}-{\bm{\lambda}}$,
avoiding on the first instance a sampling framework for evaluating the
corresponding expectation, i.e. reducing the variance of the estimate for
$\tilde{\nabla}\mathcal{L}\left({\bm{\lambda}}\right)$, regardless of the
estimator used for $\mathbb{E}_{q_{\bm{\lambda}}}\left[\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]$.

(Magris et al., 2022c) focuses on the Gaussian variational case, building the
QBVI update on the NGVI update, but without using the model’s gradient and
Hessian as for VON. Indeed, using equations (61) and (62), (Magris et al.,
2022c) recovers, for a full-covariance posterior, the following updates:

$\displaystyle\Sigma^{-1}_{t+1}$
$\displaystyle=\left(1-\beta\right)\Sigma^{-1}_{t}+\beta\left(\Sigma^{-1}_{0}+\mathbb{E}_{q_{\bm{\lambda}}}\left[\left(\Sigma^{-1}_{t}-\bm{v}_{t}\bm{v}^{\top}_{t}\right)\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\right)\text{,}$ (95)
$\displaystyle{\bm{\mu}}_{t+1}$
$\displaystyle={\bm{\mu}}_{t}+\beta\Sigma_{t+1}\left(\Sigma^{-1}_{0}\left({\bm{\mu}}_{0}-{\bm{\mu}}_{t}\right)+\mathbb{E}{q_{\bm{\lambda}}}\left[\bm{v}_{t}\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\right)\text{,}$ (96)

where $\bm{v}_{t}=\Sigma^{-1}\left({\bm{\theta}}-{\bm{\mu}}_{t}\right)$ and
${\bm{\mu}}_{0},\Sigma_{0}$ denote the mean vector and covariance matrix of
the prior distribution on the model parameter ${\bm{\theta}}$, respectively.
The following naive MC estimator provides a simple approach for tacking the
above expectations

$\displaystyle\mathbb{E}{q_{\bm{\lambda}}}$
$\displaystyle\left[\left(\Sigma^{-1}_{t}-\bm{v}_{t}\bm{v}^{\top}_{t}\right)\log
p\left(y|{\bm{\theta}}\right)\right]$ (97)
$\displaystyle\approx\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\left[(\Sigma^{-1}_{t}-\Sigma^{-1}_{t}\left({\bm{\theta}}_{s}-{\bm{\mu}}_{t}\right)\left({\bm{\theta}}_{s}-{\bm{\mu}}_{t}\right)^{\top}\Sigma^{-1}_{t})\log
p\left(\mathcal{D}|{\bm{\theta}}_{s}\right)\right]\text{,}$ (98)
$\displaystyle\mathbb{E}{q_{\bm{\lambda}}}$ $\displaystyle\left[\bm{v}_{t}\log
p\left(\mathcal{D}|{\bm{\theta}}\right)\right]\approx\frac{1}{N_{s}}\sum_{s=1}^{N_{s}}\left[\Sigma^{-1}_{t}\left({\bm{\theta}}_{s}-{\bm{\mu}}_{t}\right)\log
p\left(\mathcal{D}|{\bm{\theta}}_{s}\right)\right]\text{,}$ (99)

with ${\bm{\theta}}_{s}\sim q_{\bm{\lambda}}$, $s=1,\dots,N_{s}$. Algorithm 11
provides the pseudo-code for the QBVI implementation.

Algorithm 11 Quasi Black-Box Variational Inference

repeat

Simulate ${\bm{\theta}}_{s}\sim q_{{\bm{\lambda}}}$, $s=1,\dots,N_{s}$

$\bm{v}=\Sigma^{-1}\left({\bm{\theta}}-{\bm{\mu}}\right)$

$\hat{\bm{g}}_{\Sigma}=1/N_{s}\sum_{s}\left(\Sigma^{-1}_{t}-\bm{v}\bm{v}^{\top}\right)\log
p\left(\mathcal{D}|{\bm{\theta}}_{s}\right)$

$\hat{\bm{g}}_{\bm{\mu}}=1/N_{s}\sum_{s}\bm{v}\log
p\left(\mathcal{D}|{\bm{\theta}}_{s}\right)$

$\Sigma^{-1}=\left(1-\beta\right)\Sigma^{-1}+\beta\left[\Sigma^{-1}_{0}+\hat{\bm{g}}_{S}\right]$

${\bm{\mu}}={\bm{\mu}}+\beta\Sigma\left[\Sigma^{-1}_{0}\left({\bm{\mu}}_{0}-{\bm{\mu}}\right)+\hat{\bm{g}}_{\bm{\mu}}\right]$

until stopping condition is met

##  8 Variational Inference on manifolds

In this section, we review a class of methods that pursue a theoretically
different approach, i.e., manifold optimization. The major challenge in VI
optimization is that of guaranteeing constraints on the variational parameter.
In a Gaussian, or e.g. an Inverse Wishart, setting, this corresponds to
guaranteeing updates under which the covariance matrix is Symmetric and
Positive Definite (SPD).

We first introduce in general terms the concept and practice of Riemann
optimization. Therefore, we provide an introduction to Riemann manifolds, the
concepts of tangent vectors, tangent spaces, and Riemann gradient to finally
provide a more rigorous discussion of the specific problem of performing valid
covariance updated for Bayesian Inference under a Gaussian variational FFVI
setting. This section addresses the most crucial aspects concerning the
purpose of introducing the Manifold Gaussian Variational Bayes and Exact
Manifold Gaussian Variational Bayes optimizers. As the topic is itself broad
and quite technical, we intentionally provide a descriptive illustration
suitable for a general audience, referring to the specialized literature for
additional details and a rigorous mathematical treatment at the end of the
following section.

###  8.1 Introduction to manifold optimization

Riemann optimization is an alternative to standard SGD that well fits problems
of the kind

$\displaystyle\operatorname*{arg\,min}_{{\bm{\zeta}}\in\mathcal{M}}\mathcal{L}\left({\bm{\zeta}}\right)\text{,}$
(100)

where $\mathcal{L}$ is a real-valued function of some parameter
${\bm{\zeta}}$, defined on a Riemannian manifold $\left(\mathcal{M},g\right)$.
A manifold is a topological space that locally resembles Euclidean space near
each point, in more detail, is a set that can locally be mapped one-to-one to
$\mathbb{R}^{k}$, where $k$ is the dimension of the manifold. $g$ stands for a
metric the manifold is equipped with.

The optimization problem aims at minimizing $\mathcal{L}$ by finding the
parameter ${\bm{\zeta}}\in\mathcal{M}$ that lies on the “smooth surface” of
the Riemannian manifold $\left(\mathcal{M},g\right)$ resembling a constrained
optimization problem requiring the optimum ${\bm{\zeta}}^{*}$ to lie on the
Riemannian manifold, such as a sphere or the SPD set. As with SGD in Euclidean
vector spaces, Riemann optimization is generally tackled with gradient descent
on the surface of the manifold, based on the gradients of $\mathcal{L}$. Yet,
because of the manifold constraint, there are important differences compared
to the standard SGD approach.

The Euclidean vector space $\mathbb{R}^{n}$ can be interpreted as a Riemannian
manifold $\left(\mathbb{R}^{n},g\right)$, with $g$ the common Euclidean
metric, where the usual SGD iteratively updates the parameter ${\bm{\zeta}}$
as

${\bm{\zeta}}_{t+1}={\bm{\zeta}}_{t}+\beta\nabla_{\bm{\zeta}}\mathcal{L}\left({\bm{\zeta}}_{t}\right)\text{,}$
(101)

where

$\nabla_{\bm{\theta}}\mathcal{L}\left({\bm{\zeta}}_{t}\right)=\left.\frac{\partial}{\partial{\bm{\zeta}}}\mathcal{L}\left({\bm{\zeta}}\right)\right|_{{\bm{\zeta}}={\bm{\zeta}}_{t}}\text{.}$
(102)

It is clear that applying the above to a generic non-Euclidean manifold
$\mathcal{M}$ is not trivial as there is no guarantee that
${\bm{\zeta}}_{t+1}$ is a valid update, i.e. that ${\bm{\zeta}}_{t+1}$ lies in
$\mathcal{M}$. Consider an optimization problem where the parameter
${\bm{\zeta}}=\left(x,y,z\right)$ is required to lie on a 2-dimensional
spherical manifold of radius 1, embedded in a 3-dimensional ambient space. The
Riemannian manifold is
$\mathcal{M}=\left\\{{\bm{\zeta}}\in\mathbb{R}^{3}:\,||{\bm{\zeta}}||_{2}=1\right\\}$,
with $g$ being the Euclidean metric, and $\mathcal{L}$ corresponding to a
custom loss function for an arbitrary point ${\bm{\zeta}}$ on the sphere
$\mathcal{M}$. Though partial derivatives $\nabla_{{\bm{\zeta}}}\mathcal{L}$
are straightforward to compute or evaluate, e.g. with backpropagation, at the
current parameter value ${\bm{\zeta}}_{t}$, there is no guarantee that the
update rule for the Euclidean space
${\bm{\zeta}}_{t+1}={\bm{\zeta}}_{t}+\beta\nabla_{\bm{\zeta}}\mathcal{L}{{\bm{\zeta}}_{t}}$
would result in an updated parameter lying on sphere $\mathcal{M}$.
Intuitively, on the “curved” surfaces of Riemannian manifolds the updates
should follow the “curved” geodesics instead of straight lines as on familiar
$\mathbb{R}^{n}$ Euclidean spaces. To this end, Riemann Stochastic Gradient
Descent (RSGD)constitutes a manifold generalization of the SGD.

####  8.1.1 Elements of Riemannian manifolds

In $\mathbb{R}^{k}$, a steepest-ascent approach updates the current iterate
${\bm{\zeta}}$ in the direction where the first-order increase of the
objective function $\mathcal{L}$ is most positive. Formally, the update
direction is chosen to be the unit norm vector $\bm{\bm{\eta}}$ that minimizes
the directional derivative

$\text{D}\mathcal{L}\left({\bm{\zeta}}\right)\left[\bm{\bm{\eta}}\right]=\lim_{t\rightarrow
0}\frac{\mathcal{L}\left({\bm{\zeta}}+t\bm{\eta}\right)-\mathcal{L}\left({\bm{\zeta}}\right)}{t}\text{.}$
(103)

With the domain of $\mathcal{L}$ being the manifold $\mathcal{M}$, the
argument ${\bm{\zeta}}+t\bm{\bm{\bm{\eta}}}$ does not make much sense in
general as $\mathcal{M}$ is not necessarily a vector space. This leads to the
notion of a tangent vector. A possibility for generalizing the directional
derivative is to replace $t\mapsto{\bm{\zeta}}+t\bm{\bm{\eta}}$ by a smooth
curve $\gamma$ on $\mathcal{M}$ passing through ${\bm{\zeta}}$, i.e.
$\gamma\left(0\right)={\bm{\zeta}}$. A smooth mapping
$\gamma:\mathbb{R}\rightarrow\mathcal{M}:\,t\mapsto\gamma\left(t\right)$ is
termed as curve in $\mathcal{M}$. Defining a derivative
$\gamma^{\prime}\left(t\right)$ as
$\gamma^{\prime}\left(t\right):=\lim_{t\rightarrow
0}\frac{\gamma\left({\bm{\zeta}}+t\right)-\gamma\left({\bm{\zeta}}\right)}{t}$
fails on a general manifold as it requires a vector space structure to compute
the difference
$\gamma\left({\bm{\zeta}}+t\right)-\gamma\left({\bm{\zeta}}\right)$, however
for a smooth function $\mathcal{L}$ on $\mathcal{M}$ the function
$\mathcal{L}\circ\gamma:\,t\mapsto\mathcal{L}\left(\gamma\left(t\right)\right)$
is a smooth and well-defined function from $\mathbb{R}$ to $\mathbb{R}$ with a
well-defined classical derivative. To sum up, let ${\bm{\zeta}}$ be a point on
$\mathcal{M}$, $\gamma$ a curve such that $\gamma\left(0\right)={\bm{\zeta}}$
and $\mathcal{F}_{\bm{\zeta}}\left(\mathcal{M}\right)$ is the set of smooth
real-valued functions defined in a neighborhood of ${\bm{\zeta}}$ in
$\mathcal{M}$. The mapping $\dot{\gamma}\left(0\right)$ from