Year: 2018

Abstract:
We propose a robust Bayesian deep learning algorithm to infer complex posteriors with latent variables. Inspired by dropout, a popular tool for regularization and model ensemble, we assign sparse priors to the weights in deep neural networks (DNN) in order to achieve automatic “dropout” and avoid over-fitting. By alternatively sampling from posterior distribution through stochastic gradient Markov Chain Monte Carlo (SG-MCMC) and optimizing latent variables via stochastic approximation (SA), the trajectory of the target weights is proved to converge to the true posterior distribution conditioned on optimal latent variables. This ensures a stronger regularization on the over-fitted parameter space and more accurate uncertainty quantification on the decisive variables. Simulations from large-p-small-n regressions showcase the robustness of this method when applied to models with latent variables. Additionally, its application on convolutional neural networks (CNN) leads to state-of-the-art performance on MNIST and Fashion MNIST datasets and improved resistance to adversarial attacks.

heading: 1 INTRODUCTION
text: Bayesian deep learning, which evolved from Bayesian neural networks (Neal, 1996; Denker and leCun, 1990), provides an alternative to point estimation due to its close connection to both Bayesian probability theory and cutting-edge deep learning models. It has been shown of the merit to quantify uncertainty (Gal and Ghahramani, 2016b), which not only increases the predictive power of DNN, but also further provides a more robust estimation to enhance AI safety. Particularly, Gal and Ghahramani (2016a;b) described dropout (Srivastava et al., 2014) as a variational Bayesian approximation. Through enabling dropout in the testing period, the randomly dropped neurons generate some amount of uncertainty with almost no added cost. However, the dropout Bayesian approximation is variational inference (VI) based thus it is vulnerable to underestimating uncertainty.
MCMC, known for its asymptotically accurate posterior inference, has not been fully investigated in DNN due to its unscalability in dealing with big data and large models. Stochastic gradient Langevin dynamics (SGLD) (Welling and Teh, 2011), the first SG-MCMC algorithm, tackled this issue by adding noise to a standard stochastic gradient optimization, smoothing the transition between optimization and sampling. Considering the pathological curvature that causes the SGLD methods inefficient in DNN models, Li et al. (2016) proposed combining adaptive preconditioners with SGLD (pSGLD) to adapt to the local geometry and obtained state-of-the-art performance on MNIST dataset. To avoid SGLD’s random-walk behavior, Chen et al. (2014) proposed using stochastic gradient Hamiltonian Monte Carlo (SGHMC), a second-order Langevin dynamics with a large friction term, which was shown to have lower autocorrelation time and faster convergence (Chen et al., 2015). Saatci and Wilson (2017) used SGHMC with GANs (Goodfellow et al., 2014a) to achieve a fully probabilistic inference and showed the Bayesian GAN model with only 100 labeled images was able to achieve 99.3% testing accuracy in MNIST dataset. Raginsky et al. (2017); Zhang et al. (2017); Xu et al. (2018) provided theoretical interpretations of SGLD from the perspective of non-convex optimization, echoing the empirical fact that SGLD works well in practice.
When the number of predictors exceeds the number of observations, applying the spike-and-slab priors is particularly powerful and efficient to avoid over-fitting by assigning less probability mass on
the over-fitted parameter space (Song and Liang, 2017). However, the inclusion of latent variables from the introduction of the spike-and-slab priors on Bayesian DNN is not handled appropriately by the existing methods. It is useful to devise a class of SG-MCMC algorithms that can explore the Bayesian posterior as well as optimize the hidden structures efficiently.
In this paper, we propose a robust Bayesian learning algorithm: SG-MCMC with a stochastic approximation adaptation (SG-MCMC-SA) to infer complex posteriors with latent variables and apply this method to supervised deep learning. This algorithm has four main contributions: 1. It enables SG-MCMC method to efficiently sample from complex DNN posteriors with latent variables and is proved to converge; 2. By automatically searching the over-fitted parameter space to add more penalties, the proposed method quantifies the posterior distributions of the decisive variables more accurately; 3. The proposed algorithm is robust to various hyperparameters; 4. This method demonstrates more resistance to over-fitting in simulations, leads to state-of-the-art performance and shows more robustness over the traditional SG-MCMC methods on the real data.

heading: 2 STOCHASTIC GRADIENT MCMC
text: We denote the decaying learning rate at time k by (k), the entire data by D = {di}Ni=1, where di = (xi, yi), the log of posterior by L(β),∇ as the gradient of any function in terms of β. The minibatch of data B is of size n with indices S = {s1, s2, ..., sn}, where si ∈ {1, 2, ..., N}. Stochastic gradient∇L̃(β) from a mini-batch of data B randomly sampled from D is used to approximate the true gradient∇L(β):
∇L̃(β) = ∇ log P(β) + N n ∑ i∈S ∇ log P(di|β).
The stochastic gradient Langevin dynamics (no momentum) is formed as follows:
β(k+1) = β(k) + (k)G∇L̃(β(k)) +N (0, 2 (k)Gτ−1), (1) where τ > 0 denotes the temperature,G is a positive definite matrix to precondition the dynamics. It has been shown that SGLD asymptotically converges to a stationary distribution π(β|D) ∝ eτL(β) (Teh et al., 2015; Zhang et al., 2017). As τ increases, the algorithm tends towards optimization with underestimated uncertainty. Another variant of SG-MCMC, SGHMC (Chen et al., 2014; Ma et al., 2015), proposes to use second-order Langevin dynamics to generate samples: β (k+1) = β(k) + (k)M−1r(k)
r(k+1) = r(k) + (k)∇L̃(β(k))− (k)CM−1r +N (0, (k)(2C − (k)B̂(k))τ−1) (2)
where r is the momentum item, M is the mass, B̂ is an estimate of the error variance from the stochastic gradient, C is a user-specified friction term to counteracts the noisy gradient.

