Generalized Variational Inference in Function Spaces: Gaussian Measures meet Bayesian Deep Learning
Authors: Veit David Wild
Robert Hu
Dino Sejdinovic
Venue: NeurIPS 2022 Accept (2022)
Keywords: Function Space Inference
Variational Inference
Bayesian Learning
Bayesian Inference
Generalised Variational Inference
Gaussian Processes
Gaussian Measures

Abstract: We develop a framework for generalized variational inference in infinite-dimensional function spaces and use it to construct a method termed Gaussian Wasserstein inference (GWI). GWI leverages the Wasserstein distance between Gaussian measures on the Hilbert space of square-integrable functions in order to determine a variational posterior using a tractable optimization criterion. It avoids pathologies arising in standard variational function space inference. An exciting application of GWI is the ability to use deep neural networks in the variational parametrization of GWI, combining their superior predictive performance with the principled uncertainty quantification analogous to that of Gaussian processes. The proposed method obtains state-of-the-art performance on several benchmark datasets.

Generalized Variational Inference in Function Spaces:
Gaussian Measures meet Bayesian Deep Learning
Veit D. Wild" Department of Statistics University of Oxford weit.wild@stats ox. ac uk
Robert Hu"
Amazon
robyhu@amazon.co..
Dino Sejdinovie School Computer and Mathematical Sciences University of Adelaide
dino seejdinovvicaedllladeee. edu. au
Abstract
We develop a framework for generalized variational inference in infinite- dimensional function spaces and use it to construct a method termed Gaussian Wasserstein inference (GWI). GWI leverages the Wasserstein distance between Gaussian measures on the Hilbert space of square-intggrabl functions in order to determine variational posterior using a tractable optimization criterion. It avoids pathologies arising in standard variational function space inference. An exciting application of GWT is the ability use deep neural networks the variational parametrization of GWI, combining their superior predictive performance with the principled uncertainty quantification analogous to that of Gaussian processes. The proposed method obtains state-of-the-at performance on several benchmark datasets.
Introduction
In the past decade, considerable effort has been invested in developing Bayesian deep learning approaches [Welling and Teh 201 1. Chen al., 2014 Blundell al. 2015 Gal and Ghahramani 2016 Kendall and Gal 2017 Ritter et al. 20188] Khan et al.) 2018 MMddox et al. [019]. There are at least two key advantages Bayesian models. Firstly, Bayesian model averaging known to improve predictive performance [Komaki] 1996] even in misspecified situations Fushiki 2005 Ramamoorthi
al. 2015). The empirical success of methods such as deep ensembles [laakshminarayanan et al.
2017 may be interpreted compelling evidence for this claim [Wilson and Izmailov] (2220). 2020)) Sec- ondly, Bayesian models provide the user with a predictive distribution for an unseen data point. This can be naturally leveraged to quantifyy posterior uncertainty.
Even though impressive progress has been made, there are problems that remain unresolved. The prior distribution for the unknown function is typically induced by prior distribution over deep neural network weights (and biases). is hard to interpret the inductive bias in function space that is induced by such priors for weights and unclear how one might incorporate prior knowledge about the unknown function. Additionally, the resulting inference problem extremely high-dimensional and requires approximation techniques that are either computationally expensive [Neal, 2012) or so crude that the approximate posterior may suffer from pathological behavior [Foong al.) 2020). The difficulties of performing Bayesian inference for weights have led to the emergence of methods that
"qual contribution, order decided by coinflip Work mainly done while the authors were with the Department of Statistics, University of Oxford
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

