# Tractable Approximate Gaussian Inference for Bayesian Neural Networks

James-A. Goulet  Luong Ha Nguyen  Saeid Amiri

###### Abstract

In this paper, we propose an analytical method allowing for _tractable
approximate Gaussian inference_ (TAGI) in Bayesian neural networks. The method
enables: (1) the analytical inference of the posterior mean vector and
diagonal covariance matrix for weights and bias, (2) the end-to-end treatment
of uncertainty from the input layer to the output, and (3) the online
inference of model parameters using a single observation at a time. The method
proposed has a computational complexity of $\mathcal{O}(n)$ with respect to
the number of parameters $n$, and the tests performed on regression and
classification benchmarks confirm that, for a same network architecture, it
matches the performance of existing methods relying on gradient
backpropagation.

Machine Learning, ICML, Bayesian, Neural Networks, TAGI, Approximate Gaussian
Inference, Tractable

##  1 Introduction

The estimation of weight and bias in neural networks is currently dominated by
approaches employing point estimates when learning model parameters using
gradient backpropagation (Rumelhart et al., 1986). Although these approaches
allow for the state-of-the-art performance in many domains of applications, it
is recognized that they fall short in situations where, for instance, datasets
are small, when the task requires quantifying the uncertainty about the
prediction made, or for continual learning (Ghahramani, 2015; Kendall & Gal,
2017; Farquhar & Gal, 2019). In such situations, the Bayesian approach for
infering the parameters’ posterior probability, is known to be theoretically
better suited than a point estimate. This is in theory because applying exact
Bayesian inference on large neural networks has been considered to be
intractable (Goodfellow et al., 2016).

Many researchers have already proposed approximate inference methods for
_Bayesian Neural Networks_ (BNN). Early on, it was proposed to employ the
_Extended_ and _Unscented Kalman Filter_ to estimate the Gaussian posterior
for feedforward neural networks (Singhal & Wu, 1989; Puskorius & Feldkamp,
1991; Wan & Merwe, 2000). More recently, this Gaussian framework was extended
with the _Cubature filter_ (Arasaratnam & Haykin, 2008), which, like the
extended and unscented methods, is limited by the computational complexity of
$\mathcal{O}(n^{3})$, where $n$ is the number of parameters. In another
approach, MacKay (1992) employed the _Laplace approximation_ to describe the
posterior covariance of parameters. Later, Neal (1995) explored the potential
of BNN using _sampling methods_. In parallel, several researchers have applied
_variational inference_ for estimating the posterior distribution of neural
networks’ parameters (Hinton & Van Camp, 1993; Barber & Bishop, 1998). The
development of moment matching and variational approaches for BNN is still
nowadays an active research area (Hernández-Lobato & Adams, 2015; Blundell et
al., 2015; Louizos & Welling, 2016; Osawa et al., 2019). Recently, the
technique of using Dropout as a Bayesian approximation has received a lot of
attention in the community (Gal & Ghahramani, 2016). For all the recent
methods that are either based on variational approaches or dropout, the
inference of parameters is still treated as an optimization problem relying on
gradient backpropagation.

In this paper, we propose a _tractable approximate Gaussian inference_ method
(TAGI) for Bayesian neural networks. The approach, which does not rely on
backpropagation, allows for an analytical end-to-end treatment of uncertainty
for the weight and bias parameters, as well as in observed covariates and
system responses. The paper is organized as follows: Section 2 first
introduces the _Gaussian multiplication approximation_ (GMA) for propagating
uncertainty in feedforward neural networks, and second, it presents how to
perform tractable inference for the posterior mean vector and diagonal
covariance of weight and bias parameters. Section 3 validates the performance
of the approach on benchmark regression problems and on the MNIST
classification problem.

##  2 Gaussian Approximation for BNN

###  2.1 Gaussian Feedforward Neural Network

Figure 1 details the nomenclature for the feedforward neural network (FNN)
employed in this paper. Let us consider a vector of input covariates
$\bm{X}=[X_{1}~{}X_{2}~{}\ldots~{}X_{\mathtt{X}}]^{\intercal}:\bm{x}\in\mathbb{R}^{{\mathtt{X}}}$
that are described by random variables in order to take into account errors
potentially arising from observation uncertainties, and then suppose we have a
vector of $\mathtt{Y}$ observed system responses
$\bm{Y}=[Y_{1}~{}Y_{2}~{}\ldots~{}Y_{\mathtt{Y}}]^{\intercal}:\bm{y}\in\mathbb{R}^{{\mathtt{Y}}}$.
The relations between observed system responses and its covariates are
described by the observation model

$\bm{y}=\bm{z}^{(\mathtt{O})}+\bm{v},$ (1)

where the vector of hidden variables $\bm{z}^{(\mathtt{O})}$ corresponds to
the output layer of a neural network on which observation errors
$\bm{v}:\bm{V}\sim\mathcal{N}(\bm{v};\mathbf{0},\mathbf{\Sigma}_{\bm{V}})$ are
added. In common cases, $\mathbf{\Sigma}_{\bm{V}}$ is a diagonal covariance
matrix assuming that observation errors are independent from each other. We
can model the relations between covariates $\bm{x}$ and output hidden
variables $\bm{z}^{(\mathtt{O})}$ using a feedforward neural network
consisting of $\mathtt{L}$ hidden layers each having $\mathtt{A}$ activation
units $a^{(j)}_{i}$ and hidden variables $z^{(j)}_{i}$, $\forall
i=\\{1,2,\cdots,\mathtt{A}\\}$, where an activation unit $a^{(j)}_{i}$ is a
non-linear transformation of its associated hidden variable,
$a_{i}^{(j)}=\sigma(z_{i}^{(j)}).$

$x_{1}$$x_{2}$⋮$x_{\mathtt{X}}$$\bm{b}^{(0)}$$z^{(1)}_{1}$$z^{(1)}_{2}$$\vdots$$z^{(1)}_{\mathtt{A}}$$\bm{b}^{(1)}$$z^{(2)}_{1}$$z^{(2)}_{2}$⋮$z^{(2)}_{\mathtt{A}}$$\bm{b}^{(\cdots)}$$\cdots$$\cdots$⋮$\cdots$$\bm{b}^{(\mathtt{L}-1)}$$z^{(\mathtt{L})}_{1}$$z^{(\mathtt{L})}_{2}$⋮$z^{(\mathtt{L})}_{\mathtt{A}}$$\bm{b}^{(\mathtt{L})}$$z_{1}^{(\mathtt{O})}$$z_{2}^{(\mathtt{O})}$⋮$z_{\mathtt{Y}}^{(\mathtt{O})}$$y_{1}$$y_{2}$⋮$y_{\mathtt{Y}}$$w^{(1)}_{1,1}$$w^{(1)}_{1,1}$$w^{(1)}_{1,2}$$w^{(1)}_{2,1}$$w^{(1)}_{1,i}$$w^{(1)}_{i,1}$$w^{(1)}_{1,\mathtt{A}}$$w^{(1)}_{\mathtt{A},1}$$b^{(1)}_{1}$
(a)

$\bm{x}$$\,\bm{z}^{\\!(1)}$$\,\bm{z}^{\\!(2)}$$\cdots$$\,\bm{z}^{\\!(\mathtt{L})}$$\,\bm{z}^{\\!(\mathtt{O})}$$\bm{y}$$\\!\bm{\theta}^{(0)}\\!$$\\!\bm{\theta}^{(1)}\\!$$\\!\bm{\theta}^{(2)}\\!$$\\!\\!\bm{\theta}^{(\mathtt{L}\text{-}1)}\\!\\!$$\\!\bm{\theta}^{(\mathtt{L})}\\!$
(b)

Figure 1: Expanded and compact representations of the variable nomenclature
associated with feedforward neural networks.

We go from the input layer containing the covariates $\bm{x}$, to the
$i^{\text{th}}$ hidden variable on the first hidden layer, using an affine
function of $\bm{x}$ so that

$z_{i}^{(1)}=w_{i,1}^{(0)}x_{1}+w_{i,2}^{(0)}x_{2}+\cdots+w_{i,\mathtt{X}}^{(0)}x_{\mathtt{X}}+b^{(0)}_{i}.$
(2)

Equation 2 involves the product of weights $w_{i,j}^{(0)}$ and covariates
$x_{j}$ with an additive bias term $b^{(0)}_{i}$. In the context of a neural
network, the process of learning consists in estimating these weights and
bias. Here we consider that weights and bias parameters are described by
random variables so that our joint prior for
$\\{\bm{X},\bm{W}^{(0)},\bm{B}^{(0)}\\}$ is a multivariate Gaussian.

In Equation 2, we note that the product of Gaussian random variables is not
Gaussian. Despite this, we propose to employ moment generating functions in
order to compute analytically its expected value, its variance, as well as the
covariance between the product of Gaussian random variables and any other
Gaussian random variable.

For instance, let
$\bm{X}=[X_{1}~{}\ldots~{}X_{4}]^{\intercal}\sim\mathcal{N}(\bm{x};\bm{\mu},\mathbf{\Sigma})$
be a generic vector of Gaussian random variables, where $\bm{\mu}$ is the mean
vector and $\mathbf{\Sigma}$ is the covariance matrix, then the following
statements are held,

$\displaystyle\mathbb{E}[X_{1}X_{2}]\\!\\!\\!\\!\\!$ $\displaystyle=$
$\displaystyle\\!\\!\\!\\!\\!\mu_{1}\mu_{2}+\text{cov}(X_{1},X_{2}),$ (3)
$\displaystyle\text{cov}(X_{3},X_{1}X_{2})\\!\\!\\!\\!\\!$ $\displaystyle=$
$\displaystyle\\!\\!\\!\\!\\!\text{cov}(X_{1},\\!X_{3})\mu_{2}\\!+\\!\text{cov}(X_{2},\\!X_{3})\mu_{1},$
(4) $\displaystyle\text{cov}(X_{1}X_{2},X_{3}X_{4})\\!\\!\\!\\!\\!$
$\displaystyle=$
$\displaystyle\\!\\!\\!\\!\\!\text{cov}(X_{1},X_{3})\text{cov}(X_{2},X_{4})$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{1},X_{4})\text{cov}(X_{2},X_{3})$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{1},X_{3})\mu_{2}\mu_{4}+\text{cov}(X_{1},X_{4})\mu_{2}\mu_{3}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{2},X_{3})\mu_{1}\mu_{4}+\text{cov}(X_{2},X_{4})\mu_{1}\mu_{3},$
$\displaystyle\text{var}(X_{1}X_{2})\\!\\!\\!$ $\displaystyle=$
$\displaystyle\\!\\!\\!\sigma_{1}^{2}\sigma_{2}^{2}+\text{cov}(X_{1},X_{2})^{2}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+2\text{cov}(X_{1},X_{2})\mu_{1}\mu_{2}+\sigma_{1}^{2}\mu_{2}^{2}+\sigma_{2}^{2}\mu_{1}^{2}.$

