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