Functional Variational Inference based on Stochastic Process Generators
Year: 2022

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.

heading: 1 Introduction
text: 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-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 in the network forward pass, it is unclear the effect that a given prior distribution over the weights will have in the resulting distribution over functions.
To address these issues, there has been a recent resurgence of interest in applying the perspective of Bayesian non-parametrics to neural nets. That is, BNNs are treated as probability measures over functions, i.e., as stochastic processes, with Bayesian inference being performed now in function space.e Examples include functional BNNs (f-BNNs) [47], and variational implicit processes (VIPs) [28]. Despite their empirical advantage over weight space VI, they still suffer from a number of issues. i), f-BNNs optimize a functional KL divergence between BNNs and GPs, which is not always well-defined[5]. Also, f-BNNs rely on (spectral) stein gradient estimators, which are less efficient for high dimensional distributions [56, 14]. While VIPs do not have this problem, they use a wake-sleep procedure that does not correspond to a single unified objective. ii), F-BNNs rely on a mean-field approximation, which often lacks predictive in-between uncertainty on test data. This issue was observed both with single layer BNNs [10], and deeper BNNs (able to represent in-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 to capture non-Gaussian processes such as structured implicit priors [47].
Therefore, it is an open challenge how to 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 a new solution. Our contributions are as follows:
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
• We propose a 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-Gaussian generalizations of the VIP family [28], and can help avoid the fundamental limitation of mean field Gaussians [10] used in BNNs and functional BNNs. A theorem regarding the expressiveness of SPGs is proved (Proposition 4). Based on SPGs, our proposed functional divergence between stochastic processes can be estimated efficiently using mini-batch sampling (Proposition 5, 6 and 7), which achieves a significant speed-up against the gradient estimator approach in [47].
• We compare our methods against existing weight-space and function-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.

heading: 2 Backgrounds
text: Variational BNNs Let y = g(x,w) be a neural network, where x is the input, y is the output, and w denotes the weights. To build a Bayesian neural network (BNN), we place a prior p(w) over w, and add an observational noise ϵ ∼ N (ϵ; 0, σ2) to the output. Then, a BNN is given by
log p(y|x) = log ∫ w N (y; g(x,w), σ2)p(w)dw. (1)
Given a dataset D = {xi, yi}Ni=1, the goal of Bayesian inference is to compute the posterior p(w|D) ∝ p(D|w)p(w). Since this is 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) is usually highly multi-modal due to the non-identifiability of neural networks, VI often leads to unsatisfactory results, and suffers from over-parameterization and other pathologies [10].
Variational Implicit Processes (VIPs) An alternative way of looking at BNNs is to treat them as a prior distribution over function space, i.e., a stochastic process. Indeed, p(y|x) in Equation 1 generates a random function g(·,w), depending on what w is 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 qVIP(f |D) = ∑ s asϕs,a ∼ N (a;µ,Σ), where the basis functions {ϕs}Ss=1 are random samples drawn directly from p(f). Although VIP has shown improved performance over parameter-space VI methods, it has a few limitations. For example, qVIP(f |D) resembles a GP approximation to p(f |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(f) to be reparameterizable. i.e., can be represented in the form of f(x) = gθ(x, z), z ∼ p(z). This may limit its application to non-reparameterizable priors.

heading: 3 Problem setting and the functional KL divergence
text: We consider the problem of Bayesian inference in function space. Let p(f) be a stochastic process defined on the probability space (Ω,B). Note that we use f in its scalar form to denote a scalar function f(·) : T 7→ R. Here T is the index set of p(f) (assumed to be a compact subset of Rd). For example, p(f) could be a Gaussian process p(f) = GP(m(·), k(·, ·)), a Bayesian neural network, or any other suitable stochastic process. We use p(f) to model the uncertainty in function space. Then, a likelihood function pπ(y|f(·)) is defined on top of f to generate observable data y. Given observed data D = {xi, yi}Ni=1, our goal is to infer the posterior process p(f |D) conditioned on the observations D. If p(f) is a GP, then p(f |D) can be computed analytically. However, in most cases this is intractable. Therefore, following [32, 47], we define another stochastic process q(f) on (Ω,B) as our variational family to approximate p(f |D). q(f) can be optimized by maximizing the evidence lower bound (ELBO) in function space:
Lfunctionalq := Eq(f)[log pπ(D|f)]−DKL[q(f)||p(f)]. (2)
Note that DKL[q(f)||p(f)] is the functional KL-divergence between stochastic processes. As shown by [47], DKL[q(f)||p(f)] can be written as:
DKL[q(f)||p(f)] = sup n,Xn DKL[q(f Xn)||p(fXn)], (3)
where Xn denote a set of n measure points {xk}1≤k≤n in the domain/index set of f(·), which can be treated as an element of the product space T n. Moreover, fXn denotes the vector of function values evaluated on Xn, and DKL[q(fXn)||p(fXn)] 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 Z+ .

heading: 4 Functional Variational Inference using Stochastic Process Generators

heading: 4.1 The grid-functional KL divergence
text: ability values pn to larger n. This way, we can trim down the contributions from the KL terms
DKL[q(f Xn)||p(fXn |D)] 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 as
Dgrid[q(f)||p(f |D)] := En,Xn∼cDKL[q(fXn)||p(fXn |D)], (4)
where Xn is a set of n measurement points {xk}1≤k≤n sampled from T , according to some sampling distribution c. Note that c is defined on the product space T Z+ and the number of sampled measure points n is also random. One may recognize that [47] proposed a similar objective as an approximation to 3 in which the number of measure points n is a fixed constant instead of a random variable. However, as we show below, it is in fact critical that n is not fixed.
Proposition 1. Suppose c has full support on T Z+ . Then, Dgrid[q(f)||p(f |D)] satisfies the following conditions: i), Dgrid[q(f)||p(f |D)] ≥ 0; ii), Dgrid[q(f)||p(f |D)] = 0 if and only if q(f) = p(f |D).
In other words, if c has full support on T Z+ (we will give an example of such c later), then Dgrid[q(f)||p(f |D)] is a valid divergence in function space. Therefore, we can use Dgrid as an alternative objective for function space inference. We can show that, for certain scenarios, Dgrid can avoid some of the issues the original functional KL 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 = g + ϵ, g(x) ∼ p(g|x; Θ),Θ ∼ p(Θ), ϵ ∼ N (0, σ2) and q(f) is parameterized by:
f = g + ϵ, g(x) ∼ q(g|x; Γ),Γ ∼ q(Γ), ϵ ∼ N (0, σ2). Here, x ∈ T ⊂ Rd is the input variable of the f , g is the random latent function, and Θ ∈ RI , Γ ∈ RJ are the parameters of each random function, respectively. Suppose p(g|x; Θ), q(g|x; Θ), p(Θ) and q(Γ) all have compact supports w.r.t. g, g, Θ, and Γ, respectively. Then, there exist a sampling distribution c such that: 1) c has full support on T Z+ , and 2) Dgrid[q(f)||p(f)] is finite.
This result shows that the grid-functional KL divergence allows us to 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 KL is no longer finite. This validates our choice of grid-functional KL divergence. To use Dgrid[q(f)||p(f |D)] for VI, we further derive a new ELBO based on Dgrid (Appendix A.1): Proposition 3. Let n,Xn ∼ c be a set of random measure points such that XD ⊂ Xn. Define
Lgridq := log p(D)−Dgrid[q(f)||p(f |D)]. (5) Then we have
Lgridq = Eq(f)[log p(D|f)]−Dgrid[q(f)||p(f)] (6)
and log p(D) ≥ Lgridq ≥ Lfunctionalq .
Proposition 3 shows that Lgridq is a valid variational objective function: it is a lower bound for log p(D) and also upper-bounds Lfunctionalq . For the rest of the paper, we will discuss how to perform functional VI based on Lgridq . We will focus on how to propose a expressive variational family q(f) and how to efficiently estimate Dgrid[q(f)||p(f)]. Remark. One example of c that satisfies the requirement of Propositions 1, 2, and 3 takes the following form (which will be used throughout the paper):
(n− |D|) ∼ Geom(p),xk ∼ U(T ), ∀1 ≤ k ≤ n− |D|,Xn := XD ⋃ {xk}1≤k≤n−|D|, (7)
where we first sample n from a geometric distribution such that (n−|D|) ∼ Geom(p) with parameter p (see Appendix A.2 for more discussion). Then, (n− |D|) out of distribution (OOD) measure points are sampled independently from a uniform distribution on T .
4.2 Choosing q(f): stochastic process generators
In order to obtain good performance in FVI, it is crucial to 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 non-Gaussian extension of the variational approximation used in VIPs. Recall that in VIPs, the variational family qVIP(f) is defined by the following sampling process:
f = ∑ s asϕs(·), a ∼ N (a;µ,Σ), (8)
where {ϕs}Ss=1 are random paths sampled from the prior process p(f). In SPGs, we remove the Gaussian assumption on a. This is done by specifying a separate variational auto-encoder to represent the non-Gaussian distribution for a:
q(a) = ∫ h pθ(a|h)qη(h), (9)
where h is the latent variable of the VAE, pθ(a|h) is the decoder parameterized by θ, and qη(h) is a distribution over the latent space. Furthermore, we remove the constraint that {ϕs}Ss=1 need to be sampled from p(f). For example, we may assume that each ϕs : Rd 7→ R1 is a DNN with learnable weights ws. We can finally give the definition of our SPG variational family qSPG(f |qη(h)):
f = ∑ s asϕs(·,ws) + ν, a ∼ ∫ h pθ(a|h)qη(h), ν ∼ GP(ν; 0, δ(·, ·)σ2ν), (10)
where ν is a white noise process that models additive aleatoric uncertainty. Intuitively, qSPG serves as a generator for stochastic processes. qSPG(f |·) maps any given qη(h) to a stochastic process qSPG(f |qη(h)), hence the name. Regarding the expressiveness of SPGs, we have the following result: Proposition 4 (Expressiveness of SPGs). Let p(f) be a square-integrable stochastic process on a probability space (X ,B), and its index set T is a compact subset of Rd. Suppose X is a compact metric space, B is the Borel set on X . Then, for ∀ϵ > 0, there exists a SPG qϵSPG(f), such that:
MMD(p, qϵSPG;F) < ϵ for ∀x ∈ T , (11)
where MMD is the maximum mean discrepancy, F is the MMD function class defined to be a unit ball in a RKHS with a universal kernel [46] k(·, ·) as its reproducing kernel.
Next we will discuss how {ws}, θ and qη(h) can be estimated and how can we use SPGs to estimate the function space KL-divergence in Equation 3.