heading: 3 STOCHASTIC GRADIENT MCMC WITH A STOCHASTIC APPROXIMATION ADAPTATION
text: Dropout has been proven successful, as it alleviates over-fitting and provides an efficient way of making bagging practical for ensembles of countless sub-networks. Dropout can be interpreted as assigning the Gaussian mixture priors on the neurons (Gal and Ghahramani, 2016c). To mimic Dropout in our Bayesian CNN models, we assign the spike-and-slab priors on the most fat-tailed weights in FC1 (Fig. 1). From the Bayesian perspective, the proposed robust algorithm distinguishes itself from the dropout approach in treating the priors: our algorithm keeps updating the priors during posterior inference, rather than fix it. The inclusion of scaled mixture priors in deep learning models were also studied in Blundell et al. (2015); Li et al. (2016) with encouraging results. However, to the best of our knowledge, none of the existing SG-MCMC methods could deal with complex posterior with latent variables. Intuitively, the Bayesian formulation with model averaging and the spike-and-slab priors is expected to obtain better predictive performance through averages from a "selective" group of "good" submodels, rather than averaging exponentially many posterior probabilities roughly. For the weight priors of the rest layers (dimension u), we just assume they follow the standard Gaussian distribution, while the biases follow improper uniform priors.

heading: 3.1 CONJUGATE SPIKE-AND-SLAB PRIOR FORMULATION FOR DNN
text: Similarly to the hierarchical prior in the EM approach to variable selection (EMVS) (Rořková and George, 2014), we assume the weights β ∈ Rp in FC1 follow the spike-and-slab mixture prior
π(β|σ2,γ) = Np(0,Vσ,γ), (3)
where γ ∈ {0, 1}p, σ ∈ R, Vσ,γ = σ2 diag(a1, . . . , ap) with aj = (1− γj)v0 + γjv1 for each j and 0 < v0 < v1. By introducing the latent variable γj = 0 or 1, the mixture prior is represented as
βj |γj ∼ (1− γj)N (0, σ2v0) + γjN (0, σ2v1). (4) The interpretation is: if γj = 0, then βj is close to 0; if γj = 1, the effect of βj on the model is intuitively large. The likelihood of this model given a mini-batch of data {(xi, yi)}i∈S is
π(B|β, σ2) =  (2πσ2)−n/2 exp { − 1 2σ2 ∑ i∈S(yi − ψ(xi;β))2 } (regression) ∏ i∈S exp{ψyi(xi;β)}∑K t=1 exp{ψt(xi;β)} (classification),
(5)
where ψ(xi;β) can be a mapping for logistic regression or linear regression, or a mapping based on a series of nonlinearities and affine transformations in the deep neural network. In the classification formulation, yi ∈ {1, . . . ,K} is the response value of the i-th example. In addition, the variance σ2 follows an inverse gamma prior
π(σ2|γ) = IG(ν/2, νλ/2). (6) The i.i.d. Bernoulli prior is used since there is no structural information in the same layer.
π(γ|δ) = δ|γ|(1− δ)p−|γ|, where π(δ) ∝ δa−1(1− δ)b−1 and δ ∈ R. (7)
Finally, our posterior density follows
π(β, σ2, δ,γ|B) ∝ π(B|β, σ2)Nn π(β|σ2,γ)π(σ2|γ)π(γ|δ)π(δ). (8)
The EMVS approach is efficient in identifying potential sparse high posterior probability submodels on high-dimensional regression (Rořková and George, 2014) and classification problem (McDermott et al., 2016). These characteristics are helpful for large neural network computation, thus we refer the stochastic version of the EMVS algorithm as Expectation Stochastic-Maximization (ESM).

heading: 3.2 EXPECTATION STOCHASTIC-MAXIMIZATION
text: Due to the existence of latent variables, optimizing π(β, σ2, δ|B) directly is difficult. We instead iteratively optimize the “complete-data” posterior log π(β, σ2, δ,γ|B), where the latent indicator γ is treated as “missing data”.
More precisely, the ESM algorithm is implemented by iteratively increasing the objective function Q(β, σ, δ|β(k), σ(k), δ(k)) = Eγ|· [ log π(β, σ, δ,γ|B)|β(k), σ(k), δ(k),B ] , (9)
where Eγ|· denotes the conditional expectation Eγ|β(k),σ(k),δ(k),B(·). Given (β(k), σ(k), δ(k)) at the k-th iteration, we first compute the expectation of Q, then alter (β, σ, δ) to optimize it.
For the conjugate spike-slab hierarchical prior formulation, the objective function Q is of the form
Q(β, σ, δ|β(k), σ(k), δ(k)) = C +Q1(β, σ|β(k), σ(k), δ(k)) +Q2(δ|β(k), σ(k), δ(k)), (10)
where
Q1(β, σ|β(k), σ(k), δ(k)) =
log likelihood︷ ︸︸ ︷ N
n log π(B|β, σ2)−
spike-and-slab priors in the "sparse" layer, e.g. FC1︷ ︸︸ ︷∑p j=1 β 2 jEγ|· [ 1
v0(1− γj) + v1γj ] 2σ2
− ∑p+u j=p+1 β 2 j
2︸ ︷︷ ︸ Gaussian priors in other layers
−p+ ν + 2 2 log(σ2)− νλ 2σ2
(11)
and
Q2(δ|β(k), σ(k), δ(k)) = p∑ j=1 log ( δ 1− δ ) Eγ|·γj + (a− 1) log(δ) + (p+ b− 1) log(1− δ).
(12)

heading: 3.2.1 THE E-STEP
text: The physical meaning of Eγ|·γj in Q2 is the probability ρj , where ρ ∈ Rp, of βj having a large effect on the model. Formally, we have
ρj = Eγ|·γj = P(γj = 1|β(k), σ(k), δ(k)) = aj
aj + bj , (13)
where aj = π(β (k) j |γj = 1)P(γj = 1|δ(k)) and bj = π(β (k) j |γj = 0)P(γj = 0|δ(k)). The choice of Bernoulli prior enables us to use P(γj = 1|δ(k)) = δ(k). The other conditional expectation comes from a weighted average κj , where κ ∈ Rp.
κj = Eγ|· [
1
v0(1− γj) + v1γj
] =
Eγ|·(1− γj) v0 + Eγ|·γj v1 = 1− ρj v0 + ρj v1 . (14)

