# Sparse Uncertainty Representation in Deep Learning with Inducing Weights

Hippolyt Ritter1, Martin Kukla2, Cheng Zhang2 & Yingzhen Li3∗
1University College London 2Microsoft Research Cambridge, UK 3Imperial College
London
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS>Work done at Microsoft Research Cambridge.

###### Abstract

Bayesian Neural Networks and deep ensembles represent two modern paradigms of
uncertainty quantification in deep learning. Yet these approaches struggle to
scale mainly due to memory inefficiency, requiring parameter storage several
times that of their deterministic counterparts. To address this, we augment
each weight matrix with a small inducing weight matrix, projecting the
uncertainty quantification into a lower dimensional space. We further extend
Matheron’s conditional Gaussian sampling rule to enable fast weight sampling,
which enables our inference method to maintain reasonable run-time as compared
with ensembles. Importantly, our approach achieves competitive performance to
the state-of-the-art in prediction and uncertainty estimation tasks with fully
connected neural networks and ResNets, while reducing the parameter size to
$\leq 24.3\%$ of that of a _single_ neural network.

##  1 Introduction

Deep learning models are becoming deeper and wider than ever before. From
image recognition models such as ResNet-101 (He et al., 2016a) and DenseNet
(Huang et al., 2017) to BERT (Xu et al., 2019) and GPT-3 (Brown et al., 2020)
for language modelling, deep neural networks have found consistent success in
fitting large-scale data. As these models are increasingly deployed in real-
world applications, calibrated uncertainty estimates for their predictions
become crucial, especially in safety-critical areas such as healthcare. In
this regard, Bayesian Neural Networks (BNNs)(MacKay, 1995; Blundell et al.,
2015; Gal & Ghahramani, 2016; Zhang et al., 2020) and deep ensembles
(Lakshminarayanan et al., 2017) represent two popular paradigms for estimating
uncertainty, which have shown promising results in applications such as
(medical) image processing (Kendall & Gal, 2017; Tanno et al., 2017) and out-
of-distribution detection (Ovadia et al., 2019).

Though progress has been made, one major obstacle to scaling up BNNs and deep
ensembles is their high storage cost. Both approaches require the parameter
counts to be several times higher than their deterministic counterparts.
Although recent efforts have improved memory efficiency (Louizos & Welling,
2017; Świątkowski et al., 2020; Wen et al., 2020; Dusenberry et al., 2020),
these still use more parameters than a deterministic neural network. This is
particularly problematic in hardware-constrained edge devices, when on-device
storage is required due to privacy regulations.

Meanwhile, an _infinitely wide_ BNN becomes a Gaussian process (GP)that is
known for good uncertainty estimates (Neal, 1995; Matthews et al., 2018; Lee
et al., 2018). But perhaps surprisingly, this infinitely wide BNN is
“parameter efficient”, as its “parameters” are effectively the datapoints,
which have a considerably smaller memory footprint than explicitly storing the
network weights. In addition, sparse posterior approximations store a smaller
number of _inducing points_ instead (Snelson & Ghahramani, 2006; Titsias,
2009), making sparse GPs even more memory efficient.

Can we bring the advantages of sparse approximations in GPs— which are
infinitely-wide neural networks — to finite width deep learning models? We
provide an affirmative answer regarding memory efficiency, by proposing an
uncertainty quantification framework based on _sparse uncertainty
representations_. We present our approach in BNN context, but the proposed
approach is also applicable to deep ensembles. In detail, our contributions
are as follows:

  * •

We introduce _inducing weights_ — an auxiliary variable method with lower
dimensional counterparts to the actual weight matrices — for variational
inference in BNNs, as well as a memory efficient parameterisation and an
extension to ensemble methods (Section 3.1).

  * •

We extend Matheron’s rule to facilitate efficient posterior sampling (Section
3.2).

  * •

We provide an in-depth computation complexity analysis (Section 3.3), showing
the significant advantage in terms of parameter efficiency.

  * •

We show the connection to sparse (deep) GPs, in that inducing weights can be
viewed as _projected noisy inducing outputs_ in pre-activation output space
(Section 5.1).

  * •

We apply the proposed approach to BNNs and deep ensembles. Experiments in
classification, model robustness and out-of-distribution detection tasks show
that our inducing weight approaches achieve competitive performance to their
counterparts in the original weight space on modern deep architectures for
image classification, while reducing the parameter count to $\leq 24.3\%$ of
that of a single network.

We open-source our proposed inducing weight approach, together with baseline
methods reported in the experiments, as a PyTorch (Paszke et al., 2019)
wrapper named bayesianize: https://github.com/microsoft/bayesianize. As
demonstrated in Appendix I, our software makes the conversion of a
deterministic neural network to a Bayesian one with a few lines of code:

⬇

import bnn # our pytorch wrapper package

net = torchvision.models.resnet18() # construct a deterministic ResNet18

bnn.bayesianize_(net, inference="inducing") # convert it into a Bayesian one

##  2 Inducing variables for variational inference

Our work is built on variational inference and inducing variables for
posterior approximations. Given observations
$\mathcal{D}=\\{\mathbf{X},\mathbf{Y}\\}$ with
$\mathbf{X}=[\bm{x}_{1},...,\bm{x}_{N}]$,
$\mathbf{Y}=[\bm{y}_{1},...,\bm{y}_{N}]$, we would like to fit a neural
network $p(\bm{y}|\bm{x},{\mathbf{W}}_{1:L})$ with weights
${\mathbf{W}}_{1:L}$ to the data. BNNs posit a prior distribution
$p({\mathbf{W}}_{1:L})$ over the weights, and construct an approximate
posterior $q({\mathbf{W}}_{1:L})$ to the exact posterior
$p({\mathbf{W}}_{1:L}|\mathcal{D})\propto
p(\mathcal{D}|{\mathbf{W}}_{1:L})p({\mathbf{W}}_{1:L})$, where
$p(\mathcal{D}|{\mathbf{W}}_{1:L})=p(\mathbf{Y}|\mathbf{X},{\mathbf{W}}_{1:L})=\prod_{n=1}^{N}p(\bm{y}_{n}|\bm{x}_{n},{\mathbf{W}}_{1:L})$.

#### Variational inference

Variational inference (Hinton & Van Camp, 1993; Jordan et al., 1999; Zhang et
al., 2018a) constructs an approximation $q({\bm{\theta}})$ to the posterior
$p({\bm{\theta}}|\mathcal{D})\propto
p({\bm{\theta}})p(\mathcal{D}|{\bm{\theta}})$ by maximising a variational
lower-bound:

$\log
p(\mathcal{D})\geq\mathcal{L}(q({\bm{\theta}})):=\operatorname{\mathbb{E}}_{q({\bm{\theta}})}\left[\log
p(\mathcal{D}|{\bm{\theta}})\right]-\operatorname{\mathbb{KL}}\left[q({\bm{\theta}})||p({\bm{\theta}})\right].$
(1)

For BNNs, ${\bm{\theta}}=\\{{\mathbf{W}}_{1:L}\\}$, and a simple choice of $q$
is a Fully-factorized Gaussian (FFG):
$q({\mathbf{W}}_{1:L})=\prod_{l=1}^{L}\prod_{i=1}^{d_{out}^{l}}\prod_{j=1}^{d_{in}^{l}}\operatorname{\mathcal{N}}(m_{l}^{(i,j)},v_{l}^{(i,j)})$,
with $m_{l}^{(i,j)},v_{l}^{(i,j)}$ the mean and variance of
${\mathbf{W}}_{l}^{(i,j)}$ and $d_{in}^{l},d_{out}^{l}$ the respective number
of inputs and outputs to layer $l$. The variational parameters are then
$\bm{\phi}=\\{m_{l}^{(i,j)},v_{l}^{(i,j)}\\}_{l=1}^{L}$. Gradients of
$\mathcal{L}$ w.r.t. $\bm{\phi}$ can be estimated with mini-batches of data
(Hoffman et al., 2013) and with Monte Carlo sampling from the $q$ distribution
(Titsias & Lázaro-Gredilla, 2014; Kingma & Welling, 2014). By setting $q$ to
an BNN, a variational BNN can be trained with similar computational
requirements as a deterministic network (Blundell et al., 2015).

#### Improved posterior approximation with inducing variables

Auxiliary variable approaches (Agakov & Barber, 2004; Salimans et al., 2015;
Ranganath et al., 2016) construct the $q({\bm{\theta}})$ distribution with an
auxiliary variable ${\mathbf{a}}$: $q({\bm{\theta}})=\int
q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})d{\mathbf{a}}$, with the hope that
a potentially richer mixture distribution $q({\bm{\theta}})$ can achieve
better approximations. As then $q({\bm{\theta}})$ becomes intractable, an
auxiliary variational lower-bound is used to optimise
$q({\bm{\theta}},{\mathbf{a}})$ (see Appendix B):

$\log{}p(\mathcal{D})\geq\mathcal{L}(q({\bm{\theta}},{\mathbf{a}}))=\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}[\log
p(\mathcal{D}|{\bm{\theta}})]+\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{p({\bm{\theta}})r({\mathbf{a}}|{\bm{\theta}})}{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}\right].$
(2)

Here $r({\mathbf{a}}|{\bm{\theta}})$ is an auxiliary distribution that needs
to be specified, where existing approaches often use a “reverse model” for
$r({\mathbf{a}}|{\bm{\theta}})$. Instead, we define
$r({\mathbf{a}}|{\bm{\theta}})$ in a generative manner:
$r({\mathbf{a}}|{\bm{\theta}})$ is the “posterior” of the following
“generative model”, whose “evidence” is exactly the prior of ${\bm{\theta}}$:

$r({\mathbf{a}}|{\bm{\theta}})=\tilde{p}({\mathbf{a}}|{\bm{\theta}})\propto\tilde{p}({\mathbf{a}})\tilde{p}({\bm{\theta}}|{\mathbf{a}}),\quad\text{such
that
}\tilde{p}({\bm{\theta}}):=\int\tilde{p}({\mathbf{a}})\tilde{p}({\bm{\theta}}|{\mathbf{a}})d{\mathbf{a}}=p({\bm{\theta}}).\\\
$ (3)

Plugging Eq. 3 into Eq. 2:

$\mathcal{L}(q({\bm{\theta}},{\mathbf{a}}))=\mathbb{E}_{q({\bm{\theta}})}[\log
p(\mathcal{D}|{\bm{\theta}})]-\mathbb{E}_{q({\mathbf{a}})}\left[\operatorname{\mathbb{KL}}[q({\bm{\theta}}|{\mathbf{a}})||\tilde{p}({\bm{\theta}}|{\mathbf{a}})]\right]-\operatorname{\mathbb{KL}}[q({\mathbf{a}})||\tilde{p}({\mathbf{a}})].$
(4)

This approach yields an efficient approximate inference algorithm, translating
the complexity of inference in ${\bm{\theta}}$ to ${\mathbf{a}}$. If
$\text{dim}({\mathbf{a}})<\text{dim}({\bm{\theta}})$ and
$q({\bm{\theta}},{\mathbf{a}})=q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})$
has the following properties:

  * 1.

A “pseudo prior”
$\tilde{p}({\mathbf{a}})\tilde{p}({\bm{\theta}}|{\mathbf{a}})$ is defined such
that
$\int\tilde{p}({\mathbf{a}})\tilde{p}({\bm{\theta}}|{\mathbf{a}})d{\mathbf{a}}=p({\bm{\theta}})$;

  * 2.

The conditionals $q({\bm{\theta}}|{\mathbf{a}})$ and
$\tilde{p}({\bm{\theta}}|{\mathbf{a}})$ are in the same parametric family, so
can share parameters;

  * 3.

Both sampling ${\bm{\theta}}\sim q({\bm{\theta}})$ and computing
$\operatorname{\mathbb{KL}}[q({\bm{\theta}}|{\mathbf{a}})||\tilde{p}({\bm{\theta}}|{\mathbf{a}})]$
can be done efficiently;

  * 4.

The designs of $q({\mathbf{a}})$ and $\tilde{p}({\mathbf{a}})$ can potentially
provide extra advantages (in time and space complexities and/or optimisation
easiness).

We call ${\mathbf{a}}$ the _inducing variable_ of ${\bm{\theta}}$, which is
inspired by variationally sparse GP (SVGP)with inducing points (Snelson &
Ghahramani, 2006; Titsias, 2009). Indeed SVGP is a special case (see Appendix
C): ${\bm{\theta}}=\mathbf{f}$, ${\mathbf{a}}=\mathbf{u}$, the GP prior is
$p(\mathbf{f}|\mathbf{X})=\mathcal{GP}(\bm{0},\mathbf{K}_{\mathbf{X}\mathbf{X}})$,
$p(\mathbf{u})=\mathcal{GP}(\bm{0},\mathbf{K}_{\mathbf{Z}\mathbf{Z}})$,
$\tilde{p}(\mathbf{f},\mathbf{u})=p(\mathbf{u})p(\mathbf{f}|\mathbf{X},\mathbf{u})$,
$q(\mathbf{f}|\mathbf{u})=p(\mathbf{f}|\mathbf{X},\mathbf{u})$,
$q(\mathbf{f},\mathbf{u})=p(\mathbf{f}|\mathbf{X},\mathbf{u})q(\mathbf{u})$,
and $\mathbf{Z}$ are the optimisable inducing inputs. The variational lower-
bound is
$\mathcal{L}(q(\mathbf{f},\mathbf{u}))=\mathbb{E}_{q(\mathbf{f})}[\log
p(\mathbf{Y}|\mathbf{f})]-\operatorname{\mathbb{KL}}[q(\mathbf{u})||p(\mathbf{u})],$
and the variational parameters are $\bm{\phi}=\\{\mathbf{Z},\text{distribution
parameters of }q(\mathbf{u})\\}$. SVGP satisfies the marginalisation
constraint Eq. 3 by definition, and it has
$\operatorname{\mathbb{KL}}[q(\mathbf{f}|\mathbf{u})||\tilde{p}(\mathbf{f}|\mathbf{u})]=0$.
Also by using small $M=\text{dim}(\mathbf{u})$ and exploiting the $q$
distribution design, SVGP reduces run-time from $\mathcal{O}(N^{3})$ to
$\mathcal{O}(NM^{2}+M^{3})$ where $N$ is the number of inputs in $\mathbf{X}$,
meanwhile it also makes storing a full Gaussian $q(\mathbf{u})$ affordable.
Lastly, $\mathbf{u}$ can be whitened, leading to the “pseudo prior”
$\tilde{p}(\mathbf{f},\mathbf{v})=p(\mathbf{f}|\mathbf{X},\mathbf{u}=\mathbf{K}_{\mathbf{Z}\mathbf{Z}}^{1/2}\mathbf{v})\tilde{p}(\mathbf{v}),\tilde{p}(\mathbf{v})=\mathcal{N}(\mathbf{v};\bm{0},\mathbf{I})$
which could bring potential benefits in optimisation.

We emphasise that the introduction of “pseudo prior” does _not_ change the
_probabilistic model_ as long as the marginalisation constraint Eq. 3 is
satisfied. In the rest of the paper we assume the constraint Eq. 3 holds and
write $p({\bm{\theta}},{\mathbf{a}}):=\tilde{p}({\bm{\theta}},{\mathbf{a}})$.
It might seem unclear how to design such
$\tilde{p}({\bm{\theta}},{\mathbf{a}})$ for an arbitrary probabilistic model,
however, for a Gaussian prior on ${\bm{\theta}}$ the rules for computing
conditional Gaussian distributions can be used to construct $\tilde{p}$. In
Section 3 we exploit these rules to develop an efficient approximate inference
method for Bayesian neural networks with inducing weights.

##  3 Sparse uncertainty representation with inducing weights

###  3.1 Inducing weights for neural network parameters

Following the above design principles, we introduce to each network layer $l$
a _smaller_ inducing weight matrix ${\mathbf{U}}_{l}$ to assist approximate
posterior inference in ${\mathbf{W}}_{l}$. Therefore in our context,
${\bm{\theta}}={\mathbf{W}}_{1:L}$ and ${\mathbf{a}}={\mathbf{U}}_{1:L}$. In
the rest of the paper, we assume a factorised prior across layers
$p({\mathbf{W}}_{1:L})=\prod_{l}p({\mathbf{W}}_{l})$, and drop the $l$ indices
when the context is clear to ease notation.

#### Augmenting network layers with inducing weights

Suppose the weight ${\mathbf{W}}\in\mathbb{R}^{d_{out}\times d_{in}}$ has a
Gaussian prior
$p({\mathbf{W}})=p(\operatorname{vec}({\mathbf{W}}))=\operatorname{\mathcal{N}}(0,\sigma^{2}{\bm{I}})$
where $\operatorname{vec}({\mathbf{W}})$ concatenates the columns of the
weight matrix into a vector. A first attempt to augment
$p(\operatorname{vec}({\mathbf{W}}))$ with an inducing weight variable
${\mathbf{U}}\in\mathbb{R}^{M_{out}\times M_{in}}$ may be to construct a
multivariate Gaussian
$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$, such
that $\int
p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))d{\mathbf{U}}=\operatorname{\mathcal{N}}(0,\sigma^{2}{\bm{I}})$.
This means for the joint covariance matrix of
$(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$, it
requires the block corresponding to the covariance of
$\operatorname{vec}({\mathbf{W}})$ to match the prior covariance
$\sigma^{2}{\bm{I}}$. We are then free to parameterise the rest of the entries
in the joint covariance matrix, as long as this full matrix remains positive
definite. Now the conditional distribution $p({\mathbf{W}}|{\mathbf{U}})$ is a
function of these parameters, and the conditional sampling from
$p({\mathbf{W}}|{\mathbf{U}})$ is further discussed in Section D.1.
Unfortunately, as $\text{dim}(\operatorname{vec}({\mathbf{W}}))$ is typically
large (e.g. of the order of $10^{7}$), using a full covariance Gaussian for
$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$ becomes
computationally intractable.

We address this issue with matrix normal distributions (Gupta & Nagar, 2018).
The prior
$p(\operatorname{vec}({\mathbf{W}}))=\operatorname{\mathcal{N}}(\bm{0},\sigma^{2}{\bm{I}})$
has an equivalent matrix normal distribution form as
$p({\mathbf{W}})=\operatorname{\mathcal{MN}}(0,\sigma_{r}^{2}{\bm{I}},\sigma_{c}^{2}{\bm{I}})$,
with $\sigma_{r},\sigma_{c}>0$ the row and column standard deviations
satisfying $\sigma=\sigma_{r}\sigma_{c}$. Now we introduce the inducing
variable ${\mathbf{U}}$ in matrix space, as well as two auxiliary variables
${\mathbf{U}}_{r}\in\mathbb{R}^{M_{out}\times d_{in}}$,
${\mathbf{U}}_{c}\in\mathbb{R}^{d_{out}\times M_{in}}$, so that the full
augmented prior is:

$\displaystyle\begin{pmatrix}{\mathbf{W}}&{\mathbf{U}}_{c}\\\
{\mathbf{U}}_{r}&{\mathbf{U}}\end{pmatrix}$ $\displaystyle\sim
p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}}):=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c}),$
(5) with $\displaystyle{\bm{L}}_{r}=\begin{pmatrix}\sigma_{r}{\bm{I}}&0\\\
{\bm{Z}}_{r}&{\bm{D}}_{r}\end{pmatrix}\quad\text{s.t.}\quad\bm{\Sigma}_{r}={\bm{L}}_{r}{\bm{L}}_{r}^{\top}=\begin{pmatrix}\sigma_{r}^{2}{\bm{I}}&\sigma_{r}{\bm{Z}}_{r}^{\top}\\\
\sigma_{r}{\bm{Z}}_{r}&{\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}+{\bm{D}}_{r}^{2}\end{pmatrix}$
and $\displaystyle{\bm{L}}_{c}=\begin{pmatrix}\sigma_{c}{\bm{I}}&0\\\
{\bm{Z}}_{c}&{\bm{D}}_{c}\end{pmatrix}\quad\text{s.t.}\quad\bm{\Sigma}_{c}={\bm{L}}_{c}{\bm{L}}_{c}^{\top}=\begin{pmatrix}\sigma_{c}^{2}{\bm{I}}&\sigma_{c}{\bm{Z}}_{c}^{\top}\\\
\sigma_{c}{\bm{Z}}_{c}&{\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}+{\bm{D}}_{c}^{2}\end{pmatrix}.$

