# Regularized KL-Divergence for Well-Defined Function-Space Variational
Inference in Bayesian neural networks

Tristan Cinquin
University of Tübingen
<EMAIL_ADDRESS>
&Robert Bamler
University of Tübingen
<EMAIL_ADDRESS>

###### Abstract

Bayesian neural networks (BNN) promise to combine the predictive performance
of neural networks with principled uncertainty modeling important for safety-
critical systems and decision making. However, posterior uncertainty estimates
depend on the choice of prior, and finding informative priors in weight-space
has proven difficult. This has motivated variational inference (VI) methods
that pose priors directly on the function generated by the BNN rather than on
weights. In this paper, we address a fundamental issue with such function-
space VI approaches pointed out by Burt et al., (2020), who showed that the
objective function (ELBO) is negative infinite for most priors of interest.
Our solution builds on _generalized_ VI (Knoblauch et al.,, 2019) with the
regularized KL divergence (Quang,, 2019) and is, to the best of our knowledge,
the first well-defined variational objective for function-space inference in
BNNs with Gaussian process (GP) priors. Experiments show that our method
incorporates the properties specified by the GP prior on synthetic and small
real-world data sets, and provides competitive uncertainty estimates for
regression, classification and out-of-distribution detection compared to BNN
baselines with both function and weight-space priors.

##  1 Introduction

Neural networks have shown impressive results in many fields but fail to
provide well-calibrated uncertainty estimates, which are essential in
applications associated with risk, such as healthcare (Kompa et al.,, 2021;
Abdullah et al.,, 2022) or finance (Bew et al.,, 2019; Wong,, 2023). Bayesian
neural networks (BNNs) offer to combine the scalability and predictive
performance of neural networks with principled uncertainty modeling by
explicitly capturing epistemic uncertainty, i.e., uncertainty resulting from
finite training data. While the choice of prior in the Bayesian framework
strongly affects the uncertainty estimates later obtained from the posterior,
specifying informative priors on BNN weights has proven difficult and is
hypothesized to have limited their practical applicability (Knoblauch et al.,,
2019; Cinquin et al.,, 2021; Tran et al.,, 2022). For instance, the default
isotropic Gaussian prior, which is often chosen for tractability rather than
for the beliefs it carries (Knoblauch et al.,, 2019), is known to have
pathological behavior in some cases (Tran et al.,, 2022). A promising approach
to solve this issue is to place priors directly on the function represented by
the BNN instead of the weights. While being technically more challenging,
function-space priors allow incorporating interpretable knowledge into the
inference, for instance enabling the use of the extensive Gaussian Process
(GP) literature to improve prior selection and design (Williams and
Rasmussen,, 2006).

A recent line of work has focused on using function-space priors in BNNs with
variational inference (VI) (Sun et al.,, 2019). VI is an appealing method
because of its successful application to BNNs, its flexibility in terms of
approximate posterior parameterization, and its scalability to large datasets
and models (Hoffman et al.,, 2013; Blundell et al.,, 2015; Tomczak et al.,,
2020). Unfortunately, in the context of BNNs with function-space priors, the
Kullbach-Leibler (KL) divergence term in the VI objective (ELBO) involves two
intractabilities: (i) a supremum over infinitely many subsets and (ii) access
to the density of the distribution of function represented by the BNN, which
has no closed-form expression. Sun et al., (2019) propose to address problem
(i) by approximating the supremum in the KL divergence by an expectation, and
problem (ii) by using implicit score function estimators (which make this
method difficult to use in practice (Ma and Hernández-Lobato,, 2021)).
However, the problem is actually more severe. Not only is the KL divergence
intractable, it is infinite in many cases of interest (Burt et al.,, 2020),
such as when the prior is a non-degenerate GP or a BNN with a different
architecture. Thus, in these (and many more) situations, the KL divergence
cannot even be approximated. As a consequence, more recent work abandons using
BNNs and instead uses deterministic neural networks to parameterize basis
functions (Ma and Hernández-Lobato,, 2021) or a GP mean (Wild et al.,, 2022).
The only prior work (Rudner et al., 2022b, ) on function-space VI in BNNs that
overcomes the issue pointed out by Burt et al., (2020) does so by deliberately
limiting itself to cases where the KL divergence is known to be finite (by
defining the prior as the pushforward of a weight-space distribution).
Therefore, the method by Rudner et al., 2022b  suffers from the same issues
regarding prior specification as any other weight-space inference method for
BNNs.

In this paper, we address the argument by Burt et al., (2020) that VI does not
provide a valid objective for inference in BNNs with genuine function-space
priors, and we propose to apply the framework of generalized variational
inference (Knoblauch et al.,, 2019). We present a simple method for function-
space inference with GP priors that builds on the so-called _regularized_ KL
divergence (Quang,, 2019), which generalizes the conventional KL divergence
and is finite for any pair of Gaussian measures. We obtain a Gaussian measure
for the variational posterior by considering the linearized BNN from Rudner et
al., 2022b , and we are free to choose a prior from a large set of GPs which
have an associated Gaussian measure on the considered function space. While
the regularized KL divergence is still intractable, it can be consistently
estimated from samples with a known error bound. We find that our method
effectively incorporates the beliefs specified by GP priors into the inference
process (see figure 1) and provides competitive performance compared to BNN
baselines. To the best of our knowledge, our method is the first to provide a
well-defined objective for function-space inference in BNNs with informative
GP priors. Our contributions are summarized below:

  1. 1.

We propose to use generalized VI with the so-called _regularized_ KL
divergence to mitigate the issue of an infinite KL divergence when using VI in
BNNs with function-space priors.

  2. 2.

We present a new, well-defined objective for function-space inference in the
linearized BNN with GP priors, which results in a simple algorithm.

  3. 3.

We analyze our method empirically on synthetic and real-world data sets, and
find that it accurately captures structural properties specified by the GP
prior and provides competitive uncertainty estimates for regression,
classification, and out-of-distribution detection compared to BNN baselines
with both function-space and weight-space priors.

The paper is structured as follows. section 2 introduces function-space VI in
BNNs and the regularized KL divergence. section 3 presents our proposed method
for generalized function-space VI (GFSVI) in BNNs. section 4 presents
experimental results. We discuss related work in section 5 and conclude in
section 6.

Figure 1:  BNN inference on synthetic data (gray circles) with Gaussian
process priors encoding different properties such as smoothness (increasing
from Matern-1/2 to RBF) and periodicity (periodic).

##  2 Background

In this section, we provide background on function-space variational inference
(VI) in BNNs and discuss the fundamental issue of an infinite KL-divergence.
We then introduce the regularized KL divergence, which is the basis for our
solution presented in Section 3.

###  2.1 Function-space VI in BNNs

We consider a neural network $f({\,\cdot\,;{\bm{w}}})$ with weights
${\bm{w}}\in^{p}$, and a data set
$\mathcal{D}=\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{N}$ with features
${\bm{x}}_{i}\in\mathcal{X}\subset^{d}$ and values $y_{i}\in\mathcal{Y}$.
Bayesian Neural Networks (BNNs) are specified further by a likelihood function
$p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}=\prod_{i=1}^{N}p\mathopen{}\left(y_{i}\,|\,f({\bm{x}}_{i};{\bm{w}})\right)\mathclose{}$
and—traditionally—a prior $p\mathopen{}\left({\bm{w}}\right)\mathclose{}$ on
the weights, and one seeks the posterior distribution
$p\mathopen{}\left({\bm{w}}\,|\,\mathcal{D}\right)\mathclose{}\propto
p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}\,p\mathopen{}\left({\bm{w}}\right)\mathclose{}$.
The method proposed in this paper builds on variational inference (VI), which
approximates $p\mathopen{}\left({\bm{w}}\,|\,\mathcal{D}\right)\mathclose{}$
with a variational distribution
$q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}$, whose variational
parameters $\phi$ maximize the evidence lower bound (ELBO),

$\displaystyle\mathcal{L}(\phi):={\mathop{{}\mathbb{E}}}_{q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}}\mathopen{}\left[\log
p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}\right]\mathclose{}-D_{\mathrm{KL}}\mathopen{}\left(q_{\phi}\,\|\,p\right)\mathclose{}$
(2.1)

where $D_{\text{KL}}$ is the Kullback-Leibler (KL) divergence

$D_{\mathrm{KL}}\mathopen{}\left(q_{\phi}\,\|\,p\right)\mathclose{}:={\mathop{{}\mathbb{E}}}_{q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}}\mathopen{}\left[\log\mathopen{}\left(q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}\big{/}p\mathopen{}\left({\bm{w}}\right)\mathclose{}\right)\mathclose{}\right]\mathclose{}.$
(2.2)

At test time, we approximate the predictive distribution for given features
${\bm{x}}^{*}$ as
$p\mathopen{}\left(y^{*}\,|\,{\bm{x}}^{*}\right)\mathclose{}={\mathop{{}\mathbb{E}}}_{p\mathopen{}\left({\bm{w}}\,|\,\mathcal{D}\right)\mathclose{}}\mathopen{}\left[\big{.}p\mathopen{}\left(y^{*}\,|\,f({\bm{x}}^{*};{\bm{w}})\right)\mathclose{}\right]\mathclose{}\approx{\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\big{.}p\mathopen{}\left(y^{*}\,|\,f({\bm{x}}^{*};{\bm{w}})\right)\mathclose{}\right]\mathclose{}$.

##### Function-space variational inference

Since weights of neural networks are not interpretable, we abandon the weight-
space prior $p({\bm{w}})$ and instead pose a prior $\mathbb{P}$ directly on
the function $f\mathopen{}\left(\,\cdot\,;{\bm{w}}\right)\mathclose{}$, which
we denote simply as $f$ when there is no ambiguity. Here, the symbol
$\mathbb{P}$ denotes a probability measure that does not admit a density since
the function space is infinite-dimensional. We compute the expected log-
likelihood as in the first term of Eq. 2.1, but we calculate the KL-term (Eq.
2.2) using the pushforward of $q_{\phi}({\bm{w}})$ along the mapping
${\bm{w}}\mapsto f\mathopen{}\left(\,\cdot\,;{\bm{w}}\right)\mathclose{}$,
which defines the variational measure $\mathbb{Q}_{\phi}$, resulting in the
ELBO in function space,

$\mathcal{L}(\phi):={\mathop{{}\mathbb{E}}}_{q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}}\mathopen{}\left[\log
p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}\right]\mathclose{}-D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}$
(2.3)

with $D_{\text{KL}}$ the KL divergence between measures given by

$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}=\int\\!\log\mathopen{}\left(\\!\frac{d\mathbb{Q}_{\phi}}{d\mathbb{P}}(f)\\!\right)\mathclose{}d\mathbb{Q}_{\phi}.$
(2.4)

Here, the Raydon-Nikodym derivative $d\mathbb{Q}_{\phi}/d\mathbb{P}$
generalizes the density ratio from Eq. 2.2. Like Eq. 2.1, the ELBO in Eq. 2.3
is a lower bound on the evidence (Burt et al.,, 2020). In fact, if
$\mathbb{P}$ is the push-forward of $p({\bm{w}})$ then Eq. 2.3 is a tighter
bound than Eq. 2.1 by the data processing inequality,
$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}\leq
D_{\mathrm{KL}}\mathopen{}\left(q_{\phi}\,\|\,p\right)\mathclose{}$. However,
we motivated function-space VI to avoid weight-space priors, and in this case
the bound in Eq. 2.3 can be looser. We will indeed see below that the bound
becomes infinitely loose in practice, and we thus propose a different
objective in Section 3.

Two intractabilities prevent directly maximizing the ELBO in function space
(Eq 2.3). First, it is not obvious how to practically evaluate or estimate the
KL divergence between two measures in function space in Eq 2.4. Sun et al.,
(2019) showed that it can be expressed as a supremum of KL divergences between
finite-dimensional distributions,

$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}=\sup_{{\bm{x}}\in\mathcal{X}^{M}}D_{\mathrm{KL}}\mathopen{}\left(q_{\phi}(f({\bm{x}}))\,\|\,p(f({\bm{x}}))\right)\mathclose{}.$
(2.5)

Here, ${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}\in\mathcal{X}^{M}$ is a set of
$M$ points in feature space $\mathcal{X}$, and $q_{\phi}(f({\bm{x}}))$ and
$p(f({\bm{x}}))$ are densities of the marginals of $\mathbb{Q}_{\phi}$ and
$\mathbb{P}$ on $\\{f({\bm{x}}^{(i)})\\}_{i=1}^{M}$ respectively. To obtain a
tractable approximation of the supremum over infinitely many sets, Sun et al.,
(2019) replace it by an expectation and Rudner et al., 2022b  estimate it from
samples.

Second, we cannot express the pushforward measure $\mathbb{Q}_{\phi}$ of
$q_{\phi}({\bm{w}})$ in a closed form because the neural network is nonlinear.
Previous work has proposed to mitigate this issue using implicit score
function estimators (Sun et al.,, 2019) or a linearized BNN $f_{L}$ to obtain
a closed-form Gaussian variational measure (Rudner et al., 2022a, ; Rudner et
al., 2022b, ). Our proposal in Section 3 follows the linearized BNN approach
as it only minimally modifies the BNN, preserving most of its inductive bias
(Maddox et al.,, 2021) while considerably simplifying the problem by turning
the pushforward of $q_{\phi}({\bm{w}})$ into a Gaussian process (GP). More
specifically, we consider a Gaussian variational distribution
$q_{\phi}\mathopen{}\left({\bm{w}}\right)\mathclose{}=\mathcal{N}\mathopen{}\left({\bm{m}},{\bm{S}}\right)\mathclose{}$
with parameters $\phi=\left\\{{\bm{m}},{\bm{S}}\right\\}$, and we obtain
$f_{L}$ by linearizing $f$ in the weights around ${\bm{w}}={\bm{m}}$,

$f_{L}({\bm{x}};{\bm{w}}):=f({\bm{x}};{\bm{m}})+J({\bm{x}};{\bm{m}})({\bm{w}}-{\bm{m}})\qquad\text{where}\qquad
J({\bm{x}};{\bm{m}})=\nabla_{\\!{\bm{w}}}f({{\bm{x}};{\bm{w}}})|_{{\bm{w}}={\bm{m}}}.$
(2.6)

Thus, for ${\bm{w}}\sim q_{\phi}({\bm{w}})$ we have,
$f_{L}({\bm{x}};{\bm{w}})\sim\mathcal{N}\mathopen{}\left(f({\bm{x}};{\bm{m}}),J({\bm{x}};{\bm{m}}){\bm{S}}J({\bm{x}};{\bm{m}})^{\top}\right)\mathclose{}\;\forall{\bm{x}}$,
and the function $f_{L}(\,\cdot\,;{\bm{w}})$ is a degenerate GP (as
$\operatorname{rank}(J(\,\cdot\,;{\bm{m}}){\bm{S}}J(\,\cdot\,;{\bm{m}})^{\top})\leq\text{number\leavevmode\nobreak\
of\leavevmode\nobreak\ weights}<\infty$),

$f_{L}\sim\mathcal{GP}\mathopen{}\left(f(\,\cdot\,;{\bm{m}}),J(\,\cdot\,;{\bm{m}}){\bm{S}}J(\,\cdot\,;{\bm{m}})^{\top}\right)\mathclose{}.$
(2.7)

#####
$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}$
is infinite in most practically relevant cases.

Burt et al., (2020) point out an even more severe issue of function-space VI
in BNNs:
$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}$
is in fact infinite in many relevant cases, in particular for non-degenerate
GP-priors with BNNs. Thus, approximations in these cases are futile. The proof
in Burt et al., (2020) is somewhat involved, but the fundamental reason for
$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}=\infty$
is that $\mathbb{Q}_{\phi}$ has support on a finite-dimensional submanifold of
the infinite-dimensional function space, while the measure $\mathbb{P}$
induced by a (non-degenerate) GP prior has support on the entire function
space. That such a dimensionality mismatch can lead to infinite KL divergence
can be seen in a finite-dimensional example: consider the KL-divergence
between two Gaussians in $\mathbb{R}^{n}$ for some $n\geq 2$, one of which has
support on the entire $\mathbb{R}^{n}$ (i.e., its covariance matrix
${\bm{\Sigma}}_{1}$ has full rank) while the other one has support only on a
proper subspace of $\mathbb{R}^{n}$ (i.e., its covariance matrix
${\bm{\Sigma}}_{2}$ is singular). The KL divergence between two multivariate
Gaussians has a closed form expression (Eq 2.9 with ${\gamma=0}$) that
contains
$\log\det\mathopen{}\left({\bm{\Sigma}}_{2}^{-1}{\bm{\Sigma}}_{1}\right)\mathclose{}$,
which is infinite for singular ${\bm{\Sigma}}_{2}$.

