Year: 2018

Abstract:
In this paper, we extend the Bayesian neural network to functional Bayesian neural network with functional Monte Carlo methods that use the samples of functionals instead of samples of networks’ parameters for inference to overcome the curse of dimensionality for uncertainty quantification. Based on the previous work on Riemannian Langevin dynamics, we propose the stochastic gradient functional Riemannian dynamics for training functional Bayesian neural network. We show the effectiveness and efficiency of our proposed approach with various experiments.

heading: 1 INTRODUCTION
text: 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.

heading: 2 FUNCTIONAL BAYESIAN NEURAL NETWORK
text: 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.

heading: 3 FUNCTIONAL RIEMANNIAN LANGEVIN DYNAMICS
text: 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 cannot 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.

heading: 4 APPROXIMATE COMPUTATION

heading: 4.1 FINITE FUNCTIONAL APPROXIMATION
text: 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)).

heading: 4.2 STOCHASTIC COMPUTATION APPROXIMATION
text: 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.

heading: 4.3 STOCHASTIC GRADIENT FUNCTIONAL RIEMANNIAN DYNAMICS
text: With these approximations, we have the algorithm in Algorithm 1. After we derive samples of the coefficients λi for the mean-field Gaussian distribution, we can use Eq. 5 for inference.

heading: 5 EXPERIMENTS
text: In this section, we will first illustrate how functional MCMC method works with functional Hamiltonian dynamics and Hamiltonian dynamics on a double-well potential function. Then we will compare our proposed method with other methods including stochastic gradient Langevin dynamics (Max & Whye, 2011) (SGLD), stochastic gradient Hamiltonian Monte Carlo (Max & Whye, 2011)(SGHMC), preconditioned stochastic gradient Langevin dynamics (Li et al., 2016) (pSGLD) in various examples with neural networks. In the following experiments for neural networks, we set the number of (functional) samples T to be 5, and for functional Bayesian neural network, we set the number of samples from each functional M to be 5. For SGFuncRLD, we set β1 to be 0.9 and β2 to be 0.999. We implement the experiments with Tensorflow and Tensorflow-probability.
Algorithm 1: SGFuncRLD for training Bayesian neural networks
1: Input: β1, β2, η. 2: Initialize λ0, m0, G0. 3: for t = 1, 2, . . . do 4: Randomly sample a minibatch of the dataset to obtain Ũ(fλt): 5: Sample t ∼ N (0, I) from the standard Gaussian distribution ; 6: Compute the smoothed gradient mt;
mt = β1mt−1 + (1− β1)∇λŨ(λt), mt = mt/(1− βt1)
7: Compute the smoothed RIM Gt(θ);
Gt(λ) = β2Gt−1(λ) + (1− β2)(∇λŨ(λt))2, Gt(λ) = Gt(λ)/(1− βt2)
8: Update the parameter;
λt = λt−1 −Gt(λ)− 1 2 ηmt + √ 2Gt(λ)−1η t
9: end for