See Fig. 1(a) for a visualisation of the augmentation. Matrix normal
distributions have similar marginalisation and conditioning rules as
multivariate Gaussian distributions, for which we provide further examples in
Section D.2. Therefore the marginalisation constraint Eq. 3 is satisfied for
any ${\bm{Z}}_{c}\in\mathbb{R}^{M_{in}\times
d_{in}},{\bm{Z}}_{r}\in\mathbb{R}^{M_{out}\times d_{out}}$ and diagonal
matrices ${\bm{D}}_{c},{\bm{D}}_{r}$. For the inducing weight ${\mathbf{U}}$
we have
$p({\mathbf{U}})=\operatorname{\mathcal{MN}}(0,\bm{\Psi}_{r},\bm{\Psi}_{c})$
with $\bm{\Psi}_{r}={\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}+{\bm{D}}_{r}^{2}$ and
$\bm{\Psi}_{c}={\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}+{\bm{D}}_{c}^{2}$. In the
experiments we use _whitened_ inducing weights which transforms ${\mathbf{U}}$
so that $p({\mathbf{U}})=\operatorname{\mathcal{MN}}(0,{\bm{I}},{\bm{I}})$
(Appendix H), but for clarity we continue with the above formulas in the main
text.

The matrix normal parameterisation introduces two additional variables
${\mathbf{U}}_{r},{\mathbf{U}}_{c}$ without providing additional
expressiveness. Hence it is desirable to integrate them out, leading to a
joint multivariate normal with Khatri-Rao product structure for the
covariance:

$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}\left(0,\begin{pmatrix}\sigma_{c}^{2}{\bm{I}}\otimes\sigma_{r}^{2}{\bm{I}}&\sigma_{c}{\bm{Z}}_{c}^{\top}\otimes\sigma_{r}{\bm{Z}}_{r}^{\top}\\\
\sigma_{c}{\bm{Z}}_{c}\otimes\sigma_{r}{\bm{Z}}_{r}&\bm{\Psi}_{c}\otimes\bm{\Psi}_{r}\end{pmatrix}\right).$
(6)

As the dominating memory complexity here is
$\mathcal{O}(d_{out}M_{out}+d_{in}M_{in})$ which comes from storing
${\bm{Z}}_{r}$ and ${\bm{Z}}_{c}$, we see that the matrix normal
parameterisation of the augmented prior is memory efficient.

#### Posterior approximation in the joint space

We construct a factorised posterior approximation across the layers:
$q({\mathbf{W}}_{1:L},{\mathbf{U}}_{1:L})=\prod_{l}q({\mathbf{W}}_{l}|{\mathbf{U}}_{l})q({\mathbf{U}}_{l})$.
Below we discuss options for $q({\mathbf{W}}|{\mathbf{U}})$.

The simplest option is
$q({\mathbf{W}}|{\mathbf{U}})=p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}({\bm{\mu}}_{{\mathbf{W}}|{\mathbf{U}}},\bm{\Sigma}_{{\mathbf{W}}|{\mathbf{U}}})$,
similar to sparse GPs. A slightly more flexible variant adds a rescaling term
$\lambda^{2}$ to the covariance matrix, which allows efficient KL computation
(Appendix E):

$q({\mathbf{W}}|{\mathbf{U}})=q(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}({\bm{\mu}}_{{\mathbf{W}}|{\mathbf{U}}},\lambda^{2}\bm{\Sigma}_{{\mathbf{W}}|{\mathbf{U}}}),$
(7)
$R(\lambda):=\operatorname{\mathbb{KL}}\left[q({\mathbf{W}}|{\mathbf{U}})||p({\mathbf{W}}|{\mathbf{U}})\right]=d_{in}d_{out}(0.5\lambda^{2}-\log\lambda-0.5).$
(8)

Plugging ${\bm{\theta}}={\mathbf{W}}_{1:L},a={\mathbf{U}}_{1:L}$ and Eq. 8
into Eq. 4 returns the following variational lower-bound

$\mathcal{L}(q({\mathbf{W}}_{1:L},{\mathbf{U}}_{1:L}))=\mathbb{E}_{q({\mathbf{W}}_{1:L})}[\log
p(\mathcal{D}|{\mathbf{W}}_{1:L})]\,-\sum\nolimits_{l=1}^{L}(R(\lambda_{l})+\operatorname{\mathbb{KL}}[q({\mathbf{U}}_{l})||p({\mathbf{U}}_{l})]),$
(9)

with $\lambda_{l}$ the associated scaling parameter for
$q({\mathbf{W}}_{l}|{\mathbf{U}}_{l})$. Again as the choices of
${\bm{Z}}_{c},{\bm{Z}}_{r},{\bm{D}}_{c},{\bm{D}}_{r}$ do not change the
marginal prior $p({\mathbf{W}})$, we are safe to optimise them as well.
Therefore the variational parameters are now
$\phi=\\{{\bm{Z}}_{c},{\bm{Z}}_{r},{\bm{D}}_{c},{\bm{D}}_{r},\lambda,\text{dist.~{}params.~{}of
}q({\mathbf{U}})\\}$ for each layer.

#### Two choices of $q({\mathbf{U}})$

A simple choice is FFG
$q(\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}(\bm{m}_{u},\text{diag}(\bm{v}_{u}))$,
which performs mean-field inference in ${\mathbf{U}}$ space (c.f. Blundell et
al., 2015), and here
$\operatorname{\mathbb{KL}}[q({\mathbf{U}})||p({\mathbf{U}})]$ has a closed-
form solution. Another choice is a “mixture of delta measures”
$q({\mathbf{U}})=\frac{1}{K}\sum_{k=1}^{K}\delta({\mathbf{U}}={\mathbf{U}}^{(k)})$,
i.e. we keep $K$ distinct sets of parameters $\\{U_{1:L}^{(k)}\\}_{k=1}^{K}$
in inducing space that are projected back into the original parameter space
via the _shared_ conditionals $q({\mathbf{W}}_{l}|{\mathbf{U}}_{l})$ to obtain
the weights. This approach can be viewed as constructing “deep ensembles” in
${\mathbf{U}}$ space, and we follow ensemble methods (e.g. Lakshminarayanan et
al., 2017) to drop
$\operatorname{\mathbb{KL}}[q({\mathbf{U}})||p({\mathbf{U}})]$ in Eq. 9.

Often ${\mathbf{U}}$ is chosen to have significantly lower dimensions than
${\mathbf{W}}$, i.e. $M_{in}<<d_{in}$ and $M_{out}<<d_{out}$. As
$q({\mathbf{W}}|{\mathbf{U}})$ and $p({\mathbf{W}}|{\mathbf{U}})$ only differ
in the covariance scaling constant, ${\mathbf{U}}$ can be regarded as a
_sparse representation of uncertainty_ for the network layer, as the major
updates in (approximate) posterior belief is quantified by $q({\mathbf{U}})$.

###  3.2 Efficient sampling with extended Matheron’s rule

Figure 1: Visualisation of (a) the inducing weight augmentation, and compare
(b) the original Matheron’s rule to (c) our extended version. White blocks
represent samples from the joint Gaussian.

Computing the variational lower-bound Eq. 9 requires samples from
$q({\mathbf{W}})$, which requires an efficient sampling procedure for
$q({\mathbf{W}}|{\mathbf{U}})$. Unfortunately, $q({\mathbf{W}}|{\mathbf{U}})$
derived from Eq. 6 & Eq. 7 is not a matrix normal, so _direct_ sampling is
prohibitive. To address this challenge, we extend Matheron’s rule (Journel &
Huijbregts, 1978; Hoffman & Ribak, 1991; Doucet, 2010) to efficiently sample
from $q({\mathbf{W}}|{\mathbf{U}})$, with derivations provided in Appendix F.

The original Matheron’s rule applies to multivariate Gaussian distributions.
As a running example, consider two vector-valued random variables $\bm{w}$,
$\bm{u}$ with joint distribution
$p(\bm{w},\bm{u})=\mathcal{N}(\bm{0},\bm{\Sigma})$. Then the conditional
distribution
$p(\bm{w}|\bm{u})=\mathcal{N}(\bm{\mu}_{\bm{w}|\bm{u}},\bm{\Sigma}_{\bm{w}|\bm{u}})$
is also Gaussian, and direct sampling from it requires decomposing the
conditional covariance matrix $\bm{\Sigma}_{\bm{w}|\bm{u}}$ which can be
costly. The main idea of Matheron’s rule is that we can transform a sample
from the joint Gaussian to obtain a sample from the conditional distribution
$p(\bm{w}|\bm{u})$ as follows:

$\displaystyle\bm{w}=\bm{\bar{w}}+\bm{\Sigma}_{\bm{w}\bm{u}}\bm{\Sigma}_{\bm{u}\bm{u}}^{-1}(\bm{u}-\bm{\bar{u}}),\quad\bm{\bar{w}},\bm{\bar{u}}\sim\operatorname{\mathcal{N}}(\bm{0},\bm{\Sigma}),\quad\bm{\Sigma}=\begin{pmatrix}\bm{\Sigma}_{\bm{w}\bm{w}}&\bm{\Sigma}_{\bm{w}\bm{u}}\\\
\bm{\Sigma}_{\bm{u}\bm{w}}&\bm{\Sigma}_{\bm{u}\bm{u}}\end{pmatrix}.$ (10)

One can check the validity of Matheron’s rule by computing the mean and
variance of $\bm{w}$ above:

$\mathbb{E}_{\bm{\bar{w}},\bm{\bar{u}}}[\bm{w}]=\bm{\Sigma}_{\bm{w}\bm{u}}\bm{\Sigma}_{\bm{u}\bm{u}}^{-1}\bm{u}=\bm{\mu}_{\bm{w}|\bm{u}},\quad\mathbb{V}_{\bm{\bar{w}},\bm{\bar{u}}}[\bm{w}]=\bm{\Sigma}_{\bm{w}\bm{w}}-\bm{\Sigma}_{\bm{w}\bm{u}}\bm{\Sigma}_{\bm{u}\bm{u}}^{-1}\bm{\Sigma}_{\bm{u}\bm{w}}=\bm{\Sigma}_{\bm{w}|\bm{u}}.$

It might seem counter-intuitive at first sight in that this rule requires
samples from a higher dimensional space. However, in the case where
decomposition/inversion of $\bm{\Sigma}$ and $\bm{\Sigma}_{\bm{u}\bm{u}}$ can
be done efficiently, sampling from the joint Gaussian $p(\bm{w},\bm{u})$ can
be significantly cheaper than directly sampling from the conditional Gaussian
$p(\bm{w}|\bm{u})$. This happens e.g. when $\bm{\Sigma}$ is directly
parameterised by its Cholesky decomposition and $\dim(\bm{u})<<\dim(\bm{w})$,
so that sampling
$\bm{\bar{w}},\bm{\bar{u}}\sim\operatorname{\mathcal{N}}(\bm{0},\bm{\Sigma})$
is straight-forward, and computing $\bm{\Sigma}_{\bm{u}\bm{u}}^{-1}$ is
significantly cheaper than decomposing $\bm{\Sigma}_{\bm{w}|\bm{u}}$.

Unfortunately, the original Matheron’s rule cannot be applied directly to
sample from $q({\mathbf{W}}|{\mathbf{U}})$. This is because
$q({\mathbf{W}}|{\mathbf{U}})=q(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$
differs from
$p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$ only in
the variance scaling $\lambda$, and for
$p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$, its
joint distribution counter-part Eq. 6 does not have an efficient
representation for the covariance matrix. Therefore a naive application of
Matheron’s rule requires decomposing the covariance matrix of
$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$ which
is even more expensive than direct conditional sampling. However, notice that
for the joint distribution
$p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}})$ in an even
higher dimensional space, the row and column covariance matrices
$\bm{\Sigma}_{r}$ and $\bm{\Sigma}_{c}$ are parameterised by their Cholesky
decompositions, so that sampling from this joint distribution can be done
efficiently. This inspire us to extend the original Matheron’s rule for
efficient sampling from $q({\mathbf{W}}|{\mathbf{U}})$ (details in Appendix F,
when $\lambda=1$ it also applies to sampling from
$p({\mathbf{W}}|{\mathbf{U}})$):

${\mathbf{W}}=\lambda\bar{{\mathbf{W}}}+\sigma{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}(U-\lambda\bar{{\mathbf{U}}})\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c};\quad\bar{{\mathbf{W}}},\bar{{\mathbf{U}}}\sim
p(\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}})=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c}).$
(11)

Here $\bar{{\mathbf{W}}},\bar{{\mathbf{U}}}\sim
p(\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}})$
means we first sample
$\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}}$
from the joint then drop $\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}}$; in
fact $\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}}$ are never computed, and
the other samples $\bar{{\mathbf{W}}},\bar{{\mathbf{U}}}$ can be obtained by:

$\displaystyle\bar{{\mathbf{W}}}$ $\displaystyle=\sigma{\mathbf{E}}_{1},\
\bar{{\mathbf{U}}}={\bm{Z}}_{r}{\mathbf{E}}_{1}{\bm{Z}}_{c}^{\top}+\hat{{\bm{L}}}_{r}\tilde{{\mathbf{E}}}_{2}{\bm{D}}_{c}+{\bm{D}}_{r}\tilde{{\mathbf{E}}}_{3}\hat{{\bm{L}}}_{c}^{\top}+{\bm{D}}_{r}{\mathbf{E}}_{4}{\bm{D}}_{c},$
$\displaystyle{\mathbf{E}}_{1}$
$\displaystyle\sim\operatorname{\mathcal{MN}}(0,{\bm{I}}_{d_{out}},{\bm{I}}_{d_{in}});\quad\tilde{{\mathbf{E}}}_{2},\tilde{{\mathbf{E}}}_{3},{\mathbf{E}}_{4}\sim\operatorname{\mathcal{MN}}(0,{\bm{I}}_{M_{out}},{\bm{I}}_{M_{in}}),$
(12) $\displaystyle\hat{{\bm{L}}}_{r}$
$\displaystyle=\text{chol}({\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}),\
\hat{{\bm{L}}}_{c}=\text{chol}({\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}).$

The run-time cost is
$\mathcal{O}(2M_{out}^{3}+2M_{in}^{3}+d_{out}M_{out}M_{in}+M_{in}d_{out}d_{in})$
required by inverting $\bm{\Psi}_{r},\bm{\Psi}_{c}$, computing
$\hat{{\bm{L}}}_{r}$, $\hat{{\bm{L}}}_{c}$, and the matrix products. The
extended Matheron’s rule is visualised in Fig. 1 with a comparison to the
original Matheron’s rule for sampling from
$q(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$. Note
that the original rule requires joint sampling from Eq. 6 (i.e. sampling the
white blocks in Fig. 1(b)) which has
$\mathcal{O}((d_{out}d_{in}+M_{out}M_{in})^{3})$ cost. Therefore our recipe
avoids inverting and multiplying big matrices, resulting in a significant
speed-up for conditional sampling.

Table 1: Computational complexity per layer. We assume
${\mathbf{W}}\in\mathbb{R}^{d_{out}\times d_{in}}$,
${\mathbf{U}}\in\mathbb{R}^{M_{out}\times M_{in}}$, and $K$ forward passes for
each of the $N$ inputs. (∗ uses a parallel computing friendly vectorisation
technique (Wen et al., 2020) for further speed-up.)

Method | Time complexity | Storage complexity
---|---|---
Deterministic-${\mathbf{W}}$ | $\mathcal{O}(Nd_{in}d_{out})$ | $\mathcal{O}(d_{in}d_{out})$
FFG-${\mathbf{W}}$ | $\mathcal{O}(NKd_{in}d_{out})$ | $\mathcal{O}(2d_{in}d_{out})$
Ensemble-${\mathbf{W}}$ | $\mathcal{O}(NKd_{in}d_{out})$ | $\mathcal{O}(Kd_{in}d_{out})$
Matrix-normal-${\mathbf{W}}$ | $\mathcal{O}(NKd_{in}d_{out})$ | $\mathcal{O}(d_{in}d_{out}+d_{in}+d_{out})$
$k$-tied FFG-${\mathbf{W}}$ | $\mathcal{O}(NKd_{in}d_{out})$ | $\mathcal{O}(d_{in}d_{out}+k(d_{in}+d_{out}))$
rank-1 BNN | $\mathcal{O}(NKd_{in}d_{out})^{*}$ | $\mathcal{O}(d_{in}d_{out}+2(d_{in}+d_{out}))$
FFG-${\mathbf{U}}$ | $\mathcal{O}(NKd_{in}d_{out}+2M_{in}^{3}+2M_{out}^{3}$ | $\mathcal{O}(d_{in}M_{in}+d_{out}M_{out}+2M_{in}M_{out})$
| $+K(d_{out}M_{out}M_{in}+M_{in}d_{out}d_{in}))$ |
Ensemble-${\mathbf{U}}$ | same as above | $\mathcal{O}(d_{in}M_{in}+d_{out}M_{out}+KM_{in}M_{out})$

###  3.3 Computational complexities

In Table 1 we report the complexity figures for two types of inducing weight
approaches: FFG $q({\mathbf{U}})$ (FFG-${\mathbf{U}}$) and Delta mixture
$q({\mathbf{U}})$ (Ensemble-${\mathbf{U}}$). Baseline approaches include:
Deterministic-${\mathbf{W}}$, variational inference with FFG $q({\mathbf{W}})$
(FFG-${\mathbf{W}}$, Blundell et al., 2015), deep ensemble in ${\mathbf{W}}$
(Ensemble-${\mathbf{W}}$, Lakshminarayanan et al., 2017), as well as parameter
efficient approaches such as matrix-normal $q({\mathbf{W}})$ (Matrix-
normal-${\mathbf{W}}$, Louizos & Welling (2017)), variational inference with
$k$-tied FFG $q({\mathbf{W}})$ ($k$-tied FFG-${\mathbf{W}}$, Świątkowski et
al. (2020)), and rank-1 BNN (Dusenberry et al., 2020). The gain in memory is
significant for the inducing weight approaches, in fact with $M_{in}<d_{in}$
and $M_{out}<d_{out}$ the parameter storage requirement is smaller than a
single deterministic neural network. The major overhead in run-time comes from
the extended Matheron’s rule for sampling $q({\mathbf{W}}|{\mathbf{U}})$. Some
of the computations there are performed only once, and in our experiments we
show that by using a relatively low-dimensional ${\mathbf{U}}$ and large
batch-sizes, the overhead is acceptable.

##  4 Experiments

We evaluate the inducing weight approaches on regression, classification and
related uncertainty estimation tasks. The goal is to demonstrate competitive
performance to popular ${\mathbf{W}}$-space uncertainty estimation methods
while using significantly fewer parameters. We acknowledge that existing
parameter efficient approaches for uncertainty estimation (e.g. k-tied or
rank-1 BNNs) have achieved close performance to deep ensembles. However,
_none_ of them reduces the parameter count to be smaller than that of a
_single_ network. Therefore we decide _not_ to include these baselines and
instead focus on comparing: (1) variational inference with FFG
$q({\mathbf{W}})$ (FFG-${\mathbf{W}}$, Blundell et al., 2015) v.s. FFG
$q({\mathbf{U}})$ (FFG-${\mathbf{U}}$, ours); (2) deep ensemble in
${\mathbf{W}}$ space (Ensemble-${\mathbf{W}}$, Lakshminarayanan et al., 2017)
v.s. ensemble in ${\mathbf{U}}$ space (Ensemble-${\mathbf{U}}$, ours). Another
baseline is training a deterministic neural network with maximum likelihood.
Details and additional results are in Appendices J and K.

###  4.1 Synthetic 1-D regression

The regression task follows Foong et al. (2019), which has two input clusters
$x_{1}\sim\mathcal{U}\left[-1,-0.7\right]$,
$x_{2}\sim\mathcal{U}\left[0.5,1\right]$, and targets
$y\sim\mathcal{N}(\cos(4x+0.8),0.01)$. For reference we show the exact
posterior results using the NUTS sampler (Hoffman & Gelman, 2014). The results
are visualised in Fig. 2 with predictive mean in blue, and up to three
standard deviations as shaded area. Similar to historical results,
FFG-${\mathbf{W}}$ fails to represent the increased uncertainty away from the
data and in between clusters. While underestimating predictive uncertainty
overall, FFG-${\mathbf{U}}$ shows a small increase in predictive uncertainty
away from the data. In contrast, a per-layer Full-covariance Gaussian (FCG)in
both weight (FCG-${\mathbf{W}}$) and inducing space (FCG-${\mathbf{U}}$) as
well as Ensemble-${\mathbf{U}}$ better capture the increased predictive
variance, although the mean function is more similar to that of
FFG-${\mathbf{W}}$.

(a) FFG-${\mathbf{U}}$

(b) FCG-${\mathbf{U}}$

(c) Ensemble-${\mathbf{U}}$

(d) FFG-${\mathbf{W}}$

(e) FCG-${\mathbf{W}}$

(f) NUTS

Figure 2: Toy regression results, with observations in red dots and the ground
truth function in black.

###  4.2 Classification and in-distribution calibration

Table 2: CIFAR in-distribution metrics (in $\%$).

CIFAR10 CIFAR100

Method Acc. $\uparrow$ ECE $\downarrow$ Acc. $\uparrow$ ECE $\downarrow$

Deterministic $94.72$ $4.46$ $75.73$ $19.69$

Ensemble-W $95.90$ $1.08$ $79.33$ $6.51$

FFG-W $94.13$ $0.50$ $74.44$ $4.24$

FFG-U $94.40$ $0.64$ $75.37$ $2.29$

Ensemble-U $94.94$ $0.45$ $75.97$ $1.12$

---

Figure 3: Resnet run-times & model sizes.

As the core empirical evaluation, we train Resnet-50 models (He et al., 2016b)
on CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009). To avoid underfitting
issues with FFG-W, a useful trick is to set an upper limit $\sigma_{max}^{2}$
on the variance of $q({\mathbf{W}})$ (Louizos & Welling, 2017). This trick is
similarly applied to the ${\mathbf{U}}$-space methods, where we cap
$\lambda\leq\lambda_{max}$ for $q({\mathbf{W}}|{\mathbf{U}})$, and for
FFG-${\mathbf{U}}$ we also set $\sigma_{max}^{2}$ for the variance of
$q({\mathbf{U}})$. In convolution layers, we treat the 4D weight tensor
${\mathbf{W}}$ of shape $(c_{out},c_{in},h,w)$ as a $c_{out}\times c_{in}hw$
matrix. We use ${\mathbf{U}}$ matrices of shape $64\times 64$ for all layers
(i.e. $M=M_{in}=M_{out}=64$), except that for CIFAR-10 we set $M_{out}=10$ for
the last layer.