We find that the fact that
$D_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}\,\|\,\mathbb{P}\right)\mathclose{}=\infty$
has severe practical consequences even when the KL divergence is only
estimated from samples. It naturally explains the stability issues discussed
in Appendix D.1 of Sun et al., (2019) (we compare the authors’ solution to
this stability issue to our method in Section 3.2). Surprisingly, similar
complications arise even in the setup by Rudner et al., 2022b , which performs
VI in function space with the pushforward of a weight-space prior. While the
KL divergence is technically finite because prior and variational posterior
have the same support, numerical errors lead to mismatching supports and thus
to stability issues even there.

In summary, the ELBO for VI in BNNs is not well-defined for a large class of
interesting function-space priors. In section 3, we propose a solution by
using the so-called regularized KL divergence, which we introduce next.

###  2.2 Regularized KL divergence for Gaussian measures

Our solution to the negative infinite function-space ELBO builds on a
regularized KL divergence, which is expressed in terms of Gaussian measures
for the variational posterior and prior. We obtain these Gaussian measures
from corresponding Gaussian processes (GP). We first discuss under which
conditions a GP induces a Gaussian measure, and then present the regularized
KL divergence.

##### Gaussian measures and Gaussian processes

The regularized KL divergence is defined in terms of Gaussian measures, and
thus we need to verify that the GP variational posterior induced by the
linearized BNN (Eq 2.7) has an associated Gaussian measure. We consider the
Hilbert space $L^{2}(\mathcal{X},\rho)$ of square-integrable functions with
respect to a probability measure $\rho$ on a compact set
$\mathcal{X}\subset^{d}$, with inner product $\langle
f,g\rangle=\int_{\mathcal{X}}f(x)g(x)d\rho(x)$. This is not a restrictive
assumption as we can typically bound the region in feature space that contains
the data and any points where we might want to evaluate the BNN.

######  Definition 2.1 (Gaussian measure, Kerrigan et al., (2023), Definition
1).

Let $(\Omega,\mathcal{B},\mathbb{P})$ be a probability space. A measurable
function $F:\Omega\mapsto L^{2}(\mathcal{X},\rho)$ is called a Gaussian random
element (GRE) if for any $g\in L^{2}(\mathcal{X},\rho)$ the random variable
$\langle g,F\rangle$ has a Gaussian distribution on . For every GRE $F$, there
exists a unique mean element $m\in L^{2}(\mathcal{X},\rho)$ and a finite trace
linear covariance operator $C:L^{2}(\mathcal{X},\rho)\mapsto
L^{2}(\mathcal{X},\rho)$ such that $\langle
g,F\rangle\sim\mathcal{N}\mathopen{}\left(\langle g,m\rangle,\langle
Cg,g\rangle\right)\mathclose{}$ for all $g\in L^{2}(\mathcal{X},\rho)$. The
pushforward of $\mathbb{P}$ along $F$ denoted
$\mathbb{P}^{F}:=F_{\\#}\mathbb{P}$ is a Gaussian (probability) measure on
$L^{2}(\mathcal{X},\rho)$.

Gaussian measures generalize Gaussian distributions to infinite-dimensional
spaces where measures do not have associated densities. Following Wild et al.,
(2022), we notate the Gaussian measure obtained from the GRE $F$ with mean
element $m$ and covariance operator $C$ as
$\mathbb{P}^{F}:=\mathcal{N}\mathopen{}\left(m,C\right)\mathclose{}$.

Gaussian processes provide a convenient method to specify Gaussian measures
over functions via mean and covariance functions (Wild et al.,, 2022; Kerrigan
et al.,, 2023). A GP $f\sim\mathcal{GP}(\mu,K)$ has an associated Gaussian
measures in $L^{2}(\mathcal{X},\rho)$ if its mean function satisfies $\mu\in
L^{2}(\mathcal{X},\rho)$ and its covariance function $K$ is trace-class, i.e.,
if $\int_{\mathcal{X}}K(x,x)d\rho\mathopen{}\left(x\right)\mathclose{}<\infty$
(Wild et al.,, 2022, Theorem 1). The GP variational posterior induced by the
linearized BNN (Eq 2.7) satisfies both properties as neural networks are well-
behaved functions on the compact subset $\mathcal{X}\subset^{d}$. It thus
induces a Gaussian measure
$\mathbb{Q}_{\phi}^{F}\sim\mathcal{N}\mathopen{}\left(m_{Q},C_{Q}\right)\mathclose{}$
with mean element $m_{Q}=f(\,\cdot\,;{\bm{m}})$ and covariance operator
$C_{Q}g\mathopen{}\left(\cdot\right)\mathclose{}=\int_{\mathcal{X}}J(\,\cdot\,;{\bm{m}}){\bm{S}}J({\bm{x}}^{\prime},{\bm{m}})^{\top}g({\bm{x}}^{\prime})d\rho\mathopen{}\left({\bm{x}}^{\prime}\right)\mathclose{}$.
It turns out that the infinite KL divergence discussed in Section 2.1 is
easier to prove for Gaussian measures, and we provide the proof in section
A.1.1 for the interested reader.

######  Definition 2.2 (Regularized KL divergence, Quang, (2022) Definition
5).

Let $\nu_{1}=\mathcal{N}\mathopen{}\left(m_{1},C_{1}\right)\mathclose{}$ and
$\nu_{2}=\mathcal{N}\mathopen{}\left(m_{2},C_{2}\right)\mathclose{}$ be two
Gaussian measures with $m_{1},m_{2}\in L^{2}(\mathcal{X},\rho)$ and
$C_{1},C_{2}$ bounded, self-adjoint, positive and trace-class linear operators
on $L^{2}(\mathcal{X},\rho)$. Let $\gamma\in$, $\gamma>0$ be fixed. The
regularized KL divergence is defined as follows,

$\displaystyle
D^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\nu_{1}\,\|\,\nu_{2}\right)\mathclose{}:=$
$\displaystyle\,\frac{1}{2}\langle
m_{1}-m_{2},(C_{2}+\gamma\mathbb{I})^{-1}\mathopen{}\left(m_{1}-m_{2}\right)\mathclose{}\rangle+\frac{1}{2}\operatorname{Tr}_{X}\left[\mathopen{}\left(C_{2}+\gamma\mathbb{I}\right)\mathclose{}^{-1}\mathopen{}\left(C_{1}+\gamma\mathbb{I}\right)\mathclose{}-\mathbb{I}\right]$
$\displaystyle-\frac{1}{2}\log\operatorname{det}_{X}\left[\mathopen{}\left(C_{2}+\gamma\mathbb{I}\right)\mathclose{}^{-1}\mathopen{}\left(C_{1}+\gamma\mathbb{I}\right)\mathclose{}\right].$
(2.8)

For any $\gamma>0$, the regularized KL divergence is well-defined and finite,
even if the Gaussian measures are singular (Quang,, 2019). It converges to the
conventional KL divergence (if it is well-defined) for $\gamma\to 0$ (Quang,,
2022, Theorem 6). Furthermore, and importantly for our application to VI, if
two Gaussian measures $\nu_{1}$ and $\nu_{2}$ are induced by GPs
$f_{i}\sim\mathcal{GP}(\mu_{i},K_{i})$ for $i=1,2$, then
$D^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\nu_{1}\,\|\,\nu_{2}\right)\mathclose{}$
can be consistently estimated (Quang,, 2022) from a finite number $M$ of
samples using the estimator

$\displaystyle\hat{D}^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\nu_{1}\,\|\,\nu_{2}\right)\mathclose{}:=$
$\displaystyle\,\frac{1}{2}\mathopen{}\left({\bm{m}}_{1}-{\bm{m}}_{2}\right)\mathclose{}^{\top}({\bm{\Sigma}}^{(\gamma)}_{2})^{-1}\mathopen{}\left({\bm{m}}_{1}-{\bm{m}}_{2}\right)\mathclose{}+\frac{1}{2}\operatorname{Tr}\big{[}({\bm{\Sigma}}_{2}^{(\gamma)})^{-1}{\bm{\Sigma}}^{(\gamma)}_{1}-\mathbb{I}_{M}\big{]}$
$\displaystyle-\frac{1}{2}\log\det\big{[}({\bm{\Sigma}}_{2}^{(\gamma)})^{-1}{\bm{\Sigma}}^{(\gamma)}_{1}\big{]}$
(2.9)

with ${\bm{m}}_{i}:=\mu_{i}(\mathbf{x})$ and
${\bm{\Sigma}}_{i}^{(\gamma)}:=K_{i}(\mathbf{x},\mathbf{x})+\gamma
M\,\mathbb{I}_{M}$ where $\mu_{i}(\mathbf{x})$ and
$K_{i}(\mathbf{x},\mathbf{x})$ are the mean vector and the covariance matrix
obtained by evaluating $\mu_{i}$ and $K_{i}$ respectively, at measurement
points ${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}$,
${\bm{x}}^{(1)},\dots,{\bm{x}}^{(M)}\mathrel{\mathop{\sim}\limits^{\vbox
to1.20552pt{\kern-2.0pt\hbox{$\scriptstyle\textup{\tiny
i.i.d}$}\vss}}}\rho({\bm{x}})$. The right-hand side of Eq. 2.9 is the closed-
form expression for the KL-divergence between multivariate Gaussian
distributions $\mathcal{N}({{\bm{m}}_{1}},{{\bm{\Sigma}}_{1}^{(\gamma)}})$ and
$\mathcal{N}({{\bm{m}}_{2}},{{\bm{\Sigma}}_{2}^{(\gamma)}})$. Quang, (2022)
shows that the absolute error of the estimator is bounded by
$\mathcal{O}(1/M)$ with high probability with constants depending on $\gamma$
and properties of the GP mean and covariance functions (see section A.1.2 for
the exact bound).

##  3 Generalized function-space VI with the regularized KL divergence

This section presents our proposed generalized function-space variational
inference (GFSVI) method, which addresses the problem of the infinite KL
divergence discussed in section 2.1, which we take for an indication that VI
is too restrictive if one wants to use genuine function-space priors. We
instead consider generalized variational inference (Knoblauch et al.,, 2019),
which reinterprets the ELBO in Eq. 2.1 as a regularized expected log-
likelihood and explores alternative divergences for the regularizer.
Specifically, we propose to use the regularized KL divergence. The
presentation in this section builds heavily on tools introduced in section 2,
which turn out to fit together perfectly: the pushforward of a Gaussian
variational distribution in weight-space through the linearized neural network
(Eq. 2.6) induces a GP variational posterior (Eq. 2.7) that admits a Gaussian
measure on $L^{2}(\mathcal{X},\rho)$ (definition 2.1). Further selecting a GP
prior which has an associated Gaussian measure on $L^{2}(\mathcal{X},\rho)$
allows us to use the regularized KL divergence (Eq. 2.8). We present the GFSVI
method in section 3.1 and discuss connections to prior work in section 3.2.

###  3.1 Generalized function-space variational inference

We present a simple and well-defined objective for function-space inference,
and an algorithm for its optimization.

##### Objective function

We start from the ELBO in Eq. 2.3, where we use the Gaussian variational
measure $\mathbb{Q}_{\phi}^{F}$ induced by the pushforward of a Gaussian
variational distribution
$q_{\phi}({\bm{w}})=\mathcal{N}\mathopen{}\left({\bm{w}}\,|\,{\bm{m}},{\bm{S}}\right)\mathclose{}$
along the linearized network $f_{L}$ (Eq. 2.6). The function-space prior may
be any GP that has an associated Gaussian measure $\mathbb{P}^{F}$ on
$L^{2}(\mathcal{X},\rho)$ (see definition 2.1 and discussion below it). We now
replace the KL divergence in the ELBO with the regularized KL divergence
$D_{\text{KL}}^{\gamma}$ (Eq. 2.8), which is well-defined and finite for any
pair of Gaussian measures. For a likelihood
$p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}=\prod_{i=1}^{N}p\mathopen{}\left(y_{i}\,|\,f_{L}({\bm{x}}_{i};{\bm{w}})\right)\mathclose{}$,
we thus obtain

$\mathcal{L}(\phi)\\!:=\\!\\!\sum_{i=1}^{N}{\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\log
p\mathopen{}\left(y_{i}|f_{L}({\bm{x}}_{i};{\bm{w}})\right)\mathclose{}\right]\mathclose{}\\!-\\!D^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}^{F}\,\|\,\mathbb{P}^{F}\right)\mathclose{}.$
(3.1)

Input: linearized BNN $f_{L}$ with measure $\mathbb{Q}_{\phi}^{F}$,
Gaussian process prior $\mathcal{GP}(\mu,K)$ with measure
$\mathbb{P}^{F}$, measurement point distribution $\rho({\bm{x}})$, data
$\mathcal{D}=\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{N}$, batch size $B$,
$\gamma>0$.

for  _each minibatch
$({\bm{x}}_{\mathcal{B}},y_{\mathcal{B}})\sim\mathcal{D}$_ do

Calculate
$\hat{\ell}_{1}=\frac{N}{B}{\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\log
p\mathopen{}\left(y_{\mathcal{B}}\,|\,f_{L}({\bm{x}}_{\mathcal{B}},{\bm{w}})\right)\mathclose{}\right]\mathclose{}$;

[-1.5ex] Draw measurement points
${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}\stackrel{{\scriptstyle\\!\\!\text{i.i.d.}\\!\\!}}{{\sim}}\rho({\bm{x}})$;

[-1.5ex] Calculate
$\hat{\ell}_{2}=\hat{D}_{\text{KL}}^{\gamma}\big{(}\mathbb{Q}_{\phi}^{F}\,\big{|}\\!\big{|}\,{\mathbb{P}^{F}}\big{)}$
using ${\bm{x}}$ (Eq. 2.9);

[-1.5ex] Calculate $\hat{\mathcal{L}}(\phi)=\hat{\ell}_{1}-\hat{\ell}_{2}$;

Update $\phi$ using a step in the direction
$\nabla_{\\!\phi}\hat{\mathcal{L}}(\phi)$;

end for

Algorithm 1 Generalized function-space variational inference (GFSVI)

##### Estimation and optimization

The expected log-likelihood (first term of Eq. 3.1) can be estimated by
sampling from $q_{\phi}({\bm{w}})$. For a Gaussian likelihood, it can also be
computed in closed form because (unlike Rudner et al., 2022b ) we use the
linearized network $f_{L}$ to parameterize the likelihood function, which made
training more stable in our experiments. We estimate the regularized KL
divergence (second term of Eq. 3.1) using its consistent estimator (see Eq.
2.9), with ${\bm{m}}_{1}=f({\bm{x}};{\bm{m}})$,
${\bm{\Sigma}}_{1}^{(\gamma)}=J({\bm{x}};{\bm{m}}){\bm{S}}J({\bm{x}};{\bm{m}})^{\top}+\gamma
M\mathbb{I}_{M}$, ${\bm{m}}_{2}=\mu({\bm{x}})$, and
${\bm{\Sigma}}_{2}^{(\gamma)}=K({\bm{x}},{\bm{x}})+\gamma M\mathbb{I}_{M}$,
where $\mu$ and $K$ are the mean function and kernel of the GP prior, and
${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}$,
${\bm{x}}^{(1)},\dots,{\bm{x}}^{(M)}\mathrel{\mathop{\sim}\limits^{\vbox
to1.20552pt{\kern-2.0pt\hbox{$\scriptstyle\textup{\tiny
i.i.d}$}\vss}}}\rho({\bm{x}})$ are measurement points sampled from the
probability measure $\rho({\bm{x}})$ on $\mathcal{X}$. We maximize the
estimated objective over the mean ${\bm{m}}$ and covariance ${\bm{S}}$ of the
Gaussian variational distribution $q_{\phi}({\bm{w}})$, and over any
likelihood parameter (e.g., $\sigma_{y}$ for a Gaussian likelihood), see
Algorithm 1. section A.2 provides explicit expressions for the estimator with
Gaussian and Categorical likelihoods.

##### Technical details ($\gamma$ and $\rho$)

It turns out that increasing $\gamma$ reduces the influence of the prior on
inference (see figure 17 in Appendix). At the same time, $\gamma$ acts as
jitter that prevents numerical errors (see section 3.2). We recommend setting
$\gamma$ large enough to avoid numerical errors but sufficiently small to
strongly regularize the objective in Eq. 3.1 (see figure 15 in Appendix). We
found that GFSVI is robust to a wide range of values and fixed
$\gamma=10^{-10}$. Furthermore, the probability measure $\rho$ defined with
$L^{2}(\mathcal{X},\rho)$ has to assign non-zero probability to any open set
of $\mathcal{X}$ to regularize the BNN on all of its support. Following Rudner
et al., 2022b , we draw measurement points from a uniform distribution over
$\mathcal{X}$ when using tabular data and explore different configurations
(such as samples from other data sets) when using high-dimensional image data
(see section A.3.4).