The development of statements (3-2.1) is presented in the Appendix A. In this
paper, we define the _Gaussian multiplication approximation_ as the
approximation of the probability density function (PDF) for any product term
$X_{i}X_{j}$ by a Gaussian whose first two moments are defined by Equations
(3-2.1). With this approximation, we can now employ $X_{i}X_{j}$ along with
the random state vector ${\bm{X}}$ in affine functions. It allows propagating
the uncertainty from the input covariates and prior knowledge on weight and
bias parameter through a FNN.

Figure 2a illustrates the passage from the activation units $\bm{A}^{(j)}$, to
a subsequent hidden unit $\bm{Z}^{(j+1)}_{i}$. Figure 2b-d compare the true
theoretical PDF with those obtained using the GMA for different number of
activation unit $\mathtt{A}$, under the assumption that both
$A^{(j)}_{k}\sim\mathcal{N}(a^{(j)}_{k};\mu,\sigma^{2})$,
$W^{(j)}_{i,k}\sim\mathcal{N}(w^{(j)}_{i,k};\mu,\sigma^{2})$ and
$b^{(j)}_{i}=0$. This shows that even if the GMA is a crude approximation for
the Chi-square PDF resulting from the product of two Gaussians, when several
of these product terms are added, the result quickly tends to a Gaussian PDF
as expected from the central limit theorem.

$a^{(1)}_{1}$$a^{(1)}_{2}$$\vdots$$a^{(1)}_{\mathtt{A}}$$z^{(2)}_{1}$

(a)

(b)

(c)  (d)

Figure 2: Illustration of the effect of the GMA on the PDF of a hidden unit
${Z}^{(j+1)}_{i}$ as a function of the number of activation units per layer
$\mathtt{A}$. The blue curve represents the true PDF and the red one, the PDF
resulting from the GMA approximation.

The passage from the hidden variables $\bm{Z}$, to their corresponding
activation units $\bm{A}$ cannot be done analytically using non-linear
activation functions. In order to work around this difficulty, we propose to
employ functions that are linearized at $\mathbb{E}[\bm{Z}]=\bm{\mu_{Z}}$.
_Linearized activation functions_ $\tilde{\sigma}(\cdot)$ allow calculating
analytically the expected vector $\mathbb{E}[\bm{A}]$, the covariance
$\text{cov}(\bm{A})$, as well as the covariance between activation units and
the weight and bias $\text{cov}(\bm{A},\bm{\theta})$, where
$\bm{\theta}=\\{\bm{W},\bm{B}\\}$. Figure 3 presents an example for the
linearization of a softplus activation function. Note that a linearized
activation function is not equivalent to having a linear activation function
because for each instance of input covariates $\bm{x}_{i}$, the linearization
is done at a different value $\mu_{z}$ which maintains the non-linear
dependency between $\bm{x}_{i}$ and the output $\bm{y}_{i}$. It is known that
this linearization procedure is an approximation of the _change of variable
rule_ that would be required to obtain true theoretical PDF for $f(a)$. It is
also known that more advanced sampling-based approches exist for propagating
uncertainty through non-linear function , e.g. (Arasaratnam & Haykin, 2009).
Nevertheless, we choose to employ the linearization procedure because of its
minimal computational cost which still allows, as it will be shown in §3, to
match the state-of-the-art performance on equivalent neural networks
architectures using backpropagation. Finally, note that the linearization
procedure is compatible with all the common activation functions such as the
relu, tanh, logistic sigmoid, etc.

Figure 3: Examples of linearization for a softplus activation function
${\sigma}(\cdot)$ at the expected value $\mathbb{E}[{Z}]$ for
$Z\sim\mathcal{N}(z;\mu_{Z},0.25^{2})$ where $\mu_{Z}=\\{-1,0,1\\}$.

The transition from the knowledge of the $j^{\text{th}}$ layer’s activation
units to an activation unit on the $j+1$ layer is defined by
$a_{i}^{(j+1)}=\tilde{\sigma}(z_{i}^{(j+1)})=\tilde{\sigma}(w_{i,1}^{(j)}a_{1}^{(j)}+w_{i,2}^{(j)}a_{2}^{(j)}+\cdots+w_{i,\mathtt{A}}^{(j)}a_{\mathtt{A}}^{(j)}+b^{(j)}_{i})$.
Analogously, we go from the last hidden layer to the output layer by following

$z_{i}^{(\mathtt{O})}=w_{i,1}^{(\mathtt{L})}a_{1}^{(\mathtt{L})}+w_{i,2}^{(\mathtt{L})}a_{2}^{(\mathtt{L})}+\cdots+w_{i,\mathtt{A}}^{(\mathtt{L})}a_{\mathtt{A}}^{(\mathtt{L})}+b^{(\mathtt{L})}_{i},$

for all $i\in\\{1,2,\cdots,\mathtt{Y}\\}$. All these steps define what we call
the _approximate Gaussian feedforward neural network_ (AG-FNN), which can be
summarized by the input-output relation

$\\!\left\\{\bm{\mu}_{\bm{Z}}^{(\mathtt{O})},\mathbf{\Sigma}_{\bm{Z}}^{(\mathtt{O})},\mathbf{\Sigma}_{\bm{Z\theta}}^{(\mathtt{O})}\right\\}=\text{AG-
FNN}\left(\bm{\mu}_{\bm{X}},\mathbf{\Sigma}_{\bm{X}},\bm{\mu}_{\bm{\theta}},\mathbf{\Sigma}_{\bm{\theta}}\right).$
(7)

For the _regression setup_ where $\bm{y}\in\mathbb{R}^{\mathtt{\mathtt{Y}}}$,
the observed model output is directly defined by the last layer as described
in Equation 1 so that
$\mathcal{N}(\bm{y};\bm{\mu}_{\bm{Y}},\mathbf{\Sigma}_{\bm{Y}})$, where,
$\bm{\mu}_{\bm{Y}}=\bm{\mu}_{\bm{Z}}^{(\mathtt{O})}$, and
$\mathbf{\Sigma}_{\bm{Y}}=\mathbf{\Sigma}_{\bm{Z}}^{(\mathtt{O})}+\mathbf{\Sigma}_{\bm{V}}$.

For the _classification setup_ , we need to convert the output
$\bm{y}\in\mathbb{R}^{\mathtt{\mathtt{Y}}}$ into a class observation
${y}^{(\mathtt{C})}\in\\{1,2,\cdots,\mathtt{K}\\}$, as depicted in Figure 4.

$y_{1}$$y_{2}$$\cdots$$y_{\mathtt{Y}}$$\,{y}^{\\!(\mathtt{C})}$

(a)

$\bm{y}$$\,{y}^{\\!(\mathtt{C})}$ (b)

Figure 4: Expanded and compact representations of the output layer for the
classification setup.

Note that using the traditional _Softmax_ output layer would not allow for a
closed-form solution in order to propagate and marginalize the uncertainty
associated with the output $\bm{Y}$. Instead, we propose to employ a
_hierarchical binary decomposition_ similar to what was employed by Morin &
Bengio (2005). For that, each class is encoded in a binary tree with
$\mathtt{H}=\lceil\log_{2}(\mathtt{K})\rceil$ layers, and which is defined by
$\mathtt{Y}=\mathtt{K}-1$ hidden states when
$\log_{2}(\mathtt{K})\in\mathbb{Z}^{+}$.

Figure 5: Representation of a 3-layers hierarchical binary decomposition of
classes $y_{ijk}^{(\mathtt{C})}\in\\{1,2,\cdots,8\\}$ using the output layer
variables
$\bm{y}=[y\,y_{0}\,y_{1}\,y_{00}\,y_{01}\,y_{10}\,y_{11}]^{\intercal}\in\mathbb{R}^{7}$.

Figure 5 depicts the hierarchical decomposition for $\mathtt{K}=8$ classes and
$\mathtt{H}=3$ layers, where a given class
$y^{(\mathtt{C})}_{\mathcal{C}}:\mathcal{C}=\\{j,k,l\\}\in\\{0,1\\}^{3}$ is
uniquely described by a set of $\mathtt{H}$ indices. In a binary context where
$\mathtt{H}=1$, we can transform a regression problem into a probability for a
class $y^{\mathtt{C}}_{i},~{}i\in\\{0,1\\}$ by using

$p(y^{\mathtt{C}}_{i}|y)=\Phi\\!\left(\\!(-1)^{i}\frac{y}{\alpha}\right),$

where, $\Phi(\cdot)$ denotes the standard normal CDF, and
$\alpha\in\mathbb{R}^{+}$ is a scaling factor ensuring the compatibility with
the output layer that is intended to return values in the interval $[-1,1]$.
In the general case with $\mathtt{K}$-classes, the conditional probability of
a class given the output values $\bm{y}$ is defined by

$p(y^{(\mathtt{C})}_{\mathcal{C}}|\bm{y})=\prod_{h=1}^{\mathtt{H}}\Phi\\!\left(\\!(-1)^{\mathcal{C}_{h}}\frac{[\bm{y}_{\mathcal{C}}]_{h}}{\alpha}\\!\right),$

where
$\bm{y}_{\mathcal{C}}=[y\,y_{j}\,y_{jk}\,y_{jkl}\,\cdots]^{\intercal}\in\mathbb{R}^{\mathtt{H}}$.
For the example in Figure 5 where $\mathtt{H}=3$ layers, it simplifies to
$\bm{y}_{\mathcal{C}}=[y\,y_{j}\,y_{jk}]^{\intercal}$, so that

$p(y^{(\mathtt{C})}_{\\{ijk\\}}|\bm{y})=\Phi\\!\left(\\!(-1)^{i}\frac{y}{\alpha}\\!\right)\cdot\Phi\\!\left(\\!(-1)^{j}\frac{y_{j}}{\alpha}\\!\right)\cdot\Phi\\!\left(\\!(-1)^{k}\frac{y_{jk}}{\alpha}\\!\right).$

For the special case where
$\bm{Y}_{\mathcal{C}}|\mathcal{D}\sim\mathcal{N}(\bm{\mu}_{\bm{Y}_{\mathcal{C}}},\text{diag}(\bm{\sigma}_{\bm{Y}_{\mathcal{C}}}^{2}))$
follows a Gaussian distribution with diagonal covariance, we can employ the
development found in Rasmussen & Williams (2006) in order to obtain a closed-
form solution to marginalize the output layer’s uncertainty,

$\begin{array}[]{ll}p(y^{(\mathtt{C})}_{\mathcal{C}}|\mathcal{D})&=\displaystyle\int
p(y^{(\mathtt{C})}_{\mathcal{C}}|\bm{y}_{\mathcal{C}})\cdot
f(\bm{y}_{\mathcal{C}}|\mathcal{D})d\bm{y}\\\\[-1.0pt]
&\displaystyle=\prod_{h=1}^{\mathtt{H}}\Phi\\!\left(\\!(-1)^{\mathcal{C}_{h}}\frac{[\bm{\mu}_{\bm{Y}_{\mathcal{C}}}]_{h}}{\sqrt{\alpha^{2}+[\bm{\sigma}_{\bm{Y}_{\mathcal{C}}}^{2}]_{h}}}\\!\right),\end{array}$
(8)

