FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS

1 INTRODUCTION . Bayesian neural networks ( BNNs ) ( Hinton & Van Camp , 1993 ; Neal , 1995 ) have the potential to combine the scalability , flexibility , and predictive performance of neural networks with principled Bayesian uncertainty modelling . However , the practical effectiveness of BNNs is limited by our ability to specify meaningful prior distributions and by the intractability of posterior inference . Choosing a meaningful prior distribution over network weights is difficult because the weights have a complicated relationship to the function computed by the network . Stochastic variational inference is appealing because the update rules resemble ordinary backprop ( Graves , 2011 ; Blundell et al. , 2015 ) , but fitting accurate posterior distributions is difficult due to strong and complicated posterior dependencies ( Louizos & Welling , 2016 ; Zhang et al. , 2018 ; Shi et al. , 2018a ) . In a classic result , Neal ( 1995 ) showed that under certain assumptions , as the width of a shallow BNN was increased , the limiting distribution is a Gaussian process ( GP ) . Lee et al . ( 2018 ) recently extended this result to deep BNNs . Deep Gaussian Processes ( DGP ) ( Cutajar et al. , 2017 ; Salimbeni & Deisenroth , 2017 ) have close connections to BNNs due to similar deep structures . However , the relationship of finite BNNs to GPs is unclear , and practical variational BNN approximations fail to match the predictions of the corresponding GP . Furthermore , because the previous analyses related specific BNN architectures to specific GP kernels , it ’ s not clear how to design BNN architectures for a given kernel . Given the rich variety of structural assumptions that GP kernels can represent ( Rasmussen & Williams , 2006 ; Lloyd et al. , 2014 ; Sun et al. , 2018 ) , there remains a significant gap in expressive power between BNNs and GPs ( not to mention stochastic processes more broadly ) . In this paper , we perform variational inference directly on the distribution of functions . Specifically , we introduce functional variational BNNs ( fBNNs ) , where a BNN is trained to produce a distribution of functions with small KL divergence to the true posterior over functions . We prove that the KL divergence between stochastic processes can be expressed as the supremum of marginal KL divergences at finite sets of points . Based on this , we present functional ELBO ( fELBO ) training objective . Then we introduce a GAN-like minimax formulation and a sampling-based approximation for functional variational inference . To approximate the marginal KL divergence gradients , we adopt the recently proposed spectral Stein gradient estimator ( SSGE ) ( Shi et al. , 2018b ) . ∗Equal contribution . Our fBNNs make it possible to specify stochastic process priors which encode richly structured dependencies between function values . This includes stochastic processes with explicit densities , such as GPs which can model various structures like smoothness and periodicity ( Lloyd et al. , 2014 ; Sun et al. , 2018 ) . We can also use stochastic processes with implicit densities , such as distributions over piecewise linear or piecewise constant functions . Furthermore , in contrast with GPs , fBNNs efficiently yield explicit posterior samples of the function . This enables fBNNs to be used in settings that require explicit minimization of sampled functions , such as Thompson sampling ( Thompson , 1933 ; Russo & Van Roy , 2016 ) or predictive entropy search ( Hernández-Lobato et al. , 2014 ; Wang & Jegelka , 2017 ) . One desideratum of Bayesian models is that they behave gracefully as their capacity is increased ( Rasmussen & Ghahramani , 2001 ) . Unfortunately , ordinary BNNs don ’ t meet this basic requirement : unless the asymptotic regime is chosen very carefully ( e.g . Neal ( 1995 ) ) , BNN priors may have undesirable behaviors as more units or layers are added . Furthermore , larger BNNs entail more difficult posterior inference and larger description length for the posterior , causing degeneracy for large networks , as shown in Figure 1 . In contrast , the prior of fBNNs is defined directly over the space of functions , thus the BNN can be made arbitrarily large without changing the functional variational inference problem . Hence , the predictions behave well as the capacity increases . Empirically , we demonstrate that fBNNs generate sensible extrapolations for both explicit periodic priors and implicit piecewise priors . We show fBNNs outperform competing approaches on both small scale and large scale regression datasets . fBNNs ’ reliable uncertainty estimates enable state-of-art performance on the contextual bandits benchmark of Riquelme et al . ( 2018 ) . 2 BACKGROUND . 2.1 VARIATIONAL INFERENCE FOR BAYESIAN NEURAL NETWORKS . Given a dataset D = { ( xi , yi ) } ni=1 , a Bayesian neural network ( BNN ) is defined in terms of a prior p ( w ) on the weights , as well as the likelihood p ( D|w ) . Variational Bayesian methods ( Hinton & Van Camp , 1993 ; Graves , 2011 ; Blundell et al. , 2015 ) attempt to fit an approximate posterior q ( w ) to maximize the evidence lower bound ( ELBO ) : Lq = Eq [ log p ( D|w ) ] −KL [ q ( w ) ‖p ( w ) ] . ( 1 ) The most commonly used variational BNN training method is Bayes By Backprop ( BBB ) ( Blundell et al. , 2015 ) , which uses a fully factorized Gaussian approximation to the posterior , i.e . q ( w ) = N ( w ; µ , diag ( σ2 ) ) . Using the reparameterization trick ( Kingma & Welling , 2013 ) , the gradients of ELBO towards µ , σ can be computed by backpropagation , and then be used for updates . Most commonly , the prior p ( w ) is chosen for computational convenience ; for instance , independent Gaussian or Gaussian mixture distributions . Other priors , including log-uniform priors ( Kingma et al. , 2015 ; Louizos et al. , 2017 ) and horseshoe priors ( Ghosh et al. , 2018 ; Louizos et al. , 2017 ) , were proposed for specific purposes such as model compression and model selection . But the relationships of weight-space priors to the functions computed by networks are difficult to characterize . 2.2 STOCHASTIC PROCESSES . A stochastic process ( Lamperti , 2012 ) F is typically defined as a collection of random variables , on a probability space ( Ω , F , P ) . The random variables , indexed by some set X , all take values in the same mathematical space Y . In other words , given a probability space ( Ω , Σ , P ) , a stochastic process can be simply written as { F ( x ) : x ∈ X } . For any point ω ∈ Ω , F ( · , ω ) is a sample function mapping index space X to space Y , which we denote as f for notational simplicity . For any finite index set x1 : n = { x1 , ... , xn } , we can define the finite-dimensional marginal joint distribution over function values { F ( x1 ) , · · · , F ( xn ) } . For example , Gaussian Processes have marginal distributions as multivariate Gaussians . The Kolmogorov Extension Theorem ( Øksendal , 2003 ) shows that a stochastic process can be characterized by marginals over all finite index sets . Specifically , for a collection of joint distributions ρx1 : n , we can define a stochastic process F such that for all x1 : n , ρx1 : n is the marginal joint distribution of F at x1 : n , as long as ρ satisfies the following two conditions : Exchangeability . For any permutation π of { 1 , · · · , n } , ρπ ( x1 : n ) ( π ( y1 : n ) ) = ρx1 : n ( y1 : n ) . Consistency . For any 1 ≤ m ≤ n , ρx1 : m ( y1 : m ) = ∫ ρx1 : n ( y1 : n ) dym+1 : n . 2.3 SPECTRAL STEIN GRADIENT ESTIMATOR ( SSGE ) . When applying Bayesian methods to modern probabilistic models , especially those with neural networks as components ( e.g. , BNNs and deep generative models ) , it is often the case that we have to deal with intractable densities . Examples include the marginal distribution of a non-conjugate model ( e.g. , the output distribution of a BNN ) , and neural samplers such as GANs ( Goodfellow et al. , 2014 ) . A shared property of these distributions is that they are defined through a tractable sampling process , despite the intractable density . Such distributions are called implicit distributions ( Huszár , 2017 ) . The Spectral Stein Gradient Estimator ( SSGE ) ( Shi et al. , 2018b ) is a recently proposed method for estimating the log density derivative function of an implicit distribution , only requiring samples from the distribution . Specifically , given a continuous differentiable density q ( x ) , and a positive definite kernel k ( x , x′ ) in the Stein class ( Liu et al. , 2016 ) of q , they show ∇xi log q ( x ) = − ∞∑ j=1 [ Eq∇xiψj ( x ) ] ψj ( x ) , ( 2 ) where { ψj } j≥1 is a series of eigenfunctions of k given by Mercer ’ s theorem : k ( x , x′ ) =∑ j µjψj ( x ) ψj ( x ′ ) . The Nyström method ( Baker , 1997 ; Williams & Seeger , 2001 ) is used to approximate the eigenfunctions ψj ( x ) and their derivatives . The final estimator is given by truncating the sum in Equation ( 2 ) and replacing the expectation by Monte Carlo estimates . 3 FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS . 3.1 FUNCTIONAL EVIDENCE LOWER BOUND ( fELBO ) We introduce function space variational inference analogously to weight space variational inference ( see Section 2.1 ) , except that the distributions are over functions rather than weights . We assume a stochastic process prior p over functions f : X → Y . This could be a GP , but we also allow stochastic processes without closed-form marginal densities , such as distributions over piecewise linear functions . For the variational posterior qφ ∈ Q , we consider a neural network architecture with stochastic weights and/or stochastic inputs . Specifically , we sample a function from q by sampling a random noise vector ξ and defining f ( x ) = gφ ( x , ξ ) for some function gφ . For example , standard weight space BNNs with factorial Gaussian posteriors can be viewed this way using the reparameterization trick ( Kingma & Welling , 2013 ; Blundell et al. , 2015 ) . ( In this case , φ corresponds to the means and variances of all the weights . ) Note that because a single vector ξ is shared among all input locations , it corresponds to randomness in the function , rather than observation noise ; hence , the sampling of ξ corresponds to epistemic , rather than aleatoric , uncertainty ( Depeweg et al. , 2017 ) . Functional variational inference maximizes the functional ELBO ( fELBO ) , akin to the weight space ELBO in Equation ( 1 ) , except that the distributions are over functions rather than weights . L ( q ) : = Eq [ log p ( D|f ) ] −KL [ q||p ] . ( 3 ) Here KL [ q‖p ] is the KL divergence between two stochastic processes . As pointed out in Matthews et al . ( 2016 ) , it does not have a convenient form as ∫ log q ( f ) p ( f ) q ( f ) df due to there is no infinitedimensional Lebesgue measure ( Eldredge , 2016 ) . Since the KL divergence between stochastic processes is difficult to work with , we reduce it to a more familiar object : KL divergence between the marginal distributions of function values at finite sets of points , which we term measurement sets . Specifically , let X ∈ Xn denote a finite measurement set and PX the marginal distribution of function values at X . We equate the function space KL divergence to the supremum of marginal KL divergences over all finite measurement sets : Theorem 1 . For two stochastic processes P and Q , KL [ P‖Q ] = sup n∈N , X∈Xn KL [ PX‖QX ] . ( 4 ) Roughly speaking , this result follows because the σ-algebra constructed with the Kolmogorov Extension Theorem ( Section 2.2 ) is generated by cylinder sets which depend only on finite sets of points . A full proof is given in Appendix A. fELBO . Using this characterization of the functional KL divergence , we rewrite the fELBO : L ( q ) = Eq [ log p ( D|f ) ] − sup n∈N , X∈Xn KL [ q ( fX ) ||p ( fX ) ] = inf n∈N , X∈Xn ∑ ( xi , yi ) ∈D Eq [ log p ( yi|f ( xi ) ) ] −KL [ q ( fX ) ||p ( fX ) ] : = inf n∈N , X∈Xn LX ( q ) . ( 5 ) We also denote Ln ( q ) : = inf X∈Xn LX ( q ) for the restriction to sets of n points . This casts maximizing the fELBO as a two-player zero-sum game analogous to a generative adversarial network ( GAN ) ( Goodfellow et al. , 2014 ) : one player chooses the stochastic network , and the adversary chooses the measurement set . Note that the infimum may not be attainable , because the size of the measurement sets is unbounded . In fact , the function space KL divergence may be infinite , for instance if the prior assigns measure zero to the set of functions representable by a neural network ( Arjovsky & Bottou , 2017 ) . Observe that GANs face the same issue : because a generator network is typically limited to a submanifold of the input domain , an ideal discriminator could discriminate real and fake images perfectly . However , by limiting the capacity of the discriminator , one obtains a useful training objective . By analogy , we obtain a well-defined and practically useful training objective by restricting the measurement sets to a fixed finite size . This is discussed further in the next section .

Summary: The primary contribution of this paper is the presentation of a novel ELBO objective for training BNNs which allows for more meaningful priors to be encoded in the model rather than the less informative weight priors featured in the literature. This is achieved by way of introducing a KL measure over stochastic processes which allows for priors to take the form of GP priors and other custom variations. Two approaches are given for training the model, one inspired by GANs, and a more practical sampling-based scheme. The performance of this training scheme is validated on a variety of synthetic and real examples, choosing Bayes by Backprop as the primary competitor. An experiment on contextual bandit exploration, and an illustrative Bayesian optimisation example  provided in the supplementary material showcase the effectiveness of this method in applications where well-calibrated uncertainty is particularly pertinent.