heading: 3.2.2 THE STOCHASTIC M-STEP
text: Since there is no closed-form optimal solution for β here, to optimize Q1 with respect to β, we use Adam (Kingma and Ba, 2014), a popular algorithm with adaptive learning rates, to train the model.
In order to optimize Q1 with respect to σ, by denoting diag{κi}pi=1 as V , following the formulation in McDermott et al. (2016) and Rořková and George (2014) we have:
σ(k+1) =  √ N n ∑ i∈S ( yi − ψ(xi;β(k+1)) )2 + ||V1/2β(k+1)||2 + νλ N + p+ ν (regression),√√√√∥∥∥(V1/2β(k+1))∥∥∥2 + νλ
p+ ν + 2 (classification).
(15)
To optimize Q2, a closed-form solution can be derived from Eq.(12) and Eq.(13).
δ(k+1) = arg max δ∈R Q2(δ|β(k), σ(k), δ(k)) = ∑p j=1 ρj + a− 1 a+ b+ p− 2 . (16)
Algorithm 1 SGLD-SA with spike and slab priors
Inputs: { (k)}Tk=1, {ω(k)}Tk=1, τ , a, b, ν, λ, v0, v1 Outputs: {β(k)}Tk=1 Initialize: β(1), ρ(1), κ(1), σ(1) and δ(1) from scratch for k ← 1 : T do
Sample mini-batch B(k) Sampling β(k+1) ← β(k) + (k)∇Q(β, σ, δ|β(k), σ(k), δ(k)) +N (0, 2 (k)τ−1) Optimization Compute ρ(k+1) in Eq.(13), perform SA: ρ(k+1) ← (1− ω(k+1))ρ(k) + ω(k+1)ρ(k+1) Compute κ(k+1) in Eq.(14), perform SA: κ(k+1) ← (1− ω(k+1))κ(k) + ω(k+1)κ(k+1) Compute σ(k+1) in Eq.(15), perform SA: σ(k+1) ← (1− ω(k+1))σ(k) + ω(k+1)σ(k+1) Compute δ(k+1) in Eq.(16), perform SA: δ(k+1) ← (1− ω(k+1))δ(k) + ω(k+1)δ(k+1)
end for
3.3 SG-MCMC-SA
The EMVS algorithm is designed for linear regression models, although the idea can be extended to nonlinear models. However, when extending to nonlinear models, such as DNNs, the M-step will not have a closed-form update anymore. A trivial implementation of the M-step will likely cause a local-trap problem. To tackle this issue, we replace the E-step and the M-step by SG-MCMC with the prior hyperparameters tuned via stochastic approximation (Benveniste et al., 1990):
β(k+1) = β(k) + (k)∇L̃(β(k),θ(k)) +N (0, 2 (k)τ−1), θ(k+1) = θ(k) + ω(k+1) ( gθ(k)(β (k+1))− θ(k) ) ,
(17)
where gθ(k)(β(k+1))− θ(k) is the gradient-like function in stochastic approximation (see details in Appendix A.2), gθ(·) is the mapping detailed in Eq.(13), Eq.(14), Eq.(15) and Eq.(16) to derive the optimal θ based on the current β, the step size ω(k) can be set as A(k+Z)−α with α ∈ (0.5, 1]. The interpretation of this algorithm is that we sample β(k+1) from L̃(β(k),θ(k)) and adaptively optimize θ(k+1) from the mapping gθ(k) . We expect to obtain an augmented sequence as follows:
β(1),θ(1);β(2),θ(2);β(3),θ(3); . . . (18)
We show the (local) L2 convergence rate of SGLD-SA below and present the details in Appendix B. Theorem 1 (L2 convergence rate). For any α ∈ (0, 1] and any compact subset Ω ∈ R2p+2, under assumptions in Appendix B.1, the algorithm satisfies: there exists a constant λ such that
E [ ‖θ(k) − θ∗‖2I(k ≤ t(Ω)) ] ≤ λω(k),
where t(Ω) = inf{k : θ(k) /∈ Ω} and ω(k) = A(k + Z)−α ∼ O(k−α). Corollary 1. For any α ∈ (0, 1] and any compact subset Ω ∈ R2p+2, under assumptions in Appendix B.2, the distribution of β(k) in (18) converges weakly to the invariant distribution eL(β,θ ∗) as → 0.
The key to guaranteeing the consistency of the latent variable estimators is from stochastic approximation and the fact of ∑∞ i=1 ω
(i) = ∞. The non-convex optimization of SGLD (Raginsky et al., 2017; Xu et al., 2018), in particular the ability to escape shallow local optima (Zhang et al., 2017), ensures the robust optimization of the latent variables. Furthermore, the mappings of Eq.(13), Eq.(14), Eq.(15) and Eq.(16) all satisfy the assumptions on g in a compact subset Ω (Appendix B), which enable us to apply theorem 1 to SGLD-SA. Because SGHMC proposed by Chen et al. (2014) is essentially a second-order Langevin dynamics and yields a stationary distribution given an accurate estimation of the error variance from the stochastic gradient, the property of SGLD-SA also applies to SGHMC-SA. Corollary 2. For any α ∈ (0, 1] and any compact subset Ω ∈ R2p+2, under assumptions in Appendix B.2, the distribution of β(k) from SG-MCMC-SA converges weakly to the invariant distribution π(β,ρ∗,κ∗, σ∗, δ∗|D) as → 0.

heading: 3.4 POSTERIOR APPROXIMATION
text: The posterior average given decreasing learning rates can be approximated through the weighted sample average E[ψ(β)] = ∑T k=1 (k)ψ(β(k))∑T
k=1 (k) (Welling and Teh, 2011) to avoid over-emphasizing the
tail end of the sequence and reduce the variance of the estimator. Teh et al. (2015); Chen et al. (2015) showed a theoretical optimal learning rate (k) ∝ k−1/3 for SGLD and k−1/5 for SGHMC to achieve faster convergence for posterior average, which are used in Sec. 4.1 and Sec. 4.2 respectively.

heading: 4 EXPERIMENTS

