Title: ON PAC-BAYES BOUNDS FOR DEEP NEURAL NETWORKS USING THE LOSS CURVATURE Anonymous

Abstract: We investigate PAC-Bayes bounds for deep neural networks through the closely related IB-Lagrangian objective. We first propose to approximate the IB-Lagrangian through a second order Taylor expansion of the randomized loss around the minimum. For the case of Gaussian priors and posteriors with fixed means and diagonal covariance, we are able to derive a lower bound to this approximation that corresponds to an "invalid" PAC-Bayes prior (a prior that is training set dependent). Our lower bound depends only on the flatness of the minimum, and the distance between the prior and posterior means. Through a number of experiments our lower bound implies easy and and hard cases where one can or cannot prove generalization even through "cheating". Motivated by this result, we see that for the easy cases, a simple baseline closely matches our lower bound and is sufficient to achieve a non-vacuous PAC-Bayes bound. Crucially the baseline prior is centered on the random deep neural network initialization. This suggests that a good prior mean choice is the main innovation in recent non-vacuous bounds with further optimization of the PAC-Bayes bound through SGD having a complementary role.

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 Figure 1 : The importance of retaining the original minimum: A number of works have shown the existence of an empirical correlation between generalization error and flatness in the loss landscape around the deep neural network minimum µ 0 . Recently popular PAC-Bayes analyses can be seen as a formal way of quantifying this flatness. However, moving forward from simple empirical correlations, existing optimization based non-vacuous bounds compute implicitly a different minimum μ0 and then evaluate the flatness around that minimum. By contrast we aim to estimate the flatness and a related optimal posterior distribution around the original minimum µ 0 .

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

EQUATION

with some probability 1 -δ where A and B are constants related to the derivation. The PAC-Bayes 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. & 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. 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 cannot 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

EQUATION

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 

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 cannot 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.

3. PAC-BAYESIAN FRAMEWORK:

We consider the hypothesis class H L realized by the feedforward neural network architecture of depth L with coordinate-wise activation functions σ defined as the set of functions 1999)) For any data distribution over X × Y, we have that the following bound holds with probability at least 1 -δ over random i.i.d. samples D = {(x j , y j )} N j=1 of size N drawn form the data distribution:

f θ : X → Y (X ⊆ R p , Y ⊆ R K ) with f θ (x) = σ(σ(...σ(x T W 0 )W 1 )W 2 )..)W L ) where θ ∈ Θ L ⊆ R d is a vectorization of the weights and Θ L = R p×k1 × R k1×k2 × ... × R k L ×K .

EQUATION

)

Here Q is an arbitrary "posterior" distribution over parameter vectors, which may depend on the sample D and on the prior P .

The framework models the complexity of the randomized classifier as the KL-Divergence between the posterior Q and a prior P . The prior P must be valid in that it cannot depend in any way on the training data. On the contrary the posterior Q can be chosen to be any arbitrary distribution. This flexibility allows one to model deterministic neural networks as the mean of an arbitrary posterior distribution, thus deriving results for a stochastic but closely related classifier.

The square root in 3 will make subsequent calculations cumbersome. We will therefore be analyzing the similar objective

EQUATION

which is known as the IB-Lagrangian Achille & Soatto (2018); Tishby et al. (2000) . While the two objectives are not exactly equivalent, we argue through experiments that the differences are minimal and our results translate to the original PAC-Bayes bound 3.

4. QUADRATIC APPROXIMATION AND LOWER BOUND:

The stochastic and non-convex objective 4 is difficult to analyze theoretically. As such we first propose to expand the randomized loss using a Taylor expansion which will make the subsequent analysis tractable. We get

EQUATION

where Q is a centered version of Q. We've made a number of assumptions. In the second line we assumed that the loss L(θ) at the minimum is 0. In line 3 we assumed that the gradient ∂ L(θ)/∂θ at the minimum is also 0, the term O(||η|| 3 ) is negligible and that the loss is locally convex resulting in the quadratic approximation being an upper bound. For a well trained overparametrized DNN assuming that the loss and the gradient at the minimum are zero is reasonable. Assuming local convexity at the minimum is also often a reasonable assumption Hassibi & Stork (1993) being highly informative about parameter relevence. Finally we will be analyzing Gaussian posterior distributions with diagonal or close to diagonal covariance which can be assumed to sufficiently concentrate around the minimum making the local second order approximation meaningful.

