Generalization Bounds for Gradient Methods via Discrete and Continuous Prior
Abstract
Proving algorithm-dependent generalization error bounds for gradient-type optimization methods has attracted signiﬁcant attention recently in learning theory.
However, most existing trajectory-based analyses require either restrictive assumptions on the learning rate (e.g., fast decreasing learning rate), or continuous injected noise (such as the Gaussian noise in Langevin dynamics).
In this paper, we introduce a new discrete data-dependent prior to the PAC-Bayesian framework, and prove high probability generalization bounds of order O ( 1 n · P
[-Tt =1 ( γ t /ε t ) 
-][Equation]
2 ) for ﬂoored GD and SGD (i.e. ﬁnite precision versions of GD and SGD with precision level " t ) where, where n is the number of training samples,  t is the learning rate at step t , g t is roughly the difference between the average gradient over all samples and that over only prior samples.
k g t k is upper bounded by (typically much smaller) than the gradient norm kr f ( W t ) k .
We remark that our bounds hold for nonconvex and nonsmooth loss functions.
Moreover, our theoretical results provide numerically favorable upper bounds of testing errors ( 0 . 026 on MNIST and 0 .
198 on CIFAR10).
Furthermore, we study the generalization bounds for gradient Langevin Dynamics (GLD).
Using the same framework with a carefully constructed continuous prior, we show a new high probability generalization bound of order O ( 1 n + L 2 n P Tt {+=
+} (  t /  t ) 2 ) for GLD.
The new 1 /n 2 rate is obtained using the concentration of the difference between the gradient of training samples and that of the prior.
Introduction
Bounding generalization error of learning algorithms is one of the most important problems in machine learning theory.
Formally, for a supervised learning problem, the generalization error is deﬁned as the testing error (or population error) minus the training error (or empirical error).
In particular, we denote R ( w, ( x, y )) :
=[-
-]
[ h w ( x ) 6 = y ] as the error of a single data point ( x, y ) , where h w ( x ) is the output of a model with parameter w 2 R d .
Suppose S is the set of training data, each i.i.d.
sampled from the population distribution D , and we use R ( w, S ) :
= 1 | S | P z 2 S R ( w, z ) and
R ( w, D ) :=
E z ⇠ D [ R ( w, z )] to denote the training error and the testing error, respectively.
The generalization error of w is formally deﬁned as err gen ( w ) =
R ( w, D )  R ( w, S ) .
Proving tighter generalization bounds for general nonconvex learning and particularly deep learning has attracted signiﬁcant attention recently.
While the classical learning theory (uniform convergence theory) which bounds the generalization error by various complexity measures (e.g., the 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
VC-dimension and Rademacher complexity) of the hypothesis class has been successful in several classical convex learning models, however, they become vacuous and hence fail to explain the success of modern nonconvex over-parametrized neural networks (i.e., the number of parameters signiﬁcantly exceeds the number of training data) (see e.g., Zhang et al.
[2017], Nagarajan and Kolter [2019]).
Recently, learning theorists have tried to understand and explain generalization of deep learning from several other perspectives, such as margin theory [Bartlett et al., 2017, Wei et al., 2019], algorithmic stability [Hardt et al., 2016, Mou et al., 2018, Li et al., 2020, Bousquet et al., 2020],
PAC-bayeisan [London, 2017, Bartlett et al., 2017, Neyshabur et al., 2018, Zhou et al., 2019, Yang et al., 2019], neural tangent kernel
[Jacot et al., 2018, Du et al., 2019, Arora et al., 2019, Cao and Gu, 2019], information theory [Pensia et al., 2018, Negrea et al., 2019], model compression [Arora et al., 2018, Zhou et al., 2019], differential privacy [Oneto et al., 2017, Wu et al., 2021] and so on.
In this paper, we aim to obtain tighter generalization error bounds that depend on both the training data and the optimization algorithms (a.k.a.
gradient-type methods) for general nonconvex learning problems.
In particular, we prove algorithm-dependent generalization bounds for several gradientbased optimization algorithms such as certain variants of gradient descent (GD), stochastic gradient descent (SGD) and stochastic gradient Langevin dynamics (SGLD).
Our proofs are based on the classic Catoni’s PAC-Bayesian framework [Catoni, 2007] and also have a ﬂavor of algorithmic stability [Bousquet and Elisseeff, 2002].
Several prior works have obtained generalization bounds for SGD and SGLD by analyzing trajectory through either the PAC-Bayesian or the algorithmic stability framework (or closely related information theoretic arguments).
However, most existing results based on analyzing the optimization trajectories require either restrictive assumptions on the learning rates, or continuous noise (such as the Gaussian noise in Langevin dynamics) in order to bound the stability or the KL-divergence.
In this paper, we resolve the above restrictions by combining the PAC-Bayesian framework with a few simple (yet effective) ideas, so that we can obtain new high probability and non-vacuous generalization bounds for several gradient-based optimization methods with either discrete or continuous noises (in particular certain variants of GD and SGD, either being deterministic or with discrete noise, which cannot be handled by existing techniques).
1.1 Prior work
We ﬁrst brieﬂy mention some recent work on bounding the generalization error of gradient-based methods.
Hardt et al.
ﬁrst studied the uniform stability (hence the generalization) of stochastic gradient descent (SGD) for both convex and non-convex functions.
Their results for non-convex functions requires that the learning rate ⌘ t scales with 1 /t .
Their work motivates a long line of subsequent work on generalization error bounds of gradient-based optimization methods: Kuzborskij and Lampert [2018], London [2016], Chaudhari et al.
[2019], Raginsky et al.
[2017], Mou et al. [2018], Chen et al.
[2018], Li et al.
[2020], Negrea et al.
[2019], Wang et al.
Recently, Simsekli et al. [2020], Hodgkinson et al.
[2022] obtained generalization bound of SGD through the perspective of heavy-tailed behaviors and using the notion of Hausdorff dimension d H which depends on both the algorithm and data.
PAC-Bayesian bounds.
The PAC-Bayesian framework [McAllester, 1999] is a powerful method for proving high probability generalization bound [Bartlett et al., 2017, Zhou et al., 2019, Mou et al., 2018].
Roughly speaking, it bounds the generalization error by the KL divergence KL  Q     P  , where Q is the distribution of the learned output and P is a prior distribution which is typically independent of dataset S .
In this framework, bounding KL  Q     P  is the most crucial part for obtaining tighter PAC-Bayesian bounds.
In order to bound the KL divergence, both the prior P and posterior Q are typically chosen to be continuous distributions (mostly Gaussians so that KL can be computed in closed form).
Hence, most prior work either considered gradient methods with continuous noise (such as Gradient Langevin Dynamics)
(e.g., [Mou et al., 2018, Li et al., 2020,
Negrea et al., 2019]), or injected a Gaussian noise to the ﬁnal parameter at the end (e.g., [Neyshabur et al., 2018, Zhou et al., 2019])
(so Q is a Gaussian distribution).
We also note that designing effective prior P can be also very important.
For example, Lever et al.
[2013] proposed to use the population distribution to compute the prior.
In fact, the prior can even partially depend on the training data [Parrado-Hernández et al., 2012, Negrea et al., 2019], and our Theorem 4.1 is partially inspired by this idea.
1.2 Our contributions
First, we provide high probability generalization bounds for discrete gradient methods.
In particular, we study the generalization of Floored Gradient Descent (FGD), which is a variant of GD, and Floored
Stochastic Gradient Descent (FSGD), a variant of SGD.
We obtain our bound by an interesting construction of discrete priors.
Secondly, we consider well studied gradient methods with continuous noise, (stochastic) gradient Langvin dynamics (GLD and SGLD).
We show sharper generalization bounds by carefully bounding the concentration of the sample gradients.
Now, we summarize our results.
FGD and FSGD.
We ﬁrst study an interesting variant of GD, called Floored GD (FGD) (Algorithm 1).
The update rule of FGD is deﬁned as follows:
[Equation]
where S J is the subset of training dataset S with size m indexed by subset J ⇢ [ n ] ( J is chosen before training), r f ( W t  1 , Z ) :
= 1 | Z |
P z 2 Z r f ( W t  1 , z ) is the average gradient over the dataset Z ,  t is the learning rate, " t is the precision level, and g t := r f ( W t  1 , S )  r f ( W t  1 , S J ) is the gradient difference.
The ﬂooring operation is deﬁned by ﬂoor( x ) :
= sign( x ) b | x | c for any real number x .
FGD can viewed as GD with given precision limit " t .
We can see if we ignore the ﬂoor operation or let " t approaches 0, FGD reduces to GD (see also Appendix A).
We also study a ﬁnite precision variant of SGD, called Floored SGD (FSGD) (see Section 5 for its formal deﬁnition).
Empirically, the optimization and generalization capabilities of FGD and FSGD are very close to those of GD and SGD (see Figure 5 and 6 in Appendix H).
By constructing a discrete data-dependent prior and incorporate it into Catoni’s PAC-Bayesian framework, we prove that the following bound (Theorem 5.2) holds for FGD with high probability:
[Equation]
where d is the dimension of parameter space and c 0 can be chosen to be a small constant.
The bound for FSGD is very similar (see Theorem 5.3).
Now we make a few remarks about our results.
Our result holds for nonconvex and nonsmooth learning problems (replacing the gradients with subgradients for nonsmooth cases).
Moreover, there is no additional requirement on the learning rate  t . 2.
The gradient difference g t is typical much smaller than the worst case gradient norm.
It usually decreases when m = | J | grows (see Figure 1c in Section 7).
We obtain non-vacuous generalization bounds on commonly used datasets.
Speciﬁcally, our theoretical test error upper bounds on MNIST and CIFAR10 are 0 .
026 and 0 . 198 , respectively (see Section 7).
Both of them are tighter than the best-known MNIST bound ( 11% ) and CIFAR10 bound ( 23% ) reported in Dziugaite et al.
See Table 1 in Appendix B for more comparisons.
In order to bound the KL between P and the deterministic process of FGD, we construct the prior P from a discrete random processes..
We hope it may inspire future research on handling deterministic optimization algorithms or discrete noise.
Why study FGD/FSGD?
We would like to remark that we study FGD/FSGD, not because
FGD/FSGD have better performances than GD/SGD or other advantages. Indeed, their performances are almost the same as those of GD/SGD (see Appendix H). We use them as important stepping stones to study generalization bounds for GD and SGD.
Note that most existing trajectorybased generalization bounds require either fast decreasing learning rate, or continuous injected noise, such as the Gaussian noise in Langevin dynamics, for general non-convex loss functions.
Handling deterministic algorithms (such as GD) or discrete noises (such as SGD) is challenging and beyond the reach of existing techniques.
In fact, understanding such discrete noises and their effects on generalization has been an important research topic (see e.g., Li et al.
[2020], Zhu et al.
Ziyin et al.
In particular, Zhu et al. [2019] show that it is insufﬁcient to approximate SGD’s discrete noise by isotropic Gaussian noise.
Moreover, proving nontrivial generalization bounds for SGD-like algorithms with discrete noise has also been proposed as an open research direction in Li et al.
GLD and SGLD.
We provide a new generalization bound for Gradient Langevin Dynamics (GLD).
The update rule of GLD is deﬁned as follows.
[Equation]
In this paper, we show that the following generalization bound (Theorem 6.2) holds with high probability over the randomness of S ⇠ D n and random subset J ⇢ [ n ] ( | J | = m ):
[Equation]
where L ( W t  1 ) :
= max z 2 S k f ( W t  1 , z ) k is the longest gradient norm of any training sample in S at step t and m is the size of J .
Since W T is independent of the index set J , the ﬁrst term R ( W T , S [ n ] \ J ) is upper bounded by R ( W T , S )
+ O ( 1 p n  m ) with high probability, using standard
Hoeffding’s inequality.
By setting m = n/ 2 , our generalization bound has an O ( 1 p n + 1 n + Tn 2 ) rate.
The new 1 /n 2 rate is obtained using the concentration of the difference between the gradient of training samples and that of the prior (See Lemma 6.1).
We also prove a high probability generalization bound for Stochastic Gradient Langevin Dynamics (SGLD) (see Theorem 6.3):
[Equation]
We compare our bounds with other GLD/SGLD bounds obtained in [Mou et al., 2018, Negrea et al., 2019, Li et al., 2020] and the details can be found in Appendix B. CLD.
Using the PAC-Bayesian framework, we obtain a new generalization bound for Continuous
Langevin Dynamics (CLD), deﬁned by the stochastic differential equation d W t = r F ( W t , S ) d t + p 2   1 d B t . The main term of the generalization bound scales as O (1 /n 2 ) (by choosing m = n/ 2) and does not grow to inﬁnity as the training time T increases. See Theorem G.6 for the details.
Other Related Work
Stochastic Langevin Dynamics Stochastic Langevin dynamics is a popular sampling and optimization method in machine learning [Welling and Teh, 2011].
Zhang et al.
[2017], Chen et al. [2020] show a polynomial hitting time (hitting a stationary point) of SGLD in general non-convex setting.
Raginsky et al.
[2017] study the generalization and excess risk of SGLD in nonconvex settings and their bound depends inversely polynomially on a certain spectral gap parameter, which may be exponential small in the dimension.
Continuous Langevin dynamics (SDE) with various noise structure has also been used extensively as approximations of SGD in literature (see e.g., [Li et al., 2017, 2021]).
However, in terms of generalization, isotropic Gaussian noise is not a good approximation of the discrete noise in SGD (Zhu et al.
Nonvacuous PAC-Bayesian Generalization Bounds.
Dziugaite and Roy [2017] ﬁrst present a nonvacuous PAC-Bayesian generalization bound on MNIST (0.161 for a 1-layer MLP, see column Tof Table 1 in their paper).
They use a very different training algorithm that explicitly optimizes the PAC-Bayesian bound and the output distribution is a multivariate normal distribution.
To computing the closed form of KL, they choose a zero-mean Gaussian distribution as the prior distribution.
Zhou et al.
[2019] obtain the ﬁrst non-vacuous generalization bound for ImageNet via a different method.
Their method does not require any continuous noise injected but assumes that the network can be signiﬁcantly compressed (so that the prior distribution is supported over the set of discrete parameters with ﬁnite precision).
To our best knowledge, it is the only work that utilizes a discrete prior for proving generalization bounds of deep neural networks.
Our result for FGD/FSGD has a similar ﬂavor in a high level, that is the optimization method has a ﬁnite precision.
However, our results do not need any assumption on compressibility of the model and can be applied to nonconvex learning problems other than neural networks.
Generalization bounds via Information theory.
Raginsky et al.
[2017] ﬁrst show that the expected generalization error E S ⇠ D n [ R ( W, D )  R ( W, S )] is bounded by p 2 I ( S ; W ) /n , where I ( S ; W ) : = [Equation]
is the mutual information between the data set S and the parameter
W .
This work motivates several subsequent studies [Pensia et al., 2018, Negrea et al., 2019, Bu et al., 2020, Wang et al., 2021].
The main goal in this line of work is to obtain a tight bound on the mutual information I ( S ; W ) .
This is again reduced to bounding the KL divergence and thus typically requires continuous injected noise (e.g., Wang et al.
[2021], Negrea et al.
Preliminaries
Notations.
We assume that the training dataset S =
( z 1 , .., z n ) is sampled from D n , where D is the population distribution over the data domain ⌦ .
The model parameter w is in R d .
The risk function
R : R d ⇥ ⌦ ! [0 , 1] measures the error of a model on a datapoint.
The loss function f : R d ⇥ ⌦ ! R is a proxy of the risk.
The optimization algorithm minimizes the loss function and we assume we can compute the gradient of the loss function.
We note that the loss function may be different from the risk function (e.g., 0/1 risk vs the cross-entropy loss).
The empirical risk is R ( w, S )
= 1 | S | [Equation] ( w, z ) and population risk is R ( w, D ) =
E z ⇠ D [ R ( w, z )] .
Similarly, we can deﬁne the empirical loss f ( w, S ) and population loss f ( w, D ) .
For any J =
( j 1 , .., j m ) , we use S J to denote the sequence ( S j 1 , ..., S j m ) .
The subsequence ( A i , A i +1 , ..., A j ) is denoted by A ji .
We use ( A n 1 , B m 1 ) to denote the merged sequence ( A 1 , A 2 , ..., A n , B 1 , ..., B m ) .
When the elements in sequence J are distinct, we also use J to represent the set consisting of all of its elements.
We may also slightly abuse the notation of a random variable to denote its distribution.
For example, E x ⇠ X [ f ( x )] is a shorthand for E x ⇠ P X [ f ( x )] , and KL  X
Y  means KL
P X  P Y  .
For a random variable W , we deﬁne R ( W, S ) =
E w ⇠ W [ R ( w, S )] and R ( W, D ) =
E w ⇠ W [ R ( w, D )] .
The set { 1 , 2 , ..., n } is denoted by [ n ] .
KL-divergence.
Let P and Q be two probability distributions.
The Kullback–Leibler divergence KL
P  Q  is deﬁned only when P is absolute continuous with respect to Q (i.e., for any x , Q ( x ) = 0 implies P ( x ) = 0 ).
In particular, if P and Q are discrete distributions, then
KL  P
Q  =
P x P ( x ) ln P ( x ) Q ( x ) .
Otherwise, if P and Q are continuous distributions, it is deﬁned as
R P ( x ) ln P ( x ) Q ( x ) d x .
The following Lemma 3.1 is frequently used in this paper and is a well known property of KL divergence (see Cover [1999, Theorem 2.5.3], Li et al.
[2020], Negrea et al.
Lemma 3.1 (Chain Rule of KL) .
We are given two random sequences W =
( W 0 , ..., W T ) and
W 0 =
( W 0 0 , ..., W 0 T ) .
Then, the following equation holds (given all KLs are well deﬁned):
[Equation]
Here W t | W t  1= w denotes the distribution of W t conditioning on W t {+ 
+}
[Equation]
[Equation]
PAC-Bayesian.
In this paper, we use the PAC-Bayesian bound presented in Catoni [2007] which enjoys a tighter O (KL  Q
P  /n ) rate comparing to the traditional O ( q KL  Q  P  /n ) bound, but with a slightly larger constant factor on the empirical error.
We restate their bound as follows.
Lemma 3.2 (Catoni’s Bound) .
(see e.g., Lever et al.
For any prior distribution P independent of the training set S , any  2 (0 , 1) , and any ⌘ > 0 , the following bound holds w.p.  1   over S ⇠ D n :
[Equation]
where C ⌘ = 1 1  e  ⌘ is an absolute constant.
Concentration inequality.
We use the following variant of McDiramid inequality (Lemma 3.3) to prove the concentration of cumulative gradient difference in Section 6.
The proof is deferred to
Appendix C.
Lemma 3.3.
Suppose  : [ n ] m ! R + is order-independent 1 and |  ( J )   ( J 0 ) |  c holds for any adjacent J, J 0 2 [ n ] m satisfying | J \ J 0 |
= m  1 2 .
Let J be m indices sampled uniformly from [ n ] without replacement.
Then Pr J [  ( J )  E J [  ( J )] > ✏ ]  exp(  2 ✏ 2 mc 2 ) . [Equation] [Table]
Data-Dependent PAC-Bayesian Bound
The dominating term in the PAC-Bayesian bound (1) is KL  Q  P  /n , where P is a prior distribution independent of the training dataset S .
Typically, without knowing any information from S , the best possible bound for KL  Q  P  we can hope is at least ⇥ (1) (it should not be a function of n hence should not decrease with n ).
However, if we are allowed to see m data points from S when constructing our prior, we may produce better prediction on posterior Q S .
The following theorem enables us to use data-dependent prior in PAC-Bayesian bound.
The proof is almost the same as
Cantoni’s original proof and we provide a proof for completeness in Appendix D.
Theorem 4.1 (Data-Dependent PAC-Bayesian) .
Suppose J is a random sequence including m indices uniformly sampled from [ n ] without replacement.
For any  2 (0 , 1) and ⌘ > 0 , we have w.p.  1   over S ⇠ D n and J :
[Equation]
where I =
[ n ] \ J is the set of indices not in J , P ( S J ) is the prior distribution only depending on the information of S J ( S J is the subset of S indexed by J ), and
C ⌘ := 1 1  e  ⌘ is a constant.
Remarks.
Note that the above bound holds regardless of whether Q depends on S or not.
Also note that the ﬁrst term in the right hand side is R ( Q, S I ) , not R ( Q, S ) as in the usual generalization bounds.
We remark that for most of our learning algorithms that are independent of J (i.e., changing J does not change the output Q ), by standard Chernoff-Hoeffding inequality, R ( Q, S I ) can be bounded by R ( Q, S )
+ O (1 / p n  m ) with high probability over the randomness of J .
For example, the update rules of GLD, SGLD and CLD are independent of J , hence R ( Q, S I ) can be replaced by R ( Q, S )
+ O (1 / p n  m ) in Theorem 4.1.
However, we point out a subtle point that FGD (Algorithm 1) studied in this paper depends on J .
It may be the case that by knowing J , FGD extracts more information from S J but not much from S I , unintentionally making R ( Q, S I ) a validation error, rather than the training error as it should be.
However, from our experiment (see Figure 5 andin Appendix H, and Figure 2a), we can see that FGD is very close to GD and the S I error R ( W T , S I ) is indeed close to the training error R ( W T , S ) and both are signiﬁcantly smaller than the testing error R ( W T , D ) .
So R ( Q, S I ) can be considered as a genuine training error in our study of FGD.
FGD and FSGD
In this section, we study the generalization error of ﬁnite precision variants of gradient descent and stochastic gradient descent: Floored Gradient Descent (FGD) and Floored Stochastic Gradient
Descent (FSGD).
First we need to deﬁne the “ﬂoor” operation which is used in the deﬁnitions of FGD and FSGD.
Deﬁnition 5.1 (Floor) .
For any vector X 2 R d , let Y = ﬂoor( X ) deﬁned as:
[Equation]
for all i
[Equation]
[ d ] . FGD: The Floored Gradient Descent algorithm is formally deﬁned in Algorithm 1, where (  t ) t and ( " t ) t  0 are the step size and precision sequences, respectively.
For a subset Z ✓ S , we write r f ( W t  1 , Z ) :
= 1 | Z |
P z 2 Z r f ( W t  1 , z ) .
Note that FGD can be viewed as gradient descent with given precision limit " t .
We can see if we ignore the ﬂoor operation or let " t approach 0, FGD reduces to the ordinary GD (see Appendix A).
We also study momentum FGD, in which the 5th line of Algorithm 1 is replaced by
[Equation]
Here ↵ > 0 is a constant.
We remark that both FGD and its momentum version are deterministic algorithms.
The following theorem provides the generalization error bound for both algorithms.
Theorem 5.2.
Suppose J is a random sequence consisting of m indices uniformly sampled from [ n ] without replacement.
Then for any  2 (0 , 1) , both FGD (Algorithm 1) and its momentum version satisfy the following generalization bound w.p.
at least 1   over S ⇠ D n and J :
[Equation]
Input: Training dataset S = ( z 1 , .., z n ) . Index set J .
Result:
Parameter W T 2 R d .
1 Initialize Ww 0 ; 2 for t : 1 ! T do g t r f ( W t  1 , S ) ;
g t r f
( W t  1 , S J ) ; W t
W t  1  g 2  " t · ﬂoor(( g 1  g 2 ) / " t ) ; 6 end where d is the dimension of parameter space, I =
[ n ] \ J is the set of indices not in J , C ⌘ := 1 1  e  ⌘ is a constant, and g t :=
[Equation]
Proof.
We use Theorem 4.1 to prove our theorem for the momentum version.
The ordinary FGD is a special case of the momentum version with ↵ = 0 .
The key is to construct the prior distribution
P ( S J ) such that KL  [Equation]  P
[Equation]
S J )  is tractable.
Let p be any real number in (0 , 1 / 3) .
We ﬁrst deﬁne a stochastic process { W 0 0 , . . . , W 0 T } , by the following update rule ( W 0 0 :
= w 0 ):
[Equation]
where ⇠ t is a discrete random variable such that for all ( a 1 , .., a d ) 2 Z d :
[Equation]
It is easy to verify that the sum of the probabilities ( [Equation] to 1 . Note that W 0 t only depends on S J .
We deﬁne P ( S J ) as the distribution of W 0 T .
Recall that W t
= ( W 0 , ..., W t ) is the parameter sequence of FGD (Algorithm 1).
Applying the chain rule of KL-divergence (Lemma 3.1), we have: KL
[Equation]
The last equation holds because FGD is deterministic.
Let w = W t  1.
The distribution of
W t | W t  1= w (where w = ( w 0 , ..., w t  1 ) ) is a point mass on
[Equation]
Let vector a = ( a 1 , . .
. , a d ) =
ﬂoor(  t " t ( r f ( w t  1 , S )  r f ( w t  1 , S J ))) .
By the deﬁnition of W 0 t , we have KL
[Equation]
Since | i |  i 2 and p 2 (0 , 1 / 3) , we have ln [Equation]
⌘ is at most d ln  1 +
P 1 i =1 p i  .
It can be further bounded by d ln (1 + 3 p ) .
Moreover, it can be bounded by 3 dp as ln(1 + x )  x .
Thus, the above KL-divergence can be bounded by 3 dp + P [Equation]
1 ln(1 /p )
a 2 k . Recall that the k th entry of a is a k :=
b  t " t · ( r k f ( w t  1 , S )  r k f ( w t  1 , S J )) c , which is less than or equal to  t " t · ( r k f ( w t  1 , S )  r k f ( w t  1 , S J )) .
Therefore, we have
[Equation]
Plugging the above inequality into (2), we have
[Equation]
We conclude our proof by plugging it into Theorem 4.1 (setting p = 1 / ( Td ) ).
FSGD: We can use a similar approach to prove a generalization bound for Floored Stochastic Gradient Desent (FSGD).
Formally, FSGD is identical to Algorithm 1 except for the deﬁnitions of gand g 2 replaced with:
[Equation]
where B t ✓ [ n ] is a random batch independent of S, J and W t  1.
Formally, each B t is a set including b indices uniformly sampled from [ n ] without replacement.
The following theorem provides a generalization bound for FSGD.
The proof can be found in Appendix E.
Theorem 5.3.
Suppose J is a random sequence consisting of m indices uniformly sampled from [ n ] without replacement.
Then for any  2 (0 , 1) , " 2 (0 , 1) , FSGD satisﬁes the following generalization bound: w.p. at least 1   over S ⇠ D n and J :
[Equation]
where d is the dimension of parameter space, I =
[ n ] \ J , C ⌘ := 1 1  e  ⌘ is a constant, and g t :
= f ( W t  1 , S B t )  r f ( W t  1 , S J \ B t ) .
Gradient Langevin Dynamics
In this section, we present new generalization bounds for Gradient Langevin Dynamics (GLD) and
Stochastic Gradient Langevin Dynamics (SGLD)
based on Theorem 4.1.
Gradient Langevin Dynamics (GLD):
The GLD algorithm can be viewed as gradient descent plus a Gaussian noise.
Formally, for a given training set S ⇠ D n , the update rule of GLD is deﬁned as follows:
[Equation]
Here the gradient r f ( W t , z ) can be replaced with any gradient-like vector such as a clipped gradient.
The output of GLD is the last step parameter W T or some function of the whole training trajectory
W T 0 (e.g., the average of the sufﬁx 1 K
[Equation]
We still use the data-dependent PAC-Bayesian framework (Theorem 4.1) to prove the generalization bound for GLD.
Unlike FGD (Algorithm 1), GLD is independent of the prior indices J , which enables us to prove the following concentration bound (Lemma 6.1) for the gradient difference.
The proof is based on Lemma 3.3, which is postponed to Appendix F.
Lemma 6.1.
Let S =
( z 1 , ...
z n ) be any ﬁxed training set.
J is a random sequence including m indices uniformly sampled from [ n ] without replacement, and W =
( W 0 , ..., W T ) is any random sequence independent of J .
Then the following bound holds with probability at least 1   over the randomness of J :
[Equation]
where C  = 4 + 2 ln(1 /  )
+ 5 [-. 
 (cid:112)-].p ln(1 /  ) , and L ( w )
= max i 2 [ n ] kr f ( w, z i ) k . [Figure]
Now we are ready to present our main results.
The proofs can be found in Appendix F.
Theorem 6.2.
Suppose J is a random sequence consisting of m indices uniformly sampled from [ n ] without replacement.
Let W T be the output of GLD.
Then for any  2 (0 , 12 ) and ⌘ > 0 , we have w.p.  1  2  over S ⇠ D n and J , the following holds ( L ( w ) :
[Equation]
where C  = 4 + 2 ln(1 /  ) + 5 .[-
-]
[Equation]
Stochastic Gradient Langevin Dynamics (SGLD): For a given training data set S , the update rule of SGLD is deﬁned as:
[Equation]
where B t ⇠ uniform([ n ]) b is the mini-batch of size b at step t .
Note that B t is a sequence instead of a set, thus it may include duplicate elements.
Similar to the analysis of GLD, we can prove the following bound for SGLD.
Theorem 6.3.
Let W T be the output of SGLD when the training set is S , and J be a random sequence with m indices uniformly sampled from [ n ] without replacement.
For any  2 (0 , 1) and m  1 , we have w.p.  1  2  over S ⇠ D n and J , the following holds:
[Equation]
where L ( w ) :
= max z 2 S k f ( w, z ) k , C  = 4 + 2 ln [Equation]
1 /  ) + 5 .[-
-]
[Equation] b is the batch size, and I =
[ n ] \ J .
Remark 6.4.
If the gradient norm is bounded 3 and we use a decaying learning rate schedule such as  t / O (1 /t ) , then the summation in our bound converges.
Hence, under such a learning rate schedule, Theorem 6.2 and 6.3 imply the following test error bound for GLD or SGLD:
R ( W T , D )  ⌘ C ⌘ R ( W T , S I )
+ e
O ( 1 n  m ) which is independent of T , where e
O hides some logarithmic factors.
Experiment
In this section, we conduct experiments for FGD and FSGD on MNIST [LeCun et al., 1998] and
{+CIFAR
+}
[Krizhevsky et al., 2009] to investigate the the optimization and generalization properties of FGD and FSGD, and the numerical closeness between our theoretical bounds and true test errors.
Due to space limit, the detailed experimental setting and some additional experimental results can be found in Appendix H.
FGD/FSGD vs GD/SGD.
We ﬁrst demonstrate that the training and testing curves of FGD and GD are nearly identical (we choose precision level " = 0 .
005 or 0 . 004 ).
We also show that the same is true for FSGD vs SGD.
Due to space limit, the ﬁgures are presented in Appendix H (Figure 5 and 6).
[Figure]
Non-vacuous bounds.
For MNIST, we train a CNN ( d = 1 . 4 · 10 6 ) by FGD with  t = 0 .
9 b t 150 c and " t = 0 .
005 and momentum ↵ = 0 .
The size m = |
J | is set to n/ 2 = 30000 .
As shown in Figure 1a and 1b, our bound (Theorem 5.2 with ⌘ = 1 . 5 ,  = 0 . 1 ) tracks the testing error closely.
At step T = 990 , our bound is 0 .
026 while the testing error is 0 .
This is non-vacuous and tighter than best known 11% MNIST bound reported in Dziugaite et al.
For CIFAR10, we train a SimpleNet [Hasanpour et al., 2016] without BatchNorm and Dropout.
The number of parameters d is nearly 18 · 10 6 .
We use FSGD to train our model.
The learning rate  t is set to 0 .
001 · 0 . 9 b t/ 200 c , the precision " t is set to 0 .
004 , and the momentum ↵ is set to 0 .
The batch size is 2000 .
m =
| J | is set to n/ 5 = 10000 .
The result is shown in Figure 2b.
We stop training at step t = 8000 when the testing error is and 0 .
At that time, our testing error bound is 0 .
198 which is non-vacuous and tighter than best known 0 .
23 CIFAR10 bound reported in Dziugaite et al.
Decrease of the gradient difference.
Intuitively, the cumulative squared norm of gradient difference g t :=
r f ( W t , S )  r f ( W t , S J ) should decrease as m = | J | increases.
Although we cannot prove a concentration like Lemma 6.1 (i.e., k g t k 2 scales as O (1 /m ) ), we can still observe that k g t kdecreases when m increases.
The results are depicted in Figure 1c and Figure 2c.
Random labels.
We conduct the random label experiment designed in Zhang et al.
Our theoretical bounds can distinguish the datasets with different portion ( p ) of random labels.
See
Appendix H.  Conclusion In this paper, we prove new generalization bounds for several gradient-based methods with either discrete or continuous noises based on carefully constructed data-dependent priors.
Recall that FGD requires to compute the gradient difference for technical reasons.
It would be more natural and desirable if we only need to compute the full gradient and rounded to the nearest grid point.
An intriguing future direction is to free FGD/FSGD from the dependence of the prior subset J so that we can apply the concentration on the gradient difference to obtain a tighter bound.
Of course, a major further direction is to obtain similar generalization bounds for vanilla GD and SGD, which remains to be an important open problem in this line of work.
Our technique can be useful for handling deterministic algorithms and discrete noises, but it seems that new technical ideas or assumptions are needed for tackling GD or SGD.  Acknowledgements
The authors would like to thank the anonymous reviewers for their constructive comments.
The authors are supported in part by the National Natural Science Foundation of China Grant 62161146004,
Turing AI Institute of Nanjing and Xi’an Institute for Interdisciplinary Information Core Technology.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] • Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See the conclusion section.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] It can be found in our supplemental material.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A]
We only use the standard MNIST and CIFAR10 datasets.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]