where $\alpha$ is a scaling factor that enable working with either normalized
or unnormalized data. Note that in the case where number of classes
$\mathtt{K}$ does not correspond to an integer to the power $2$, it is
required to normalize the marginal probabilities obtained in Equation 8 in
order to account for the unused leaves from the binary tree. During the
training phase where we infer the network’s parameters from observations
$y^{(\mathtt{C})}\in\\{1,2,\cdots,\mathtt{K}\\}$, we convert each class into a
$\mathtt{H}$-component vector $\bm{y}_{\mathcal{C}}\in\\{-1,1\\}^{\mathtt{H}}$
so that $[\bm{y}_{\mathcal{C}}]_{i}=(-1)^{\mathcal{C}_{i}}$.

###  2.2 Linear Algebra for AG-FNN

In this section, we describe how the steps involved in the evaluation of
Equation 7 can be performed using linear algebra. Our first hypothesis
supposes that the knowledge for covariates, hidden units, as well as the
weights and bias, is described by Gaussian random variables. We can then
generalize the operations for going from the $\mathtt{A}$ activation units at
a layer $j$ to the subsequent $\mathtt{A}$ hidden units at layer $j+1$ using
linear algebra so that

$\underbrace{\begin{bmatrix}Z_{1}^{(j+1)}\\\ Z_{2}^{(j+1)}\\\ \vdots\\\
Z_{\mathtt{A}}^{(j+1)}\end{bmatrix}}_{\bm{Z}^{(j+1)}}\\!=\\!\underbrace{\begin{bmatrix}W_{1,1}^{(j)}&W_{1,2}^{(j)}&\cdots&W_{1,\mathtt{A}}^{(j)}\\\
W_{2,1}^{(j)}&W_{2,2}^{(j)}&\cdots&W_{2,\mathtt{A}}^{(j)}\\\
\vdots&\vdots&\ddots&\vdots\\\
W_{\mathtt{A},1}^{(j)}&W_{\mathtt{A},2}^{(j)}&\cdots&W_{\mathtt{A},\mathtt{A}}^{(j)}\\\
\end{bmatrix}}_{\bm{W}^{(j)}}\\!\times\\!\underbrace{\begin{bmatrix}A_{1}^{(j)}\\\
A_{2}^{(j)}\\\ \vdots\\\
A_{\mathtt{A}}^{(j)}\end{bmatrix}}_{\bm{A}^{(j)}}\\!+\\!\underbrace{\begin{bmatrix}B^{(j)}_{1}\\\
B^{(j)}_{2}\\\ \vdots\\\ B^{(j)}_{\mathtt{A}}\end{bmatrix}}_{\bm{B}^{(j)}}.$
(9)

Our prior knowledge for activation units is described by
$\bm{A}^{(j)}\sim\mathcal{N}(\bm{a}^{(j)};\bm{\mu}^{(j)}_{\bm{A}},\mathbf{\Sigma}^{(j)}_{\bm{A}})$
as well as by the covariance $\text{cov}(\bm{\theta},\bm{A}^{(j)})$ between
the activation units and the vector $\bm{\theta}\in\mathbb{R}^{\mathtt{P}}$
containing all the weights and bias parameters defined for all the layers in
the network. We can re-write Equation 9 by breaking down the matrix-vector
product $\bm{W}\times\bm{A}$, into an operation-wise equivalent vector
$(\bm{W\\!\\!A})\in\mathbb{R}^{\mathtt{A}^{2}\times 1}$. We can then employ
Equation 3 in order to compute the expected vector
$\bm{\mu}^{(j)}_{\bm{W\\!\\!A}}\equiv\mathbb{E}[(\bm{W\\!\\!A})^{(j)}]$,
Equation 2.1-2.1 for the covariance matrix
$\mathbf{\Sigma}^{(j)}_{\bm{W\\!\\!A}}\equiv\text{cov}\left((\bm{W\\!\\!A})^{(j)}\right)\in\mathbb{R}^{\mathtt{A}^{2}\times\mathtt{A}^{2}}$
associated with the product terms
$(\bm{W\\!\\!A})^{(j)}\in\mathbb{R}^{\mathtt{A}^{2}}$, and Equation 4 for the
covariance matrix
$\mathbf{\Sigma}^{(j)}_{\bm{W\\!\\!A\theta}}\equiv\text{cov}\left((\bm{W\\!\\!A})^{(j)},\bm{\theta}\right)\in\mathbb{R}^{\mathtt{A}^{2}\times\mathtt{P}}$.
We then introduce two new deterministic matrices
$\mathbf{F}_{\bm{wa}}^{(j)}\in\\{0,1\\}^{\mathtt{A}\times\mathtt{A}^{2}}$ and
$\mathbf{F}_{\bm{b}}^{(j)}\in\\{0,1\\}^{\mathtt{A}\times\mathtt{A}}$, which
allow rewriting Equation 9 as a system of linear equations involving the
product-terms vector $(\bm{W\\!\\!A})$,

$\bm{Z}^{(j+1)}=\mathbf{F}_{\bm{wa}}^{(j)}(\bm{W\\!\\!A})^{(j)}+\mathbf{F}_{\bm{b}}^{(j)}\bm{B}^{(j)}.$
(10)

Note that $\mathbf{F}_{\bm{wa}}^{(j)}$ and $\mathbf{F}_{\bm{b}}^{(j)}$ are
non-unique as their specific definition depend on the ordering of variables in
the problem. An exemple of structure for these matrices is presented in
Appendix B. Using the properties of linear functions of Gaussian random
variables, we obtain

$\\!\\!\begin{array}[]{rrcl}\bm{\mu}^{(j+1)}_{\bm{Z}}\equiv&\\!\\!\\!\\!\\!\mathbb{E}[\bm{Z}^{(j+1)}]&\\!\\!\\!\\!\\!=\\!\\!\\!\\!\\!&\mathbf{F}_{\bm{wa}}^{(j)}\bm{\mu}^{(j)}_{\bm{W\\!\\!A}}\\!\\!+\\!\mathbf{F}_{\bm{b}}^{(j)}\\!\bm{\mu}^{(j)}_{\bm{B}},\\\\[4.0pt]
\mathbf{\Sigma}^{(j+1)}_{\bm{Z}}\equiv&\\!\\!\\!\\!\\!\text{cov}(\bm{Z}^{(j+1)})&\\!\\!\\!\\!\\!=\\!\\!\\!\\!\\!&\mathbf{F}_{\bm{wa}}^{(j)}\mathbf{\Sigma}^{(j)}_{\bm{W\\!\\!A}}\mathbf{F}_{\bm{wa}}^{(j)\intercal}\\\\[4.0pt]
\\!&\\!&\\!\\!\\!&\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\mathbf{F}_{\bm{b}}^{(j)}\mathbf{\Sigma}^{(j)}_{\bm{B}}\mathbf{F}_{\bm{b}}^{(j)\intercal}\\\\[4.0pt]
&&\\!\\!\\!&\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+2\mathbf{F}_{\bm{wa}}^{(j)}\text{cov}(\bm{W\\!\\!A}^{(j)}\\!,\\!\bm{B}^{(j)})\mathbf{F}_{\bm{b}}^{(j)\intercal},\\\\[4.0pt]
\mathbf{\Sigma}^{(j+1)}_{\bm{Z\theta}}\equiv&\\!\\!\\!\\!\text{cov}(\bm{Z}^{(j+1)}\\!\\!,\bm{\theta})&\\!\\!\\!\\!\\!=\\!\\!\\!\\!\\!&\mathbf{F}_{\bm{wa}}^{(j)}\mathbf{\Sigma}^{(j)}_{\bm{W\\!\\!A\theta}}+\mathbf{\Sigma}^{(j)}_{\bm{B\theta}},\end{array}$
(11)

where
$\mathbf{\Sigma}^{(j)}_{\bm{B\theta}}\equiv\text{cov}(\bm{B}^{(j)},\bm{\theta})\in\mathbb{R}^{\mathtt{A}\times\mathtt{P}}$
is the covariance between the bias parameters from the $j^{\text{th}}$ layer
and all the other parameters. In order to apply the linearized activation
function $\bm{A}^{(j+1)}=\tilde{\sigma}(\bm{Z}^{(j+1)})$,

$\bm{A}^{(j+1)}=\mathbf{J}^{(j+1)}\left(\bm{Z}^{(j+1)}-\bm{\mu}^{(j+1)}_{\bm{Z}}\right)+{\sigma}(\bm{\mu}^{(j+1)}_{\bm{Z}}),$
(12)

we need to define the diagonal Jacobian matrix of the transformation evaluated
at $\bm{\mu}^{(j+1)}_{\bm{Z}}$,
$\mathbf{J}^{(j+1)}=\text{diag}\big{(}\nabla_{\\!\\!\bm{z}}{\sigma}(\bm{\mu}^{(j+1)}_{\bm{Z}})\big{)}$.
Using again the properties of linear functions of Gaussian random variables,
we obtain

$\begin{array}[]{rrcl}\bm{\mu}^{(j+1)}_{\bm{A}}\equiv&\mathbb{E}[\bm{A}^{(j+1)}]&=&\tilde{\sigma}(\bm{\mu}^{(j+1)}_{\bm{Z}}),\\\\[4.0pt]
\mathbf{\Sigma}^{(j+1)}_{\bm{A}}\equiv&\text{cov}(\bm{A}^{(j+1)})&=&\mathbf{J}^{(j+1)}\mathbf{\Sigma}^{(j+1)}_{\bm{Z}}\mathbf{J}^{(j+1)\intercal},\\\\[4.0pt]
\mathbf{\Sigma}^{(j+1)}_{\bm{A\theta}}\equiv&\text{cov}(\bm{A}^{(j+1)},\bm{\theta})&=&\mathbf{J}^{(j+1)}\mathbf{\Sigma}^{(j+1)}_{\bm{Z\theta}}.\end{array}$
(13)

Equations 11 & 13 allow propagating the information about the covariance of
activation units and its dependence on parameters through any pairs of
successive layers. For the input layer, the steps described in Equations 9-13
remain the same excepted that the activation units $\bm{A}^{(j)}$ are replaced
by the covariates
$\bm{X}\sim\mathcal{N}(\bm{x};\bm{\mu}_{\bm{X}},\mathbf{\Sigma}_{\bm{X}})$.

###  2.3 Tractable Approximate Gaussian Inference (TAGI)