In Table 2 we report test accuracy and test expected calibration error
(ECE)(Guo et al., 2017) as a first evaluation of the uncertainty estimates.
Overall, Ensemble-${\mathbf{W}}$ achieves the highest accuracy, but is not as
well-calibrated as variational methods. For the inducing weight approaches,
Ensemble-${\mathbf{U}}$ outperforms FFG-${\mathbf{U}}$ on both datasets;
overall it performs the best on the more challenging CIFAR-100 dataset (close-
to-Ensemble-${\mathbf{W}}$ accuracy and lowest ECE). Appendices K and K in
Appendix K show that increasing the ${\mathbf{U}}$ dimensions to $M=128$
improves accuracy but leads to slightly worse calibration.

In Fig. 3 we show prediction run-times for batch-size $=500$ on an NVIDIA
Tesla V100 GPU, relative to those of an ensemble of deterministic nets, as
well as relative parameter sizes to a single ResNet-50. The extra run-times
for the inducing methods come from computing the extended Matheron’s rule.
However, as they can be calculated once and cached for drawing multiple
samples, the overhead reduces to a small factor when using larger number of
samples $K$, especially for the bigger Resnet-$50$. More importantly, when
compared to a _deterministic_ ResNet-50, the inducing weight models reduce the
parameter count by over $75\%$ ($5,710,902$ vs. $23,520,842$) for $M=64$.

Figure 4: Ablation study: average CIFAR-10 accuracy ($\uparrow$) and ECE
($\downarrow$) for the inducing weight methods on ResNet-18. In the first two
columns $M=128$ for ${\mathbf{U}}$ dimensions. For
$\lambda_{max},\sigma_{max}=0$ we use point estimates for
${\mathbf{U}},{\mathbf{W}}$ respectively.

#### Hyper-parameter choices

We visualise in Fig. 4 the accuracy and ECE results for computationally
lighter inducing weight ResNet-18 models with different hyper-parameters (see
Appendix J). Performance in both metrics improves as the ${\mathbf{U}}$ matrix
size $M$ is increased (right-most panels), and the results for $M=64$ and
$M=128$ are fairly similar. Also setting proper values for
$\lambda_{max},\sigma_{max}$ is key to the improved results. The left-most
panels show that with fixed $\sigma_{max}$ values (or
Ensemble-${\mathbf{U}}$), the preferred conditional variance cap values
$\lambda_{max}$ are fairly small (but still larger than $0$ which corresponds
to a point estimate for ${\mathbf{W}}$ given ${\mathbf{U}}$). For
$\sigma_{max}$ which controls variance in ${\mathbf{U}}$ space, we see from
the top middle panel that the accuracy metric is fairly robust to
$\sigma_{max}$ as long as $\lambda_{max}$ is not too large. But for ECE, a
careful selection of $\sigma_{max}$ is required (bottom middle panel).

###  4.3 Robustness, out-of-distribution detection and pruning

Figure 5: Mean${\pm}$two errs. for Acc$\uparrow$ and ECE $\downarrow$ on
corrupted CIFAR (Hendrycks & Dietterich, 2019). Table 3: OOD detection metrics
for Resnet-$50$ trained on CIFAR$10$/$100$.

In-dist $\rightarrow$ OOD C10 $\rightarrow$ C100 C10 $\rightarrow$ SVHN C100
$\rightarrow$ C10 C100 $\rightarrow$ SVHN
Method / Metric AUROC AUPR AUROC AUPR AUROC AUPR AUROC AUPR
Deterministic $.84{\scriptstyle{\pm}.00}$ $.80{\scriptstyle{\pm}.00}$
$.93{\scriptstyle{\pm}.01}$ $.85{\scriptstyle{\pm}.01}$
$.74{\scriptstyle{\pm}.00}$ $.74{\scriptstyle{\pm}.00}$
$.81{\scriptstyle{\pm}.01}$ $.72{\scriptstyle{\pm}.02}$
Ensemble-W $.89{\scriptstyle\phantom{{\pm}.00}}$
$.89{\scriptstyle\phantom{{\pm}.00}}$ $.95{\scriptstyle\phantom{{\pm}.00}}$
$.92{\scriptstyle\phantom{{\pm}.00}}$ $.78{\scriptstyle\phantom{{\pm}.00}}$
$.79{\scriptstyle\phantom{{\pm}.00}}$ $.86{\scriptstyle\phantom{{\pm}.00}}$
$.78{\scriptstyle\phantom{{\pm}.00}}$
FFG-W $.88{\scriptstyle{\pm}.00}$ $.90{\scriptstyle{\pm}.00}$
$.90{\scriptstyle{\pm}.01}$ $.86{\scriptstyle{\pm}.01}$
$.76{\scriptstyle{\pm}.00}$ $.79{\scriptstyle{\pm}.00}$
$.80{\scriptstyle{\pm}.01}$ $.69{\scriptstyle{\pm}.01}$
FFG-U $.89{\scriptstyle{\pm}.00}$ $.91{\scriptstyle{\pm}.00}$
$.94{\scriptstyle{\pm}.01}$ $.91{\scriptstyle{\pm}.01}$
$.77{\scriptstyle{\pm}.00}$ $.79{\scriptstyle{\pm}.00}$
$.83{\scriptstyle{\pm}.01}$ $.74{\scriptstyle{\pm}.01}$
Ensemble-U $.90{\scriptstyle{\pm}.00}$ $.91{\scriptstyle{\pm}.00}$
$.93{\scriptstyle{\pm}.00}$ $.91{\scriptstyle{\pm}.00}$
$.77{\scriptstyle{\pm}.00}$ $.79{\scriptstyle{\pm}.00}$
$.82{\scriptstyle{\pm}.01}$ $.72{\scriptstyle{\pm}.02}$

To investigate the models’ robustness to distribution shift, we compute
predictions on corrupted CIFAR datasets (Hendrycks & Dietterich, 2019) after
training on clean data. Fig. 5 shows accuracy and ECE results for the
ResNet-50 models. Ensemble-${\mathbf{W}}$ is the most accurate model across
skew intensities, while FFG-${\mathbf{W}}$, though performing well on clean
data, returns the worst accuracy under perturbation. The inducing weight
methods perform competitively to Ensemble-${\mathbf{W}}$ with
Ensemble-${\mathbf{U}}$ being slightly more accurate than FFG-${\mathbf{U}}$
as on the clean data. For ECE, FFG-${\mathbf{U}}$ outperforms
Ensemble-${\mathbf{U}}$ and Ensemble-${\mathbf{W}}$, which are similarly
calibrated. Interestingly, while the accuracy of FFG-${\mathbf{W}}$ decays
quickly as the data is perturbed more strongly, its ECE remains roughly
constant.

Section 4.3 further presents the utility of the maximum predicted probability
for out-of-distribution (OOD)detection. The metrics are the area under the
receiver operator characteristic (AUROC) and the precision-recall curve
(AUPR). The inducing-weight methods perform similarly to
Ensemble-${\mathbf{W}}$; all three outperform FFG-${\mathbf{W}}$ and
deterministic networks across the board.

Figure 6: CIFAR100 pruning accuracy($\uparrow$) & ECE ($\downarrow$).
Rightmost points are w/out pruning.

#### Parameter pruning

We further investigate pruning as a pragmatic alternative for more parameter-
efficient inference. For FFG-${\mathbf{U}}$, we prune entries of the $Z$
matrices, which contribute the largest number of parameters to the inducing
methods, with the smallest magnitude. For FFG-${\mathbf{W}}$ we follow Graves
(2011) in setting different fractions of ${\mathbf{W}}$ to $0$ depending on
their variational mean-to-variance ratio and repeat the previous experiments
after fine-tuning the distributions on the remaining variables. We stress
that, unlike FFG-${\mathbf{U}}$, the FFG-${\mathbf{W}}$ pruning corresponds to
a post-hoc change of the probabilistic model and no longer performs inference
in the original weight-space.

For FFG-${\mathbf{W}}$, pruning $90\%$ of the parameters (leaving $20\%$ of
parameters as compared to its deterministic counterpart) worsens the ECE, in
particular on CIFAR100, see Fig. 6. Further pruning to $1\%$ worsens the
accuracy and the OOD detection results as well. On the other hand, pruning
$50\%$ of the ${\bm{Z}}$ matrices for FFG-${\mathbf{U}}$ reduces the parameter
count to $13.2\%$ of a deterministic net, at the cost of only slightly worse
calibration. See Appendix K for the full results.

##  5 Discussions

###  5.1 A function-space perspective on inducing weights

Although the inducing weight approach performs approximate inference in weight
space, we present in Appendix G a function-space inference perspective of the
proposed method, showing its deep connections to sparse GPs. Our analysis
considers the function-space behaviour of each network layer’s output and
discusses the corresponding interpretations of the ${\mathbf{U}}$ variables
and ${\bm{Z}}$ parameters.

Figure 7: Visualising ${\mathbf{U}}$ variables in pre-activation spaces.

The interpretations are visualised in Fig. 7. Similar to sparse GPs, in each
layer, the ${\bm{Z}}_{c}$ parameters can be viewed as the (transposed)
inducing input locations which lie in the same space as the layer’s input. The
${\mathbf{U}}_{c}$ variables can also be viewed as the corresponding (noisy)
inducing outputs that lie in the pre-activation space. Given that the output
dimension $d_{out}$ can still be high (e.g. $>1000$ in a fully connected
layer), our approach performs further dimension reduction in a similar spirit
as probabilistic PCA (Tipping & Bishop, 1999), which projects the _column
vectors_ of ${\mathbf{U}}_{c}$ to a lower-dimensional space. This returns the
inducing weight variables ${\mathbf{U}}$, and the projection parameters are
$\\{{\bm{Z}}_{r},{\bm{D}}_{r}\\}$. Combining the two steps, it means the
column vectors of ${\mathbf{U}}$ can be viewed as collecting the “noisy
projected inducing outputs” whose corresponding “inducing inputs” are row
vectors of ${\bm{Z}}_{c}$ (see the red bars in Fig. 7).

In Appendix G we further derive the resulting variational objective from the
function-space view, which is almost identical to Eq. 9, except for scaling
coefficients on the $R(\lambda_{l})$ terms to account for the change in
dimensionality from weight space to function space. This result nicely
connects posterior inference in weight- and function-space.

###  5.2 Related work

#### Parameter-efficient uncertainty quantification methods

Recent research has proposed Gaussian posterior approximations for BNNs with
efficient covariance structure (Ritter et al., 2018; Zhang et al., 2018b;
Mishkin et al., 2018). The inducing weight approach differs from these in
introducing structure via a hierarchical posterior with low-dimensional
auxiliary variables. Another line of work reduces the memory overhead via
efficient parameter sharing (Louizos & Welling, 2017; Wen et al., 2020;
Świątkowski et al., 2020; Dusenberry et al., 2020). The third category of work
considers a hybrid approach, where only a selective part of the neural network
receives Bayesian treatments, and the other weights remain deterministic
(Bradshaw et al., 2017; Daxberger et al., 2021). However, both types of
approaches maintain a “mean parameter” for the weights, making the memory
footprint at least that of storing a deterministic neural network. Instead,
our approach shares parameters via the augmented prior with efficient low-rank
structure, _reducing_ the memory use compared to a deterministic network. In a
similar spirit, Izmailov et al. (2019) perform inference in a $d$-dimensional
sub-space obtained from PCA on weights collected from an SGD trajectory. But
this approach does not leverage the layer-structure of neural networks and
requires $d\times$ memory of a single network.

#### Network pruning in uncertainty estimation context

There is a large amount of existing research advocating network pruning
approaches for parameter-efficient deep learning, e.g. see Han et al. (2016);
Frankle & Carbin (2018); Lee et al. (2019). In this regard, mean-field VI
approaches have also shown success in network pruning, but only in terms of
maintaining a minimum accuracy level (Graves, 2011; Louizos et al., 2017;
Havasi et al., 2019). To the best of our knowledge, our empirical study
presents the first evaluation for VI-based pruning methods in maintaining
uncertainty estimation quality. Deng et al. (2019) considers pruning BNNs with
stochastic gradient Langevin dynamics (Welling & Teh, 2011) as the inference
engine. The inducing weight approach is orthogonal to these BNN pruning
approaches, as it leaves the prior on the network parameters intact, while the
pruning approaches correspond to a post-hoc change of the probabilistic model
to using a sparse weight prior. Indeed our parameter pruning experiments
showed that our approach can be combined with network pruning to achieve
further parameter efficiency improvements.

#### Sparse GP and function-space inference

As BNNs and GPs are closely related (Neal, 1995; Matthews et al., 2018; Lee et
al., 2018), recent efforts have introduced GP-inspired techniques to BNNs (Ma
et al., 2019; Sun et al., 2019; Khan et al., 2019; Ober & Aitchison, 2021).
Compared to weight-space inference, function-space inference is appealing as
its uncertainty is more directly relevant for predictive uncertainty
estimation. While the inducing weight approach performs computations in
weight-space, Section 5.1 establishes the connection to function-space
posteriors. Our approach is related to sparse deep GP methods with
${\mathbf{U}}_{c}$ having similar interpretations as inducing outputs in e.g.
Salimbeni & Deisenroth (2017). The major difference is that ${\mathbf{U}}$
lies in a low-dimensional space, projected from the pre-activation output
space of a network layer.

The original Matheron’s rule (Journel & Huijbregts, 1978; Hoffman & Ribak,
1991; Doucet, 2010) for sampling from conditional multivariate Gaussian
distributions has recently been applied to speed-up sparse GP inference
(Wilson et al., 2020, 2021). As explained in Section 3.2, direct application
of the original rule to sampling ${\mathbf{W}}$ conditioned on ${\mathbf{U}}$
still incurs prohibitive cost as
$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$ does
not have a convenient factorisation form. Our extended Matheron’s rule
addresses this issue by exploiting the efficient factorisation structure of
the joint matrix normal distribution
$p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}})$, reducing the
dominating factor of computation cost from cubic
($\mathcal{O}(d_{out}^{3}d_{in}^{3})$) to linear
($\mathcal{O}(d_{out}d_{in})$). We expect this new rule to be useful for a
wide range of models/applications beyond BNNs, such as matrix-variate Gaussian
processes (Stegle et al., 2011).

#### Priors on neural network weights

Hierarchical priors for weights has also been explored (Louizos et al., 2017;
Krueger et al., 2017; Atanov et al., 2019; Ghosh et al., 2019; Karaletsos &
Bui, 2020). However, we emphasise that $\tilde{p}({\mathbf{W}},{\mathbf{U}})$
is a _pseudo prior_ that is constructed to _assist posterior inference_ rather
than to _improve model design_. Indeed, parameters associated with the
inducing weights are optimisable for improving posterior approximations. Our
approach can be adapted to other priors, e.g. for a Horseshoe prior
$p(\theta,\nu)=p(\theta|\nu)p(\nu)=\mathcal{N}(\theta;0,\nu^{2})C^{+}(\nu;0,1)$,
the pseudo prior can be defined as
$\tilde{p}(\theta,\nu,a)=\tilde{p}(\theta|\nu,a)\tilde{p}(a)p(\nu)$ such that
$\int\tilde{p}(\theta|\nu,a)\tilde{p}(a)da=p(\theta|\nu)$. In general, pseudo
priors have found broader success in Bayesian computation (Carlin & Chib,
1995).

##  6 Conclusion

We have proposed a parameter-efficient uncertainty quantification framework
for neural networks. It augments each of the network layer weights with a
small matrix of inducing weights, and by extending Matheron’s rule to matrix-
normal related distributions, maintains a relatively small run-time overhead
compared with ensemble methods. Critically, experiments on prediction and
uncertainty estimation tasks show the competence of the inducing weight
methods to the state-of-the-art, while reducing the parameter count to under a
quarter of a deterministic ResNet-50 before pruning. This represents a
significant improvement over prior Bayesian and deep ensemble techniques,
which so far have not managed to go below this threshold despite various
attempts of matching it closely.

Several directions are to be explored in the future. First, modelling
correlations across layers might further improve the inference quality. We
outline an initial approach leveraging inducing variables in Appendix H.
Second, based on the function-space interpretation of inducing weights, better
initialisation techniques can be inspired from the sparse GP and dimension
reduction literature. Similarly, this interpretation might suggest other
innovative pruning approaches for the inducing weight method, thereby
achieving further memory savings. Lastly, the run-time overhead of our
approach can be mitigated by a better design of the inducing weight structure
as well as vectorisation techniques amenable to parallelised computation.
Designing hardware-specific implementations of the inducing weight approach is
also a viable alternative for such purposes.