heading: 4.3 Efficient estimation of grid-functional KL
text: In order to estimate the grid-functional KL-divergence in Equation 4, we propose to use a two-step method for a scalable approximation. In the first step, we distill p(f) by fitting a stochastic process generator p̃SPG(f) to p(f). In the second step, we calculate the KL-divergence between qSPG(f) and p̃SPG(f) as the surrogate for the KL-divergence between qSPG(f) and p(f).
Distilling p(f) via a stochastic process generator. Assume that we can draw M random functions [f1(·), f2(·), ..., fm(·), ..., fM (·)] from the prior process p(f). Let p̃SPG(f) be an SPG given by
f = ∑ s asϕs(·,ws) + ν, a ∼ ∫ h pθ(a|h)p0(h), ν ∼ GP(ν; 0, δ(·, ·)σ2ν), (12)
where p0(h) is a fixed standard normal distribution. We denote the above process by p̃SPG(f |p0(h)). We can train p̃SPG(f) on f1, f2, ..., fm, ..., fM , by optimizing the aggregated ELBO on p̃SPG(f):
max {ws},θ,λ EXO ∑ m log p̃SPG(f XO m ) ≥ max{ws},θ,λ EXO ∑ m E q̃λ(h|f XO m ) log p̃SPG(f XO m |h)p0(h) q̃λ(h|fXOm ) , (13)
where fXm are the function values of fm evaluated on XO, XO are |O| ≤ |D| measure points independently sampled from the training set XD = {xi}Ni=1. fXm are the function values of fm evaluated on XO, and q̃λ(h|fXOm ) is an encoder network that approximates the true posterior p̃SPG(h|fXOm ). Product of Experts (PoE) encoder. When sampling XO, since its size might vary each time, we would need to set up to 2N 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
q̃λ(h|fXOm ) ∝ p0(h) |O|∏ i=1 q̃λ(h|fm(xi),xi), (14)
where q̃λ(h|fm(xi),xi) is an inference network representing the expert associated with the i-th measurement point. Now the we can use a single encoder q̃λ(h|fXOm ) to handle all the possible inputs fXOm . In practice, we let q̃λ(h|fm(xi),xi) to be a Gaussian expert that maps [fm(xi),xi] to a factorized Gaussian in latent space. Since the product of Gaussian experts is still Gaussian, q̃λ(h|fXOm ) is a Gaussian distribution whose statistics can be computed analytically. Estimating the grid-functional KL-divergence given Xn. In order to estimate the grid-functional KL divergence between qSPG(f) and p(f), we first discuss how this divergence can be estimated on measurement points Xn, i.e., DKL[q(fXn)||p(fXn)] where fXn is the vector of function values
evaluated on Xn. We then discuss how this can be used to estimate the grid-functional divergence in Equation 4. To begin with, as in Section 4.2, our variational family is given by
f = ∑ s asϕs(·,ws) + ν, a ∼ q(a) = ∫ h pθ(a|h)qη(h), ν ∼ GP(ν; 0, σ2ν). (15)
We denote the above variational family by qSPG(f |qη(h)). The key ingredient of our estimation method is that we force qSPG(f |qη(h)) and p̃SPG(f |p0(h)) to share the same basis functions (or weights {ws}) and decoder parameters θ. That is, once optimal {ws} and θ are obtained by fitting p̃SPG(f) to p(f), these are frozen and reused in qSPG(f). This makes sense, since according to our definition in Section 3, qSPG(f) and p̃SPG(f) share the same measurable space (RT ,BTR ). Therefore, the only difference between qSPG(f) and p̃SPG(f) is the choice of the prior distributions on h, which is qη(h) and p0(h), respectively. Given this property, we can compute the KL divergence between qSPG(f) and p̃SPG(f) given measurement points Xn (Appendix A.5):
Proposition 5 (KL divergence on measurement points between SPGs). Let qSPG(f) and p̃SPG(f) be the SPGs defined in Equation 12 and 15. Then we have
DKL[qSPG(f Xn)||p̃SPG(fXn)] = Ef∼qSPG(f) logZ(f Xn), (16)
where Z(fXn) is the partition function, Z(fXn) = ∫ h p̃SPG(h|fXn) qη(h)p0(h)dh.
Note that Z(fXn) is intractable to compute due to the intractability of the posterior p̃SPG(h|fXn). Fortunately, this is already approximated by the PoE inference net q̃λ(h|fXn) given by Equation 14:
Z(fXn) ≈ Z̃(fXn) := ∫ h q̃λ(h|fXn) qη(h) p0(h) dh. (17)
Since qη(h), p0(h), and q̃λ(h|fXn) are all Gaussian distributions, Z̃(fXn) can be computed using analytic solutions. Note also that thanks to the VAE-like structure in SPGs, all the calculations are performed in the latent space, whose dimensionality is much lower than fX. With the additional help of analytic solutions for Z̃(fX), the estimation of (17) is very efficient and scalable.