4.1. OPTIMAL POSTERIOR:

We make the following modeling choices Q = N (µ 0 , Σ 0 ) and P = N (µ 1 , λΣ 1 ) which are popular in VI and PAC-Bayes literature. We can then show that the optimal posterior covariance of the above objective for fixed prior and posterior means has a closed form solution.

Lemma 4.1. The convex optimization problem min Σ0 E η∼Q [ 1 2 η T H l η] + βKL(Q||P ) where Q = N (µ 0 , Σ 0 ) and P = N (µ 1 , λΣ 1 ) has the global minimum:

EQUATION

where H l captures the curvature in the directions of the parameters, while Σ 1 is a chosen prior covariance.

This can been seen as a Laplace approximation Bishop ( 2006) to the posterior around the MAP solution. In practice we will be using 6 by performing a grid search over the parameters β and λ. Then we will try to find Pareto optimal pairs balancing the accuracy of the randomized classifier and the KL complexity term.

4.2. OPTIMAL PRIOR:

The above solution is not optimal with respect to the prior covariance in that we have up to now chosen it arbitrarily. Furthermore given that the choice of the random initialization as the prior mean 

2N:

using 7, 8. We compute the accuracy of the stochastic classifier with MC for 5 samples. We plot this lower bound with the solid black line. All points above it are feasible. We see that for the Mnist problem the two regions intersect suggesting that we might be able to prove generalization using Gaussians with diagonal covariance. By contrast in the Cifar case the two regions do not intersect suggesting that the prior and posterior means have a prohibitive distance between them, and we cannot prove generalization with diagonal covariances. We ignore these concerns for the moment and optimize the prior covariance directly. The objective is non-convex, however for the case of diagonal prior and posterior covariances we can find the global minimum. Lemma 4.2. The optimal prior and posterior for C β (D;

P, Q) = E η∼Q [ 1 2 η T H l η] + βKL(Q||P )

with Q = N (µ 0 , Σ 0 ) and P = N (µ 1 , λΣ 1 ) and assuming that Σ -1

1 = Λ 1 = diag(Λ 11 , Λ 21 , ..., Λ k1

) and H l = diag(h 1l , h 2l , ..., h kl ) have:

EQUATION

EQUATION

)

where H l encodes the local curvature at the the minimum, µ 1 corresponds to the random initialization (by design) of the DNN, and µ 0 corresponds to the minimum obtained after optimization.

For our choice of Gaussian prior and posterior, the following is a lower bound on the IB-Lagrangian under any Gaussian prior covariance:

EQUATION

where

a il a il (β, µ i0 , µ i1 , h il ) = 1 2 [ h 2 il + 4βh il (µi0-µi1) 2 -h il ].

The above result is intuitively pleasing setting a lower bound to what we can achieve which depends only on the initialization (by design), obtained minimum, curvature at the minimum and the regularization parameter β. In particular the scaling factor λ has disappeared.

We now make some important notes about what one can and cannot prove using these results, by stressing that the above result is a necessary but not a sufficient condition for generalization under our prior and posterior modeling.

• Given a deterministic deep neural network and it's initialization (or other prior mean) one can rule out being able to prove generalization using any choice of diagonal covariances when modeling the priors and posteriors as multivariate Gaussians with fixed means. Modeling with other distributions may give different results † .

• One cannot prove generalization using this result, even in the case when the prior mean is valid (only distribution dependent) and the feasible and non-vacuous sets intersect. One still has to compute the prior and posterior covariances in a valid manner. As such a computationally feasible region given finite data and computational resources as well as privacy constraints, might be much smaller than the one we derive here.

We plot our lower bound for simplified 2 class Mnist and Cifar problems in Figure 2 . We see that while for the Mnist problem the feasible and non-vacuous regions intersect the same is not true for the Cifar problem. What remains is to see if our results apply for the case of valid priors. First we detail a number of computational issues in section 5.

5. COMPUTATIONAL ASPECTS:

