On PAC-Bayes Bounds for Deep Neural Networks using the Loss Curvature

1 INTRODUCTION . Recently two works Dziugaite & Roy ( 2017 ) Zhou et al . ( 2018 ) have made significant progress in proving generalization for deep neural networks . They manage to prove non-vacuous generalization bounds for a simplified Mnist dataset Dziugaite & Roy ( 2017 ) and the Imagenet dataset respectively Zhou et al . ( 2018 ) . This stands in stark contrast with previous works Bartlett et al . ( 2017 ) Neyshabur et al . ( 2017b ) Golowich et al . ( 2017 ) which yield bounds that are vacuous by several orders of magnitude and are usually motivated by simple empirical correlation with the generalization error . Both works rely crucially on applying the PAC-Bayes framework McAllester ( 1999 ) which typically addresses the generalization error of stochastic classifiers . Given the randomized empirical error Eθ∼Q [ L̂ ( θ ) ] , the randomized population error Eθ∼Q [ L ( θ ) ] , a prior weight distribution P , a posterior weight distribution Q and N training samples the PAC-Bayes bound gives a guarantee of the form E θ∼Q [ L ( θ ) ] ≤ E θ∼Q [ L̂ ( θ ) ] +A √ ( KL ( Q||P ) +B ) /N , ( 1 ) with some probability 1 − δ where A and B are constants related to the derivation . The PACBayes bound models the complexity of a classifier as the KL-Divergence between a prior P and a posterior Q weight distribution . Apart from yielding state of the art bounds for SVMs , a further motivation for this framework , is it ’ s Bayesian nature i.e . is the existance of the prior P . Complexity is not measured with respect to an arbitrary reference point , but as in the luckiness framework Shawe-Taylor et al . ( 1998 ) , with respect to a reference point that can potentially incorporate our prior knowledge about good solutions to the classification problem , and can lead to possibly tighter bounds . In Dziugaite & Roy ( 2017 ) the authors model the weights as originating from a Gaussian posterior distribution with diagonal covariance θ = µ + ξ σ where ξ ∼ N ( 0 , I ) . Then they optimize directly the stochastic objective resulting from the PAC-Bayes bound , by approximating stochastic quantities with MC sampling . Similarly in Zhou et al . ( 2018 ) the authorts first compress a DNN with an off the shelf compession algorithm , removing redundant parameters and applying the PAC-Bayes approach to the remaining weights . This can be seen as explicitly minimizing the length of a code describing the DNN . While compression in Zhou et al . ( 2018 ) is done explicitly , Dziugaite & Roy ( 2017 ) can also be seen under this light , the random variables form a variational code whose length is explicitly minimized Blier & Ollivier ( 2018 ) . While these two works result in non-vacuous bounds they have a number of important limitations . • Importantly they provide generalization error guarantees for a different classifier than the original . Compressing the neural network or modelling it as originating from a Gaussian distribution whose mean is modified , results in finding a completely different point in the parameter space than the original . The function that the new weights describe might or might not be close to the original . This might seem like an academic problem . However we argue that analyzing networks that typically result from vanilla SGD is equally important potentially leading to better initialization and regularisation of SGD , as well as discovering inherent limitions of our generalization proof toolbox . • The bounds provided by Dziugaite & Roy ( 2017 ) Zhou et al . ( 2018 ) are non-vacuous but loose . What is the source of this looseness ? Both methods rely in non-convex optimisation that might have simply not converged properly . This is a particular problem in Dziugaite & Roy ( 2017 ) where it is well known that VI techniques require laborious hyperparameter tuning , something that has hindered significantly the wide applicability of Bayesian neural networks Wu et al . ( 2018 ) . Even when care is taken the resuls are widely considered suboptimal in terms of uncertainty estimation for prediction tasks and code length description of the DNN Wu et al . ( 2018 ) Blier & Ollivier ( 2018 ) . As such we argue that one stands to benefit from circumventing these non-convex optimisation procedures when possible . The PAC-Bayes bound can also be seen under the light of flat minima . Flat minima have been empirically shown to correlate with better generalization Neyshabur et al . ( 2017a ) Keskar et al . ( 2016 ) . The PAC-Bayes bound can be interpreted as balancing two terms , the randommized empirical loss and the KL complexity term Neyshabur et al . ( 2017a ) . As the posterior can be arbitrarily chosen increasing the variance of the noise added to the parameters will typically decrease the KL term while increasing the empirical loss . If the neural network solution corresponds to a flat minimum , more noise can be added to the parameters without affecting the empirical performance Neyshabur et al . ( 2017a ) . In this work we propose to tackle the above problems by using a second order Taylor expansion of the loss around the minimum , in conjunction with the PAC-Bayes framework . The second order term parameterized by the Hessian matrix corresponds to the curvature of the loss around the minimum . In line with the works linking flat minima to better generalization Neyshabur et al . ( 2017a ) Keskar et al . ( 2016 ) , we are effectively estimating the flatness of the minimum along all parameter dierections . Thus we are able to add more noise to the flat directions and less noise to the curved ones , something that will typically decrease significantly the KL term in the PAC-Bayes framework while ensuring that the loss of the stochastic classifier remains small . Crucially for the case of Gaussian posteriors the optimal posterior covariance can be found in closed form , allowing us to circumvent a non-convex optimization procedure by incuring an approximation penalty due to using a second order approximation of the loss . The resulting approximation is closely linked to the Laplace approximation in Bayesian statistics Bishop ( 2006 ) . It and similar approximations to the posterior have a appeared a number of times in the literature of bayesian neural networks yielding good results in a number of tasks Maddox et al . ( 2019 ) ; Khan et al . ( 2019 ) ; Ritter et al . ( 2018 ) ; Zhang et al . ( 2017 ) ; Khan et al . ( 2018 ) . Similarly second order approximations of the loss around a minimum have a long history in the literature of DNN compression Dong et al . ( 2017 ) ; Wang et al . ( 2019 ) ; Peng et al . ( 2019 ) ; LeCun et al . ( 1990 ) ; Hassibi & Stork ( 1993 ) often yielding state of the art results in parameter reduction . Even though the resulting posterior is optimal with respect to our approximation , being able to chose a data driven informative prior still leaves room for looseness . Importantly the prior in the PAC-Bayes framework can depend on the data generating distribution but not on the training set used to train the evaluated classifier . Workarounds include choosing data driven priors by training a separate classifier on a different training set , or using the same training set but enforcing that the training set is not too informative about each individual training signal . The later can be formalized through the framework of differential privacy . Both of the above will usually involve some nonconvex optimisation procedure leaving again doubt as to the optimality of the classifier complexity estimate . We show that for the case of Gaussian priors and posteriors with diagonal covariance we can derive the optimal invalid prior covariance in closed form . The invalid prior can not be used to prove generalization , however it can be used as a sanity check to see whether proving generalization is possible in principle . The optimal solution with respect to both prior and posterior covariance results in a lower bound on a function closely related to the PAC-Bayes bound ( but not exactly equal given that we make a number of approximations ) . Through experiments we find that , depending on the hardness of the dataset , one can find cases where the set of feasible solutions implied from the lower bound and the set of non-vacuous PAC-Bayes bound solutions don ’ t intersect . One is unable to prove generalization , even through choosing a prior in an invalid manner . A number of works Achille & Soatto ( 2018 ) ; Dziugaite & Roy ( 2017 ) ; Germain et al . ( 2016 ) have noted the similarity between the stochastic PAC-Bayes objective and the objective Cβ ( D ; P , Q ) = E θ∼Q [ L̂ ( θ ) ] + βKL ( Q||P ) . ( 2 ) For β = 1 this is known as the Evidence Lower Bound ( ELBO ) objective in Variational Inference literature Kingma et al . ( 2015 ) Bishop ( 2006 ) . In the Information Bottleneck framework Achille & Soatto ( 2018 ) Tishby et al . ( 2000 ) it is know as the IB-Lagrangian . More recently the same objective has been interpreted as the ” task complexity ” Achille et al . ( 2019 ) . Then β has the role of regulating the amount of information in the randomized neural network Achille & Soatto ( 2018 ) , smaller values correspond to more information and potential to overfit . While the objectives 1 and 2 are not entirely equivalent due to a square root term over the KL divergence in the PAC-Bayes case , we will be using the IB formulation as removing the square root will ease our derivations . We note that Dziugaite & Roy ( 2017 ) have used the two objectives intercheangably with no significant difference in results . A number of preprints have appeared on Arxiv trying to link PAC-Bayes to flat minima Wang et al . ( 2018 ) ; Li et al . ( 2019 ) ; Tsuzuku et al . ( 2019 ) ; Yang et al . ( 2019 ) . These typically involve heuristic choices for the optimal posterior , do not analyze the role of the prior , and focus on quantities that simply correlate with the generalization error . Motivating vacuous generalization bounds on the basis of empirical correlations with generalization error has been criticised in a number of recent works Kawaguchi et al . ( 2017 ) ; Nagarajan & Kolter ( 2019b ) ; Pitas et al . ( 2019 ) . 2 CONTRIBUTIONS . • We propose to analyze a PAC-Bayes bound through a second order Taylor expansion of the randomized empirical loss in the closely related IB-Lagrangian objective . We thus pay an error resulting from the second order approximation in order to circumvent nonconvex optimisation procedures , which might require extensive hyperparameter tuning and not converge properly . • We are able to find a lower bound to this second order taylor expansion that corresponds to an invalid PAC-Bayes prior ( the prior is training set dependent ) . Experimentally we find cases where the set of feasible solutions implied by the lower bound does not intersect with the set of non-vacuous PAC-Bayes accuracy-complexity pairs . This in turn implies easy and hard cases were one can and can not prove generalization using Gaussian priors and posteriors with diagonal covariance . • While we rely on approximations , we see empirically that in easy cases one can find nonvacuous generalization bounds using a very simple baseline where the prior is crucially centered at the random DNN initialization . This in turn highlights the insightful prior mean choice as the main innovation in previous non-vacous bound works Dziugaite & Roy ( 2017 ) with bound optimization through SGD having a complementary role . • We motivate a layerwise method to optimise PAC-Bayes bounds with respect to the posterior variance , reducing the required computations significantly .

Summary: The authors  replace the empirical risk term in a PAC-Bayes bound by its second-order Taylor series approximation, obtaining an approximate (?)  PAC-Bayes bound that depends on the Hessian. Note that the bound is likely overoptimistic unless the minimum is quadratic. They purpose to study SGD by centering the posterior at the weights learned by SGD. The posterior variance that minimizes this approximate PAC-Bayes bound can then be found analytically. They also solve for the optimal prior variance (assuming diagonal Gaussian priors/posteriors), producing a hypothetical "best possible bound" (at least under the particular choices of priors/posteriors, and under this approximation of the empirical risk term). The authors evaluate their approximate bound and "best bound possible" empirically on MNIST and CIFAR. This requires  computing approximations of the Hessian for small fully connected neural networks trained on MNIST and CIFAR10. There are some nice visualizations (indeed, these may be one of the most interesting contributions.)