heading: 4.1 SIMULATION OF LARGE-P-SMALL-N REGRESSION
text: SGLD-SA can be applied to the (logistic) linear regression cases, as long as u = 0 in Eq.(11). We conduct the linear regression experiments with a dataset containing n = 100 observations and p = 1000 predictors. Np(0,Σ) is chosen to simulate the predictor values X (training set) where Σ = (Σ)pi,j=1 with Σi,j = 0.6
|i−j|. Response values y are generated from Xβ + η, where β = (β1, β2, β3, 0, 0, ..., 0)
′ and η ∼ Nn(0, 3In). To make the simulation in Rořková and George (2014) more challenging, we assume β1 ∼ N (3, σ2c ), β2 ∼ N (2, σ2c ), β3 ∼ N (1, σ2c ), σc = 0.2. We introduce some hyperparameters, but most of them are uninformative, e.g. ν ∈ {10−3, 1, 103} makes little difference in the test set performance. Sensitivity analysis shows that three hyperparameters are important: v0, a and σ, which are used to identify and regularize the over-fitted space. We fix τ = 1, λ = 1, ν = 1, v1 = 100, δ = 0.5, b = p and set a = 1. The learning rates for SGLD-SA and SGLD are set to (k) = 0.001× k− 13 , and the step size ω(k) = 0.1× (k+ 1000)− 34 . We vary v0 and σ to show the robustness of SGLD-SA with respect to different initializations.
To implement SGLD-SA, we perform stochastic gradient optimization by randomly selecting 50 observations and calculating the corresponding gradient in each iteration. We simulate 500, 000 samples from the posterior distribution and at the same time keep optimizing the latent variables. EMVS is implemented with β directly optimized each time. We also simulate a group of the test set with 1000 observations (display 50 in Fig. 2(e)) in the same way as generating the training set to evaluate the generalizability of our algorithm. Tab.1 shows that EMVS frequently fails given bad initializations, while SGLD-SA is fairly robust to the hyperparameters. In addition, from Fig. 2(a), Fig. 2(b) and Fig.2(c), we can see SGLD-SA is the only algorithm among the three that quantifies the uncertainties of β1, β2 and β3 reasonably and always gives more accurate posterior average (Fig.2(f)); by contrast, the estimated response values y from SGLD is close to the true values in the training set (Fig.2(d)), but are far away from them in the testing set (2(e)), indicating the over-fitting problem of SGLD without proper regularization.
For the simulation of SGLD-SA in logistic regression to demonstrate the advantage of SGLD-SA over SGLD and ESM, we leave the results in Appendix C.

heading: 4.2 CLASSIFICATION PROBLEM
text: We implement all the algorithms in Pytorch (Paszke et al., 2017) and run the experiments on GeForce GTX 1080 GPUs. The first DNN we use is a standard 2-Conv-2-FC CNN model (Fig.1) of 670K parameters (see details in Appendix D.1). The first set of experiments is to compare methods on the same model without using other complication, such as data augmentation (DA) or batch normalization (BN) (Ioffe and Szegedy, 2015). We refer to the general CNN without dropout as Vanilla, with 50% dropout rate applied to the green neurons (Fig.1) as Dropout. Vanilla and Dropout models are trained
with Adam (Kingma and Ba, 2014) with Pytorch default parameters (with learning rate 0.001). We use SGHMC as a benchmark method as it is also sampling-based and has a close relationship with the popular momentum based optimization approaches in DNN. SGHMC-SA differs from SGHMC in that SGHMC-SA applies the spike-and-slab priors to the FC1 layer while SGHMC just uses the standard normal priors. The hyperparameters v1 = 1, v0 = 1× 10−3 and σ = 0.1 in SGHMC-SA are used to regularize the over-fitted space, and a, b are set to p to obtain a moderate "sparsity" to resemble dropout, the step size is ω(k) = 0.1× (k + 1000)− 34 . We use training batch size 1000 and a thinning factor 500 to avoid a cumbersome system, and the posterior average is applied to each Bayesian model. Temperatures are tuned to achieve better results (Appendix D.2).
The four CNN models are tested on the classical MNIST and the newly introduced Fashion MNIST (FMNIST) (Xiao et al., 2017) dataset. Performance of these models is shown in Tab.2. Compared with SGHMC, our SGHMC-SA outperforms SGHMC on both datasets. We notice the posterior averages from SGHMC-SA and SGHMC obtain much better performance than Vanilla and Dropout. Without using either DA or BN, SGHMC-SA achieves 99.60% which even outperforms some state-of-the-art models, such as stochastic pooling (99.53%) (Zeiler and Fergus, 2013), Maxout Network (99.55%) (Goodfellow et al., 2013) and pSGLD (99.55%) (Li et al., 2016) . In F-MNIST, SGHMC-SA obtains 93.01% accuracy, outperforming all other competing models.
To further test the maximum performance of SGHMC-SA, we apply DA and BN to the following experiments (see details in Appendix D.3) and refer the datasets with DA as aMNIST and aFMNIST. All the experiments are conducted using a 2-Conv-BN-3-FC CNN of 490K parameters. Using this model, we obtain 99.75% on aMNIST (300 epochs) and 94.38% on aFMNIST (1000 epochs). The results are noticeable because posterior model averaging is essentially conducted on a single Bayesian neural network. We also conduct the experiments based on the ensemble of five networks and refer them as aMNIST-5 and aFMNIST-5 in Tab. 2. We achieve 99.79% on aMNIST-5 using 5 small Bayesian neural networks each with 2 thinning samples (4 thinning samples in aFMNIST-5), which is comparable with the state-of-the-art performance (Wan et al., 2013).

heading: 4.3 DEFENSES AGAINST ADVERSARIAL ATTACKS
text: Continuing with the setup in Sec. 4.2, the third set of experiments focus on evaluating model robustness. We expect less robust models perform considerably well on a certain dataset due to
over-tuning; however, as the degree of adversarial attacks increases, the performance decreases sharply. In contrast, more robust models should be less affected by these adversarial attacks.
We apply the Fast Gradient Sign method (Goodfellow et al., 2014b) to generate the adversarial examples with one single gradient step as in Papernot et al. (2016)’s study:
xadv ← x− ζ · sign{δxmax y log P(y |x)},
where ζ ranges from 0.1, 0.2, . . . , 0.5 to control the different levels of adversarial attacks.
Similar to the setup in the adversarial experiments by Li and Gal (2017), we normalize the adversarial images by clipping to the range [0, 1]. As shown in Fig. 3(b) and Fig.3(d), there is no significant difference among all the four models in the early phase. As the degree of adversarial attacks arises, the images become vaguer as shown in Fig.3(a) and Fig.3(c). In this scenario the performance of Vanilla decreases rapidly, reflecting its poor defense against adversarial attacks, while Dropout performs better than Vanilla. But Dropout is still significantly worse than the sampling based methods SGHMC-SA and SGHMC. The advantage of SGHMC-SA over SGHMC becomes more significant when ζ > 0.25.
In the case of ζ = 0.5 in MNIST where the images are hardly recognizable, both Vanilla and Dropout models fail to identify the right images and their predictions are as worse as random guesses. However, SGHMC-SA model achieves roughly 11% higher than these two models and 1% higher than SGHMC, which demonstrates the strong robustness of our proposed SGHMC-SA. Overall, SGHMC-SA always yields the most robust performance.

