On the Importance of Gradient Norm in PAC-Bayesian Bounds
Abstract
Generalization bounds which assess the difference between the true risk and the empirical risk have been studied extensively.
However, to obtain bounds, current techniques use strict assumptions such as a uniformly bounded or a Lipschitz loss function.
To avoid these assumptions, in this paper, we follow an alternative approach: we relax uniform bounds assumptions by using on-average bounded loss and on-average bounded gradient norm assumptions.
Following this relaxation, we propose a new generalization bound that exploits the contractivity of the logSobolev inequalities.
These inequalities add an additional loss-gradient norm term to the generalization bound, which is intuitively a surrogate of the model complexity.
We apply the proposed bound on Bayesian deep nets and empirically analyze the effect of this new loss-gradient norm term on different neural architectures.
Introduction
Deep neural networks are ubiquitous across disciplines and often achieve state-of-the-art results.
Despite the fact that deep nets are able to encode highly complex input-output relations, in practice, they do not tend to overfit [Zhang et al., 2016].
This tendency to not overfit has been investigated in numerous works on generalization bounds.
Indeed, many generalization bounds apply to composite functions specified by deep nets.
However, most of these results assume that the loss function satisfies various assumptions, such as uniformly bounded [Alquier et al., 2016], Lipschitz [Alquier et al., 2019], sub-Gaussian [Bégin et al., 2016, Alquier and Guedj, 2018], sub-gamma [Germain et al., 2016], or self-bounding [Haddouche et al., 2021].
Unfortunately, these assumptions exclude many state-of-the-art deep nets, whose properties cannot be statistically determined per data distribution.
In this work, we take a different route and use the deep net gradient-norm as a measure for PACBayesian generalization.
For this purpose, we treat the loss function as a random function that is generated by a high-dimensional probability space, which is governed by the generating process of the training data.
This view allows us to introduce new tools from high-dimensional probability theory that relate measure concentration to expansion of the loss function, as determined by its gradient-norm [van Handel, 2016].
We begin by adjusting a theorem (Herbst theorem) to estimate the measure concentration via the entropy of the loss function (Lemma 3.1).
We then derive a term for the entropy tailored to the multiclass classification setting (discrete labels which condition Gaussian data), (Lemma 3.2).
Finally we bound the entropy using a log-Sobolev inequality (Lemma 3.3).
Importantly, these steps result in a PAC-Bayesian bound which depends on the norm of the gradient of the loss.
Intuitively this norm measures the complexity of the loss function, i.e., the model.
Different from prior work, the bound hence depends on the structure of the employed model and its gradient norm.
This result admits a bound for multi-class classification with linear models and Lipschitz loss (Theorem 3.4), extending a result of Alquier et al.
[2016], and multi-class classification with deep nets when the loss function and its gradient are on-average bounded (Theorem 3.5 and Theorem 3.6).
The assumptions for the derived bound (Theorem 3.5 and Theorem 3.6) are closer to present-day deep net practice emphasizing the importance of the gradients in learning: in addition to the critical importance of controlling the loss gradient to better optimize a deep net, the gradient-norm also controls the generalization of the loss function (as measured by the deep net).
Our contributions: (1) We present a new PAC-Bayesian generalization bound that depends on the gradient-norm of the loss function.
Consequently, we replace the uniform bound assumptions with an on-average bounded loss and an on-average bounded gradient norm.
This answers an open problem raised by Bartlett et al.
[2017a] (cf.
Sec. 4) in the PAC-Bayesian setting.
This result is presented in
Theorem 3.5.
(2) We extend the PAC-Bayesian bounds of Alquier et al.
[2016] for the hinge-loss to any linear model with a Lipschitz loss function.
This result appears in Theorem 3.4.
(3) We empirically demonstrate that our bounds produce tighter generalization performance than the baseline methods.
This helps to bridge the gap between theory and practice towards tighter bounds that have more realistic assumptions that match modern deep nets.
Background
Generalization bounds provide statistical guarantees on learning algorithms.
They assess how the learned parameters w of a model perform on test data given the model’s result on the training data
S = { ( x 1 , y 1 ) , . . .
, ( x m , y m ) } , where x i is the data instance and y i is the corresponding label.
The performance of the parametric model is measured by a loss function ℓ ( w, x, y ) .
The risk
L D ( w ) =
E ( x,y ) ∼ D ℓ ( w, x, y ) of this model is its average loss, when the data instance and its label are sampled from the true but unknown distribution D .
The empirical risk is the average training set loss L S ( w )
= 1 m (cid:80) mi =1 ℓ ( w, x i , y i ) .
2.1 PAC-Bayesian bounds
PAC-Bayesian theory bounds the expected risk E w ∼ q L D ( w ) of a model when its parameters are averaged over the learned posterior distribution q .
The posterior distribution is learned from the training data S .
In our work, we start from the following PAC-Bayesian bound:
Theorem 2.1 (Alquier et al.
Let KL ( q || p )
=
(cid:82) q ( w ) log( q ( w ) /p
( w ))
dw be the KLdivergence between two probability density functions p, q .
For any λ > 0 , for any δ ∈ (0 , 1) and for any prior distribution p , with probability at least 1 − δ over the draw of the training set S , the following holds simultaneously for any posterior distribution q :
[Equation]
where C ( λ, p ) ≜ log (cid:16) E w ∼ p,S ∼ D m [ e λ ( L D ( w ) − L S ( w )) ] (cid:17) .
Unfortunately, the complexity term C ( λ, p ) of this bound is challenging to compute exactly in general: it requires to integrate (and maximize if the log-sum-exp trick is used) over large parameter and data spaces.
Using different assumptions, e.g., learning with sub-Gaussian or sub-gamma loss functions, learning with specific losses, etc., it is possible to analytically bound the complexity term C ( λ, p ) as we briefly illustrate next.
We say that a loss function ℓ is sub-Gaussian with variance
ν 2 if it can be described by a sub-Gaussian random variable, i.e., if its log-moment generating function is upper bounded by ν 2 / 2 :
[Equation]
Alquier et al.
[2016] use Hoeffding’s lemma to show that a uniformly bounded loss function, namely 0 ≤ ℓ ( w, x, y ) ≤ B is a sub-Gaussian loss function with variance B 2 and hence derive from
Theorem 2.1 and Equation (2) a PAC-Bayesian bound for bounded loss functions.
2.2 Measure concentration in high-dimension
However, the Hoeffding lemma, which asserts the sub-Gaussian property to a bounded loss function, treats the value of the loss as a real-valued random variable ℓ ( w, x, y ) while entirely ignoring the high-dimensional probability space ( x, y ) that generates this random value.
Said differently: any properties of a function F transforming the data ( x, y ) are ignored, and all functions F are treated identically.
This is sub-optimal.
In the following, we describe an entropy method that utilizes the measure concentration phenomena that exist in the high-dimensional random space of ( x, y ) , cf.
Chapter 3 of van Handel [2016].
The entropy of a non-negative random variable F is
[Equation]
The Herbst theorem connects Ent[ F ] to measure concentration by providing a bound on the logmoment generating function, as summarized next.
Theorem 2.2 (Herbst) .
Suppose that for all λ > 0 the following bound for the entropy holds:
[Equation]
Then the scaled log-moment generating function which we also refer to as ψ ( λ ) /λ is bounded as follows:
[Equation]
Proof.
First we note that due to the fundamental theorem of calculus
[Equation]
Using l’Hopital rule, we verify that lim α → 0 ψ ( α ) α
= E [ F ] , which yields the first term of the bound.
The following identity
[Equation]
which is obtained from the definition of ψ ( α ) /α and where ( · ) ′ refers to the derivative w.r.t.
α , derives the theorem after plugging Equation (7) into Equation (6) and using Equation (4) to bound the integral.
Importantly, note that this is a step forward.
Different from classical PAC-Bayesian bounds discussed in Section 2.1, which use the fact that a uniformly bounded loss function is sub-Gaussian, the Herbst theorem asserts that any random variable that satisfies Equation (4) is σ 2 sub-Gaussian.
To take advantage of this step, in this work, we benefit from the (modified) log-Sobolev inequality (LSI) for Gaussian random variables, which bounds the entropy.
Theorem 2.3 (Gaussian log-Sobolev inequality (LSI), Gross [1975]) .
Let Z 1 , ..., Z d be independent
Gaussian random variables, i.e., Z i ∼ N ( µ i , σ 2 i ) .
We then have
[Equation]
where σ · ∇ Z f ( Z 1 , ..., Z d ) is an element-wise multiplication.
Proof.
See supplementary material, Appendix A.
One can apply the LSI to the Herbst theorem to attain measure concentration bounds for highdimensional Lipschitz functions: a differentiable function f ( z 1 , ..., z d ) is said to be g -Lipschitz if | f ( z 1 , ..., z d ) − f
( z ′ 1 , ..., z ′ d ) | ≤ g ∥ z − z ′ ∥ , or equivalently, if ∥∇ f ( z 1 , ..., z d ) ∥ ≤ g for any z 1 , ..., z d .
The LSI for g -Lipschitz functions with s ≜ (cid:80) di =1 σ 2 i variance reduces to the bound
Ent[ e λf ( z 1 ,...,z d ) ] ≤ λ 2 g 2 s
E [ e λf ( z 1 ,...,z d ) ] , which in turn provides a bound on the log-moment generating function of f ( z 1 , ..., z d ) using Theorem 2.2.
PAC-Bayesian bounds with log-Sobolev inequalities
Different from prior work, we suggest to measure the generalization of neural nets by taking into account the expected gradient-norm of the loss function with respect to the data generation process.
For this we note that classical PAC-Bayesian generalization bounds depend on measure concentration, as described in Theorem 2.1 and Equation (2.1).
However, the sub-Gaussian assumption used in classical PAC-Bayesian generalization bounds seamlessly bypasses the measure concentration phenomena.
This is too restrictive and enforces all loss functions, and consequently all neural nets, to have the same measure concentration phenomena.
Instead, we first use the Herbst theorem to estimate the measure concentration by taking into account the entropy.
In a second step we use Log-Sobolev inequalities (LSI) to bound the entropy, which allows us to consider the high-dimensional probability space ( x, y ) ∼ D that controls the measure concentration of the loss function ℓ ( w, x, y ) .
This differs from the use of a crude worst-case bound on the loss function (e.g., ℓ ( w, x, y ) ≤ B for any w, x, y ) in classical PAC-Bayesian generalization bounds, which entirely ignores the neural net that generates the loss function.
1) Herbst theorem to estimate measure concentration: As just discussed, we first adjust the Herbst theorem to our setting:
Lemma 3.1 (Herbst) .
For any λ > 0 we have
[Equation]
Proof.
First, we use the statistical independence of the training samples to decompose the moment generating function
[Equation]
as follows:
[Equation]
We apply the differential representation of Theorem 2.2 with ψ ( λ ) =
log M ( λ ) and obtain: ψ ( λ )
[Equation]
Since M ( λ/m ) =
e ψ ( λ/m ) and E ( x,y )
∼D [ ℓ ( w, x, y )]
=
L D ( w ) , we obtain from Equation (14) the identity
M
[Equation]
which completes the proof when being combined with Equation (13).
2) Entropy bound:
In the second step we now aim to bound the entropy Ent D [ e − αℓ ( w,x,y ) ] .
Since we focus on the classification setting, the label y is discrete.
We further assume that for any (discrete) label y its corresponding data x = ( x 1 , ..., x d ) is generated from a d -dimensional
Gaussian distribution that consists of d i.i.d.
Gaussian random variables x i ∼ N ( µ y,i , σ 2 y,i ) .
We denote this Gaussian distribution by x ∼ N ( µ y , σ 2 y ) and abbreviate it by x ∼ N y .
The data generation distribution of x is the Gaussian mixture model
D = (cid:80)
y D y ·N y , where D y is the marginal distributions of D w.r.t y .
We begin by decomposing the entropy of the (exponent) of the loss function according to the labels:
Lemma 3.2.
Let D =
(cid:80)
y D y · N y be a mixture model, whose label marginal probabilities are D y and set f w ( x, y ) =
e − αℓ ( w,x,y ) .
One can show that
[Equation]
Proof.
From the definition of the entropy of a function ℓ ( w, x, y ) with respect to its measure ( x
[Equation]
y ) ∼
D we have
[Equation]
The result is attained by adding and subtracting the quantity
[Equation]
The result then follows from the definitions of entropies of the mixed components, while setting g ( y ) =
E x
∼N y [ f w ( x, y )] .
Specifically we have
[Equation]
It is important to note that the number of components in the mixture model is not limited.
Such a
Gaussian mixture model can approximate any smooth density [Titterington et al., 1985, Scott, 1992,
Goodfellow et al., 2016].
Next, we use the LSI to bound the loss function entropy using its expected gradient-norm with respect to the data generation process.
Lemma 3.3.
Assume that the loss per label is balanced, namely E x
∼N y [ e − αℓ ( w,x,y ) ]
=
c for every w and y , then under the conditions of Lemma 3.2.
[Equation]
Proof.
Lemma 3.2 asserts that Ent D [ e − αℓ ( w,x,y ) ] is composed of the entropies Ent N y
[ e − αℓ ( w,x,y ) ] and Ent D y
[ E x ∼N y [ e − αℓ ( w,x,y ) ]] .
The assumption E x
∼N y [ e − αℓ ( w,x,y ) ]
= c implies the identity Ent D
y
[ E x ∼N y [ e − αℓ ( w,x,y ) ]] =
Ent D y [ c ] = 0 .
Thus Ent D [ e − αℓ ( w,x,y ) ]
= (cid:80) y D y Ent N y
[ e − αℓ ( w,x,y ) ] .
We use Theorem 2.3 to bound
[Equation] [Equation]
Combination with Lemma 3.1 concludes the proof.
Note, the per-label loss balance originates from bounding the term Ent D y
[ E x ∼N y [ f w ]] in Lemma 3.2.
If the loss is per-label balanced, the entropy equals zero.
In fact, the per-label balance can be relaxed by assuming that the loss is per-label bounded within an amplitude (Sec.
3.16 in van Handel [2016]).
Notably, the above lemma is more theoretical than practical: to estimate it in practice, one needs to avoid integrating over α .
Nevertheless, it serves as an essential generalization of the Herbst theorem and the LSI that is useful to derive PAC-Bayesian bounds for various settings.
We look at these different settings next.
3.1 Linear models
In the following we consider differentiable and Lipschitz loss functions 1 over linear models in a multi-class setting of the form ℓ ( w, x, y ) ≜ ˆ ℓ ( Wx, y ) .
Included in these assumptions are the popular
NLL loss − log p ( y | x, w )
=
− ( Wx ) y + log( (cid:80)
ˆ y e
( Wx ) ˆ y ) that is used in logistic regression and the multi-class hinge loss max ˆ y { ( Wx )
ˆ y − ( Wx ) y +
[ y ̸
= ˆ y ] } that is used in support vector machines.
Theorem 3.4.
Let ℓ ( w, x, y ) ≜ ˆ ℓ
( Wx, y ) be a differentiable loss function over x = ( x 1 , ..., x d ) and k classes y ∈ { 1 , ..., k } , with Lipschitz constant g , i.e., ∥∇ t ˆ ℓ ( t, y )
∥ ≤ g .
Consider a Gaussian prior distribution p ∼ N (0 , σ 2 p ) .
Under the conditions of Lemma 3.3, for any 0 < λ ≤ (cid:112) m
/gσ p σ y and δ ∈ (0 , 1) , with probability at least 1 − δ over the draw of the training set S , we have
[Equation]
Proof.
This bound is derived by applying Lemma 3.3.
We begin by realizing the gradient of ˆ ℓ ( Wx, y ) with respect to x .
Using the chain rule, ∇ x ˆ ℓ ( Wx, y )
=
W ⊤ ∇ Wx ˆ ℓ ( Wx, y ) .
Hence, we obtain for the gradient norm ∥∇
[Equation]
Plugging this result into Lemma 3.3, we obtain the following bound:
[Equation]
Since M ( α ) ≜ E D (cid:2)
e − αℓ ( w,x,y ) (cid:3) , the ratio in the integral equals one and the integral (cid:
) λm
dα =
λm .
Finally, whenever λgσ p σ y ≤ (cid:112) m/ 4 we follow the Gaussian integral and derive the bound
[Equation]
Finally, we obtain Equation (22) by plugging the above into Theorem 2.1.
We provide a detailed proof in Appendix B.
Notice, in the linear case, the gradients of the model are a function of W .
Thus they are not uniformly bounded for every W .
Since we assume the loss to be Lipschitz, we can bound C ( λ, p ) , and by further assuming bounded λ , and Gaussian W , we obtain kd · log(4 / 3) .
Theorem 3.4 provides a PAC-Bayesian bound for classification using the NLL loss.
This extends the result of Alquier et al.
[2016] for binary hinge-loss to the multi-class hinge loss (cf.
[Alquier et al., 2016], Sec. 6).
While the above result can be applied to deep nets, obtaining a value for the bound requires to compute the Lipschitz constant, which is an NP-hard problem even for two-layer neural nets [Virmaux and
Scaman, 2018].
Moreover, common Lipschitz constant approximation algorithms tend to overestimate the constant and are exponential in the network’s depth [Virmaux and Scaman, 2018, Combettes and
Pesquet, 2019].
3.2 Non-linear models
Sadly, the bound proposed in the previous sub-section cannot be applied to non-linear loss functions since their gradient-norm is not bounded, as evident by the exploding gradients property in deep nets.
To mitigate this, we suggest utilizing on-average bounded losses and gradients.
Formally, one can show the following:
Theorem 3.5.
Consider smooth loss functions that are on-average bounded, i.e., for every w the following holds: E D ℓ ( w, x, y ) ≤ b and E D (cid:2) ∥∇ x ℓ ( w, x, y ) ∥ 2 (cid:3) ≤ g .
Under the conditions of
Lemma 3.3 for any 0 < λ ≤ m and δ ∈ (0 , 1) , with probability at least 1 − δ over the draw of the training set S , we obtain
[Equation]
Proof.
This bound is derived by applying Lemma 3.3 and bounding (cid:82) 1
e − αℓ ( w,x,y ) M
( α ) dα ≤ e b .
We derive this bound in three steps: First, from ℓ ( w, x, y ) ≥
we obtain e − αℓ
( x,x,y ) ≤ 1 .
Then, we lower bound M ( α ) ≥ M (1) for any 0 ≤ α ≤ λ/m .
Also, since ℓ ( w, x, y ) ≥
the function e − αℓ ( w,x,y ) is monotone in α within the unit interval and M ( α ) ≥ M (1) for any α ≤ 1 .
Lastly, the assumption
[Figure]
E D [ − ℓ ( w, x, y )] ≥ − b and the monotonicity of the exponential function result in the lower bound e E D [ − ℓ ( w,x,y )] ≥ e − b .
From convexity of the exponential function we have E D e − ℓ ( w,x,y ) ≥ e E D [ − ℓ ( w,x,y )] .
Combining these bounds and replacing E D (cid:2)
∥ σ y ∇ ℓ
[Equation]
x
[Equation]
y
[Equation]
2 (cid:3) with gσ 2 y , we obtain the upper bound
[Equation]
Finally, we obtain Equation (24) by plugging Equation (25) into Theorem 2.1.
We provide a detailed proof in Appendix B.
In contrast to prior work that assumes the loss function to be uniformly bounded, i.e., ℓ ( w, x, y ) ≤ B for every w, x, y
[Alquier et al., 2016], in Theorem 3.5, we assume an on average bounded loss and an on average bounded gradient norm, i.e., E D ℓ ( w, x, y ) ≤ b and E D (cid:2) ∥∇ ℓ ( w, x, y ) ∥ 2 (cid:3) ≤ g for all w .
The above bound corresponds to an open problem raised by Bartlett et al.
[2017a], wondering about the existence of generalization bounds that assume on-average bounds on the loss and/or the gradient norm.
Nevertheless, this bound is more theoretical than practical, as it enforces global on-average bounds b ≥ E D ℓ
[Equation]
w, x, y ) and g ≥
[Equation]
For a more practical perspective, in the following, we consider the on-average loss and gradient norm to be functions of the model parameters.
Theorem 3.6.
Consider smooth loss functions and the conditions of Lemma 3.3.
For any 0 < λ ≤ m we obtain
[Equation]
The proof on this Theorem follows the proof of Theorem 3.5.
However, in this case, we cannot remove the expectation over the prior distribution since we assume the bound over the loss and gradients to depend on w .
The above derivation upper bounds the complexity term C ( λ, p ) by the expected gradient-norm of the loss function, i.e., the flow of its gradients through the model’s architecture.
We show empirically that the rate of the bound λ can be as high as m , dependent on the gradient-norm.
This is a favorable property since the convergence of the bound scales as 1 /λ .
Therefore, one would like to avoid exploding gradient-norms, which effectively harm the true risk bound.
While one may achieve a fast rate bound by forcing the gradient-norm to vanish rapidly, practical experience shows that vanishing gradients prevent the deep net from fitting the model to the training data when minimizing the empirical risk.
In our experimental evaluation, we demonstrate the influence of the expected gradient-norm on the bound of the true risk.
Experiments
In this section, we evaluate our PAC-Bayesian bounds experimentally, both for linear and non-linear models.
We begin by verifying our assumptions, comparing the proposed bound to prior work, and
[Figure]
estimating its predictive generalization capabilities.
Next, we study the behavior of the complexity term C ( λ, p ) for different architectures, both for linear models and deep nets.
We conclude the section with an evaluation of the effectiveness of the proposed bound at predicting generalization performance and analyzing its different components during optimization.
All reported results were averaged over three runs using different seeds.
Complete experimental setup can be found in Appendix C.
Verifying assumptions: In Lemma 3.3 we assume that the loss per label is balanced.
To verify that this assumption holds, we use ten different architectures (ResNet18, PreActResnet18, GoogLeNet,
VGG11, VGG13, VGG16, VGG19, DenseNet121, MobileNet, EfficientNetB0) on CIFAR10 and
CIFAR100 [Krizhevsky, 2009, Simonyan and Zisserman, 2014, Szegedy et al., 2015, He et al., 2016,
Huang et al., 2017, Howard et al., 2017, Tan and Le, 2019].
The maximum standard deviation across the labels is 0.022, while the mean value is 4.605.
Hence, it is evident that this assumption holds in practice.
In Theorem 3.6 we assume that the loss is unbounded but it is on-average bounded by a function depending on w , i.e., E D ℓ ( w, x, y ) ≤ b ( w ) .
The results to verify this for MNIST [LeCun et al., 1998] and CIFAR10 are provided in Fig.
1a and Fig.
1b.
We observed that until σ 2 p = 0 . 1 , the loss is on-average bounded by ∼ 2 .
Moreover, for σ 2 p ≤ 0 . 01 , the on-average loss bound is aboutand its effect on the complexity term C ( λ, p ) is minor.
This validates empirically our assumptions that the on-average bounds E D ℓ ( w, x, y ) are small although max w,x,y ℓ ( w, x, y ) is much larger (see
Tab. 1 for its impact on the generalization).
Complexity of neural nets: We now turn to estimate our bounds over C ( λ, p ) , both for linear and non-linear models.
Recall that λ determines the rate of the generalization bound, and one would like to set 0 < λ ≤ m as large as possible while C ( p, λ ) to be as small as possible.
As the bound on
C ( λ, p ) is controlled by the expected gradient-norm E D (cid:2) ∥∇ x ℓ ( w, x, y ) ∥ 2 (cid:3) , we present in Fig.
1c and Fig.
1d the expected squared gradient-norm as a function of different variance levels for the prior distribution p over the model parameters, again using both MNIST and CIFAR10 data.
We note that the linear model has the largest expected gradient-norm.
Further, note that the deeper the network, the smaller its gradient-norm.
As a result, deeper nets can use larger values of λ for the generalization bound.
Thus, we present our complexity bound as a function of λ in Fig.
2a and
Fig.
2b.
Note, the complexity term C ( λ, p ) reaches its minimum around λ
= m .
Fig.
2c studies the effect of a network’s depth on the complexity term.
We observe that, as expected based on the earlier plots, deeper networks have a lower complexity term.
We attribute this phenomenon to vanishing gradients which create a contractivity property that stabilizes the loss function, i.e., reduces its variability.
However, this comes at the expense of expressivity of the deep net, since deep nets with vanishing gradients cannot fit the training data in the learning phase.
Comparison to prior PAC-Bayesian bounds: We compare the proposed bound to Alquier et al.
[2016] (bounded version), Germain et al.
[2016], and Dziugaite and Roy [2017] on both MNIST and
CIFAR.
Unlike our bound which assumes the expected loss to be bounded, both Alquier et al.
[2016] and Germain et al.
[2016] assume the loss to be uniformly bounded by a constant B .
Hence, for a fair comparison we estimated B to be the maximum and average loss over the training set and use it to bound the maximum and expected loss respectively.
Results are summarized in Tab. 1.
Notice, our bound produces tighter generalization bounds compared to the baselines methods for both datasets.
[Table]
Related work
Generalization bounds for deep nets were explored in various settings.
VC-theory provides both upper bounds and lower bounds to the network’s VC-dimension, which are linear in the number of network parameters [Bartlett et al., 2017b, 2019].
While VC theory asserts that such a model should overfit as it can learn any random labeling (e.g., [Zhang et al., 2016]), surprisingly, deep nets generally do not overfit.
Rademacher complexity allows to apply data-dependent bounds to deep nets [Bartlett and Mendelson, 2002, Neyshabur et al., 2015, Bartlett et al., 2017a, Golowich et al., 2017, Neyshabur et al., 2018].
These bounds rely on the loss and the Lipschitz constant of the network and consequently depend on a product of norms of weight matrices, which scales exponentially in the network depth.
Wei and
Ma [2019] developed a bound over the gradient-norm of training examples.
In contrast, our bound depends on average quantities of the gradient-norm.
Thus we answer an open question of Bartlett et al.
[2017a] about the existence of bounds that depend on average loss and average gradient-norm, albeit in a PAC-Bayesian setting.
PAC-Bayesian bounds that use Rademacher complexity have also been studied [Kakade et al., 2009, Yang et al., 2019].
Stability bounds may be applied to unbounded loss functions and in particular to the negative loglikelihood (NLL) loss [Bousquet and Elisseeff, 2002, Rakhlin et al., 2005, Shalev-Shwartz et al., 2009,
Hardt et al., 2015, Zhang et al., 2016].
However, stability bounds for convex loss functions, e.g., for logistic regression, do not apply to deep nets since they require the NLL loss to be a convex function of the parameters.
Alternatively, Hardt et al.
[2015] and Kuzborskij and Lampert [2017] estimate the stability of stochastic gradient descent dynamics, which strongly relies on early stopping.
This approach results in weaker bounds for the non-convex setting.
Stability PAC-Bayesian bounds for bounded and Lipschitz loss functions were developed by London [2017].
Li et al.
[2019] utilized the log-Sobolev inequality to bound the KL divergence under the PAC-Bayesian setting while assuming a bounded loss function.
In contrast, we assume an unbounded loss function and a Gaussian prior distribution over the model weights.
The latter allows us to compute the KL divergence.
Holland [2019] studies PAC-Bayesian learning guarantees for heavy-tailed losses.
Haddouche et al.
[2021] relax the boundness assumption by modifying the range of the loss to depend on each predictor.
Differently, through the log-Sobolev inequality, our bound relies on the expected gradient-norm of the loss function, which may be unbounded as well.
PAC-Bayesian bounds were recently applied to deep nets [McAllester, 2013, Dziugaite and Roy, 2017, Neyshabur et al., 2017, Pérez-Ortiz et al., 2021].
In contrast to our work, those related works all consider bounded loss functions.
An excellent survey on PAC-Bayesian bounds was provided by Germain et al.
Alquier et al.
[2016] introduced PAC-Bayesian bounds for linear classifiers trained with a hinge-loss by explicitly bounding its moment generating function.
Dziugaite et al.
[2021] investigate the role of data in learning a PAC-Bayesian prior.
Closely, Foong et al.
[2021] study the tightness of PAC-Bayesian bounds for small datasets.
Alquier et al.
[2012] provide an analysis for PAC-Bayesian bounds with Lipschitz functions.
Our work differs as we derive PAC-Bayesian bounds for non-Lipschitz functions.
Work by Germain et al.
[2016] is closer to our setting and considers PAC-Bayesian bounds for linear models and quadratic loss functions.
In contrast, our work considers the multi-class setting and non-linear models.
PAC-Bayesian bounds for the NLL loss in the online setting were put forward too
[Takimoto and Warmuth, 2000, Banerjee, 2006, Bartlett et al., 2013, Grünwald and Mehta, 2017].
The online setting does not consider the whole sample space and therefore is simpler to analyze in the Bayesian setting.
Recently, a PAC-Bayesian generalization bound for meta-learning was proposed [Rothfuss et al., 2020, Amit and Meir, 2018, Farid and
Majumdar, 2021, Ding et al., 2021].
In contrast to our work, Amit and Meir [2018] and Farid and
Majumdar [2021] focus on bounded loss functions while Rothfuss et al.
[2020] assume the loss function to be sub-gamma.
PAC-Bayesian bounds for the NLL loss function are intimately related to learning Bayesian inference Germain et al.
Recently, many works applied various posteriors in Bayesian deep nets.
Gal and Ghahramani [2015], Gal [2016] introduce a Bayesian inference approximation using Monte
Carlo (MC) dropout, which approximates a Gaussian posterior using Bernoulli dropout.
Srivastava et al.
[2014], Kingma et al.
[2015] introduced Gaussian dropout, which effectively created a Gaussian posterior that couples the mean and the variance of the learned parameters and explored the relevant log-uniform priors.
Blundell et al.
[2015], Louizos and Welling [2016] suggest taking a complete
Bayesian perspective and learning the mean and the variance of each parameter separately.
Stochastic gradient Langevin dynamics (SGLD) bounds [Pensia et al., 2018, Mou et al., 2018, Negrea et al., 2019, Haghifam et al., 2020]
suggest bounding the generalization gap by investigating the dynamics of the parameters through Langevin dynamics and Fokker-Planck equations.
Broadly, the Fokker-Planck equation defines a Markov process, and these works measure the KL-divergence between a prior and posterior distribution of this process in the parameter space using gradient norms of the parameters.
Differently, our work utilizes the Herbst theorem to measure the KL-divergence (in the form of functional entropy) between prior and posterior in the input space.
This distinction leads to a computation of the gradients with respect to different quantities, i.e., parameters vs. input data
Discussion and future work
In this study, we present new PAC-Bayesian generalization bounds that depend on the gradient-norm of the loss function.
We explore their properties in our experimental validation in various deep learning settings, reinforcing the importance of gradients in contemporary deep nets.
Our framework allows for new PAC-Bayesian bounds that assume on-average bounded loss and an on-average bounded gradient norm assumptions.
These findings answer an open problem raised by Bartlett et al.
[2017a] under the PAC-Bayesian setting.
Moreover, we extend the current generalization bounds proposed for the hinge-loss to any linear model with a Lipschitz loss function.
Finally, we empirically show that our bounds produce tighter generalization performance than the baseline methods.
These results are another step towards bridging the gap between theory and practice while having more realistic assumptions for modern deep nets.
Our framework can be extended in different directions.
Log-Sobolev inequalities (LSIs) are intimately related to hypercontractivity.
Our generalization bounds also imply that deep nets hyper-contracts their input and therefore generalize.
While the relation between generalization and hypercontractivity is under explored, with the small-ball theorem as a notable exception [Lecué and Mendelson, 2017,
Mendelson, 2014], we believe it is a fundamental concept in contemporary deep nets that require further research.
Additionally, our framework focuses on bounding the complexity term.
It is interesting future work to also consider the KL term.
We assume that the labels are balanced with respect to the loss, or equivalently, the influence of all labels is equal.
The theory of influence of discrete functions is well-developed and very relevant to LSIs, hypercontractivity, and measure concentration, and further investigation on its relation to generalization is desired [O’Donnell, 2014].
A Log-Sobolev inequalities
For completeness we derive Theorem 2.3.
We begin with proving the LSI for a Rademacher random variable, i
[Equation]
e., a discrete random variable that takes the values {− 1 , +1 } with equal probability.
Theorem A.1 (LSI for Rademacher distribution, Gross [1975]) .
let Z be Rademacher random variable that take values in {− 1 , +1 } with probability of 12 .
Set the discrete gradient of a function f : {− 1 , +1 } →
R to be
[Equation]
Then the following LSI holds:
[Equation]
Proof.
The quantities that constitute the LSI for Rademacher random variables take the form:
[Equation]
Setting a = f (1) and b = f ( − 1) the log-Sobolev inequality implies that:
[Equation]
In the following, we assume without loss of generality that a ≤ b .
Let’s also define
[Equation]
Given this notation, proving the LSI for Rademancher random variables amounts to proving that g ( b ) ≤ 0 .
To prove the inequality g ( b ) ≤ 0 we follow these steps:
g ′ ( a ) = 0 , 2.
for any a < r ≤ b holds g ′ ( r ) ≤ 0 .
First, we have:
[Equation]
One can easily verify that g ′ ( a ) = 0 .
To prove that g ′ ( r ) ≤ 0 we show that g ′′ ( r ) ≤ 0 for any a < r ≤ b , thus guaranteeing that g ′ ( r ) is non-positive.
[Equation]
Using the fact that 1 + log α − α ≤ 0 for any α we are able to verify that g ′′ ( r ) ≤ 0 .
The LSI for Rademacher complexity extends to Gaussian distributions using the central limit theorem, which asserts that an average of i.i.d.
Rademacher random variables converges to Gaussian distribution.
This extension relies on tensorization of the functional entropy, which effortlessly scales to highdimensional settings.
Theorem A.2 (Tensorization of Rademacher random variables, Gross [1975], Boucheron et al.
Let Z 1 , ..., Z d be i.i.d.
Rademacher random variables and let
Ent i [
[Equation]
Tensorization of independent random variables Z 1 , ..., Z d implies the bound:
[Equation]
Proof.
We denote the d -th tuple by z = ( z 1 , ..., z d ) and set z [ d ] \ i =
( z 1 , ..., z i − 1 , z i +1 , ..., z d ) .
Following the definitions:
[Equation]
We assume without loss of generality 2 that (cid:80) z
2 − d f ( z ) = 1 and set q ( z ) ≜ 2 − d f ( z ) .
Since f ( z ) ≥ 0 , then q ( z ) is a distribution and its entropy is H ( q )
= − (cid:80)
z q ( z ) log q ( z ) .
Then:
[Equation] [Equation]
We denote the marginal distribution by q ( z [ d ] \ i ) =
(cid:80) z i q ( z [ d ] \ i , z i ) .
Then
[Equation] [Equation]
E Z 1 ,...
[Equation]
Z
[Equation]
(44) where q [ d ] \ i is the distribution of q ( z [ d ] \ i ) .
With these equalities in mind,
Ent
[Equation] [Equation]
The right hand side of the above equation is know as Han’s inequality.
We use tensorization and the central limit theorem to extend the LSI of Rademacher random variables to derive the LSI for Gaussian random variables.
Theorem A.3 (Gaussian log-Sobolev inequality (LSI), Gross [1975]) .
Let Z be a Gaussian random variables Z ∼ N (0 , 1) .
Then
[Equation]
Proof.
In the following we denote by Z 1 , ..., Z d Rademacher random variables.
We set
[Equation]
Using tensorization and the LSI for Rademacher random variables:
[Equation]
The theorem follows using the central limit theorem as lim d →∞
[Equation]
di
[Equation]
d = Z , when Z ∼ N (0 , 1) and f ( Z ) =
[Equation]
The above provides a LSI for the standard Gaussian distribution.
While we can use the same technique to prove a LSI for Z ∼ N ( µ, σ 2 ) , we separate these two theorems for clarity and turn to prove the general case as a corollary.
Theorem A.4 (Gaussian log-Sobolev inequality (LSI), Gross [1975]) .
Let Z ∼ N ( µ, σ 2 ) be a
Gaussian random variable with mean µ and variance σ 2 .
Then
[Equation]
Proof.
The proof follows via a simple change of variable.
Let ˆ Z ∼ N (0 , 1) , then µ + σ ˆ Z ∼ N ( µ, σ 2 ) .
We set g ( ˆ Z ) ≜ f ( µ + σZ ) .
We use the LSI for the standard Gaussian: Ent[ e g ( ˆ Z ) ] ≤ 12 E [ g ′ ( ˆ Z ) 2 e g ( ˆ Z ) ] and note that g ′ (ˆ z ) =
σf ′
( z ) .
Theorem 2.3 for multivariate Gaussians holds when applying tensorization for each Gaussian independently.
B Proofs
Theorem B.1.
Let ℓ ( w, x, y ) ≜ ˆ ℓ ( Wx, y ) be a differentiable loss function over x = ( x 1 , ..., x d ) and k classes y ∈ { 1 , ..., k } , with Lipschitz constant g , i.e., ∥∇ t ˆ ℓ ( t, y ) ∥ ≤ g .
Consider a Gaussian prior distribution p ∼ N (0 , σ 2 p ) .
Under the conditions of Lemma 3.3, for any 0 < λ ≤ (cid:112) m
/gσ p σ y and δ ∈ (0 , 1) , with probability at least 1 − δ over the draw of the training set S , we have
[Equation]
Proof.
This bound is derived by applying Lemma 3.3.
We begin by realizing the gradient of ˆ ℓ ( Wx, y ) with respect to x .
Using the chain rule, ∇ x ˆ ℓ ( Wx, y )
=
W ⊤ ∇ Wx ˆ ℓ ( Wx, y ) .
Hence, we obtain for the
[Equation]
Plugging this result into Lemma 3.3, we obtain the following bound:
[Equation]
Since M ( α ) ≜ E D (cid:2)
e − αℓ ( w,x,y ) (cid:3) , the ratio in the integral equals one and the integral (cid:
) λm
dα =
λm .
Combining these results we obtain:
[Equation]
Finally, whenever λgσ p σ y ≤ (cid:112) m/ 16 we follow the Gaussian integral and derive the bound log
[Equation]
Finally, we obtain Equation (53) by plugging the above into Theorem 2.1.
Theorem B.2.
Consider smooth loss functions that are on-average bounded, i.e., for every w the following holds: E D ℓ ( w, x, y ) ≤ b and E D (cid:2) ∥∇ x ℓ ( w, x, y ) ∥ 2 (cid:3) ≤ g .
Under the conditions of
Lemma 3.3 for any 0 < λ ≤ m and δ ∈ (0 , 1) , with probability at least 1 − δ over the draw of the training set S , we obtain
[Equation]
Proof.
This bound is derived by applying Lemma 3.3 and bounding (cid:82) 1
e − αℓ ( w,x,y ) M
( α ) dα ≤ e b .
We derive this bound in three steps: First, from ℓ ( w, x, y ) ≥
we obtain e − αℓ
( x,x,y ) ≤ 1 .
[Equation]
Then, we lower bound M ( α ) ≥ M (1) for any 0 ≤ α ≤ λ/m : we note that 0 < λ ≤ m , therefore we consider 0 ≤ α ≤ 1 .
Also, since ℓ ( w, x, y ) ≥
the function e − αℓ ( w,x,y ) is monotone in α within the unit interval, i.e., for 0 ≤ α 1 ≤ α 2 ≤ 1 there holds
[Equation]
and consequently M ( α ) ≥ M (1) for any α ≤ 1 .
Lastly, the assumption E D [ − ℓ ( w, x, y )] ≥ − b and the monotonicity of the exponential function result in the lower bound
[Equation] [Equation]
y )]
[Equation] [Equation]
From convexity of the exponential function, M (1) =
E D e − ℓ ( w,x,y ) ≥
e E D [ − ℓ ( w,x,y )] , and the lower bound M (1) ≥ e − b follows.
Combining these bounds we derive the upper bound
[Equation]
and the result follows.
Next, by replacing E D (cid:2) ∥ σ y ∇ ℓ ( w, x, y ) ∥ 2 (cid:3) with gσ 2 y , we obtain
[Equation]
Finally, we obtain Equation (57) by plugging Equation (62) into Theorem 2.1.
C Additional results C.1 Connection between generalization and bound
To further demonstrate the use of the bound, we performed a new experiment.
We study the effect of important components in ResNet using our bound.
For this, we train four variations of the ResNet model: 1) a standard model (ResNet); 2) a model without skip connections (ResNetNoSkip); 3) a model without batch normalization layers (ResNetNoBN); and 4) a model without both skip connections and batch normalization layers (ResNetNoSkipNoBN).
We optimize all models on the
CIFAR10 data: We observe that ResNet and ResNetNoSkip achieve comparable performance in all
[Table]
metrics.
Additionally, removing the batch normalization layers and including the skip connections achieves comparable performance to ResNet and ResNetNoSkip.
Similar to Zhang et al.
[2019], this finding suggests that even without batch normalization, models can converge using precise initialization.
Interestingly, by removing batch normalization and skip connection layers, the model gets to a rate of λ = m and achieves a good generalization bound.
However, this comes at the expense of poor model fitting to the train set due to gradient vanishing.
These results are consistent with prior findings in which batch normalization improves optimization [Santurkar et al., 2018].
To conclude, we were able to obtain a tight upper bound of the generalization gap with our proposed bound.
However, it is important to note that when using any generalization bound, one should care about the training loss as well as the complexity term.
C.2 Gradient statistics
We further study the gradient norm statistics, we report the max and mean values of the gradient norm (not squared) using the same networks described in Section.
4 in the main paper: # LAYERS
[Table]
We observe that the maximum value of the gradient norm is not higher than twice the mean value.
D Experimental setup
For all MNIST experiments we use MLPs of depth d ∈ { 1 , . . .
For CIFAR10 experiments we use CNNs of depth d ∈ { 2 , . . .
For the CNN models, d denotes the number of convolutional layers.
We also added a max-pooling layer after each convolutional layer.
We included two additional fully connected layers in all CNN models to fix the target output dimension.
In all models, we use the ReLU activation function.
We optimize the negative log-likelihood (NLL) loss function using stochastic gradient descent (SGD) with a learning rate of 0.01 and a momentum value of 0.9 in all settings for 50 epochs.
We use mini-batches of size 128 and did not use any learning rate scheduler.
To span the possible weights, we sampled from a normal prior distribution with different variances.
For a fair comparison, we set the layers’ width to reach roughly the same number of parameters in each model (except for the linear case).
All reported results use δ = 0 .