Year: 2019

Abstract:
We investigate PAC-Bayes bounds for deep neural networks through the closely related IB-Lagrangian objective. We first propose to approximate the IBLagrangian 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.

heading: 1 INTRODUCTION
text: 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 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
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).

heading: 2 CONTRIBUTIONS
text: • 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.

heading: 3 PAC-BAYESIAN FRAMEWORK
text: We consider the hypothesis class HL realized by the feedforward neural network architecture of depth L with coordinate-wise activation functions σ defined as the set of functions fθ : X → Y (X ⊆ Rp,Y ⊆ RK) with fθ(x) = σ(σ(...σ(xTW0)W1)W2)..)WL) where θ ∈ ΘL ⊆ Rd is a vectorization of the weights and ΘL = Rp×k1 × Rk1×k2 × ...× RkL×K . 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 = {(xj ,yj)}Nj=1 the empirical loss L̂(θ) := 1N ∑N i=1 `(fθ(xi),yi).
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 (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 = {(xj ,yj)}Nj=1 of size N drawn form the data distribution:
E θ∼Q [L(θ)] ≤ E θ∼Q [L̂(θ)] +
√ KL(Q||P ) + ln 2(N−1)δ
2N . (3)
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
Cβ(D;P,Q) = E θ∼Q [L̂(θ)] + βKL(Q||P ), (4)
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.

heading: 4 QUADRATIC APPROXIMATION AND LOWER BOUND
text: 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
Cβ(D;P,Q) = E θ∼Q [L̂(θ)] + βKL(Q||P )
≈ E η∼Q′ [
( ∂L̂(θ)
∂θ
)T η + 1
2 ηTHη +O(||η||3)] + βKL(Q||P )
≤ E η∼Q′
[ 1
2 ηTHη] + βKL(Q||P ),
(5)
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 termO(||η||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 Sagun et al. (2017); Li et al. (2018). Regarding the validity of the quadratic approximation we note that it often provides state of the art results in DNN compression Dong et al. (2017); Wang et al. (2019); Peng et al. (2019); LeCun et al. (1990); 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.

heading: 4.1 OPTIMAL POSTERIOR
text: We make the following modeling choicesQ = 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′ [ 12η THlη] + βKL(Q||P ) where Q =

heading: N (µ0,Σ0) and P = N (µ1, λΣ1) has the global minimum:
text: Σ∗0 = β(Hl + β λ Σ−11 ) −1, (6)
where Hl 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.

heading: 4.2 OPTIMAL PRIOR
text: 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 has been independently shown to result in much tighter bounds in a variety of settings Dziugaite & Roy (2017); Nagarajan & Kolter (2019a) one would wish to isolate the effects of the prior mean on the bound tightness from the prior covariance.
PAC-Bayesian theory allows one to choose an informative prior, however the prior can only depend on the data generating distribution and not the training set. A number of previous works ParradoHernández et al. (2012); Catoni (2003); Ambroladze et al. (2007) have used this insight mainly on simpler linear settings and usually by training a classifier on a separate training set and using the
terms
√
KL(Q||P )+ln 2(N−1)δ 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.
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.
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′ [ 12η THlη] + βKL(Q||P )
with Q = N (µ0,Σ0) and P = N (µ1, λΣ1) and assuming that Σ−11 = Λ1 = diag(Λ11,Λ21, ...,Λk1) and Hl = diag(h1l, h2l, ..., hkl) have:
Λ∗i1 = λ
2β [
√ h2il +
4βhil (µi0 − µi1)2 − hil], (7)
Λ∗i0 = 1
2β [hil +
√ h2il +
4βhil (µi0 − µi1)2 ]. (8)
where Hl 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:
min Σ0,Σ1
Cβ(D;P,Q) & 1 2 ( ∑ i ail(µi0 − µi1)2 + β ∑ i ln( hil + ail ail )), (9)
where ail , ail(β, µi0, µi1, hil) = 12 [ √ h2il + 4βhil (µi0−µi1)2 − hil].
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.

heading: 5 COMPUTATIONAL ASPECTS
text: 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 ∈ Rd×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.

heading: 5.1 APPROXIMATING THE FULL HESSIAN
text: As noted in Kunstner et al. (2019) the generalized Gauss-Newton approximation of the Hessian H(θ) coincides with the Fisher matrix F(θ) = ∑ n Epθ(y|xn)[∇θ log pθ(y|xn)∇θ log pθ(y|xn)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)
F(θ) ≈ ∑ n [∇θ log pθ(ỹn|xn)∇θ log pθ(ỹn|xn)T], (10)
where care must be taken to sample ỹn from the model predictive distribution ỹn ∼ pθ(y|xn). 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.
(a) (b)
Figure 3: Accuracy vs Complexity for different bounds: We plot √
KL(Q||P )+ln 2(N−1)δ 2N and train-
ing 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.

heading: 5.2 LAYERWISE APPROXIMATION
text: 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 = [x0, ...,xN ] and Y = [y0, ...,yN ] the following is an upper bound on the IB Lagrangian given that we are at a local minimum:
Cβ(D;P,Q) . ∑ l √∑ j clj E η∼Q′lj [ 1 2 ηTHljη] + β ∑ l,j KL(Qlj ||Plj), (11)
where l denotes different layers, j denotes the different neurons at each layer (we assume the same number for simplicity), Hlj denotes the local Hessian, and Q′lj is a centered version of Qlj . The
local Hessian can be computed efficiently as Hlj = 1N ∑N i=1 z i l−1z i l−1
T and zil−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.

heading: 6 EXPERIMENTS
text: 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-3002 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)δ 2N .
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.

heading: 7 CONCLUSION
text: We have presented a lower bound on an approximation of the IB-Lagrangian for the case of multivariate Gaussian priors and posteriors with diagonal covariance. This coincides with a lower bound on a PAC-Bayesian generalization bound for an invalid (training set dependent) prior. For cases where the feasible and non-vacuous regions intersect we have seen that it is possible to reach the lower bound and achieve non-vacuous bounds by using valid non-informative priors. We have also presented closed form solutions for the optimal posteriors given fixed means under our modeling assumptions, and motivated theoretically breaking the estimation into layerwise subproblems. Crucially all results depend on high quality estimates of the Hessian which remains an open topic of research for large scale modern deep neural networks.

heading: A. ADDITIONAL EXPERIMENTS
text: We test the architectures
input→ 300FC→ 300FC→ #classesFC→ output
on Mnist and input→ 200FC→ 200FC→ #classesFC→ output
on Cifar. We conduct additional experiments on the original Cifar10 and Mnist10 datasets, as well as Cifar5 and Mnist5 where we merge the 10 classes into 5. The results are consistent across datasets, with more improvement when incorporating the Hessian for more difficult datasets.
(a) (b) (c)
(d) (e) (f)
Figure 4: Accuracy vs Complexity for different bounds: We plot √
KL(Q||P )+ln 2(N−1)δ 2N and train-
ing 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.

heading: B. PROOFS
text: Lemma 4.1. The convex optimization problem minΣ0 E η∼Q′ [ 12η THlη] + βKL(Q||P ) where Q = N (µ0,Σ0) and P = N (µ1, λΣ1) has the global minimum:
Σ∗0 = β(Hl + β λ Σ−11 ) −1, (12)
where Hl captures the curvature in the directions of the parameters, while Σ1 is a chosen prior covariance.
Proof.
Cβ(D;P,Q) = E η∼Q′
[ 1
2 ηTHlη] + βKL(Q||P ) =
E η∼Q′
[ 1
2 tr(Hlηη
T )] + βKL(Q||P ) =
1 2 tr(Hl E η∼Q′ [ηηT ]) + βKL(Q||P ) =
1 2 tr(HlΣ0) + β 2 (tr( 1 λ Σ−11 Σ0)− k + 1 λ (µ0 − µ1)TΣ−11 (µ0 − µ1)
+ ln ( detλΣ1 det Σ0 ) )
(13)
The gradient with respect to Σ0 is
∂Cβ(D;P,Q) ∂Σ0 = [ 1 2 Hl + β 2λ Σ−11 − β 2 Σ−10 ]. (14)
Setting it to zero, we obtain the minimizer Σ∗0 = β(Hl + β λΣ −1 1 ) −1.
Lemma 4.2. The optimal prior and posterior forCβ(D;P,Q) = E η∼Q′ [ 12η THlη]+βKL(Q||P ) with

heading: Q = N (µ0,Σ0) and P = N (µ1, λΣ1) and assuming that Σ−11 = Λ1 = diag(Λ11,Λ21, ...,Λk1)
text: and Hl = diag(h1l, h2l, ..., hkl) have:
Λ∗i1 = λ
2β [
√ h2il +
4βhil (µi0 − µi1)2 − hil], (15)
Λ∗i0 = 1
2β [hil +
√ h2il +
4βhil (µi0 − µi1)2 ]. (16)
where Hl 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:
min Σ0,Σ1
Cβ(D;P,Q) & 1 2 ( ∑ i ail(µi0 − µi1)2 + β ∑ i ln( hil + ail ail )), (17)
where ail , ail(β, µi0, µi1, hil) = 12 [ √ h2il + 4βhil (µi0−µi1)2 − hil].
Proof. Setting Λ1 = Σ−11 We can then see that the minimizer is equal to Σ ∗ 0 = β(Hl + β λΛ1) −1. Substituting Σ0 = Σ∗0 in Cβ(D;P,Q) we obtain:
Cβ(D;P,Q)|Σ0=Σ∗0 = Eη∼Q[ 1 2 ηTHlη] + βKL(Q||P )|Σ0=Σ∗0 =
1 2 tr(Hlβ(Hl + β λ Λ1) −1) + β 2 (tr( 1 λ Λ1β(Hl + β λ Λ1) −1)
+ 1
λ (µ0 − µ1)TΛ1(µ0 − µ1)− k + ln
( detλΛ−11
detβ(Hl + β λΛ1) −1
) )
= β
2 tr(Hl(Hl +
β λ Λ1) −1) + β2 2λ (tr(Λ1(Hl + β λ Λ1) −1))
+ β
2 (+
1 λ (µ0 − µ1)TΛ1(µ0 − µ1)− k + ln
( detλΛ−11
detβ(Hl + β λΛ1) −1
) )
= β
2 (tr((Hl +
β λ Λ1)(Hl + β λ Λ1) −1)
1 λ (µ0 − µ1)TΛ1(µ0 − µ1)− k + ln
( detλΛ−11
detβ(Hl + β λΛ1) −1
) )
= β
2 [+
1 λ (µ0 − µ1)TΛ1(µ0 − µ1) + ln
( detλΛ−11
detβ(Hl + β λΛ1) −1
) ]
(18)
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 Hl = diag(h1l, h2l, ..., hkl) in the above expression we get
Cβ(D;P,Q)|Σ0=Σ∗0 = β 2 ( 1 λ ∑ i Λi1(µi0 − µi1)2 − ∑ i ln( Λi1 λ ) + ∑ i ln( hil + β λΛi1 β )) (19)
The above expression is easy to optimize. We see that the sole stationary point exists at
Λ∗i1 = λ
2β [
√ h2il +
4βhil (µi0 − µi1)2 − hil]. (20)
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.
Cβ(D;P,Q) = ∑ i hil 2 νi0 + ∑ i β 2λ νi0 νi1 − ∑ i β 2 + ∑ i β(µi0 − µi1)2 2λ 1 νi1
+ β 2 [ ∑ i ln(λνi1)− ∑ i ln(νi0)]
= ∑ i Aiνi0 + ∑ i Bi νi0 νi1 − ∑ i β 2 + ∑ i Ci 1 νi1 +Di[ ∑ i ln(λνi1)− ∑ i ln(νi0)]
(21)
where we have set Ai = hil2 , Bi = β 2λ , Ci = β(µi0−µi1)2 2λ , Di = β 2 . Denoting Ciβ(D;P,Q) one function from this sum we calculate
∂Ciβ(D;P,Q) ∂νi0 = Ai + Bi ν1i − Di νi0 , ∂Ciβ(D;P,Q) ∂νi1 = −Biνi0 ν2i1 − Ci ν2i1 + Di νi1
(22)
and ∂Ciβ(D;P,Q)
∂2νi0 = Di ν2i0 , ∂Ciβ(D;P,Q) ∂2νi1 = 2(Biνi0 + Ci) 1 ν3i1 − Di ν2i1
(23)
∂Ciβ(D;P,Q) ∂νi0∂νi1 = −Bi ν2i1 , ∂Ciβ(D;P,Q) ∂νi1∂νi0 = −Bi ν2i1
(24)
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.
∇2Ciβ(νi0, νi1) =
[ Di ν2i0
− Bi ν2i1
− Bi ν2i1 2(Biνi0 + Ci) 1 ν3i1 − Di ν2i1
] (25)
We see easily that det(Di ν2i0 ) > 0. While
det(∇2Ciβ(νi0, νi1)) = Di ν2i0
( 2(Biνi0 + Ci) 1
ν3i1 − Di ν2i1
) − B 2 i
ν4i1
= 1
ν2i0ν 4 i1
( 2CiDiνi1 − (Diνi1 −Biνi0)2 ) = ( 1
ν2i0ν 4 i1
β2
2
)( (µi0 − µi1)2
λ νi1 −
1 2 (νi1 − νi0 λ
)2 ) (26)
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 = β(hil + β λ 1 νi1 )−1 in the multiplicand of 26 as the multiplier is positive by definition
det(∇2Ciβ(ν?i0, νi1)) = 1
ν?i0 2ν4i1
β2
2
( (µi0 − µi1)2
λ νi1 −
1 2 (νi1 − β λ (hil + β λ 1 νi1 )−1)2 ) = 1
ν?i0 2ν4i1
β2
2
( (µi0 − µi1)2
λ νi1 −
1 2 (νi1 − β λ (
νi1λ hilλνi1 + β ))2 )
= 1
ν?i0 2ν4i1
β2
2
( (µi0 − µi1)2
λ νi1 − ν2i1 2 (1− ( β hilλνi1 + β
))2 )
= 1
ν?i0 2ν3i1
β2
2
( (µi0 − µi1)2
λ − νi1 2 ( hilλνi1 hilλνi1 + β
)2 )
= 1
ν?i0 2ν3i1
β2
2
( (µi0 − µi1)2
λ − λ
2h2ilν 3 i1
2(hilλνi1 + β)2 ) = 1
ν?i0 2ν3i12λ(hilλνi1 + β)
2 (2(µi0 − µi1)2(hilλνi1 + β)2 − λ3h2ilν3i1)
= 1
ν?i0 22λ(hilλΛ −1 i1 + β)
2 (2Λi1(µi0 − µi1)2(hilλ+ Λi1β)2 − λ3h2il)
(27)
Where we substituted νi1 = Λ−1i1 as this will make the calculations easier. We now show a useful identity for Λ?i1 = λ 2β [ √ h2il + 4βhil (µi0−µi1)2 − hil]
(Λ?i1) 2
= λ2
4β2
( h2il +
4βhil (µi0 − µi1)2 − 2hil
√ h2il +
4βhil (µi0 − µi1)2 + h2il
)
= λ2
4β2
( 2hil ( hil − √ h2il +
4βhil (µi0 − µi1)2
) +
4βhil (µi0 − µi1)2
)
= hilλ
β
λ
2β
(( hil − √ h2il +
4βhil (µi0 − µi1)2
) +
2β
(µi0 − µi1)2
)
= hilλ
β
( λ
(µi0 − µi1)2 − Λ?i1
) (28)
We substitute Λi1 = Λ?i1 in 27 and again develop only the multiplicand
det(∇2Ciβ(ν?i0, ν?i1)) = 1
ν?i0 22λ(hilλΛ?i1
−1 + β)2 (2Λ?i1(µi0 − µi1)2(hilλ+ Λ?i1β)2 − λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(hilλ+ Λ?i1β)2 − λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + 2hilλΛ?i1β + (Λ?i1) 2 β2)− λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + 2hilλΛ?i1β +
hilλ
β
( λ
(µi0 − µi1)2 − Λ?i1
) β2)− λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + hilλΛ?i1β + βλ2hil (µi0 − µi1)2 )− λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + βλ2hil (µi0 − µi1)2 ) + 2(Λ?i1) 2(µi0 − µi1)2hilλβ − λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + βλ2hil (µi0 − µi1)2 )
+ 2 hilλ
β
( λ
(µi0 − µi1)2 − Λ?i1
) (µi0 − µi1)2hilλβ − λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + βλ2hil (µi0 − µi1)2 ) + 2λ3h2il − 2h2ilλ2(µi0 − µi1)2Λ?i1 − λ3h2il)
= Ai(2Λ ? i1(µi0 − µi1)2(h2ilλ2 + βλ2hil (µi0 − µi1)2 ) + λ3h2il − 2h2ilλ2(µi0 − µi1)2Λ?i1)
= Ai(2Λ ? i1βλ 2hil + λ 3h2il)
> 0 (29)
where we have set Ai = 1ν?i022λ(hilλ(Λ?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’Hôpital’s rule it’s easy to show that
lim νi0→0 νi1=ct Ciβ(νi0, νi1) = lim νi0→+∞ νi1=ct Ciβ(νi0, νi1)
= lim νi0=ct νi1→0 Ciβ(νi0, νi1) = lim νi0=ct νi1→+∞
Ciβ(νi0, νi1) = +∞ (30)
this concludes the proof.
Lemma 5.1. Assuming the following empirical loss L̂(θ) = ||fθ(X)−Y||F with X = [x0, ...,xN ] and Y = [y0, ...,yN ] the following is an upper bound on the IB Lagrangian given that we are at a local minimum:
Cβ(D;P,Q) . ∑ l √∑ j clj E η∼Q′lj [ 1 2 ηTHljη] + β ∑ l,j KL(Qlj ||Plj), (31)
where l denotes different layers, j denotes the different neurons at each layer (we assume the same number for simplicity), Hlj denotes the local Hessian, and Q′lj is a centered version of Qlj . The
local Hessian can be computed efficiently as Hlj = 1N ∑N i=1 z i l−1z i l−1
T and zil−1 is the latent representation input to layer l for signal i. Proof. We start by defining a layerwise empirical error Êl(θl) := 1N ∑N i=1 ||Wlzil−1 − zil ||22. One
can then easily show that L̂(θ) ≤ ∑L−1 k=1 √ Êl(θl) ∏L k=l+1 ||θ̂k||F + √ ÊL(θL) Dong et al. (2017) substituting this in the IB Lagrangian we get
Cβ(D;P,Q) = E θ∼Q [L̂(θ)] + βKL(Q||P )
≤ E θ∼Q [ L−1∑ l=1 √ Êl(θl) L∏ k=l+1 ||θ̂k||F + √ ÊL(θL)] + βKL(Q||P )
≤ L−1∑ l=1 √ E θ∼Q [Êl(θl)] L∏ l=k+1 E θ∼Q [||θ̂l||F ] + √ E θ∼Q [ÊL(θL)] + βKL(Q||P )
≤ L∑ l=1 cl √ E θ∼Q [Êl(θl)] + βKL(Q||P ) ≤ L∑ l=1 cl √√√√ E η∼Q′ [ ( ∂Êl(θl) ∂θl )T η + 1 2 ηTHlη +O(||η||3)] + βKL(Q||P )
≈ L∑ l=1 cl √ E η∼Q′ [ 1 2 ηTHlη] + βKL(Q||P )
(32)
were in line 3 we use the linearity of expectation, Hö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 cl. Each error term Ê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 cl as constant. In practice we will just consider all cl as unknown scaling factors. In line 5 we expand each Ê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 Wl we get
∂El(θ)
∂Wl =
1
N N∑ i=1 ∂ ∂Wl ||Wlzil−1 − zil ||22 = 1 N N∑ i=1 (Wlz i l−1 − zil )2zil−1 T (33)
∂2El(θ)
∂Wl∂W (j,:) l
= 1
N N∑ i=1 zil−1z i l−1 T (34)
Where the second derivative is with respect to any row W(j,:)l of the weight matrix Wl. We see that the full Hessian matrix Hl = ∂2El(θ) ∂2Wl then has a block diagonal structure where each block is
equal to Hlj = 1N ∑N i=1 z i l−1z 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) . ∑ l √∑ j clj E η∼Q′lj [ 1 2 ηTHljη] + β ∑ l,j KL(Qlj ||Plj) (35)
this completes the proof.

Review summaries:
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.)