We now present a number of computational and memory issues associated with the Hessian of a modern deep neural network. There is ambiguity in the literature about the size of the Hessians that can be computed exactly Kunstner et al. (2019) . There have been few results in this area and the main problem seems to be that the relevant computations are not well supported from common autodifferentiation libraries, such as Tensorflow and Keras. However there is certainty on the fact that storing and manipulating the full Hessian of a number of modern deep neural network architectures would be infeasible as the matrix is of size H ∈ R d×d where d is the number of parameters. As a point of reference a dense uncompressed Numpy matrix for d = 50000 takes up 20GB of memory.

As such we detail in the next section a number of approximations that make both computing and storing the Hessian feasible.

5.1. APPROXIMATING THE FULL HESSIAN:

As noted in Kunstner et al. (2019) the generalized Gauss-Newton approximation of the Hessian H(θ) coincides with the Fisher matrix

F(θ) = n E p θ (y|xn) [∇ θ log p θ (y|x n )∇ θ log p θ (y|x n ) T ]

in modern deep learning architectures. While the Fisher matrix is difficult to compute exactly one can compute an unbiased but noisy estimate as Martens & Grosse ( 2015)

EQUATION

where care must be taken to sample ỹn from the model predictive distribution ỹn ∼ p θ (y|x n ).

Additionally we note that the interpretation of the outputs after the softmax as probabilities is not well grounded theoretically Gal & Ghahramani (2016) . Determining the true predictive distribution requires MC sampling for example by taking multiple dropout samples Gal & Ghahramani (2016) .

We now make two additional notes regarding computational aspects of the above. The approximation of the Hessian can be computed efficiently as the outer product of large but manageable gradient vectors. The main computational burden after we approximate the Hessian, and given that we choose a standard normal prior, is inverting a matrix of the form H + αI. This problem can be tackled in a few different ways. The simplest would be to consider only the diagonal elements of H and the resulting diagonal matrix can be efficiently inverted. However inversion of the full matrix H + αI is also possible recursively using the Sherman-Morrison formula. † While our result is a formal lower bound on what is achievable by 6, it's applicability on direct minimization of the IB-Lagrangian 4 depends on the tightness of the second order approximation. and training accuracy (of the randomized classifier) for different architectures and datasets. Points to the right of the dashed line correspond to non-vacuous pairs. All Mnist bounds are non-vacuous. All Cifar bounds are vacuous. We are able to progressively get tighter bounds by using the diagonal Hessian and then the full layerwise Hessian. The improvement is larger over the more difficult Cifar dataset.

Further issues exist with computing the KL-Divergence of large multivariate Gaussians with nondiagonal covariances in closed form which includes a determinant term that has to be computed with the matrix determinant lemma, as well as sampling efficiently from these distributions. As such we have used the diagonal variant of approximation 10 for our lower bound, but perform a layerwise approximation of the Hessian for all other experiments. We detail this layerwise approximation in the next section, and motivate it theoretically.

5.2. LAYERWISE APPROXIMATION:

For a specific case of the empirical loss we will now derive an upper bound on 4 which is more suitable for optimization.

Lemma 5.1. Assuming the following empirical loss L(θ) = ||f θ (X) -Y|| F with X = [x 0 , ..., x N ] and Y = [y 0 , ..., y N ] the following is an upper bound on the IB Lagrangian given that we are at a local minimum:

EQUATION

where l denotes different layers, j denotes the different neurons at each layer (we assume the same number for simplicity), H lj denotes the local Hessian, and Q lj is a centered version of Q lj . The local Hessian can be computed efficiently as

H lj = 1 N N i=1 z i l-1 z i l-1

T and z i l-1 is the latent representation input to layer l for signal i.

We see that we have managed to upper bound the empirical randomized loss by a scaled sum of quadratic terms involving layerwise Hessian matrices and centered random noise vectors. Intuitively we have reduced the complexity of our optimization problem simply by turning it into a number of separate subproblems. The local Hessians can be computed efficiently from outer products of a forward pass of the dataset. Apart from avoiding using backpropagation, breaking the Hessian into subproblems in this manner allows us to move beyond the simplistic diagonal approximation. Implicitly the Hessian now has a block diagonal structure and the blocks are small enough to be inverted directly for the architectures used in this paper. For architectures with larger latent representations the Sherman-Morrison formula can be used instead. While 5.1 holds for a specific loss which is uncommon in practice, we have found empirically that solving the posterior optimization problem in a layerwise manner gives good results with the more common loss functions.