Let us assume we have a set of joint observations for covariates and system
responses so that
$\mathcal{D}=\\{\mathcal{D}_{x},\mathcal{D}_{y}\\}=\\{(\bm{x}_{i},\bm{y}_{i}),\forall
i\in\\{1:\mathtt{D}\\}\\}$. Given that our prior knowledge for the neural
network’s parameter is
$f(\bm{\theta})=\mathcal{N}(\bm{\theta};\bm{\mu}_{\bm{\theta}},\mathbf{\Sigma}_{\bm{\theta}})$,
the method presented in §2.1 supposes that the joint PDF
$f(\bm{\theta},\bm{y})$ for parameters $\bm{\theta}$ and observations $\bm{y}$
is Gaussian with mean vector and covariance

$\bm{\mu}=\left[\begin{array}[]{c}\bm{\mu}_{\bm{\theta}}\\\
\bm{\mu}_{\bm{Y}}\end{array}\right],~{}\mathbf{\Sigma}=\left[\begin{array}[]{cc}\mathbf{\Sigma}_{\bm{\theta}}&\mathbf{\Sigma}_{\bm{Y\\!\theta}}^{\intercal}\\\
\mathbf{\Sigma}_{\bm{Y\\!\theta}}&\mathbf{\Sigma}_{\bm{Y}}\end{array}\right].$

The conditional PDF for the vector $\bm{\theta}$ given observations
$\bm{Y}=\bm{y}$ is described by the Gaussian conditional equations
$f(\bm{\theta}|\bm{y})=\mathcal{N}(\bm{\theta};\bm{\mu}_{\bm{\theta}|\bm{y}},\mathbf{\Sigma}_{\bm{\theta}|\bm{y}})$
defined by its conditional mean vector and covariance matrix,

$\begin{array}[]{l}\bm{\mu}_{\bm{\theta}|\bm{y}}=\bm{\mu}_{\bm{\theta}}+\mathbf{\Sigma}_{\bm{Y\\!\theta}}^{\intercal}\mathbf{\Sigma}_{\bm{Y}}^{-1}(\bm{y}-\bm{\mu}_{\bm{Y}})\\\\[4.0pt]
\mathbf{\Sigma}_{\bm{\theta}|\bm{y}}=\mathbf{\Sigma}_{\bm{\theta}}-\mathbf{\Sigma}_{\bm{Y\\!\theta}}^{\intercal}\mathbf{\Sigma}_{\bm{Y}}^{-1}\mathbf{\Sigma}_{\bm{Y\\!\theta}}.\end{array}$
(14)

This 1-step network-wise inference procedure is computationally prohibitive
because the forward propagation of uncertainty depicted in Figure 6a involves
large-sized densely populated matrices, and the inference using Equation 14
again involves full matrices.

$\bm{x}$$\,\bm{z}^{\\!(1)}$$\cdots$$\,\bm{z}^{\\!(\mathtt{L})}$$\,\bm{z}^{\\!(\mathtt{O})}$$\bm{y}$$\\!\bm{\theta}^{(0)}\\!$$\\!\bm{\theta}^{(1)}\\!$$\\!\\!\bm{\theta}^{(\mathtt{L}\text{-}1)}\\!\\!$$\\!\bm{\theta}^{(\mathtt{L})}\\!$$\begin{array}[]{c}f(\bm{x})\\\
f(\bm{\theta})\end{array}$$f(\bm{\theta},\\!\bm{z}^{\\!(1)}\\!)$$\cdots$$f(\bm{\theta},\\!\bm{z}^{\\!(\mathtt{L})}\\!)$$f(\bm{\theta},\\!\bm{z}^{\\!(\mathtt{O})}\\!)$$f(\bm{\theta},\\!\bm{y})$$f(\bm{\theta}|\bm{y})$$f(\bm{y})$
(a)

$\bm{x}$$\,\bm{z}^{\\!(1)}$$\cdots$$\,\bm{z}^{\\!(\mathtt{L})}$$\,\bm{z}^{\\!(\mathtt{O})}$$\bm{y}$$\\!\bm{\theta}^{(0)}\\!$$\\!\bm{\theta}^{(1)}\\!$$\\!\\!\bm{\theta}^{(\mathtt{L}\text{-}1)}\\!\\!$$\\!\bm{\theta}^{(\mathtt{L})}\\!$$\begin{array}[]{c}f(\bm{x})\\\
f(\bm{\theta})\end{array}$$\begin{array}[]{c}f(\bm{x},\\!\bm{z}^{\\!(1)}\\!)\\\
f(\bm{\theta}^{(0)}\\!\\!,\\!\bm{z}^{\\!(1)}\\!)\end{array}$$\cdots$$\begin{array}[]{c}f(\bm{z}^{\\!(\mathtt{L}\text{-}1)}\\!\\!,\\!\bm{z}^{\\!(\mathtt{L})}\\!)\\\
f(\bm{\theta}^{(\mathtt{L}\text{-}1)}\\!\\!,\\!\bm{z}^{\\!(\mathtt{L})}\\!)\end{array}$$\begin{array}[]{c}f(\bm{z}^{\\!(\mathtt{L})}\\!\\!,\\!\bm{z}^{\\!(\mathtt{O})}\\!)\\\
f(\bm{\theta}^{(\mathtt{L})}\\!\\!,\\!\bm{z}^{\\!(\mathtt{O})}\\!)\end{array}$$f(\bm{z}^{\\!(\mathtt{O})}\\!\\!,\\!\bm{y})$$f(\bm{x}|\bm{y})$$f(\bm{z}^{\\!(\mathtt{1})}\\!|\bm{y})$$\cdots$$f(\bm{z}^{\\!(\mathtt{L})}\\!|\bm{y})$$f(\bm{z}^{\\!(\mathtt{O})}\\!|\bm{y})$$f(\bm{y})$$f(\bm{\theta}^{(0)}\\!|\bm{y})$$f(\bm{\theta}^{(1)}\\!|\bm{y})$$f(\bm{\theta}^{(\mathtt{L\text{-}1})}\\!|\bm{y})$$f(\bm{\theta}^{(\mathtt{L})}\\!|\bm{y})$
(b)

Figure 6: Representation of the _forward_ propagation of uncertainty (magenta
arrows) and _inference_ procedures (cyan arrows).

The solution we propose to overcome these challenges is twofold: (1) employ a
diagonal covariance structure for both the parameters $\bm{\theta}$, and
hidden units $\bm{Z}^{(j)}$, and (2) use the inherent conditional independence
of hidden units between layers, that is,
$\bm{Z}^{(j-1)}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\bm{Z}^{(j+1)}|\bm{z}^{(j)}$, in order to
perform recursive layer-wise inference. As depicted in Figure 6b, the first
step consists in inferring the posterior mean vector and diagonal covariance
for the output layer following

$\begin{array}[]{rcl}f(\bm{z}^{(\mathtt{O})}|\bm{y})&=&\mathcal{N}(\bm{z}^{(\mathtt{O})};\bm{\mu}_{\bm{Z}^{(\mathtt{O})}|\bm{y}},\bm{\Sigma}_{\bm{Z}^{(\mathtt{O})}|\bm{y}})\\\\[6.0pt]
\bm{\mu}_{\bm{Z}^{(\mathtt{O})}|\bm{y}}&=&\bm{\mu}_{\bm{Z}^{(\mathtt{O})}}+\mathbf{\Sigma}_{\bm{Y\\!Z}^{(\mathtt{O})}}^{\intercal}\mathbf{\Sigma}_{\bm{Y}}^{-1}(\bm{y}-\bm{\mu}_{\bm{Y}})\\\\[4.0pt]
\mathbf{\Sigma}_{\bm{Z}^{(\mathtt{O})}|\bm{y}}&=&\mathbf{\Sigma}_{\bm{Z}^{(\mathtt{O})}}-\mathbf{\Sigma}_{\bm{Y\\!Z}^{(\mathtt{O})}}^{\intercal}\mathbf{\Sigma}_{\bm{Y}}^{-1}\mathbf{\Sigma}_{\bm{Y\\!Z}^{(\mathtt{O})}}.\end{array}$
(15)

Note that for classification problems, because of the hierarchical formulation
described in §2.1, even if there are $\mathtt{Y}$ classes, only
$\mathtt{H}=\lceil\log_{2}(\mathtt{Y})\rceil$ hidden units from the output
layer are updated for each observation.

In order to perform layer-wise inference for hidden units and parameters, we
define the short-hand notation
$\\{\bm{\theta}^{\text{+}},\bm{Z}^{\text{+}}\\}\equiv\\{\bm{\theta}^{(j+1)},\bm{Z}^{(j+1)}\\}$
and $\\{\bm{\theta},\bm{Z}\\}\equiv\\{\bm{\theta}^{(j)},\bm{Z}^{(j)}\\}$. As
depicted in Figure 6b, the inference for the $j^{\text{th}}$ layer’s weights
$\bm{\theta}$ and hidden units $\bm{Z}$ is done using the Rauch-Tung-Striebel
recursive procedure (Rauch et al., 1965) where

$\begin{array}[]{rcl}f(\bm{z}|\bm{y})&=&\mathcal{N}(\bm{z};\bm{\mu}_{\bm{Z}|\bm{y}},\bm{\Sigma}_{\bm{Z}|\bm{y}})\\\\[4.0pt]
\bm{\mu}_{\bm{Z}|\bm{y}}&=&\bm{\mu}_{\bm{Z}}+\mathbf{J}_{\bm{Z}}\left(\bm{\mu}_{\bm{Z}^{\text{+}}|\bm{y}}-\bm{\mu}_{\bm{Z}^{\text{+}}}\right)\\\\[4.0pt]
\bm{\Sigma}_{\bm{Z}|\bm{y}}&=&\bm{\Sigma}_{\bm{Z}}+\mathbf{J}_{\bm{Z}}\left(\bm{\Sigma}_{\bm{Z}^{\text{+}}|\bm{y}}-\bm{\Sigma}_{\bm{Z}^{\text{+}}}\right)\mathbf{J}_{\bm{Z}}^{\intercal}\\\\[4.0pt]
\mathbf{J}_{\bm{Z}}&=&\bm{\Sigma}_{\bm{Z}\\!\bm{Z}^{\text{+}}}\bm{\Sigma}_{\bm{Z}^{\text{+}}}^{-1},\end{array}$
(16)
$\begin{array}[]{rcl}f(\bm{\theta}|\bm{y})&=&\mathcal{N}(\bm{\theta};\bm{\mu}_{\bm{\theta}|\bm{y}},\bm{\Sigma}_{\bm{\theta}|\bm{y}})\\\\[4.0pt]
\bm{\mu}_{\bm{\theta}|\bm{y}}&=&\bm{\mu}_{\bm{\theta}}+\mathbf{J}_{\bm{\theta}}\left(\bm{\mu}_{\bm{Z}^{\text{+}}|\bm{y}}-\bm{\mu}_{\bm{Z}^{\text{+}}}\right)\\\\[4.0pt]
\bm{\Sigma}_{\bm{\theta}|\bm{y}}&=&\bm{\Sigma}_{\bm{\theta}}+\mathbf{J}_{\bm{\theta}}\left(\bm{\Sigma}_{\bm{Z}^{\text{+}}|\bm{y}}-\bm{\Sigma}_{\bm{Z}^{\text{+}}}\right)\mathbf{J}_{\bm{\theta}}^{\intercal}\\\\[4.0pt]
\mathbf{J}_{\bm{\theta}}&=&\bm{\Sigma}_{\bm{\theta}\bm{Z}^{\text{+}}}\bm{\Sigma}_{\bm{Z}^{\text{+}}}^{-1}.\end{array}$
(17)