heading: 5.1 ILLUSTRATION OF FUNCTIONAL MCMC
text: To illustrate how functional MCMC methods work, we show an example of sampling double-well potential function. Double-well potential function is widely used for evaluating MCMC methods, the target distribution function used is: U(θ) = −2θ2 + 0.2θ4, similar to the one used in Chen et al. (2014). Note that the sampling iteration is set to be 10 instead of 80000× 50 in the original example (Chen et al. (2014)) to test the methods’ ability to efficiently sample the parameter space in a limited amount of time. For Functional MCMC, we set the initialization functional to be a standard normal distribution. We use the Hamiltonian dynamics for the Functional MCMC method-FuncHMC, which is the functional extension of the widely-used Hamiltonian Monte Carlo (HMC). For fair comparison, we use the same parameter for FuncHMC and HMC (learning rate is 0.2, momentum is 0.5). We show the sampled functional from FuncHMC at each step and HMC’s samples in Figure 1 (Left). From Figure 1 (Left), we can observe that different from HMC that can only generate one single sample of parameter at each time, FuncHMC generates one sample of distribution function at each time. Even with very few iterations, FuncHMC is able to sample the multi-modal distribution while HMC’s samples are concentrated near θ = 1. This property can help overcome the curse of dimensionality when sampling for high dimension distributions especially for Bayesian neural networks.
To compare quality of samples generated by FuncHMC or HMC, we use the KLIEP 1 to estimate the ratio distribution of benchmark samples to FuncHMC samples or HMC samples. We use HMC with a learning rate of 0.02 and a momentum of 0.5 and run it for 106 iterations and sample every 100 iterations which gives a total of 1e4 samples as the benchmark samples.
From Figure 1, we can observe that ratio distribution of FuncHMC varies close to the line of y = 1. Though HMC achieves good approximation close to θ = 1, HMC fails to correlate well with the benchmark samples within a broader range. For FuncHMC, we generate 100 samples from each posterior functional sampled and this gives us 1000 samples in total from 10 iterations. The plot of the ratio distribution is shown in Figure 1 (Right). From Figure 1 (Right), we can observe that ratio distribution of FuncHMC varies closely to the black horizontal line y = 1. Besides, HMC can achieve good approximation close to θ = 1, but it fails to correlate well with the benchmark samples across the distribution range in a limited amount of steps.
Next, we will compare our proposed method SGFuncRLD with other popular SG-MCMC methods for Bayesian neural networks on uncertainty estimation.

heading: 5.2 UNCERTAINTY ESTIMATION ON OUT-OF-DISTRIBUTION DATA
text: Accurate uncertainty estimation on unseen data is key to many safety-critical applications, such as autonomous driving and disease diagnosis. To measure the networks’ ability on uncertainty estimation for out-of-distribution data, we use a similar experiment setting as in Ritter et al. (2018). We train Bayesian neural networks and the functional Bayesian neural network on the MNIST dataset and use the trained models to predict on the unseen NoMNIST dataset2 that consists of images of letters from “A” to “J” with various fonts. We compare the entropy of predictions on the NoMNIST dataset. The minimum value of the predictive entropy is 0 and the maximum value of the predictive entropy is log(10) ≈ 2.30. We use the neural network architecture of conv (6 filters with a kernel size of 5)-pool-conv (16 filters with a kernel size of 5)-pool-conv (120 filters with a kernel siz of 5)-fully connected (84)-fully connected (10), which is the same as in the Tensorflow probability example 3. We train the network for 5000 iterations and use a batch size of 128, which is the same as in the official example. In the training phase, we set the learning rate of SGFuncRLD, pSGLD and Adam(Deterministic method) to be 0.001, SGLD and SGHMC to be 0.1, and the momentum for
1Sugiyama et al. (2008) http://www.ms.k.u-tokyo.ac.jp/software.html 2https://www.kaggle.com/lubaroli/notmnist 3https://github.com/tensorflow/probability/blob/r0.3/tensorflow probability/examples/bayesian neural
network.py
SGHMC to be 0.1. In the test phase, we draw five samples (functionals) for MCMC methods such as SGFuncRLD, pSGLD, SGHMC and SGLD; For SGFuncRLD, we further draw five samples for each functional; For deterministic training method, one sample is used. We average the predicted probability from samples as the final predicted probability. The empirical cumulative distribution function of predictive entropy on NoMNIST dataset is shown in Figure 2. From Figure 2, we can observe that all MCMC methods can achieve better uncertainty estimation performance than the deterministic method which indicates that Bayesian neural network is necessary for safety-critical applications requiring accurate uncertainty estimation. SGFuncRLD achieves the best performance by predicting the unseen examples with high entropy.