approach the problem in function space directly (Ma et al. 2019 Sun et al. 2019 Rudner et al. 2020
Ma and Hernändez- Lobato, 202T].
The theory of constructing prior distributions in function spaces is well developed and the most famous class of prior distributions are Gaussian processes. They have been commonly used for decades in the machine learning community to elicit interpretable functional priors and are known to have well-calibrated predictive uncertainties [Rasmussen 2003].
In separate thread of research, new powerful inference framework called Generalized Variational Inference (GVT) has been recently developed [Knoblauch et al.. 2019). The authors argue that standard assumptions of Bayesian inference such as well-specified priors, well-speeified likelihoods and infinite computing power are often violated in practice. They therefore propose a generalized view on Bayesian inference that takes these points into consideration. We extend the work of Knoblauch al. [019 to situations where probability density functions for the prior exist and are thus able to use generalized variational inference in infinite-dimensional function spaces directly. We then specify both the prior and variational measures as Gaussian measures and measure their dissimilarity using the Wasserstein distance. This results in the method which we call Gaussian Wasserstein Inference in Function Spaces (GWI-FS). An exciting application of our method is the ability to equip deep neural networks with uncertainty quantification using the framework analogous to that of Gaussian processes, resulting in state-of- of-tthe-ar art method termed GWI-net. Our main contributions are:
We create general framework for inference in function space based on Gaussian measures on the space of square-integrable functions,
We derive an objective function that can be expressed in terms of the parameters of the Gaussian measures,
We derive tractable approximation to our objective function that is valid for (almost) arbitrary kernels and mean functions, We demonstrate the utility of our method by obtaining state-of-the-art results on the UCI regression datasets and on Fashion MNIST and CIFAR 1
2 Related Work
GWI-FS draws on the work developed in the Gaussian process literature, but can be used to equip traditional neural network architectures with uncertainty. We therefore give a brief overview of the relevant related methods in both the Bayesian neural network (BNNs) and Gaussian process community.
Bayesian neural networks Traditionally Bayesian neural networks have been assigned priors in weight space. The effects of various priors on inference and uncertainty quantification are still not well understood [Fortuin et al. [021]. As the posterior (over weights) is intractable, sampling algorithms such as Hamiltonian Monte Carlo (HMC) were initially proposed Neal [2012]. J. Due to the unfavorable scaling properties of standard HMC which requires the full gradient, batch-size approximations of HMC evolved [Chen et al., 2014). Another line of research exploits Langevin dynamics to generate posterior samples [Welling and Teh 201T) in weight space.
Variational methods for BNNs in weight space In variational inference, the true posterior is approximated by more tractable so-called variational distribution. The user specifies a class of approximate posterior measures and selects the best posterior approximation by maximizing the so-called evidence lower bound (ELBO). The Bayes by Backprop [Blundell et al. 2015 method is one such variational mien-field approximation of the weight- space posterior. variational dropout [Gal and Ghahramani 2016), specific approximation is chosen to reinterpret dropout (Srivastaval cct al. 2014] at test time as a variational procedure.
Variational methods for BNNs in function spaces Inference in weight space is challenging, as the problem is typically high- -igh--imensional and the posterior distribution over weights multi-modal. This led to a line of research in which inference algorithms are formulated in function spaces. Variational implicit processes [Ma et al., 2019] approximate the BNN posterior as a linear combination of draws from the prior. Functional-BNN [Sun et al., 2019) matches BNN to functional prior (for
example a GP) and performs inference by optimising functional Kullback-Leiblee (KI.) divergence
exploiting score function estimators [Li and Turner 2017 Shi et al. 2018). Rudner et al. (2020) use
a local approximation the prior and variational posterior processes obtain a tractable functional
Cooebaase: https//githbb.om//Mtttufff
2

---

Kullback- Leibler divergence. Ma and Hernández-Lobato [2021] generalise the variational family in Ma et al. [2099] and obtain a more scalable procedure by using a different approximation the functional KI -iv-eivegencee Recent work has also proposed to adapt BNN priors to interpretable functional priors by minimizing the Wasserstein distance between a BNN prior and Gaussian process (Tran et al. 2020]. Another line of research exploits the Wasserstein gradient flow and tries
to encourage diversity in the function space (D Angelo et al. 2021 A'Aggolo and Fortuin 2021].
Gaussian processes Standard Gaussian process regression [Rasmussen 2003] allows interpretable prior specification but scales poorly with respect to the number of data points, As result, plethora of approximation techniques are introduced. On one hand, there are variational approximations to the true posterior (Titsias 2009 [Hensman et al. 2013] and several extensions [Hensman al. 2017 Salimbeni et al.] [018 Dutordoir et al. [2020] On the other hand, GPU utilization is combined with
Krylov subspace methods to obtain scalability [Gardner et al. 2018 Wang et al. 2019
3 Background
In this section we give some background on generalized variational inference in infinite dimensions and introduce Gaussian measures in Hilbert spaces. We further discuss their relation the more familiar Gaussian processes at the end.
3.1 Generalized Variational Inference in Function Spaces
In functional variational inference, we assign a prior p(f) ) to the unknown function f E E. where is a function space" The prior is combined with the likelihood p(ylf) to give the posterior p(fly). The posterior is often intractable which is why variational inference we specify a tractable variational approximation (()) to p(fly) and train model by maximising the evidence lower bound (ELBO)
f = EE) log p(u]f)] - Dxx ((()),(()))
(1)
where DRi. denotes the KL divergence. Note that in the case where E is infinite dimensional p(/) and q(f) cannot be probability density functions with respect to the Lebesgue measure [see e.g. Hunt et al. 1992 for discussion]. which is why the above notation, although commonly used, is imprecise. What we in fact mean are the probability measures over E associated with the prior and variational approximation. We will denote these measures as pf and Q" from now on to make this difference explicit. The ELBO in this notation reads as
== Eel log p(y) P)] DKI. (QF,P.)
Note that the KL divergence (for measures) is defined as
DRL (Q (Pf = f log
dQP dPF (f) dQ (f).
(3)
where we assume that QF is dominated by the measure PF which guarantees the existence of the Radon-Nikodym derivative dQF /dIF". A number of papers focus on obtaining tractable approxima- tions of (3) [Sn et al. 2019, Rudner et al., 2020 Ma and Hernández- Lobato 202T]. However, the function spaces can be a delicate task, since benign constructions of priors and variational approximations may not satisfy that "" is dominated by 1'"
use of KL- divergence in anfinit-dimension
which leads DxL ( pF) = oo [Burt et al. 2020). This often renders the objective useless or at least problematic.
A true Bayesian is committed to the use of the KL divergence in maximizing L is equivalent to minimizing the KL divergence between the true posterior measure and the variational measure. This equivalence is typically demonstrated using pdfs but the argument generalizes to infinite dimensions as is shown for GPs in Matthews et al. 2016) or in a more measure theoretic formulation in Theorem 4 of Wild and Wynne 2021).
However, Knoblauch et al. [019 argue that given the problems of prior and likelihood specification as well available compute, an axiomatically justified way of moving from prior to posterior
belief' is by solving a more general optimization problem [Knoblauch et al.] 2019 Theorem 15]. Crucially it is valid replace the KL-divergence by an arbitrary measure of dissimilarity ID satisfying
D(QF, F. IPF) 2 and (Q" PF) = 0 = Q = PF. The arguments in Knoblauch et al. 2019 are
'We assume E' to be Polish space, which avoids technical difficulties in defining the posterior measure Ghosal and Van der Vaart 2017 Chapter 1.3
3

---

made assuming the existence of pdf for the prior, but they rely solely on a reformulation of Bayesian inference as optimization problem [Knoblauch et al. 2019 Chapter 2). We show in Appendix A.. that this reformulation can also be made for infniit-dimenniooaal prior measures and therefore consider the generalized loss
c log F) + DQQ PP))
valid optimization objective for arbitrary dissimilarity measure D. This is merely an infinite- dimensional version of equation (10) in Knoblauch et al.] (2019). We refer to inference targeting the objective Generalized variational inference in function space (GVI-FS).
Generalised variational inference can be interpreted as regularised loss minimisation lifted into the space of probability measures. The first term in 4 is understood as loss which we want to minimise on average, while the second term punishes strong deviations from the prior.
The particular instance of GVI-FS that we explore is where both F and Q are Gaussian measures (on an dimenioadi Hilbert space) and MD chosen to be the Wasserstein metrie [Kantorovich] 1960]. We will refer to this setting as Gaussian Wasserstein Inference in Function Space (GWI-FS) or more consciously as Gaussian Wasserstein Inference (GWI)
3.2 Gaussian Random Elements and Gaussian Measures in Hilbert spaces
In this section introduce Gaussian random elements (GRE) and Gaussian measures in Hilbert spaces - these concepts are somewhat technical but crucial in the construction of our method. We then describe their close relationship to the more familiar Gaussian process notions in the next section.
Let (2, A, P) be the underlying (physical) probability space and (H, ())) be Hilbert space.
Gaussian random elements
A measurable function F : 52 + H is called GRE (in H) if and only if (F, h) : Ω + E has scalar Gaussian distribution for all h E / | Every GRE F has a mean element m E H defined by
m. / F(w) dP(w)
(5)
and a (linear) covariance operator C: H - H defined by
Ch(-) = (F(W), P(w) (m, h)m.
(6)
for h E H. Both integrals are to be understood as Bochner integrals Kukush 2020 Chapter 3]. The
Bochner integral has the property that (J F(w) dw) h) = (F h) dl((w) w for all h E H. This combined with Fubini's 's theorem and the definition of a GRE implies that
(F, h) ~ (mm, h), ()),
for any h E H with o²) denoting the normal distribution with mean l E R and variance a² > O. Similarly we denote F (mm,, C) for a GRE in H with mean element m and covariance operator C. It can be shown that the covariance operator C' of GRE is positive self-adjoint trace-class operator. Conversely, for every positive self-adjoint trace class operator and every m E H, there exists GRE with F ~ (m, C') [Bogachey 1998 Theorem 2.3.1].
Gaussian measures The push- forward measure of P through F is defined as P (A) ((F- for all Borel-measurable A H. If F A(m, C') is a GRE, call P = Pp GM and write P = (m, C). Note that GMs equivalently GREs allow us to specify probability distributions over (infinite-dinn dimens-ddmmesionaa) Hilbert spaces by using a given mean element and a given covariance operator.
Details about Gaussian Measures in Hilbert spaces can be found Chapter 2 Da Prato and Zabezyk (214)] or in Kukush 2020 In fact, Gaussian measures can be defined on even more general linear spaces such as Banach or Fréchet spaces [Bogachev] 1998].
3.3 Gaussian Processes and Their Corresponding Measures
In this section describe how Gaussian processes standard tool to assign functional priors in Bayesian machine learning are related to Gaussian measures.
'e allow for the degenerate case where the variance of (F, h) is zero. This means we interpret Gaussian with variance zero Dirac measure.

---

Let (22, A. P) be the underlying (physical) probability space and X' C R be measurable, The (product-) measurable mapping G X + R is called Gaussian process (GP) if and only for all N E N and all X = {In In C X' the random vector G(X) = (G(,,z1),....
G(,,NN) is multivariate Gaussian. For GP G we define mean function m(x) := E [[]())) a E X, and covariance function by k(x,, x²) == C[G(x), G(a')] for I, x' E X. Here E denotes the expected value and + ] the covariance. follows from the definition that G(X) N (m((), (X, X)) for any J C X, where we define m(X) == (m((,))), N and k(X, == (k(Zn, N n,n'=1' We write G ~ GP k) for GP with mean function m and covariance function k. Note that by the properties of the covariance we know that {_X, X) is (symmetric) positive semi-definite matrix for all {_,] C X' and N E N. A function with this property is called kernel, terminology that we adopt henceforth. Kolmogorov's 's existence theorem (Billingsley 2008 Section 36] guarantees the existence of Gaussian process for any kernel k and any mean function m. The standard reference for Gaussian processes in machine learning is Rasmussen (2003]
The main advantage of Gaussian processes in specifying priors over a function space is that the kernel k allows to incorporate readily interpretable prior assumptions, such as smoothness periodicity. For example, choosing the squared exponential kernel [Rasmussen] (2033 implies that the unknown function is infinitely differentiable and that the correlation of the functional output higher the closer the inputs are.
In order insert the Gaussian process prior into our generalized loss in (4)
we need to know the probability measure that is associated to the Gaussian process. In general, we can associate more than one Gaussian measure with given Gaussian process. For example:
If the GP has continuous sample paths we can associate a Gaussian measure on the space E of continuous functions with it [Lifshits 2012, Example 2.4].
If the GP has -quare-integrall sample paths we can associate Gaussian measure on the Hilbert space of qquare-integrain functions with it (cf. Theorem D
These sample path properties can be guaranteed under additional assumptions on the kernel. The next theorem discusses one such kernel condition which guarantees the GP to have sample paths in the Hilbert space of square integrable functions, denoted T,2 R), with inner product (g. h)2 Sx ((x))(x) dp(x).
Theorem 1. Let F ~ GP(m, k) be a GP with mean m E L2 p. R) and kernel i such that
J k(z,z)
dplax) dp(x)) c0.
We call a kernel satisfying K) trace-class kernel. Then the mapping F: (2 p. TR) defined as F(w)= = F(w,-) is a Gaussian random element with mean m and covariance operator C given as
Cg(-) == {f -f (K(,22)) (')g(z)))))))))))))))))) dp(x')
(9)
for any E (2, Consequently p= P ~ N(m, C) is a Gaussian measure.
Proof. The fact that F as defined above is GRE follows immediately from Example 2.3.16 in Bogachev [1998]. The fact that m is its mean and C' as defined in is its covariance operator tollows from Fubini's theorem.
shall be noted that there is no need to appeal to GPs in order to justify the use of GMs. In fact, it has recently been demonstrated that variational inference for GPs can be formulated purely in terms of GMs [iild and Wynne 2021). In the following sections we will therefore deploy GMs without any reference GPs, but it of course always possible to think of them as the measures that correspond to GPs where the kernel satisfies an additional assumption such as (8)
4 Gaussian Wasserstein Inference in Function Spaces
This section describes how the Wasserstein distance between Gaussian measures be used to obtain a tractable optimization target for inference in function spaces. In the end, we discuss several parametrizations of GWI and introduce our main inference method the GWI-net.
5

---

4.1 Model description
Let (annya)))x )a) C X X be N N paired observations. We assume that X' C CR.DEN and further that y = R for regression and y {I.....J} for classification with J e N classes. We focus in our exposition here on the regression case but have given the relevant derivations for classification in Appendix A.6
GVI in function space minimises the generalized loss L We make the mild assumption that the unknown function square integrable with respect to the data distribution p on X' which means f E I, R)). The prior P := P is described by Gaussian measure with mean mp (X, RR) and covari- ance operator Cp described by a trace-class kernel k X x X + R which means it is given as (C) == fx k(z, xz) xf)f(x)) dp(x') for all f E L p, R). We assume Gaussian likelihood for y = (y),-... MN) given as p(yls) = II n=1 p(un 1 with
As pointed out in section 3.1
F(] -Eq log ply
+ D(Q IPF)
p(un f] = NN Nyn (xx)),
(10)
where (- μu, o²) denotes the pdf of a normal distribution with mean j E RR and variance o² > 0. This prior and likelihood are natural choices they mimic the standard formulation of Gaussian process regression. The variational approximation of the posterior is chosen to be another Gaussian measure @ with arbitrary mean mQ I.3 (2,, p, R) and arbitrary covariance operator Cq induced by trace-class kernel r: (C)f)(x) == fx r(z, x") ff' dp(x') for all f L2 R).
remains for us to select a dissimilarity measure DD. As already pointed out in the introduction we decide to use the Wasserstein distance W2 (a formal definition is given in Appendix A.3) This choice was guided by two considerations:
1. The Wasserstein metric was proven to be useful metric for probability distributions in machine learning applications [Arjovsky al.] [017]] Tran et al.] 2020]. Furthermore the Wasserstein metric is known to have desirable statistical properties [Panaretos and Zemel 2019).
2. The Wasserstein distance is tractable for arbitrary Gaussian measures on (separable) Hilbert spaces [Gelbrich, 1990] and given as
(P, = [mp mellă + tr(Cp) + tr((C)) - to
1/21
(11)
where tr denotes the trace of an operator and C is the square root of the positive, self- adjoint operator Cp. This is in stark contrast to the KL-divergence that is infinite whenever @ is not dominated by PF and even in the case where is finite there exists no explicit formula for the KL-divergence in infinite dimensions.
The generalized loss for our model is therefore given as
NV
L
Σ
EQ
log (yn F(zx)) + (2.QQQ)
(12)
1=1
Note that the expected log- likeliiid in 12 can be calculated analytically as
N
N
(yn mQ (2r))² + r(5n,
Eg log Mogg((s,)) ( ((xx), = - (og(2mo?)
(13)
202
==1
remains to produce an approximation of [IT) in order to obtain a tractable inference procedure. To this end, note that by definition [|mp melé = (mp(z) - m dp(x) and further tr(Cp) = f k(æ,, x) dp(x) Brislawn 1991]. We now replace the true input distribution with the empirical data distribution p== / Dmm Se,, where d, denotes the Dirac measure in x E X.
This gives I|mp- mQll x + EN (mp(z,)) m²(xx))) tr(Cp)) ~ NV Σm', k(zns un) and
Astute readers may notice that the definition of the likelihood contains a pointwise evaluation f(x.) which may not be well defined operation on L² R). We detail in Appendix 30 how that problem can be circumvented and that in fact F(x) (m((z)) k( z)) as one would expected.

---

tr(Co) = & E Tn). It remains to provide an approximation of tr [[ck
CoC 3 1/2
The key idea is to approximate the spectrum of Cp Cop by that of appropriate kernel matrix. Details are discussed in Appendix A.4 This leads to the following final approximation for the Wasserstein metric
W
(mp(Tn) mQ(en))
k(In Xn)
(14)
N
Ng
r r(an) - 2
v Ao (r(Xs
(15)
N
VNNS
n=1
E RD being subsampled from the in- put data X. Further r r(Xs. = (() and Xx), = (((Zn:Zs,s) for I....,NN Ns and da X) X)), denotes the s-th eigenvalue of the matrix r r(Xs, X)s(X,Xs) E RRNssNs The approximation quality of W is related to the spectral decay of the operator CpCq. which in turn determined by the kernels and r. For the choices made in Section 4.2 we empirically observe rapid spectral decay (cp. Appendix (A.13) and therefore are confident that the Wasserstein distance estimated reliably for our method.
where Xg == (0S.,....,,.NN..) with
== = N. I.....,...
The combination of (3). 14 and 15 gives generalized loss that is tractable in terms of mp, mQ; k. and r. we disregard computation time mp, mq; and r, the generalized loss can be evaluated in O(N N N + N where typically Ns < N, e.g. Ns 100. We provide a batch version of our loss in Appendix A.5 which reduces the computations to O (N + N where NB < N is the batch-size. Note, however, that the final computation time for our method will be determined by the complexity hidden in the evaluation of mQ, mp, k, and as we need Ng evaluations of mQ and mp and Ns NB evaluations of and per iteration.
4.2 Parameterisations of Prior and Variational Measure
The prior for our model is given P -N (m,,, Cp) with Cp induced by a trace-class kernel k. One of the advantages the proposed approach is that any trace-class kernel allowed and this is where one can incorporate specific assumptions and domain expertise. This is thoroughly studied topic: the prior kernel can encode periodicity (Durrande et al. 2016]. geometric intuition [van der Wilk et al. 2018), and even model linear constraints for the unknown function [Jidling et al. 2017]. In order to keep the exposition simple and maintain focus on the inference, however, and in line with using simple priors network weights in standard Bayesian deep learning, opt for simple zero mean prior m. = and a standard ARD kernel i given as
-x,x')=0f = of exp (-2
Σ
1
D
(Id Ra--x')² x'd)2
(16)
2
a
d=1
for z, x' E X R... We refer to o > as kernel scaling factor and to aa > 0 as length-scale for dimension d. The parameters oy and := (a1...., aD) are called prior hyperparameters.
The rest of the section explores various choices for the variational mean mQ and the variational kernel r. The parameters appearing in the specification of mQ and r are referred to variational parameters.
GWI: Stochastic variational Gaussian process X with M < N. We define the posterior mean
Let ,I...., E A' be a subsample of the data
2M
m²(#) == mp(x)) +
Bm km (2:)
(17)
m==
with 3m E IR and (x))x) == k(x, m. m = 1 1...., M where l is the prior kernel i and B (Br,... Bas) RR variational parameters. Define further the variational kernel
=
r(z,x') 2') - kz)
k(Z, kz - kz (x) + kz(x)* Ekz((x)
(18)
where Σ MxM is the symmetric and positive definite variational covariance matrix that parame- terises This choice of mg and essentially recovers the stochastic variational Gaussian processes
7

---

sinix) 0.1x
Y0d sinix) 0.1x²
-200
-23
-20
0.0
-100 -0.75 -2.50 -023 0.00
120
Figure 1:
Training data
Unseen data
Inducing points
We query the above functions N 1000 equidistant points and add white noise with (0, 0.5²). We use M = 30 inducing points and train our method as described in Appendix A.7 The plot shows (x) = 1.96 /VY" where V[Y*(x) is the posterior predictive variance given rxx, x) + o².
(SVGP) model of Titsias 2009). Note that in our framework it is straightforward to use all (or just more) basis functions for the mean mo(#) =- mp(x) + >>> K, (a) where (x) == k( z.), 3n E R, nl .,... N. This mirrors the construction in Cheng and Boots 2017 where we allow more parameters learn the mean than SVGP. However, both Titsias (2009 and Cheng and Boots 2017 use different objective function than GWI to learn the unknown parameters.
GWI: deep neural network with SVGP An interesting approach is to parameterise the posterior mean as a deep neural network (DNN). We assume the DNN has E hidden layers and the width of layer l 1,..., L is denoted D3 with Do = D and D
1. This means (x) = W =+ and further h² == (6g (x)), g g (z) = W (x)+b++1 ( for = 1,...,
define
Here W is D++1 x De matrix, b E Peia is a bias vector for layer I and á an activation function. We can then define the variational mean as m²(x)) == mp(x) + -x++(xx). If we choose the SVGP kernel r in (T8), we essentially predict with neural network and quantify uncertainty with a (sparse) Gaussian process, capturing the beneficial properties of both.
Neural networks have been combined in several ways with GPs [Wilson et al. 2016 Tran et al. 2020]. However, the best of our knowledge they were not used directly parametrize posterior in the context of generalized variational inference in function space. The spirit of our approach is fundamentally different: rather than thinking neural network as a model which needs be made Bayesian, we use it as parametrisation of a variational posterior.
We note that we do not here provide an exhaustive study how to best parameterize the variational measure. This paper is focused on demonstrating the ability of the proposed method to obtain valid uncertainty quantification. An exploratory study how properties and quality of uncertainty quantification relate to different choices of mQ and r is reserved for future work. We mention potential problems that can occur from misspecification in Appendix A..
5 Experiments
We show results for GWI with the SVGP mean [77 and the SVGP kernel (18). We use the shorthand GWI: SVGP for this approach. Additionally we implement the DNN mean with the SVGP kernel (18). This combination achieves impressive results on various regression and classification tasks. We call this method GWI: DNN-SVGP or simply GWI-net.
Illustrative Examples In Figure we illustrate GWI-net on few toy examples. One can clearly see that the posterior predictive variance expands for regions lacking observations which demonstrates the ability our method quantify uncertainty. We provide an additional graphic comparison with SVGP in Appendix 12 and an example for two-dimensional inputs in Appendix A.9
There show that the pathologies regarding the quantification of in-between uncertainty discussed in Foong et al. [2020] are not present for our method.
UCI Regression In Table T we report the average test negative log-likelihood (NLL) (ef. Appendix A.7 for details) of GWI: SVGP and GWI-net (GWI: DNN-SVGP) and the results of several weight- space approaches for BNNs: bayes-by-Backprrg (BBB) [Blundell et al. 2015), variational dropout

---

(VDO) IGal and Ghahramani, 2016], and variational alpha dropout (a = 0.5) [Li and Gal 2017]. We also compare with four function space BNN inference methods: functional variational inference with BNN prior (FVT) (M and Hernández-Lobant 2021], variationally implicit processes (VIP) with
BNNs, VIP-Neural processes |M et al. | 2019]. and functional BNNs (FBNNs) [Sun et al. 2019] In
order to ensure fair comparison we matched neural network architectures and training procedures for the different methods. Detailed explanations are given in Appendix A.7
Damaset
PV1
VDO
0.5
EXACT
BOSTON
506
28:00.11
2.77:3.06
233520.04
2,4550004
14550005
1.76E004
263-0010
24553002
1,3030.10
2.4620.04
CONCRRTE
1050
K
524+0009
2.6-0.06
2.58-0.06
102-4002
5.55+0.02
1.24-0011
5.25-0.11
3,06+0005
3,06+6001
326-0.00
ENERGY
708
5
1.811:.19
0.1110.12
0.58230.5
0.56:8.04
0.0020.05
2.172:002
1.13-0.02
0.9510.09
0,5810.02
0.54.0.02
KINENM
8192
-0.56-0.38
-1.2-0.8
-135+00.1
-1112-0m
-1.05-0000
-081+000
-088-0.01
-0220002
N/A+000
NA4+0000
POWER
0568
4
1.351:0:22
259.20002
2.9-10.00
29210000
250:.000
2.8310011
2.88
288100.0
N/A10000
NA4D00
PROTEIN
WINE
1558
0.97:0002
0.76:1.00
0:97:0.06
0.9760002
1.20:00044
1.01:0002
0:97-0.22
1.0120002
1.4-0001
0.36:0000
YACHI
106
6
2,5710055
0.29:10.1
0.59.0..11
4.002:0.07
0.59.00.5
1.1110004
1:22-00.0
0.2966011
1,001000
0:10:00,0
NAVAL
11934
15
-7.25-4088
.67640.11
2..1
-5.6260.04
-..1:00.00
-2.9 2.80000000
2,80-0000
-2.7700144
-.1340002
NAA40.00
"S5"""
2Jfh
277
'1:9
----7---
5.1
'001
Table I: The table shows the average test NLL on several UCI regression datasets. We train on
random 90% of the data and predict on 10%. This is repeated 10 times and we report mean and
standard deviation. The results for our competitors are taken from Ma and Hernaddez-
[202T]
One can see that GWI-net obtains the best mean rank of all methods being best model on 4/9 datasets and performing competitively on all datasets. Note that we exclude FBNN and exact Gaussian processes from the comparison because their computational complexity is often prohibitively large.
Classification and OOD Detection We demonstrate the ability of GWI to perform image classifi- cations on Fashion MNIST Xiao et al., 2017] and CIFAR-10 (Krizhessyy et al. 2009]. We compare to FVI, mean-field variational inference (MVFI) [Blundell et al., 2015|. maximum a posteriori ap- proximation (MAP), K-FAC Laplace-GNN [Martens and Grosse 2015] and its dampened version [Ritter et al. 2018). Implementation details are discussed A.8
We also assess the ability of our model to perform out-of-distribution detection using in-distribution (ID) / out of-distribution (OOD) pairs given as FastioonMNISTTMMIISTT and CIFAR 10/S Fol- lowing the setting os(aa al. 2019), [mmer et al. T
calculate the area under the curve (AUC) of a binary out-of-distribution classifier based on predictive entropies. Results are shown in
Table 2
FMNIST NLL
CIFAR 10
Model
Accuracy
OOD-AUC
Accuracy
NL.
OOD-AUC
GWI-net
93.25 10.09
0.250 10.00
0.959 10.01
10.00
0,553 10.00
0.618 10.00
FVI
91.60=t0.14
0.254++0005
0.956:10.06
77.69 +0.64
0.675++0.03
0.883+0.04
MFVI
91.20:10.10
0.343-10.01
0.78240.02
76.40:10.52
1.372:10.02
0.58910.01
MAP
91.39+-0.11
0.258+0.00
0.864+0.00
77.41+0.06
0.690+0.00
0.809+0.01
KFAC-LAPLACE
84.42++0.12
0.942++0.0
0.945++0.00
72.49++0.20
1.274++000
0.548+0001
RITTER al.
91.20=0.007
0.265=00000
0.947=00.00
00.0
0.661+0.00
0.796=#0.00
Table 2: We report average accuracy, NLL and OOD-AUC on test data for 10 different train/test splits.
The results for FVI obtained from Ma and Hernández-Lobato
[021 and for MAP, KFAC and
Ritter et al. results are taken from Immer et al. [2021]
Our method performs best in all categories on the Fashion MNIST dataset achieving state-of-the-an results. On CIFAR10 we obtain the highest accuracy and best NLL. by a significant margin and perform competitively in the OOD detection task.
Limitations
In this section we discuss some of the shortcomings and difficulties which are related to our method.
The GVI-FS framework allows the specification of function space inference via infinite dimensional parameters such as mean and kernel functions. This great flcxibility essentially allows the specification of mismatched prior and posterior parameters. We illustrate such case in Appendix AA.
GWI-net relies on the SVGP kernel defined in 18 for its posterior approximation. It therefore inherits numerical instabilities associated with the inversion of the kernel matrix. For the data sets discussed in this paper it was possible to overcome these issues by smart initialisation of the optimiser (cf. Appendix (A.7), but it may be an interesting research avenue to come up with a kernel that avoids these instabilities.
Our method approximates the Wasserstein distance in function space via the spectrum of kernel matrices (cf. Appendix A.4)) These approximations require quick spectral decay of the composition

---

of prior and variational covariance operator to be accurate and computationally tractable. The prior SE kernel combined with the variational SVGP kernel did have this property (cf. A.133) which allowed for cheap and accurate approximations. However, other parameterisations may result in less accurate estimation. theoretical investigation of how the approximation quality relates kernel properties is an interesting topic for further research.
The proposed framework models prior and variational distribution with Gaussian measure on the space of square integrable functions. As a consequence the posterior distribution for the functional output Gaussian as well. This means it is unimodal and concentrated around the posterior mean. Although this constrains the form of functional posterior significantly the authors would argue that the empirical success of GWI-net demonstrates that the approach is flexible meaningfully quantify uncertainty.
7 Conclusion
In this paper, developed framework for generalized variational inference in function spaces. We leveraged the function space perspective to develop a new inference approach combining Gaussian measures and Wasserstein distance with predictive performance of deep neural networks, yielding principled uncertainty quantification. The value of our method was demonstrated several benchmark datasets.
10

---

References
B. Adlam, J. Snoek, and S. L. Smith. Cold posteriors and aleatoric uncertainty. arXiv preprint arXiv:2008. 00029, 2020.
L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
P. Billingsley. Probability and measure. John Wiley & Sons, 2008.
C. Blundell, J. Cormebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural network.
In International conference on machine learning, pages 1613-1622. PMLR, 2015.
Bogachev. Gaussian Measures. American Mathematical Society, 1998.
C. Brislawn. Traceable integral kernels on countably generated measure spaces. Pacific Journal of Mathematics, 150(2):229-244 1991.
D. R. Burt, S. W. Ober, A. Garriga-Alonso, and M. der Wilk. Understanding variational inference in function-space. arXiv preprint arXiv:201 1.09421, 2020.
T. Chen, F. Fox, and C. Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pages 1683-1691. PMLR, 2014.
C.AA. Cheng and B. Boots. Variational inference for gaussian process models with linear complexity. Advances in Neural Information Processing Systems, 30, 2017.
G. Da Prato and J. Zabezyk. Stochastic equations in infinite dimensions. Cambridge university press, 2014.
D' Angelo and Fortuin. Repulsive deep ensembles are bayesian. Advances in Neural Information Processing Systems, 34:3451-3465 2021.
F. D' Angelo, V. Fortuin, and F. Wenzel. On stein variational neural network ensembles. arXiv preprint arXiv:2 106. 10760, 2021.
N. Durrande, J. Hensman, M. Rattray, and N. D. Lawrence. Detecting periodicities with gaussian processes. PeerJ Computer Science, 2:e50, 2016.
V. Dutordoir, N. Durrande, and J. Hensman. Sparse gaussian processes with spherical harmonic features. In International Conference on Machine Learning, pages 2793-2802. PMLR, 2020.
D. Duvenaud. The kernel cookbook: Advice on covariance functions. URL https://www. cs. toronto. elu/duvenaud/conthook; 2014.
A. Foong, D. Burt, Li, and R. Turner. On the expressiveness of approximate inference in bayesian neural networks. Advances in Neural Information Processing Systems, 33: 33:15897-15908, 2020.
V. Fortuin, A. Garriga-Alonso, F. Wenzel, G. Rätsch, R. Turner, M. der Wilk, and L. Aitchison. Bayesian neural network priors revisited. arXiv preprint arXiv:2102.06571 2021.
T. Fushiki. Bootstrap prediction and bayesian prediction under misspecified models. Bernoulli, 11 (4):747-758, 2005.
Y. Gal and Z. Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference machine learning, pages 1050-1059, PMLR, 2016.
J. Gardner, G. Pleiss, K. Q Weinberger, D. Bindel, and A. G. Wilson. Gpytorch: Blackbox matrix- matrix gaussian process inference with gpu acceleration. Advances neural procoessnnnnooosssnnnn systems, 31, 2018.
processing

---

M. Gelbrich. On formula for the 12 wasserstein metric between measures on euclidean and hilbert spaces. Mathematische Nachrichten, 147, 1990.
Ghosal and A. Van der Vaart, Fundamentals of nonparametric Bayesian inference, volume 44. Cambridge University Press, 2017.
1. Goodfellow, Y. Bengio, and A. Courville. Deep learning. MIT press, 2016.
Hensman, N. Fusi, and N. D. Lawrence. Gaussian processes for big data. arXiv preprint arXiv:1309.6835, 2013.
J. Hensman, N. Durrande, A. Solin, et al. Variational fourier features for gaussian processes. J. Mach. Learn. Res., 18(1):5537-5588, 2017.
G. E. Hinton and D. Van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the sixth annual conference Computational learning theory, pages 5-13, 1993.
M. Hladnik and M. Omladič. Spectrum of the product of operators. Proceedings of the American Mathematical Society, 10202)300-332,, 1988.
B. Hunt, T. Sauer, and A. Yorke. Prevalence: translation-invariant "almost every" on infinite- dimensional spaces. Bulletin of the American mathematical society, 27(2):217-238, 1992.
A. Immer, M. Korzepa, and M. Bauer. Improving predictions of bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics, pages 703-711. PMLR, 2021.
C. Jidling, N. Wahlström, A. Wills, and T. B. Schön. Linearly constrained gaussian processes. Advances in Neural Information Processing Systems, 30, 2017.
L. V. Kantorovich. Mathematical methods of organizing and planning production. Management science, 6(4):366 422, 1960.
A. Kendall and Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
M. Khan, D. Nielsen, V. Tangkaratt, W. Lin, Y. Gal, and A. Srivastava. Fast and scalable bayesian deep learning by weight-perturbatinn in adam. In International Conference on Machine Learning, pages 2-12-2. PMLR, 2018.
Knoblauch, J. Jewson, and T. Damoulas. Generalized variational inference: Three arguments for deriving new posteriors. arXiv preprint SKKiv:1904.0206.3 2019.
F. Komaki. On asymptotic properties of predictive distributions. Biometrika, 83(2):299-313, 1996.
A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
A. Kukush. Gaussian measures in Hilbert space: construction and properties. John Wiley & Sons, 2020.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances neural information processing systems, 30, 2017.
Y. Li and Y. Gal. Dropout inference bayesian neural networks with alpha-divergences. In International conference machine learning, pages 2052-2061 PMLR, 2017.
Y. Li and R. E. Turner Gradient estimators for implicit models. arXiv preprint arXiv: 2017.
M. Lifshits. Lectures on gaussian processes. In Lectures on Gaussian Processes, pages 1-117. Springer, 2012.
C. Ma and J. M. Hernández-Lobann Functional variational inference based on stochastic process generators. Advances in Neural Information Processing Systems, 34, 2021.
12

---

C. Ma, Y. Li, and J. M. Hernández-Lobato. Variational implicit processes. In International Conference on Machine Learning. pages 4222-422244 4233. PMLR, 2019.
J. Maddox, P. Izmailov, T. Garipov, D. P. Vetrov, and A. G. Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32. 2019.
J. Martens and R. Grosse. Optimizing neural networks with kroneccer-factored approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
A. G. G. Matthews. Scalable Gaussian process inference using variational methods. PhD thesis, University of Cambridge, 2017.
A. G. G. Matthews, J. Hensman, R. Turner, and Z. Ghahramani. On sparse variational methods and the kullback-leible divergence between stochastic processes. In Artificial Intelligence and Statistics, pages 231 1-19.9. PMLR, 2016.
A. G. d. G. Matthews, M. Rowland, J. Hron, R. E. Turner, and Z. Ghahramani. Gaussian process behaviour wide deep neural networks. arXiv preprint arXivv:804.1 1201, 2018.
R. M. Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
K. Osawa, S. Swaroop, M. E. E. Khan, A. Jain, R. Eschenhagen, R. E. Turner, and R. Yokota. Practical deep learning with bayesian principles. Advances neural information processing systems, 32, 2019.
V. M. Panaretos and Y. Zemel. Statistical aspects of wasserstein distances. Annual review statistics and its application, 6:405-431, 2019.
R. Ramamoorthi, K. Sriram, and R. Martin. On posterior concentration in misspecified models. Bayesian Analysis, 10(4):759-789, 2015.
C. E. Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, pages 63-71. Springer, 2003.
H. Ritter, A. Botev, and D. Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018- 2818-Conference Track Proceedings, volume 6. International Conference Representation Learning, 2018.
G. Rudner, Z. Chen, and Y. Gal. Rethinking function-space variational inference in bayesian neural networks. In Third Symposium on Advances in Approximate Bayesian Inference, 2020.
H. Salimbeni, C.-A. Cheng, Boots, and M. Deisenroth. Orthogonally decoupled variational gaussian processes. Advances in neural information processing systems, 31, 2018.
Schneider, L. Balles, and Hennig. Deepobs: A deep learning optimizer benchmark suite. arXiv preprint arXiv:/903.054900 2019.
Shi, Sun, and J. Zhu. spectral approach gradient estimation for implicit distributions. In International Conference on Machine Learning, pages 4644-4653. PMLR, 2018
N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1): 1929-1958, 2014.
S. Sun, G. Zhang, J. Shi, and R. Grosse. Functional variational bayesian neural networks. arXiv preprint arXiv: arXiv:/903.0 09039993707777 2019.
M. Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial intelligence and statistics, pages 567-574. PMLR, 2009.
B.-H. Tran, S. Rossi, Milios, and M. Filippone. All you need good functional prior for bayesian deep learning. arXiv preprint arXiv:2011.12829, 2020.
13

---

M. Van der Wilk, C. E. Rasmussen, and J. Hensman. Convolutional gaussian processes. Advances in Neural Information Processing Systems, 30, 2017.
M. van der Wilk, M. Bauer, S. John, and Hensman. Learning invariances using the marginal likelihood. Advances in Neural Information Processing Systems, 31, 2018.
K. Wang, G. Pleiss, J. Gardner, S. Tyree, K. Q. Weinberger, and A. G. Wilson. Exact gaussian processes on a million data points. Advances in Neural Information Processing Systems, 32, 2019.
M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688. Citeseer, 2011.
F. Wenzel, K. Roth, B. S. Veeling, J. Swiątkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, and S. Nowozin. How good is the bayes posterior deep neural networks really? arXiv preprint arXiv:2002.02405, 2020.
V. Wild and G. Wynne. Variational gaussian processes: A functional analysis view. arXiv preprint 1r.19:210..12798, 2021.
A. G. Wilson and P. Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-4708, 2020.
A. G. Wilson, Z. Hu, R. Salakhutdinov, and E. P. Xing. Deep kernel learning. In Artificial intelligence and statistics, pages 370-378. PMLR, 2016.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:/708.07740 2017.
14

---

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yss] [No], or [N/A] You are strongly encouraged to include a justification to your answer, either by referencing the appropriate section of your paper or providing brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section ??.
Did you include the license to the code and datasets? [No] The code and the data are proprietary.
Did you include the license to the code and datasets? [N/]]
Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Appendix AA.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Nothing discuss
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proof' of all theoretical results? [Yes] See Appendix A.1-A.6
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (cither in the supplemental material or as a URL)? [Yes] See footnote in introduction
(b) Did you specify all the training details (e.g., g.g data splits, hyperparameters, how they were chosen)? [Yes] See Appendix A.7 and A.8
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix [A.IT
4. If you using cxisting assets (e.g., code, data, models) or curating/releusing
assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "d5sPuG_X0k", "reviewer": "Reviewer_N2sS", "summary": "The paper proposes a new probabilistic regression and classification framework called \"generalized variational inference in function space\" (GVI-FS), which involves taking the standard function-space ELBO formulation and regularizing $\\mathbb{E}_\\mathbb{Q} [\\log p(y|F)]$ with the Wasserstein-2 metric, $\\mathcal{W}_2(\\mathbb{Q}^F, \\mathbb{P}^F)$, instead of the standard KL-divergence, $KL(\\mathbb{Q}^F \\left| \\right| \\mathbb{P}^F)$. This substitution is motivated by the generalized variational inference framework of Knoblauch et al. (2019) and the difficulties of working with KL divergences. The paper proposes to use Gaussian measures (GMs) to parameterize the predictive functions, although GPs with additional assumptions provide an equivalent approach. The paper proposes two specific variants of this model, a model in which both the predictive mean and the predictive covariance are parameterized by a sparse variational GP and another in which the predictive mean is parameterized by a neural network and the predictive covariance is parameterized by a sparse variational GP. The paper then presents results from regression, classification, and predictive variance-based out-of-distribution detection, comparing to existing probabilistic methods.", "questions": "1. In the experiments section, the performance numbers for baseline models are taken from previous work. Are the model architectures and training procedures for the proposed methods comparable? I didn't notice any notes on this issue.\n\n2. An approximate Wasserstein-2 metric is given (Eq. 15). How good is this approximation in practice? In what situations is it more or less accurate. Some content along these lines would help readers understand the pros and cons of method.\n\n3. (minor) It would be interesting to rely on the GP to influence both the predictive mean and uncertainty in the DNN-SVGP case. This could be achieved by using the NN to output \"pseudo-observations\" with uncertainty, which are then conditioned using the GP mean and covariance to produce a predictive mean.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1. The work is technically sound.\n\n2. The work tackles problems of broad interest: classification, regression, and OOD detection using function-space methods. \n\n3. The work demonstrates substantially improved performance compared to the baseline models tested.\n\n\nWeaknesses\n\n1. The motivation for the proposed approach in terms of \"generalized\" variational inference is very weak. Specifically, I find the generalized variational inference framework of Knoblauch et al. (2019) to be so general as to be mostly divorced from the intentions and benefits of Bayesian inference, other than being an explicitly probabilistic method. Here are two alternative interpretations/motivations I find more convincing.\n\n1a) The proposed method is GP regression/classification with Wasserstein-2 regularization. This is very simple and, at least to a first approximation, accurate.\n\n1b) Another connection relates to Bayesian inference more directly. Here is a sketch. Consider the objective\n\n$$\\max_Q \\mathbb{E}_{F \\sim Q} [\\log p(y|F) - \\log q(F)]  - \\lambda \\mathcal{W}_2(Q,P)$$\n\nwhere $\\lambda \\in \\mathbb{R}^+$. This is a close variant of Eq. 12, only differing in the $ \\mathbb{E} [- \\log q(F)]$ entropy term and the $lambda$ multiplier. With an unrestricted variational family, the optimal approximate posterior in the unregularized $\\lambda=0$ case is $q(F) = p(F|y)$, the true posterior.\n\nConsider Langevin diffusion for this unregularized case. We start with a single particle $F_0$ and update in continuous time $t \\geq 0$ according to\n\n$$\\text{d}F_t = -\\nabla V(F_t) \\text{d}t + \\sqrt{2} \\text{d}B_t $$\n\nwhere $V(\\cdot) = - \\log p(y|\\cdot)$ and $B_t$ denotes Brownian motion. The stationary distribution is $\\propto p(F|y) \\propto \\exp(-V)$, the true posterior. The work of Jordan, Kinderlehrer, & Otto (1998) shows that the marginal law of this diffusion process obeys gradient flow of the functional $KL(\\cdot,\\pi)$ with respect to the Wasserstein-2 metric.\n\nThis implies the resulting marginal law of $F_t$ follows gradient flow of $KL(q(F), p(F|y)) = \\mathbb{E}_{F \\sim Q} [\\log p(y|F) - \\log q(F)]$, the first term in the original objective.\n\nNow, back to the original objective, we can interpret $\\lambda$ as a Lagrange multiplier, giving the dual formulation\n\n$$\\max_Q \\mathbb{E}_{F \\sim Q} [\\log p(y|F) - \\log q(F)] \\quad \\text{s.t. } \\mathcal{W}_2(Q,P) \\leq C$$\n\nfor some $C = C(\\lambda) > 0$. this gives the original objective a nice interpretation. If we start the Langevin diffusion process with our variational posterior equal to the prior, $Q_0(F) = P(F)$ and stop the process once the diffusion process has covered a distance of $C$ w.r.t the Wasserstein-2 metric, then we will have approximately optimized our objective. As $C \\leftarrow \\infty$ (corresponding to $\\lambda \\leftarrow 0$), the approximate posterior $Q(F)$ approaches $p(F|y)$, the solution to the full Bayesian inference problem. This objective can then be interpreted as Bayesian inference with early stopping.\n\nOne wrinkle is that $Q(F)$ is restricted to be a GP. See the recent work of Lambert et al. (2022), \"Variational inference via Wasserstein gradient flows\" for a good discussion of this issue.\n\n2. The experiments section should contain an additional comparison with a DNN with the same architecture and objective as DNN-SVGP, but with both predictive mean and uncertainty output by the network. This will give the reader a sense of how much the uncertainty handling by the sparse GP is helping performance. An additional comparison without the Wasserstein-2 regularization would also be informative. Additionally, a sparse GP method in addition to the last column in Table 1 (Exact GP) would be helpful for interpreting how good these results are.\n\n3. (minor) l.296-8, VAEs use neural networks to parameterize variational posteriors. I would not consider this approach fundamentally different.\n\n4. (minor) I don't believe $\\Sigma$ in Eq. 18 is defined.\n\n5. (minor) The supplied code is almost completely uncommented, making it impractical read.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "b74DgukvHAA", "reviewer": "Reviewer_gB9j", "summary": "The paper presents a new methodology to apply variational inference in function space. Starting from the classic ELBO, the authors replace the KL divergence for measures with the Wasserstein distance and, by choosing Gaussian measures for the prior and the approximate posterior, the authors show how to approximate this quantity to have a tractable loss to optimize. Also, the authors present two ways to parameterize the variational posterior, one related to sparse Gaussian processes and the second based on deep neural networks. Finally, the paper is concluded with some experiments, including 1D toy regression, the classic UCI regression benchmark and some image classification problems (with OOD detection).", "questions": "- For me, the example shown in Figure 1 is not very illustrative. Fig 1 shows only one of the proposed models with 3 different function estimation problems. For example, it would be interesting to see the comparison between GWI-SVGP and GWI-DNN-SVGP and SVGP alone (SVGP is particularly interesting to see, given that it should be the same model as GWI-SVGP trained on a different loss)\n- How are you treating the prior parameters? Are they fixed, cross-validated or optimized? And what about the inducing points? Are these just a sub-sample of the training set or are they optimized?", "limitations": "Some additional comments on the limitations would have been welcomed (especially regarding the practical implementation). \nNo comments on potential negative societal impact, as expected for this kind of work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n- The paper is technically very sound and promises to improve upon standard methods for weight-space and function-space inference of neural networks (and not only).\n- The paper is generally easy to follow and it is sufficiently self-contained \n\n## Weaknesses\n\n- The original assumption of working with Gaussian posteriors could be a limiting factor, even if we are working in function-space. For example, a recent work [Pleiss and Cunningham, 2021] on the behaviour of wide parametric (and non) models seems to suggest that the limiting Gaussian posterior performs worse than the BNN/DGP posterior. This is common for all variational inference methods in function space with Gaussian posterior, but I would still appreciate a comment from the Authors on this. \n- The experimental evaluation is a bit limited (with nonetheless interesting results). My only concern is on the comparison with other methods, given that those numbers have been copied from various papers. Did you use the same models, same architectures, and same setup?\n\n\n### Minor\n- The reference section can be improved (e.g. miss-capitalizations and arxiv citations when proceedings are available).\n- In the experiment, the second parameterization is sometimes referred to as GWI-DNN-SVGP and other times as GWI-Net. Maybe a uniform notation is easier to follow.\n\nGeoff Pleiss, John Patrick Cunningham. The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective. NeurIPS 2021", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "59QwKQ5fne", "reviewer": "Reviewer_6Xua", "summary": "The authors propose to use the Wasserstein distance between Gaussian measures on infinite-dimensional function spaces to perform generalized variational inference over distributions over functions, with applications to Gaussian processes and Bayesian neural networks.", "questions": "Major:\n\n- If we can think of the $X_S$ like inducing points in sparse GPs, it seems like their selection should impact the performance and the estimated distances in the proposed method. Could you comment on how they are chosen? Is there any theoretical guidance on what would be a good choice?\n- In the proposed GWInet, it seems that the model ultimately just uses a deterministic NN for the mean function and then wraps a normal GP around that. This seems like a rather complicated method to ultimately \"just\" perform GP inference. How would one use the proposed method to perform actual BNN inference in function space (that is, using a neural network with stochastic weights), so that one could actually use the learned features to capture the uncertainties instead of having to rely on a GP kernel? This seems crucial, since BNNs outperform GPs in many practical applications.\n- State-of-the-art performance on FashionMNIST is 97% and on CIFAR10 around 95%, so the reported performances in Tab. 2 all seem suspiciously low. Could you comment on why that is? I think using baselines that get closer to the performances we actually see in practice would make the comparison stronger and more trustworthy.\n\nMinor:\n\n- l. 75: Regarding variational BNNs in function space, functional SVGD [1] and functional repulsive ensembles [2] seem to be relevant methods to mention.\n- Tab. 1: I understand that both SVGP and DNN-SVGP are fit using GWI, right? If so, it seems like a comparison against a standard SVGP implementation would be interesting to see.\n- Overall, it's unclear to me whether the proposed method is faster or slower than the baselines. Could you report runtimes (at least rough estimates)?\n\n\n[1] https://arxiv.org/abs/2106.10760\n\n[2] https://proceedings.neurips.cc/paper/2021/hash/1c63926ebcabda26b5cdb31b5cc91efb-Abstract.html", "limitations": "There seems to be no negative societal impact, but the authors could do a slightly better job at addressing the technical limitations, especially regarding runtimes, performance, and how to integrate BNNs (see above).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well written and is mathematically rigorous.\n- The idea of using Wasserstein inference in function spaces in this way seems novel.\n- It is plausible that the proposed method might have advantages over standard variational GPs.\n\nWeaknesses:\n\n- The performance in the experiments seems suspiciously low and a few ablation studies seem to be missing.\n- Given that the title advertises \"Bayesian Deep Learning\", it is unclear how one would exactly apply this to Bayesian neural networks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept