Bayesian Deep Learning via Stochastic Gradient MCMC with a Stochastic Approximation Adaptation

1 INTRODUCTION . 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 . 2 STOCHASTIC GRADIENT MCMC . 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 . 3 STOCHASTIC GRADIENT MCMC WITH A STOCHASTIC APPROXIMATION ADAPTATION . 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 . 3.1 CONJUGATE SPIKE-AND-SLAB PRIOR FORMULATION FOR DNN . 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 ) . 3.2 EXPECTATION STOCHASTIC-MAXIMIZATION . 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 )

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