The key aspect of this approach is that for each layer, we only need to store
the mean vectors $\\{\bm{\mu}_{\bm{\theta}},\bm{\mu}_{\bm{Z}}\\}$ and the
covariances
$\\{\bm{\Sigma}_{\bm{\theta}},\bm{\Sigma}_{\bm{Z}},\bm{\Sigma}_{\bm{\theta}\bm{Z}^{\text{+}}},\bm{\Sigma}_{\bm{Z}\\!\bm{Z}^{\text{+}}}\\}$,
where in addition to the relations already given in §2.2,

$\bm{\Sigma}_{\bm{Z}\\!\bm{Z}^{\text{+}}}=\mathbf{F}_{\bm{wa}}^{(j)}\text{cov}(\bm{W\\!\\!A}^{(j)}\\!,\\!\bm{Z}^{(j)})+\mathbf{F}_{\bm{b}}^{(j)}\text{cov}(\bm{B}^{(j)}\\!,\\!\bm{Z}^{(j)}).$

With a diagonal covariance structure for both $\bm{Z}$ and $\bm{\theta}$, the
covariance matrices defining each layer contain at most
$\mathtt{A}^{2}+\mathtt{A}$ non-zero terms, i.e., the number of weights
($\mathtt{A}^{2}$) and bias ($\mathtt{A}$) per layer; Because of the diagonal
structures of covariance matrices, equations 15–17 have a computational
complexity $\mathcal{O}(\mathtt{A}^{2})$, which scales linearly with the
number of hidden layers $\mathtt{L}$.

###  2.4 Hyper-parameter estimation

There are typically tens of thousands, if not millions, of parameters in
$\bm{\theta}$, for which we typically have little or no prior information for
defining the hyper-parameters
${\bm{\eta}}^{(0)}=\\{\bm{\mu}_{\bm{\theta}}^{(0)},\mathbf{\Sigma}_{\bm{\theta}}^{(0)}\\}$.
In the case where we have small datasets, the weakly informative prior
combined with limited data will lead to a weakly informative posterior. One
solution to go around this difficulty while avoiding overfitting is to learn
the model parameters over multiple epochs, $\mathtt{E}>1$, using a training
$\mathcal{D}_{\mathtt{T}}$ and validation set $\mathcal{D}_{\mathtt{V}}$.
Here, we propose to employ the posterior’s hyper-parameter values at the
$i^{th}$ iteration
${\bm{\eta}}^{(i)}=\\{\bm{\mu}_{\bm{\theta}|\mathcal{D}_{\mathtt{T}}}^{(i)},\mathbf{\Sigma}_{\bm{\theta}|\mathcal{D}_{\mathtt{T}}}^{(i)}\\}$
and use them as the prior’s hyper-parameters at the next iteration $i+1$. This
recursive procedure is stopped when the marginal likelihood
$f(\mathcal{D}_{y,\mathtt{V}}|\mathcal{D}_{x,\mathtt{V}},{\bm{\eta}}^{(i)})\,=\,\mathcal{N}\big{(}\mathcal{D}_{\mathtt{V}};\bm{\mu}_{\bm{y}_{\mathtt{v}}|\mathcal{D}_{\mathtt{T}}}^{(i)},\mathbf{\Sigma}_{\bm{y}_{\mathtt{v}}|\mathcal{D}_{\mathtt{T}}}^{(i)}\big{)}$
for the validation set $\mathcal{D}_{\mathtt{V}}$, has reached its maximal
value. This procedure is analogous to the empirical Bayes approach (Efron,
2012) where the prior knowledge’s hyper-parameters are learnt through the
maximization problem

$\hat{\bm{\eta}}=\underset{\bm{\eta}}{\arg\,\max}~{}\int
f(\mathcal{D}_{y,\mathtt{V}}|\mathcal{D}_{x,\mathtt{V}},\bm{\theta})\cdot
f(\bm{\theta}|\bm{\eta})d\bm{\theta}.$ (18)

Note that unlike in Equation 18 where the maximization is explicit, in our
case the maximization is implicitly performed by updating over multiple
epochs.

##  3 Experiments

###  3.1 1D toy problem

We apply TAGI to the 1D regression problem $y=x^{3}+v$,
$v:V\sim\mathcal{N}(0,9)$ taken from (Hernández-Lobato & Adams, 2015), using
an AG-FNN having a single hidden layer with 100 units, and a ReLU activation
function. The objective of this case study is to showcase how TAGI can be
applied on small datasets ($\mathtt{D}=20$ points), and to compare the results
obtained by considering either diagonal or full covariance matrices. In this
example, the inference is performed using one observation at a time, where
both the covariates $x$ and observations $y$ were normalized in the range
$[-1,1]$. The optimal number of epochs is identified from a validation set
$\mathcal{D}_{\mathtt{V}}$ consisting in 20 additional points. The prior
covariance for bias is initialized to
$\Sigma_{\bm{B}}^{0}=0.01\cdot\mathbf{I}$, and for weights
$\Sigma_{\bm{W}}^{0}$, by multiplying the Xaviers’s approach (Glorot & Bengio,
2010) by a factor 0.25. The prior mean vector is randomly sampled from
$\bm{\mu}_{\bm{\theta}}^{0}\sim\mathcal{N}(\mathbf{0},\Sigma_{\bm{\theta}}^{0})$.

Figure 7 compares the true function employed to generate the data, with the
AG-FNN (with diagonal covariances) predictions described by their expected
values and $\pm 3\sigma$ confidence regions. We can see in (a) that the prior
predictive obtained before updating with observations ($\mathtt{E}=0$) is
weakly informative, and that the posterior predictive obtained after the first
epoch (b, $\mathtt{E}=1$) is still a poor approximation of the true function.
The log-likelihood reported in (d) for the validation set allows identifying
that the optimal number of epochs is $\mathtt{E}=24$. The log-likelihood
values reported in Figure 7d confirms that employing a training set to
identify the number of epoch would not be able to prevent overfitting as
depicted in (c).

(a)  (b)  (c)  (d)

Figure 7: Application of AG-FNN with diagonal covariances to a toy regression
problem where (a-c) describe the evolution of the predictive distribution with
respect to the number of epochs $\mathtt{E}$. In (d), we compare the training
and validation log-likelihood in order to identify the optimal number of
epochs, i.e., $\mathtt{E}=24$.

(a)

(b)

Figure 8: Representation of the sorted correlation coefficients extracted from
upper-triangular posterior covariance matrices for the parameters
$\bm{\theta}^{(j)}$ and hidden units $\bm{z}^{(j)}$. The left-most graphs
present the correlations for each of the 20 observations within the first
epoch, and the center graphs present the correlations at the end of each
epoch.

In a second experiment, we now apply TAGI to the same dataset while
considering the full covariance matrices for $\bm{\theta}^{(j)}$ and
$\bm{z}^{(j)}$. Two networks are studied: $\mathtt{L}\\!=\\!1$,
$\mathtt{A}\\!=\\!100$ and $\mathtt{L}\\!=\\!2$, $\mathtt{A}\\!=\\!50$. Figure
8 displays the sorted correlation coefficients extracted from the upper-
triangular posterior covariances, either for each observation from the first
epoch, or for the last observation from subsequent epochs. In Figure 8a, for a
single hidden layer of 100 units, we can see that the correlation is close to
zero for most pairs of parameters and hidden units. In Figure 8b, for a
network made of two hidden layers of 50 units, we can see that the correlation
is again negligible for most pairs of parameters, and for the hidden units
from the first layer, $\bm{z}^{(1)}$. For the hidden units on the second
layer, $\bm{z}^{(2)}$, the correlation is up to $\approx 0.5$ during the
firsts epochs. This positive correlation on the second hidden layer is caused
by the forward propagation of uncertainty from $\bm{z}^{(1)}\to\bm{z}^{(2)}$.
This positive correlation is not present while going from $x\to\bm{z}^{(1)}$
because here $x$ is treated as a constant. We can see in Figure 8b that as the
number of epochs increases and the parameters’ posterior allow capturing the
pattern in the data, the positive correlation between the pairs of hidden
units in $\bm{z}^{(2)}$ vanishes.

These results provide empirical evidence explaining why considering only
diagonal covariance matrices may lead to results that are comparable to those
obtained while considering the full covariance structure.

Table 1: Comparison TAGI’s results with PBP (Hernández-Lobato & Adams, 2015),
VMG (Louizos & Welling, 2016), and MC-dropout (Gal & Ghahramani, 2016), for
$\mathtt{E}=40$ epochs (Rank legend:  first, second).

| Root mean square error (RMSE) | Average log-likelihood (LL)
---|---|---
Datasets | PBP | VMG | MC-Dropout | TAGI | PBP | VMG | MC-Dropout | TAGI
Boston |  3.01$\pm$0.18 | 2.70$\pm$0.13 | 2.97$\pm$0.85 |  3.02$\pm$0.83 | -2.56$\pm$0.12 | -2.46$\pm$0.09 | -2.46$\pm$0.25 | -2.56$\pm$0.33
Concrete |  5.67$\pm$0.09 | 4.89$\pm$0.12 | 5.23$\pm$0.53 |  5.84$\pm$0.57 |  -3.14$\pm$0.11 | -3.01$\pm$0.03 | -3.04$\pm$0.09 |  -3.19$\pm$0.10
Energy |  1.80$\pm$0.05 | 0.54$\pm$0.02 |  1.66$\pm$0.19 | 1.54$\pm$0.16 |  -2.04$\pm$0.02 | -1.06$\pm$0.03 |  -1.99$\pm$0.09 | -1.87$\pm$0.11
Kin8nm |  0.10$\pm$0.00 |  0.10$\pm$0.00 |  0.10$\pm$0.00 |  0.10$\pm$4E-3 |  0.90$\pm$0.01 | 1.10$\pm$0.01 | 0.95$\pm$0.03 |  0.86$\pm$0.04
Naval |  0.01$\pm$0.00 | 0.00$\pm$0.00 |  0.01$\pm$0.00 |  7E-3$\pm$5E-4 | 3.73$\pm$0.01 |  2.46$\pm$0.12 | 3.80$\pm$0.05 |  3.27$\pm$0.20
Power |  4.12$\pm$0.03 | 4.04$\pm$0.04 | 4.02$\pm$0.18 |  4.13$\pm$0.15 |  -2.84$\pm$0.01 | -2.82$\pm$0.01 | -2.80$\pm$0.05 |  -2.84$\pm$0.04
Protein |  4.73$\pm$0.01 | 4.13$\pm$0.02 | 4.36$\pm$0.04 |  4.64$\pm$0.15 |  -2.97$\pm$0.00 | -2.84$\pm$0.00 | -2.89$\pm$0.01 |  -2.96$\pm$6E-3
Wine |  0.64$\pm$0.01 | 0.63$\pm$0.01 | 0.62$\pm$0.04 | 0.63$\pm$0.04 |  -0.97$\pm$0.01 | -0.95$\pm$0.01 | -0.93$\pm$0.06 |  -0.96$\pm$0.06
Yacht |  1.02$\pm$0.05 | 0.71$\pm$0.05 |  1.11$\pm$0.38 | 0.86$\pm$0.25 |  -1.63$\pm$0.02 | -1.30$\pm$0.02 | -1.55$\pm$0.12 | -1.30$\pm$0.23