###  3.2 Connections to prior work

TFSVI (Rudner et al., 2022b, ) and FVI (Sun et al.,, 2019) solve stability
issues by introducing jitter/noise, which has a similar effect as the
regularization in Eq. 2.8. However, as mentioned in Section 2.1, TFSVI
introduces jitter only to overcome numerical issues and is fundamentally
restricted to prior specification in weight space since its function-space
prior is the pushforward of a weight-space prior. FVI does not linearize the
BNN, and therefore does not have access to an explicit variational measure in
function space. This severely complicates the estimation of (gradients of) the
KL divergence, and the authors resort to implicit score function estimators,
which make their method difficult to use in practice (Ma and Hernández-
Lobato,, 2021). Our proposed GFSVI does not suffer from these difficulties as
the variational posterior is an explicit (Gaussian) measure. This allows us to
estimate the (regularized) KL divergence without sampling any noise or using
implicit score function estimators.

##  4 Experiments

Figure 2: Inference on synthetic data (gray circles) using an RBF prior for
function-space methods GFSVI and FVI. The proposed GFSVI provides the best
approximation of the exact GP posterior.

In this section, we evaluate our generalized function-space variational
inference (GFSVI) method qualitatively on synthetic data (Section 4.1) and
quantitatively on real-world data (Section 4.2). We find that GFSVI accurately
captures structural properties specified by the GP prior, that it approximates
the exact Gaussian process (GP) posterior more faithfully than BNN baselines,
and that it performs competitively on regression, classification and out-of-
distribution detection tasks. We also discuss the influence of the BNN’s
inductive biases.

##### Baselines

We compare the proposed GFSVI method to two weight-space inference methods
(mean-field variational inference (Blundell et al.,, 2015) and linearized
Laplace (Immer et al.,, 2021)) and to three function-space inference methods
(FVI (Sun et al.,, 2019), TFVSI (Rudner et al., 2022b, ) and VIP (Ma et al.,,
2019), where TFSVI performs inference in function space but with the
pushforward of a weight-space prior and VIP using a BNN prior). All BNNs have
the same architecture and fully-factorized Gaussian approximate posterior. We
also include results for a sparse GP with a posterior mean parameterized by a
neural network (GWI) (Wild et al.,, 2022), as well as for a Gaussian Process
(GP) (Williams and Rasmussen,, 2006) (when the size of the dataset allows it),
and for a sparse GP (Hensman et al.,, 2013) for regression tasks. Ẇe consider
the GP, sparse GP and GWI as gold standards as they represent the exact (or
near exact) posterior for models with GP priors.

###  4.1 Qualitative results on synthetic data

##### Regression

We consider synthetic regression data with a one-dimensional feature-space
$\mathcal{X}\subset$, where the values $y_{i}$ are randomly sampled around
$\sin(2\pi x_{i})$ (circles in Figures 1-11, see Appendix A.3.1). The green
lines in Figures 1-8 and 11 show functions sampled from the (approximate)
posteriors, and the red lines are the inferred mean functions. Figure 2
compares GFSVI with an RBF GP prior to baselines and to the exact posterior.
We find that GFSVI visually matches the true posterior best (the same holds
for a Matern-1/2 prior, see Figure 5 in the Appendix). Figure 1 in the
Introduction and Figure 6 in the Appendix show that GFSVI notably adapts to
varying prior assumptions (varying smoothness and length scale, respectively).
In addition, Figures 4 and 8 in the Appendix show that GFSVI provides strong
regularization when the data generative process is noisy, and that it can be
trained with fewer measurement points $M$ than FVI without significant
degradation.

##### Classification

We further consider a binary classification task derived from the
2-dimensional two moons data set (Pedregosa et al.,, 2011), see red and blue
dots in Figures 9 and 10 in the Appendix. The first and second row in both
figures show the inferred mean probability of class 1 (blue dots) and its
2-standard deviations with respect to posterior samples. Just like for
regression, we find that GFSVI captures the beliefs of the RBF and Matern-1/2
GP priors better than all BNN-baselines, and shows greater uncertainty outside
of the support of the data than the BNN baselines.

##### Inductive biases

figure 11 in the Appendix compares GFSVI to the exact GP-posterior across two
different priors and three model architectures (details in Appendix A.3.1). We
find that, with ReLU activations, small models are prone to underfitting for
smooth priors (RBF), and to collapsing uncertainty for rough priors
(Matern-1/2). By contrast, with smooth activations (Tanh), smaller models
suffice, and they are compatible with most standard GP priors (the results
shown in figure 11 extend to RBF, Matern, and Rational Quadratic in our
experiments). We also analyzed how the number $M$ of measurement points
affects performance. figures 7 and 14 in Appendix show that capturing the
beliefs of rough GP priors and estimating $D_{\text{KL}}^{\gamma}$ with these
priors requires larger $M$.

Table 1: Test expected log-likelihood (higher is better) of evaluated methods
on regression datasets. GFSVI performs competitively compared to all BNN
baselines and obtains the best mean rank.

Dataset | Function-space priors | Weight-space priors | Gaussian Processes (Gold Standards)
---|---|---|---
| GFSVI (ours) | FVI | TFSVI | MFVI | VIP | Laplace | GWI | Sparse GP | GP
Boston | -0.733 $\pm$ 0.144 | -0.571 $\pm$ 0.113 |  -1.416 $\pm$ 0.046 |  -1.308 $\pm$ 0.052 | -0.722 $\pm$ 0.196 | -0.812 $\pm$ 0.205 |  -0.940 $\pm$ 0.145 |  -0.884 $\pm$ 0.182 |  -1.594 $\pm$ 0.556
Concrete | -0.457 $\pm$ 0.041 | -0.390 $\pm$ 0.017 |  -0.983 $\pm$ 0.012 |  -1.353 $\pm$ 0.018 | -0.427 $\pm$ 0.050 |  -0.715 $\pm$ 0.025 |  -0.744 $\pm$ 0.079 |  -0.966 $\pm$ 0.025 |  -2.099 $\pm$ 0.421
Energy |  -1.319 $\pm$ 0.052 |  -1.377 $\pm$ 0.042 |  -0.797 $\pm$ 0.098 |  -0.926 $\pm$ 0.197 |  -1.046 $\pm$ 0.378 |  -1.304 $\pm$ 0.043 |  -0.461 $\pm$ 0.093 |  -0.206 $\pm$ 0.027 |  -0.205 $\pm$ 0.022
Kin8nm |  -0.136 $\pm$ 0.013 |  -0.141 $\pm$ 0.023 |  -0.182 $\pm$ 0.011 |  -0.641 $\pm$ 0.225 | -0.102 $\pm$ 0.013 |  -0.285 $\pm$ 0.014 |  -0.708 $\pm$ 0.054 |  -0.443 $\pm$ 0.014 | (infeasible)
Naval |  -3.637 $\pm$ 0.132 |  -2.165 $\pm$ 0.194 |  -2.758 $\pm$ 0.044 |  -1.034 $\pm$ 0.160 |  -1.502 $\pm$ 0.061 |  -3.404 $\pm$ 0.084 |  -0.301 $\pm$ 0.254 |  -4.951 $\pm$ 0.014 | (infeasible)
Power |  -0.044 $\pm$ 0.011 |  -0.031 $\pm$ 0.021 |  -0.007 $\pm$ 0.013 |  -0.003 $\pm$ 0.015 |  -0.036 $\pm$ 0.018 |  -0.002 $\pm$ 0.019 |  -0.043 $\pm$ 0.009 |  -0.100 $\pm$ 0.010 | (infeasible)
Protein |  -1.036 $\pm$ 0.005 |  -1.045 $\pm$ 0.005 |  -1.010 $\pm$ 0.004 |  -1.112 $\pm$ 0.007 |  -0.994 $\pm$ 0.007 |  -1.037 $\pm$ 0.006 |  -1.050 $\pm$ 0.009 |  -1.035 $\pm$ 0.002 | (infeasible)
Wine |  -1.289 $\pm$ 0.040 | -1.215 $\pm$ 0.007 |  -2.138 $\pm$ 0.221 |  -1.248 $\pm$ 0.018 |  -1.262 $\pm$ 0.025 |  -1.249 $\pm$ 0.025 |  -1.232 $\pm$ 0.038 |  -1.240 $\pm$ 0.037 |  -1.219 $\pm$ 0.035
Yacht |  -1.058 $\pm$ 0.080 |  -0.545 $\pm$ 0.735 |  -1.187 $\pm$ 0.064 |  -1.638 $\pm$ 0.030 |  -0.062 $\pm$ 1.378 |  -0.680 $\pm$ 0.171 |  -0.441 $\pm$ 0.138 |  -0.976 $\pm$ 0.092 |  -0.914 $\pm$ 0.045
Wave |  -5.521 $\pm$ 0.036 |  -6.612 $\pm$ 0.008 |  -5.148 $\pm$ 0.117 |  -6.883 $\pm$ 0.008 |  -4.043 $\pm$ 0.093 |  -4.658 $\pm$ 0.027 |  -1.566 $\pm$ 0.123 |  -4.909 $\pm$ 0.001 | (infeasible)
Denmark | -0.487 $\pm$ 0.013 |  -0.801 $\pm$ 0.005 |  -0.513 $\pm$ 0.013 |  -0.675 $\pm$ 0.007 |  -0.583 $\pm$ 0.021 |  -0.600 $\pm$ 0.008 |  -0.841 $\pm$ 0.026 |  -0.768 $\pm$ 0.001 | (infeasible)
Mean rank | 1.545 | 2.000 | 2.727 | 3.455 | 2.091 | 2.455 | - | - | -

Table 2:  Test expected log-likelihood, accuracy, expected calibration error
and OOD detection accuracy on MNIST and Fashion MNIST.

Data | Metric | Function-space priors | Weight-space priors | GP-based
---|---|---|---|---
| GFSVI (rnd) | GFSVI (kmnist) | FVI (rnd) | FVI (kmnist) | TFSVI (rnd) | TFSVI (kmnist) | MFVI | VIP | Laplace | GWI
MNIST |  Log-like. ($\uparrow$) | -0.033 $\pm$ 0.000 |  -0.041 $\pm$ 0.000 |  -0.145 $\pm$ 0.005 |  -0.238 $\pm$ 0.006 |  -0.047 $\pm$ 0.003 |  -0.041 $\pm$ 0.001 |  -0.078 $\pm$ 0.001 | -0.033 $\pm$ 0.001 |  -0.108 $\pm$ 0.002 |  -0.090 $\pm$ 0.003
Acc. ($\uparrow$) |  -0.992 $\pm$ 0.000 |  -0.991 $\pm$ 0.000 |  -0.976 $\pm$ 0.001 |  -0.943 $\pm$ 0.001 |  -0.989 $\pm$ 0.000 |  -0.989 $\pm$ 0.000 |  -0.990 $\pm$ 0.000 |  -0.989 $\pm$ 0.000 |  -0.984 $\pm$ 0.000 |  -0.971 $\pm$ 0.001
|  ECE ($\downarrow$) |  -0.002 $\pm$ 0.000 |  -0.006 $\pm$ 0.000 |  -0.064 $\pm$ 0.001 |  -0.073 $\pm$ 0.003 |  -0.007 $\pm$ 0.000 |  -0.006 $\pm$ 0.000 |  -0.021 $\pm$ 0.000 |  -0.002 $\pm$ 0.001 |  -0.048 $\pm$ 0.001 |  -0.003 $\pm$ 0.000
|  OOD acc. ($\uparrow$) |  -0.921 $\pm$ 0.008 |  -0.980 $\pm$ 0.004 |  -0.894 $\pm$ 0.010 |  -0.891 $\pm$ 0.006 |  -0.887 $\pm$ 0.011 |  -0.893 $\pm$ 0.005 |  -0.928 $\pm$ 0.002 |  -0.871 $\pm$ 0.012 |  -0.903 $\pm$ 0.007 |  -0.829 $\pm$ 0.007
FMNIST |  Log-like. ($\uparrow$) |  -0.260 $\pm$ 0.003 |  -0.294 $\pm$ 0.006 |  -0.300 $\pm$ 0.002 |  -0.311 $\pm$ 0.005 |  -0.261 $\pm$ 0.001 |  -0.261 $\pm$ 0.002 |  -0.290 $\pm$ 0.002 | -0.252 $\pm$ 0.001 |  -0.426 $\pm$ 0.009 |  -0.260 $\pm$ 0.001
Acc. ($\uparrow$) |  -0.910 $\pm$ 0.001 |  -0.909 $\pm$ 0.001 |  -0.910 $\pm$ 0.002 |  -0.906 $\pm$ 0.002 |  -0.909 $\pm$ 0.001 |  -0.907 $\pm$ 0.001 |  -0.913 $\pm$ 0.001 |  -0.911 $\pm$ 0.001 |  -0.886 $\pm$ 0.001 |  -0.906 $\pm$ 0.000
|  ECE ($\downarrow$) |  -0.020 $\pm$ 0.003 |  -0.042 $\pm$ 0.002 |  -0.027 $\pm$ 0.005 |  -0.024 $\pm$ 0.002 |  -0.022 $\pm$ 0.002 |  -0.021 $\pm$ 0.002 |  -0.010 $\pm$ 0.001 |  -0.024 $\pm$ 0.001 |  -0.060 $\pm$ 0.004 |  -0.016 $\pm$ 0.001
|  OOD acc. ($\uparrow$) |  -0.853 $\pm$ 0.005 |  -0.997 $\pm$ 0.001 |  -0.925 $\pm$ 0.005 |  -0.975 $\pm$ 0.002 |  -0.802 $\pm$ 0.006 |  -0.779 $\pm$ 0.010 |  -0.805 $\pm$ 0.010 |  -0.790 $\pm$ 0.010 |  -0.826 $\pm$ 0.006 |  -0.792 $\pm$ 0.005

###  4.2 Quantitative results on real-world data

We study the fidelity of GFSVI compared to the true GP-posterior on real data,
and evaluate GFSVI on regression, classification, and out-of-distribution
detection tasks. In all tables, we bold the highest score and any score whose
error bar (standard error) overlaps with the highest score’s error bar.

##### Ocean current modeling

Figure 3: Results for the ocean current modeling experiment. We report the
mean velocity vectors and the squared errors of compared methods. Unlike
TFSVI, we find that GFSVI accurately captures ocean current dynamics.

We test the extend to which GFSVI can incorporate prior knowledge specified
via a GP prior by considering the problem of modeling ocean currents in the
Gulf of Mexico. We are given $20$ two-dimensional velocity vectors collected
from ocean drifter buoys, and we are interested in estimating ocean currents.
We use the GulfDrifters dataset (Lilly and Pérez-Brunius,, 2021) and we follow
the setup by Shalashilin, (2024). We compare our GFSVI to a GP and to TFSVI.
We incorporate known physical properties of fluid motions by applying the
Helmholtz decomposition to the GP prior (Berlinghieri et al.,, 2023) as well
as to the neural networks. More details can be found in section A.3.2.

Table 3: Results for the ocean current modeling task. We find that
incorporating prior knowledge via a Gaussian process prior in GFSVI improves
score over weight-space priors in TFSVI.

Dataset | GFSVI (ours) | TFSVI | GP
---|---|---|---
Log-like. |  -6.627 $\pm$ 0.753 |  -22.651 $\pm$ 2.947 |  -0.507 $\pm$ 0.000
MSE |  -0.021 $\pm$ 0.002 |  -0.034 $\pm$ 0.003 |  -0.013 $\pm$ 0.000

Results are shown in table 3 and in figure 3. We find that incorporating prior
knowledge via an informative GP prior in GFSVI improves expected log-
likelihood (Log-like.) and mean squared-error (MSE) over weight-space priors
in TFSVI. However, for this relatively small dataset, exact GP inference is
still possible, and it outperforms both BNN-based predictions. This suggests
that the physically motivated GP describes the true fluid dynamics well enough
that the additional inductive bias introduced by a neural network hurts
performance rather than helping it. In the following, we consider experiments
with larger datasets (making exact GP inference computationally infeasible in
many cases), and where structural prior knowledge in function space exists but
is not derived from strict laws of nature.

##### Regression