heading: 6 UNCERTAINTY ESTIMATION ON ADVERSARIAL ATTACK DATA
text: Though neural networks have achieved great success in many applications, they are found to be vulnerable to maliciously manipulated perturbation over the original data (Szegedy et al., 2013). To test the neural networks’ uncertainty estimates to adversarial attack, we generate the adversarial examples from one instance of each Bayesian neural network trained by different MCMC methods and compute the predictive entropy of the adversarial attack examples from other instances to quantify the uncertainty of the neural network’s predictions. We use the same architecture and configuration as the previous experiment but test with MNIST data instead. We use the fast gradient sign method (FGSM) to generate the adversarial data with different attack strength-epsilon. We plot the predictive entropy and accuracy on the attacked data in Figure 3. From Figure 3, we can observe that with the increase of attack strength epsilon, the predictive entropy of SGFuncRLD increases quickly while the accuracy decreases smoothly. When the FGSM attack epsilon is below 0.05, the accuracy remains close to 98% . When the FGSM attack epsilon is above 0.05, the predictive entropy increases quickly while the accuracy decreases smoothly. Note that for other methods, when the attack epsilon is larger than 0.2, the predictive entropy does not increase, which might be dangerous in real practice as the strongly-perturbed data still have high confidence for Bayesian neural network. This indicates that SGFuncRLD trained Functional Bayesian neural network can achieve the best performance on uncertainty estimation on adversarial attack data. We speculate that this is because sampling in the high dimensional space is hard for common MCMC method and the proposed SGFuncRLD can explore the parameter space much more efficiently to capture the multi-modality.
Next, we will test our proposed method’s generalization abilities on a practical application.

heading: 6.1 TRAFFIC SIGN RECOGNITION
text: To test methods’ generalization abilities, we use the German traffic sign recognition benchmarks (GTSRB) dataset. The GTSRB dataset consists of 39209 training images and 12630 test images. We
further split the training images into two part-34799 images for training, 4410 images for validation. Then We use a neural network with similar architecture as LeNet Lecun et al. (1998). We use the gray version of the traffic sign images and do the image local normalization and normalize the images for pre-processing the data. The code is adapted from the repository 4. The network architecture is conv (6 filters with a kernel size of 5)-pool-conv (16 filters with a kernel size of 5)-pool-fully connected (128)-fully connected (43). We do a grid search to determine the best parameter for each method. The best parameter setting for each method is: SGFuncRLD(learning rate is 0.001), pSGLD(learning rate is 0.0001), SGHMC(learning rate is 0.1, momentum is 0.6), SGLD(learning rate is 0.1). The training and validation curves are shown in Figure 4. From Figure 4, we can observe that the SGLD converges faster among all methods with SGFuncRLD only slower than SGLD. However, different from all other methods, the training and validation curves are more smooth for SGFuncRLD. All methods can achieve similar results after mixing. We speculate that this is because the landscape of loss function in functional space is smoother than the loss function in parameter space because of the mean-field Gaussian approximation. We test the trained models on the test dataset. As we observe that there is some small numerical difference for each run, we run the experiment five times and get the mean and the standard deviation of test accuracy as shown in Table 1. From Table 1, we can observe that the SGFuncRLD can achieve the best generalization performance on the large test dataset.

heading: 7 CONCLUSION
text: In this work, we extend the Bayesian neural networks to functional Bayesian neural networks for model uncertainty quantification. We propose a novel stochastic functional Riemannian dynamics for training functional Bayesian neural networks. We show the superiority of functional Bayesian neural networks over various examples.

Review summaries:
The authors propose an approximate MCMC method for sampling a posterior distribution of weights in a Bayesian neural network.  They claim that existing MCMC methods are limited by poor scaling with dimensionality of the weights, and they propose a method inspired by HMC on finite-dimensional approximations of measures on an infinite-dimensional Hilbert space (Beskos et al, 2011).  In short, the idea is to use a low dimensional approximation to the parameters (i.e. weights) of the neural network, representing them instead as a weighted combination of basis functions in neural network parameter space.  Then the authors propose to use HMC on this lower dimensional representation.  While the idea is intriguing, there are a number of flaws in the presentation, notational inconsistencies, and missing experiments that prohibit acceptance in the current form.

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.