###  3.2 Benchmark regression problems

The performance of TAGI is now compared with PBP (Hernández-Lobato & Adams,
2015), VMG (Louizos & Welling, 2016), and MC-dropout (Gal & Ghahramani, 2016)
using benchmark regression datasets. For the purpose of comparison with the
results from other methods, all the datasets are analyzed for a fixed number
of epochs, that is $\mathtt{E}=40$. For all cases, the data is normalized, the
activation function is a ReLU, and the batch size is one, i.e. the inference
is performed using one observation at a time; The prior covariance for bias is
initialized to $\Sigma_{\bm{B}}^{0}=0.01\cdot\mathbf{I}$, and for weights
$\Sigma_{\bm{W}}^{0}$, by multiplying the Xaviers’s approach (Glorot & Bengio,
2010) by a factor 0.25; The initial value for the observation error’s standard
deviation is set to $\sigma_{V}=1$, and this value is optimized using a 5-fold
cross-validation setup.

The results reported in Table 1 indicate that TAGI matches the performance of
existing methods in term of root mean square error (RMSE) and log-likelihood
(LL). Even if VMG displays the best predictive performance, its computational
time is two order of magnitude greater than TAGI, PBP and MC-Dropout (Sun et
al., 2017). The current TAGI’s implementation has a computation time that is
comparable to PBP and MC-Dropout (Gal & Ghahramani, 2016). The timing details
for each dataset is reported in the Appendix C. Because TAGI had to be
implemented from scratch, it is currently not yet fully optimized for
computational efficiency. The same is true for the optimization of hyper-
parameters such as $\sigma_{V}$ which currently rely on a gradient-based
approach. Furthermore, the results reported for TAGI did not employ dropout
(Srivastava et al., 2014). Also, as it is the case for the other methods
reported in Table 1, the number of epochs employed was not optimized.

Note as presented in Appendix C, TAGI is not limited to using ReLU and can
employ any common activation function. Moreover, the results presented in
Appendix C confirm that TAGI is also able to handle deeper architectures,
where having more than one hidden layer leads to a performance improvement in
7 out of 9 dataset.

###  3.3 Application on MNIST

We apply TAGI to the MNIST classification problem (LeCun et al., 1998)
consisting of $\mathtt{D}=70\,000$ $(28\times 28)$ greyscale images for
$\mathtt{K}=10$ classes ($60\,000$ training and $10\,000$ test). Here, we
compare the performance of two AG-FNN configurations, each having
$\mathtt{L}=2$ hidden layers with a number of hidden units equal to
$\mathtt{A}\in\\{100,800\\}$. Each AG-FNN has the same structure for the input
($\mathtt{X}=784$ nodes) and the output layer ($\mathtt{Y}=11$ nodes). The
ReLU activation function is used for the two hidden layers. For each digit,
the vector of covariates $\mathbf{x}_{i}\in(0,1)^{784}$ is assumed to be
deterministic so that $\bm{\mu}_{\bm{X}_{i}}=\mathbf{x}_{i}$ and
$\mathbf{\Sigma}_{\bm{X}_{i}}=\mathbf{0}$. The prior covariance for bias is
initialized to $\Sigma_{\bm{B}}^{0}=0.01\cdot\mathbf{I}$, and by using the
Xaviers’s initialization approach (Glorot & Bengio, 2010) for weights
$\Sigma_{\bm{W}}^{0}$. The prior mean vector is randomly sampled from
$\bm{\mu}_{\bm{\theta}}^{0}\sim\mathcal{N}(\mathbf{0},\Sigma_{\bm{\theta}}^{0})$.
The hyper-parameter associated with the output layer is set to $\alpha=1/3$.
The posterior mean vector as well as the main diagonal of the posterior
covariance are learnt using two setups: (1) a single observation per batch,
that is, $\mathtt{B}=1$, and (2) $10$ observations per batch, $\mathtt{B}=10$.
Each network is evaluated for $\sigma_{V}=\\{0.1,0.2,0.3,0.4\\}$ and the
optimal value for $\sigma_{V}$ is selected using a randomly selected
validation set corresponding to 5% of the training set. The optimal number of
epochs $\mathtt{E}$ is identified using an early-stop procedure evaluated on
the validation set.

Table 2 presents the average test error evaluated on $10\,000$ images for the
different AG-FNN configurations for $\mathtt{E}=1$ epoch, and for the optimal
number of epochs found using early stop. In order to factor in the effect of
random weight initialization, the results reported are the average and
standard deviations from five runs. The performance achieved with respect to
the average classification errors matches the reported state-of-the-art
results of approximately 1.6% for FNNs having a same architecture with 2
layers and 800 hidden units and trained using gradient backpropagation (Simard
et al., 2003; Wan et al., 2013). We can see by comparing the results for 100
and 800 hidden units that increasing the hidden layer’s size improves the
performance both for the first and the optimal number of epochs. Moreover, the
results in Table 2 indicate that TAGI’s classification accuracy is not
significantly affected by the usage of batch sizes greater than one.
Nevertheless, we noticed though our experiments that using large batch sizes
makes the learning phase sensitive to the network initialization as well as
the observation noise parameter $\sigma_{V}$.

Table 2: MNIST test-set average classification error [%] for the first and
last epoch. $\mathtt{A}$: number of hidden units on each layer; $\mathtt{B}$:
number of observations per batch; $\mathtt{E}$: number of epochs; $e$: optimal
number of epoch found using early-stop. The results reported are the average
of five runs along with $\pm$ one standard deviation.

  $\mathtt{B}=1$ $\mathtt{B}=10$ $\mathtt{A}$ $\mathtt{E}=1$ $\mathtt{E}=e$
$e$ $\sigma_{V}$ $\mathtt{E}=1$ $\mathtt{E}=e$ $e$ $\sigma_{V}$ $100$
$3.45\\!\pm\\!0.19$ $2.35\\!\pm\\!0.12$ $12\\!\pm\\!5$ $0.2$
$3.57\\!\pm\\!0.09$ $2.17\\!\pm\\!0.10$ $24\\!\pm\\!5$ $0.4$ $800$
$3.21\\!\pm\\!0.11$ $1.51\\!\pm\\!0.05$ $8\\!\pm\\!1$ $0.1$
$3.26\\!\pm\\!0.10$ $1.53\\!\pm\\!0.05$ $24\\!\pm\\!11$ $0.4$

(a)

(b)

(c)

(d)

(e)

(f)

Figure 9: Leftmost graphs describe the posterior predictive probabilities for
each class from the test set. The height of each region on rightmost graphs
describe the probability of correct (white), incorrect (black), and unknown
(gray) classes as a function of the threshold classification probability
employed $\phi\in(0.1,0.999)$. All cases are for a single epoch, i.e.,
$\mathtt{E}\\!=\\!1$.

We perform a second experiment for an AG-FNN with 2 hidden layers, each having
100 activation units, and where $\sigma_{V}=0.2$. Figure 9 presents on the
leftmost plots, the posterior predictive probability of each class for the
test set, where $\mathtt{D}=\\{0,60,600,6\,000,60\,000\\}$ observations have
been seen during a single epoch, $\mathtt{E}=1$. The rightmost plots present
the probability of correct, incorrect, and unknown classes as a function of
the threshold classification probability employed $\phi\in(0.1,0.999)$. A
_correct class_ correspond to the case where the true label has the highest
probability among all classes and its probability is $\geq\phi$. An _incorrect
class_ corresponds to the case where an incorrect label has the highest
probability among all classes and its probability is $\geq\phi$. An _unknown
class_ correspond to the case where no class has a probability $\geq\phi$.
Figure 9a-e show a gradual transition from uniform posterior predictive
probabilities for $\mathtt{D}\\!=\\!0$ observations, to a low entropy one as
$\mathtt{D}$ increases. This confirms that TAGI is suited for online learning
whereas it can learn continuously using a single observation at a time and
over a single epoch. For reference, Figure 9f shows a how standard neural
network relying on backpropagation is unable to learn from a single epoch.

##  4 Conclusion

The tractable approximate Gaussian inference method proposed in this paper
allows for: (1) the analytical inference of the posterior mean vector and
diagonal covariance matrix for the parameters of Bayesian neural networks, (2)
the end-to-end treatment of uncertainty from the input layer to the output
layer, and (3) the online Bayesian estimation of model parameters using a
single observation at a time. The applications on the regression and
classification datasets validate that the approach matches the performance of
existing methods with respect to computational efficiency and accuracy. TAGI’s
performance and its linear complexity with respect to the number of parameters
makes it a viable alternative to gradient backpropagation. By allowing the
end-to-end treatment of uncertainty and online inference, we foresee that the
approach will enable transformative developments in supervised, unsupervised,
and reinforcement learning.

Through our experiments, we noticed that the current architectures and
theories developed for backpropagation and employed for defining the network
components and initializing its parameters may be sub-optimal for TAGI. Future
work in this direction may further improve the performance of the approach.

## Acknowledgements

The second author was financially supported by research grants from Hydro-
Quebec, and the Natural Sciences and Engineering Research Council of Canada
(NSERC). The third author was supported by a research grant from the Institute
for Data Valorization (IVADO).

## References

  * Arasaratnam & Haykin (2008) Arasaratnam, I. and Haykin, S.  Nonlinear bayesian filters for training recurrent neural networks.  In _Mexican International Conference on Artificial Intelligence_ , pp. 12–33. Springer, 2008.
  * Arasaratnam & Haykin (2009) Arasaratnam, I. and Haykin, S.  Cubature kalman filters.  _IEEE Transactions on automatic control_ , 54(6):1254–1269, 2009.
  * Barber & Bishop (1998) Barber, D. and Bishop, C. M.  Ensemble learning in Bayesian neural networks.  _NATO ASI Series F Computer and Systems Sciences_ , 168:215–238, 1998.
  * Blundell et al. (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D.  Weight uncertainty in neural network.  In _International Conference on Machine Learning_ , pp. 1613–1622, 2015.
  * Efron (2012) Efron, B.  _Large-scale inference: empirical Bayes methods for estimation, testing, and prediction_ , volume 1.  Cambridge University Press, 2012.
  * Farquhar & Gal (2019) Farquhar, S. and Gal, Y.  A unifying Bayesian view of continual learning.  _arXiv preprint arXiv:1902.06494_ , 2019.
  * Gal & Ghahramani (2016) Gal, Y. and Ghahramani, Z.  Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.  In _International Conference on Machine Learning_ , pp. 1050–1059, 2016.
  * Ghahramani (2015) Ghahramani, Z.  Probabilistic machine learning and artificial intelligence.  _Nature_ , 521(7553):452, 2015.
  * Glorot & Bengio (2010) Glorot, X. and Bengio, Y.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth International Conference on Artificial Intelligence and Statistics_ , pp. 249–256, 2010.
  * Goodfellow et al. (2016) Goodfellow, I., Bengio, Y., and Courville, A.  _Deep learning_.  MIT Press, 2016.
  * Hernández-Lobato & Adams (2015) Hernández-Lobato, J. M. and Adams, R.  Probabilistic backpropagation for scalable learning of Bayesian neural networks.  In _International Conference on Machine Learning_ , pp. 1861–1869, 2015.
  * Hinton & Van Camp (1993) Hinton, G. E. and Van Camp, D.  Keeping the neural networks simple by minimizing the description length of the weights.  In _Proceedings of the sixth annual conference on Computational learning theory_ , pp. 5–13. ACM, 1993.
  * Kendall & Gal (2017) Kendall, A. and Gal, Y.  What uncertainties do we need in Bayesian deep learning for computer vision?  In _Advances in neural information processing systems_ , pp. 5574–5584, 2017.
  * LeCun et al. (1998) LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Louizos & Welling (2016) Louizos, C. and Welling, M.  Structured and efficient variational deep learning with matrix Gaussian posteriors.  In _International Conference on Machine Learning_ , pp. 1708–1716, 2016.
  * MacKay (1992) MacKay, D.  A practical Bayesian framework for backpropagation networks.  _Neural computation_ , 4(3):448–472, 1992.
  * Morin & Bengio (2005) Morin, F. and Bengio, Y.  Hierarchical probabilistic neural network language model.  In _Aistats_ , volume 5, pp. 246–252. Citeseer, 2005.
  * Neal (1995) Neal, R. M.  _Bayesian learning for neural networks_.  PhD thesis, University of Toronto, 1995.
  * Osawa et al. (2019) Osawa, K., Swaroop, S., Jain, A., Eschenhagen, R., Turner, R. E., Yokota, R., and Khan, M. E.  Practical deep learning with bayesian principles.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Puskorius & Feldkamp (1991) Puskorius, G. V. and Feldkamp, L. A.  Decoupled extended kalman filter training of feedforward layered networks.  In _IJCNN-91-Seattle International Joint Conference on Neural Networks_ , volume 1, pp. 771–777. IEEE, 1991.
  * Rasmussen & Williams (2006) Rasmussen, C. E. and Williams, C. K.  _Gaussian processes for machine learning_.  the MIT Press, 2006.
  * Rauch et al. (1965) Rauch, H. E., Striebel, C. T., and Tung, F.  Maximum likelihood estimates of linear dynamic systems.  _AIAA journal_ , 3(8):1445–1450, 1965.
  * Rumelhart et al. (1986) Rumelhart, D. E., Hinton, G. E., and Williams, R. J.  Learning representations by back-propagating errors.  _nature_ , 323(6088):533, 1986.
  * Simard et al. (2003) Simard, P. Y., Steinkraus, D., and Platt, J. C.  Best practices for convolutional neural networks applied to visual document analysis.  In _Proceedings of Seventh International Conference on Document Analysis and Recognition_ , 2003.
  * Singhal & Wu (1989) Singhal, S. and Wu, L.  Training multilayer perceptrons with the extended kalman algorithm.  In Touretzky, D. S. (ed.), _Advances in Neural Information Processing Systems 1_ , pp. 133–140. Morgan-Kaufmann, 1989.
  * Srivastava et al. (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ , 15(1):1929–1958, 2014.
  * Sun et al. (2017) Sun, S., Chen, C., and Carin, L.  Learning structured weight uncertainty in bayesian neural networks.  In _Artificial Intelligence and Statistics_ , pp. 1283–1292, 2017\.
  * Wan & Merwe (2000) Wan, E. A. and Merwe, R. v. d.  The unscented Kalman filter for nonlinear estimation.  In _Adaptive Systems for Signal Processing, Communications, and Control Symposium_ , pp. 153–158. IEEE, 2000.
  * Wan et al. (2013) Wan, L., Zeiler, M., Zhang, S., Cun, Y. L., and Fergus, R.  Regularization of neural networks using dropconnect.  In _Proceedings of the 30th International Conference on Machine Learning_ , volume 28 of _Proceedings of Machine Learning Research_ , pp. 1058–1066, 2013.

## Appendix

##  Appendix A Derivation of moments for the Gaussian multiplicative
approximation

The proof of statement (3) can be obtained directly from the definition of
covariance. To prove the statement (4), one needs the moments of variables and
their products. Because the underlying distribution is Gaussian and its moment
generating function is known,
$M_{\bm{X}}(\mathbf{t}^{\intercal})=\mathbb{E}[e^{\mathbf{t}^{\intercal}\bm{X}}]=e^{\mathbf{t}^{\intercal}\bm{\mu}+\frac{1}{2}\mathbf{t}^{\intercal}\bm{\Sigma}\mathbf{t}},\mathbf{t}=[t_{1}~{}\ldots~{}t_{4}]^{\intercal}$,
all moments can be extracted using the moment-generating function. Using the
derivatives of the moment-generating function,

$\displaystyle\mathbb{E}[X_{1}X_{2}X_{3}]$ $\displaystyle=$
$\displaystyle\frac{\partial^{3}}{\partial t_{1}\partial t_{2}\partial
t_{3}}\mathbb{E}[e^{\mathbf{t}^{\intercal}\bm{x}}]|_{t_{1}=t_{2}=t_{3}=t_{4}=0}$
(19) $\displaystyle=$ $\displaystyle\frac{\partial^{3}}{\partial t_{1}\partial
t_{2}\partial
t_{3}}e^{\mathbf{t}^{\intercal}\bm{\mu}+\frac{1}{2}\mathbf{t}^{\intercal}\mathbf{\Sigma}\mathbf{t}}|_{t_{1}=t_{2}=t_{3}=t_{4}=0}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!=\frac{\partial^{3}}{\partial
t_{1}\partial t_{2}\partial
t_{3}}e^{\sum_{i=1}^{4}\\!t_{i}\mu_{i}+\frac{1}{2}\sum_{i,j=1}^{4}\\!t_{i}t_{j}\text{cov}(X_{i},X_{j})}|_{t_{1,2,3,4}=0}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!=\text{cov}(X_{1},X_{2})\mu_{3}+\text{cov}(X_{1},X_{3})\mu_{2}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{2},X_{3})\mu_{1}+\mu_{1}\mu_{2}\mu_{3}.$

Using the definition of covariance and substituting (19), the statement (4)
can be established:

$\displaystyle\text{cov}(X_{3},X_{1}X_{2})$ $\displaystyle=$
$\displaystyle\mathbb{E}[X_{1}X_{2}X_{3}]-\mathbb{E}(X_{1}X_{3})\mathbb{E}(X_{2})$
(20) $\displaystyle=$
$\displaystyle\text{cov}(X_{1},X_{2})\mu_{3}+\text{cov}(X_{1},X_{3})\mu_{2}$
$\displaystyle+\text{cov}(X_{2},X_{3})\mu_{1}+\mu_{1}\mu_{2}\mu_{3}$
$\displaystyle-\mu_{3}\big{(}\mu_{1}\mu_{2}+\text{cov}(X_{1},X_{2})\big{)}.$

The expansion of the right side for the last statement leads to Equation (2.1)
so that

$\text{cov}(X_{1}X_{2},X_{3}X_{4})=\mathbb{E}[X_{1}X_{2}X_{3}X_{4}]-\mathbb{E}[X_{1}X_{2}]\mathbb{E}[X_{3}X_{4}],$

where the expected value of the product two random variables is given in (3)
and the expectation for the product of four random variables is a
generalization of (19) that can be obtained using the derivatives of the
moment-generating function:

$\displaystyle\mathbb{E}[X_{1}X_{2}X_{3}X_{4}]$ $\displaystyle=$
$\displaystyle\frac{\partial^{4}}{\partial t_{1}\partial t_{2}\partial
t_{3}\partial
t_{4}}\mathbb{E}[e^{\mathbf{t}^{\intercal}\bm{X}}]|_{t_{1,2,3,4}=0}$
$\displaystyle=$ $\displaystyle\frac{\partial^{4}}{\partial t_{1}\partial
t_{2}\partial t_{3}\partial
t_{4}}e^{\mathbf{t}^{\intercal}\bm{\mu}+\frac{1}{2}\mathbf{t}^{\intercal}\mathbf{\Sigma}\mathbf{t}}|_{t_{1,2,3,4}=0}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!=\text{cov}(X_{1}X_{2})\big{(}\text{cov}(X_{3},X_{4})+\mu_{3}\mu_{4}\big{)}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{1}X_{3})\big{(}\text{cov}(X_{2},X_{4})+\mu_{2}\mu_{4}\big{)}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{2}X_{3})\big{(}\text{cov}(X_{1},X_{4})$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\mu_{1}\mu_{4}\big{)}+\text{cov}(X_{1},X_{4})\mu_{2}\mu_{3}+\text{cov}(X_{2},X_{4})\mu_{1}\mu_{3}$
$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!+\text{cov}(X_{3},X_{4})\mu_{1}\mu_{2}+\mu_{1}\mu_{2}\mu_{3}\mu_{4}.$

Using the definition of variance

$\displaystyle\text{var}(X_{1}X_{2})=\mathbb{E}[(X_{1}X_{2})^{2}]-\mathbb{E}[X_{1}X_{2}]^{2}$
(21)

The elements of variance can be expended as below

$\displaystyle\mathbb{E}[(X_{1}X_{2})^{2}]$ $\displaystyle=$
$\displaystyle\frac{\partial^{4}}{\partial t_{1}^{2}\partial
t_{2}^{2}}\mathbb{E}[e^{\mathbf{t}^{\intercal}\bm{X}}]|_{t_{1}=t_{2}=t_{3}=t_{4}=0}$
(22) $\displaystyle=$ $\displaystyle\frac{\partial^{4}}{\partial
t^{2}_{1}\partial
t^{2}_{2}}e^{\mathbf{t}^{\intercal}\bm{\mu}+\frac{1}{2}\mathbf{t}^{\intercal}\mathbf{\Sigma}\mathbf{t}}|_{t_{1}=t_{2}=t_{3}=t_{4}=0}$
$\displaystyle=$
$\displaystyle\sigma_{1}^{2}\sigma_{2}^{2}+2\text{cov}(X_{1},X_{2})^{2}$
$\displaystyle+4\text{cov}(X_{1},X_{2})\mu_{1}\mu_{2}$
$\displaystyle+\sigma_{1}^{2}\mu_{2}^{2}+\sigma_{2}^{2}\mu_{1}^{2}+\mu_{1}^{2}\mu_{2}^{2}.$
$\displaystyle\mathbb{E}[X_{1}X_{2}]^{2}$ $\displaystyle=$
$\displaystyle\big{(}\text{cov}(X_{1},X_{2})+\mathbb{E}(X_{1})\mathbb{E}(X_{2})\big{)}^{2}$
(23) $\displaystyle=$
$\displaystyle\text{cov}(X_{1},X_{2})^{2}+2\text{cov}(X_{1},X_{2})\mu_{1}\mu_{2}$
$\displaystyle+\mu_{1}^{2}\mu_{2}^{2}.$

