Functional Bayesian Neural Networks for Model Uncertainty Quantification

1 INTRODUCTION . Though deep neural networks have achieved tremendous success in various domains recently , there still exists some fundamental issues largely unclear and unsolved . One of them is that the training of neural networks usually ignores the uncertainty over the parameters since typically maximum a posteriori ( MAP ) is adopted for learning the parameters and only point estimates could be obtained . This problem can cause neural networks to make overconfident predictions , particularly in the region poorly covered by training data or far away from the underlying data manifold . Besides , the ignorance of model uncertainty also make neural networks vulnerable to adversarially perturbed data , so-called “ adversarial examples ” ( Szegedy et al. , 2013 ; Gal & Smith , 2018 ) . The Bayesian perspective provides a principled way of accounting for the uncertainty of model ’ s prediction by integrating over the posterior distribution over parameters . In the scenario of neural networks , Bayesian neural networks ( MacKay , 1992 ; Hinton & Van Camp , 1993 ; Neal , 1995 ) was proposed . However , for modern neural networks , the integral over posterior distribution is highly intractable due to the large network size and complex multi-modality . There are mainly two groups of approaches for handling the intractability of the complex highdimensional distribution . The first one is to employ various Markov Chain Monte Carlo ( MCMC ) ( Neal ( 1993 ) ; Max & Whye ( 2011 ) ; Chen et al . ( 2014 ) ; Shang et al . ( 2015 ) ; Ye et al . ( 2017 ) ) to obtain the samples from the distribution ; and then the Monte Carlo approximation can be used . However , the high dimensionality and a large number of modes in Bayesian neural networks prohibit the scaling of current MCMC methods , i.e . the mixing is extremely time-consuming . The other method is variational inference , typically approximating the original complex posterior distribution with independent Gaussian distributions over each weight in networks ( Graves ( 2011 ) ; Hernández-Lobato & Adams ( 2015 ) ; Blundell et al . ( 2015 ) ; Gal & Ghahramani ( 2016 ) ) . The limitation of this strategy is that the introduced approximation might underestimate the posterior when optimizing the variational lower bound . In this paper , we extend the Bayesian neural network to functional Bayesian neural network for MCMC methods . Instead of sampling a single parameter each step , we propose to sample a function which can represents the posterior distribution each step . Based on the previous work for functional Hamiltonian method , we extend a previous work on Riemannian Langevin dynamics and propose its functional extension with mean-field Gaussian approximation . 2 FUNCTIONAL BAYESIAN NEURAL NETWORK . In this section , we will first explain the preliminary of Bayesian neural network and then extend it into the functional space . Neural networks are most commonly trained in a maximum a posteriori ( MAP ) setting , where only point estimates of the network parameters can be obtained . This ignores any uncertainty about the parameters that often result in overconfident predictions , especially in regimes that are weakly covered by training data or far away from the data manifold . Bayesian neural networks Neal ( 1995 ) incorporates neural networks into the Bayesian framework for accounting for the parameter uncertainty through modeling the posterior distribution . Firstly , one need to provide a prior distribution over the weights , p0 ( θ ) = N ( 0 , σ20I ) , where σ20 is the variance magnitude . Assuming the likelihood function has the form , p ( yi|xi , θ ) = N ( yi|u ( xi ; θ ) , σ2I ) , ( 1 ) where u ( x ; θ ) represents the output of the neural network . Then the posterior of the weights θ is p ( θ|D ) ∝ ∏N i=1 p ( yi|xi , θ ) p0 ( θ ) . The uncertainty of the model , typically formulated as the expectation of a specific statistics g ( θ ; x , y ) . could be computed based on the posterior distribution of the weights , µg = E [ g ( θ ; x , y ) ] = ∫ g ( θ ; x , y ) p ( θ|D ) dθ . ( 2 ) Due to the analytic intractability of the integral , the expectation of g ( θ ; x , y ) is estimated using Monte Carlo integration µg ≈ 1 T T∑ i=1 g ( θi ; x , y ) , ( 3 ) where { θi } Mi=1 is drawn from the posterior p ( θ|D ) . To derive the functional Bayesian neural networks , we first define a functional f : θ → [ 0 , 1 ] which maps a parameter setting θ to a probability , or more formally , we define a separable Hilbert space H with inner product and norm operation and the functional f : θ → [ 0 , 1 ] is inH . Note that f is a probability distribution of weights θ . Similar to Bayesian neural networks , suppose we get posterior distribution of functionals given the data : p ( f |D ) ∝ ∏N i=1 p ( yi|xi , f ) p0 ( f ) , where p0 ( f ) is the prior distribution of f , the uncertainty of the model can be formulated as the expectation of a specific statistics g ( θ ; x , y ) : µg = E [ g ( f ; x , y ) ] = ∫ g ( f ; x , y ) p ( f |D ) df , ( 4 ) where g ( f ; x , y ) = ∫ g ( θ ; x , y ) f ( θ ) dθ is the expectation of the g ( θ ; x , y ) given the distribution functional f . Due to the analytic intractability of the integral , similarly , the expectation of g ( θ ; x , y ) can be estimated using a two-loop Monte Carlo integration , µg ≈ 1 T T∑ i=1 M∑ j=1 g ( θij ; x , y ) , Sample θij from fi , ( 5 ) where we sample M parameters from each posterior fi and T posteriors is used for computing the expectation . Through this formulation , we can perform inference with functional Bayesian neural networks . Note that with the same number of iterations T , the functional Bayesian neural network can have M times more samples than Bayesian neural networks . We provide a functional sampling example in section 5.1 for illustration . In the next parts , we will propose the stochastic gradient functional Riemannian Langevin dynamics ( SGFuncRLD ) for training functional Bayesian neural network . 3 FUNCTIONAL RIEMANNIAN LANGEVIN DYNAMICS . Beskos et al . ( 2011 ) first introduced the concept of functional Monte Carlo method on Hilbert space with ( second-order ) Langevin dynamics . Beskos et al . proved that if the target distribution U ( θ ) can be represented by the sum of weighted sequence of functionals f with the coefficients of functionals decreasing faster than a polynomial rate , along with other necessary assumptions , sampling the functionals can be equivalent to sampling the posterior U ( θ ) directly . However , the HMC method can not adapt to the local geometry and need to tune parameters carefully and in real practice , tuning hyper-parameter can be costly . We want the training algorithm to be scalable and robust to different hyper-parameter settings , we thus extend the Riemannian Langevin dynamics to functional space to make the functional MCMC method more robust . Similar to Beskos et al . ( 2011 ) , we represent the posterior U ( θ ) with its finite approximation UD ( θ ) = ∑D i λiui . We denote the parameter of UD ( θ ) to be λ ∈ RD thus given λ ∈ RD , we have a mapping from λ to UD ( θ ) . Note that this mapping is continuous and bounded . Then by sampling λ , we sample a functional f equivalently . The Riemannian Langevin dynamics on the functional space can thus be written as : dλ = b ( λ ) dt+ √ 2G ( λ ) dW ( 6 ) where dW represents the standard Brownian motion , and the drift force b ( λ , α ) has the following form , b ( λ ) = −G ( λ ) ∂U ( λ ) ∂λ − ξ ( λ ) ( 7 ) ξi ( λ ) = − d∑ j=1 ∂ ∂λj Gij ( λ ) , ( 8 ) where Gij is the Riemannian information metric . Similar to its counterpart defined in real number space , we have the following theorem , Theorem 1. p ( λ ) ∝ exp ( −U ( λ ) ) is a stationary distribution of the dynamics of Eq . ( 6 ) , if G ( λ ) is a positive semidefinite matrix . Proof . The Fokker-Planck equation of the dynamics in Eq . ( 6 ) , ∂tρ ( λ , t ) =− ∑ i ∂ ∂λi ∑ j bj ( λ , α ) ρ ( λ , t ) +∑ i , j ∂2 ∂λiλj 2Gijρ ( λ , t ) ( 9 ) We insert b ( λ ) = −G ( λ ) ∂U ( λ ) ∂λ − ξ ( λ ) into the above equation , the Fockker-Planck equation can be simplified as , ∂tρ ( λ , t ) = ∑ i ∂ ∂λi ∑ j { Gij ∂U ( λ ) ∂λj − ∂Gij ∂λj } ρ ( λ , t ) +∑ i , j ∂2 ∂λiλj 2Gijρ ( λ , t ) ( 10 ) When ρ ( λ , t ) = p ( λ ) , the stationary distribution will be reached ( ∂tρ ( λ , t ) = 0 ) . Then we observe that right hand side of the Fokker Planck equation vanishes . Thus , when p ( λ , t ) ∝ exp ( −U ( λ ) ) , ∂tρ ( λ , t ) = 0 . To simulate Functional Riemannian Langevin dynamics for neural networks , we use the following approximate computation for stochastic functional Riemannian Langevin dynamics ( SGFuncRLD ) for neural networks . 4 APPROXIMATE COMPUTATION . 4.1 FINITE FUNCTIONAL APPROXIMATION . We use the mean-field Gaussian approximation to approximate U ( θ ) . Mean-field Gaussian approximation has been widely used for neural networks and shows good performance for Bayesian inference for neural networks , such as in variational inference ( Blei et al . ( 2017 ) ) . We use this approximation and thus U ( θ ) is represented in a simple one-line form , p ( θ|D ) ≈ p ( f ) = N∏ i=1 p ( θi|D ) ( 11 ) where p ( θi|D ) = N ( µi , σi|D ) and N ( µi , σi ) is a normal distribution with mean µi and standard deviation σi . Note that unlike variational inference , we do not have the problem of zero forcing or zero avoiding ( Murphy ( 2013 ) ) . 4.2 STOCHASTIC COMPUTATION APPROXIMATION . We use the Fisher information metric as the Riemman information metric ( RIM ) as in the previous work ( Li et al. , 2016 ) : G ( λ ) = E [ ∂U ( λ ) ∂λ ∂U ( λ ) ∂λ T ] ( 12 ) We only use the diagonal part of the RIM and set the remaining part to be zero for computationally purpose , Gii ( λ ) = E [ ( ∂U ( λ ) ∂λi ) 2 ] ( 13 ) We approximate the gradients on the dataset with the smoothed gradients on mini-batches for scalable computation . Besides , as we only need to approximate form of RIM , we do not consider ξ ( θ ) for scalable computation on big data . We use the gradient computed on the k-th mini-batch { xk1 , . . . , xkm } to approximate the gradient , Ũ ( λ ) = −N m m∑ j=1 log p ( xkj |λ ) − log p0 ( λ ) . ( 14 ) However , previous work in optimization literature shows that the smoothed and bias-corrected form of gradient shows more stable performance than this unbiased stochastic approximation when scaled to large datasets ( Kingma & Ba ( 2014 ) ) , we then approximate the gradient with smoothed and bias-corrected version , mt = β1mt−1 + ( 1− β1 ) ∇λŨ ( λt ) ( 15 ) mt = mt/ ( 1− βt1 ) ( 16 ) where mt is the smoothed gradient at step t , β1 is the smoothing coefficient . We also apply the same technique for computing the RIM . Gt ( λ ) = β2Gt−1 ( λ ) + ( 1− β2 ) ( ∇λŨ ( λt ) ) 2 ( 17 ) Gt ( λ ) = Gt ( λ ) / ( 1− βt2 ) ( 18 ) where Gt is the smoothed gradient at step t , β2 is the smoothing coefficient .

Summary: The idea of extending  Riemannian Langevin dynamics to functional spaces is elegant, however it is extremely hard to follow the proposed method as details are kept to a minimum. The finite approximation of the posterior distribution is a function of the parameters theta, however it displays parameters lambda. The couple of sentences: "Then by sampling λ, we sample a functional f equivalently. The Riemannian Langevin dynamics on the functional space can thus be written as: (6)" come without a single explanation.