## References

  * Agakov & Barber (2004) Agakov, F. V. and Barber, D.  An auxiliary variational method.  In _ICONIP_ , 2004.
  * Atanov et al. (2019) Atanov, A., Ashukha, A., Struminsky, K., Vetrov, D., and Welling, M.  The deep weight prior.  In _ICLR_ , 2019.
  * Bingham et al. (2019) Bingham, E., Chen, J. P., Jankowiak, M., Obermeyer, F., Pradhan, N., Karaletsos, T., Singh, R., Szerlip, P., Horsfall, P., and Goodman, N. D.  Pyro: Deep universal probabilistic programming.  _JMLR_ , 2019.
  * Blundell et al. (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D.  Weight uncertainty in neural networks.  In _ICML_ , 2015.
  * Bradshaw et al. (2017) Bradshaw, J., Matthews, A. G. d. G., and Ghahramani, Z.  Adversarial examples, uncertainty, and transfer testing robustness in Gaussian process hybrid deep networks.  _arXiv preprint arXiv:1707.02476_ , 2017.
  * Brown et al. (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  In _NeurIPS_ , 2020.
  * Carlin & Chib (1995) Carlin, B. P. and Chib, S.  Bayesian model choice via Markov chain Monte Carlo methods.  _JRSS B_ , 1995.
  * Daxberger et al. (2021) Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., and Hernández-Lobato, J. M.  Expressive yet tractable Bayesian deep learning via subnetwork inference.  In _ICML_ , 2021.
  * Deng et al. (2019) Deng, W., Zhang, X., Liang, F., and Lin, G.  An adaptive empirical bayesian method for sparse deep learning.  In _NeurIPS_ , 2019.
  * Doucet (2010) Doucet, A.  A note on efficient conditional simulation of Gaussian distributions.  _Technical report, University of British Columbia_ , 2010.
  * Dusenberry et al. (2020) Dusenberry, M. W., Jerfel, G., Wen, Y., Ma, Y.-a., Snoek, J., Heller, K., Lakshminarayanan, B., and Tran, D.  Efficient and scalable Bayesian neural nets with rank-1 factors.  In _ICML_ , 2020.
  * Foong et al. (2019) Foong, A. Y., Li, Y., Hernández-Lobato, J. M., and Turner, R. E.  ’In-between’ uncertainty in Bayesian neural networks.  _arXiv preprint arXiv:1906.11537_ , 2019.
  * Frankle & Carbin (2018) Frankle, J. and Carbin, M.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  In _ICLR_ , 2018.
  * Gal & Ghahramani (2016) Gal, Y. and Ghahramani, Z.  Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.  In _ICML_ , 2016.
  * Ghosh et al. (2019) Ghosh, S., Yao, J., and Doshi-Velez, F.  Model selection in Bayesian neural networks via horseshoe priors.  _JMLR_ , 2019.
  * Graves (2011) Graves, A.  Practical variational inference for neural networks.  In _NeurIPS_ , 2011.
  * Guo et al. (2017) Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q.  On calibration of modern neural networks.  In _ICML_ , 2017.
  * Gupta & Nagar (2018) Gupta, A. K. and Nagar, D. K.  _Matrix variate distributions_.  CRC Press, 2018.
  * Han et al. (2016) Han, S., Mao, H., and Dally, W. J.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.  In _ICLR_ , 2016.
  * Havasi et al. (2019) Havasi, M., Peharz, R., and Hernández-Lobato, J. M.  Minimal random code learning: Getting bits back from compressed model parameters.  In _ICLR_ , 2019.
  * He et al. (2016a) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _CVPR_ , 2016a.
  * He et al. (2016b) He, K., Zhang, X., Ren, S., and Sun, J.  Identity mappings in deep residual networks.  In _ECCV_ , 2016b.
  * Hendrycks & Dietterich (2019) Hendrycks, D. and Dietterich, T.  Benchmarking neural network robustness to common corruptions and perturbations.  In _ICLR_ , 2019.
  * 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 _COLT_ , 1993.
  * Hoffman & Gelman (2014) Hoffman, M. D. and Gelman, A.  The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo.  _JMLR_ , 2014.
  * Hoffman et al. (2013) Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J.  Stochastic variational inference.  _JMLR_ , 2013.
  * Hoffman & Ribak (1991) Hoffman, Y. and Ribak, E.  Constrained realizations of Gaussian fields-a simple algorithm.  _ApJ_ , 1991.
  * Huang et al. (2017) Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q.  Densely connected convolutional networks.  In _CVPR_ , 2017.
  * Izmailov et al. (2019) Izmailov, P., Maddox, W. J., Kirichenko, P., Garipov, T., Vetrov, D., and Wilson, A. G.  Subspace inference for Bayesian deep learning.  In _UAI_ , 2019.
  * Jordan et al. (1999) Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, L. K.  An introduction to variational methods for graphical models.  _Machine Learning_ , 1999.
  * Journel & Huijbregts (1978) Journel, A. G. and Huijbregts, C. J.  _Mining geostatistics_.  Academic press London, 1978.
  * Karaletsos & Bui (2020) Karaletsos, T. and Bui, T. D.  Hierarchical Gaussian process priors for Bayesian neural network weights.  In _NeurIPS_ , 2020.
  * Kendall & Gal (2017) Kendall, A. and Gal, Y.  What uncertainties do we need in Bayesian deep learning for computer vision?  In _NeurIPS_ , 2017.
  * Khan et al. (2019) Khan, M. E. E., Immer, A., Abedi, E., and Korzepa, M.  Approximate inference turns deep networks into Gaussian processes.  In _NeurIPS_ , 2019.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _ICLR_ , 2015.
  * Kingma & Welling (2014) Kingma, D. P. and Welling, M.  Auto-encoding variational Bayes.  In _ICLR_ , 2014.
  * Krizhevsky et al. (2009) Krizhevsky, A., Nair, V., and Hinton, G.  CIFAR-10 and CIFAR-100 datasets.  _URL: https://www. cs. toronto. edu/kriz/cifar. html_ , 2009.
  * Krueger et al. (2017) Krueger, D., Huang, C.-W., Islam, R., Turner, R., Lacoste, A., and Courville, A.  Bayesian hypernetworks.  _arXiv preprint arXiv:1710.04759_ , 2017.
  * Lakshminarayanan et al. (2017) Lakshminarayanan, B., Pritzel, A., and Blundell, C.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _NeurIPS_ , 2017.
  * Lee et al. (2018) Lee, J., Sohl-Dickstein, J., Pennington, J., Novak, R., Schoenholz, S., and Bahri, Y.  Deep neural networks as Gaussian processes.  In _ICLR_ , 2018.
  * Lee et al. (2019) Lee, N., Ajanthan, T., and Torr, P.  SNIP: Single-shot network pruning based on connection sensitivity.  In _ICLR_ , 2019.
  * Leibfried et al. (2020) Leibfried, F., Dutordoir, V., John, S., and Durrande, N.  A tutorial on sparse Gaussian processes and variational inference.  _arXiv preprint arXiv:2012.13962_ , 2020.
  * Louizos & Welling (2017) Louizos, C. and Welling, M.  Multiplicative normalizing flows for variational Bayesian neural networks.  In _ICML_ , 2017.
  * Louizos et al. (2017) Louizos, C., Ullrich, K., and Welling, M.  Bayesian compression for deep learning.  In _NeurIPS_ , 2017.
  * Ma et al. (2019) Ma, C., Li, Y., and Hernández-Lobato, J. M.  Variational implicit processes.  In _ICML_ , 2019.
  * MacKay (1995) MacKay, D. J.  Bayesian neural networks and density networks.  _NIMPR A_ , 1995.
  * Matthews et al. (2018) Matthews, A. G. d. G., Hron, J., Rowland, M., Turner, R. E., and Ghahramani, Z.  Gaussian process behaviour in wide deep neural networks.  In _ICLR_ , 2018.
  * Mishkin et al. (2018) Mishkin, A., Kunstner, F., Nielsen, D., Schmidt, M., and Khan, M. E.  SLANG: Fast structured covariance approximations for Bayesian deep learning with natural gradient.  In _NeurIPS_ , 2018.
  * Neal (1995) Neal, R. M.  _Bayesian Learning for Neural Networks_.  PhD thesis, University of Toronto, 1995.
  * Ober & Aitchison (2021) Ober, S. W. and Aitchison, L.  Global inducing point variational posteriors for Bayesian neural networks and deep Gaussian processes.  In _ICML_ , 2021.
  * Ovadia et al. (2019) Ovadia, Y., Fertig, E., Ren, J., Nado, Z., Sculley, D., Nowozin, S., Dillon, J., Lakshminarayanan, B., and Snoek, J.  Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift.  In _NeurIPS_ , 2019.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.  PyTorch: An imperative style, high-performance deep learning library.  In _NeurIPS_ , 2019.
  * Ranganath et al. (2016) Ranganath, R., Tran, D., and Blei, D.  Hierarchical variational models.  In _ICML_ , 2016.
  * Ritter et al. (2018) Ritter, H., Botev, A., and Barber, D.  A scalable Laplace approximation for neural networks.  In _ICLR_ , 2018.
  * Salimans et al. (2015) Salimans, T., Kingma, D., and Welling, M.  Markov chain Monte Carlo and variational inference: Bridging the gap.  In _ICML_ , 2015.
  * Salimbeni & Deisenroth (2017) Salimbeni, H. and Deisenroth, M.  Doubly stochastic variational inference for deep Gaussian processes.  In _NeurIPS_ , 2017.
  * Snelson & Ghahramani (2006) Snelson, E. and Ghahramani, Z.  Sparse Gaussian processes using pseudo-inputs.  In _NeurIPS_ , 2006.
  * Stegle et al. (2011) Stegle, O., Lippert, C., Mooij, J. M., Lawrence, N. D., and Borgwardt, K.  Efficient inference in matrix-variate gaussian models with iid observation noise.  In _NeurIPS_ , 2011.
  * Sun et al. (2019) Sun, S., Zhang, G., Shi, J., and Grosse, R.  Functional variational Bayesian neural networks.  In _ICLR_ , 2019.
  * Świątkowski et al. (2020) Świątkowski, J., Roth, K., Veeling, B. S., Tran, L., Dillon, J. V., Mandt, S., Snoek, J., Salimans, T., Jenatton, R., and Nowozin, S.  The k-tied Normal distribution: A compact parameterization of Gaussian mean field posteriors in Bayesian neural networks.  In _ICML_ , 2020.
  * Tanno et al. (2017) Tanno, R., Worrall, D. E., Ghosh, A., Kaden, E., Sotiropoulos, S. N., Criminisi, A., and Alexander, D. C.  Bayesian image quality transfer with CNNs: exploring uncertainty in dMRI super-resolution.  In _MICCAI_ , 2017.
  * Tipping & Bishop (1999) Tipping, M. E. and Bishop, C. M.  Probabilistic principal component analysis.  _JRSS B_ , 1999.
  * Titsias (2009) Titsias, M.  Variational learning of inducing variables in sparse Gaussian processes.  In _AISTATS_ , 2009.
  * Titsias & Lázaro-Gredilla (2014) Titsias, M. and Lázaro-Gredilla, M.  Doubly stochastic variational Bayes for non-conjugate inference.  In _ICML_ , 2014.
  * Welling & Teh (2011) Welling, M. and Teh, Y. W.  Bayesian learning via stochastic gradient langevin dynamics.  In _ICML_ , 2011.
  * Wen et al. (2020) Wen, Y., Tran, D., and Ba, J.  Batchensemble: an alternative approach to efficient ensemble and lifelong learning.  In _ICLR_ , 2020.
  * Wilson et al. (2020) Wilson, J. T., Borovitskiy, V., Terenin, A., Mostowsky, P., and Deisenroth, M. P.  Efficiently sampling functions from Gaussian process posteriors.  In _ICML_ , 2020.
  * Wilson et al. (2021) Wilson, J. T., Borovitskiy, V., Terenin, A., Mostowsky, P., and Deisenroth, M. P.  Pathwise conditioning of Gaussian processes.  _JMLR_ , 2021.
  * Xu et al. (2019) Xu, H., Liu, B., Shu, L., and Yu, P.  BERT post-training for review reading comprehension and aspect-based sentiment analysis.  In _ACL_ , 2019.
  * Zhang et al. (2018a) Zhang, C., Bütepage, J., Kjellström, H., and Mandt, S.  Advances in variational inference.  _TPAMI_ , 2018a.
  * Zhang et al. (2018b) Zhang, G., Sun, S., Duvenaud, D., and Grosse, R.  Noisy natural gradient as variational inference.  In _ICML_ , 2018b.
  * Zhang et al. (2020) Zhang, R., Li, C., Zhang, J., Chen, C., and Wilson, A. G.  Cyclical stochastic gradient MCMC for Bayesian deep learning.  In _ICLR_ , 2020.

##  Appendix A Notation

Generally, we use regular font letters $a,A$ to denote scalars, lower-case
bold symbols ${\mathbf{a}}$ to denote vectors and upper-case bold symbols
${\mathbf{A}}$ for matrices.

Table 4: Overview of notation and matrix shapes Random variables: |
---|---
${\mathbf{W}}:d_{out}\times d_{in}$ |  Weight matrix of a layer
${\mathbf{U}}:M_{out}\times M_{in}$ |  Inducing weight matrix
${\mathbf{U}}_{r}:M_{out}\times d_{in}$ |  Inducing row matrix
${\mathbf{U}}_{c}:d_{out}\times M_{in}$ |  Inducing column matrix
Parameters: |
${\bm{Z}}_{r}:M_{out}\times d_{out},{\bm{Z}}_{c}:M_{in}\times d_{in}$ |  Inducing covariance parameters
${\bm{D}}_{r}:M_{out}\times M_{out},{\bm{D}}_{c}:M_{in}\times M_{in}$ |  Inducing precision parameters (we use diagonal matrices in our approach)
$\lambda^{2}$ |  Conditional rescaling factor ($<1$ reduces variance of weight distribution)
Hyperparameters: |
$\sigma^{2}$ |  Prior variance
$\sigma^{2}_{max}$ |  Maximum approximate posterior variance
$\lambda^{2}_{max}$ |  Maximum conditional rescaling
Other variables: |
$\bm{\Sigma}_{r}:(d_{out}+M_{out})\times(d_{out}+M_{out})$ |  Joint row covariance of ${\mathbf{W}}$ and ${\mathbf{U}}$
$\bm{\Sigma}_{c}:(d_{in}+M_{in})\times(d_{in}+M_{in})$ |  Joint column covariance of ${\mathbf{W}}$ and ${\mathbf{U}}$
${\bm{L}}_{r},{\bm{L}}_{c}$ |  Lower Cholesky factor of the joint row/column covariance
$\bm{\Psi}_{r}:M_{out}\times M_{out},\bm{\Psi}_{c}:M_{in}\times M_{in}$ |  Marginal row/column inducing covariance

##  Appendix B Derivations of the auxiliary variational objective

When the variational distribution is constructed as a mixture distribution
$q({\bm{\theta}})=\int
q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})d{\mathbf{a}}$, the original
variational lower-bound becomes intractable:

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

However, notice that we can also rewrite the “marginal” $q({\bm{\theta}})$
using Bayes’ rule:
$q({\bm{\theta}})=\frac{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}{q({\mathbf{a}}|{\bm{\theta}})}$,
meaning that the variational lower-bound can be re-formulated as

$\mathcal{L}(q({\bm{\theta}}))=\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}[\log
p(\mathcal{D}|{\bm{\theta}})]+\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{p({\bm{\theta}})q({\mathbf{a}}|{\bm{\theta}})}{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}\right].$
(14)

For many flexible mixture distributions, $q({\mathbf{a}}|{\bm{\theta}})$
remains intractable. Fortunately, notice that for _any_ distribution
$r({\mathbf{a}}|{\bm{\theta}})$ with the same support as
$q({\mathbf{a}}|{\bm{\theta}})$, we have:

$\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{p({\bm{\theta}})q({\mathbf{a}}|{\bm{\theta}})}{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}\right]=\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{p({\bm{\theta}})r({\mathbf{a}}|{\bm{\theta}})}{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}\right]+\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{q({\mathbf{a}}|{\bm{\theta}})}{r({\mathbf{a}}|{\bm{\theta}})}\right],$
(15)

and more importantly, the second term on the RHS of the above equation
satisfies

$\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{q({\mathbf{a}}|{\bm{\theta}})}{r({\mathbf{a}}|{\bm{\theta}})}\right]=\mathbb{E}_{q({\bm{\theta}})}[\operatorname{\mathbb{KL}}[q({\mathbf{a}}|{\bm{\theta}})||r({\mathbf{a}}|{\bm{\theta}})]]\geq
0.$ (16)

This means we can remove this KL term and construct a lower-bound to the
variational lower-bound:

$\log{}p(\mathcal{D})\geq\mathcal{L}(q({\bm{\theta}}))\geq\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}[\log
p(\mathcal{D}|{\bm{\theta}})]+\mathbb{E}_{q({\bm{\theta}},{\mathbf{a}})}\left[\log\frac{p({\bm{\theta}})r({\mathbf{a}}|{\bm{\theta}})}{q({\bm{\theta}}|{\mathbf{a}})q({\mathbf{a}})}\right]:=\mathcal{L}(q({\bm{\theta}},{\mathbf{a}})),$
(17)

which corresponds to the auxiliary variational lower-bound Eq. 2 presented in
the main text. This auxiliary bound can be improved by optimising
$r({\mathbf{a}}|{\bm{\theta}})$ towards better approximating
$q({\mathbf{a}}|{\bm{\theta}})$, and it recovers the original variational
lower-bound iff.
$r({\mathbf{a}}|{\bm{\theta}})=q({\mathbf{a}}|{\bm{\theta}})$. Still we
emphasise that the auxiliary bound is valid for _any_
$r({\mathbf{a}}|{\bm{\theta}})$ with the same support as
$q({\mathbf{a}}|{\bm{\theta}})$, which enables our design presented in the
main text to improve memory efficiency.

##  Appendix C An introduction to SVGP

This section provides a brief introduction to sparse variational approximation
for variationally sparse GP (SVGP). We use regression as a running example,
but the principles of SVGP also apply to other supervised learning tasks such
as classification. Readers are also referred to e.g. Leibfried et al. (2020)
for a modern tutorial.

Assume we have a regression dataset $\mathcal{D}=\\{\mathbf{X},\mathbf{Y}\\}$
where $\mathbf{X}=[\bm{x}_{1},...,\bm{x}_{N}]$ and
$\mathbf{Y}=[\bm{y}_{1},...,\bm{y}_{N}]$. In GP regression we build the
following probabilistic model to address this regression task:

$\bm{y}_{n}=f(\bm{x}_{n})+\bm{\epsilon}_{n},\quad\bm{\epsilon}_{n}\sim\mathcal{N}(\bm{0},\sigma^{2}\mathbf{I}),\quad
f(\cdot)\sim\mathcal{GP}(0,\mathcal{K}(\cdot,\cdot)).$ (18)

Here we put on the regression function $f(\cdot)$ a GP prior with zero mean
function and covariance function defined by kernel $\mathcal{K}(\cdot,\cdot)$.
In practice we can only evaluate the function $f(\cdot)$ on finite number of
inputs, but fortunately by construction, GP allows sampling of function values
from a joint Gaussian distribution. In detail, we can sample from the
following Gaussian distribution to get the function value samples
$\mathbf{f}=[f(\bm{x}_{1}),...,f(\bm{x}_{N})]$ given the input locations
$\mathbf{X}$:

