Functional Variational Inference based on Stochastic Process Generators
Authors: Chao Ma
José Miguel Hernández-Lobato
Venue: NeurIPS 2021 Poster (2021)
Keywords: functional variational inference
Bayesian neural networks
stochastic processes

Abstract: Bayesian inference in the space of functions has been an important topic for Bayesian modeling in the past. In this paper, we propose a new solution to this problem called Functional Variational Inference (FVI). In FVI, we minimize a divergence in function space between the variational distribution and the posterior process. This is done by using as functional variational family a new class of flexible distributions called Stochastic Process Generators (SPGs), which are cleverly designed so that the functional ELBO can be estimated efficiently using analytic solutions and mini-batch sampling. FVI can be applied to stochastic process priors when random function samples from those priors are available. Our experiments show that FVI consistently outperforms weight-space and function space VI methods on several tasks, which validates the effectiveness of our approach.

Functional Variational Inference based on Stochastic Process Generators
Chao Ma University of Cambridge Cambridge, UK cm905@cam. ac uk
José Miguel Hernández-Lobato University of Cambridge Cambridge, UK
j jmh233@cam.c ac uk
Abstract
Bayesian inference in the space of functions has been an important topic for Bayesian modeling the past. this paper, we propose a new solution to this problem called Functional Variational Inference (FVI). In FVI, we minimize a divergence function space between the variational distribution and the poste- rior process. This is done by using as functional variational family new class fiexible distributions called Stochastic Process Generators (SPGs), which are cleverly designed so that the functional ELBO can be estimated efficiently using an- alytic solutions and mini-batch sampling. FVI can be applied to stochastic process priors when random function samples from those priors available. Our experi- ments show that FVI consistently outperforms weight-space and function space VI methods several tasks, which validates the effectiveness of our approach.
1 Introduction
As an important approach to Bayesian deep learning, Bayesian neural networks (BNNs) have been proposed and studied for decades [29, 34, 19]. Despite some successful applications in specific cases [55, 44, 23). BNNs still suffer from several shortcomings. As over- wree-parameterized models, BNNs often have multiple posterior modes in its weight space, that generate identical predictions [30]. Therefore, Bayesian inference in weight space can be very difficult [50]. Furthermore, due to the complexity of interactions between weights the network forward pass, it is unclear the effect that given prior distribution over the weights will have in the resulting distribution over functions.
To address these issues, there has been recent resurgence of interest in applying the perspective of Bayesian non-pan con-paamettiis to neural nets. That is, BNNs are treated as probability measures over func- tions, i.c., as stochastic processes, with Bayesian inference being performed now in function space.e Examples include functional BNNs (f- (f-BNNs) [47], and variational implicit processes (VIPs) [28]. Despite their empirical advantage over weight space VI, they still suffer from number of issues. i), f-BNNs optimize functional KI. divergence between BNNs and GPs, which is not always well-defined(5). Also, f-BNNs rely (spectral) stein gradient estimators, which are less efficient for high dimensional distributions [56. 14]. While VIPs do not have this problem, they use wake-slcep single unified objective. ii), F-BNNs rely mean-field
procedure that does not correspond approximation, which often lacks predictive in-between uncertainty on test data. This issue was ob- served both with single layer BNNs [10], and deeper BNNs (able to represent in- n--between uncertainty, but more over-confident than HMC empirically for certain settings [11, 9]), On the other hand, VIPs use as variational family a Gaussian process, which is not able capture non-Gaussian processes such structured implicit priors [47].
Therefore, it is open challenge how improve and justify (variational) inference in the space of functions using priors given by stochastic processes. In this paper, we investigate this old but important problem and propose new solution. Our contributions are follows:
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