Substituting (22) and (23) in (21) establishes (2.1).

##  Appendix B Example of $\mathbf{F}_{\bm{wa}}^{(j)}$ and in
$\mathbf{F}_{\bm{b}}^{(j)}$ matrices

Figure 10 presents an example of two successive hidden layers each comprising
only two hidden units.

$a^{(1)}_{1}$$a^{(1)}_{2}$$\bm{b}^{(1)}$$z^{(2)}_{1}$$z^{(2)}_{2}$$w^{(1)}_{1,1}$$w^{(1)}_{1,1}$$w^{(1)}_{1,2}$$w^{(1)}_{2,1}$$b^{(1)}_{1}$
Figure 10: Example of trivial network configuration employed to illustrate the
configuration for the matrices $\mathbf{F}_{\bm{wa}}^{(j)}$ and
$\mathbf{F}_{\bm{b}}^{(j)}$.

The formulation of the the $\mathbf{F}_{\bm{wa}}^{(j)}$ and in
$\mathbf{F}_{\bm{b}}^{(j)}$ matrices corresponding to the network in Figure 10
is

$\underbrace{\small\begin{bmatrix}z^{(2)}_{1}\\\\[4.0pt]
z^{(2)}_{2}\end{bmatrix}}_{\bm{z}^{(j+1)}}=\underbrace{\small\begin{bmatrix}1&1&0&0\\\
0&0&1&1\end{bmatrix}}_{\mathbf{F}_{\bm{wa}}^{(j)}}\times\underbrace{\small\begin{bmatrix}w^{(1)}_{1,1}a^{(1)}_{1}\\\\[4.0pt]
w^{(1)}_{1,2}a^{(1)}_{2}\\\\[4.0pt] w^{(1)}_{2,1}a^{(1)}_{1}\\\\[4.0pt]
w^{(1)}_{2,2}a^{(1)}_{2}\end{bmatrix}}_{(\bm{w\\!\\!a})^{(j)}}+\underbrace{\small\begin{bmatrix}1&0\\\
0&1\end{bmatrix}}_{\mathbf{F}_{\bm{b}}^{(j)}}\times\underbrace{\small\begin{bmatrix}b^{(1)}_{1}\\\\[4.0pt]
b^{(1)}_{2}\end{bmatrix}}_{\bm{b}^{(j)}}.$

Note that the structure of $\mathbf{F}_{\bm{wa}}^{(j)}$ depends on the
ordering of variables.

##  Appendix C Experiment configurations for benchmark regression datasets

Table 3 presents the details for the experiments conducted for the benchmark
regression datasets. Note that the times presented in the last columns are the
average parameter (i.e. weights and bias) inference time in second per folds
and the average hyper parameter (i.e. $\sigma_{V}$) optimization time in
second. All these experiments were conducted using CPU.

Table 3: Experiment details for the benchmark regression datasets using BLNN.
$\mathtt{X}$: number of covariates, $\mathtt{L}$: number of layers,
$\mathtt{A}$: number of activation units per layer, $\mathtt{F}$ number of
random training/test folds.

|  | Train | Test |  |  |  Average inference ($\bm{\theta}$) |  Average optimization ($\sigma_{V}$) |  Average $\sigma_{V}$
---|---|---|---|---|---|---|---|---
Datasets | $\mathtt{X}$ | #obs. | #obs. | $\mathtt{L}\times\mathtt{A}$ | $\mathtt{F}$ | time per fold (s) | time per fold (s) | (20 folds)
Boston | 13 | 455 | 51 |  1$\times$50 | 20 | 4 | 14 | 0.32
Concrete | 8 | 927 | 103 |  1$\times$50 | 20 | 8 | 22 | 0.33
Energy | 8 | 691 | 77 |  1$\times$50 | 20 | 6 | 24 | 0.15
Kin8nm | 8 | 7373 | 819 |  1$\times$50 | 20 | 61 | 167 | 0.35
Naval | 16 | 11934 | 1193 |  1$\times$50 | 20 | 90 | 280 | 0.31
Power | 4 | 8611 | 957 |  1$\times$50 | 20 | 53 | 124 | 0.24
Protein | 9 | 41157 | 4373 |  1$\times$100 | 5 | 363 | 466 | 0.72
Wine | 11 | 1439 | 160 |  1$\times$50 | 20 | 12 | 20 | 0.75
Yacht | 6 | 277 | 31 |  1$\times$50 | 20 | 2 | 23 | 0.063

Table 4: Regression benchmark results using the ReLU activation function and different number of hidden layer. | Number of hidden layers
---|---
Datasets | 1 | 2 | 3 | 4
| Root mean square error (RMSE)
Boston | 3.02$\pm$0.83 | 2.67$\pm$0.73 | 2.75$\pm$0.7 | 2.78$\pm$0.66
Concrete | 5.84$\pm$0.57 | 5.10$\pm$0.45 | 5.06$\pm$0.54 | 5.16$\pm$0.47
Energy | 1.54$\pm$0.16 | 0.50$\pm$0.08 | 0.50$\pm$0.07 | 0.52$\pm$0.08
Kin8nm | 0.10$\pm$4E-3 | 0.07$\pm$3E-3 | 0.07$\pm$1E-3 | 0.07$\pm$2E-3
Naval | 7E-3$\pm$5E-4 | 1.4E-3$\pm$2E-5 | 6.5E-4$\pm$7E-5 | 5.2E-4$\pm$6E-5
Power | 4.13$\pm$0.15 | 3.97$\pm$0.16 | 3.88$\pm$0.16 | 3.85$\pm$0.18
Protein | 4.64$\pm$0.15 | 4.09$\pm$0.01 | 3.76$\pm$0.02 | 3.65$\pm$0.04
Wine | 0.63$\pm$0.04 | 0.63$\pm$0.04 | 0.64$\pm$0.04 | 0.64$\pm$0.04
Yacht | 0.86$\pm$0.25 | 0.70$\pm$0.36 | 0.74$\pm$0.29 | 4.46$\pm$6.56
| Average log-likelihood (LL)
Boston | -2.56$\pm$0.33 | -2.43$\pm$0.20 | -2.44$\pm$0.18 | -2.46$\pm$0.19
Concrete | -3.19$\pm$0.10 | -3.06$\pm$0.07 | -3.06$\pm$0.08 | -3.07$\pm$0.07
Energy | -1.87$\pm$0.11 | -1.44$\pm$0.02 | -1.45$\pm$0.02 | -1.45$\pm$0.02
Kin8nm | 0.86$\pm$0.04 | 1.16$\pm$0.02 | 1.20$\pm$0.01 | 1.19$\pm$0.02
Naval | 3.27$\pm$0.20 | 4.42$\pm$0.01 | 4.46$\pm$3E-3 | 4.46$\pm$2E-3
Power | -2.84$\pm$0.04 | -2.80$\pm$0.04 | 2.78$\pm$0.04 | -2.77$\pm$0.04
Protein | -2.96$\pm$6E-3 | -2.83$\pm$0.00 | 2.78$\pm$0.00 | -2.75$\pm$0.00
Wine | -0.96$\pm$0.06 | -0.96$\pm$0.07 | -0.97$\pm$0.06 | -0.97$\pm$0.07
Yacht | -1.30$\pm$0.23 | -1.19$\pm$0.21 | -1.27$\pm$0.16 | -34.65$\pm$60.77
Table 5: Regression benchmark results using the softplus activation function and different number of hidden layer. | Number of hidden layers
---|---
Datasets | 1 | 2 | 3 | 4
| Root mean square error (RMSE)
Boston | 3.08$\pm$0.80 | 2.84$\pm$0.70 | 2.91$\pm$0.60 | 2.99$\pm$0.62
Concrete | 5.89$\pm$0.47 | 5.21$\pm$0.50 | 5.31$\pm$0.49 | 5.16$\pm$0.47
Energy | 1.78$\pm$0.25 | 0.52$\pm$0.07 | 0.52$\pm$0.07 | 0.61$\pm$0.13
Kin8nm | 0.10$\pm$3E-3 | 0.07$\pm$3E-3 | 0.08$\pm$3E-3 | 0.1$\pm$4E-3
Naval | 6E-3$\pm$3E-4 | 9E-4$\pm$1E-4 | 5.4E-4$\pm$1E-4 | 4.6E-4$\pm$7E-5
Power | 4.20$\pm$0.15 | 4.04$\pm$0.16 | 3.99$\pm$0.16 | 3.94$\pm$0.18
Protein | 4.69$\pm$0.03 | 4.18$\pm$0.05 | 3.84$\pm$0.01 | 3.79$\pm$0.02
Wine | 0.63$\pm$0.03 | 0.63$\pm$0.05 | 0.65$\pm$0.05 | 0.65$\pm$0.05
Yacht | 0.98$\pm$0.27 | 1.14$\pm$1.02 | 2.63$\pm$1.80 | 3.85$\pm$1.69
| Average log-likelihood (LL)
Boston | -2.56$\pm$0.28 | -2.45$\pm$0.17 | -2.48$\pm$0.17 | -2.52$\pm$0.20
Concrete | -3.20$\pm$0.09 | -3.08$\pm$0.08 | -3.09$\pm$0.08 | -3.07$\pm$0.07
Energy | -2.05$\pm$0.19 | -1.43$\pm$0.02 | -1.44$\pm$0.02 | -1.45$\pm$0.03
Kin8nm | 0.88$\pm$0.03 | 1.14$\pm$0.02 | 1.08$\pm$0.03 | 0.88$\pm$0.05
Naval | 3.70$\pm$0.09 | 4.45$\pm$7E-3 | 4.46$\pm$3E-3 | 4.46$\pm$2E-3
Power | -2.85$\pm$0.04 | -2.82$\pm$0.04 | -2.80$\pm$0.04 | -2.79$\pm$0.04
Protein | -2.97$\pm$7E-3 | -2.85$\pm$0.01 | -2.78$\pm$3E-3 | -2.77$\pm$4E-3
Wine | -0.97$\pm$0.06 | -0.96$\pm$0.07 | -1.00$\pm$0.08 | -1.00$\pm$0.08
Yacht | -1.42$\pm$0.26 | -1.78$\pm$1.34 | -3.9$\pm$4.6 | -6.55$\pm$0.70