We assess the predictive performance of GFSVI on data sets from the UCI
repository (Dua and Graff,, 2017) (described in table 6 in the Appendix).
Table 1, and Table 7 in the Appendix, show expected log-likelihood and mean
squared error, respectively. We perform 5-fold cross validation and report
means and standard errors across the test folds. We also rank the methods for
each dataset and report the mean rank of each method across all datasets. See
section A.3.3 for more details. We find that GFSVI performs competitively
compared to baselines and obtains the best mean rank for both metrics,
matching the top performing methods on nearly all datasets. In particular, we
find that using GP priors in the linearized BNN with GFSVI yields improvements
over the weight-space priors used in TFSVI and that GFSVI performs slightly
better than FVI.

Table 4: Out-of-distribution accuracy (higher is better) of evaluated methods
on regression datasets. GFSVI (ours) performs competitively on OOD detection
and obtains the highest mean rank.

Dataset | Function-space priors | Weight-space priors | Gaussian Processes (Gold Standards) |
---|---|---|---|---
| GFSVI (ours) | FVI | TFSVI | MFVI | VIP | Laplace | GWI | Sparse GP | GP |
Boston | 0.893 $\pm$ 0.011 |  0.594 $\pm$ 0.024 |  0.705 $\pm$ 0.107 |  0.563 $\pm$ 0.013 |  0.628 $\pm$ 0.010 |  0.557 $\pm$ 0.009 |  0.817 $\pm$ 0.017 |  0.947 $\pm$ 0.011 |  0.952 $\pm$ 0.003 |
Concrete | 0.656 $\pm$ 0.016 |  0.583 $\pm$ 0.022 |  0.511 $\pm$ 0.003 |  0.605 $\pm$ 0.012 |  0.601 $\pm$ 0.024 |  0.578 $\pm$ 0.015 |  0.730 $\pm$ 0.020 |  0.776 $\pm$ 0.006 |  0.933 $\pm$ 0.004 |
Energy | 0.997 $\pm$ 0.002 |  0.696 $\pm$ 0.017 | 0.997 $\pm$ 0.001 |  0.678 $\pm$ 0.014 |  0.682 $\pm$ 0.037 |  0.782 $\pm$ 0.020 |  0.998 $\pm$ 0.001 |  0.998 $\pm$ 0.001 |  0.998 $\pm$ 0.001 |
Kin8nm |  0.588 $\pm$ 0.007 | 0.604 $\pm$ 0.023 |  0.576 $\pm$ 0.008 |  0.570 $\pm$ 0.009 |  0.563 $\pm$ 0.015 | 0.606 $\pm$ 0.009 |  0.602 $\pm$ 0.011 |  0.608 $\pm$ 0.014 | (infeasible) |
Naval | 1.000 $\pm$ 0.000 | 1.000 $\pm$ 0.000 | 1.000 $\pm$ 0.000 |  0.919 $\pm$ 0.017 |  0.621 $\pm$ 0.059 | 1.000 $\pm$ 0.000 |  1.000 $\pm$ 0.000 |  1.000 $\pm$ 0.000 | (infeasible) |
Power | 0.698 $\pm$ 0.006 |  0.663 $\pm$ 0.021 |  0.676 $\pm$ 0.008 |  0.636 $\pm$ 0.019 |  0.514 $\pm$ 0.004 |  0.654 $\pm$ 0.013 |  0.754 $\pm$ 0.004 |  0.717 $\pm$ 0.004 | (infeasible) |
Protein | 0.860 $\pm$ 0.011 |  0.810 $\pm$ 0.022 | 0.841 $\pm$ 0.018 |  0.693 $\pm$ 0.020 | 0.549 $\pm$ 0.020 |  0.629 $\pm$ 0.013 |  0.942 $\pm$ 0.002 |  0.967 $\pm$ 0.001 | (infeasible) |
Wine |  0.665 $\pm$ 0.013 |  0.517 $\pm$ 0.004 |  0.549 $\pm$ 0.015 |  0.542 $\pm$ 0.009 | 0.706 $\pm$ 0.028 |  0.531 $\pm$ 0.007 |  0.810 $\pm$ 0.008 |  0.781 $\pm$ 0.014 |  0.787 $\pm$ 0.007 |
Yacht |  0.616 $\pm$ 0.030 |  0.604 $\pm$ 0.025 | 0.659 $\pm$ 0.043 | 0.642 $\pm$ 0.035 | 0.688 $\pm$ 0.040 |  0.612 $\pm$ 0.024 |  0.563 $\pm$ 0.014 |  0.762 $\pm$ 0.018 |  0.787 $\pm$ 0.011 |
Wave | 0.975 $\pm$ 0.005 |  0.642 $\pm$ 0.004 |  0.835 $\pm$ 0.034 |  0.658 $\pm$ 0.026 |  0.500 $\pm$ 0.000 |  0.529 $\pm$ 0.005 |  0.903 $\pm$ 0.001 |  0.513 $\pm$ 0.001 | (infeasible) |
Denmark |  0.521 $\pm$ 0.006 | 0.612 $\pm$ 0.008 |  0.519 $\pm$ 0.006 |  0.513 $\pm$ 0.003 |  0.500 $\pm$ 0.000 |  0.529 $\pm$ 0.008 |  0.688 $\pm$ 0.003 |  0.626 $\pm$ 0.002 | (infeasible) |
Mean rank | 1.455 | 2.364 | 1.909 | 2.909 | 3.364 | 2.909 | - | - | - |

Table 5:  Average point-wise Wasserstein-2 distance (lower is better) between
exact and approximate posterior of reported methods. GFSVI (ours) provides a
more accurate approximation than FVI.

Dataset | Boston | Concrete | Energy | Wine | Yacht | Mean rank
---|---|---|---|---|---|---
GFSVI (ours) | 0.0259 $\pm$ 0.0040 | 0.0499 $\pm$ 0.0029 | 0.0035 $\pm$ 0.0004 | 0.0571 $\pm$ 0.0097 | 0.0036 $\pm$ 0.0006 | 1.0
FVI |  0.0469 $\pm$ 0.0044 |  0.0652 $\pm$ 0.0037 | 0.0037 $\pm$ 0.0004 |  0.1224 $\pm$ 0.0167 | 0.0052 $\pm$ 0.0013 | 1.6
GP sparse |  0.0074 $\pm$ 0.0022 |  0.0125 $\pm$ 0.0016 |  0.0042 $\pm$ 0.0003 |  0.0170 $\pm$ 0.0035 |  0.0035 $\pm$ 0.0008 | -

##### Classification

We further evaluate the classification performance of our method on the MNIST
(LeCun et al.,, 2010) and FashionMNIST (Xiao et al.,, 2017) image data sets.
We fit the models on a random subset of 90% of the training set, use the
remaining 10% as validation data, and evaluate on the provided test split. We
repeat with 5 different random seeds and report the mean and standard error of
the expected log-likelihood, accuracy and expected calibration error (ECE) of
the predictive distribution’s mean in table 2. For GFSVI, FVI, and TFSVI, we
tested measurement points from both a random (rnd) distribution
$\rho({\bm{x}})$ and from kmnist (see Algorithm 1 and section A.3.4). We find
that GFSVI performs competitively on MNIST, exceeding the expected log-
likelihood and accuracy of top-scoring baselines (see table 2) and similarly
to best baselines on FashionMNIST. Additionally, GFSVI yields well-calibrated
models with low ECE.

##### Out-of-distribution detection

We further evaluate our method by testing if its epistemic uncertainty is
predictive of out-of-distribution (OOD) data. We consider two settings: (i)
with tabular data and a Gaussian likelihood following Malinin et al., (2020),
and (ii) with image data and a categorical likelihood following Osawa et al.,
(2019). We report the accuracy of classifying OOD from in-distribution (ID)
data using a (learned) threshold on the predictive uncertainty. Additional
details are provided in section A.3.5. In setting (i) (table 4), we find that
GFSVI performs competitively and obtains the highest mean rank. Likewise, in
setting (ii) (table 8 in the Appendix), GFSVI strongly outperforms all
baselines when using the kmnist measurement point distribution
$\rho({\bm{x}})$. We find that with high-dimensional image data, the choice of
measurement point distribution highly influences OOD detection accuracy. We
provide more details on this point in Appendix A.4.4. In both settings, we
find that using GP priors with GFSVI rather than weight-space priors with
TFSVI is beneficial, and that GFSVI improves over FVI.

##### Variational measure evaluation

table 5 evaluates our inference method by comparing samples drawn from the
exact posterior (where computationally feasible) with the approximate
posterior obtained with our method (GFSVI). More details are provided in
section A.3.6. We find that GFSVI approximates the exact posterior more
accurately that FVI, obtaining a higher mean rank, but worse than the gold
standard sparse GP, which demonstrates to be most accurate.

##  5 Related work

In this section, we review related work on function-space variational
inference with neural networks, and on approximating functions-space measures
with weight-space priors.

##### Function-space inference with neural networks

Prior work on function-space VI in BNNs has addressed issues (i) and (ii)
mentioned in Section 2.1. Sun et al., (2019) address (i) (intractable
variational posterior in function space) by using implicit score function
estimators, and (ii) (intractable KL divergence) by replacing the supremum
with an expectation. Rudner et al., 2022a ; Rudner et al., 2022b  address (i)
by using a linearized BNN (Khan et al.,, 2020; Immer et al.,, 2021; Maddox et
al.,, 2021), and (ii) by replacing the supremum with a maximum over a finite
set of samples. Other work abandons approximating the posterior over neural
network weights altogether and instead uses a BNN only to specify a prior (Ma
et al.,, 2019), or deterministic neural networks to fit basis functions for
Bayesian linear regression (Ma and Hernández-Lobato,, 2021) or the mean of a
generalized sparse GP with Wasserstein-2 metric (Wild et al.,, 2022). Our work
combines linearized BNNs with generalized variational inference, but we use
the regularized KL divergence (Quang,, 2019), which naturally generalizes the
conventional KL divergence and allows for informative GP priors.

##### Approximating function-space measures with weight-space priors in BNNs

Flam-Shepherd et al., (2017); Tran et al., (2022) minimize a divergence
between the BNN’s prior predictive and a GP before performing inference on
weights, while Wu et al., (2023) directly incorporate the bridging divergence
inside the inference objective. Alternatively, Pearce et al., (2020) derive
BNN architectures mirroring GPs, and Matsubara et al., (2022) use the Ridgelet
transform to design weight-spaces priors approximating a GP in function space.
Another line of work considers weight-space priors which regularize in
function space by comparing the model’s predictions to those of a reference
model (Nalisnick et al.,, 2020) and using an empirical prior (Rudner et al.,,
2023).

##  6 Discussion

We proposed a simple inference method with a well-defined variational
objective function for Bayesian neural networks with Gaussian process (GP)
priors in function-space. As standard VI with functions-space priors suffers
from an infinite KL divergence problem, we propose to follow the generalized
VI framework. Specifically, we substitute the conventional KL divergence in
the ELBO by the regularized KL divergence, which is always finite, and which
can be estimated consistently within the linearized BNN approximation. We
demonstrated that our method allows to incorporate interpretable structural
properties via a GP prior, accurately approximates the true GP posterior on
synthetic and small real-world data sets, and provides competitive uncertainty
estimates for regression, classification and out-of-distribution detection
compared to BNN baselines with both function-space and weight-space priors.

Future work should investigate the use of more expressive variational
distributions, such as Gaussian with low-rank plus diagonal covariance
proposed by Tomczak et al., (2020), which is compatible with our proposed
method.

## Acknowledgments

This work was funded by the Deutsche Forschungsgemeinschaft (DFG, German
Research Foundation) under Germany’s Excellence Strategy – EXC number 2064/1 –
Project number 390727645. Additional support was provided by the German
Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ:
01IS18039A. Robert Bamler acknowledges funding by the German Research
Foundation (DFG) for project 448588364 of the Emmy Noether Programme. The
authors extend their gratitude to the International Max Planck Research School
for Intelligent Systems (IMPRS-IS) for supporting Tristan Cinquin. Finally, we
thank Vincent Fortuin and Marvin Pförtner for useful discussions and feedback.