(a) functional BNN
(b) Mean field VI BNN
(c) HMC BNN
(d) Ours (FVI on BNN)
Figure I: regression task on a synthetic dataset (red crosses) from [10]. We plot predictive mean and uncertainties for each algorithm. This tasks is used to demonstrate the theoretical pathologies of weight-space VI for single-layer BNNs: there is no setting of the variational parameters that can model the in-between uncertainty between two data clusters. The functional BNNs [47] also have this problem, since mean-field BNNs are used as part of the model. On the contrary, our FVI method can produce sensible uncertainty estimates. See Appendix C.2 for more details.
We propose new objective function for variational inference in function space, as an alternative to functional KL divergence between stochastic prcoesses [47]. We show that this new objective is a valid divergence, and can avoid some of the problems that the functional KL divergence has.
We propose a new class of flexible variational distributions in function space, called stochastic process generators (SPGs). SPGs are non- non-Gausiian generalizations of the VIP family [28], and can help avoid the fundamental limitation of mean field Gaussians [10] used in BNNs and functional BNNs. theorem regarding the expressiveness of SPGs proved (Proposition 4). Based on SPGs, our proposed functional divergence between stochastic processes can be estimated efficiently using mini-batch sampling (Proposition 5, and 7), which achieves a significant speed-up up against the gradient estimator approach in [47].
We compare our methods against existing weight-space and function-spant space inference methods in several tasks. Our method consistently outperforms the baselines, and is much faster than f-BNN, which validates the effectiveness of our approach.
2 Backgrounds
Variational BNNs Lety g(x, w) be neural network, where is the input, y the output, and denotes the weights. To build a Bayesian neural network (BNN), we place prior p( (w) over w, and add an observational noise (e::0, o²) to the output. Then, a BNN is given by
logp(ulx)= = log / N g(x, ())))))
Given dataset D fx; Vi )/ the goal of Bayesian inference is to compute the posterior p(w] (D) o p(D(w) ((())))(w)) Since this intractable, we often resort to approximate inference methods such as variational inference (VI), e.g. Bayes by Backprop (BBB) [4], variational dropout [12], or other extensions [20, 26, 41]. Since the posterior p( (W)(D) ) usually highly multi-modal due to the non-identifiability of neural networks, VI often leads to unsatisfactory results, and suffers from over-parameterrizaton and other pathologies [10].
Variational Implicit Processes (VIPs) An alternative way of looking BNNs is to treat them as a prior distribution over function space, i.e., a stochastic process. Indeed, p(y|x) in Equation generates a random function g(,, w). depending on what w actually sampled from the prior. This idea is further explored by the variational implicit process (VIPs). Given a prior p(f) on random functions, VIP approximates the posterior in function space p(f(D) by using a Bayesian linear regression model (f (D) = E. a,Q., a ~ N(a: μ4, E), where the basis functions {Ø. 1 are random samples drawn directly from p(f). Although VIP has shown improved performance over parameter-spe space VI methods. it has a few limitations. For example, (vir(f(D) resembles GP approximation to (D) whereas the true posterior in function space might be arbitrarily complex. Also, the wake-sleep training procedure of VIP does not optimize any coherent objective function. Lastly, VIP requires the prior p(J) to reparameterizable. i.e., can be represented the form of fx) = go (x, z), z ~ p(z). This may limit its application non-teparameteriiaale priors.

---

3 Problem setting and the functional KL divergence
We consider the problem of Bayesian inference in function space. Let p(f) f) be stochastic process defined on the probability space (92, B). Note that we use in its scalar form to denote a scalar function f(.)= 7 R. Here 7 is the index set of p(f) (assumed to be a compact subset of R2), For example, p(f) could be a Gaussian process p(f) GP((m))))) :)) a Bayesian neural network, or any other suitable stochastic process. We use pff) to model the uncertainty in function space. Then, a likelihood function ps (v))(-)) is defined on top of f to generate observable data y.
Given observed data D = {x, 3/ } our goal is to infer the posterior process P(IlD) conditioned on the observations D. If pff) is GP, then p(f(D) can be computed analytically. However, in most cases this intractable. Therefore, following [32, 47], we define another stochastic process q(f) on (22, B) our variational family approximate p (D) q(f) can be optimized by maximizing the evidence lower bound (ELBO) in function space:
L fffunctional == B P- -
Note that DKL (f) is the functional KL-divergence between stochastic processes. As shown by [47], DEL (l)((() can be written as:
= sup DKI(a(P*) M1p((*x))) n,X,,
where x, denote a set of n measure points {x*) in the domain/index set of f()), which can be treated as an element of the product space T Moreover, px.. denotes the vector of function values evaluated on X,, and D (g(xx, ((f=>>ll [[]][x is the KL-divergence over random vectors typically used by the machine learning community. In other words, the KL-divergence between stochastic processes is the supreme of the relative entropies obtained on all possible measure points in T
Functional Variational Inference using Stochastic Process Generators
4.1 The grid-functional KL divergence
As pointed out by [5], the functional KL diver- gence defined in Eq is not always well-defined. For example, the functional KL divergence be- tween two BNNs with different network architec- tures can be infinite. To address this issue, we propose to optimize a new divergence measure called the grid-functional KL divergence. Figure 2 illustrates the idea of grid functional KL diver- gence, which we now proceed to describe. Instead of taking the supremum as in the original func- tional KL divergence, we take expectation over both n and X We specify a probability distri- bution {pn 1inc<co over assigning lower prob- ability values Ps. to larger n. This way, we can trim down the contributions from the KL terms D (Ip((X that correspond to large n, and hopefully arrive at a finite expectation value. Following this idea, we give the formal definition of grid-functional KL divergence
acn
((x2)
f(x))
p(
((xz)
((x2)
f(r
X1
FF
Ps
Pa
D gri Eyoope
De I 14( (P
00
Figure 2: Illustration of Dgrid
D (grialg(ff))ln (())((())) = En.x, DKLLq( (Ip(PX.
(4)
where X,, is set of measurement points ( sampled from T, according some sampling distribution c. Note that e is defined on the product space T and the number of sampled measure points is also random. One may recognize that [47] proposed similar objective approximation to 3 in which the number of measure points n is a fixed constant instead of random variable. However, as we show below, it is in fact critical that n is not fixed.
Proposition 1. Suppose has full support on T Then, D Doria ((ll((() satisfies the following conditions: i). D (flp((I 2 0; ii). D Doria ((siig())lp(( =0 if and only ifa(f) pD)
3

---

In other words, if è has full support on 7 (we will give an example of such later), then D (enulg(f))lp((((1))) is a valid divergence in function space. Therefore, we can use D grid as an alternative objective for function space inference. We can show that, for certain scenarios, D grid can avoid some of the issues the original functional KI. divergence has (see Appendix A.2 for details):
Proposition 2. Let p(f) and q(f) be two distributions for random functions with different parametric forms. Assume that p(f) is parameterized by the following sampling processes:
f f= g + gx) ~ p(g]x: 0), θ e
and q(f) is parameterized by:
J = y + e, g(x( ( ((T),
Here, Xx XETC C. Fd is the input variable of the f. 9 is the random latent function, and θ E R', Te R are the parameters of each random function, respectively. Suppose p(g)x: x; O), q(g) x: θ), p()) and q(I") all have compact supports w.r.. 9. 9. θ and I, respectively. Then, there exist a sampling distribution such that: 1)e has full support on T .andd 2) D (grraa(())lpp()) is finite.
This result shows that the prid-functional KL divergence allows us perform VI between p and q even if they have different parametric forms. In Appendix A.2, we further show that the compactness assumptions can be relaxed under certain assumptions. Moreover, in Appendix A.3 Corollary 1, we have shown that if one of the distributions is replaced by a Gaussian process (of certain kernel function), then under some additional assumptions, the grid-functional KL is still finite. In those cases, the original functional KI. is no longer finite. This validates our choice of grid-functional KL divergence. To use D (f)lal()))l(((())))) for VI, we further derive new ELBO based on D Porid (Appendix A.I):
Proposition 3. Lel n. Xn c be a of random measure points such that X C Xn. Define
Carid logp(D) Porrala(f)lll(f))
(5)
Then we have
L = (oggl(() - D (grral(())ll((()
(6)
and log p(D) 2 Carid 2 C Fumctiional
Proposition shows that Carid is a valid variational objective function: it is a lower bound for log)(D) and also upper-bounds f C.fumetional, For the rest of the paper, we will discuss how to perform functional VI based on Lerid, We will focus on how to propose a expressive variational family q(f) and how to efficiently estimate D Dgrialg(f)ll) (l)l(()(
Remark. One example of that satisfies the requirement of Propositions 1, 2, and 3 takes the following form (which will be used throughout the paper):
(n PP) ~ Geom(p), xx U(T), VI VI < n- 1 Xn := XD Uixphsess-1p (7)
where first sample from geometric distribution such that (n- IDI) Geom(p) with parameter p (see Appendix A.2 for more discussion). Then, (n- (D)) out of distribution (OOD) measure points sampled independently from uniform distribution T.
4.2 Choosing q(f): stochastic process generators
In order to obtain good performance in FVI, it is crucial pick an expressive variational family for q(f). Here, we propose a new class of variational distributions called stochastic process generator (SPG). SPGs can be seen as the Gaussian extension of the variational approximation used in VIPs. Recall that in VIPs, the variational family (f) is defined by the following sampling process:
Σ=>on aa)
Ni E);
where { are random paths sampled from the prior process p(f). In SPGs, we remove the Gaussian assumption a. This is done by specifying a separate variational auto-encoder to represent the non-Gaussian distribution for a:
(aa) = po(a(b)4,

---

where h is the latent variable of the VAE, (a(a(h) is the decoder parameterized by 8, and 9n (h) is distribution over the latent space. Furthermore, we remove the constraint that { [ need to be sampled from p(f). For example, we may assume that each Φ, : IR + R is a DNN with learnable weights wa. We can finally give the definition of our SPG variational family qsea (fl9,()))
f f= Σ (aoo(:.we) +u, a po(al (h), u ~ (P(,,,,0(,))))))
(10)
where v is white noise process that models additive aleatoric uncertainty. Intuitively, QSPG serves as a generator for stochastic processes. (spa(fl-) maps any given 9,, (h) to stochastic process QSPG qspG(fl (fll)1(()))) hence the name. Regarding the expressiveness SPGs, we have the following result: Proposition 4 (Expressiveness of SPGs). Let p(f) be probability space (X, B), and its index set T is compuct subset of Rd Suppose X is compact metric space, the Borel set on X. Then, for VVe > 0, there exists a SPG (sra()) such that:
iiturr--ttegrable stochastic process
MMD(p, dsPGi F) < for VX E
(11)
where MMD) is the maximum mean discrepancy, F is the MMD function class defined to be a unit ball in RKHS with a universal kernel [46] k(.,-) its reproducing kernel.
Next we will discuss how { and gn (h) can be estimated and how can we use SPGs to estimate the function space KL-divergence in Equation 3.
4.3 Efficient estimation of grid-functione KI.
In order to estimate the grid-functions KL-divergence in Equation 4, we propose to use a two-step method for scalable approximation. In the first step, we distill p(f)) by fitting stochastic process generator PsPG(J (f') to p(f). In the second step, we calculate the Kl--diergencc between (fpa(f) and PSPG (f) as the surrogate for the KL-divergence between Qspa(f) and p(f).
Distilling p(f) via a stochastic process generator. Assume that we can draw M random functions (2(),,(()))))))
from the prior process p(f). Let pspa( be an SPG given by
f f= -=aadss wo) +v, a pol)))) -~ GP(x:0,0(..)))
(12)
where (o) is a fixed standard normal distribution. We denote the above process by DSPG (f Ipo (h)). We can train (fP on fi, f2. Sm....... fMr by optimizing the aggregated ELBO on (f):
PSPG (fXo (h)
max (w.),0,)
Exo
Σ 272
log pspG (Pxp
(PXO)
max Exo Σ m
(b)E log
(13)
[w.) 3,8A
where fx are the function values of fm evaluated on Xo. Xo are |0| (D) measure points inde-
pendently sampled from the training set Xp {x, ) fx are the function values of fm evaluated on Xo. and q (h| is an encoder network that approximates the true posterior ĐSPG ( (h o)))
Product of Experts (PoE) encoder. When sampling Xo. since its size might vary each time, we would need to set up to 2d inference nets, one for each possible subsets. To overcome this issue, we adopt the Product of Experts encoder [53], a simple and flexible approach for such scenario given by
101
d (2// (X)) x (h) II (b)
(14)
where 7A (x)lfmxx;) x;) is an inference network representing the expert associated with the i-th measurement point. Now the can use single encoder T (h|PX0) ) to handle all the possible inputs pXo, In practice, we let " fm (x,), x;) to be a Gaussian expert that maps [fm (xx), x,] to factorized Gaussian in latent space. Since the product of Gaussian experts is still Gaussian, TA (x(b(XX)) is a Gaussian distribution whose statistics can be computed analytically.
Estimating the grid-functional KI.-diverennce given X,. In order to estimate the grid-functional KI. divergence between (fpG( and p(J), we first discuss how this divergence can be estimated on measurement points X i.c. D (pp((((x)))) }] where fx,, is the vector of function values
5

---

evaluated on x We then discuss how this can be used to estimate the prid-functional divergence in Equation 4. To begin with, as in Section 4.2, our variational family is given by
S Σ asda(s w.) + us a~ g(a) = (a
-~ -~gP(v;D.o.) (15)
We denote the above variational family by qsva (f 4n (h)). The key ingredient of our estimation method is that we force (fl9, (h)) and Psra (flp)(h) to share the same basis functions (or weights { (w,)) ) and decoder parameters O. That is, once optimal {w,) and 0 are obtained by fitting (J) to p(J), these are frozen and reused in QSPG (J'. This makes sense, since according to our definition in Section (fpa( and ĀSPG (f) share the same measurable space (TR. ,B
Therefore, the only difference between (f and DSPG(f) is the choice of the prior distributions on h, which is On (h) and po po()), respectively. Given this property, we can compute the KL divergence between QspG (fP and pfra(f) given measurement points X,, (Appendix A.5):
Proposition 5 (KL divergence on measurement points between SPGs). Let (f) and pSPG( be the SPGs defined Equation 12 and 15. Then we have
DKL [gspa | (xx
By-awwer)) log 2(Pxx),
(16)
where Z(fXn) ) is the partition function, Z(PXm)= Sn Pspre (h|PXn) (n((h) dh.
Po(h)
Note that 2ffXn) ) is intractable to compute due the intractability of the posterior ūspa ffx Fortunately, this is already approximated by the PoE inference net @x (h|Px.) given by Equation 14:
Z(PXM) = 2(PX.) S (h/p2
2n (h) (n) dh.
(17)
Po (h(h) (h)
Since (h)), (h),,, and q (h|Px.) are all Gaussian distributions, (FFx.. can be computed using analytic solutions. Note also that thanks to the VAE-like structure in SPGs, all the caleulations are performed in the latent space, whose dimensionality is much lower than fx With the additional help of analytic solutions for ((xx), ), the estimation of (17) is very efficient and scalable.
4.4 The final algorithm: Mini-batching and de-biasing
So far we have been discussing the estimation of the KL-divergence meusurement points X.,. Next, we derive a practical algorithm based on the grid-functional ELBO Carid defined Equation Applying the approximation (17) to Equation 6, our final variational objective becomes
D
logp(D)
Egen (v.l((x)) So..x.DDDKLll (lp(().)))
(18)
D (fn - En,x,, Brroar log 2(Px.).
To make Equation 18 scalable large data, we can apply mini-batch sampling to the likelihood term E!P ((x)) Then the only bottleneck of Equation 18 is that the input X,, to the inference net ē (used in (((x))) can be very high dimensional due the condition X X,, required by Proposition 3. Fortunately, we can derive the following mini-batch estimators for En,x,,~lli frgunl))) log ((rx.) (Appendix A.6 and A.7):
Proposition 6. En.x, NoE f log (fxx)) can be estimated by the mini-batch estimator
JK := 1 Σ LE (m [log oogo + log x, log (o + 0x 1) -1,,0 0 -M2, a
1
2 1=1
(19)
+ (o 2 Pm,
- Px,) (0 2 1) 1)-1

---

Algorithm 1 Functional Variational Inference (FVI)
Require:
data = (x,, } j= prior p(f);; surrogate PSra (f), variational process (s) likeli- hood function pr (ylyf) mini-batch sizes I and K
while not converged do
2: Sample (-) f-)) [x]( from p(f).
3: Improve PSPG by optimizing the aggregated ELBO in Equation 13 w.r.t [wa]0,,
4: end while
5: while not converged do
6: sample mini-batch I from (1.... [D]} and set of measure points X,, via c.
7: Optimize Lfvi in Equation 20 w.r.t. 77 and via reparameterization tricks
8: end while
where H is the dimensionality of h, (b; Pap o (n N ((h (x(1/1xx). 0%2 and Pxx the mini-batch approximators for px; and ,3 respectively:
D
Σ
E
KEK
XXE
A4: 02
D PallssaMall
Σ
T 2 Palpai
A4
KEK
x/EX,
where K is mini-batch of size K sampled from (1,.... (D), X1 X XD is set of OOD samples sampled from TT using in Eq. 7, and Ph,15* and of helfsn are the mean and variance parameter returned from qx (xx)))(xxx))
The estimation in Eq. 19 is biased (but consistent). To remove the bias, we propose to debias Eq. 19 based on the Russian Roulette estimator [22]:
Proposition 7 (Debiasing). Let R be random integer from a distribution P(N) with sup- port over the integers larger than K. Xg is a random location sampled from T. Then X,, En.x,,~oEE g log Z(PX.) can be estimated by E JK + =k=K F(2E) where As Jk+1 -Tk and the expectation is taken over R, X and all mini-batches used by each Ju terms.
A.
This will enable us to also perform mini-batch sampling on the measurement points when performing FVI. Our final optimization objective function is
D Σ T
A P(N
LIVI
Pogn JR
Σ k=B
(20)
2)
where I is a mini-batch size for the likelihood terms, is an integer sampled from P(N), which is set to be (R K) G0.m( Finally, the full algorithm is sketched in Algorithm We call this proposed method Functional Variational Inference (FVI).
Scalability. Our method is empirically much faster than f-BNN (Appendix C.3). When estimating (prialg(S) (l))()) our method scales as ((M)), where M, is the number of samples sampled from ISPG(]] f' that are used in Jx. In practice, we use M, 1. On the contrary, the SSGE estimator used in f-BNN scales as O(M3 + M² D), where much larger value needs to used (e.g., M2 100).
5 Related Works
Since BNNs and VIPs are discussed in Section and 2, here we only address the rest of the related works including f-BNNs, functional priors, and Bayesian paraaaaeeteeis..
F-BNNs and F-PO. The functional BNN [47] proposed to address the issue of specifying mean- ingful priors to BNN weights. It matches a BNN to GP prior by minimizing the functional KL. divergence estimated by score function estimators [27]. As discussed [5], this objective is not
well-defined for wide class of distributions. Also, the score function estimators used in f-BNNs
7

---

Prior
Ground truth
VTP
FBNN
Precenise Liner Precewise Constank
Figure 3: Implicit function prior and posterior samples from ground truth, FVI, VIP, and f-BNN. respectively. The first row corresponds to piecewise constant prior, and the second row corresponds piecewise linear prior. The leftmost column shows 5 prior samples. From the second column to the rightmost column we show posterior samples generated by ground truth (returned by SIR), FVI. VIP and f-BNN, respectively. Red dots denote the training data. We plot 10 posterior samples in black lines and show predictive uncertainty grey shaded areas.
often perform poorly in high dimensional spaces and are less efficient. In practice, we found that the f-BNN computational time prohibitive. On the contrary, our FVI estimator avoids these issues by making use of the grid-functional divergence, which can be efficiently estimated using the latent representation of SPGs. More recently, the concurrent work of [39] proposes a tractable function space VI method for BNNs, in which the functional KL divergence approximated via lincarization. In their problem setting, the functional KL divergence well defined since both p and q are BNNs BNN priors
with the same neural network structures. Similarly, f-POVI [50] performs inference by performing particle optimization in function space. One limitation of f-POVI is that it requires the prior to be reparameterizable and differentiable while our method does not have this issue.
Function space priors. Another line of work directly defines distributions over functions by combining stochastic processes with neural networks. For example, neural processes (NPs) [13] and its variants (24, 15] focus on meta- lecaalca scenarios and propose to use set encoders to model all possible posterior distributions of the form {p(f[C)]C D}. where C is the so-called "context points" in neural processes. This could be inefficient for large datasets, since it needs to feed all data points to the set encoder, which scales linearly w.r.t. the dataset size. More importantly, NPs still use for learning and inference an ELBO defined on parameter space instead function space. On the contrary,
method focuses the functional VI for supervised learning scenarios and does not need to model all possible conditionals. When computing predictive distribution, we only need to evaluate gr (h), which is simple Gaussian distribution (no encoders involved).
Bayesian non-parametries. In the field of Bayesian non-parameteries, Gaussian Processes (GPs) [36] and deep GPs [8] are great examples of using function-space priors to produce models that are reliable under uncertainty. To reduce the prohibitive computation cost of exact/deep GPs, various VI methods [49, 17, 32, 40] have been studied. These methods share similar principle with our work, that is, to minimize the functional divergence between the posterior and variational processes. Nevertheless, the GP components of the functional prior play critical role in this line of work, which makes them less applicable to general non-GP based priors.
6 Experiments
In this section, we evaluate the performance of FVI using a number of tasks, including interpolation with structured implicit priors, multivariate regression with BNN priors, contextual bandits, and image classification. We mainly compare FVI with other weight-space and function-space Bayesian inference methods using the same priors. For more implementation details, please refer to Appendix B. Additional experiments can be found in Appendix C.
6.1 Interpolation with non-Gaussian priors: structured implicit priors
advantage of FVT is that can be applied to implicit (and non-Gaussian) priors over functions, where typical GPs do not apply. In this experiment, we evaluate the interpolation task as [47]. We
8

---

Table I: Regression experiment: Average test negative log likelihood
DNTASET
N
D
FVI
VIP-ENN
VIP-NP
B88
VDO
a = 0.5
FBNN
EXACT GP
BOSTON
306
13
2.3310.04
2.4520.04
2.4510.03
2.7610.04
2.6310.10
2.4310.02
1.3010.10
2.4610.04
CONCRETE
0.88:0.06
3.02+0.02
3.13+0.02
3.28+0001
3.2340,01
3.06:4.03
3.09-0.01
3.05+0..00 02
ENERGY
768
0.58=00.05
0.56:00.0
0.60=0003
2.17+0..2
1.13+0.02
0.95=00.99
0.68=00022
0.34+0.02
KINSNM
8192
-1.15:0.01
-1.1220.01
-1.05=0.00
-0.8120.01
-0.8340.01
-0.9210.02
NIA40.00
N/A20.00
POWER
9568
2.69:0.00
2.92:0.00
2.90-00.0
2.8340.011
2.88:0.000
2.81
N/A00,000
N/A10.00
PROTEIN
457.00
2.85++0.00
2.87+0000
2.96+0002
3.00+0.00
2.99+0.00
2,90+0.00
N/A00.00
N/A+0.00
RED WINE
1588
11
0.97:0.06
0.97=0002
1.20=0.04
1.0110.02
0.97:0.02
1.0120.02
10420.01
0.2610.03
YACHT
.088
0.5940.11
-0.02-0.077
0.59.0.13
1.1140.04
1.22
0.7940.11
1.3
0.1010.05
NAVAL
11934
16
-7.2110.06
-5.6220.04
4..1120.00
2.80.00.00
-2.8010.00
-2.9720.14
7.13-20.02 NIA.0.00
MEAN KANK
N/A
N7A"
1.33
2.11
3.56
3,22
4.36
3:33
NIA
NNAA"
Table 2: Contextual bandits performance comparison. Results are relative to the cumulative regret of the worst algorithm on each dataset. Numbers after the algorithm are the network sizes. The best methods are boldfaced, and the second best methods are highlighted in brown.
MEAN RANK
MUSHROOM
STATLOG
COVERTYPE
JESTER
ADULT
WHEEL
CENSUS
PVIZX 50 UNIFORM
2.11
16.46
21 1:04
7,95
= 2.92
49.59 = 1.61
68.59 = 6.87
90.33 : 0.36
41:44 + 9.28
51.77 = 3.16
10.45
100.0
00
99.85 : 0.36
99.49
0.62
100.0 : 0.00
99.60
0.53
94.04
11.9
99.30 : 0.55
RMS
17.74
16:36
2:51
69.72
7.23
5.50
97.65
1.48
94.35 +
DROPOUT
67.72
2.32
4,66
97.44
BBB 2
4.88
23.18 5.90
30.90 2 3.29
63.91
1.96
72,93
1 5.69
95.49
2.03
56.38
11.3
70.68
2.32
BBB 1 50
8.22
15.52
4.40
80.25
= 18.6
94.80
4.84
83.30 2: 5.26
99.24
1 0.66
58.12
18.0
99.46
1 0.37
NEURALLINEAR
6.91
19.04
1 2.96
21.22
1.98
75.34
1.00
$6.55 1. 3.61
97.93
1.37
37.41
1 8.86
83.75
1 L.44
BoorRMS
4.51
17.11 # 5.99
9.47 2.03
63.27
1.35
74.66
+ 3.87
96.1
1.02
63.15
+ 25.9
90.47
+ 3.40
PARAMNOISE
17.76
20.95
+ 3.07
78.08
5.66
76.95
2 5.84
96.23
1,81
41.26
96.34
4.56
RBo
95.22
4.88
98.60
1.45
2.69
98.72
1.28
97
2.01
FBNN x
16.55 2.41
10:01
1.39
50.
5.70
70.82
3.27
90.72
3.18
77.70 + 21.2
51.22
consider two 1-D implicit priors on [0, 1]: 1), piecewise constant random functions and 2), piecewise linear random functions. Please refer to Appendix B.2 for details. For each prior, we first sample a random function from the prior; then, 100 observed data points
sampled as D, half of which are sampled from [0, 0.2] and the other half are sampled from [0.8, 1]. Finally, we ask the algorithms to perform inference using the prior, i.c., producing samples from (D).
We compare the performance of FVI with ground truth, f-BNN and VIPs. The ground truth posterior samples are generated by sampling importance re-sampling. F-BNNs are based on the code kindly open- open-sourced by [47]. As we found that the training time required by f-BNN is prohibitive, we only trained f-BNN for 100 epochs for fairness (additional results for fully trained f-BNNs are provided in Appendix C.5). For VIPs (Gaussian approximations), we use an empirical covariance kernel, which is estimated from random function samples of the implicit priors. For FVI, implementation details can be found in Appendix B.2.
Results are displayed in Figure 3. FVI can successfully generate samples that mimic the piccewise constant/linear behaviors. The posterior uncertainty returned by FVI is also close to the ground truth estimates. On the other hand, f-BNNs severely under-fit the data and provide very poor in-between uncertainties. Note that, although f-BNNs are only trained for 100 epochs, their running time is still 100x higher than that of FVI (Appendix C.3). VIP performs better than f- f-BNNs, but fails to mimic the behaviour of the priors: the posterior samples from VIP are very noisy. This is due to the prior function samples violating the Gaussian assumption, with the correlation level between points being lower than expected. This results in very noisy VIP posterior samples that are hard to interpret.
6.2 Multivariate regression with BNNs priors
In this experiment, we test if the proposed FVI can perform accurate posterior inference with BNNs as functional priors. We consider multivariate regression tasks based on 9 different UCI datasets. We mainly compare with the following weight-space VI baselines for BNNs:
[4], variational dropout [12], and variational alpha dropout [26] (a 0.5). We also compare with three function-space BNN inference methods: VIP-BNNs, VIP-Neural processes [28], and f-BNNs. Finally, we include comparisons to function space particle optimization [50] in Appendix C.7 for reference purpose. All inference methods are based on the same BNN priors whenever applicable. For experimental settings, we follow [28]. Each dataset was randomly split into train (90%) and test sets (10%). This was repeated 10 times and results were averaged.
Results are shown in Table 1. Overall, FVI consistently outperforms other VI-based inference methods for BNNs and achieves the best result in datasets (out of 9). FVI also outperforms f-BNNs (in 5 datasets out of 6), despite the fact that they are more expensive to train. Note that exact GPs and

---

Table 3: Image classification and OOD detection performance. Accuracy, negative log- rogllikelhoodd (NLL) and rrea-under-the-curee (AUC) of OOD detection are reported. Our method outperforms all baselines in terms of classification accuracy and OOD-AUC and performs competitively on NLL for CIFARIO. Results for MAP, KFAC and Ritter et al. are obtained from [21].
FMNIST
CIFARTO
MODEL
ACCURACY
NLL
OOD-AUC
ACCURACY
NLL
OOD-AUC
FVI
91.60=0.144
0.254=0.05
0.956=0.06
77.69 =0.64
0.673=0.03
0.883=0004
MFVT
91.20+00.1
0.343+0.01
0.782+0.02
76.40+0.52
1.372++0.02
0.589+0.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+00.12
0.942+0.01
0.945+0.00
72.49+0.20
1.274=0.01
0.548+0.01
RITTER ET AL..
91.20+0.07
0.265+0.00
0.947+0.00
77.38+0.06
0.661+0.00
0.796+0.00
f-BNNs are not directly comparable to other methods, since i), they perform inference over different priors; and ii), they are much more expensive as they require the evaluation of the exact GP likelihood. Thus, their results are only available for smaller datasets, and are not included for ranking.
6.3 Contextual Bandits
Uncertainty estimates important for downstream eeeiiion-making scenarios, since exploration- exploitation is a common dilemma that must be addressed. In this section, we consider a classic task called contextual bandits, where the agent is asked make decisions that maximize the reward given some contexts (inputs). For this, Thompson sampling [48] is an elegant approach to guide exploration, where model configuration is first sampled from the posterior and then an optimal action is choosen based on the sampled configuration.
We compare FVI with several NN-related baselines on datasets benchmarked by [37]. The hyperpa- rameter settings are consistent with [47] except that we used a smaller batchsize (32). The learning rates for each baseline mean ranks. Experiments are repeated for 10 runs. As shown in Table 2. no single algorithm always outperforms the others in all bandit problems. However, FVI tends to give better performance than the baselines (ranks the first overall, performs the best on 4 out of 7 datasets, and ranked top 2 on of datasets), indicating that FVI can provide reliable uncertainty estimates for decision making. Moreover, FVI is much more efficient than f-BNN (nearly 500 times faster, c.f. Appendix C.4).
tuned from (0.001, 0.05]. We report the cumulative regret well as the
6.4 Image classification and out-of-distribution detection
To demonstrate the scalability of our method higher dimensional data, we consider image classifica- tion tasks on Fashion MINIST [54] and CIFAR 10 [25] with BNN priors. We compare our method to the following baselines: mean field VI (MFVI), maximum posteriori (MAP), KFAC Laplace- GNN approximation [31] and its dampened version [38]. For all models, we use Bayesian CNNs with the same mixed CNN-fully connected structure in [21, 43). Apart from test accuracy and negative log likelihood (NLL), we also perform out-of-distributioo detection using in-distribution (ID) / out- of-distribution (OOD) pairs including FashionMNIST/MNISTT and CIFARIO/SVNH, Following the settings of [35, 21J. we calculate the area under the curve (AUC) of -iittr-iisttiiuttion detection based on predictive entropies. Results are shown in Table 3. On both datasets, our proposed FVI method consistently outperforms all baselines in terms of (in-distribution) classification accuracy and OOD detection AUC. Although FVI does not achieve the best NLL on CIFAR10, it still performs competitively to MAP and dampened KFAC. This demonstrates that our method is able to scale to high dimensional data and produce accurate predictions with well-calibrated uncertainties.
7 Conclusion
In this paper, we propose Functional Variational Inference (FVI), new VI method in function space. optimizes a grid-based functional divergence, which can estimated based on our proposed SPG model. We demonstrated that FVI works well with implicit priors, scales well to high dimensional data and provides reliable uncertainty estimates. Possible directions for future work might include developing grid-function KL estimation methods without surrogate models and improving the theoretical understanding of functional space VI.
10

---

Acknowledgements
We thank Javier Antoran and Yingzhen for helpful discussions and comments during the early stage of this work.
References
[1] Arthur Asuncion and David Newman. UCI machine learning repository, 2007.
[2] Keith Ball. Eigenvalues of Euclidean distance matrices. Journal of Approximation Theory, 68(1):74-82, 1992.
[3] BJC Baxter. Norm estimates for inverses of Toeplitz distance matrices. Journal of Approximation Theory, 79(2)222-2242, 1994.
[4] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. urXiv:/505.05422 2015.
[5] David R Burt, Sebastian W Ober, Adrià Garriga-Alonso, and Mark van der Wilk. Understanding variational inference in function-space. arXiv preprini urXiv:2011.09422 2020.
[6] Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henrik Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pages 9916-9926, 2019.
[7] Bin Dai and David Wipf. Diagnosing and enhancing VAE models. arXiv: VrXii:I903.055 1903.0579, 2019.
arXiv preprint
[8] Andreas Damianou and Neil Lawrence. Deep Gaussian processes. In Artificial Intelligence and Statistics, pages 207-215, 2013.
[9] Sebastian Farguhar, Lewis Smith, and Yarin Gal. Liberty or depth: Deep Bayesian neural nets do not need complex weight posterior approximations. arXiv e-prints, pages arXiv-2002, 2020.
[10] Andrew YK Foong, David R Burt, Yingzhen Li, and Richard E Turner On the expressiveness of approximate inference in Bayesian neural networks. arXiv, pages arXiv-1909, 2019.
[11] Andrew YK Foong, Yingzhen Li, José Miguel Hermández-Lobato, and Richard Turner. 'in-between uncertainty in Bayesian neural networks. arXiv preprint 1537, 2019,
[12] Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning, pages 1050-1059, 2016.
[13] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.01622, 2018.
[14] Wenbo Gong, Yingzhen Li, and José Miguel Hernández-Lobato. Sliced kernelized Stein discrepancy. arXiv preprint arXiv:2006. 16531, 2020.
[15] Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard E Turner. Convolutional conditional neural processes. arXiv preprint 10X0v:19/0.13555 2019.
[16] Arthur Gretton, Karsten Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
[17] James Hensman, Nicolo Fusi, and Neil D Lawrence. Gaussian processes for big data. arXiv:1309.66. 6835..855.... 2013.
[18] José Miguel Hernández-Looato. Balancing flexibility and robustness machine learning: semi-parametric methods and sparse linear models. 2010.

---

[19] José Miguel Hernández- and Ryan P Adams. Probabilistic backpropagation for scalable learning of Bayesian neural networks. arXiv:1502.055336 2015.
[20] José Miguel Hernández-Lobato, Yingzhen Li, Mark Rowland, Daniel Hernández-Lobato, Thang Bui, and Richard Eric Turner. Black-box a-divergence minimization. 2016.
[21] Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of Bayesian neural nets via local linearization. In International Conference Artificial Intelligence and Statistics, pages 703-711. 1. PMLR, 2021.
[22] Herman Kahn. Use of different Monte Carlo sampling techniques. 1955.
[23] Alex Kendall and Roberto Cipolla. Modelling uncertainty in deep learning for camera relocal- ization. In 2016 IEEE international conference on Robotics and Automation (ICRA), pages 4762-4769, IEEE, 2016.
[24] Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv:1901.08 005611, 2019.
[25] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[26] Yingzhen Li and Yarin Gal. Dropout inference in Bayesian neural networks with alpha- divergences. arXiv: 7003:02904 2017.
[27] Yingzhen Li and Richard E Turner Gradient estimators for implicit models. arXiv:700.007107, 2017.
[28] Chao Ma, Yingzhen Li, and José Miguel Hernández-Lobato. Variational implicit processes, In International Conference on Machine Learning, pages
2019.
[29] David JC MacKay, A practical Bayesian framework for backpropagation networks. Neural computation, (3):448-4722 1992.
[30] Wesley Maddox, Gregory Benton, and Andrew Gordon Wilson. Rethinking parameter counting in deep models: Effective dimensionality revisited. arXiv preprint arXiv:2003.021139, 2020.
[31] James Martens and Roger Grosse. Optimizing neural networks with Cronecker-factoree approx- imate curvature. In International conference machine learning, pages 4408-2417. PMLR. 2015.
[32] Alexander G de G Matthews, James Hensman, Richard Turner, and Zoubin Ghahramani. On sparse variational methods and the Kullback-Leibler divergence between stochastic processes. Journal of Machine Learning Research, 51:231-239 2016.
[33] Francis Narcowich and Joseph D Ward. Norm estimates for the inverses of a general class of scattered-data radial-function interpolation matrices. Journal of Approximation Theory, 69(1):84-109, 1992.
[34] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
[35] Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard E Turner, Rio Yokota, and Mohammad Emtiyaz Khan. Practical deep learning with Bayesian principles. arXiv preprint arXiv: 095600 2019.
[36] Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning, volume 1. MIT press Cambridge, 2006.
[37] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep Bayesian bandits showdown: An empirical comparison of Bayesian deep networks for Thompson sampling. arXiv preprint arXiv:1802.00127, 2018.
2018.
12

---

[38] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018- Conference Track Proceedings, volume 6. International Conference on Representation Learning, 2018.
[39] Tim G. Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal. Rethinking Function-Sye SSace Variational Inference in Bayesian Neural Networks. In Third Symposium Advances in Approximate Bayesian Inference, 2021.
[40] Hugh Salimbeni and Marc Peter Deisenroth. Doubly stochastic variational inference for deep Gaussian processes. In Proceedings of the Ist International Conference Nearal Information Processing Systems, pages
2017.
[41] Simón Rodríguez Santana and Daniel Hernández-Lobato. Adversarial a-divergence minimiza- tion for Bayesian approximate inference. Neurocomputingg 2020.
[42] Robert Schaback. Lower bounds for norms of inverses of interpolation matrices for radial basis functions. Journal of Approximation Theory, 79(2):287-306, 1994.
[43] Frank Schneider, Lukas Balles, and Philipp Hennig. Deepobs: A deep learning optimizer benchmark suite. arXiv preprint arXiv:1903.05445 2019.
[44] Lewis Smith and Yarin Gal. Understanding measures of uncertainty for adversarial example detection. arXiv preprint 1rX3vv1803.0000 2018.
[45] Bharath Sriperumbudur, Arthur Gretton, Kenji Fukumizu, Bernhard Schölkopf. and Gert RG Lanckriet. Hilbert space embeddings and metrics on probability measures. The Journal of Machine Learning Research, 11:1517-1561, 2010.
[46] Ingo Steinwart. On the influence of the kernel on the consistency of support vector machines. Journal of machine learning research, 2(Nov):67-93, 2001.
[47] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational Bayesian neural networks. 2018.
[48] William Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3/4):285-294 1933.
[49] Michalis K Titsias. Variational learning of inducing variables in sparse Gaussian processes. In International Conference on Artificial Intelligence and Statistics, pages 567-574, 2009.
[50] Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for Bayesian neural networks. arXiv preprint urXiv:1902.099
2019.
[51] Andrew Wathen and Shengxin Zhu. On spectral distribution of kernel matrices related to radial basis functions. Numerical Algorithms, 70(4))709-726, 2015.
[52] Holger Wendland. Scattered data approximation, volume 17. Cambridge university press, 2004.
[53] Mike Wu and Noah Goodman. Multimodal generative models for scalable weakly-supervise learning. In Advances in Neural Information Processing Systems, pages 3575-5585, 2018.
[54] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 1eXXiv1708.0777 2017.
[55] Ruiyi Zhang, Chunyuan Li, Changyou Chen, and Lawrence Carin. Learning structural weight uncertainty for sequential decision-maken making. In International Conference on Artificial Intelli- gence and Statistics, pages 1137-1146, PMLR, 2018
[56] Yuhao Zhou, Jiaxin Shi, and Jun Zhu. Nonparametric score estimators. In International Conference Machine Learning, pages 11513-11522. PMLR, 2020.
13


Reviews:
[{"review_id": "xVoEP6w_73", "reviewer": "Reviewer_w44o", "summary": "This paper proposes a new framework for conducting approximate Bayesian inference in the context of function-space optimization. The method described is fairly general and with good theoretical justifications to some of the most important claims here. To achieve their goal, the authors define a new approximate objective function with an alternative divergence to the regular KL  and propose the SPGs as a generalization over previous work. Even though there is some important information missing in the experimentation phase, the performance of the method seems overall much better than previous state-of-the-art models in many contexts. This needs to be addressed since it constitutes an important issue, but the method is promising and interesting. \n\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "In this paper, the authors propose a new approach to inference in functional space, which has been an important issue of research for the past few years. This is done by introducing an alternative objective function with different properties to the ones of the regular functional-ELBO. The change in the objective function is effectively a newly proposed divergence (grid-functional KL divergence), which could be interpreted as the expected value of the KL divergence between functions sampled from two stochastic processes, evaluated at a random number of measurement locations. The key difference here is that the number of sampling points is a random variable instead of remaining fixed, proposing a grid approximation to the KL divergence. Combining this with the SPG, the authors extend on the formulation by previous work avoiding some of the restrictions of previous methods. \n\nThe results presented in the paper show that the method behaves well and performs better overall than the rest of the methods compared it is being against in the experiments conducted. It represents a fairly general approach (although with certain important similarities to previous literature work) that could help further the comprehension and usage of function space optimization methods.  However, although the results are promising, there are some important experimental concerns that should be addressed in order to have a clearer picture of the capabilities of the proposed framework. \n\n# Pros\n\n* There is considerable theoretical work being conducted to justify a lot of the properties mentioned of this new approach. Although a bit tough at times, the descriptions in the main text as well as in the appendix help a lot with the comprehension of the proposal.\n\n* The framework constructed is general enough to be used with several methods. This flexibility could be interesting to extend the formulation with the usage of new methods. \n\n* The performance of the method seems to improve on the previous state-of-the-art, both when compared with some of the best methods in weight and function space optimization. \n\n* The paper, as well as the appendix material, is well written and easily comprehensible for the most part. Although some sections need a careful read, in most cases it is written in a very clear manner. \n\n# Cons\n\n* If L^grid >= L^functional, then optimizing it does not mean we are also optimizing L^functional since it makes a weaker lower bound definition (since it is higher, maximizing it would not necessarily imply maximizing L^grid). Although the method seems to perform well, this point needs to be addressed more carefully due to its crucial importance in the development of the paper and the training of the model.  \n\n* The experiments are conducted using the same prior, which is good as a reference. Since this approach also supports other priors, it would be interesting to see how other methods would perform in comparison given their respectively \"best\" prior (as described in their respective papers). For example, for f-BNN, the usage of a piece-wise implicit prior may not be optimal since that method is not able to train such type of priors, and thus a GP (or a sparse GP) could help improve its performance w.r.t. what is seen in the experiments. \n\n* Regarding the regression experiments, there is important information not being reported in the main text nor in the appendix. It would be very helpful to have a table with the RMSE (or some other alternative error metric) for completeness, while also having another alternative measure other than the log-likelihood to address the goodness-of-fit of the results (see [1]). \n\n* Through the paper there are some claims about the scalability of the method. However, the material provided lacks some kind of information on the convergence rate and time of this method in comparison to the other ones. To better support these claims, please provide a convergence study for the methods here. \n\n\n\n\n# Notes\n\n* Does this formulation generalize over implicit stochastic processes (as described in the paper of Variational Implicit Processes), or is it a more restricted description form for the stochastic processes involved? Specifically, clarify this w.r.t. the description of the functions sampled from both stochastic processes in Proposition 2 and the implications made in the following paragraph (those very same properties could be assigned to implicit processes as well). \n\n* Also regarding the experiments, although an extensive comparison with other methods is conducted, in the case of weight-space optimization there are some methods that currently perform better than the ones present here, such as the one described in [2]. For a more comprehensive description of the performance, it would be convenient to include these approaches.  \n\n* How flexible could be the resulting predictive distributions that the method is able to produce? Could you try to provide any tests to showcase this? As far as the formulation goes, the predictive distributions could be pretty general, but it would help if there was any experimental confirmation of this fact. \n\n* Are there any concerns with the problems that may arise due to using a sampling-based approach for the estimation of the KL? For f-BNN it could be seen that they would be somewhat affected by the dimensionality of the data precisely due to this fact. \n\n* Please provide the information on the computational resources employed to conduct the experiments.  \n\n*Minor details:*\n\n* Although it is comprehensible as is, maybe the notation in the equations of Proposition 2 could be a little more different so that there is a clearer distinction between the functions sampled from $q$ and $p$\n\n* Double \":\" in line 85\n\n* The enumeration in the second paragraph of the introduction could appear differently so it helps reading the text, either by enclosing properly the (i) and (ii) or by separating each point with a different item in an enumerated list. \n\n\n\n##  References\n\n[1] Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 102(477), 359-378.\n\n[2] Santana, S. R., & Hernández-Lobato, D. (2020). Adversarial α-divergence minimization for Bayesian approximate inference. Neurocomputing.\n\n\n ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The societal impact of research such as this one is mostly restricted to the community of approximate inference, especially to those interested in function space optimization. Therefore, I deem there is no potential negative societal impact. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "S5Iqa6UW5LP", "reviewer": "Reviewer_Jd3R", "summary": "The authors propose \"Functional Variational Inference\" (FVI), a novel variational approach to approximate Bayesian inference in function space. The approach is based on the \"grid-functional KL divergence\", a novel form of KL divergence between stochastic processes. The authors show that this is a proper KL-divergence and that it is well-defined in cases where previous work (fBNN) is not.\n\nFurthermore, the authors propose \"Stochastic Process Generators\" (SPGs), a novel family of variational distributions which are more flexible in comparison to distributions used in previous work (VIP). Using SPGs, the authors derive an efficient approximation of the ELBO corresponding to the grid-functional KL divergence. The authors demonstrate that this approximation scales better than the one used in fBNN. \n\nThe authors evaluate FVI on a range of experiments (interpolation with implicit priors, regression and classification with BNN priors, Thompson sampling). FVI tends to perform favourably in comparison to fBNN, VIP, and a range of other existing methods.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Originality: The proposed FVI is a novel variational approach in function space. The authors explain in which respect their work extends and improves upon prior approaches (fBNN and VIP). As described below (Sec. \"Clarity\"), the paper could be strengthened by providing a bit more intuition on how FVI remedies the weaknesses of prior work.\n\nQuality: The proposed method appears to be formulated in a technical sound way (I did not check all of the proofs). The authors establish the well-definedness of FVI and provide both theoretical and empirical evidence that their method improves upon prior work. \n\nNevertheless, there are some weaknesses, on which I ask the authors to elaborate:\n\n- The authors use the Product of Experts (PoE) architecture for the encoder used to distill the prior into an SPG. While I think this is a valid approach, I do not understand the author's comments on the \"Deep Set\"-like approach used for the Neural Process (NP) family. Why should this method be inferior to the PoE approach, i.p., why should it be more \"difficult in high-dimensional spaces\" (l. 219). Also, what do the authors mean with \"need to model all possible conditionals\" (l. 222). Did the authors try the \"Deep Set\"-like approach for their method?\n- The results for fBNN in Fig. 2 seem to be way worse than those shown by the original authors on a similar experiment, cf. Sun et al., \"Functional Variational Bayesian Neural Networks\" Fig. 3. Could the authors elaborate on this discrepancy? What do the authors mean with \"for fairness\" in l. 247?\n- Why do the authors consider an additive white noise process $\\nu$ in their definition of the SPG? Isn't noise already considered through the likelihood $p_\\pi$?\n\nClarity: Overall, the writing style and paper organisation is good. Nevertheless, I think some points could be improved. As FVI is closely related to fBNN and VIP, I would like to see a bit more background on those methods and how FVI remedies their shortcomings. In particular, it would help the reader if the authors could provide some intuition for the following (I would suggest moving some of the more technical details to the appendix):\n\n- What is the intuition behind using Eq. (3) instead of Eq. (2)? Why is it crucial to have a random number $n$ of measurement points in Eq. (3)? What is the influence of the sampling distribution $c$ in Eq. (3) and of the particular choice Eq. (6)? I am aware that the authors provide formal arguments to support these choices, but I think that their presentation lacks intuitive motivation for them.\n- Why does it make sense to freeze the basis functions and decoder parameters after fitting the prior using an SPG? Is this merely a choice required to make the KL divergence between SPGs tractable? How does this choice relate to Proposition 4, i.e., how does this \"parameter freezing\" affect SPG's expressiveness?\n\nMinor concerns/typos (do not influence my score):\n\n- l. 26: missing space\n- l. 28: F-BNN → f-BNN\n- l. 56: to generated → to generate\n- l. 68: supreme → supremum\n- l. 88: exist → exists\n- l. 97: $X_\\mathcal D$ is undefined\n- Eq. (8): the integral is missing a closing $\\mathrm d h$. The same goes for all other appearances of this VAE integral.\n- l. 126: Then, for $\\forall$ → Then, $\\forall$\n- Text around Eq. (11) states twice that Eq. (11) is called $\\tilde p_{SPG}$.\n- ll. 140-142: $f_m^X$ is explained twice\n- l. 157: extra comma\n- Eq. (19): The likelihood misses a subscript $\\pi$\n\nSignificance: The authors propose an interesting approach to an important problem. The results show improvements compared to previous work. Further, the method seems to be theoretically well-founded. I therefore think the proposed FVI can be a valuable contribution. However, in its current form, I judge the paper to be borderline because of the concerns I raised above. If the authors could elaborate and improve upon those points, I am willing to raise my score.\n\n------------------------------------------------------------------------------\nUpdate after reading the author's response:\n\nI thank the authors for their response. I raise my score because the authors addressed my concerns and will clarify some of the points I raised in the updated manuscript.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No potential negative societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Qm6xUKqnFJF", "reviewer": "Reviewer_4izH", "summary": "This paper worked on the approximation for the posterior distribution, especially when we perform Bayesian inference in the space of functions. The authors provided a new Kullback Leibler divergence for stochastic processes. The number of the measurement points is a random variable. This new KL divergence is well defined for stochastic processes which are generated from different network architectures. Based on the new KL divergence, the authors developed a scalable variational inference using the VAE to generate approximate posterior stochastic processes.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Overall, I think the paper could be accepted, but requires a revision for the algorithm part in Section 3.2 to 3.4.\n\n# Pros\n-  This paper provides a new KL divergence, which is valid for different network architectures.\n-  Numerical experiments, especially for the structured implicit priors show strong evidence that the proposed method is useful numerically.\n\n# Cons\n-  The theoretical part of the paper is hard to follow for a non-expert. Especially, in the proof of Proposition 4, I could not follow the discussion of the sequence of VAEs.\n-  Algorithmic part is hard to follow at first since each detailed technique is presented one by one.\n\n# Comments and Questions:\n- How strong are assumptions in Proposition 2 ? In the proposition, the compactness of the supports of $p(\\theta)$ and $q(\\Gamma)$  are assumed. I thought these p and q are Gaussian distributions or Inverse Gamma distributions in standard Bayesian neural networks. Then, their supports are not compact, thus it violates the assumption.\n- Is there any other example for the distribution of $n$ other than a geometry distribution ? And is there any reason that using a geometry distribution is important numerically ? \n- In numerical experiments, how did authors construct the compact space $\\mathcal{T}$ for the input space from the data set ?\n- Why should we distill the prior distribution $p(f)$ ? \n- Is it possible to compare the KL divergence defined in (2) and (3) numerically ? \n- In proposition 4, using VAE is justified regarding MMD distance. On the other hand, in variational inference, we are working on KL divergence, not MMD. Is it possible to justify using VAE based on KL divergence under some mild conditions?\n- I think experiments for the structured implicit priors in Sec 5.1 strongly support the usefulness of the proposed method. I recommend using the real dataset for this task to support your proposed method. \n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they are shown in the checklist.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 to 4 hours"}, "extra_text": {}}, {"review_id": "CMzuKCt466y", "reviewer": "Reviewer_YNdn", "summary": "The paper presents an approach to probabilistic regression. The posterior distribution is estimated using a variational approach. A variational class of stochastic processes is presented, which is parameterized by neural networks - amongst other things by a product of experts variational autoencoder. For optimizing the variational distribution a variational bound and its sample approximation is derived. The proposed approach is evaluated on regression, classification and reinforcement learning tasks.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Overall, I enjoyed reading the paper and the results seem to be very competitive and impressive.  \n\nOriginality: The approach seems like a new approach, which combines and extends ideas from previous work, such as [44], [27]  and [50]. The techniques are as far as I could tell well known. The paper differentiates well from the related work, which is sufficiently discussed. \n\nQuality: As far as I could tell the submission is technically sound. The resulting algorithm is extensively evaluated empirically and some theoretical inside are given. \n\nClarity: The paper is nicely written and organized. I think that it adequately informs the reader. \n\nSignificance: The empirical results seem pretty impressive and could be of much use to the machine learning community. \n\n\nOne thing I missed was:\n- An experiment where the true posterior is a Gaussian process and the proposed approach is evaluated against it. In such an example the ground truth is not only given by samples from the posterior as in Sec. 5.1, but the whole posterior distribution can be evaluated in closed form. As such, one could evaluate how accurate the variational approach is.\n\nThings I did not understand:\n- Why is in eq. (12) the prior only fitted using regressors $\\mathbf X_O \\subseteq \\mathbf X_{\\mathcal{D}}$ from the data and not some arbitrary regressors $\\mathbf X_O \\subseteq \\mathcal T^{\\mathbb Z_+}$?\n- Why is FBNN performing so poorly in Figure 2? The example seems extremely simple and Figure 3 in [44] suggest different results. If it is just because of the training time needed, I do not that such a comparison is fair. \n\nRemarks:\n\n-Reference [44] is missing the venue.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for approximate Bayesian inference in function space using a novel form of the KL divergence between stochastic processes, the grid-functional KL divergence, where the number of measurement points is a random variable. This divergence has better properties than those used in previous work. All the reviewers recommend acceptance and I agree. 
Decision: Accept (Poster)