$\mathbf{f}\sim\mathcal{N}(\bm{0},\mathbf{K}_{\mathbf{X}\mathbf{X}}),\quad\mathbf{K}_{\mathbf{X}\mathbf{X}}(m,n)=\mathcal{K}(\bm{x}_{m},\bm{x}_{n}).$
(19)

Therefore we can rewrite the probabilistic model in finite dimension as

$p(\mathbf{Y}|\mathbf{X},f(\cdot))=p(\mathbf{Y}|\mathbf{f})=\prod_{n=1}^{N}\mathcal{N}(\bm{y}_{n};f(\bm{x}_{n}),\sigma^{2}\mathbf{I}),\quad
p(\mathbf{f}|\mathbf{X})=\mathcal{N}(\mathbf{f};\bm{0},\mathbf{K}_{\mathbf{X}\mathbf{X}}).$
(20)

Note that the GP prior can be extended to a larger set of inputs
$\mathbf{X}\cup\mathbf{X}^{*}$ as
$p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{X}^{*})=\mathcal{N}([\mathbf{f},\mathbf{f}^{*}];\bm{0},\mathbf{K}_{[\mathbf{X},\mathbf{X}^{*}],[\mathbf{X},\mathbf{X}^{*}]})$,
where $\mathbf{f}^{*}$ denotes the function values given the inputs
$\mathbf{X}^{*}$, and

$\displaystyle\mathbf{K}_{[\mathbf{X},\mathbf{X}^{*}],[\mathbf{X},\mathbf{X}^{*}]}$
$\displaystyle=\begin{pmatrix}\mathbf{K}_{\mathbf{X}\mathbf{X}}&\mathbf{K}_{\mathbf{X}\mathbf{X}^{*}}\\\
\mathbf{K}_{\mathbf{X}^{*}\mathbf{X}}&\mathbf{K}_{\mathbf{X}^{*}\mathbf{X}^{*}}\end{pmatrix}.$

Importantly, this definition leaves the marginal prior unchanged ($\int
p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{X}^{*})d\mathbf{f}^{*}=p(\mathbf{f}|\mathbf{X})$),
and the conditional distribution
$p(\mathbf{f}|\mathbf{f}^{*},\mathbf{X},\mathbf{X}^{*})$ is also Gaussian.
This means predictive inference can be done in the following way: given test
inputs $\mathbf{X}^{*}$, the posterior predictive for $\mathbf{f}^{*}$ is

$p(\mathbf{Y}|\mathbf{X},\mathbf{X}^{*})=\int
p(\mathbf{Y}|\mathbf{f})p(\mathbf{f}|\mathbf{X})p(\mathbf{f}^{*}|\mathbf{f},\mathbf{X},\mathbf{X}^{*})d\mathbf{f}d\mathbf{f}^{*}=\int
p(\mathbf{Y}|\mathbf{f})p(\mathbf{f}|\mathbf{X})d\mathbf{f}=p(\mathbf{Y}|\mathbf{X}),$
(21) $\displaystyle p(\mathbf{f}^{*}|\mathbf{X},\mathbf{Y},\mathbf{X}^{*})$
$\displaystyle=\int\frac{p(\mathbf{Y}|\mathbf{f})p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{X}^{*})}{p(\mathbf{Y}|\mathbf{X},\mathbf{X}^{*})}d\mathbf{f}$
(22)
$\displaystyle=\int\frac{p(\mathbf{Y}|\mathbf{f})p(\mathbf{f}|\mathbf{X})p(\mathbf{f}^{*}|\mathbf{f},\mathbf{X},\mathbf{X}^{*})}{p(\mathbf{Y}|\mathbf{X})}d\mathbf{f}$
$\displaystyle=\int\underbrace{p(\mathbf{f}|\mathbf{X},\mathbf{Y})p(\mathbf{f}^{*}|\mathbf{f},\mathbf{X},\mathbf{X}^{*})}_{=p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{Y},\mathbf{X}^{*})}d\mathbf{f}.$

Unfortunately the exact posterior $p(\mathbf{f}|\mathbf{X},\mathbf{Y})$ is
intractable even for GP regression. Although in such case
$p(\mathbf{f}|\mathbf{X},\mathbf{Y})$ is Gaussian, evaluating this posterior
requires inverting/decomposing an $N\times N$ covariance matrix which has time
complexity $\mathcal{O}(N^{3})$. Also the storage cost of $\mathcal{O}(N^{2})$
for the posterior covariance is prohibitively expensive when $N$ is large.

To address the intractability issue, we seek to define an _approximate
posterior_ $q(f(\cdot))$, so that we can evaluate it on
$\mathbf{X}\cup\mathbf{X}^{*}$ and approximate the posterior as
$p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{Y},\mathbf{X}^{*})\approx
q(\mathbf{f},\mathbf{f}^{*})$. SVGP (Snelson & Ghahramani, 2006; Titsias,
2009) defines such approximate posterior by introducing _inducing_ inputs and
outputs. Again this is done by noticing that we can extend the GP prior to an
even larger set of inputs $\mathbf{X}\cup\mathbf{X}^{*}\cup\mathbf{Z}$, where
$\mathbf{Z}=[\bm{z}_{1},...,\bm{z}_{M}]$ are called _inducing inputs_ , and
the corresponding function values
$\mathbf{u}=[f(\bm{z}_{1}),...,f(\bm{z}_{M})]$ are named _inducing outputs_ :

$\displaystyle
p(\mathbf{f},\mathbf{f}^{*},\mathbf{u}|\mathbf{X},\mathbf{X}^{*},\mathbf{Z})$
$\displaystyle=\mathcal{N}([\mathbf{f},\mathbf{f}^{*},\bm{u}];\bm{0},\mathbf{K}_{[\mathbf{X},\mathbf{X}^{*},\mathbf{Z}],[\mathbf{X},\mathbf{X}^{*},\mathbf{Z}]})$
(23)
$\displaystyle=p(\mathbf{u}|\mathbf{Z})p(\mathbf{f},\mathbf{f}^{*}|\mathbf{u},\mathbf{X},\mathbf{X}^{*},\mathbf{Z}).$

Importantly, marginal consistency still holds for any $\mathbf{Z}$: $\int
p(\mathbf{f},\mathbf{f}^{*},\mathbf{u}|\mathbf{X},\mathbf{X}^{*},\mathbf{Z})d\mathbf{u}=p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{X}^{*})$
(c.f. eq.(3) in the main text). Furthermore, the conditional distribution
$p(\mathbf{f},\mathbf{f}^{*}|\mathbf{u},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})$
is a Gaussian distribution. Observing these, the SVGP approach defines the
approximate posterior as

$q(\mathbf{f},\mathbf{f}^{*},\mathbf{u})=p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{X}^{*},\mathbf{Z}^{*})q(\mathbf{u}),\quad
q(\mathbf{f},\mathbf{f}^{*})=\int
q(\mathbf{f},\mathbf{f}^{*},\mathbf{u})d\mathbf{u},$ (24)

and minimises an upper-bound of the KL divergence to find the optimal
$q(\mathbf{f},\mathbf{f}^{*})$:

$\displaystyle\operatorname{\mathbb{KL}}[q(\mathbf{f},\mathbf{f}^{*})||p(\mathbf{f},\mathbf{f}^{*}|\mathbf{X},\mathbf{Y},\mathbf{X}^{*})]$
$\displaystyle=\mathbb{E}_{q(\mathbf{f},\mathbf{f}^{*})}\left[\log\frac{q(\mathbf{f},\mathbf{f}^{*})p(\mathbf{Y}|\mathbf{X})}{p(\mathbf{f},\mathbf{f}^{*},\mathbf{Y}|\mathbf{X},\mathbf{X}^{*})}\right]$
(25) $\displaystyle=\log p(\mathbf{Y}|\mathbf{X})+$
$\displaystyle\mathbb{E}_{q(\mathbf{f},\mathbf{f}^{*},\mathbf{u})}\left[\log\frac{q(\mathbf{f},\mathbf{f}^{*},\mathbf{u})}{p(\mathbf{Y}|\mathbf{f})p(\mathbf{f},\mathbf{f}^{*},\mathbf{u}|\mathbf{X},\mathbf{X}^{*},\mathbf{Z})}+\frac{p(\mathbf{u}|\mathbf{f},\mathbf{f}^{*},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})}{q(\mathbf{u}|\mathbf{f},\mathbf{f}^{*})}\right]$
$\displaystyle=\log p(\mathbf{Y}|\mathbf{X})+$
$\displaystyle\mathbb{E}_{q(\mathbf{f},\mathbf{f}^{*},\mathbf{u})}\left[\log\frac{\cancel{p(\mathbf{f},\mathbf{f}^{*}|\mathbf{u},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})}q(\mathbf{u})}{p(\mathbf{Y}|\mathbf{f})\cancel{p(\mathbf{f},\mathbf{f}^{*}|\mathbf{u},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})}p(\mathbf{u}|\mathbf{Z})}\right]$
$\displaystyle-$
$\displaystyle\mathbb{E}_{q(\mathbf{f},\mathbf{f}^{*})}\operatorname{\mathbb{KL}}[q(\mathbf{u}|\mathbf{f},\mathbf{f}^{*})||p(\mathbf{u}|\mathbf{f},\mathbf{f}^{*},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})]$
$\displaystyle\leq\log p(\mathbf{Y}|\mathbf{X})-$
$\displaystyle\\{\underbrace{\mathbf{E}_{q(\mathbf{f})}[\log
p(\mathbf{Y}|\mathbf{f})]-\operatorname{\mathbb{KL}}[q(\mathbf{u})||p(\mathbf{u}|\mathbf{Z})]}_{:=\mathcal{L}(q(\mathbf{f},\mathbf{u})}\\}.$

As KL divergences are non-negative, the above derivations also means $\log
p(\mathbf{Y}|\mathbf{X})\geq\mathcal{L}(q(\mathbf{f},\mathbf{u}))$, and we can
optimise the parameters in $q$ to tighten the lower-bound. These variational
parameters include the distributional parameters of $q(\mathbf{u})$ (e.g. mean
and covariance if $q(\mathbf{u})$ is Gaussian), as well as the inducing inputs
$\mathbf{Z}$, since the variational lower-bound
$\mathcal{L}(q(\mathbf{f},\mathbf{u}))$ is valid for any settings of
$\mathbf{Z}$. The lower-bound requires evaluating

$q(\mathbf{f})=\int
p(\mathbf{f},\mathbf{f}^{*}|\mathbf{u},\mathbf{X},\mathbf{X}^{*},\mathbf{Z})q(\mathbf{u})d\mathbf{f}^{*}d\mathbf{u}=\int
p(\mathbf{f}|\mathbf{u},\mathbf{X},\mathbf{Z})q(\mathbf{u})d\mathbf{u},$

which can be done efficiently when $M<<N$, as evaluating the conditional
Gaussian $p(\mathbf{f}|\mathbf{u},\mathbf{X},\mathbf{Z})$ has
$\mathcal{O}(NM^{2}+M^{3})$ run-time cost. Similarly, once $q$ is optimised,
in prediction time one can directly sample from $q(\mathbf{f}^{*})$ by
computing $q(\mathbf{f}^{*})=\int
p(\mathbf{f}^{*}|\mathbf{u},\mathbf{X}^{*},\mathbf{Z})q(\mathbf{u})d\mathbf{u}$,
and by caching the inverse/decomposition of both the covariacne matrix of
$q(\mathbf{u})$ and $\mathbf{K}_{\mathbf{Z}\mathbf{Z}}$, predictive inference
can be approximated efficiently.

Using shorthand notations by dropping $\mathbf{Z}$, e.g.
$p(\mathbf{u})=p(\mathbf{u}|\mathbf{Z})$ and
$p(\mathbf{f}|\mathbf{X},\mathbf{u})=p(\mathbf{f}|\mathbf{u},\mathbf{X},\mathbf{Z})$,
returns the desired results discussed in Section 2 of the main text, if we set
$a=\mathbf{u}$ and $\theta=\mathbf{f}$. In fact from the above discussions, we
see that $\theta$ in GP inference is infinite dimensional:
$\theta=f_{\neq\mathbf{u}}$, as we can extend the finite collection of
function values to include both $\mathbf{f}$ and $\mathbf{f}^{*}$ for any
$\mathbf{X}^{*}\neq\mathbf{Z}$. By tying the conditional distribution given
$\mathbf{u}$ in both the GP prior and the approximate posterior $q(f(\cdot))$,
the posterior belief updates are “compressed” into $\mathbf{u}$ space, which
is also reflected by the name “sparse approximation” of the approach.

##  Appendix D Derivations of the augmented (pseudo) prior

###  D.1 Inducing auxiliary variables: multivariate Gaussian case

Suppose each weight matrix has an isotropic Gaussian prior with zero mean,
i.e.
$\operatorname{vec}({\mathbf{W}})\sim\operatorname{\mathcal{N}}(0,\sigma^{2}{\bm{I}})$
where $\operatorname{vec}$ concatenates the columns of a matrix into a vector
and $\sigma$ is the standard deviation. Augmenting this Gaussian with an
auxiliary variable ${\mathbf{U}}$ that also has a mean of zero and some
covariance that we are free to parameterise, the joint distribution is

$\displaystyle\begin{pmatrix}\operatorname{vec}({\mathbf{W}})\\\
\operatorname{vec}({\mathbf{U}})\end{pmatrix}$
$\displaystyle\sim\operatorname{\mathcal{N}}(0,\bm{\Sigma})\quad\text{with}\quad
L=\begin{pmatrix}\sigma{\bm{I}}&0\\\ {\bm{Z}}&{\bm{D}}\end{pmatrix}$
$\displaystyle\quad\text{s.t.}\quad\bm{\Sigma}$
$\displaystyle={\bm{L}}{\bm{L}}^{\top}=\begin{pmatrix}\sigma^{2}{\bm{I}}&\sigma{\bm{Z}}^{\top}\\\
\sigma{\bm{Z}}&{\bm{Z}}{\bm{Z}}^{\top}+{\bm{D}}^{2}\end{pmatrix}$

where ${\bm{D}}$ is a positive diagonal matrix and ${\bm{Z}}$ a matrix with
arbitrary entries. Through defining the Cholesky decomposition of
$\bm{\Sigma}$ we ensure its positive definiteness. By the usual rules of
Gaussian marginalisation, the augmented model leaves the marginal prior on
${\mathbf{W}}$ unchanged. Further, we can analytically derive the conditional
distribution on the weights given the inducing weights:

$\displaystyle
p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$
$\displaystyle=\operatorname{\mathcal{N}}({\bm{\mu}}_{{\mathbf{W}}|{\mathbf{U}}},\bm{\Sigma}_{{\mathbf{W}}|{\mathbf{U}}}),$
(26) $\displaystyle{\bm{\mu}}_{{\mathbf{W}}|{\mathbf{U}}}$
$\displaystyle=\sigma{\bm{Z}}^{\top}\bm{\Psi}^{-1}\operatorname{vec}(U),$ (27)
$\displaystyle\bm{\Sigma}_{{\mathbf{W}}|{\mathbf{U}}}$
$\displaystyle=\sigma^{2}({\bm{I}}-{\bm{Z}}^{\top}\bm{\Psi}^{-1}{\bm{Z}}),$
(28) $\displaystyle\bm{\Psi}$
$\displaystyle={\bm{Z}}{\bm{Z}}^{\top}+{\bm{D}}^{2}.$

For inference, we now need to define an approximate posterior over the joint
space $q({\mathbf{W}},{\mathbf{U}})$. We will do so by factorising it as
$q({\mathbf{W}},{\mathbf{U}})=q({\mathbf{W}}|{\mathbf{U}})q({\mathbf{U}})$.
Factorising the prior in the same way leads to the following KL term in the
ELBO:

$\operatorname{\mathbb{KL}}\left[q({\mathbf{W}},{\mathbf{U}})||p({\mathbf{W}},{\mathbf{U}})\right]=\operatorname{\mathbb{E}}_{q({\mathbf{U}})}\left[\operatorname{\mathbb{KL}}\left[q({\mathbf{W}}|{\mathbf{U}})||p({\mathbf{W}}|{\mathbf{U}}\right]\right]+\operatorname{\mathbb{KL}}\left[q({\mathbf{U}})||p({\mathbf{U}})\right]$
(29)

###  D.2 Inducing auxiliary variables: matrix normal case

Now we introduce the inducing variables in matrix space, and, in addition to
the inducing weight ${\mathbf{U}}$, we pad in two inducing matrices
${\mathbf{U}}_{r}$, ${\mathbf{U}}_{c}$, such that the full augmented prior is:

$\displaystyle\begin{pmatrix}{\mathbf{W}}&{\mathbf{U}}_{c}\\\
{\mathbf{U}}_{r}&{\mathbf{U}}\end{pmatrix}$ $\displaystyle\sim
p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}}):=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c}),$
(30) $\displaystyle\text{with}\quad{\bm{L}}_{r}$
$\displaystyle=\begin{pmatrix}\sigma_{r}{\bm{I}}&0\\\
{\bm{Z}}_{r}&{\bm{D}}_{r}\end{pmatrix}\quad\text{and}\quad{\bm{L}}_{c}=\begin{pmatrix}\sigma_{c}{\bm{I}}&0\\\
{\bm{Z}}_{c}&{\bm{D}}_{c}\end{pmatrix}$
$\displaystyle\text{s.t.}\quad\bm{\Sigma}_{r}$
$\displaystyle={\bm{L}}_{r}{\bm{L}}_{r}^{\top}=\begin{pmatrix}\sigma_{r}^{2}{\bm{I}}&\sigma_{r}{\bm{Z}}_{r}^{\top}\\\
\sigma_{r}{\bm{Z}}_{r}&{\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}+{\bm{D}}_{r}^{2}\end{pmatrix},$
$\displaystyle\text{and}\quad\bm{\Sigma}_{c}$
$\displaystyle={\bm{L}}_{c}{\bm{L}}_{c}^{\top}=\begin{pmatrix}\sigma_{c}^{2}{\bm{I}}&\sigma_{c}{\bm{Z}}_{c}^{\top}\\\
\sigma_{c}{\bm{Z}}_{c}&{\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}+{\bm{D}}_{c}^{2}\end{pmatrix}.$

Matrix normal distributions have similar marginalisation and conditioning
properties as multivariate Gaussians. As such, the marginal both over some set
of rows and some set of columns is still a matrix normal. Hence,
$p({\mathbf{W}})=\operatorname{\mathcal{MN}}(0,\sigma_{r}^{2}{\bm{I}},\sigma_{c}^{2}{\bm{I}})$,
and by choosing $\sigma_{r}\sigma_{c}=\sigma$ this matrix normal distribution
is equivalent to the multivariate normal
$p(\operatorname{vec}({\mathbf{W}}))=\operatorname{\mathcal{N}}(0,\sigma^{2}{\bm{I}})$.
Also
$p({\mathbf{U}})=\operatorname{\mathcal{MN}}(0,\bm{\Psi}_{r},\bm{\Psi}_{c})$,
where again $\bm{\Psi}_{r}={\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}+{\bm{D}}_{r}^{2}$
and $\bm{\Psi}_{c}={\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}+{\bm{D}}_{c}^{2}$.
Similarly, the conditionals on some rows or columns are matrix normal
distributed:

$\displaystyle{\mathbf{U}}_{c}|{\mathbf{U}}$
$\displaystyle\sim\operatorname{\mathcal{MN}}(\sigma_{r}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}},\sigma_{r}^{2}({\bm{I}}-{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}),\bm{\Psi}_{c}),$
$\displaystyle{\mathbf{U}}_{r}|{\mathbf{U}}$
$\displaystyle\sim\operatorname{\mathcal{MN}}({\mathbf{U}}\bm{\Psi}^{-1}_{c}\sigma_{c}{\bm{Z}}_{c},\bm{\Psi}_{r},\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c})),$
$\displaystyle{\mathbf{W}}|{\mathbf{U}}_{c}$
$\displaystyle\sim\operatorname{\mathcal{MN}}\left({\mathbf{U}}_{c}\bm{\Psi}_{c}^{-1}\sigma_{c}{\bm{Z}}_{c},\sigma_{r}^{2}{\bm{I}},\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}_{c}^{-1}{\bm{Z}}_{c})\right),$
$\displaystyle{\mathbf{W}},{\mathbf{U}}_{r}|{\mathbf{U}}_{c},{\mathbf{U}}$
$\displaystyle\sim\operatorname{\mathcal{MN}}(\begin{pmatrix}{\mathbf{U}}_{c}\\\
{\mathbf{U}}\end{pmatrix}\bm{\Psi}^{-1}_{c}\sigma_{c}{\bm{Z}}_{c},\bm{\Sigma}_{r},\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c})),$
(31)
$\displaystyle{\mathbf{W}}|{\mathbf{U}}_{r},{\mathbf{U}}_{c},{\mathbf{U}}$
$\displaystyle\sim\operatorname{\mathcal{MN}}(M_{\mathbf{W}},\sigma_{r}^{2}({\bm{I}}-{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}{\bm{Z}}_{r}),\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c})),$
$\displaystyle M_{\mathbf{W}}$
$\displaystyle=\sigma({\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}}_{r}+{\mathbf{U}}_{c}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}-{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}).$

##  Appendix E KL divergence for rescaled conditional weight distributions

For the conditional distribution on the weights, in the simplest case we set
$q({\mathbf{W}}|{\mathbf{U}})=p({\mathbf{W}}|{\mathbf{U}})$, hence the KL
divergence would be zero. For the most general case of arbitrary Gaussian
distributions with
$q=\operatorname{\mathcal{N}}({\bm{\mu}}_{q},\bm{\Sigma}_{q})$ and
$p=\operatorname{\mathcal{N}}({\bm{\mu}}_{p},\bm{\Sigma}_{p})$, the KL
divergence is:

$\operatorname{\mathbb{KL}}\left[q||p\right]=\frac{1}{2}(\log\frac{\det\bm{\Sigma}_{p}}{\det\bm{\Sigma}_{q}}-d+\operatorname{tr}(\bm{\Sigma}_{p}^{-1}\bm{\Sigma}_{q})+({\bm{\mu}}_{p}-{\bm{\mu}}_{q})^{\top}\bm{\Sigma}_{p}^{-1}({\bm{\mu}}_{p}-{\bm{\mu}}_{q})),$
(32)

where $d$ is the number of elements of ${\bm{\mu}}$. As motivated, it is
desirable to make $q({\mathbf{W}}|{\mathbf{U}})$ similar to
$p({\mathbf{W}}|{\mathbf{U}})$. We consider a scalar rescaling of the
covariance, i.e. for $p=\operatorname{\mathcal{N}}({\bm{\mu}},\bm{\Sigma})$ we
set $q=\operatorname{\mathcal{N}}({\bm{\mu}},\lambda^{2}\bm{\Sigma})$. This
leads to the final term, which is the Mahalanobis distance between the means
under $p$, cancelling out entirely and the log determinant and trace terms
becoming a function of $\lambda$ only: with
$d=\text{dim}(\operatorname{vec}({\mathbf{W}}))$,

$\displaystyle\operatorname{\mathbb{KL}}\left[q||p\right]$
$\displaystyle=\frac{1}{2}(\log\frac{\det\bm{\Sigma}}{\det\lambda^{2}\bm{\Sigma}}-d+\operatorname{tr}(\bm{\Sigma}^{-1}\lambda^{2}\bm{\Sigma}))$
$\displaystyle=\frac{1}{2}(\log\frac{\det\bm{\Sigma}}{\lambda^{2d}\det\bm{\Sigma}}-d+\operatorname{tr}(\lambda^{2}{\bm{I}}))$
$\displaystyle=\frac{1}{2}(-2d\log\lambda-d+d\lambda^{2})$
$\displaystyle=d(\frac{1}{2}\lambda^{2}-\log\lambda-\frac{1}{2}).$

##  Appendix F The extended Matheron’s rule to matrix normal distributions

The original Matheron’s rule (Journel & Huijbregts, 1978; Hoffman & Ribak,
1991; Doucet, 2010) for sampling conditional Gaussian variables states the
following. If the joint multivariate Gaussian distribution is

$\displaystyle\begin{pmatrix}\operatorname{vec}({\mathbf{W}})\\\
\operatorname{vec}({\mathbf{U}})\end{pmatrix}$ $\displaystyle\sim
p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}})):=\operatorname{\mathcal{N}}(\bm{0},\bm{\Sigma}),$
$\displaystyle\bm{\Sigma}$
$\displaystyle=\begin{pmatrix}\bm{\Sigma}_{{\mathbf{W}}{\mathbf{W}}}&\bm{\Sigma}_{{\mathbf{W}}{\mathbf{U}}}\\\
\bm{\Sigma}_{{\mathbf{U}}{\mathbf{W}}}&\bm{\Sigma}_{{\mathbf{U}}{\mathbf{U}}}\end{pmatrix},$

then, conditioned on ${\mathbf{U}}$, sampling ${\mathbf{W}}\sim
p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))$ can be
done as

$\displaystyle\operatorname{vec}({\mathbf{W}})$
$\displaystyle=\operatorname{vec}(\bar{{\mathbf{W}}})+\bm{\Sigma}_{{\mathbf{W}}{\mathbf{U}}}\bm{\Sigma}_{{\mathbf{U}}{\mathbf{U}}}^{-1}(\operatorname{vec}({\mathbf{U}})-\operatorname{vec}(\bar{{\mathbf{U}}})),$
$\displaystyle\operatorname{vec}(\bar{{\mathbf{W}}}),\operatorname{vec}(\bar{{\mathbf{U}}})$
$\displaystyle\sim\operatorname{\mathcal{N}}(\bm{0},\bm{\Sigma}).$

Matheron’s rule can provide significant speed-ups if
$\operatorname{vec}({\mathbf{U}})$ has significantly smaller dimensions than
that of $\operatorname{vec}({\mathbf{W}})$, and the Cholesky decomposition of
$\bm{\Sigma}$ can be computed with low costs (e.g. due to the specific
structure in $\bm{\Sigma}$). Recall from the main text that the augmented
prior is

$p(\operatorname{vec}({\mathbf{W}}),\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}\left(0,\begin{pmatrix}\sigma_{c}^{2}{\bm{I}}\otimes\sigma_{r}^{2}{\bm{I}}&\sigma_{c}{\bm{Z}}_{c}^{\top}\otimes\sigma_{r}{\bm{Z}}_{r}^{\top}\\\
\sigma_{c}{\bm{Z}}_{c}\otimes\sigma_{r}{\bm{Z}}_{r}&\bm{\Psi}_{c}\otimes\bm{\Psi}_{r}\end{pmatrix}\right),$
(33)

and the corresponding conditional distribution is:

$p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))=\operatorname{\mathcal{N}}(\sigma_{c}\sigma_{r}\operatorname{vec}({\bm{Z}}_{r}\bm{\Psi}^{-1}_{r}{\mathbf{U}}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}^{\top}),\sigma_{c}^{2}\sigma_{r}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}\otimes{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r})).$
(34)

Therefore, while $\text{dim}(\operatorname{vec}({\mathbf{U}}))$ is indeed
significantly smaller than of $\text{dim}(\operatorname{vec}({\mathbf{W}}))$
by construction, the joint covariance matrix does not support fast Cholesky
decompositions, meaning that Matheron’s rule for efficient sampling does not
directly apply here.

However, in the full augmented space, the joint distribution does have an
efficient matrix normal form:
$p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}}_{c})=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c})$.
Furthermore, the row and column covariance matrices $\bm{\Sigma}_{r}$ and
$\bm{\Sigma}_{c}$ are parameterised by their Cholesky decompositions, meaning
that sampling from the joint distribution
$p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}})$ can be done
in a fast way. Importantly, Cholesky decompositions for $p({\mathbf{U}})$’s
row and column covariance matrices $\bm{\Psi}_{r}$ and $\bm{\Psi}_{c}$ can be
computed in $\mathcal{O}(M_{out}^{3})$ and $\mathcal{O}(M_{in}^{3})$ time,
respectively, which are much faster than the multi-variate Gaussian case that
requires $\mathcal{O}(M_{in}^{3}M_{out}^{3})$ time. Observing these, we extend
Matheron’s rule to sample $p({\mathbf{W}}|{\mathbf{U}})$ where
$p({\mathbf{W}},{\mathbf{U}})$ is the marginal distribution of
$p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}}_{c})=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c})$.

In detail, for drawing a sample from $p({\mathbf{W}}|{\mathbf{U}})$ we need to
draw a sample from the joint $p({\mathbf{W}},{\mathbf{U}})$. To do so, we
sample from the augmented prior
$\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}}\sim
p(\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}})=\operatorname{\mathcal{MN}}(0,\bm{\Sigma}_{r},\bm{\Sigma}_{c})$,
computed using the Cholesky decompositions of $\bm{\Sigma}_{r}$ and
$\bm{\Sigma}_{c}$:

$\begin{pmatrix}\bar{{\mathbf{W}}}&\bar{{\mathbf{U}}}_{c}\\\
\bar{{\mathbf{U}}}_{r}&\bar{{\mathbf{U}}}\end{pmatrix}=\begin{pmatrix}\sigma_{r}{\bm{I}}&0\\\
{\bm{Z}}_{r}&{\bm{D}}_{r}\end{pmatrix}\begin{pmatrix}{\mathbf{E}}_{1}&{\mathbf{E}}_{2}\\\
{\mathbf{E}}_{3}&{\mathbf{E}}_{4}\end{pmatrix}\begin{pmatrix}\sigma_{c}{\bm{I}}&{\bm{Z}}_{c}^{\top}\\\
0&{\bm{D}}_{c}\end{pmatrix},$

where ${\mathbf{E}}_{1}\in\mathbb{R}^{d_{out}\times d_{in}}$,
${\mathbf{E}}_{2}\in\mathbb{R}^{d_{out}\times M_{in}}$,
${\mathbf{E}}_{3}\in\mathbb{R}^{M_{out}\times d_{in}}$,
${\mathbf{E}}_{4}\in\mathbb{R}^{M_{out}\times M_{in}}$ are standard Gaussian
noise samples, and $\bar{{\mathbf{W}}}\in\mathbb{R}^{d_{out}\times d_{in}}$
and $\bar{{\mathbf{U}}}\in\mathbb{R}^{M_{out}\times M_{in}}$. Then we
construct the conditional sample ${\mathbf{W}}\sim
p({\mathbf{W}}|{\mathbf{U}})$ as follows, similar to Matheron’s rule in the
multivariate Gaussian case:

${\mathbf{W}}=\bar{{\mathbf{W}}}+\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}(U-\bar{U})\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}.$
(35)

From the above equations we see that $\bar{{\mathbf{U}}_{r}}$ and
$\bar{{\mathbf{U}}_{c}}$ do not contribute to the final ${\mathbf{W}}$ sample.
Therefore we do not need to compute $\bar{{\mathbf{U}}}_{r}$ and
$\bar{{\mathbf{U}}}_{c}$, and we write the separate expressions for
$\bar{{\mathbf{W}}}$ and $\bar{{\mathbf{U}}}$ as:

$\bar{{\mathbf{W}}}=\sigma_{r}\sigma_{c}{\mathbf{E}}_{1},\quad\bar{{\mathbf{U}}}=\underbrace{{\bm{Z}}_{r}{\mathbf{E}}_{1}{\bm{Z}}_{c}^{\top}}_{\bar{{\mathbf{U}}}_{1}}+\underbrace{{\bm{Z}}_{r}{\mathbf{E}}_{2}{\bm{D}}_{c}}_{\bar{{\mathbf{U}}}_{2}}+\underbrace{{\bm{D}}_{r}{\mathbf{E}}_{3}{\bm{Z}}_{c}^{\top}}_{\bar{{\mathbf{U}}}_{3}}+\underbrace{{\bm{D}}_{r}{\mathbf{E}}_{4}{\bm{D}}_{c}}_{\bar{{\mathbf{U}}}_{4}}.$
(36)

Note that $\bar{{\mathbf{U}}}$ is a sum of four samples from matrix normal
distributions. In particular, we have that:

$\bar{{\mathbf{U}}}_{2}\overset{d}{\sim}\operatorname{\mathcal{MN}}(0,{\bm{Z}}_{r}{\bm{Z}}_{r}^{\top},{\bm{D}}_{c}^{2})\quad\text{and}\quad\bar{{\mathbf{U}}}_{3}\overset{d}{\sim}\operatorname{\mathcal{MN}}(0,{\bm{D}}_{r}^{2},{\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}).$

Hence instead of sampling the “long and thin” Gaussian noise matrices
${\mathbf{E}}_{2}$ and ${\mathbf{E}}_{3}$, we can reduce variance by sampling
standard Gaussian noise matrices
$\tilde{E}_{2},\tilde{E}_{3}\in\mathbb{R}^{M_{out}\times M_{in}}$, and
calculate $\bar{{\mathbf{U}}}$ as

$\bar{{\mathbf{U}}}={\bm{Z}}_{r}{\mathbf{E}}_{1}{\bm{Z}}_{c}^{\top}+\hat{L}_{r}\tilde{E}_{2}{\bm{D}}_{c}+{\bm{D}}_{r}\tilde{E}_{3}\hat{L}_{c}^{\top}+{\bm{D}}_{r}{\mathbf{E}}_{4}{\bm{D}}_{c}.$
(37)

This is enabled by calculating the Cholesky decompositions
$\hat{{\bm{L}}}_{r}\hat{{\bm{L}}}_{r}^{\top}={\bm{Z}}_{r}{\bm{Z}}_{r}^{\top}$
and
$\hat{{\bm{L}}}_{c}\hat{{\bm{L}}}_{c}^{\top}={\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}$,
which have $\mathcal{O}(M_{out}^{3})$ and $\mathcal{O}(M_{in}^{3})$ run-time
costs, respectively. As a reminder, the Cholesky factors are _square_
matrices, i.e. $\hat{{\bm{L}}}_{r}\in\mathbb{R}^{M_{out}\times M_{out}}$,
$\hat{{\bm{L}}}_{c}\in\mathbb{R}^{M_{in}\times M_{in}}$). We name the approach
the _extended Matheron’s rule_ for sampling conditional Gaussians when the
full joint has a matrix normal form.

To verify the proposed approach, we compute the mean and the variance of the
random variable ${\mathbf{W}}$ defined in Eq. 35, and check if they match the
mean and variance of Eq. 34. First as $\bar{{\mathbf{W}}},\bar{{\mathbf{U}}}$
have zero mean, it is straightforward to verify that
$\mathbb{E}[W]=\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}U\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}$
which matches the mean of Eq. 34. For the variance of
$\operatorname{vec}({\mathbf{W}})$, it requires computing the following terms:

$\displaystyle\mathbb{V}(\operatorname{vec}({\mathbf{W}}))=$
$\displaystyle\mathbb{V}(\operatorname{vec}(\bar{{\mathbf{W}}}))+\mathbb{V}(\operatorname{vec}(\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}\bar{U}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}))$
$\displaystyle-2\text{Cov}[\operatorname{vec}({\mathbf{W}}),\operatorname{vec}(\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}\bar{U}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c})]$
$\displaystyle=:$ $\displaystyle{\bm{A}}_{1}+{\bm{A}}_{2}-2{\bm{A}}_{3}.$ (38)

First it can be shown that

$\displaystyle{\bm{A}}_{1}$
$\displaystyle=\sigma_{r}^{2}\sigma_{c}^{2}{\bm{I}}\quad\text{since
}\bar{{\mathbf{W}}}\sim\operatorname{\mathcal{MN}}(0,\sigma_{r}^{2}{\bm{I}},\sigma_{c}^{2}{\bm{I}}),$
$\displaystyle{\bm{A}}_{2}$
$\displaystyle=\sigma_{r}^{2}\sigma_{c}^{2}{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}\otimes{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}$
$\displaystyle\text{as}\quad{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}\bar{U}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}$
$\displaystyle\sim\operatorname{\mathcal{MN}}(0,{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r},{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}).$

For the correlation term ${\bm{A}}_{3}$, we notice that $\bar{{\mathbf{W}}}$
and $\bar{{\mathbf{U}}}$ only share the noise matrix ${\mathbf{E}}_{1}$ in the
joint sampling procedure Eq. 36. This also means

$\displaystyle{\bm{A}}_{3}$
$\displaystyle=\sigma_{r}^{2}\sigma_{c}^{2}\text{Cov}[\operatorname{vec}({\mathbf{E}}_{1}),\operatorname{vec}({\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}{\mathbf{E}}_{1}{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c})]$
$\displaystyle=\sigma_{r}^{2}\sigma_{c}^{2}{\bm{Z}}_{c}^{\top}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}\otimes{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}.$

Plugging in ${\bm{A}}_{1},{\bm{A}}_{2},{\bm{A}}_{3}$ into Eq. 38 verifies that
$\mathbb{V}(\operatorname{vec}({\mathbf{W}}))$ matches the variance of the
conditional distribution
$p(\operatorname{vec}({\mathbf{W}})|\operatorname{vec}({\mathbf{U}}))$,
showing that the proposed extended Matheron’s rule indeed draws samples from
the conditional distribution.

As for sampling ${\mathbf{W}}$ from $q({\mathbf{W}}|{\mathbf{U}})$, since it
has the same mean but a rescaled covariance as compared with
$p({\mathbf{W}}|{\mathbf{U}})$, we can compute the samples by adapting the
extend Matheron’s rule as follows. Notice that the mean of ${\mathbf{W}}$ in
Eq. 35 is
$\mathbb{E}[{\mathbf{W}}|{\mathbf{U}}]=\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}U\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}$,
therefore by rearranging terms, Eq. 35 can be re-written as

$\displaystyle{\mathbf{W}}=$
$\displaystyle{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}+[\bar{{\mathbf{W}}}-\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}\bar{{\mathbf{U}}}\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c}]$
$\displaystyle:=$ $\displaystyle\text{mean}+\text{noise}.$

So sampling from $q({\mathbf{W}}|{\mathbf{U}})$ can be done by rescaling the
noise term in the above equation with the scale parameter $\lambda$. In
summary, the extended Matheron’s rule for sampling
$q({\mathbf{W}}|{\mathbf{U}})$ is as follows:

$\begin{split}{\mathbf{W}}&=\lambda\bar{{\mathbf{W}}}+\sigma_{r}\sigma_{c}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}(U-\lambda\bar{{\mathbf{U}}})\bm{\Psi}^{-1}_{c}{\bm{Z}}_{c},\\\
\bar{{\mathbf{W}}},\bar{{\mathbf{U}}}&\sim
p(\bar{{\mathbf{W}}},\bar{{\mathbf{U}}_{c}},\bar{{\mathbf{U}}_{r}},\bar{{\mathbf{U}}}).\end{split}$
(39)

Plugging in $\sigma_{r}\sigma_{c}=\sigma$ here returns the conditional
sampling rule Section 3.2 in the main text.

##  Appendix G Function-space view of inducing weights

Here we present the detailed derivations of Section 5.1. Assume a neural
network layer with weight ${\mathbf{W}}$ computes the following transformation
of the input
$\mathbf{X}=[\bm{x}_{1},...,\bm{x}_{N}],\bm{x}_{i}\in\mathbb{R}^{d_{in}\times
1}$:

$\mathbf{F}={\mathbf{W}}\mathbf{X},\
\mathbf{H}=g(\mathbf{F}),\quad{\mathbf{W}}\in\mathbb{R}^{d_{out}\times
d_{in}},\mathbf{X}\in\mathbb{R}^{d_{in}\times N},$