## References

  * Abdullah et al., (2022) Abdullah, A., Hassan, M., and Mustafa, Y. (2022).  A review on Bayesian deep learning in healthcare: Applications and challenges.  IEEE Access, 10:1–1.
  * Belkin, (2018) Belkin, M. (2018).  Approximation beats concentration? an approximation view on inference with smooth radial kernels.
  * Berlinghieri et al., (2023) Berlinghieri, R., Trippe, B. L., Burt, D. R., Giordano, R., Srinivasan, K., Özgökmen, T., Xia, J., and Broderick, T. (2023).  Gaussian processes at the helm(holtz): A more fluid model for ocean currents.
  * Bertin-Mahieux et al., (2011) Bertin-Mahieux, T., Ellis, D. P., Whitman, B., and Lamere, P. (2011).  The million song dataset.  In Proceedings of the 12th International Conference on Music Information Retrieval (ISMIR 2011).
  * Bew et al., (2019) Bew, D., Harvey, C. R., Ledford, A., Radnor, S., and Sinclair, A. (2019).  Modeling analysts’ recommendations via Bayesian machine learning.  The Journal of Financial Data Science, 1(1):75–98.
  * Biewald, (2020) Biewald, L. (2020).  Experiment tracking with weights and biases.  Software available from wandb.com.
  * Blundell et al., (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015).  Weight uncertainty in neural networks.
  * Bradbury et al., (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018).  JAX: composable transformations of Python+NumPy programs.
  * Burt et al., (2020) Burt, D. R., Ober, S. W., Garriga-Alonso, A., and van der Wilk, M. (2020).  Understanding variational inference in function-space.  arXiv preprint arXiv:2011.09421.
  * Chen et al., (2021) Chen, H., Zheng, L., Kontar, R. A., and Raskutti, G. (2021).  Gaussian process inference using mini-batch stochastic gradient descent: Convergence guarantees and empirical benefits.
  * Cinquin et al., (2021) Cinquin, T., Immer, A., Horn, M., and Fortuin, V. (2021).  Pathologies in priors and inference for bayesian transformers.  In I (Still) Can’t Believe It’s Not Better! NeurIPS 2021 Workshop.
  * Clanuwat et al., (2018) Clanuwat, T., Bober-Irizar, M., Kitamoto, A., Lamb, A., Kazuaki, Y., and Ha, D. (2018).  Deep learning for classical japanese literature.
  * Dua and Graff, (2017) Dua, D. and Graff, C. (2017).  Uci machine learning repository.
  * Flam-Shepherd et al., (2017) Flam-Shepherd, D., Requeima, J., and Duvenaud, D. (2017).  Mapping gaussian process priors to bayesian neural networks.
  * Flamary et al., (2021) Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. (2021).  Pot: Python optimal transport.  Journal of Machine Learning Research, 22(78):1–8.
  * Hennigan et al., (2020) Hennigan, T., Cai, T., Norman, T., Martens, L., and Babuschkin, I. (2020).  Haiku: Sonnet for JAX.
  * Hensman et al., (2013) Hensman, J., Fusi, N., and Lawrence, N. D. (2013).  Gaussian processes for big data.
  * Hoffman et al., (2013) Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. (2013).  Stochastic variational inference.  Journal of Machine Learning Research, 14(40):1303–1347.
  * Immer et al., (2021) Immer, A., Korzepa, M., and Bauer, M. (2021).  Improving predictions of bayesian neural nets via local linearization.
  * Kerrigan et al., (2023) Kerrigan, G., Ley, J., and Smyth, P. (2023).  Diffusion generative models in infinite dimensions.
  * Khan et al., (2020) Khan, M. E., Immer, A., Abedi, E., and Korzepa, M. (2020).  Approximate inference turns deep networks into gaussian processes.
  * Kingma and Ba, (2017) Kingma, D. P. and Ba, J. (2017).  Adam: A method for stochastic optimization.
  * Knoblauch et al., (2019) Knoblauch, J., Jewson, J., and Damoulas, T. (2019).  Generalized variational inference: Three arguments for deriving new posteriors.  arXiv preprint arXiv:1904.02063.
  * Kompa et al., (2021) Kompa, B., Snoek, J., and Beam, A. (2021).  Second opinion needed: communicating uncertainty in medical machine learning.  npj Digital Medicine, 4.
  * LeCun et al., (2010) LeCun, Y., Cortes, C., and Burges, C. (2010).  Mnist handwritten digit database.  ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2\.
  * Lilly and Pérez-Brunius, (2021) Lilly, J. M. and Pérez-Brunius, P. (2021).  GulfDrifters: A consolidated surface drifter dataset for the Gulf of Mexico.
  * Ma and Hernández-Lobato, (2021) Ma, C. and Hernández-Lobato, J. M. (2021).  Functional variational inference based on stochastic process generators.  In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 21795–21807. Curran Associates, Inc.
  * Ma et al., (2019) Ma, C., Li, Y., and Hernández-Lobato, J. M. (2019).  Variational implicit processes.
  * Maddox et al., (2021) Maddox, W., Tang, S., Moreno, P., Gordon Wilson, A., and Damianou, A. (2021).  Fast adaptation with linearized neural networks.  In Banerjee, A. and Fukumizu, K., editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 2737–2745. PMLR.
  * Malinin et al., (2020) Malinin, A., Prokhorenkova, L., and Ustimenko, A. (2020).  Uncertainty in gradient boosting via ensembles.
  * Matsubara et al., (2022) Matsubara, T., Oates, C. J., and Briol, F.-X. (2022).  The ridgelet prior: A covariance function approach to prior specification for bayesian neural networks.
  * Milios et al., (2018) Milios, D., Camoriano, R., Michiardi, P., Rosasco, L., and Filippone, M. (2018).  Dirichlet-based gaussian processes for large-scale calibrated classification.  In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.
  * Nalisnick et al., (2020) Nalisnick, E., Gordon, J., and Hernández-Lobato, J. M. (2020).  Predictive complexity priors.
  * Osawa et al., (2019) Osawa, K., Swaroop, S., Jain, A., Eschenhagen, R., Turner, R. E., Yokota, R., and Khan, M. E. (2019).  Practical deep learning with bayesian principles.
  * Pearce et al., (2020) Pearce, T., Tsuchida, R., Zaki, M., Brintrup, A., and Neely, A. (2020).  Expressive priors in bayesian neural networks: Kernel combinations and periodic functions.  In Adams, R. P. and Gogate, V., editors, Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pages 134–144. PMLR.
  * Pedregosa et al., (2011) Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M., and Duchesnay, E. (2011).  Scikit-learn: Machine learning in Python.  Journal of Machine Learning Research, 12:2825–2830.
  * Pinder and Dodd, (2022) Pinder, T. and Dodd, D. (2022).  Gpjax: A gaussian process framework in jax.  Journal of Open Source Software, 7(75):4455.
  * Quang, (2019) Quang, M. H. (2019).  Regularized divergences between covariance operators and gaussian measures on hilbert spaces.
  * Quang, (2022) Quang, M. H. (2022).  Kullback-leibler and renyi divergences in reproducing kernel hilbert space and gaussian process settings.
  * (40) Rudner, T. G. J., Bickford Smith, F., Feng, Q., Teh, Y. W., and Gal, Y. (2022a).  Continual learning via sequential function-space variational inference.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S., editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 18871–18887. PMLR.
  * (41) Rudner, T. G. J., Chen, Z., Teh, Y. W., and Gal, Y. (2022b).  Tractabe Function-Space Variational Inference in Bayesian Neural Networks.  In Advances in Neural Information Processing Systems.
  * Rudner et al., (2023) Rudner, T. G. J., Kapoor, S., Qiu, S., and Wilson, A. G. (2023).  Function-space regularization in neural networks: A probabilistic perspective.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 29275–29290. PMLR.
  * Santin and Schaback, (2016) Santin, G. and Schaback, R. (2016).  Approximation of eigenfunctions in kernel-based spaces.  Advances in Computational Mathematics, 42(4):973–993.
  * Sensoy et al., (2018) Sensoy, M., Kaplan, L., and Kandemir, M. (2018).  Evidential deep learning to quantify classification uncertainty.
  * Shalashilin, (2024) Shalashilin, I. (2024).  Gaussian processes for vector fields and ocean current modelling.
  * Simpson, (2022) Simpson, D. (2022).  Priors for the parameters in a Gaussian process.
  * Sun et al., (2019) Sun, S., Zhang, G., Shi, J., and Grosse, R. (2019).  FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS.  In International Conference on Learning Representations.
  * Tomczak et al., (2020) Tomczak, M., Swaroop, S., and Turner, R. (2020).  Efficient low rank gaussian variational inference for neural networks.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 4610–4622. Curran Associates, Inc.
  * Tran et al., (2022) Tran, B.-H., Rossi, S., Milios, D., and Filippone, M. (2022).  All you need is a good functional prior for bayesian deep learning.  Journal of Machine Learning Research, 23(74):1–56.
  * Wild et al., (2022) Wild, V. D., Hu, R., and Sejdinovic, D. (2022).  Generalized variational inference in function spaces: Gaussian measures meet bayesian deep learning.
  * Williams and Rasmussen, (2006) Williams, C. K. and Rasmussen, C. E. (2006).  Gaussian processes for machine learning, volume 2.  MIT press Cambridge, MA.
  * Wilson et al., (2022) Wilson, A. G., Izmailov, P., Hoffman, M. D., Gal, Y., Li, Y., Pradier, M. F., Vikram, S., Foong, A., Lotfi, S., and Farquhar, S. (2022).  Evaluating approximate inference in bayesian deep learning.  In Kiela, D., Ciccone, M., and Caputo, B., editors, Proceedings of the NeurIPS 2021 Competitions and Demonstrations Track, volume 176 of Proceedings of Machine Learning Research, pages 113–124. PMLR.
  * Wong, (2023) Wong, Y. K. (2023).  Machine learning in portfolio management.  PhD thesis, University of Sydney.
  * Wu et al., (2023) Wu, M., Xuan, J., and Lu, J. (2023).  Indirect functional bayesian neural networks.  In Fifth Symposium on Advances in Approximate Bayesian Inference.
  * Xiao et al., (2017) Xiao, H., Rasul, K., and Vollgraf, R. (2017).  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.

##  Appendix A Appendix

###  A.1 Divergences between Gaussian measures

####  A.1.1 The KL divergence is infinite

In this section, we show that the Kullbach-Liebler (KL) divergence between the
Gaussian measures
$\mathbb{Q}_{\phi}^{F}\sim\mathcal{N}\mathopen{}\left(m_{Q},C_{Q}\right)\mathclose{}$
and
$\mathbb{P}^{F}\sim\mathcal{N}\mathopen{}\left(m_{P},C_{P}\right)\mathclose{}$,
respectively induced by the linearized BNN in Eq 2.7 and by a non-degenerate
Gaussian process satisfying conditions given in section 2.2, is infinite.
While this has already been shown by Burt et al., (2020), the proof is easier
for Gaussian measures. We first need the Feldman-Hàjek theorem which tells us
when the KL divergence between two Gaussian measures is well-defined.

######  Theorem A.1 (Feldman-Hàjek, Quang, (2022) Theorem 2, Simpson, (2022)
Theorem 7).

Consider two Gaussian measures
$\nu_{1}=\mathcal{N}\mathopen{}\left(m_{1},C_{1}\right)\mathclose{}$ and
$\nu_{2}=\mathcal{N}\mathopen{}\left(m_{2},C_{2}\right)\mathclose{}$ on
$L^{2}(\mathcal{X},\rho)$. Then $\nu_{1}$ and $\nu_{2}$ are called equivalent
if and only if the following holds:

  1. 1.

$m_{1}-m_{2}\in\text{Im}(C_{2}^{1/2})$

  2. 2.

The operator $T$ such that $C_{1}=C_{2}^{1/2}(I-T)C_{2}^{1/2}$ is Hilbert-
Schmidt, that is $T$ has a countable set of eigenvalues $\lambda_{i}$ that
satisfy $\lambda_{i}<1$ and $\sum_{i=1}^{\infty}\lambda_{i}^{2}<\infty$.

otherwise $\nu_{1}$ and $\nu_{2}$ are singular. If $\nu_{1}$ and $\nu_{2}$ are
equivalent, then the Radon-Nikodym dervative exists and
$D_{\mathrm{KL}}\mathopen{}\left(\nu_{1}\,\|\,\nu_{2}\right)\mathclose{}$
admits an explicit formula. Otherwise,
$D_{\mathrm{KL}}\mathopen{}\left(\nu_{1}\,\|\,\nu_{2}\right)\mathclose{}=\infty$.

Let us now show that the KL divergence between $\mathbb{Q}_{\phi}^{F}$ and
$\mathbb{P}^{F}$ is indeed infinite.

######  Proposition 1.

The Gaussian measures $\mathbb{Q}_{\phi}^{F}$ and $\mathbb{P}^{F}$ are
mutually singular and
$D_{\text{KL}}(\mathbb{Q}_{\phi}^{F}||\mathbb{P}^{F})=\infty$.

###### Proof.

The proof follows from the Feldman-Hàjek theorem (theorem A.1). In our case,
$C_{Q}$ has at most $p$ non-zero eigenvalues as the covariance function of the
GP induced by the BNN is degenerate, while $C_{P}$ has a set of (countably)
infinite non-zeros eigenvalues (prior is non-degenerate as per assumption).
Hence, for the equality in condition (2) to hold, $T$ must have eigenvalue $1$
which violates the requirement that $T$ is Hilbert-Schmidt i.e. that its
eigenvalues $\left\\{\lambda_{i}\right\\}_{i=1}^{\infty}$ satisfy
$\lambda_{i}<1$ and $\sum_{i=1}^{\infty}\lambda_{i}^{2}<\infty$. Therefore,
$\mathbb{Q}_{\phi}^{F}$ and $\mathbb{P}^{F}$ are mutually singular and
$D_{\text{KL}}(\mathbb{Q}_{\phi}^{F}||\mathbb{P}^{F})=\infty$. ∎

####  A.1.2 The regularized KL divergence

In this section, we provide the bound describing the asymptotic convergence of
the regularized KL divergence estimator.

######  Theorem A.2 (Convergence of estimator, Quang, (2022) Theorem 45).

Assume the following:

  1. 1.

Let T be a $\sigma-\text{compact}$ metric space, that is
$T=\cup_{i=1}^{\infty}T_{i}$, where $T_{1}\subset T_{2}\subset\cdots$ with
each $T_{i}$ being compact.

  2. 2.

$\rho$ is a non-degenerate Borel probability measure on T, that is
$\rho\mathopen{}\left(B\right)\mathclose{}>0$ for each open set $B\subset T$.

  3. 3.

$K_{1},K_{2}:T\times T\to$ are continuous, symmetric, positive definite
kernels and there exists $\kappa_{1}>0,\kappa_{2}>0$ such that
$\int_{T}K_{i}(x,x)d\rho\mathopen{}\left(x\right)\mathclose{}\leq\kappa_{i}^{2}\;$
for $i=1,2$.

  4. 4.

$\operatorname{sup}_{x\in T}K_{i}(x,x)\leq\kappa_{i}^{2}\;$ for $i=1,2$.

  5. 5.

$f_{i}\sim GP\mathopen{}\left(\mu_{i},K_{i}\right)\mathclose{}$, where
$\mu_{i}\in L^{2}(T,\rho)$ for $i=1,2$.

  6. 6.

$\exists B_{i}>0$ such that $\left\lVert\mu_{i}\right\rVert_{\infty}\leq
B_{i}$ for $i=1,2$.

Let ${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}$,
${\bm{x}}^{(1)},\dots,{\bm{x}}^{(M)}\mathrel{\mathop{\sim}\limits^{\vbox
to1.20552pt{\kern-2.0pt\hbox{$\scriptstyle\textup{\tiny
i.i.d}$}\vss}}}\rho({\bm{x}})$. If Gaussian measures
$\mathcal{N}\mathopen{}\left(m_{i},C_{i}\right)\mathclose{}$ are induced by
GPs $f_{i}\sim\mathcal{GP}(\mu_{i},K_{i})$ for $i=1,2$, then for any
$0<\delta<1$, with probability at least $1-\delta$,

$|D_{\mathrm{KL}}\mathopen{}\left(\mathcal{N}\mathopen{}\left(\mu_{1}(\mathbf{x}),K_{1}(\mathbf{x},\mathbf{x})+M\gamma\mathbb{I}_{M}\right)\mathclose{}\,\|\,{\mathcal{N}\mathopen{}\left(\mu_{2}(\mathbf{x}),K_{2}(\mathbf{x},\mathbf{x})+M\gamma\mathbb{I}_{M}\right)\mathclose{}}\right)\mathclose{}\\\
-D^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\mathcal{N}\mathopen{}\left(m_{1},C_{1}\right)\mathclose{}\,\|\,\mathcal{N}\mathopen{}\left(m_{2},C_{2}\right)\mathclose{}\right)\mathclose{}|\\\
\leq\frac{1}{2\gamma}(B_{1}+B_{2})^{2}[1+\kappa_{2}^{2}/\gamma]\mathopen{}\left(\frac{2\log\frac{48}{\gamma}}{M}+\sqrt{\frac{2\log\frac{48}{\gamma}}{M}}\right)\mathclose{}\\\
+\frac{1}{2\gamma^{2}}[\kappa_{1}^{4}+\kappa_{2}^{4}+\kappa_{1}^{2}\kappa_{2}^{2}(2+\kappa_{2}^{2}/\gamma)]\mathopen{}\left(\frac{2\log\frac{12}{\gamma}}{M}+\sqrt{\frac{2\log\frac{12}{\gamma}}{M}}\right)\mathclose{}$
(A.1)

###  A.2 Additional details on the GFSVI objective estimator

In this section, we present details on the estimation of the generalized
function-space variational inference (GFSVI) objective. Let
$f_{L}({\,\cdot\,;{\bm{w}}})$ be the linearized BNN (Eq 2.6) with weights
${\bm{w}}\in^{p}$, and $\mathcal{D}=\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{N}$ a
data set with features ${\bm{x}}_{i}\in\mathcal{X}\subset^{d}$ and associated
values $y_{i}\in\mathcal{Y}$. Assuming a likelihood
$p\mathopen{}\left(\mathcal{D}\,|\,{\bm{w}}\right)\mathclose{}=\prod_{i=1}^{N}p\mathopen{}\left(y_{i}\,|\,f({\bm{x}}_{i};{\bm{w}})\right)\mathclose{}$
and a Gaussian variational distribution on model weights
$q_{\phi}({\bm{w}})=\mathcal{N}\mathopen{}\left({\bm{w}}\,|\,{\bm{m}},{\bm{S}}\right)\mathclose{}$,
the GFSVI objective function is

$\mathcal{L}(\phi)=\sum_{i=1}^{N}{\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\log
p\mathopen{}\left(y_{i}\,|\,f_{L}({\bm{x}}_{i};{\bm{w}})\right)\mathclose{}\right]\mathclose{}-D^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}^{F}\,\|\,\mathbb{P}^{F}\right)\mathclose{}$
(A.2)

where $\mathbb{Q}_{\phi}^{F}$ and $\mathbb{P}^{F}$ are the Gaussian measures
induced by the linearized BNN and a Gaussian process prior respectively.

##### Expected log-likelihood

When considering a Gaussian likelihood, we use the closed form expression
available due to the Gaussian variational measure over functions induced by
the linearized BNN

${\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\log\mathcal{N}\mathopen{}\left(y_{i}\,|\,f_{L}({\bm{x}}_{i};{\bm{w}}),\sigma_{y}^{2}\right)\mathclose{}\right]\mathclose{}=-\frac{1}{2}\log\mathopen{}\left(2\pi\sigma_{y}^{2}\right)\mathclose{}-\frac{\mathopen{}\left(y_{i}-f\mathopen{}\left({\bm{x}}_{i};{\bm{m}}\right)\mathclose{}\right)\mathclose{}^{2}+J\mathopen{}\left({\bm{x}}_{i};{\bm{m}}\right)\mathclose{}{\bm{S}}J\mathopen{}\left({\bm{x}}_{i};{\bm{m}}\right)\mathclose{}^{\top}}{2\sigma_{y}^{2}}.$
(A.3)

When considering a Categorical likelihood with $C$ different classes, we
estimate the expected log-likelihood term using Monte-Carlo integration as

$\begin{split}{\mathop{{}\mathbb{E}}}_{q_{\phi}({\bm{w}})}\mathopen{}\left[\log\text{Cat}(y_{i}\,|\,\sigma(f_{L}({\bm{x}}_{i};{\bm{w}})))\right]\mathclose{}&=\frac{1}{K}\sum_{k=1}^{K}\sum_{c=1}^{C}\mathbb{I}[y_{i}=c]\Bigl{[}f^{c}_{L}({\bm{x}}_{i};{\bm{w}}^{(k)})\\\
&\qquad-\log\left[\sum_{c^{\prime}=1}^{C}\exp\left(f^{c^{\prime}}_{L}({\bm{x}}_{i};{\bm{w}}^{(k)})\right)\right]\Bigr{]}\end{split}$
(A.4)

where ${\bm{w}}^{(k)}\sim q_{\phi}({\bm{w}})$ for $k=1,\dots,K$,
$\mathbb{I}[\cdot]$ is the indicator function, $\sigma(\cdot)$ is the softmax
function and $f^{c}_{L}(\,\cdot\,;{\bm{w}})$ is the logit for class $c$
obtained from $f_{L}$.