heading: 4.4 The final algorithm: Mini-batching and de-biasing
text: So far we have been discussing the estimation of the KL-divergence on measurement points Xn. Next, we derive a practical algorithm based on the grid-functional ELBO Lgridq defined in Equation 6. Applying the approximation (17) to Equation 6, our final variational objective becomes
log p(D) ≥ |D|∑ i Eq(f)[log pπ(yi|f(xi))]− En,Xn∼cDKL[q(fXn)||p(fXn)]
≈ |D|∑ i Eq(f)[log pπ(yi|f(xi))]− En,Xn∼cEf∼qSPG(f) log Z̃(f Xn).
(18)
To make Equation 18 scalable to large data, we can apply mini-batch sampling to the likelihood term∑|D| i Eq(f)[log p(yi|f(xi))]. Then the only bottleneck of Equation 18 is that the input Xn to the inference net q̃λ (used in Z̃(fXn)) can be very high dimensional due to the condition XD ⊂ Xn required by Proposition 3. Fortunately, we can derive the following mini-batch estimators for En,Xn∼cEf∼qSPG(f) log Z̃(fXn) (Appendix A.6 and A.7):
Proposition 6. En,Xn∼cEf∼qSPG(f) log Z̃(fXn) can be estimated by the mini-batch estimator
JK := 1
2 H∑ i=1 Ef∼qSPG(f) [ log σ−2ηi + log σ̂ −2 λi − log (σ−2ηi + σ̂ −2 λi − 1)− µ̂2λi σ̂ −2 λi − µ2ηiσ −2 ηi
+ (σ̂−2ηi µ̂ηi + σ̂ −2 λi µ̂λi) 2(σ−2ηi + σ̂ −2 λi
− 1)−1 ] ,
(19)
Algorithm 1 Functional Variational Inference (FVI) Require: data D = {xi, yi}Ni=1; prior p(f); surrogate p̃SPG(f), variational process qSPG(f), likelihood function pπ(y|f), mini-batch sizes I and K
1: while not converged do 2: Sample [f1(·), f2(·), ..., fM (·)] from p(f). 3: Improve p̃SPG(f) by optimizing the aggregated ELBO in Equation 13 w.r.t. {ws}, θ, λ. 4: end while 5: while not converged do 6: sample mini-batch I from {1, ..., |D|} and a set of measure points Xn via c. 7: Optimize L̂FVI in Equation 20 w.r.t. η and π via reparameterization tricks 8: end while
where H is the dimensionality of h, N (h;µηi , σ2ηi) = qη(hi), N (h;µλi , σ 2 λi ) = q̃λ(hi|fX). σ̂−2λi and µ̂λi are the mini-batch approximators for µλi and σ 2 λi , respectively:
σ̂−2λi := ∑ k∈K |D| K σ−2hi|fxk + ∑
xl∈Xn\XD
σ−2hi|fxl
µ̂λi σ̂2λi := ∑ k∈K |D| K σ−2hi|fxkµhi|f xb + ∑ xl∈Xn\XD σ−2hi|fxlµhi|f xl ,
where K is a mini-batch of size K sampled from {1, ..., |D|}, xl ∈ Xn \XD is a set of OOD samples sampled from T using c in Eq. 7, and µhi|fxk and σ2hi|fxk are the mean and variance parameter returned from q̃λ(hi|f(xk)).
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 a random integer from a distribution P(N) with support over the integers larger than K. x0 is a random location sampled from T . Then En,Xn∼cEf∼qSPG(f) log Z̃(fXn) can be estimated by E [ JK + ∑R k=K ∆k P(N≥k) ] , where ∆k = Jk+1 −Jk, and the expectation is taken over R, n, Xn, and all mini-batches used by each Jk terms.
This will enable us to also perform mini-batch sampling on the measurement points when performing FVI. Our final optimization objective function is
L̂FVI := |D| I ∑ i∈I Eq(f)[log p(yi|f(xi))]− JK − R∑ k=B ∆k P(N ≥ k) , (20)
where I is a mini-batch of size I for the likelihood terms, R is an integer sampled from P(N), which is set to be (R−K) ∼ Geom(0.5). Finally, the full algorithm is sketched in Algorithm 1. We call this proposed method Functional Variational Inference (FVI).
Scalability. Our method is empirically much faster than f-BNN (Appendix C.3). When estimating Dgrid[q(f)||p(f)], our method scales as O(Mq), where Mq is the number of samples sampled from qSPG(f) that are used in JK . In practice, we use Mq = 1. On the contrary, the SSGE estimator used in f-BNN scales as O(M3q +M2qD), where much larger value needs to used (e.g., Mq = 100).

