PAC-Bayesian Neural Network Bounds

Bayesian neural networks , which both use the negative log-likelihood loss function and average their predictions using a learned posterior over the parameters , have been used successfully across many scientific fields , partly due to their ability to ‘ effortlessly ’ extract desired representations from many large-scale datasets . However , generalization bounds for this setting is still missing . In this paper , we present a new PAC-Bayesian generalization bound for the negative log-likelihood loss which utilizes the Herbst Argument for the log-Sobolev inequality to bound the moment generating function of the learners risk . We explore the generalization and calibration properties of the learned posterior on several image classification benchmarks , showing that the proposed approach provides better generalization and uncertainty estimates . 1 INTRODUCTION . Deep neural networks are ubiquitous across disciplines and often achieve state of the art results ( e.g. , Krizhevsky et al . ( 2012 ) ; Simonyan & Zisserman ( 2014 ) ; He et al . ( 2016 ) ) . Albeit neural networks are able to encode highly complex input-output relations , in practice , they do not tend to overfit ( Zhang et al. , 2016 ) . This tendency to not overfit has been investigated in numerous works on generalization bounds ( Langford & Shawe-Taylor , 2002 ; Langford & Caruana , 2002 ; Bartlett et al. , 2017a ; 2019 ; McAllester , 2003 ; Germain et al. , 2016 ; Dziugaite & Roy , 2017 ) . Indeed , many generalization bounds apply to neural networks . However , most of these bounds assume that the loss function is bounded ( Bartlett et al. , 2017a ; Neyshabur et al. , 2017 ; Dziugaite & Roy , 2017 ) . Unfortunately , this assumption excludes the popular negative log-likelihood ( NLL ) loss , which is instrumental to Bayesian neural networks that have been used extensively to calibrate model performance and provide uncertainty measures to the model prediction . In this work we introduce a new PAC-Bayesian generalization bound for NLL loss of deep neural networks . Our work utilizes the Herbst argument for the logarithmic-Sobolev inequality ( Ledoux , 1999 ) in order to bound the moment-generating function of the model risk . Broadly , our PACBayesian bound is comprised of two terms : The first term is dominated by the norm of the gradients with respect to the input and it describes the expressivity of the model over the prior distribution . The second term is the KL-divergence between the learned posterior and the prior , and it measures the complexity of the learning process . In contrast , bounds for linear models or bounded loss functions lack the term that corresponds to the expressivity of the model over the prior distribution and therefore are the same when applied to shallow and deep models . We empirically show that our PAC-Bayesian bound is tightest when we learn the mean and variance of each parameter separately , as suggested by Blundell et al . ( 2015 ) in the context of Bayesian neural networks ( BNNs ) . We also show that the proposed bound holds different insights regarding model architecture , optimization and prior distribution selection . We demonstrate that such optimization minimizes the gap between risk and the empirical risk compared to the standard Bernoulli dropout and other Bayesian inference approximation while being consistent with the theoretical findings . Additionally , we explore in-distribution and out-of-distribution examples to show that such optimization produces better uncertainty estimates than the baseline . 2 RELATED WORK . Generalization bounds for neural networks were explored in various settings . VC-theory provides both upper bounds and lower bounds to the network VC-dimension , which are linear in the number of network parameters ( Bartlett et al. , 2017b ; 2019 ) . While VC-theory asserts that such a model should overfit as it can learn any random labeling ( e.g. , Zhang et al . ( 2016 ) ) , surprisingly , neural networks generally do not overfit . Rademacher complexity allows to apply data dependent bounds to neural networks ( Bartlett & Mendelson , 2002 ) . However , they rely on bounds of the loss function ( Wan et al. , 2013 ; Gao & Zhou , 2016 ) . Rademacher complexity bounds may also take into account the Lipschitz constant of the network . While this presents an improvement when its gradient-norms are bounded , these results also require the loss to be bounded . Stability bounds may be applied to unbounded loss functions and in particular to the negative log-likelihood ( NLL ) loss ( Bousquet & Elisseeff , 2002 ; Rakhlin et al. , 2005 ; Shalev-Shwartz et al. , 2009 ; Hardt et al. , 2015 ; Zhang et al. , 2016 ) . However , stability bounds for convex loss functions , e.g. , for logistic regression , do not apply to neural networks since they require the NLL loss to be a convex function of the parameters . Alternatively , Hardt et al . ( 2015 ) estimate the stability of stochastic gradient descent dynamics , which strongly relies on early stopping . This approach results in weaker bounds for the non-convex setting . PAC-Bayesian bounds were recently applied to neural networks ( Langford & Caruana , 2002 ; McAllester , 2013 ; Dziugaite & Roy , 2017 ; Neyshabur et al. , 2017 ) . In contrast to our work , those related works all consider bounded loss functions . PAC-Bayesian bounds for the NLL loss in the online setting were put forward by Banerjee ( 2006 ) . The online setting does not to consider the whole sample space and therefore is simpler to analyze in the Bayesian setting . An excellent survey on PAC-Bayesian bounds was provided by Germain et al . ( 2016 ) . Alquier et al . ( 2016 ) introduced PAC-Bayesian bounds for linear classifiers with the hinge-loss . Germain et al . ( 2016 ) is closer to our setting and considers PAC-Bayesian bounds for the quadratic loss function , while proving that the quadratic loss function has exponential decay ( sub-gamma ) when the data is sampled from a Gaussian distribution . This result is a PAC-Bayesian generalization bound for the NLL loss of a regressor which assumes Gaussian model over a continuous label space y . Our work differs from this work in two important aspects : ( 1 ) our work does not consider linear classifiers and apply to any deep neural networks ; ( 2 ) our bound considers the NLL loss for classification tasks , for which the label space y is discrete . PAC-Bayesian bounds for the NLL loss function are intimately related to learning Bayesian inference ( Germain et al. , 2016 ) . Recently many works applied various posteriors in Bayesian neural networks . Gal & Ghahramani ( 2015 ) ; Gal ( 2016 ) introduce a Bayesian inference approximation using Monte Carlo ( MC ) dropout , which approximates a Gaussian posterior using Bernoulli dropout . Srivastava et al . ( 2014 ) introduced Gaussian dropout which effectively creates a Gaussian posterior that couples between the mean and the variance of the learned parameters . Kingma et al . ( 2015 ) explored the relation of this posterior to log-uniform priors , while Blundell et al . ( 2015 ) suggests to take a full Bayesian perspective and learn separately the mean and the variance of each parameter . Our work uses the bridge between PAC-Bayesian bounds and Bayesian inference , as described by Germain et al . ( 2016 ) , to find the optimal prior parameters in PAC-Bayesian setting and apply it in the Bayesian setting . Most of the literature regarding Bayesian modeling involves around a two-step formalism ( Bernardo & Smith , 2009 ) : ( 1 ) a prior is specified for the parameters of the deep net ; ( 2 ) given the training data , the posterior distribution over the parameters is computed and used to quantify predictive uncertainty . Since exact Bayesian inference is computationally intractable for neural networks , approximations are used , including MacKay ( 1992 ) ; Hernández-Lobato & Adams ( 2015 ) ; Hasenclever et al . ( 2017 ) ; Balan et al . ( 2015 ) ; Springenberg et al . ( 2016 ) . In this study we follow this two-step formalism , particularly we follow a similar approach to Blundell et al . ( 2015 ) in which we learn the mean and standard deviation for each parameter of the model using variational Bayesian practice . Our experimental validation emphasizes the importance of learning both the mean and the variance . 3 BACKGROUND . Generalization bounds provide statistical guarantees on learning algorithms . They measure how the learned parameters w perform on test data given their performance on the training data S = { ( x1 , y1 ) , . . . , ( xm , ym ) } , where xi is the data instance and yi is its corresponding label . The performance of the learning algorithm is measured by a loss function ` ( w , x , y ) . The risk of a learner is its average loss , when the data instance and its label are sampled from their true but unknown distribution D. We denote the risk by LD ( w ) = E ( x , y ) ∼D ` ( w , x , y ) . The empirical risk is the average training set loss LS ( w ) = 1m ∑m i=1 ` ( w , xi , yi ) . PAC-Bayesian theory bounds the risk of a learner Ew∼qLD ( w ) when the parameters are averaged over the learned posterior distribution q . The parameters of the posterior distribution are learned from the training data S. In our work we focus on the following PAC-Bayesian bound : Theorem 1 ( Alquier et al . ( 2016 ) ) . Let KL ( q||p ) = ∫ q ( w ) log ( q ( w ) /p ( w ) ) dw be the KLdivergence between two probability density functions p , q . For any λ > 0 and for any δ ∈ ( 0 , 1 ] and for any prior distribution p , with probability at least 1 − δ over the draw of the training set S , the following holds simultaneously for any posterior distribution q : Ew∼q [ LD ( w ) ] ≤ Ew∼q [ LS ( w ) ] + logEw∼p , S∼Dm [ eλ ( LD ( w ) −LS ( w ) ) ] +KL ( q||p ) + log ( 1/δ ) λ . ( 1 ) PAC-Bayesian theory is intimately connected to Bayesian inference when considering the negative log-likelihood loss function ` ( w , x , y ) = − log p ( y|x , w ) and λ = m. Germain et al . ( 2016 ) proved that the optimal posterior in this setting is q ( w ) = p ( w|S ) . Bayesian inference considers the posterior p ( y|x , S ) = ∫ p ( w|S ) p ( y|x , w ) dw , at test time for a data instance x , which corresponds to the risk of the optimal posterior . Unfortunately , the optimal posterior is rarely available , and PAC-Bayes relies on the approximated posterior q. Coincidently , the approximated posterior and its KL-divergence from the prior distribution are instrumental to the evidence lower bound ( ELBO ) , which is extensively used in Bayesian neural networks ( BNNs ) to bound the log-likelihood ∑m i=1 log p ( yi|xi ) : − m∑ i=1 log p ( yi|xi ) ≤ − m∑ i=1 Ew∼q log p ( yi|xi , w ) +KL ( q||p ) ( 2 ) While the right hand side of a PAC-Bayesian bound , with the negative log-likelihood loss and λ = m , is identical to the right hand side of the ELBO bound in term of learning , they serve different purposes . One is used for bounding the risk while the other is used for bounding the marginal loglikelihood . Nevertheless , the same algorithms can be used to optimize BNNs and PAC-Bayesian intuitions and components can influence the practice of Bayesian neural networks .

Summary: The paper offers PAC-generalization bounds for Bayesian Neural Networks based on a previous result by Alquier et al. (Theorem 1) which connects the generalization gap to the log partition function of the same gap for the prior distribution on the learned parameters (which is identical to the ELBO bound used in Bayesian neural networks for NLL loss). Due to the fact that the optimal bound occurs for the true posterior, the PAC-bayesian bounds offer a novel interpretation as an objective for BNNs.