Bayesian Inference for Large Scale Image Classification

Bayesian inference promises to ground and improve the performance of deep neural networks . It promises to be robust to overfitting , to simplify the training procedure and the space of hyperparameters , and to provide a calibrated measure of uncertainty that can enhance decision making , agent exploration and prediction fairness . Markov Chain Monte Carlo ( MCMC ) methods enable Bayesian inference by generating samples from the posterior distribution over model parameters . Despite the theoretical advantages of Bayesian inference and the similarity between MCMC and optimization methods , the performance of sampling methods has so far lagged behind optimization methods for large scale deep learning tasks . We aim to fill this gap and introduce ATMC , an adaptive noise MCMC algorithm that estimates and is able to sample from the posterior of a neural network . ATMC dynamically adjusts the amount of momentum and noise applied to each parameter update in order to compensate for the use of stochastic gradients . We use a ResNet architecture without batch normalization to test ATMC on the Cifar10 benchmark and the large scale ImageNet benchmark and show that , despite the absence of batch normalization , ATMC outperforms a strong optimization baseline in terms of both classification accuracy and test log-likelihood . We show that ATMC is intrinsically robust to overfitting on the training data and that ATMC provides a better calibrated measure of uncertainty compared to the optimization baseline . 1 Introduction . In contrast to optimization approaches in machine learning that derive a single estimate for the weights of a neural network , Bayesian inference aims at deriving a posterior distribution over the weights of the network . This makes it possible to sample model instances from the distribution over the weights and offers unique advantages . Multiple model instances can be aggregated to obtain robust uncertainty estimates over the network ’ s predictions ; uncertainty estimates are crucial in domains such as medical diagnosis and autonomous driving where following a model ’ s incorrect predictions can result in catastrophe ( Kendall & Gal , 2017 ) . Sampling a distribution , as opposed to optimizing a loss , is less prone to overfitting and more training doesn ’ t decrease test performance . Bayesian inference can also be applied to differential privacy , where each individual sample has increased privacy guarantees ( Wang et al. , 2015 ) , and to reinforcement learning , where one can leverage model uncertainty to balance between exploration and exploitation ( Osband & Van Roy , 2017 ) . Traditional Markov Chain Monte Carlo ( MCMC ) methods like HMC ( Neal et al. , 2011 ) are a standard class of methods for generating samples from the posterior distribution over model parameters . These methods are seldom applied in deep learning because they have traditionally failed to scale well with large datasets and many parameters ( Rajaratnam & Sparks , 2015 ) . Stochastic Gradient MCMC ( SG-MCMC ) methods have fared somewhat better in scaling to large datasets due to their close relationship to stochastic optimization methods . For example the SGLD sampler ( Welling & Teh , 2011 ) amounts to performing stochastic gradient descent while adding Gaussian noise to each parameter update . Despite these improvements , samplers like SGLD are only guaranteed to converge to the correct Algorithm 1 The ATMC sampler . The algorithm accepts the initialized model parameters θ0 , step size h , pre-conditioner m , and momentum noise Dc . 1 : procedure atmc training ( θ0 , h , m , D ) 2 : p0 ← 0 3 : ξ0 ← 0 4 : while t < T do 5 : Gt ← minibatch gradient ( θt ) 6 : ηt ← random normal ( ) 7 : αt ← max ( D − ξt , 0 ) 8 : βt ← αt + ξt 9 : pt+h ← eβth [ pt − exp [ βth ] −1βt Gt + √ exp [ 2βth ] −1 βt αt ηt ] 10 : θt+h ← θt + hpt+hm 11 : ξt+h ← ξt + h [ p2t+h m − 1 ] 12 : t← t+ h. distribution when the step size is annealed to zero ; additional control variates have been developed to mitigate this to some extent ( Ahn et al. , 2012 ; Ding et al. , 2014 ) . The objective of this work is to make Bayesian inference practical for deep learning by making SG-MCMC methods scale to large models and datasets . The contributions described in this work fall in three categories . We first propose the Adaptive Thermostat Monte Carlo ( ATMC ) sampler that offers improved convergence and stability . ATMC dynamically adjusts the amount of momentum and noise applied to each model parameter . Secondly , we improve an existing second order numerical integration method that is needed for the ATMC sampler . Third , since ATMC , like other SG-MCMC samplers , is not directly compatible with stochastic regularization methods such as batch normalization ( BatchNorm ) and Dropout ( see Sect . 4 ) , we construct the ResNet++ network by taking the original ResNet architecture ( He et al. , 2016 ) , removing BatchNorm and introducing SELUs ( Klambauer et al. , 2017 ) , Fixup initialization ( Zhang et al. , 2019a ) and weight normalization ( Salimans & Kingma , 2016 ) . We design ResNet++ so that its parameters are easy to sample from and the gradients are well-behaved even in the absence of BatchNorm . We show that the ATMC sampler is able to outperform optimization methods in terms of accuracy , log-likelihood and uncertainty calibration in the following settings . First , when using the ResNet++ architecture for both the ATMC sampler and the optimization baseline , the ATMC sampler significantly outperforms the optimization baseline on both Cifar-10 and ImageNet . Secondly , when using the standard ResNet for the optimization baseline and the ResNet++ for the ATMC sampler , multiple samples of the ATMC that approximate the predictive posterior of the model are still able to outperform the optimization baseline on ImageNet . Using the ResNet++ architecture , the ATMC sampler reduces the need for hyper-parameter tuning since it does not require early stopping , does not use stochastic regularization , is not prone to over-fitting on the training data and avoids a carefully tuned learning rate decay schedule . 2 ATMC Sampler . In this section we define the Stochastic Differential Equation ( SDE ) that gives rise to the ATMC sampler described in Algorithm 1 . A detailed background and framework for constructing SDEs that converge to a target distribution can be found in ( Ma et al. , 2015 ) . 2.1 General form of the SDE . Our starting point for constructing the ATMC sampler is the framework of Stochastic Differential Equations . We are interested in SDEs that converge to a distribution p ( z ) over the vector z ∈ Rd for which we can evaluate ∇ log p ( z ) . Because only the gradient of log p ( z ) is required , it is sufficient to define an energy function H ( z ) = − log p ( z ) + C up to a constant C. As a consequence , we can sample from the posterior distribution p ( θ|x ) by only evaluating the energy function gradient ∇H ( θ ) = −∇ log p ( x , θ ) . The general form of SDEs converging to p ( z ) for which only the gradient of p ( z ) is required is as follows ( Ma et al. , 2015 ) : dz = − [ D ( z ) +Q ( z ) ] ∇H ( z ) dt+ Γ ( z ) dt+ √ 2D ( z ) dWt , Γi ( z ) = d∑ j ∂ [ Dij ( z ) +Qij ( z ) ] ∂zj , ( 1 ) where D ( z ) is a positive-definite matrix that determines the amount of noise , Q ( z ) is a skew-symmetric matrix that mixes energy between variables , Wt is a Wiener process , and Γ ( z ) is a correction factor that compensates for dynamics that depend on the current state z . The ATMC sampler that we propose is an instance of ( 1 ) for specific definitions of H ( z ) , D ( z ) , and Q ( z ) . 2.2 Energy Function . We start by defining the energy function H ( z ) . The energy function for the model posterior p ( θ|x ) is defined by the loss function L ( θ ) = − log p ( x , θ ) . Because the dataset x is generally large , we would like to only evaluate a mini-batch loss L̃ ( θ ) . However , naively using a stochastic gradient in ( 1 ) will result in significant bias ( Chen et al. , 2014 ) . Motivated by the Central Limit Theorem , the stochastic gradient is assumed to follow a Gaussian distribution ∇L̃ ( θ ) ∼ N ( ∇L ( θ ) , B ) where the covariance B is additionally assumed to be diagonal and constant w.r.t . θ . The energy function for the ATMC sampler is defined as : H ( θ , p , ξ ) = L ( θ ) +K ( p ) + 1 2 ( ξ − diag ( B ) 2m ) 2 , ( 2 ) where p is the momentum , K ( p ) defines the momentum distribution , and ξ is a control variate referred to as the temperature . Both p and ξ have the same dimensionality as θ . The hyper-parameter m controls the strength of the coupling between ξ and p. The distribution of the control variate p ( ξ ) depends on the amount of noise B in the stochastic gradient estimate L̃ ( θ ) . 2.3 Noise robust dynamics . Next we define the dynamics Q ( z ) and D ( z ) such that the SDE that results from ( 1 ) can be simulated without the need to evaluate B : D ( θ , p , ξ ) = 0 0 00 α ( ξ ) m+ 12B 0 0 0 0 , Q ( θ , p , ξ ) = ( 0 −I 0I 0 m∇K ( p ) 0 −m∇K ( p ) 0 ) , ( 3 ) where α ( ξ ) is a non-negative function that determines how the temperature ξ affects the amount of noise added to the momentum update . We first illustrate the resulting SDE by using a simpler Gaussian momentum distribution K ( p ) = ‖p‖2 / ( 2m ) . Note that the variance of the momentum Var ( p ) = m is reused in ( 2 ) and ( 3 ) to control the strength of the coupling between ξ and p. This will result in a temperature control with a momentum friction term proportional to ξ , unlike previously reported thermostat MCMC methods ( Ding et al. , 2014 ; Lu et al. , 2016 ) where the friction term is proportional to ξ/m . We substitute the dynamics Q ( z ) and D ( z ) defined in ( 3 ) and energy function H ( z ) defined in ( 2 ) into ( 1 ) : ( dθ dp dξ ) = p/m−∇L̃ ( θ ) − β ( ξ ) p p2/m− 1 dt+ 0 0 00 √2α ( ξ ) m 0 0 0 0 dWt , β ( ξ ) = α ( ξ ) + ξ , ( 4 ) where we use ∇L̃ ( θ ) dt = ∇L ( θ ) dt+ √ BdWt to replace the gradient of the loss with the minibatch estimate . The momentum p is dampened by a friction term β ( ξ ) that depends on the choice of α ( ξ ) . The stochastic gradient noise B does not show up in ( 4 ) due to the particular choice of energy function H ( z ) and dynamics Q ( z ) , D ( z ) . Note however this analysis relies on the assumption that the covariance of the stochastic gradient noise B is constant in θ and a single temperature variable per parameter can only correct for a diagonal covariance B . We do not expect that this assumption will hold in practice and the approximation will therefore lead to bias in the samples . However , annealing the step size h will reduce the error due to mini-batching together with other sources of discretization error ( Welling & Teh , 2011 ) .

Summary: This paper proposes a novel MCMC algorithm (ATMC) that estimates and samples from the posterior distribution of neural network weights. The motivation for this approach is that applying Bayesian inference to deep learning should lead to less overfitting and better uncertainty-calibrated models. Unlike previous work, the proposed method scales to large models (ResNet) and data sets (ImageNet).