where $g(\cdot)$ is the non-linearity. Therefore the Gaussian prior
$p({\mathbf{W}})=\operatorname{\mathcal{N}}(0,\sigma^{2}{\bm{I}})$ induces a
Gaussian distribution on the linear transformation output $\mathbf{F}$, in
fact each of the rows in
$\mathbf{F}=[\mathbf{f}_{1},...,\mathbf{f}_{d_{out}}]^{\top},\mathbf{f}_{i}\in\mathbb{R}^{N\times
1}$ has a Gaussian process form with linear kernel:

$\mathbf{f}_{i}|\mathbf{X}\sim\mathcal{GP}(\bm{0},\mathbf{K}_{\mathbf{X}\mathbf{X}}),\quad\mathbf{K}_{\mathbf{X}\mathbf{X}}(m,n)=\sigma^{2}\bm{x}_{m}^{\top}\bm{x}_{n}.$
(40)

Performing inference on $\mathbf{F}$ directly has
$\mathcal{O}(N^{3}+d_{out}N^{2})$ cost, so a sparse approximation is needed.
Slightly different from the usual approach, we introduce “scaled noisy
inducing outputs”
${\mathbf{U}}_{c}=[\mathbf{u}_{1}^{c},...,\mathbf{u}^{c}_{d_{out}}]^{\top}\in\mathbb{R}^{d_{out}\times
M_{in}}$ in the following way, using shared inducing inputs
${\bm{Z}}_{c}^{\top}\in\mathbb{R}^{d_{in}\times M_{in}}$:

$\displaystyle p(\mathbf{f}_{i},\hat{\mathbf{u}}_{i}|\mathbf{X})$
$\displaystyle=\mathcal{GP}\left(\bm{0},\begin{pmatrix}\mathbf{K}_{\mathbf{X}\mathbf{X}}&\mathbf{K}_{\mathbf{X}{\bm{Z}}_{c}}\\\
\mathbf{K}_{{\bm{Z}}_{c}\mathbf{X}}&\mathbf{K}_{{\bm{Z}}_{c}{\bm{Z}}_{c}}\end{pmatrix}\right),$
$\displaystyle p(\mathbf{u}^{c}_{i}|\hat{\mathbf{u}}_{i})$
$\displaystyle=\operatorname{\mathcal{N}}\left(\frac{\hat{\mathbf{u}}_{i}}{\sigma_{c}},\sigma_{r}^{2}{\bm{D}}_{c}^{2}\right),$

with $\mathbf{K}_{{\bm{Z}}_{c}\mathbf{X}}=\sigma^{2}{\bm{Z}}_{c}\mathbf{X}$
and
$\mathbf{K}_{{\bm{Z}}_{c}{\bm{Z}}_{c}}=\sigma^{2}{\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}$.
By marginalising out the “noiseless inducing outputs” $\hat{\mathbf{u}}_{i}$,
we have the joint distribution $p(\mathbf{f}_{i},\mathbf{u}_{i})$ as

$\displaystyle p(\mathbf{u}^{c}_{i})$
$\displaystyle=\operatorname{\mathcal{N}}(\bm{0},\sigma_{r}^{2}\bm{\Psi}_{c}),\
\bm{\Psi}_{c}={\bm{Z}}_{c}{\bm{Z}}_{c}^{\top}+{\bm{D}}_{c}^{2},$
$\displaystyle p(\mathbf{f}_{i}|\mathbf{X},\mathbf{u}^{c}_{i})$
$\displaystyle=\operatorname{\mathcal{N}}(\sigma_{c}\sigma^{-2}\mathbf{K}_{\mathbf{X}{\bm{Z}}_{c}}\bm{\Psi}_{c}^{-1}\mathbf{u}^{c}_{i},\mathbf{K}_{\mathbf{X}\mathbf{X}}-\sigma^{-2}\mathbf{K}_{\mathbf{X}{\bm{Z}}_{c}}\bm{\Psi}_{c}^{-1}\mathbf{K}_{{\bm{Z}}_{c}\mathbf{X}}).$

Collecting all the random variables in matrix forms, this leads to

$\displaystyle p({\mathbf{U}}_{c})$
$\displaystyle=\operatorname{\mathcal{MN}}(\bm{0},\sigma_{r}^{2}{\bm{I}},\bm{\Psi}_{c}),$
$\displaystyle p(\mathbf{F}|\mathbf{X},{\mathbf{U}}_{c})$
$\displaystyle=\operatorname{\mathcal{MN}}(\sigma_{c}\sigma^{-2}{\mathbf{U}}_{c}\bm{\Psi}_{c}^{-1}\mathbf{K}_{{\bm{Z}}_{c}\mathbf{X}},\sigma_{r}^{2}{\bm{I}},\sigma_{r}^{-2}(\mathbf{K}_{\mathbf{X}\mathbf{X}}-\sigma^{-2}\mathbf{K}_{\mathbf{X}{\bm{Z}}_{c}}\bm{\Psi}_{c}^{-1}\mathbf{K}_{{\bm{Z}}_{c}\mathbf{X}}))$
(41)
$\displaystyle=\operatorname{\mathcal{MN}}({\mathbf{U}}_{c}\bm{\Psi}_{c}^{-1}\sigma_{c}{\bm{Z}}_{c}\mathbf{X},\sigma_{r}^{2}{\bm{I}},\mathbf{X}^{\top}\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}_{c}^{-1}{\bm{Z}}_{c})\mathbf{X}).$

Also recall from conditioning rules of matrix normal distributions, we have
that

$p({\mathbf{W}}|{\mathbf{U}}_{c})=\operatorname{\mathcal{MN}}\left({\mathbf{U}}_{c}\bm{\Psi}_{c}^{-1}\sigma_{c}{\bm{Z}}_{c},\sigma_{r}^{2}{\bm{I}},\sigma_{c}^{2}({\bm{I}}-{\bm{Z}}_{c}^{\top}\bm{\Psi}_{c}^{-1}{\bm{Z}}_{c})\right).$

Since for
${\mathbf{W}}\sim\operatorname{\mathcal{MN}}(\mathbf{M},\bm{\Sigma}_{1},\bm{\Sigma}_{2})$
we have
${\mathbf{W}}\mathbf{X}\overset{d}{\sim}\operatorname{\mathcal{MN}}(\mathbf{M}\mathbf{X},\bm{\Sigma}_{1},\mathbf{X}^{\top}\bm{\Sigma}_{2}\mathbf{X})$,
this immediately shows that $p(\mathbf{F}|\mathbf{X},{\mathbf{U}}_{c})$ is the
push-forward distribution of $p({\mathbf{W}}|{\mathbf{U}}_{c})$ for the
operation $\mathbf{F}={\mathbf{W}}\mathbf{X}$. In other words:

$\mathbf{F}\sim
p(\mathbf{F}|\mathbf{X},{\mathbf{U}}_{c})\quad\Leftrightarrow\quad{\mathbf{W}}\sim
p({\mathbf{W}}|{\mathbf{U}}_{c}),\ \mathbf{F}={\mathbf{W}}\mathbf{X}.$

This confirms the interpretation of ${\mathbf{U}}_{c}$ as "scaled noisy
inducing outputs" that lie in the same space as $\mathbf{F}$. Notice that in
the main text we provide a pictorial visualisation of ${\mathbf{U}}_{c}$ by
selecting $\sigma_{c}=1$. As the inducing weights ${\mathbf{U}}$ are the focus
of our analysis here, we conclude that this specific choice of $\sigma_{c}$ is
without loss of generality.

So far the ${\mathbf{U}}_{c}$ variables assist the posterior inference to
capture correlations across functions values of different inputs. Up to now
the function values remain independent across output dimensions, which is also
reflected by the diagonal row covariance matrices in the above matrix normal
distributions. As in neural networks the output dimension can be fairly large
(e.g. $d_{out}=1000$), to further improve memory efficiency, we proceed to
project the _column vectors_ of ${\mathbf{U}}_{c}$ to an $M_{out}$ dimensional
space with $M_{out}<<d_{out}$. This dimension reduction step is done with a
generative approach, similar to probabilistic PCA (Tipping & Bishop, 1999):

$\begin{split}{\mathbf{U}}&\sim\operatorname{\mathcal{MN}}(0,\bm{\Psi}_{r},\bm{\Psi}_{c}),\\\
{\mathbf{U}}_{c}|{\mathbf{U}}&\sim\operatorname{\mathcal{MN}}(\sigma_{r}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}},\sigma_{r}^{2}({\bm{I}}-{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}),\bm{\Psi}_{c}).\end{split}$
(42)

Note that the column covariance matrices in the above two matrix normal
distributions are the same, and the conditional sampling procedure is done by
a linear transformation of the columns in ${\mathbf{U}}$ plus noise terms.
Again from the marginalisation and conditioning rules of matrix normal
distributions, we have that the full joint distribution Eq. 5, after proper
marginalisation and conditioning, returns

$\displaystyle p({\mathbf{U}})$
$\displaystyle=\operatorname{\mathcal{MN}}(0,\bm{\Psi}_{r},\bm{\Psi}_{c}),$
$\displaystyle p({\mathbf{U}}_{c}|{\mathbf{U}})$
$\displaystyle=\operatorname{\mathcal{MN}}(\sigma_{r}{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\mathbf{U}},\sigma_{r}^{2}({\bm{I}}-{\bm{Z}}_{r}^{\top}\bm{\Psi}^{-1}_{r}{\bm{Z}}_{r}),\bm{\Psi}_{c}).$

This means ${\mathbf{U}}$ can be viewed as “projected noisy inducing points”
for the GPs $p(\mathbf{F})$, whose corresponding “inducing inputs” are row
vectors in ${\bm{Z}}_{c}$. Similarly, column vectors in
${\mathbf{U}}_{r}\mathbf{X}$ can be viewed as the noisy projections of the
column vectors in $\mathbf{F}$, in other words ${\mathbf{U}}_{r}$ can also be
viewed as “neural network weights” connecting the data inputs $\mathbf{X}$ to
the projected output space that ${\mathbf{U}}$ lives in.

As for the variational objective, since $q({\mathbf{W}}|{\mathbf{U}})$ and
$p({\mathbf{W}}|{\mathbf{U}})$ only differ in the scale of the covariance
matrices, it is straightforward to show that the push-forward distribution
$q({\mathbf{W}}|{\mathbf{U}})\rightarrow
q(\mathbf{F}|\mathbf{X},{\mathbf{U}})$ has the same mean as
$p(\mathbf{F}|\mathbf{X},{\mathbf{U}})$, but with a different covariance
matrix that scales $p(\mathbf{F}|\mathbf{X},{\mathbf{U}})$’s covariance matrix
by $\lambda^{2}$. As the operation $\mathbf{F}={\mathbf{W}}\mathbf{X}$ maps
${\mathbf{W}}\in\mathbb{R}^{d_{out}\times d_{in}}$ to
$\mathbf{F}\in\mathbb{R}^{d_{out}\times N}$, this means the conditional KL is
scaled up/down, depending on whether $N\geq d_{in}$ or not:

$\displaystyle\operatorname{\mathbb{KL}}[q(\mathbf{F}|\mathbf{X},{\mathbf{U}})||p(\mathbf{F}|\mathbf{X},{\mathbf{U}})]$
$\displaystyle=\frac{N}{d_{in}}R(\lambda),$ $\displaystyle R(\lambda)$
$\displaystyle:=\operatorname{\mathbb{KL}}[q({\mathbf{W}}|{\mathbf{U}})||p({\mathbf{W}}|{\mathbf{U}})].$

In summary, the push-forward distribution of
$q({\mathbf{W}}_{1:L},{\mathbf{U}}_{1:L})\rightarrow
q(\mathbf{F}_{1:L},{\mathbf{U}}_{1:L})$ is

$q(\mathbf{F}_{1:L},{\mathbf{U}}_{1:L})=\textstyle\prod\nolimits_{l=1}^{L}q(\mathbf{F}_{l}|\mathbf{F}_{l-1},{\mathbf{U}}_{l})q({\mathbf{U}}_{l}),\quad\mathbf{F}_{0}:=\mathbf{X},$

and the corresponding variational lower-bound for
$q(\mathbf{F}_{1:L},{\mathbf{U}}_{1:L})$ becomes (with
$\mathcal{D}=(\mathbf{X},\mathbf{Y})$)

$\mathcal{L}(q(\mathbf{F}_{1:L},{\mathbf{U}}_{1:L}))=\mathbb{E}_{q(\mathbf{F}_{1:L})}[\log
p(\mathbf{Y}|\mathbf{F}_{L})]-\textstyle\sum\nolimits_{l=1}^{L}\left(\frac{N}{d^{l}_{in}}R(\lambda_{l})+\operatorname{\mathbb{KL}}[q({\mathbf{U}}_{l})||p({\mathbf{U}}_{l})]\right),$
(43)

with $d^{l}_{in}$ the input dimension of layer $l$.

Note that

$\mathbb{E}_{q(\mathbf{F}_{1:L})}[\log
p(\mathbf{Y}|\mathbf{F}_{L})]=\mathbb{E}_{q({\mathbf{W}}_{1:L})}[\log
p(\mathbf{Y}|\mathbf{X},{\mathbf{W}}_{1:L})]=\mathbb{E}_{q({\mathbf{W}}_{1:L})}[\log
p(\mathcal{D}|{\mathbf{W}}_{1:L})].$ (44)

Comparing equations eqs. 9 and 43, we see that the only difference between
weight-space and function-space variational objectives comes in the scale of
the conditional KL term. Though not investigated in the experiments, we
conjecture that it could bring potential advantage to optimise the following
variational lower-bound:

$\displaystyle\tilde{\mathcal{L}}(q(\mathbf{F}_{1:L},{\mathbf{U}}_{1:L}))=$
$\displaystyle\mathbb{E}_{q(\mathbf{F}_{1:L})}[\log
p(\mathbf{Y}|\mathbf{F}_{L})]-\textstyle\sum\nolimits_{l=1}^{L}\left(\beta_{l}R(\lambda_{l})+\operatorname{\mathbb{KL}}[q({\mathbf{U}}_{l})||p({\mathbf{U}}_{l})]\right),$
(45) $\displaystyle\beta_{l}=$ $\displaystyle\min(1,\frac{N}{d^{l}_{in}}).$

The intuition is that, as uncertainty is expected to be lower when $N\geq
d_{in}$, it makes sense to use $\beta=1\leq N/d_{in}$ to reduce the
regularisation effect introduced by the KL term. In other words, this allows
the variational posterior to focus more on fitting the data, and in this
“large-data” regime over-fitting is less likely to appear. On the other hand,
function-space inference approaches (such as GPs) often return better
uncertainty estimates when trained on small data ($N<d_{in}$). So choosing
$\beta=N/d_{in}<1$ in this case would switch to function-space inference and
thereby improving uncertainty quality potentially. In the CIFAR experiments,
the usage of convolutional filters leads to the fact that $N\geq d_{in}^{l}$
for all ResNet layers. Therefore in those experiments $\beta_{l}=1$ for all
layers, which effectively falls back to the weight-space objective Eq. 9.

##  Appendix H Whitening and hierarchical inducing variables

The inducing weights ${\mathbf{U}}_{1:L}$ further allow for introducing a
single inducing weight matrix ${\mathbf{U}}$ that is shared across the
network. By doing so, correlations of weights between layers in the
approximate posterior are introduced. The inducing weights are then sampled
jointly conditioned on the global inducing weights. This requires that all
inducing weight matrices are of the same size along at least one axis, such
that they can be concatenated along the other one.

The easiest way of introducing a global inducing weight matrix is by
proceeding similarly to the introduction of the per-layer inducing weights. As
a pre-requisite, we need to work with “whitened” inducing weights, i.e. set
the covariance of the marginal $p({\mathbf{U}}_{l})$ to the identity and pre-
multiply the covariance block between ${\mathbf{W}}_{l}$ and
${\mathbf{U}}_{l}$ with the inverse Cholesky of $\bm{\Psi}_{l}$. In this
whitened model, the full augmented prior per-layer is:

$\displaystyle\begin{pmatrix}{\mathbf{W}}&{\mathbf{U}}_{c}\\\
{\mathbf{U}}_{r}&{\mathbf{U}}\end{pmatrix}$ $\displaystyle\sim
p({\mathbf{W}},{\mathbf{U}}_{c},{\mathbf{U}}_{r},{\mathbf{U}}):=\operatorname{\mathcal{MN}}(0,\tilde{\bm{\Sigma}}_{r},\tilde{\bm{\Sigma}}_{r}),$
(46) $\displaystyle\text{with}\quad\tilde{{\bm{L}}}_{r}$
$\displaystyle=\begin{pmatrix}\sigma_{r}{\bm{I}}&0\\\
{\bm{L}}_{r}^{-1}{\bm{Z}}_{r}&{\bm{L}}_{r}^{-1}{\bm{D}}_{r}\end{pmatrix}$ (47)
$\displaystyle\text{s.t.}\quad\tilde{\bm{\Sigma}}_{r}$
$\displaystyle=\tilde{{\bm{L}}}_{r}\tilde{{\bm{L}}}_{r}^{\top}=\begin{pmatrix}\sigma_{r}^{2}{\bm{I}}&\sigma_{r}{\bm{Z}}_{r}^{\top}{\bm{L}}_{r}^{-\top}\\\
\sigma_{r}{\bm{L}}_{r}^{-1}{\bm{Z}}_{r}&{\bm{I}}\end{pmatrix}$
$\displaystyle\text{and}\quad\tilde{{\bm{L}}}_{c}$
$\displaystyle=\begin{pmatrix}\sigma_{c}{\bm{I}}&0\\\
{\bm{L}}_{c}^{-1}{\bm{Z}}_{c}&{\bm{L}}_{c}^{-1}{\bm{D}}_{c}\end{pmatrix}$ (48)
$\displaystyle\text{s.t.}\quad\tilde{\bm{\Sigma}}_{c}$
$\displaystyle=\tilde{{\bm{L}}}_{c}\tilde{{\bm{L}}}_{c}^{\top}=\begin{pmatrix}\sigma_{c}^{2}{\bm{I}}&\sigma_{c}{\bm{Z}}_{c}^{\top}{\bm{L}}_{c}^{-\top}\\\
\sigma_{c}{\bm{L}}_{c}^{-1}{\bm{Z}}_{c}&{\bm{I}}\end{pmatrix}.$

One can verify that this whitened model leads to the same conditional
distribution $p({\mathbf{W}}|{\mathbf{U}})$ as presented in the main text.
After whitening, for each ${\mathbf{U}}_{l}$ we have that
$p(\operatorname{vec}({\mathbf{U}}_{l}))=\operatorname{\mathcal{N}}(0,{\bm{I}})$,
therefore we can also write their joint distribution as
$p(\operatorname{vec}({\mathbf{U}}_{1:L}))=\operatorname{\mathcal{N}}(0,{\bm{I}})$.
In order to construct a matrix normal prior
$p({\mathbf{U}}_{1:L})=\operatorname{\mathcal{MN}}(0,{\bm{I}},{\bm{I}})$, the
inducing weight matrices ${\mathbf{U}}_{1:L}$ needs to be stacked either along
the rows or columns, requiring the other dimension to be matching across all
layers. Then, As the covariance is the identity with
$\sigma=\sigma_{r}=\sigma_{c}=1$, we can augment $p({\mathbf{U}}_{1:L})$ in
the exact same way as we previously augmented the prior $p({\mathbf{W}}_{l})$
with ${\mathbf{U}}_{l}$.

##  Appendix I Open-source code

We open-source our approach as a PyTorch wrapper bayesianize:
https://github.com/microsoft/bayesianize

bayesianize is a lightweight Bayesian neural network (BNN) wrapper, and the
goal is to allow for easy conversion of neural networks in existing scripts to
BNNs with minimal changes to the code. Currently the wrapper supports the
following uncertainty estimation methods for feed-forward neural networks and
convolutional neural networks:

  * •