heading: 5 Related Works
text: Since BNNs and VIPs are discussed in Section 1 and 2, here we only address the rest of the related works including f-BNNs, functional priors, and Bayesian non-parameterics.
F-BNNs and F-PO. The functional BNN [47] is proposed to address the issue of specifying meaningful priors to BNN weights. It matches a BNN to a GP prior by minimizing the functional KL divergence estimated by score function estimators [27]. As discussed in [5], this objective is not well-defined for a wide class of distributions. Also, the score function estimators used in f-BNNs
often perform poorly in high dimensional spaces and are less efficient. In practice, we found that the f-BNN computational time is 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 is approximated via linearization. In their problem setting, the functional KL divergence is well defined since both p and q are BNNs with the same neural network structures. Similarly, f-POVI [50] performs inference on BNN priors 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-learning 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 of function space. On the contrary, our method focuses on 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 qη(h), which is a simple Gaussian distribution (no set encoders involved).
Bayesian non-parametrics. In the field of Bayesian non-parameterics, 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 a 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 a critical role in this line of work, which makes them less applicable to general non-GP based priors.

heading: 6 Experiments
text: 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.

heading: 6.1 Interpolation with non-Gaussian priors: structured implicit priors
text: An advantage of FVI is that it 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
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 are 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.e., producing samples from p(f |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-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 piecewise 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-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.

heading: 6.2 Multivariate regression with BNNs priors
text: 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: Bayes-by-Backprop [4], variational dropout [12], and variational alpha dropout [26] (α = 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 7 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
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.

heading: 6.3 Contextual Bandits
text: Uncertainty estimates are important for downstream decision-making scenarios, since explorationexploitation is a common dilemma that must be addressed. In this section, we consider a classic task called contextual bandits, where the agent is asked to make decisions that maximize the reward given some contexts (inputs). For this, Thompson sampling [48] is an elegant approach to guide exploration, where a 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 hyperparameter settings are consistent with [47] except that we used a smaller batchsize (32). The learning rates for each baseline are tuned from [0.001, 0.05]. We report the cumulative regret as well as the 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 6 out of 7 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).

heading: 6.4 Image classification and out-of-distribution detection
text: To demonstrate the scalability of our method to higher dimensional data, we consider image classification tasks on Fashion MINIST [54] and CIFAR10 [25] with BNN priors. We compare our method to the following baselines: mean field VI (MFVI), maximum a 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 as in [21, 43]. Apart from test accuracy and negative log likelihood (NLL), we also perform out-of-distribution detection using in-distribution (ID) / outof-distribution (OOD) pairs including FashionMNIST/MNIST and CIFAR10/SVNH. Following the settings of [35, 21], we calculate the area under the curve (AUC) of out-of-distribution 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.

heading: 7 Conclusion
text: In this paper, we propose Functional Variational Inference (FVI), a new VI method in function space. It optimizes a grid-based functional divergence, which can be 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.

heading: Acknowledgements
text: We thank Javier Antoran and Yingzhen Li for helpful discussions and comments during the early stage of this work.

Review summaries:
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.  Furthermore, 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.   The 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.

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.    

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.

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.