##### Regularized KL divergence

We estimate the regularized KL divergence using its consistent estimator (Eq.
2.9)

$\hat{D}^{\gamma}_{\mathrm{KL}}\mathopen{}\left(\mathbb{Q}_{\phi}^{F}\,\|\,\mathbb{P}^{F}\right)\mathclose{}=\frac{1}{2}\mathopen{}\left(f({\bm{x}};{\bm{m}})-\mu({\bm{x}})\right)\mathclose{}^{\top}(K({\bm{x}},{\bm{x}})+\gamma
M\mathbb{I}_{M})^{-1}\mathopen{}\left(f({\bm{x}};{\bm{m}})-\mu({\bm{x}})\right)\mathclose{}\\\
+\frac{1}{2}\operatorname{Tr}\big{[}(K({\bm{x}},{\bm{x}})+\gamma
M\mathbb{I}_{M})^{-1}(J({\bm{x}};{\bm{m}})SJ({\bm{x}};{\bm{m}})^{\top}+\gamma
M\mathbb{I}_{M})-\mathbb{I}_{M}\big{]}\\\
-\frac{1}{2}\log\det\big{[}(K({\bm{x}},{\bm{x}})+\gamma
M\mathbb{I}_{M})^{-1}(J({\bm{x}};{\bm{m}})SJ({\bm{x}};{\bm{m}})^{\top}+\gamma
M\mathbb{I}_{M})\big{]}$ (A.5)

with measurement points ${\bm{x}}=\\{{\bm{x}}^{(i)}\\}_{i=1}^{M}$,
${\bm{x}}^{(1)},\dots,{\bm{x}}^{(M)}\mathrel{\mathop{\sim}\limits^{\vbox
to1.20552pt{\kern-2.0pt\hbox{$\scriptstyle\textup{\tiny
i.i.d}$}\vss}}}\rho({\bm{x}})$ sampled from a probability measure on
$\mathcal{X}$.

###  A.3 Additional details on the experimental setup

####  A.3.1 Experiments on synthetic data

##### Regression

We consider the following generative model for the toy data

$y_{i}=\operatorname{sin}(2\pi
x_{i})+\epsilon\quad\text{with}\quad\epsilon\sim\mathcal{N}\mathopen{}\left(0,\sigma_{n}^{2}\right)\mathclose{}$
(A.6)

and draw $x_{i}\sim\mathcal{U}([-1,-0.5]\cup[0.5,1])$. When not otherwise
specified, we use $\sigma_{n}=0.1$. On the plots, the data points are shown as
gray circles, inferred mean functions as red lines, their 2-standard-
deviations interval around the mean in light green, and functions sampled from
the approximate posterior as green lines. In general, we consider two hidden-
layer BNNs with 30 neurons per layer and hyperbolic tangent activation (Tanh)
functions. Specifically in figure 11, the small BNN has the same architecture
as above while the large BNN has 100 neurons per layer. All the BNN baselines
have the same architecture and fully-factorized Gaussian approximate
posterior. The prior scale of TFSVI (Rudner et al., 2022b, ) is set to
$\sigma_{p}=0.2$ and $\sigma_{p}=0.75$ for MFVI (Blundell et al.,, 2015) and
Laplace (Immer et al.,, 2021). For the Gaussian process posterior baseline, we
fit the prior parameters by maximizing the log-marginal likelihood (Williams
and Rasmussen,, 2006). Apart from the cases where the parameters of the GP
prior used for GFSVI (our method) and FVI (Sun et al.,, 2019) are explicitly
stated, we consider a constant zero-mean function and find the parameters of
the covariance function by maximizing the log-marginal likelihood from mini-
batches (Chen et al.,, 2021). Except where otherwise stated, we estimate the
functional KL divergences with 500 measurement points and use the regularized
KL divergence with $\gamma=10^{-10}$.

##### Classification

We sample 100 data points perturbed by Gaussian noise with $\sigma_{n}=0.1$
from the two moons data (Pedregosa et al.,, 2011). On the plots, the data
points are shown as red (class 0) and blue (class 1) dots. We plot the mean
and 2-standard-deviations of the probability that ${\bm{x}}$ belongs to class
1 with respect to the posterior (i.e. $p(y=1\,|\,{\bm{w}}^{(k)},{\bm{x}})$)
which we estimate from $K=100$ samples ${\bm{w}}^{(k)}\sim q_{\phi}({\bm{w}})$
for $k=1,\dots,K$. We consider two hidden-layer BNNs with 100 neurons per
layer and hyperbolic tangent activation (Tanh) functions. All the BNN
baselines have the same architecture and fully-factorized Gaussian approximate
posterior. The prior scale of MFVI (Blundell et al.,, 2015) is set to
$\sigma_{p}=0.8$ and $\sigma_{p}=1.0$ for TFSVI (Rudner et al., 2022b, ) and
Laplace (Immer et al.,, 2021). For the Gaussian process posterior baseline, we
approximate the intractable posterior using the Laplace approximation and find
the prior parameters by maximizing the log-marginal likelihood (Williams and
Rasmussen,, 2006). The GP prior for GFSVI (our method) and FVI (Sun et al.,,
2019) has a constant zero-mean function and we find the parameters of the
covariance function by maximizing the log-marginal likelihood from mini-
batches (Chen et al.,, 2021) using the method to transform classifications
labels into regression targets from Milios et al., (2018). We estimate the
functional KL divergences with 500 measurement points and use the regularized
KL divergence with $\gamma=10^{-10}$.

####  A.3.2 Ocean current modeling experiment

We use $20$ two-dimensional velocity vectors from the GulfDrifters dataset
(Lilly and Pérez-Brunius,, 2021) and we follow the setup by Shalashilin,
(2024). We apply the Helmholtz decomposition to the neural network $f$ using
the following parameterization

$f(\cdot,{\bm{w}})=\operatorname{grad}\Phi(\cdot,{\bm{w}}_{1})+\operatorname{rot}\Psi(\cdot,{\bm{w}}_{2})$
(A.7)

where ${\bm{w}}=\left\\{{\bm{w}}_{1},{\bm{w}}_{2}\right\\}$ and,
$\Phi(\cdot,{\bm{w}}_{1})$ and $\Psi(\cdot,{\bm{w}}_{2})$ are two 2-layered
fully-connected neural networks with hyperbolic tangent activation functions
and $50$ hidden units per layer. GFSVI and TFSVI both use $160$ fixed context
points. The prior scale of TFSVI is set to $\sigma_{p}=0.5$. We fit the neural
networks on the entire dataset and average the scores with respect to five
different random seeds.

####  A.3.3 Regression experiments with tabular data

##### Datasets and pre-processing

We evaluate the predictive performance of our model on regression datasets
from the UCI repository (Dua and Graff,, 2017) described in table 6. These
datasets are also considered in Sun et al., (2019); Wild et al., (2022) but we
include two additional larger ones (Wave and Denmark). We perform 5-fold cross
validation, leave out one fold for testing, consider 10% of the remaining 4
folds as validation data and the rest as training data. We report mean and
standard-deviation of the average expected log-likelihood and average mean
square error on the test fold. We also report the mean rank of the methods
across all datasets by assigning rank 1 to the best scoring method as well as
any method who’s error bars overlap with the highest score’s error bars, and
recursively apply this procedure to the methods not having yet been assigned a
rank. The expected log-likelihood is estimated by Monte Carlo integration when
it is not available in closed form (MFVI, TFSVI and FVI) with 100 posterior
samples. We preprocess the dataset by encoding categorical features as one-hot
vectors and standardizing the features and labels.

##### Baseline specification

We compare our GFSVI method to two weight-space inference methods (mean-field
variational inference (Blundell et al.,, 2015) and linearized Laplace (Immer
et al.,, 2021)) and two function-space inference methods (FVI (Sun et al.,,
2019) and TFSVI (Rudner et al., 2022b, )). While FVI uses GP priors, TFSVI
performs inference in function space but with the pushforward to function
space of the variational distribution and prior on the weights. We compute the
function-space (regularized) KL divergence using a set of 500 measurement
points sampled from a uniform distribution for GFSVI and TFSVI, and 50 points
drawn from a uniform distribution along with 450 samples from the training
batch for FVI as specified in Sun et al., (2019). All the BNN baselines have
the same architecture and fully-factorized Gaussian approximate posterior. We
also provide results with a GP (Williams and Rasmussen,, 2006) when the size
of the dataset allows it, and a sparse GP (Hensman et al.,, 2013). As we
restrict our comparison to BNNs, we do not consider the GP and sparse GP as
baselines but rather as gold-standards. All models have a Gaussian
heteroskedastic noise model with a learned scale parameter. All the BNNs are
fit using the Adam optimizer (Kingma and Ba,, 2017) using a mini-batch size of
2000 samples. We also perform early stopping when the validation loss stops
decreasing.

##### Model selection

Hyper-parameter optimization is conducted using the Bayesian optimization tool
provided by Wandb (Biewald,, 2020). BNN parameters are selected to maximize
the average validation expected log-likelihood across the 5 cross-validation
folds. We optimize over prior parameters (kernel and prior scale), learning-
rate and activation function. We select priors for GFSVI, FVI, sparse GP and
GP among the RBF, Matern-1/2, Matern-3/2, Matern-5/2, Linear and Rational
Quadratic covariance functions. The GP prior parameters used with GFSVI and
FVI are selected by maximizing the log-marginal likelihood from batches as
proposed by Chen et al., (2021) and done in Sun et al., (2019). Hyper-
parameters for GPs and sparse GPs (kernel parameters and learning-rate) are
selected to maximize the mean log-marginal likelihood of the validation data
across the 5 cross-validation folds.

Table 6: UCI regression dataset description

Dataset | Boston | Naval | Power | Protein | Yacht | Concrete | Energy | Kin8nm | Wine | Wave | Denmark
---|---|---|---|---|---|---|---|---|---|---|---
Number samples | 506 | 11 934 | 9 568 | 45 730 | 308 | 1 030 | 768 | 8 192 | 1 599 | 288 000 | 434 874
Number features | 13 | 16 | 4 | 9 | 6 | 8 | 8 | 8 | 11 | 49 | 2

####  A.3.4 Classification experiments with image data

##### Datasets and pre-processing

We further evaluate the predictive performance of our model on classification
tasks with the MNIST (LeCun et al.,, 2010) and Fashion MNIST (Xiao et al.,,
2017) image data sets. We fit the models on a random subset of 90% of the
provided training split, consider the remaining 10% as validation data and
evaluate on the provided test split. We repeat this procedure 5 times with
different random seeds and report the mean and standard-deviation of the
average expected log-likelihood, accuracy and expected calibration error (ECE)
of the mean of the predictive distribution on the test set. The expected log-
likelihood is estimated by Monte Carlo integration with 100 posterior samples
when it is not available in closed form (MFVI, TFSVI and FVI). We estimate the
mean of the predictive distribution to compute the accuracy and the ECE with
100 posterior samples. We preprocess the dataset by standardizing the images.

##### Baseline specification

We compare our GFSVI method to the same baselines as for the regression
experiments (see A.3.3). All the BNN baselines have the same architecture and
fully-factorized Gaussian approximate posterior. More specifically, we
consider a CNN with three convolutional layers (with output channels 16, 32
and 64) before two fully connected layers (with output size 128 and 10). The
convolutional layers use $3\times 3$ shaped kernels. Each pair of
convolutional layers is interleaved with a max-pooling layer. We consider
three different measurement point distributions $\rho$ to estimate the
(regularized) KL divergence in GFSVI, FVI and TFSVI: random, random pixel and
kmnist. The random measurement point distribution is sampled from by drawing
50% of the samples from the training data batch and 50% of the samples from a
uniform distribution over $[p_{min},p_{max}]^{H\times W\times C}$, where $H$,
$W$ and $C$ are respectively the height, width and number of channels of the
images, and $p_{min}=v_{min}-0.5\times\Delta$ and
$p_{max}=v_{max}+0.5\times\Delta$ where $\Delta=v_{max}-v_{min}$ is the
difference between the minimal ($v_{min}$) and maximal ($v_{max}$) pixel
values of the data set. The random pixel measurement point distribution is
taken from Rudner et al., 2022b  and is sampled from by randomly choosing each
pixel value among the ones available from the training data batch at the same
position in the $28\times 28$ pixel grid. Finally, the kmnist measurement
point distribution is also taken from Rudner et al., 2022b  and is drawn from
by randomly sampling data points from the Kuzushiji-MNIST (KMNIST) dataset
(Clanuwat et al.,, 2018). The KMNIST dataset is a collection of 70’000 gray-
scale images of size $28\times 28$ which we preprocess by standardizing the
images. We sample 10 measurement points when using random, 25 measurement
points when using random pixel and 20 when using kmnist. All the BNNs are
trained using the Adam optimizer (Kingma and Ba,, 2017) using a mini-batch
size of 100. We also perform early stopping when the validation loss stops
decreasing.

##### Model selection

Hyper-parameter optimization is conducted just like for the regression tasks
(see A.3.3). The Gaussian process prior parameters used with GFSVI and FVI are
selected by maximizing the log-marginal likelihood from batches (Chen et al.,,
2021) using the method to transform classifications labels into regression
targets from Milios et al., (2018). We optimize the same hyper-parameters as
for the regression experiments with the exception of the additional
$\alpha_{\epsilon}$ parameter introduced by Milios et al., (2018) for the
function-space VI methods with GP priors (FVI and GFSVI).

####  A.3.5 OOD detection

##### Tabular data with a Gaussian likelihood

Following the setup from Malinin et al., (2020) we take epistemic uncertainty
to be the variance of the mean prediction with respect to samples from the
posterior. We consider the test data to be in-distribution (ID) data and a
subset of the song dataset (Bertin-Mahieux et al.,, 2011) of equal length and
with an equal number of features as out-of-distribution (OOD) data. We use the
same preprocessing as for regression as well as the same baselines with the
same hyper-parameters (see section A.3.3). We first fit a model, then evaluate
the extend by which the epistemic uncertainty under the model is predictive of
the ID and OOD data using a single threshold obtained by a depth-1 decision
tree fit to minimize the classification loss. We report the mean and standard
error of the accuracy of the threshold to classify OOD from ID data based on
epistemic uncertainty across the 5 folds of cross-validation. We also provide
results obtained using a GP and sparse GP as gold standard.

##### Image data with a Categorical likelihood

Following the setup by Osawa et al., (2019), we take the epistemic uncertainty
to be the entropy of the mean of the predictive distribution with respect to
samples from the posterior. We evaluate models trained on MNIST using MNIST’s
test split as ID data and a subset of the training set of Fashion MNIST as OOD
data. Likewise, we evaluate models trained on Fashion MNIST using Fashion
MNIST’s test split as ID data and a subset of the training set of MNIST as OOD
data. We use the same preprocessing as for classification, as well as the same
baselines with the same hyper-parameters (see section A.3.4). We first fit a
model, then evaluate the extend by which the epistemic uncertainty under the
model is predictive of the ID and OOD data using a single threshold obtained
by a depth-1 decision tree fit to minimize the classification loss. We
estimate mean of the predictive distribution by Monte-Carlo integration using
100 posterior samples. We report the mean and standard error of the accuracy
of the threshold to classify OOD from ID data based on epistemic uncertainty
for the 5 models trained on different random seeds (see section A.3.4).

####  A.3.6 Variational measure evaluation

We evaluate our inference method by comparing the samples drawn from the exact
posterior over functions with the approximate posterior obtained with our
method (GFSVI). We follow the setup by Wilson et al., (2022) and we compute
the average Wasserstein-2 metric between 1000 samples drawn from a GP
posterior with a RBF kernel evaluated at the test points, and samples from the
approximate posterior of GFSVI, sparse GP and FVI evaluated at the same points
and with the same prior. We consider the Boston, Concrete, Energy, Wine and
Yacht datasets for which the exact GP posterior can be computed and use the
same preprocessing as for regression (see section A.3.3). We report the mean
and standard error of the average Wasserstein-2 metric across the 5 folds of
cross-validation. The Wasserstein-2 metric is computed using the Python
Optimal Transport library (Flamary et al.,, 2021).

##### Baseline specification

FVI and GFSVI have the same two hidden layer neural network architecture with
100 neurons each and hyperbolic tangent activation. These models are fit with
the same learning rate and set of 500 measurement points jointly sampled from
a uniform distribution over the feature-space and mini-batch of training
samples. We use $\gamma=10^{-15}$ for the regularized KL divergence. We
further consider a sparse GP with 100 inducing points.

####  A.3.7 Software