Mean-field variational inference (MFVI) with fully factorised Gaussian (FFG)
approximation, i.e. FFG-${\mathbf{W}}$ in abbreviation.

  * •

Variational inference with full-covariance Gaussian approximation
(FCG-${\mathbf{W}}$).

  * •

Inducing weight approaches, including FFG-${\mathbf{U}}$, FCG-${\mathbf{U}}$
and Ensemble-${\mathbf{U}}$.

The main workhorse of our library is the bayesianize_ function, which turns
deterministic nn.Linear and nn.Conv layers into their Bayesian counterparts.
For example, to construct a Bayesian ResNet-18 that uses the variational
inducing weight method, run:

⬇

import bnn

net = torchvision.models.resnet18()

bnn.bayesianize_(net, inference="inducing", inducing_rows=64,
inducing_cols=64)

In the above code, inducing_rows corresponds to $M_{out}$ and inducing_cols
corresponds to $M_{in}$. In other words, they specify the dimensions of
${\mathbf{U}}\in\mathbb{R}^{M_{out}\times M_{in}}$. Then the converted BNN can
be trained in almost identical way as one would train a deterministic net:

⬇

yhat = net(x_train)

nll = F.cross_entropy(yhat, y_train)

kl = sum(m.kl_divergence() for m in net.modules()

if hasattr(m, "kl_divergence"))

loss = nll + kl / dataset_size

loss.backward()

optim.step()

Note that while the call to the forward method of the net looks the same, it
is no longer deterministic because the weights are sampled, so subsequent
calls will lead to different predictions. Therefore, when testing, an average
of multiple predictions is needed. For example, in BNN classification:

⬇

net.eval()

with torch.no_grad():

logits = torch.stack([net(x_test) for _ in range(num_samples)])

probs = logits.softmax(-1).mean(0)

In the above code, probs computes

$p(\bm{y}|\bm{x},\mathcal{D})\approx\frac{1}{K}\sum_{k=1}^{K}p(\bm{y}|\bm{x},{\mathbf{W}}_{1:L}^{(k)}),\quad{\mathbf{W}}_{1:L}^{(k)}\sim
q,$

where $K$ is the number of Monte Carlo samples num_samples.

bayesianize also supports using different methods or arguments for different
layers, by passing in a dictionary for the inference argument. This way we
can, for example, take a _pre-trained_ ResNet and only perform (approximate)
Bayesian inference over the weights of the final, linear layer:

⬇

net = … # load pre-trained network with net.fc as the last layer

bnn.bayesianize_(net, inference={

net.fc: {"inference": "fcg"}

})

optim = torch.optim.Adam(net.fc.parameters(), 1e-3)

# then train the last Bayesian layer accordingly

…

For more possible ways of configuring the BNN settings, see example config
json files in the open-source repository.

##  Appendix J Experimental details

###  J.1 Regression experiments

Following (Foong et al., 2019), we sample $50$ inputs each from
$\mathcal{U}\left[-1,-0.7\right]$ and $\mathcal{U}\left[0.5,1\right]$ as
inputs and targets $y\sim\operatorname{\mathcal{N}}(\cos(0.4x+0.8),0.01)$. As
a prior we use a zero-mean Gaussian with standard deviation
$\frac{4}{\sqrt{d_{in}}}$ for the weights and biases of each layer. Our
network architecture has a single hidden layer of $50$ units and uses a
$\tanh$-nonlinearity. All three variational methods are optimised using Adam
(Kingma & Ba, 2015) for $20,000$ updates with an initial learning rate of
$10^{-3}$. We average over $32$ MC samples from the approximate posterior for
every update. For Ensemble-U and FCG-U we decay the learning rate by a factor
of $0.1$ after $10,000$ updates and the size of the inducing weight matrix is
$2\times 25$ for the input layer (accounting for the bias) and $25\times 1$
for the output layer. Ensemble-U uses an ensemble size of $8$.

For NUTS we use the implementation provided in Pyro (Bingham et al., 2019). We
draw a total of $25,000$ samples, discarding the first $5000$ as burn-in and
using $1000$ randomly selected ones for prediction.

###  J.2 Classification experiments

We base our implementation on the Resnet-18 class in torchvision (Paszke et
al., 2019), replacing the input convolutional layer with a $3\times 3$ kernel
size and removing the max-pooling layer. We train the deterministic network on
CIFAR-10 using Adam with a learning rate of $3\times 10^{-4}$ for 200 epochs.
On CIFAR-100 we found SGD with a momentum of $0.9$ and initial learning rate
of $0.1$ decayed by a factor of $0.1$ after $60$, $120$ and $160$ epochs to
lead to better accuracies. The ensemble is formed of the five deterministic
networks trained with different random seeds.

For FFG-${\mathbf{W}}$ we initialised the mean parameters using the default
initialisation in pytorch for the corresponding deterministic layers. The
initial standard deviations are set to $10^{-4}$. We train using Adam for
$200$ epochs on CIFAR-10 with a learning rate of $3\times 10^{-4}$, and $300$
epochs on CIFAR-100 with an initial learning rate of $10^{-3}$, decaying by a
factor of $0.1$ after $200$ epochs. On both datasets we only use the negative
log likelihood (NLL)part of the variational lower bound for the first $100$
epochs as initialisation to the maximum likelihood parameter and then anneal
the weight of the kl term linearly over the following $50$ epochs. For the
prior we use a standard Gaussian on all weights and biases and restrict the
standard deviation of the posterior to be at most $\sigma_{max}=0.1$. We also
experimented with a larger upper limit of $\sigma_{max}=0.3$, but found this
to negatively affect both accuracy and calibration.

All the ${\mathbf{U}}$-space approaches use Gaussian priors
$p(\operatorname{vec}({\mathbf{W}}_{l}))=\operatorname{\mathcal{N}}(0,1/\sqrt{d_{in}})$,
motivated by the connection to GPs. Hyperparameter and optimisation details
for the inducing weight methods on CIFAR-10 are discussed below in the details
on the ablation study. We train all methods using Adam for $300$ epochs with a
learning rate of $10^{-3}$ for the first $200$ epochs and then decay by a
factor of $10$. For the initial $100$ epochs we train without the KL-term of
the ELBO and then anneal its weight linearly over the following $50$ epochs.
For the tables and figures in the main text, we set $\lambda_{max}=0.1$ for
Ensemble-${\mathbf{U}}$ on both datasets, and
$\sigma_{max}=0.1,\lambda_{max}=0.03$ on CIFAR-10 for FFG-${\mathbf{U}}$. We
initialise the entries of the $Z$ matrices by sampling from a zero-mean
Gaussian with variance $\frac{1}{M}$ and set the diagonal entries of the $D$
matrices to $10^{-3}$. For FFG-${\mathbf{U}}$ we initialise the mean of the
variational Gaussian posterior in ${\mathbf{U}}$-space by sampling from a
standard Gaussian and set the initial variances to $10^{-3}$. For
Ensemble-${\mathbf{U}}$ initialisation, we draw an $M\times M$ shaped sample
from a standard Gaussian that is shared across ensemble members and add
independent Gaussian noise with a standard deviation of $0.1$ for each member.
We use an ensemble size of $5$. During optimisation, we draw $1$ MC samples
per update step for both FFG-${\mathbf{U}}$ and Ensemble-${\mathbf{U}}$ (such
that each ensemble member is used once). For testing we use $20$ MC samples
for all variational methods. We fit BatchNorm parameters by minimising the
NLL.

#### The study of hyper-parameter selection on CIFAR-10

We run the inducing weight method with the following options:

  * •

Row/column dimensions of ${\mathbf{U}}_{l}$ ($M$): $M\in\\{16,32,64,128\\}$.
We set $M=M_{in}=M_{out}$ except for the last layer, where we use $M_{in}=M$
and $M_{out}=10$.

  * •

$\lambda_{max}$ values for FFG-${\mathbf{U}}$ and Ensemble-${\mathbf{U}}$:
$\lambda_{max}\in\\{0,0.03,0.1,0.3\\}$.
When $\lambda_{max}=0$ it means $q({\mathbf{W}}|{\mathbf{U}})$ is a delta
measure centered at the mean of $p({\mathbf{W}}|{\mathbf{U}})$.

  * •

$\sigma_{\max}$ values for FFG-${\mathbf{U}}$:
$\sigma_{\max}=\\{0,0.1,0.3\\}$.
When $\sigma_{max}=0$ we use a MAP estimate for ${\mathbf{U}}$.

Each experiment is repeated with 5 random seeds to collect the averaged
results on a single NVIDIA RTX 2080TI. The models are trained with 100 epochs
in total. We first run 50 epochs of maximum likelihood to initialise the
model, then run 40 epochs training on the modified variational lower-bound
with KL annealing (linear scaling schedule), finally we run 10 epochs of
training with the variational lower-bound (i.e. no KL annealing). We use Adam
optimiser with learning rate $3e-4$ and the default $\beta_{1},\beta_{2}$
parameters in PyTorch’s implementation.

##  Appendix K Additional Results

Below in Appendices K and K we provide extended versions of Table 2. This
table contains standard errors across the random seeds for the corresponding
metrics and we additionally report NLLs and Brier scores. The error bar
results are not available for Ensemble-${\mathbf{W}}$, as it is constructed
from the 5 independently trained deterministic neural network with maximum
likelihood.

Table 5: Complete in-distribution results for Resnet-50 on CIFAR10

Method Acc. $\uparrow$ NLL $\downarrow$ ECE $\downarrow$ Brier $\downarrow$
Deterministic $94.72{\scriptstyle{\pm}0.08}$ $0.43{\scriptstyle{\pm}0.01}$
$4.46{\scriptstyle{\pm}0.08}$ $0.10{\scriptstyle{\pm}0.00}$
Ensemble-W $95.90{\scriptstyle\phantom{{\pm}0.00}}$
$0.20{\scriptstyle\phantom{{\pm}0.00}}$
$1.08{\scriptstyle\phantom{{\pm}0.00}}$
$0.06{\scriptstyle\phantom{{\pm}0.00}}$
FFG-W $94.13{\scriptstyle{\pm}0.08}$ $0.18{\scriptstyle{\pm}0.00}$
$0.50{\scriptstyle{\pm}0.06}$ $0.09{\scriptstyle{\pm}0.00}$
FFG-U (M=64) $94.40{\scriptstyle{\pm}0.05}$ $0.17{\scriptstyle{\pm}0.00}$
$0.64{\scriptstyle{\pm}0.06}$ $0.08{\scriptstyle{\pm}0.00}$
FFG-U (M=128) $94.66{\scriptstyle{\pm}0.09}$ $0.17{\scriptstyle{\pm}0.00}$
$1.59{\scriptstyle{\pm}0.06}$ $0.08{\scriptstyle{\pm}0.00}$
Ensemble-U (M=64) $94.94{\scriptstyle{\pm}0.07}$ $0.16{\scriptstyle{\pm}0.00}$
$0.45{\scriptstyle{\pm}0.06}$ $0.08{\scriptstyle{\pm}0.00}$
Ensemble-U (M=128) $95.34{\scriptstyle{\pm}0.05}$
$0.17{\scriptstyle{\pm}0.00}$ $1.29{\scriptstyle{\pm}0.05}$
$0.07{\scriptstyle{\pm}0.00}$

Table 6: Complete in-distribution results for Resnet-50 on CIFAR100

Method Acc. $\uparrow$ NLL $\downarrow$ ECE $\downarrow$ Brier $\downarrow$
Deterministic $75.73{\scriptstyle{\pm}0.16}$ $2.14{\scriptstyle{\pm}0.01}$
$19.69{\scriptstyle{\pm}0.15}$ $0.43{\scriptstyle{\pm}0.00}$
Ensemble-W $79.33{\scriptstyle\phantom{{\pm}0.00}}$
$1.23{\scriptstyle\phantom{{\pm}0.00}}$
$6.51{\scriptstyle\phantom{{\pm}0.00}}$
$0.31{\scriptstyle\phantom{{\pm}0.00}}$
FFG-W $74.44{\scriptstyle{\pm}0.27}$ $1.01{\scriptstyle{\pm}0.01}$
$4.24{\scriptstyle{\pm}0.10}$ $0.35{\scriptstyle{\pm}0.00}$
FFG-U (M=64) $75.37{\scriptstyle{\pm}0.09}$ $0.92{\scriptstyle{\pm}0.01}$
$2.29{\scriptstyle{\pm}0.39}$ $0.34{\scriptstyle{\pm}0.00}$
FFG-U (M=128) $75.88{\scriptstyle{\pm}0.13}$ $0.91{\scriptstyle{\pm}0.00}$
$6.66{\scriptstyle{\pm}0.15}$ $0.34{\scriptstyle{\pm}0.00}$
Ensemble-U (M=64) $75.97{\scriptstyle{\pm}0.12}$ $0.90{\scriptstyle{\pm}0.00}$
$1.12{\scriptstyle{\pm}0.06}$ $0.33{\scriptstyle{\pm}0.00}$
Ensemble-U (M=128) $77.61{\scriptstyle{\pm}0.11}$
$0.94{\scriptstyle{\pm}0.00}$ $6.00{\scriptstyle{\pm}0.12}$
$0.32{\scriptstyle{\pm}0.00}$

The results of pruning different fractions of the weights can be found in
Appendix K for the in-distribution uncertainty evaluation for Resnet-50 and
the OOD detection in Appendix K. For the pruning experiments, we take the
parameters from the corresponding full runs, set a fixed percentage of the
weights to be deterministically $0$ and fine-tune the remaining weights with a
new optimizer for $50$ epochs. We use Adam with a learning rate of $10^{-4}$.
For FFG-${\mathbf{W}}$ we select the weights with the smallest ratio of
absolute mean to standard deviation in the approximate posterior, and for
FFG-${\mathbf{U}}$ the $Z$ parameters with the smallest absolute value.

For FFG-${\mathbf{W}}$ we find that pruning up to $90\%$ of the weights only
worsens ECE and NLL on the more difficult CIFAR100 datasets. Pruning $99\%$ of
the weights worsens accuracy and OOD detection, but interestingly improves ECE
on CIFAR100, where accuracy is noticeably worse.

Pruning $25\%$ and $50\%$ of the $Z$ parameters in FFG-${\mathbf{U}}$ results
in a total parameter count of $4,408,790$ and $3,106,678$, i.e. $18.7$ and
$13.2\%$ of the deterministic parameters respectively on ResNet-50. Up to
pruning $50\%$ of the $Z$ parameters, we find that only ECE becomes slightly
worse, although on CIFAR100 it is still better than the ECE for
FFG-${\mathbf{W}}$ at $100\%$ of the weights. Other metrics are not affected
neither on the in-distribution uncertainty or OOD detection, except for a
minor drop in accuracy.

Table 7: Pruning in-distribution uncertainty results for Resnet-50. The
percentage refers to the weights left for FFG-${\mathbf{W}}$ and the number of
$Z$ parameters for FFG-${\mathbf{U}}$.

CIFAR10 CIFAR100
Method Acc. $\uparrow$ NLL $\downarrow$ ECE $\downarrow$ Brier $\downarrow$
Acc. $\uparrow$ NLL $\downarrow$ ECE $\downarrow$ Brier $\downarrow$
FFG-W (100%) $94.13{\scriptstyle{\pm}0.08}$ $0.18{\scriptstyle{\pm}0.00}$
$0.50{\scriptstyle{\pm}0.06}$ $0.09{\scriptstyle{\pm}0.00}$
$74.44{\scriptstyle{\pm}0.27}$ $1.01{\scriptstyle{\pm}0.01}$
$4.24{\scriptstyle{\pm}0.10}$ $0.35{\scriptstyle{\pm}0.00}$
FFG-W (50%) $94.07{\scriptstyle{\pm}0.03}$ $0.18{\scriptstyle{\pm}0.00}$
$0.40{\scriptstyle{\pm}0.04}$ $0.09{\scriptstyle{\pm}0.00}$
$74.38{\scriptstyle{\pm}0.21}$ $1.02{\scriptstyle{\pm}0.01}$
$4.15{\scriptstyle{\pm}0.09}$ $0.36{\scriptstyle{\pm}0.00}$
FFG-W (10%) $94.17{\scriptstyle{\pm}0.06}$ $0.18{\scriptstyle{\pm}0.00}$
$0.58{\scriptstyle{\pm}0.02}$ $0.09{\scriptstyle{\pm}0.00}$
$74.42{\scriptstyle{\pm}0.23}$ $1.10{\scriptstyle{\pm}0.01}$
$6.86{\scriptstyle{\pm}0.08}$ $0.36{\scriptstyle{\pm}0.00}$
FFG-W (1%) $93.60{\scriptstyle{\pm}0.09}$ $0.19{\scriptstyle{\pm}0.00}$
$0.80{\scriptstyle{\pm}0.06}$ $0.09{\scriptstyle{\pm}0.00}$
$67.08{\scriptstyle{\pm}0.33}$ $1.19{\scriptstyle{\pm}0.01}$
$1.36{\scriptstyle{\pm}0.18}$ $0.44{\scriptstyle{\pm}0.00}$
FFG-W (0.1%) $58.59{\scriptstyle{\pm}0.75}$ $1.15{\scriptstyle{\pm}0.02}$
$7.33{\scriptstyle{\pm}0.23}$ $0.55{\scriptstyle{\pm}0.01}$
$10.66{\scriptstyle{\pm}0.43}$ $3.97{\scriptstyle{\pm}0.03}$
$2.80{\scriptstyle{\pm}0.15}$ $0.96{\scriptstyle{\pm}0.00}$
FFG-U (100%) $94.40{\scriptstyle{\pm}0.05}$ $0.17{\scriptstyle{\pm}0.00}$
$0.64{\scriptstyle{\pm}0.06}$ $0.08{\scriptstyle{\pm}0.00}$
$75.37{\scriptstyle{\pm}0.09}$ $0.92{\scriptstyle{\pm}0.01}$
$2.29{\scriptstyle{\pm}0.39}$ $0.34{\scriptstyle{\pm}0.00}$
FFG-U (75%) $94.45{\scriptstyle{\pm}0.05}$ $0.18{\scriptstyle{\pm}0.00}$
$2.19{\scriptstyle{\pm}0.11}$ $0.09{\scriptstyle{\pm}0.00}$
$75.26{\scriptstyle{\pm}0.07}$ $0.93{\scriptstyle{\pm}0.00}$
$3.74{\scriptstyle{\pm}0.67}$ $0.35{\scriptstyle{\pm}0.00}$
FFG-U (50%) $94.31{\scriptstyle{\pm}0.07}$ $0.18{\scriptstyle{\pm}0.00}$
$2.31{\scriptstyle{\pm}0.09}$ $0.09{\scriptstyle{\pm}0.00}$
$74.89{\scriptstyle{\pm}0.07}$ $0.94{\scriptstyle{\pm}0.00}$
$5.04{\scriptstyle{\pm}0.77}$ $0.35{\scriptstyle{\pm}0.00}$
FFG-U (25%) $93.34{\scriptstyle{\pm}0.03}$ $0.22{\scriptstyle{\pm}0.00}$
$4.83{\scriptstyle{\pm}0.09}$ $0.11{\scriptstyle{\pm}0.00}$
$71.32{\scriptstyle{\pm}0.16}$ $1.09{\scriptstyle{\pm}0.01}$
$12.44{\scriptstyle{\pm}0.90}$ $0.42{\scriptstyle{\pm}0.00}$

Table 8: Pruning OOD detection metrics for Resnet-$50$ trained on
CIFAR$10$/$100$.

In-dist $\rightarrow$ OOD C10 $\rightarrow$ C100 C10 $\rightarrow$ SVHN C100
$\rightarrow$ C10 C100 $\rightarrow$ SVHN
Method / Metric AUROC AUPR AUROC AUPR AUROC AUPR AUROC AUPR