heading: 5 CONCLUSION AND FUTURE RESEARCH
text: We propose a mixed sampling-optimization method called SG-MCMC-SA to efficiently sample from complex DNN posteriors with latent variables and prove its convergence. By adaptively searching and penalizing the over-fitted parameter space, the proposed method improves the generalizability of deep neural networks. This method is less affected by the hyperparameters, achieves higher prediction accuracy over the traditional SG-MCMC methods in both simulated examples and real applications and shows more robustness towards adversarial attacks.
Interesting future directions include applying SG-MCMC-SA towards popular large deep learning models such as the residual network (He et al., 2015) on CIFAR-10 and CIFAR-100, combining active learning and uncertainty quantification to learn from datasets of smaller size and proving posterior consistency and the consistency of variable selection under various shrinkage priors concretely.

heading: A REVIEW

heading: A.1 FOKKER-PLANCK EQUATION
text: The Fokker-Planck equation (FPE) can be formulated from the time evolution of the conditional distribution for a stationary random process. Denoting the probability density function of the random process at time t by q(t, β), where β is the parameter, the stochastic dynamics is given by
∂tq(t, β) = ∂β(∂β(−L(β))q(t, β)) + ∂2βq(t, β). (19)
Let q(β) = limt→∞ q(t, β). If limt→∞ ∂tq(t, β) = 0, then
lim t→∞
∂β(∂β(−L(β))q(t, β)) + ∂2βq(t, β) = 0
⇒ ∂β (∂β (−L(β)q(β))) + ∂2βq(β) = 0 ⇔ ∂β (∂β(−L(β))q(β) + ∂βq(β)) = 0 ⇔ ∂β (q(β) (∂β(−L(β) + ∂β log q(β))) = 0.
(20)
Therefore, ∂β(−L(β)) + ∂β log(q(β)) = 0, which implies q(β) ∝ eL(β). In other words, limt→∞ q(t, β) ∝ eL(β), i.e. q(t, β) gradually converges to the Bayesian posterior eL(β).

heading: A.2 STOCHASTIC APPROXIMATION

heading: A.2.1 SPECIAL CASE: ROBBINS–MONRO ALGORITHM
text: Robbins–Monro algorithm is the first stochastic approximation algorithm to deal with the root finding problem. Given the random output of H(θ,β) with respect to β, our goal is to find θ∗ such that
h(θ∗) = Eθ∗ [H(θ∗,β)] = 0, (21)
where Eθ∗ denotes the expectation with respect to the distribution of β given parameter θ∗. To implement the Robbins–Monro Algorithm, we can generate iterates of the form as follows :
θ(k+1) = θ(k) + ω(k+1)H(θ(k),β(k+1)).
Note that in this algorithm, H(θ,β) is the unbiased estimator of h(θ), that is
Eθ(k) [H(θ(k),β(k+1))− h(θ(k))|Fk] = 0. (22)
If there exists an antiderivativeQ(θ,β) that satisfiesH(θ,β) = ∂Q(θ,β)∂θ andEθ[Q(θ,β)] is concave, it is equivalent to solving the stochastic optimization problem maxθ∈ΘEθ[Q(θ,β)].

heading: A.2.2 GENERAL STOCHASTIC APPROXIMATION
text: In contrast to Eq.(21) and Eq.(22), the general stochastic approximation algorithm is intent on solving the root of the integration equation
h(θ) = lim k→∞
Eθ[H(θ,β(k+1))|Fk]
= lim k→∞
∫ H(θ,β)Πθ(β (k), dβ)
= ∫ H(θ,β)fθ(dβ) = 0,
(23)
where θ ∈ Θ, β ∈ B, for a subset B ∈ B, the transition kernel Πθ(β(k),B), which converges to the invariant distribution fθ(β), satisfies that P[β(k+1) ∈ B|Fk] = Πθ(β(k),B). The stochastic approximation algorithm is an iterative recursive algorithm consisting of two steps:
(1) Generate β(k+1) from the transition kernel Πθ(k)(β(k), ·), which admits fθ(k)(β) as the invariant distribution,
(2) Update θ(k+1) = θ(k) + ω(k+1)H(θ(k),β(k+1)).
Then we have Eθ(k) [H(θ(k),β(k+1))|Fk] = ∫ H(θ(k),β)Πθ(β
(k), dβ). Compared with Eq.(22), h(θ(k))− Eθ(k) [H(θ(k),β(k+1))|Fk] is not equal to 0 but decays to 0 as k →∞. To summarise, H(θ,β) is a biased estimator of h(θ) in finite steps, but as k → ∞, the bias decreases to 0. In the SG-MCMC-SA algorithm (17), H(θ(k),β(k+1)) is defined by gθ(k)(β (k+1))− θ(k).

heading: B CONVERGENCE ANALYSIS

heading: B.1 CONVERGENCE OF HIDDEN VARIABLES
text: The stochastic gradient Langevin Dynamics with a stochastic approximation adaptation (SGLD-SA) is a mixed half-optimization-half-sampling algorithm to handle complex Bayesian posterior with latent variables, e.g. the conjugate spike-slab hierarchical prior formulation. Each iteration of the algorithm consists of the following steps:
(1) Sample β(k+1) using SGLD based on the current θ(k), i.e.
β(k+1) = β(k) + ∂
∂β L̃(β(k),θ(k)) +N (0, 2 τ−1); (24)
(2) Optimize θ(k+1) from the following recursion θ(k+1) = θ(k) + ω(k+1) ( gθ(k)(β (k+1))− θ(k) )
= (1− ω(k+1))θ(k) + ω(k+1)gθ(k)(β(k+1)), (25)
where gθ(k)(·) is some mapping to derive the optimal θ based on the current β.
Remark: Define H(θ(k),β(k+1)) = gθ(k)(β(k+1))− θ(k). In this formulation, our target is to find θ∗ such that h(θ∗) = 0, where h(θ) := ∫ H(θ,β)fθ(dβ). Eθ(k) [H(θ(k),β(k+1))|Fk] converges to h(θ(k)) as k →∞ , this algorithm falls to the category of the general stochastic approximation.

heading: GENERAL ASSUMPTIONS
text: To provide the local L2 upper bound for SGLD-SA, we first lay out the following assumptions:
Assumption 1 (Step size and Convexity). {ω(k)}k∈N is a positive decreasing sequence of real numbers such that
ω(k) → 0, ∞∑ k=1 ω(k) = +∞. (26)
There exist constant δ > 0 and θ∗ such that for all θ ∈ Θ
〈θ − θ∗, h(θ)〉 ≤ −δ‖θ − θ∗‖2, (27)
with additionally
lim k→∞
inf 2δ ω(k) ω(k+1) + ω(k+1) − ω(k) ω(k+1) 2 > 0. (28)
Then for any α ∈ (0, 1] and suitable A and B, a practical ω(k) can be set as
ω(k) = A(k +B)−α (29)
Assumption 2 (Existence of Markov transition kernel). For any θ ∈ Θ, in the mini-batch sampling, there exists a family of noisy Kolmogorov operators {Πθ} approximating the operator (infinitesimal) of the Ito diffusion, such that every Kolmogorov operator Πθ corresponds to a single stationary distribution fθ, and for any Borel subsetA of β, we have
P[β(k+1) ∈ A|Fk] = Πθ(k)(β(k),A).
Assumption 3 (Compactness). For any compact subset Ω of Θ, we only consider θ ∈ Ω such that
‖θ‖ ≤ C0(Ω) (30)
Note that the compactness assumption of the latent variable θ is not essential, the assumption that the variable is in the compact domain is not only reasonable, but also simplifies our proof.

heading: In addition, there exists constants C1(Ω) and C2(Ω) so that
text: Eθ(k) [‖H(θ(k),β(k+1))‖2|Fk] ≤ C1(Ω)(1 + ‖θ(k)‖2) ≤ C2(Ω)(1 + ‖θ(k) − θ∗ + θ∗‖2) ≤ C2(Ω)(1 + ‖θ(k) − θ∗‖2)
(31)
Assumption 4 (Solution of Poisson equation). For all θ ∈ Θ, there exists a function µθ on β that solves the Poisson equation µθ(β)−Πθµθ(β) = H(θ,β)− h(θ), which follows that
H(θ(k),β(k+1)) = h(θ(k)) + µθ(k)(β (k+1))−Πθ(k)µθ(k)(β(k+1)). (32)
For any compact subset Ω of Θ, there exist constants C3(β,Ω) and C4(β,Ω) such that for all θ,θ′ ∈ Ω,
(1) ‖Πθµθ‖ ≤ C3(Ω), (2) ‖Πθµθ −Πθ′µθ′‖ ≤ C4(Ω)‖θ − θ′‖,
(33)
Remark: For notation simplicity, we write C1(Ω) as C1, C2(Ω) as C2, . . . in the following context.
Lemma 1 is a restatement of Lemma 25 (page 447) from Benveniste et al. (1990). Lemma 1. Suppose k0 is an integer which satisfies with
inf k≥k0
ω(k+1) − ω(k)
ω(k)ω(k+1) + 2δ − ω(k+1)C2 > 0.
Then for any k > k0, the sequence {ΛKk }k=k0,...,K defined below is increasing and bounded by 2ω(K)
ΛKk =  2ω (k) ∏K−1 j=k (1− 2ωj+1δ + ω2j+1C2) if k < K,
2ω(K) if k = K. (34)
Lemma 2 is an extension of Lemma 23 (page 245) from Benveniste et al. (1990)
Lemma 2. There exist λ0 and k0 such that for all λ ≥ λ0 and k ≥ k0, the sequence u(k) = λω(k) satisfies
u(k+1) ≥ (1− 2ω(k+1)δ + ω(k+1) 2 C2)u (k) + ω(k+1) 2 C2 + ω (k+1)C1. (35)
Proof. Replace u(k) = λω(k) in (35), we have
λω(k+1) ≥ (1− 2ω(k+1)δ + ω(k+1) 2 C2)λω (k) + ω(k+1) 2 C2 + ω (k+1)C1. (36)
From (28), we can denote a positive constant ∆+ as limk→∞ inf 2δω(k+1)ω(k) + ω(k+1) − ω(k). Then (36) can be simplified as
λ(∆+ − ω(k+1) 2 ω(k)C2) ≥ ω(k+1) 2 C2 + ω (k+1)C1. (37)
There exist λ0 and k0 such that for all λ > λ0 and k > k0, (37) holds. Note that in practical case when C1 is small, finding a suitable λ0 will not be a problem.
Theorem 1 (L2 convergence rate). Suppose that Assumptions 1-4 hold, for any compact subset Ω ∈ Θ, the algorithm satisfies: there exists a constant λ such that
E [ ‖θ(k) − θ∗‖2I(k ≤ t(Ω)) ] ≤ λω(k),
where t(Ω) = inf{k : θ(k) /∈ Ω}.
Proof. Denote T (k) = θ(k) − θ∗, with the help of (25) and Poisson equation (32), we deduce that
‖T (k+1)‖2 = ‖T (k)‖2 + ω(k+1) 2 ‖H(θ(k),β(k+1))‖2 + 2ω(k+1)〈T (k), H(θ(k),β(k+1))〉
= ‖T (k)‖2 + ω(k+1) 2 ‖H(θ(k),β(k+1))‖2
+ 2ω(k+1)〈T (k), h(θ(k))〉+ 2ω(k+1)〈T (k), µθ(k)(β(k+1))−Πθ(k)µθ(k)(β(k+1))〉 = ‖T (k)‖2 + D1 + D2 + D3.
First of all, according to (31) and (27), we have
ω(k+1) 2 ‖H(θ(k),β(k+1))‖2 ≤ ω(k+1) 2 C2(1 + ‖T (k)‖2), (D1)
2ω(k+1)〈T (k), h(θ(k))〉 ≤ −2ω(k+1)δ‖T (k)‖2, (D2)
Conduct the decomposition of D3 similar to Theorem 24 (p.g. 246) from Benveniste et al. (1990) and Lemma A.5 (Liang, 2010).
µθ(k)(β (k+1))−Πθ(k)µθ(k)(β(k+1)) = ( µθ(k)(β (k+1))−Πθ(k)µθ(k)(θ(k)) )
+ (
Πθ(k)µθ(k)(θ (k))−Πθ(k−1)µθ(k−1)(θ(k)) ) + (
Πθ(k−1)µθ(k−1)(θ (k))−Πθ(k)µθ(k)(β(k+1))
) ,
= D3-1 + D3-2 + D3-3.
(i) µθ(k)(β(k+1))−Πθ(k)µθ(k)(θ(k)) forms a martingale difference sequence since E [ µθ(k)(β (k+1))−Πθ(k)µθ(k)(θ(k))|Fk ] = 0. (D3-1)
(ii) From (33) and (30) respectively, we deduce that
‖Πθ(k)µθ(k)(θ(k))−Πθ(k−1)µθ(k−1)(θ(k))‖ ≤ C4‖θ(k) − θ(k−1)‖ ≤ 2C4C0
Thus there exists C4 = 2C4C0 such that
2ω(k+1)〈T (k),Πθ(k)µθ(k)(θ(k))−Πθ(k−1)µθ(k−1)(θ(k))〉 ≤ ω(k+1)C4. (D3-2)
(iii) D3-3 can be further decomposed to D3-3a and D3-3b
〈T (k),Πθ(k−1)µθ(k−1)(θ(k))−Πθ(k)µθ(k)(β(k+1))〉 = ( 〈T (k),Πθ(k−1)µθ(k−1)(θ(k))〉 − 〈T (k+1),Πθ(k)µθ(k)(β(k+1))〉 ) + ( 〈T (k+1),Πθ(k)µθ(k)(β(k+1))〉 − 〈Tk,Πθ(k)µθ(k)(β(k+1))〉
) = (z(k) − z(k+1)) + 〈T (k+1) − T (k),Πθ(k)µθ(k)(β(k+1))〉 = D3-3a + D3-3b,
where z(k) = 〈T (k),Πθ(k−1)µθ(k−1)(θ(k))〉 with a constant C3 = 4C0C3 which satisfies that 2ω(k+1)〈T (k+1) − T (k),Πθ(k)µθ(k)(β(k+1))〉
= 2ω(k+1)〈θ(k+1) − θ(k),Πθ(k)µθ(k)(β(k+1))〉 (D3-3b) ≤ 2ω(k+1)‖θ(k+1) − θ(k)‖ · ‖Πθ(k)µθ(k)(β(k+1))‖ ≤ 4ω(k+1)C0C3 = C3ω(k+1)
Finally, add all the items D1, D2 and D3 together, for some C1 = C3 + C4, we have E [ ‖T (k+1)‖2 ] ≤ (1− 2ω(k+1)δ + ω(k+1) 2 C2)E [ ‖Tk‖2 ] + ω(k+1) 2 C2 + ω (k+1)C1 + 2ω (k+1)E[z(k) − z(k+1)].
Moreover, from (33) and (30), there exists a constant C5 such that
E[|z(k)|] ≤ C5. (38)
Lemma 3 is an extension of Lemma 26 (page 248) from Benveniste et al. (1990).
Lemma 3. Let {u(k)}k≥k0 as a sequence of real numbers such that for all k ≥ k0, some suitable constants C1 and C2
u(k+1) ≥ u(k) ( 1− 2ω(k+1)δ + ω(k+1) 2 C2 ) + ω(k+1) 2 C2 + ω (k+1)C1, (39)
and assume there exists such k0 that E [ ‖T (k0)‖2 ] ≤ u(k0). (40)
Then for all k > k0, we have E [ ‖T (k)‖2 ] ≤ u(k) + k∑ j=k0+1 Λkj (z (j−1) − z(j)).
Proof of Theorem 1 (Continued). From Lemma 2, we can choose λ0 and k0 which satisfy the conditions (39) and (40)
E[‖T (k0)‖2] ≤ u(k0) = λ0ω(k0).
From Lemma 3, it follows that for all k > k0
E [ ‖Tk‖2 ] ≤ u(k) + E  k∑ j=k0+1 Λkj ( z(j−1) − z(j) ) . (41) From (38) and the increasing property of Λkj in Lemma 1, we have
E ∣∣∣∣∣∣ k∑
j=k0+1
Λkj
( z(j−1) − z(j) )∣∣∣∣∣∣  = E ∣∣∣∣∣∣ k−1∑
j=k0+1
(Λkj+1 − Λkj )z(j) − 2ω(k)z(k) + Λkk0+1z (k0) ∣∣∣∣∣∣ 
≤ k−1∑
j=k0+1
(Λkj+1 − Λkj )C5 + E[|2ω(k)z(k)|] + ΛkkC5
≤ (Λkk − Λkk0)C5 + Λ k kC5 + Λ k kC5
≤ 3ΛkkC5 = 6C5ω(k). (42)
Therefore, given the sequence u(k) = λ0ω(k) that satisfies conditions (39), (40) and Lemma 3, for any k > k0, from (41) and (42), we have
E[‖T (k)‖2] ≤ u(k) + 3C5Λkk = (λ0 + 6C5)ω(k) = λω(k), where λ = λ0 + 6C5.

heading: B.2 CONVERGENCE OF SAMPLES
text: In addition to the previous assumptions, we make one more assumption on the stochastic gradients to guarantee that the samples converge to the posterior conditioned on the optimal latent variables: Assumption 5 (Gradient Unbiasedness and Smoothness). For all β ∈ B and θ ∈ Θ, the mini-batch of data B, the stochastic noise ξ, which comes from ∇̃L(β,θ) − ∇L(β,θ), is a white noise and independent with each other. In addition, there exists a constant l ≥ 2 such that the following conditions hold:
EB[ξ] = 0 and EB[ξl] <∞. (43)
For all θ,θ′ ∈ Θ, there exists a constant M > 0 such that the gradient is M-smooth: ∇L(β,θ)−∇L(β,θ′) ≤M‖θ − θ′‖ (44)
Corollary 1. For all α ∈ (0, 1], under assumptions 1-5, the distribution of β(k) converges weakly to the invariant distribution eL(β,θ ∗) as → 0.
Proof. The proof framework follows from section 4 of Sato and Nakagawa (2014). In the context of stochastic noise ξ(k), we ignore the subscript of and only consider the case of τ = 1. Since θ(k) converges to θ∗ in SGLD-SA and the gradient is M-smooth (44), we transform the stochastic gradient from ∇L̃(β(k),θ(k)) to ∇L(β(k),θ∗) + ξ(k) +O( α), therefore Eq.(24) can be written as
β(k+1) = β(k) + ( ∇L(β(k),θ∗) + ξ(k) +O( α) ) + √ 2 η(k), where η(k) ∼ N (0, I). (45)
Using Eq.(43), the characteristic function of ξ(k) +O( α+1) is
E [ exp ( is( ξ(k) +O( α+1) )] = E [ ∞∑ l=0 1 l! ( is( ξ(k) +O( α+1) )l] = 1 +O( α+1). (46)
Then the characteristic function of √ 2 η(k) is exp(− s2).
Rewrite β(k+1) as β(k+ ), the characteristic function of β(t+ ) is the characteristic function of β(k) + ( ∂βL(β (k),θ∗) + ξ(k) +O( α) ) + √ 2 η(k), which is
φt+ (s) =
∫ exp(isβ + is ∇L(β,θ∗)− s2) ( 1 +O( α+1) ) q(k,β)dβ
= ∫ exp(isβ + is ∇L(β,θ∗)− s2)q(k,β)dβ +O( α+1).
(47)
With the fact exp(x) = 1 + x+O(x2), we can get
φ(k+ )(s)− φ(k)(s)
= ∫ exp(isβ) ( exp(is ∇L(β,θ∗)− s2)− 1 ) q(k,β)dβ +O( α+1)
= ∫ exp(isβ) ( is ∇L(β,θ∗)− s2 +O( 2) ) q(k,β)dβ +O( α+1)
= ∫ exp(isβ) ( is ∇L(β,θ∗)− s2 ) q(k,β)dβ +O( α+1).
(48)
Therefore, φ(k+ )(s)− φ(k)(s)
=− is ∫ exp(isβ)∇(−L(β,θ∗))q(k,β)dβ
+ (−is)2 ∫ exp(isβ)q(k,β)dβ +O( α). (49)
For any integrable function f , set F as the Fourier transform defined by
F[f(x)](s) = 1√ 2π
∫ f(x) exp(isx)dx. (50)
The inverse Fourier transform of F[f(x)] and the l-th order derivatives of f (l)(x) is
F−1[f(x)](s) = f(x) = 1√ 2π
∫ F[f(x)](s) exp(−isx)dx,
F(f (l))(s) = (−is)l(F(f))(s). (51)
Combine Eq.(49), Eq.(50) and Eq.(51), we arrive at the following simplified equation:
φ(t+ )(s)− φ(k)(s)√ 2π
=− isF∂β(−L(β,θ∗))q(k,β) + (−is)2Fq(k,β) +O( α) =F∂β (∂β(−L(β,θ∗))q(k,β)) + F∂2βq(k,β) +O( α).
(52)
Since F−1[φ(k)(s)] = √ 2πq(k,β) and α ∈ (0, 1],
lim →0
q(k + ,β)− q(k,β)
= lim →0
F−1 [ φ(k+ )(s)− φ(k)(s)√
2π ] = lim →0 ∂β(∂β(−L(β,θ∗))q(k,β)) + ∂2βq(k,β) +O( α) =∂β(∂β(−L(β,θ∗))q(k,β)) + ∂2βq(k,β).
(53)
Finally, we have proved that the distribution of β(k) converges weakly to the invariant distribution eL(β,θ ∗) as → 0.

heading: C SIMULATION OF LOGISTIC REGRESSION
text: Now we conduct the experiments on binary logistic regression. The setup is similar as before, except n is set to 500, Σi,j = 0.3|i−j| and η ∼ N (0, I/2). We set the learning rates in SGLD-SA and SGLD to 0.01× k− 13 and step size ω(k) to 0.1× (k + 1000)− 34 . The binary response values are simulated from Bernoulli(p) where p = 1/(1 + e−Xβ−η). Fig.4(a), Fig.4(b) and Fig.4(c) demonstrate the posterior distribution of SGLD-SA is significantly better than that of SGLD. As shown in Fig.4(f), SGLD-SA is the best method to regulate the over-fitting space and provides the most reasonable posterior mean. Table.3 illustrates the predictive power of SGLD-SA is overall better than the other methods and robust to different initializations. Fig.4(d) and Fig.4(e) show that the over-fitting problem of SGLD when p > n in logistic regression and the algorithm fails to regulate the over-fitting space; We observe SGLD-SA is able to resist over-fitting and always yields reproducible results.

heading: D EXPERIMENTAL SETUP

heading: D.1 NETWORK ARCHITECTURE
text: The first DNN we use is a standard 2-Conv-2-FC CNN: it has two convolutional layers with a 2 × 2 max pooling after each layer and two fully-connected layers. The filter size in the convolutional layers is 5 × 5 and the feature maps are set to be 32 and 64, respectively (Jarrett et al., 2009). The fully-connected layers (FC) have 200 hidden nodes and 10 outputs. We use the rectified linear unit (ReLU) as activation function between layers and employ a cross-entropy loss.
The second DNN is a 2-Conv-BN-3-FC CNN: it has two convolutional layers with a 2 × 2 max pooling after each layer and three fully-connected layers with batch normalization applied to the first FC layer. The filter size in the convolutional layers is 4× 4 and the feature maps are both set to 64. We use 256× 64× 10 fully-connected layers.

heading: D.2 TEMPERATURE STRATEGY
text: In practice, we observe a suitable temperature setting is helpful to improve the classification accuracy. For example, by setting τ = 100 in the second DNN (see Appendix D.1) we obtain 99.70% on aMNIST. To account for the scale difference of weights in different layers, we apply different temperatures to different layers based on different standard deviations of the gradients in each layer and obtain the results in Tab. 2.

heading: D.3 DATA AUGMENTATION
text: The MNIST dataset is augmented by (1) randomCrop: randomly crop each image with size 28 and padding 4, (2) random rotation: randomly rotate each image by a degree in [−15◦,+15◦], (3) normalization: normalize each image with empirical mean 0.1307 and standard deviation 0.3081.
The FMNIST dataset is augmented by (1) randomCrop: same as MNIST, (2) randomHorizontalFlip: randomly flip each image horizontally, (3) normalization: same as MNIST.

Review summaries:
The authors describe a new method of posterior sampling with latent variables based on SG-MCMC and stochastic approximation (SA). The new method uses a spike and slab prior on the weights of the deep neural networks to encourage sparsity. Experiments on toy regressions, classification and adversarial attacks demonstrate the superiority over SG-MCMC and EMSV.