We use the JAX (Bradbury et al.,, 2018) and DM-Haiku (Hennigan et al.,, 2020)
Python libraries to implement our Bayesian neural networks. MFVI, linearized
Laplace and TFSVI were implemented based on the information in the papers, and
code for FVI was adapted to the JAX library from the implementation provided
by the authors. We further use the GPJAX Python library for experiments
involving Gaussian processes (Pinder and Dodd,, 2022).

####  A.3.8 Hardware

All models were fit using a single NVIDIA RTX 2080Ti GPU with 11GB of memory.

###  A.4 Additional experimental results

In this section, we present additional figures for our qualitative uncertainty
evaluation as well as further experimental results on regression, out-of-
distribution detection and robustness under input distribution shift tasks. We
also provide plots illustrating the eigenvalue decay of different kernels, and
figures showing the influence of $\gamma$ in the regularized KL divergence.

####  A.4.1 Qualitative uncertainty evaluation

Figure 4: Our method (GFSVI) effectively regularizes functions generated by
the Bayesian neural network (BNN) both in settings where the generative
process is very noisy ($\sigma_{n}=1$) or not ($\sigma_{n}=0.1$).

##### Regression

We further find that our method (GFSVI) provides strong regularization when
the data generative process is noisy (see figure 4) and is more robust than
FVI to situations where ones computational budget constrains the number of
measurement points $M$ to be small (figure 8). In contrast to FVI, GFSVI
accurately approximates the exact GP posterior under rough (Matern-1/2) GP
priors effectively incorporating prior knowledge defined by the GP prior to
the inference process (see figure 5). Likewise, GFSVI adapts to the
variability of the functions specified by the kernel (see figure 6). We also
find that GFSVI requires a larger number of measurement points to capture the
behavior of a rougher prior (see figure 7).

Figure 5: Our method (GFSVI) with a Matern-1/2 Gaussian process (GP) prior
accurately approximates the exact GP posterior unlike the function-space prior
baseline (FVI). Weight-space prior baselines do not provide a straight-forward
mechanism to incorporate prior assumptions regarding the functions generated
by BNNs and underestimate the epistemic uncertainty (MFVI, Laplace). The lower
row is identical to the one in figure 2 in the main text and is reproduced
here to make comparison easier.

Figure 6: Our method (GFSVI) allows to incorporate prior beliefs in terms of
function variability using the characteristic length-scale parameter of the
Gaussian process (GP) prior. GFSVI was fit using a GP prior with RBF
covariance function.

Figure 7: Our method (GFSVI) captures the smooth behavior of a Gaussian
process (GP) prior with RBF covariance function even if the number of
measurement points is small (M=10). However, in that setting GFSVI fails to
reproduce the rough effect of a GP prior with a Matern-1/2 covariance
function, and requires a larger amount of measurement points to do so (M=100).

Figure 8: Our method (GFSVI) already provides a reasonable approximation to
the exact posterior with small numbers of measurement points (M=10) while
function-space baseline FVI requires many more (M=100).

##### Classification

We find that GFSVI better captures the beliefs induced by the smooth RBF and
rough Matern-1/2 Gaussian process priors compared to FVI (see Figures 9 and
10). Moreover, GFVSI both accurately fits the training data and shows greater
uncertainty outside of its support relative to BNNs baselines with weight-
space and function-space priors. Unlike for the toy data regression
experiments where the GP posterior was the ground truth, the Laplace
(approximate) GP posterior in Figures 9 and 10 only represents a possible
approximation to the now in-tractable posterior (due to the softmax inverse
link function). Thus the GP should not be considered as the ground truth nor
as the optimal approximation in the classification setting, but is
nevertheless useful to give a idea of the level of uncertainty a BNN with a GP
prior should provide outside of the support of the data.

Figure 9: Our method (GFSVI) with a RBF Gaussian process (GP) prior accurately
captures the smooth decision boundary induced by the prior and shows high
uncertainty outside of the data support. Weight-space baselines do not provide
a straight-forward mechanism to incorporate prior assumptions regarding the
functions generated by BNNs and underestimate the epistemic uncertainty
(TFSVI, Laplace) or underfit the data (MFVI).

Figure 10: Our method (GFSVI) with a Matern-1/2 Gaussian process (GP) prior
accurately captures the rough decision boundary unlike the function-space
baseline (FVI). Weight-space baselines do not provide a straight-forward
mechanism to incorporate prior assumptions regarding the functions generated
by BNNs and underestimate the epistemic uncertainty (TFSVI, Laplace) or
underfit the data (MFVI).

##### Inductive biases

figure 11 compares GFSVI to the exact posterior across two different priors
and three model architectures (details in A.3.1). We find that the BNN’s
ability to incorporate the beliefs introduced by the GP prior depends on its
size and activation function. When using piece-wise linear activations (ReLU),
small models are prone to underfitting for smooth priors (RBF), and to
collapsing uncertainty for rough priors (Matern-1/2). By contrast, when using
smooth activations (Tanh), smaller models suffice, and they are compatible
with most standard GP priors (the results shown in figure 11 extend to RBF,
Matern family, and Rational Quadratic in our experiments). We also analyzed
how the number $M$ of measurement points affects performance. figures 7 and 14
show that capturing the properties of rough GP priors and estimating
$D_{\text{KL}}^{\gamma}$ with these priors requires larger $M$.

Figure 11:  Our method requires that the Bayesian neural network (BNN) and
Gaussian process (GP) prior share similar inductive biases to provide an
accurate approximation to the exact posterior.

####  A.4.2 Regression on tabular data

Table 7: Test mean square error (MSE) of evaluated methods on regression
datasets. We find that GFSVI (ours) also performs competitively in terms of
MSE compared to baselines and obtains the best mean rank, matching best the
performing methods on nearly all datasets.

Dataset | Function-space priors | Weight-space priors | Gaussian Processes (Gold Standards)
---|---|---|---
| GFSVI (ours) | FVI | TFSVI | MFVI | VIP | Laplace | GWI | Sparse GP | GP
Boston |  0.123 $\pm$ 0.021 | 0.136 $\pm$ 0.022 | 0.995 $\pm$ 0.092 | 0.532 $\pm$ 0.072 |  0.201 $\pm$ 0.056 | 0.203 $\pm$ 0.047 |  0.273 $\pm$ 0.069 | 0.122 $\pm$ 0.014 | 0.115 $\pm$ 0.020
Concrete | 0.114 $\pm$ 0.008 | 0.116 $\pm$ 0.004 |  0.389 $\pm$ 0.015 |  0.698 $\pm$ 0.046 | 0.109 $\pm$ 0.008 | 0.116 $\pm$ 0.007 |  0.145 $\pm$ 0.017 | 0.399 $\pm$ 0.020 | 0.116 $\pm$ 0.007
Energy | 0.003 $\pm$ 0.000 | 0.003 $\pm$ 0.000 | 0.003 $\pm$ 0.000 |  0.152 $\pm$ 0.024 | 0.043 $\pm$ 0.036 |  0.002 $\pm$ 0.000 | 0.003 $\pm$ 0.001 | 0.087 $\pm$ 0.005 | 0.087 $\pm$ 0.004
Kin8nm |  0.071 $\pm$ 0.001 | 0.075 $\pm$ 0.003 | 0.073 $\pm$ 0.001 | 0.290 $\pm$ 0.111 |  0.068 $\pm$ 0.002 | 0.083 $\pm$ 0.001 |  0.071 $\pm$ 0.001 | 0.088 $\pm$ 0.002 | (infeasible)
Naval |  0.000 $\pm$ 0.000 |  0.001 $\pm$ 0.001 |  0.000 $\pm$ 0.000 | 0.007 $\pm$ 0.003 |  0.002 $\pm$ 0.000 |  0.000 $\pm$ 0.000 |  0.197 $\pm$ 0.174 | 0.000 $\pm$ 0.000 | (infeasible)
Power |  0.052 $\pm$ 0.001 |  0.054 $\pm$ 0.002 | 0.054 $\pm$ 0.001 | 0.058 $\pm$ 0.002 |  0.054 $\pm$ 0.002 |  0.054 $\pm$ 0.002 | 0.052 $\pm$ 0.001 | 0.071 $\pm$ 0.001 | (infeasible)
Protein | 0.459 $\pm$ 0.005 | 0.466 $\pm$ 0.004 |  0.429 $\pm$ 0.004 |  0.537 $\pm$ 0.008 |  0.421 $\pm$ 0.005 | 0.446 $\pm$ 0.006 | 0.425 $\pm$ 0.003 | 0.408 $\pm$ 0.002 | (infeasible)
Wine | 0.652 $\pm$ 0.022 | 0.663 $\pm$ 0.009 | 1.297 $\pm$ 0.093 |  0.655 $\pm$ 0.023 | 0.627 $\pm$ 0.013 | 0.637 $\pm$ 0.031 |  0.682 $\pm$ 0.048 | 0.607 $\pm$ 0.033 | 0.585 $\pm$ 0.032
Yacht |  0.003 $\pm$ 0.001 | 0.004 $\pm$ 0.001 | 0.221 $\pm$ 0.037 |  0.682 $\pm$ 0.140 | 0.004 $\pm$ 0.001 |  0.002 $\pm$ 0.001 | 0.008 $\pm$ 0.003 | 0.399 $\pm$ 0.064 | 0.355 $\pm$ 0.030
Wave |  0.000 $\pm$ 0.000 |  0.000 $\pm$ 0.000 |  0.000 $\pm$ 0.000 |  0.000 $\pm$ 0.000 |  0.000 $\pm$ 0.000 |  0.000 $\pm$ 0.000 | 0.001 $\pm$ 0.001 | 0.000 $\pm$ 0.000 | (infeasible)
Denmark |  0.155 $\pm$ 0.004 | 0.287 $\pm$ 0.003 | 0.163 $\pm$ 0.004 | 0.225 $\pm$ 0.003 |  0.189 $\pm$ 0.008 | 0.194 $\pm$ 0.003 | 0.197 $\pm$ 0.004 | 0.260 $\pm$ 0.001 | (infeasible)
Mean rank | 1.364 | 2.000 | 2.182 | 3.182 | 1.636 | 1.727 | - | - | -

We present additional regression results reporting the mean square error (MSE)
of evaluated methods across the considered baselines, see table 7. We find
that GFSVI also performs competitively in terms of MSE compared to baselines
and obtains the best mean rank, matching best the performing methods on nearly
all datasets. In particular, we find that using GP priors in the linearized
BNN setup with GFSVI yields improvements over the weight-space priors used in
TFSVI and that GFSVI performs slightly better than FVI. Function-space VI
methods (TFSVI, GFSVI, FVI) significantly improves over weight-space VI mostly
performing similarly to the linearized Laplace approximation. Further
improvement over baselines are obtained when considering GP priors with GFSVI
and FVI. Finally, GFSVI compares favorably to the GP and sparse GP.

####  A.4.3 Out-of-distribution detection with image data

We here show an additional plot from our out-of-distribution detection
experiment with image data (details in A.3.5). figure 12 shows the
(normalized) histograms of the entropy of the mean prediction produced by each
model on the in-distribution (blue) and out-of-distribution (red) data sets
considered in our OOD detection experiment. Methods which estimate the
(regularized) KL-divergence in function-space (GFSVI, FVI and TFSVI) use the
kmnist measurement distribution. We find that the entropy produced by GFSVI on
in-distribution data highly peaks around 0 while the entropy produced from
out-of-distribution data strongly concentrates around its maximum $\ln(10)$.
GFSVI best partitions ID and OOD data based on predictive entropy improving
over the function-space prior (FVI) and weight-space prior (TFSVI, MFVI,
Laplace) BNN baselines (see table 2).

Figure 12: Histograms of the entropy of the mean predictive distribution
produced by evaluated methods in the out-of-distribution detection with image
data experiment. GFSVI (ours) best partitions in-distribution and out-of-
distribution data based on the entropy of its mean predictive distribution.

####  A.4.4 Influence of measurement point distribution for image data.

We present additional results evaluating the influence of the measurement
point distribution $\rho$ on the the performance of function-space inference
methods when using high-dimensional image data. The measurement point
distribution are described in section A.3.4. Just like in Rudner et al., 2022b
, we find that the choice of measurement point distribution may highly
influence the OOD detection accuracy. While the expected log-likelihood,
accuracy and expected calibration error (ECE) of a model generally remains
comparable across measurement point distributions, the OOD accuracy of GFSVI
is greatly improved by using samples from kmnist to evaluate the (regularized)
KL divergence. The measurement point distribution determines where the BNN is
regularized and thus should be carefully selected especially for high
dimensional data.

Table 8: Influence of the measurement point distribution $\rho$ on expected
log-likelihood (log-like.), accuracy (acc.), expected calibration error (ECE)
and out-of-distribution detection accuracy (OOD acc.). $\rho$ determines where
the BNN will be regularized and strongly influences the out-of-distribution
performance of the BNN.

Data | Metric | GFSVI | FVI | TFSVI
---|---|---|---|---
| Random | Random Pixel | KMNIST | Random | Random Pixel | KMNIST | Random | Random Pixel | KMNIST
MNIST |  Log-like. ($\uparrow$) | -0.033 $\pm$ 0.000 |  -0.034 $\pm$ 0.000 |  -0.041 $\pm$ 0.000 |  -0.145 $\pm$ 0.005 |  -0.038 $\pm$ 0.000 |  -0.238 $\pm$ 0.006 |  -0.047 $\pm$ 0.003 | -0.032 $\pm$ 0.001 |  -0.041 $\pm$ 0.001
Acc. ($\uparrow$) |  -0.992 $\pm$ 0.000 |  -0.989 $\pm$ 0.000 |  -0.991 $\pm$ 0.000 |  -0.976 $\pm$ 0.001 |  -0.988 $\pm$ 0.000 |  -0.943 $\pm$ 0.001 |  -0.989 $\pm$ 0.000 |  -0.989 $\pm$ 0.000 |  -0.989 $\pm$ 0.000
|  ECE ($\downarrow$) |  -0.002 $\pm$ 0.000 |  -0.004 $\pm$ 0.000 |  -0.006 $\pm$ 0.000 |  -0.064 $\pm$ 0.001 |  -0.003 $\pm$ 0.000 |  -0.073 $\pm$ 0.003 |  -0.007 $\pm$ 0.000 |  -0.003 $\pm$ 0.000 |  -0.006 $\pm$ 0.000
|  OOD acc. ($\uparrow$) |  -0.921 $\pm$ 0.008 |  -0.868 $\pm$ 0.010 |  -0.980 $\pm$ 0.004 |  -0.894 $\pm$ 0.010 |  -0.863 $\pm$ 0.003 |  -0.891 $\pm$ 0.006 |  -0.887 $\pm$ 0.011 |  -0.861 $\pm$ 0.008 |  -0.893 $\pm$ 0.005
FMNIST |  Log-like. ($\uparrow$) | -0.260 $\pm$ 0.003 | -0.258 $\pm$ 0.002 |  -0.294 $\pm$ 0.006 |  -0.300 $\pm$ 0.002 |  -0.293 $\pm$ 0.003 |  -0.311 $\pm$ 0.005 |  -0.261 $\pm$ 0.001 | -0.258 $\pm$ 0.001 |  -0.261 $\pm$ 0.002
Acc. ($\uparrow$) |  -0.910 $\pm$ 0.001 |  -0.908 $\pm$ 0.001 |  -0.909 $\pm$ 0.001 |  -0.910 $\pm$ 0.002 |  -0.900 $\pm$ 0.001 |  -0.906 $\pm$ 0.002 |  -0.909 $\pm$ 0.001 |  -0.908 $\pm$ 0.001 |  -0.907 $\pm$ 0.001
|  ECE ($\downarrow$) |  -0.020 $\pm$ 0.003 |  -0.022 $\pm$ 0.001 |  -0.042 $\pm$ 0.002 |  -0.027 $\pm$ 0.005 |  -0.018 $\pm$ 0.002 |  -0.024 $\pm$ 0.002 |  -0.022 $\pm$ 0.002 |  -0.018 $\pm$ 0.001 |  -0.021 $\pm$ 0.002
|  OOD acc. ($\uparrow$) |  -0.853 $\pm$ 0.005 |  -0.867 $\pm$ 0.005 |  -0.997 $\pm$ 0.001 |  -0.925 $\pm$ 0.005 |  -0.842 $\pm$ 0.006 |  -0.975 $\pm$ 0.002 |  -0.802 $\pm$ 0.006 |  -0.800 $\pm$ 0.007 |  -0.779 $\pm$ 0.010

####  A.4.5 Input distribution shift with rotated image data

