Functional Variational Inference based on Stochastic Process Generators

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-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 . 2 Backgrounds . 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 . 3 Problem setting and the functional KL divergence . 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+ . 4 Functional Variational Inference using Stochastic Process Generators . 4.1 The grid-functional KL divergence . 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 .

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.  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.