6. EXPERIMENTS:

We now make a number of experiments on the simplified 2 class Mnist and Cifar datasets. Specifically we test the architecture

input → 300FC → 300FC → #classesFC → output on Mnist and input → 200FC → 200FC → #classesFC → output

on Cifar. We note that the Mnist architecture corresponds exactly to the architecture T-300 2 p.7 in Dziugaite & Roy (2017) . We train each configuration to 100% accuracy and derive the layerwise Hessians. We model the prior and posteriors as multivariate Gaussians centered at the initialization and deterministic solution respectively. For the prior we choose the uninformative unit diagonal covariance, scaled by the free parameter λ. The baseline posterior that we use has the same diagonal covariance as the prior. For the baseline we perform a grid search over λ which increases the complexity negligibly. For the optimized posterior we initially test a diagonal approximation of the Hessian "Diag Hessian" which results in an optimal diagonal covariance. We perform a grid search for the parameters λ and β using formula 6 to derive candidates for the optimal posterior covariance. For each point on the grid we calculate the empirical accuracy over the training set using Monte Carlo sampling and 5 samples, as well as the complexity term KL(Q||P )+ln 2(N -1) δ

. We then choose the Pareto optimal points from all candidates. We plot the results in Figure 3 .

Interestingly we see that for the case of Mnist the baseline is tight with respect to our lower bound and provides non-vacuous bounds. Therefore not much improvement can be achieved using the Hessian approach. This implies a more careful interpretation of the results in Dziugaite & Roy (2017). We see that non-vacuity can also be achieved as a result of the problem being very simple, and the choice of the prior mean being the random initialization. The optimization techniques employed in Dziugaite & Roy (2017) should simply tighten the bound further, mainly by moving the posterior mean closer to the prior mean (the random initialization). For the case of Cifar we see that we can significantly tighten the bound. However we cannot manage to turn a vacuous bound to a non-vacuous one in line with our lower bound. We also test a block-diagonal approximation to the Hessian "Full L Hessian" where each neuron of the network has it's own block. Our nondiagonal layerwise approximation however crude seems to improve significantly over the diagonal case and slightly crosses our diagonal lower bound. This suggests that better approximations of the Hessian as well as better prior means apart from the random initialization might be needed to prove generalization in complex datasets and architectures.

Under review as a conference paper at ICLR 2020 B. PROOFS Lemma 4.1. The convex optimization problem min Σ0 E η∼Q [ 1 2 η T H l η] + βKL(Q||P ) where Q = N (µ 0 , Σ 0 ) and P = N (µ 1 , λΣ 1 ) has the global minimum:

EQUATION

where H l captures the curvature in the directions of the parameters, while Σ 1 is a chosen prior covariance.

Proof.

EQUATION

The gradient with respect to Σ 0 is

EQUATION

Setting it to zero, we obtain the minimizer

Σ * 0 = β(H l + β λ Σ -1 1 ) -1 .

Lemma 4.2. The optimal prior and posterior for C β (D;

P, Q) = E η∼Q [ 1 2 η T H l η]+βKL(Q||P ) with Q = N (µ 0 , Σ 0 ) and P = N (µ 1 , λΣ 1 ) and assuming that Σ -1 1 = Λ 1 = diag(Λ 11 , Λ 21 , ..., Λ k1

) and H l = diag(h 1l , h 2l , ..., h kl ) have:

EQUATION

EQUATION

where H l encodes the local curvature at the the minimum, µ 1 corresponds to the random initialization (by design) of the DNN, and µ 0 corresponds to the minimum obtained after optimization.

For our choice of Gaussian prior and posterior, the following is a lower bound on the IB-Lagrangian under any Gaussian prior covariance:

EQUATION

where

a il a il (β, µ i0 , µ i1 , h il ) = 1 2 [ h 2 il + 4βh il (µi0-µi1) 2 -h il ].

Proof. Setting Λ 1 = Σ -1 1 We can then see that the minimizer is equal to

EQUATION