We here provide an experiment evaluating our method’s (GFSVI) robustness in
detecting input distribution shift. We expect the predictive uncertainty of a
well-calibrated Bayesian model to be low for in-distribution data and to
gradually increase as the input distribution shifts further away from the
training data distribution. To test this property, we follow the setup by
Sensoy et al., (2018); Rudner et al., 2022b  and assume like the related work
that increasing the rotation angle of images gradually increases the level of
input "distribution shift". We report the mean and standard-deviation of the
average mean predictive entropy of models fit on MNIST (LeCun et al.,, 2010)
and Fashion MNIST (Xiao et al.,, 2017) for increasingly large angles of
rotation of their respective test data partition. We find that GFSVI is
confident (low predictive entropy) for images with small rotation angles, and
that its predictive entropy increases with the angle. GFSVI therefore exhibits
the expected behavior of a well-calibrated Bayesian model. We note that FVI,
Laplace and MFVI tend to be under-confident (high predictive entropy) for
small rotation angles, which might be a symptom of underfitting further
supported by the results in table 2. Also, with the exception of TFSVI, the
predictive entropy of baselines across different rotation angles is generally
higher than the one produced by GFSVI.

Figure 13: Average predictive entropy of models trained on MNIST and Fashion
MNIST and evaluated for different rotation angles of their respective test
data partitions. We see that our method (GFSVI) exhibits the behavior of a
well-calibrated Bayesian model.

####  A.4.6 Additional plots for kernel eigenvalue decay

figure 14 shows a plot demonstrating the decay rate of the eigenvalues of RBF
and Matern-1/2 kernels evaluated at points sampled uniformly over
$\mathcal{X}$ . The rate of decay of covariance operator’s eigenvalues gives
important information about the smoothness of stationary kernels (Williams and
Rasmussen,, 2006) and that increased smoothness of the kernel leads to faster
decay of eigenvalues Santin and Schaback, (2016). For instance, RBF covariance
operator eigenvalues decay at near exponential rate independent of the
underlying measure (Belkin,, 2018) and Matern kernels eigenvalues decay
polynomialy (Chen et al.,, 2021). We find that the kernel evaluated at points
sampled from a uniform distribution over $\mathcal{X}$ share this same
behavior (see figure 14).

|
---|---

(a) $D=1$ (b) $D=100$

Figure 14: Mean eigenvalues of the Gram matrix obtained for different kernels
and for varying length-scales over $10$ draws from a uniform distribution on
$[-2,2]^{D}$. The mean eigenvalues are arranged in increasing order. The
eigenvalues of the Gram matrix associated with the smooth RBF kernel decays
much faster than those of the Matern-1/2. Furthermore, the eigenvalues decay
at a slower rate in high dimensions (D=100).

####  A.4.7 Additional plots for choosing $\gamma$ in $D_{\text{KL}}^{\gamma}$

The $\gamma$ parameter controls the magnitude of the regularized KL divergence
(see figure 17) and adjusts the relative weight of the regularized KL
divergence and expected log-likelihood term in the training objective (see
figure 15). Furthermore, $\gamma$ also acts as "jitter" preventing numerical
errors. We recommend choosing $\gamma$ large enough to avoid numerical errors
while remaining small enough to provide strong regularization.

Figure 15: The $\gamma$ parameter of the regularized KL divergence controls
the magnitude of the regularizer in the objective and should be small enough
to provide strong regularization.

Figure 16: The BNN’s covariance adaptation to the prior’s covariance rank
depends on its activation function. BNNs fit with a RBF prior (full) show
lower rank than with a Matern-1/2 (dotted).

Figure 17: $\gamma$ explicitly controls the magnitude of the regularized KL-
divergence $D_{\text{KL}}^{\gamma}$. Rougher priors (Matern-1/2) require more
measurement points to accurately estimate $D_{\text{KL}}^{\gamma}$ than smooth
priors (RBF).

## NeurIPS Paper Checklist

  1. 1.

Claims

  2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

  3. Answer: [Yes]

  4. Justification: The well-definiteness of our variational objective follows from the fact that the regularized KL divergence is finite for any pair of Gaussian measures Quang, (2019). tables 1, 2, 4 and 5 as well as Figures 2 to 10 support our claim about the performance of our method.

  5. Guidelines:

     * •

The answer NA means that the abstract and introduction do not include the
claims made in the paper.

     * •

The abstract and/or introduction should clearly state the claims made,
including the contributions made in the paper and important assumptions and
limitations. A No or NA answer to this question will not be perceived well by
the reviewers.

     * •

The claims made should match theoretical and experimental results, and reflect
how much the results can be expected to generalize to other settings.

     * •

It is fine to include aspirational goals as motivation as long as it is clear
that these goals are not attained by the paper.

  6. 2.

Limitations

  7. Question: Does the paper discuss the limitations of the work performed by the authors?

  8. Answer: [Yes]

  9. Justification: The inductive biases paragraph in section 4.1 discusses limitations.

  10. Guidelines:

     * •

The answer NA means that the paper has no limitation while the answer No means
that the paper has limitations, but those are not discussed in the paper.

     * •

The authors are encouraged to create a separate "Limitations" section in their
paper.

     * •

The paper should point out any strong assumptions and how robust the results
are to violations of these assumptions (e.g., independence assumptions,
noiseless settings, model well-specification, asymptotic approximations only
holding locally). The authors should reflect on how these assumptions might be
violated in practice and what the implications would be.

     * •

The authors should reflect on the scope of the claims made, e.g., if the
approach was only tested on a few datasets or with a few runs. In general,
empirical results often depend on implicit assumptions, which should be
articulated.

     * •

The authors should reflect on the factors that influence the performance of
the approach. For example, a facial recognition algorithm may perform poorly
when image resolution is low or images are taken in low lighting. Or a speech-
to-text system might not be used reliably to provide closed captions for
online lectures because it fails to handle technical jargon.

     * •

The authors should discuss the computational efficiency of the proposed
algorithms and how they scale with dataset size.

     * •

If applicable, the authors should discuss possible limitations of their
approach to address problems of privacy and fairness.

     * •

While the authors might fear that complete honesty about limitations might be
used by reviewers as grounds for rejection, a worse outcome might be that
reviewers discover limitations that aren’t acknowledged in the paper. The
authors should use their best judgment and recognize that individual actions
in favor of transparency play an important role in developing norms that
preserve the integrity of the community. Reviewers will be specifically
instructed to not penalize honesty concerning limitations.

  11. 3.

Theory Assumptions and Proofs

  12. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

  13. Answer: [Yes]

  14. Justification: Our only theoretical result is proved in Appendix and is the direct application of the Feldman-Hàjek theorem. We specify all assumptions.

  15. Guidelines:

     * •

The answer NA means that the paper does not include theoretical results.

     * •

All the theorems, formulas, and proofs in the paper should be numbered and
cross-referenced.

     * •

All assumptions should be clearly stated or referenced in the statement of any
theorems.

     * •

The proofs can either appear in the main paper or the supplemental material,
but if they appear in the supplemental material, the authors are encouraged to
provide a short proof sketch to provide intuition.

     * •

Inversely, any informal proof provided in the core of the paper should be
complemented by formal proofs provided in appendix or supplemental material.

     * •

Theorems and Lemmas that the proof relies upon should be properly referenced.

  16. 4.

Experimental Result Reproducibility

  17. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

  18. Answer: [Yes]

  19. Justification: The experimental setup is rigorously described in Appendix A.3.

  20. Guidelines:

     * •

The answer NA means that the paper does not include experiments.

     * •

If the paper includes experiments, a No answer to this question will not be
perceived well by the reviewers: Making the paper reproducible is important,
regardless of whether the code and data are provided or not.

     * •

If the contribution is a dataset and/or model, the authors should describe the
steps taken to make their results reproducible or verifiable.

     * •

Depending on the contribution, reproducibility can be accomplished in various
ways. For example, if the contribution is a novel architecture, describing the
architecture fully might suffice, or if the contribution is a specific model
and empirical evaluation, it may be necessary to either make it possible for
others to replicate the model with the same dataset, or provide access to the
model. In general. releasing code and data is often one good way to accomplish
this, but reproducibility can also be provided via detailed instructions for
how to replicate the results, access to a hosted model (e.g., in the case of a
large language model), releasing of a model checkpoint, or other means that
are appropriate to the research performed.

     * •

While NeurIPS does not require releasing code, the conference does require all
submissions to provide some reasonable avenue for reproducibility, which may
depend on the nature of the contribution. For example

       1. (a)

If the contribution is primarily a new algorithm, the paper should make it
clear how to reproduce that algorithm.

       2. (b)

If the contribution is primarily a new model architecture, the paper should
describe the architecture clearly and fully.

       3. (c)

If the contribution is a new model (e.g., a large language model), then there
should either be a way to access this model for reproducing the results or a
way to reproduce the model (e.g., with an open-source dataset or instructions
for how to construct the dataset).

       4. (d)

We recognize that reproducibility may be tricky in some cases, in which case
authors are welcome to describe the particular way they provide for
reproducibility. In the case of closed-source models, it may be that access to
the model is limited in some way (e.g., to registered users), but it should be
possible for other researchers to have some path to reproducing or verifying
the results.

  21. 5.

Open access to data and code

  22. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

  23. Answer: [Yes]

  24. Justification: The pseudo code, experiment descriptions in appendix as well as the code provides all the relevant information to reproduce experiments.

  25. Guidelines:

     * •

The answer NA means that paper does not include experiments requiring code.

     * •

Please see the NeurIPS code and data submission guidelines
(https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.

     * •

While we encourage the release of code and data, we understand that this might
not be possible, so “No” is an acceptable answer. Papers cannot be rejected
simply for not including code, unless this is central to the contribution
(e.g., for a new open-source benchmark).

     * •

The instructions should contain the exact command and environment needed to
run to reproduce the results. See the NeurIPS code and data submission
guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more
details.

     * •

The authors should provide instructions on data access and preparation,
including how to access the raw data, preprocessed data, intermediate data,
and generated data, etc.

     * •

The authors should provide scripts to reproduce all experimental results for
the new proposed method and baselines. If only a subset of experiments are
reproducible, they should state which ones are omitted from the script and
why.

     * •

At submission time, to preserve anonymity, the authors should release
anonymized versions (if applicable).

     * •

Providing as much information as possible in supplemental material (appended
to the paper) is recommended, but including URLs to data and code is
permitted.

  26. 6.

Experimental Setting/Details

  27. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

  28. Answer: [Yes]

  29. Justification: The experimental setup is rigorously described in Appendix A.3.

  30. Guidelines:

     * •

The answer NA means that the paper does not include experiments.

     * •

The experimental setting should be presented in the core of the paper to a
level of detail that is necessary to appreciate the results and make sense of
them.

     * •

The full details can be provided either with the code, in appendix, or as
supplemental material.

  31. 7.

Experiment Statistical Significance

  32. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

  33. Answer: [Yes]

  34. Justification: We report mean and standard-error of our method across five cross-validation folds (regression) or random partitions of the data (classification).

  35. Guidelines:

     * •

The answer NA means that the paper does not include experiments.

     * •

The authors should answer "Yes" if the results are accompanied by error bars,
confidence intervals, or statistical significance tests, at least for the
experiments that support the main claims of the paper.

     * •

The factors of variability that the error bars are capturing should be clearly
stated (for example, train/test split, initialization, random drawing of some
parameter, or overall run with given experimental conditions).

     * •

The method for calculating the error bars should be explained (closed form
formula, call to a library function, bootstrap, etc.)

     * •

The assumptions made should be given (e.g., Normally distributed errors).

     * •

It should be clear whether the error bar is the standard deviation or the
standard error of the mean.

     * •

It is OK to report 1-sigma error bars, but one should state it. The authors
should preferably report a 2-sigma error bar than state that they have a 96%
CI, if the hypothesis of Normality of errors is not verified.

     * •

For asymmetric distributions, the authors should be careful not to show in
tables or figures symmetric error bars that would yield results that are out
of range (e.g. negative error rates).

     * •

If error bars are reported in tables or plots, The authors should explain in
the text how they were calculated and reference the corresponding figures or
tables in the text.

  36. 8.

Experiments Compute Resources

  37. Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

  38. Answer: [Yes]

  39. Justification: We specify computational resources in Appendix A.3.7.

  40. Guidelines:

     * •

The answer NA means that the paper does not include experiments.

     * •

The paper should indicate the type of compute workers CPU or GPU, internal
cluster, or cloud provider, including relevant memory and storage.

     * •

The paper should provide the amount of compute required for each of the
individual experimental runs as well as estimate the total compute.

     * •

The paper should disclose whether the full research project required more
compute than the experiments reported in the paper (e.g., preliminary or
failed experiments that didn’t make it into the paper).

  41. 9.

Code Of Ethics

  42. Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

  43. Answer: [Yes]

  44. Justification: To the best of our knowledge, we have fully conformed with the code of ethics.

  45. Guidelines:

     * •

The answer NA means that the authors have not reviewed the NeurIPS Code of
Ethics.

     * •

If the authors answer No, they should explain the special circumstances that
require a deviation from the Code of Ethics.

     * •

The authors should make sure to preserve anonymity (e.g., if there is a
special consideration due to laws or regulations in their jurisdiction).

  46. 10.

Broader Impacts

  47. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

  48. Answer: [No]

  49. Justification: This paper presents work whose goal is to advance the field of Bayesian deep learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

  50. Guidelines:

     * •

The answer NA means that there is no societal impact of the work performed.

     * •

If the authors answer NA or No, they should explain why their work has no
societal impact or why the paper does not address societal impact.

     * •

Examples of negative societal impacts include potential malicious or
unintended uses (e.g., disinformation, generating fake profiles,
surveillance), fairness considerations (e.g., deployment of technologies that
could make decisions that unfairly impact specific groups), privacy
considerations, and security considerations.

     * •

The conference expects that many papers will be foundational research and not
tied to particular applications, let alone deployments. However, if there is a
direct path to any negative applications, the authors should point it out. For
example, it is legitimate to point out that an improvement in the quality of
generative models could be used to generate deepfakes for disinformation. On
the other hand, it is not needed to point out that a generic algorithm for
optimizing neural networks could enable people to train models that generate
Deepfakes faster.

     * •

The authors should consider possible harms that could arise when the
technology is being used as intended and functioning correctly, harms that
could arise when the technology is being used as intended but gives incorrect
results, and harms following from (intentional or unintentional) misuse of the
technology.

     * •

If there are negative societal impacts, the authors could also discuss
possible mitigation strategies (e.g., gated release of models, providing
defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms
to monitor how a system learns from feedback over time, improving the
efficiency and accessibility of ML).

  51. 11.

Safeguards

  52. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

  53. Answer: [N/A]

  54. Justification: The methods developed in the paper have no risk for misuse.

  55. Guidelines:

     * •

The answer NA means that the paper poses no such risks.

     * •

Released models that have a high risk for misuse or dual-use should be
released with necessary safeguards to allow for controlled use of the model,
for example by requiring that users adhere to usage guidelines or restrictions
to access the model or implementing safety filters.

     * •

Datasets that have been scraped from the Internet could pose safety risks. The
authors should describe how they avoided releasing unsafe images.

     * •

We recognize that providing effective safeguards is challenging, and many
papers do not require this, but we encourage authors to take this into account
and make a best faith effort.

  56. 12.

Licenses for existing assets

  57. Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

  58. Answer: [Yes]

  59. Justification: Datasets used for the experiments are cited in the main text and their license and terms of use are properly respected.

  60. Guidelines:

     * •

The answer NA means that the paper does not use existing assets.

     * •

The authors should cite the original paper that produced the code package or
dataset.

     * •

The authors should state which version of the asset is used and, if possible,
include a URL.

     * •

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

     * •

For scraped data from a particular source (e.g., website), the copyright and
terms of service of that source should be provided.

     * •

If assets are released, the license, copyright information, and terms of use
in the package should be provided. For popular datasets,
paperswithcode.com/datasets has curated licenses for some datasets. Their
licensing guide can help determine the license of a dataset.

     * •

For existing datasets that are re-packaged, both the original license and the
license of the derived asset (if it has changed) should be provided.

     * •

If this information is not available online, the authors are encouraged to
reach out to the asset’s creators.

  61. 13.

New Assets

  62. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

  63. Answer: [N/A]

  64. Justification: The paper does not release new assets.

  65. Guidelines:

     * •

The answer NA means that the paper does not release new assets.

     * •

Researchers should communicate the details of the dataset/code/model as part
of their submissions via structured templates. This includes details about
training, license, limitations, etc.

     * •

The paper should discuss whether and how consent was obtained from people
whose asset is used.

     * •

At submission time, remember to anonymize your assets (if applicable). You can
either create an anonymized URL or include an anonymized zip file.

  66. 14.

Crowdsourcing and Research with Human Subjects

  67. Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

  68. Answer: [N/A]