The above matrix equation 18 is difficult to deal with directly. We will therefore use the common diagonal approximation of the Hessian which is more amenable to manipulation. Substituting Λ 1 = diag(Λ 11 , Λ 21 , ..., Λ k1 ) and H l = diag(h 1l , h 2l , ..., h kl ) in the above expression we get

C β (D; P, Q)| Σ0=Σ * 0 = β 2 ( 1 λ i Λ i1 (µ i0 -µ i1 ) 2 - i ln( Λ i1 λ ) + i ln( h il + β λ Λ i1 β )) (19)

The above expression is easy to optimize. We see that the sole stationary point exists at

EQUATION

We now turn to the original objective and calculate it's second derivatives. For our diagonal approximation the original objective turns into a sum of separable functions. We will analyze the behavior of one of them for simplicity. The result applies to all other functions in the sum.

EQUATION

where we have set

A i = h il 2 , B i = β 2λ , C i = β(µi0-µi1) 2 2λ

, D i = β 2 . Denoting C iβ (D; P, Q) one function from this sum we calculate

EQUATION

and

EQUATION

)

We need to check whether the Hessian matrix is PSD so that the stationary point we found is a local minimum and the function is convex. We do that by calculating whether all principal minors of the Hessian are positive.

EQUATION

)

We see easily that det( Di

EQUATION

A first observation is that this determinant is not always positive and the function is not convex everywhere. However we observe that it is not highly non convex either and the non convexity mainly results from the function tending to infinity logarithmically on one of the boundaries. We now check whether the sole stationary point is always a local minimum. We start by substituting ν i0 = β(h il + β λ 1 νi1 ) -1 in the multiplicand of 26 as the multiplier is positive by definition

EQUATION

Where we substituted ν i1 = Λ -1 i1 as this will make the calculations easier. We now show a useful identity for

Λ i1 = λ 2β [ h 2 il + 4βh il (µi0-µi1) 2 -h il ] (Λ i1 ) 2 = λ 2 4β 2 h 2 il + 4βh il (µ i0 -µ i1 ) 2 -2h il h 2 il + 4βh il (µ i0 -µ i1 ) 2 + h 2 il = λ 2 4β 2 2h il h il -h 2 il + 4βh il (µ i0 -µ i1 ) 2 + 4βh il (µ i0 -µ i1 ) 2 = h il λ β λ 2β h il -h 2 il + 4βh il (µ i0 -µ i1 ) 2 + 2β (µ i0 -µ i1 ) 2 = h il λ β λ (µ i0 -µ i1 ) 2 -Λ i1 (28)

We substitute Λ i1 = Λ i1 in 27 and again develop only the multiplicand

EQUATION

where we have set

A i = 1 ν i0 2 2λ(h il λ(Λ i1 ) -1 +β) 2 > 0.

We have used 28 in lines 4 and 7.

Indeed the stationary point is always a local minimum. What remains is to show that there are no other local minima at the boundaries of the domain. From 21 we see that we only need to evaluate expressions of the form f (ν i0 ) = ν i0 -ln(ν i0 ) and g(ν i1 ) = 1 νi0 + ln(ν i0 ). By application of L'Hôpital's rule it's easy to show that lim T and z i l-1 is the latent representation input to layer l for signal i. 

were in line 3 we use the linearity of expectation, Hölder's inequality due to the non-negativity of the random variables, and Jensen's inequality for the concave square root. In line 4 we hide the Frobenius terms into constants c l . Each error term Êl (θ l ) is only multiplied with Frobenius norm terms || θl || F from the deeper layers. Therefore one can start optimizing from the final layer and proceed to the first while considering c l as constant. In practice we will just consider all c l as unknown scaling factors. In line 5 we expand each Êl (θ l ) term using a Taylor expansion, and subsequently ignore the first term as the DNN is assumed to be well trained and the first derivative will be zero, while terms with order higher than 2 are unimportant. We also use Q to denote the centered version of distribution Q.

Taking the first and second derivatives of the layerwise error with respect to W l we get

∂E l (θ) ∂W l = 1 N N i=1 ∂ ∂W l ||W l z i l-1 -z i l || 2 2 = 1 N N i=1 (W l z i l-1 -z i l )2z i l-1 T (33) ∂ 2 E l (θ) ∂W l ∂W (j,:) l = 1 N N i=1 z i l-1 z i l-1 T (34)

Where the second derivative is with respect to any row W (j,:) l of the weight matrix W l . We see that the full Hessian matrix H l = ∂ 2 E l (θ) ∂ 2 W l then has a block diagonal structure where each block is equal to

H lj = 1 N N i=1 z i l-1 z i l-1

T . Each row W (j,:) l corresponds to a neuron of the layer and for an appropriate choice of prior and posterior with block diagonal covariances it is easy to see that the final form of expression 32 factorizes as C β (D; P, Q)

EQUATION

this completes the proof.

:



Figures and Tables:

Given the loss function (•, •) we can define the population loss: L(θ) := E (x,y)∼P (f θ (x), y) and given a training set of N instances D = {(x j , y j )} N j=1 the empirical loss L(θ) := 1 N N i=1 (f θ (x i ), y i ). The PAC-Bayesian framework McAllester (1999) provides generalization error guarantees for randomized classifiers drawn from a posterior distribution Q. We will use the following form of the PAC-Bayes bound. Theorem 3.1. (PAC-Bayesian theorem McAllester (

Figure 2: Feasible solutions vs Non-Vacuous solutions: We merge the 10 classes in Mnist and Cifar to create simpler 2 class problems. For different values of β we compute the optimal complexity terms KL(Q||P )+ln 2(N -1) δ

result as a prior. Recently Dziugaite & Roy (2018) have proposed to use the original training set to derive valid priors by imposing differential privacy constraints.

Figure 3: Accuracy vs Complexity for different bounds: We plot KL(Q||P )+ln 2(N -1) δ

iβ (ν i0 , ν i1 ) = lim

iβ (ν i0 , ν i1 ) = lim νi0=ct νi1→+∞ C iβ (ν i0 , ν i1 ) = +∞(30)    this concludes the proof.Lemma 5.1. Assuming the following empirical loss L(θ) = ||f θ (X)-Y|| F with X = [x 0 , ..., x N ] and Y = [y 0 , ..., y N ] the following is an upper bound on the IB Lagrangian given that we are at a local minimum:C β (D; P, Q) lj ||P lj ),(31)where l denotes different layers, j denotes the different neurons at each layer (we assume the same number for simplicity), H lj denotes the local Hessian, and Q lj is a centered version of Q lj . The local Hessian can be computed efficiently as H lj = 1

We start by defining a layerwise empirical error Êl (θ l ) := 1 θk || F + ÊL (θ L )Dong et al. (2017) substituting this in the IB Lagrangian we getC β (D; P, Q) = E θ∼Q [ L(θ)] + βKL(Q||P ) θk || F + ÊL (θ L )] + βKL(Q||P ) l η + O(||η|| 3 )] + βKL(Q||P ) l η] + βKL(Q||P )

The resulting approximation is closely linked to the Laplace approximation in Bayesian statisticsBishop (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 tasksMaddox et al.

Table content: <table><tr><td>Similarly</td></tr><tr><td>second order approximations of the loss around a minimum have a long history in the literature of</td></tr><tr><td>DNN compression Dong et al. (2017); Wang et al. (2019); Peng et al. (2019); LeCun et al. (1990);</td></tr></table>

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 worksKawaguchi et al. (2017);Nagarajan & Kolter (2019b);Pitas et al. (2019).

Table content: <table/>

Peer review (iclr2020):
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.) The direction taken by the authors is potentially interesting. However, there are a few issues that would have to be addressed carefully for me to recommend acceptance. First, the comparison to (some very) related work is insufficient, and so the actual novelty is misrepresented (see detailed comments below). Further, the paper is full of questionable vague claims and miscites/attributes other work. At the moment, I think the paper is below the acceptance threshold: the authors need to read and understand (!) related work, and expand their theoretical and/or empirical results to produce a contribution of sufficient novelty/impact. DETAILED FEEDBACK. I believe the authors missed some related work by Tsuzuki, Sato and Sugiyama (2019), where a PAC-Bayes bound was derived in terms of the Hessian, via a second-order approximation. How are the results presented in this submission relate to Tsuzuki et al approach? When the posterior, Q, is a Gaussian (or any other symmetric distribution), eta^T H eta is the so-called Skilling-Hutchinson trace estimator. Thus E(eta^T H eta) is the Trace(H) scaled by the variance of eta. The authors seem to have completely missed this connection, which simplifies the final expression considerably. Why is the assumption that the higher order terms are negligible reasonable? Citation or experiments required. Regarding the off-diagonal Hessian approximation: how does the proposed layer-wise approximation relate to k-FAC (Martens and Grosse 2015)? IB Lagrangian: I am not sure why the authors state the result in Thm 4.2 as a lower bound on the IB Lagrangian. What’s the significance of having a lower bound on IB Lagrangian? Other comments: Introduction: “At the same time neither the non-convex optimization problem solved in .. nor the compression schemes employed in … are guaranteed to converge to a global minimum.”. This is true but it is really not clear what the point being made is. Essentially, so what? Note that PAC-Bayes bounds hold for all posteriors, even ones not centered at the global minimum (of any objective). The claims made in the rest of the paragraph are also questionable and their purposes are equally unclear. I would be grateful if the authors could clarify. First sentence of Section 3.1: “As the analytical solution for the KL term in 1 obviously underestimates the noise robustness of the deep neural network around the minimum...”. I have no idea what is being claimed here. The statement needs to be made much less vague. Please explain. Section 4: “..while we will be minimizing an upper bound on our objective we will be referring with a slight abuse of terminology to our results as a lower bound.”. I would appreciate if the authors could clarify what they mean here. Section 4.1 beginning: “We make the following model assumptions...”. Choosing a Gaussian prior and posterior is not an assumption. It*s simply a choice. The PAC-Bayes bound is valid for any choices of Gibbs classifiers. On the other hand, it is an assumption that such distributions will yield *tight* bounds, related to the work of Alquier et al. Section 4.1 “In practice we perform a grid search over the parameters..”. The authors should mention that such a search should be accounted for via a union bound (or otherwise). The *cost* of such a union bound should be discussed. The empirical risk of Q is computed using 5 MCMC samples. This seems like a very low number, as it would not even give you one decimal point of accuracy with reasonable confidence! The authors should either use more samples, or account for the error in the upper bound using a confidence interval derived from a Chernoff bound. Section 4.2: “The concept of a valid prior has been formalized under the differential privacy setting...”. I am not sure what the authors mean by that. Section 5: “There is ambiguity about the size of the Hessians that can be computed exactly.” What kind of ambiguity? Same paragraph in Section 5 discusses why there are few articles on Hessian computation. The authors claim that “the main problem seems to be that the relevant computations are not well supported...”. This is followed by another comment that is supposed to contrast the previous claim, saying that storing the Hessian is infeasible due to memory requirements. I am not sure how this claim about memory requirements shows a contrast with the claim on computation not being supported. First sentence in Section 5.1: I believe this is only true under some conditions. Section 5.1: The authors should explain why they add a damping term, alpha, to the Hessian, and how alpha affects the results. *** Additional citation issues: The connections between variational inference, PAC-Bayes and IB Lagrangian have been pointed out in previous work (e.g. Germain, Bach, Lacoste, Lacoste-Julien (2016); Achille and Soatto 2017). In the introduction, the authors say “...have been motivated simply by empirical correlations with generalization error; an argument which has been criticized …” (followed by a few citations). Note, that this was first criticized in Dziugaite and Roy (2017). “Both objectives in … are however difficult to optimize for anything but small scale experiments.”. It seems peculiar to highlight this, since the approach that the authors are presenting is actually more computationally demanding. Citations for MNIST and CIFAR10 are missing. *** Minor: Theorem 3.1 “For any data distribution over..”, I think it was meant to be mathcal{X} 	imes (and not in ) Theorem 4.2: “For our choice of Gaussian prior and posterior, the following is a lower bound on the IB-Lagrangian under any Gaussian prior covariance”. I assume only the mean of the Gaussian prior is fixed. Citations are misplaced (breaking the sentences, unclear when the paper of the authors are cited). There are many (!) missing commas, which makes some sentences hard to follow. *** Positive feedback: I thought the visualizations in Figure 2 and 3 were quite nice.