Generalization Bounds for Gradient Methods via Discrete and Continuous Prior
Abstract
Proving algorithm-dependent generalization error bounds for gradient-type optimization methods has attracted significant 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 injectednoise (such as the Gaussian noise in Langevin dynamics).
In this paper, we introduce a new discrete data-dependent prior to the PAC-Bayesian framework, andprove a high probability generalization bound of order O ( 1 n · (cid:80)
Tt =1 ( γ t /ε t ) {+
 ∥+} g t ∥ 2 )for Floored GD (i.e. a version of gradient descent with precision level ε t ), wheren is the number of training samples, γ t is the learning rate at step t , g t is roughlythe difference of the gradient computed using all samples and that using only priorsamples.
∥ g t ∥ is upper bounded by and and typical much smaller than the gradientnorm ∥∇ f ( W t ) ∥ .
We remark that our bound holds for nonconvex and nonsmoothscenarios.
Moreover, our theoretical results provide numerically favorable upperbounds of testing errors (e.g., 0 . 037 on MNIST).
Using similar technique, we canalso obtain new generalization bounds for certain variant of SGD.
Furthermore, westudy the generalization bounds for gradient Langevin Dynamics (GLD).
Using thesame framework with a carefully constructed continuous prior, we show a new highprobability generalization bound of order O ( 1 n + L 2 n 2 (cid:80) Tt =1 ( γ t /σ t ) 2 ) for GLD.
Thenew 1 /n 2 rate is due to the concentration of the difference between the gradient oftraining samples and that of the prior.
Introduction
Bounding generalization error of learning algorithms is one of the most central problems in machinelearning theory.
Formally, for a supervised learning problem, the generalization error is defined asthe testing error (or population error) minus the training error (or empirical error).
In particular,we denote R ( w, ( x, y )) :
=
[ h w ( x ) ̸ = y ] as the error of a single data point ( x, y ) , where h w ( x )is the output of a model with parameter w ∈ R d .
Suppose S is the set of training data, eachi.i.d.
sampled from the population distribution D , and we use R ( w, S ) :
= 1 | S | (cid:80) z ∈ S R ( w, z ) and
R ( w, D ) :=
E z ∼ D
[ R ( w, z )] to denote the training error and the testing error, respectively.
Thegeneralization error of w is formally defined as err gen ( w ) =
R ( w, D ) − R ( w, S ) .
Proving tighter generalization bounds for general nonconvex learning and in particular deep learninghas attracted significant attention recently.
While the classical learning theory (uniform convergence theory) which bounds the generalization error by various complexity measures (e.g., the
VC-dimension and Rademacher complexity) of the hypothesis class has been successful in severalclassical convex learning models, however, they become vacuous and hence fail to explain the successof modern nonconvex over-parametrized neural networks (i.e., the number of parameters significantlyexceeds 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 the generalization ability of deeplearning 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 [Aroraet 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 trainingdata and the optimization algorithms (a.k.a.
gradient-type methods) for general nonconvex learningproblems.
In particular, we prove algorithm-dependent generalization bounds for variants of themost commonly used algorithms such as 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 has a flavor of algorithmic stability [Bousquetand Elisseeff, 2002].
Several prior works have obtained generalization bounds for SGD and SGLDby analyzing trajectory through either the PAC-Bayesian or the algorithmic stability framework (orclosely related information theoretic arguments).
However, most existing results based on analyzingthe optimization trajectories require either restrictive assumptions on the learning rates, or continuousnoise (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-Bayesianframework with a few simple (yet effective) ideas, so that we can obtain new high probability andnon-vacuous generalization bounds for several gradient-based optimization methods with eitherdiscrete or continuous noises (in particular certain variants of GD and SGD, either being deterministicor with discrete noise, which cannot be handled by existing techniques).
1.1 Prior work
We first briefly mention some recent work on bounding the generalization error of gradient-basedmethods.
Hardt et al.
[2016] first studied the uniform stability (hence the generalization) of stochasticgradient descent (SGD) for both convex and non-convex functions.
Their results for non-convexfunctions requires that the learning rate η t scales with 1 /t .
Their work motivates a long line ofsubsequent work on generalization error bounds of gradient-based optimization methods: Kuzborskijand 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.
Simsekli et al.[2020] provided another generalization error bound of SGD using the notion of Hausdorff dimensiond H which depends on both the algorithm and data.
PAC-Bayesian bounds.
The PAC-Bayesian framework [McAllester, 1999] is a powerful methodfor proving high probability generalization bound [Bartlett et al., 2017, Zhou et al., 2019, Mou et al.
Roughly speaking, it bounds the generalization error by the KL divergence KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) ,where Q is the distribution of the learned output and P is a prior distribution which is typicallyindependent of dataset S .
In this framework, bounding KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) is the most crucial part forobtaining tighter PAC-Bayesian bounds.
In order to bound the KL divergence, both the prior Pand posterior Q are required to be continuous distributions (typically Gaussians so that KL canbe computed in closed form).
Hence, most prior work either considered gradient methods withcontinuous 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 final parameter at the end (e.g., [Neyshaburet al., 2018, Zhou et al., 2019]) (so Q is a Gaussian distribution).
We also note that designingeffective prior P can be also very important.
For example, Lever et al.
[2013] proposed to use thepopulation distribution to compute the prior.
In fact, the prior can even partially depend on thetraining data [Parrado-Hernández et al., 2012, Negrea et al., 2019], and our Theorem 4.1 is partiallyinspired 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 interestingconstruction of discrete priors.
Secondly, we consider well studied gradient methods with continuous noise, (stochastic) gradient Langvin dynamics (GLD and SGLD).
We show sharper generalizationbounds by carefully bounding the concentration of the sample gradients.
Now, we summarize ourresults.
FGD and FSGD.
We first study a close variant of GD called FGD (Algorithm 1).
The update ruleof FGD is defined as follows:
[Equation]
where S J is the subset of training dataset S with size m indexed by J ( J is chosen before the training),∇ f ( W t − 1 , Z ) : = 1 | Z | (cid:80) z ∈ Z ∇ f ( W t − 1 , z ) is the mean gradient over the dataset Z , γ t is the learningrate, ε t is the precision level, and g t := ∇ f ( W t − 1 , S ) − ∇ f ( W t − 1 , S J ) is the gradient difference.
The flooring operation on a real number is defined by floor( x ) :
= sign( x ) ⌊| x |⌋ .
It can viewed asgradient descent with given precision limit ε t .
We can see if we ignore the floor operation or let ε tapproaches 0, FGD reduces to GD.
Empirically, its behavior and optimization capability is very closeto GD (See Figure 2a and 2b).
By constructing a discrete data-dependent prior and incorporate it into Catoni’s PAC-Bayesianframework, we prove that the following bound (Theorem 5.2) holds with high probability:
[Equation]
where d is the dimension of parameter space and c 0 can be chosen as a small constant.
Now we makea few remarks about our results.1.
Our result holds for nonconvex and nonsmooth learning problems (replacing the gradientswith subgradients for nonsmooth cases).
There is no additional requirement on the learningrate γ t .2.
The gradient difference g t is typical much smaller than the worst case gradient norm.
Itusually decreases when m = | J | grows (see Figure 1c).
Thus, it enables us to obtain nonvacuous bound on classical datasets.
Specifically, our theoretical test error upper bounds on MNIST and CIFAR10 are 0 . 0288 and 0 . 196 , respectively (see Section 5.1).
Both of themare strictly tighter than the best-known MNIST bound ( 11% ) and CIFAR10 bound ( 23% )reported in Dziugaite et al.
The prior P is a carefully constructed discrete random processes, so that we can bound the
KL between P and the deterministic process of FGD.
We hope it may inspire future researchon tackling deterministic optimization algorithms or discrete noise.4.
The result can be extended to a finite precision variant of SGD, called FSGD (see Theorem 5.3 for details).
Why study FGD/FSGD?
We would like to remark that we study FGD/FSGD, not because
FGD/FSGD have better performances than GD/SGD (in fact, their performances are almost thesame as GD/SGD) or other advantages, but to use them as important stepping stones to extendthe current technique for proving tighter generalization bounds for GD and SGD.
In particular, allexisting trajectory-based analyses require either fast decreasing learning rate, or continuous injectednoise, such as the Gaussian noise in Langevin dynamics, for general non-convex loss functions.
Hence, handling deterministic algorithms (such as GD) or discrete noises (such as SGD) is beyondexisting techniques. Our new technique significantly extends the applicability of existing techniques In fact, understanding such discrete noises and their effect on generalization has been an importantresearch topic (see e.g., Li et al. [2020], Zhu et al.
[2019], Ziyin et al.
In particular, Zhu et al.[2019] shows that it is not sufficient to approximate SGD’s discrete noise by isotropic Gaussian noise.
Moreover, proving nontrivial generalization bounds for SGD-like algorithms with discrete noise hasalso been proposed as an open research direction in Li et al.
GLD, SGLD and CLD.
We also provide new a generalization bound for Gradient Langevin
Dynamics (GLD) whose update rule is defined below.
[Equation]
In this paper, we show that the following generalization bound (Theorem 6.2) holds with highprobability over the randomness of S ∼ D n and random subset J :
[Equation]
where L ( W t − 1 ) :
= max z ∈ S ∥
f ( W t − 1 , z ) ∥ is the longest gradient norm of any training sample in Sat step t and m is the size of J .
Since W T is independent of the prior distribution (defined by theindex set J ), the first term R ( W T , S [ n ] \ J ) is upper bounded by R ( W T , S )
+ O ( 1 √ n − m ) with highprobability.
By setting m = n/ 2 , our generalization bound has an O ( 1 √ n + 1 n + Tn 2 ) rate.
We alsoprove a high probability generalization bound for 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 A.
Other Related Work
Stochastic Langevin Dynamics Stochastic Langevin dynamics is a popular sampling and optimization method used 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-convexsetting.
Raginsky et al.
[2017] study the generalization and excess risk of SGLD in nonconvexsettings and their bound depends polynomially on a certain spectral gap parameter which may beexponential in the dimension.
Continuous Langevin dynamics (SDE) with various noise structure hasalso been used extensively as approximations of SGD (see e.g., [Li et al., 2017, 2021]).
Nonvacuous PAC-Bayesian Generalization Bounds.
Dziugaite and Roy [2017] first present anon-vacuous PAC-Bayesian generalization bound on MNIST (0.20 for a 3 layer MLP).
They usea somewhat different training algorithm that explicitly optimizes the PAC-Bayesian bound and theoutput distribution is a multivariate normal distribution N ( w, diag( s )) .
To facilitate computing theclosed form of KL, they chose a zero-mean Gaussian distribution as the prior.
Zhou et al.
[2019]obtain non-vacuous bound for ImageNet via a quite different approach.
Their method does not requireany continuous noise injected but assumes that the network can be significantly compressed (so thatthe prior distribution is supported over the set of discrete parameters with finite precision).
To ourbest knowledge, it is the only work that use a discrete prior for proving generalization bounds of deepneural networks.
Our result for FGD/FSGD has a similar flavor in a high level, that is the optimizationmethod has a finite precision.
However, our results do not need any assumption on compressibility ofthe model and can be applied to other nonconvex learning problems.
Generalization bound via Information theory.
Raginsky et al.
[2017] first show that the expectedgeneralization error E S ∼D n [ R ( W, D ) −R ( W, S )] is bounded by (cid:112)
2 I ( S ; W ) /n , where I ( S ; W ) :=
KL (cid:0) P ( S, W ) (cid:12)(cid:12)(cid:12)(cid:12) P ( S ) ⊗ P ( W ) (cid:1) 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, Buet al., 2020, Wang et al., 2021].
The main goal in this line of work is to obtain a tight bound on themutual information I ( S ; W ) .
This is again reduced to bounding the KL divergence and thus typicallyrequires continuous injected noise (e.g., Wang et al.
[2021], Negrea et al.
Preliminaries
Notations.
We assume the training dataset S =
( z 1 , .., z n ) is sampled from D n , where D is thepopulation distribution on 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 × Ω → Ris a proxy of the risk.
The optimization algorithm minimizes the loss function and we assume we cancompute the gradient of the loss function.
We note that the loss function may be different from the riskfunction (e.g., 0/1 risk vs the cross-entropy loss).
The empirical risk is R ( w, S )
= 1 | S | (cid:80) z ∈ S R ( w, z )and population risk is R ( w, D ) =
E z
∼D
[ R ( w, z )] .
Similarly, we can define the empirical lossf ( 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 denotethe merged sequence ( A 1 , A 2 , ..., A n , B 1 , ..., B m ) .
When the elements in sequence J are differentfrom each other, we also use J to represent the set consisting of all of its elements.
We may alsoslightly 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 (cid:0) X (
cid:12)(cid:12)(cid:12)(cid:12) Y (cid:1) means KL (cid:0) P X (cid:12)(cid:12)(cid:12)(cid:12)
P Y (cid:1) .
For a random variable W , we define R ( W, S ) =
E w ∼ W
[ R ( w, S )] and R ( W, D ) =
E w ∼ W
[ R ( w, Q )] .
The set { 1 , 2 , ..., n } isdenoted by [ n ] .
KL-divergence.
Let P and Q be two probability distributions.
The Kullback–Leibler divergence
KL (cid:0) P (cid:12)(cid:12)(cid:12)(cid:12) Q (cid:1) from P to Q is defined 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 (cid:0) P (cid:12)(cid:12)(cid:12)(cid:12) Q (cid:1)
=
(cid:80) x P ( x ) ln P ( x ) Q ( x ) .
Otherwise, if P and Q are continuous distributions, it is definedas (cid:82) P ( x ) ln P ( x ) Q ( x ) d x .
The following Lemma 3.1 is frequently used in this paper and is a well knownproperty 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 ′ =
( W ′ 0 , ..., W ′ T ) .
Then, the following equation holds (given all KLs are well defined):
[Equation]
E w ∼ W t − [-
-]1(cid:104) KL
[Equation]
Here W t | W t − 1= w denotes the distribution of W t conditioning on W t − 1= ( W 0 , . .
, W t − 1 )
= w .
PAC-Bayesian.
In this paper, we use the PAC-Bayesian bound presented in Catoni [2007] whichenjoys a tighter O (KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) /n ) rate comparing to the traditional O ( (cid:113) KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) /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 independentof S , any δ ∈ (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 B.
Lemma 3.3.
Suppose Φ : [ n ] m → R + is order-independent 1 and | Φ( J ) − Φ( J ′ )
| ≤ c holds for anyadjacent J, J ′ satisfying | J (cid:84) J ′ |
= m − 1 2 .
Let J be m indices sampled uniformly from [ n ] withoutreplacement.
Then Pr J [Φ( J ) − E J
[Φ( J )] > ϵ ] ≤ exp( − 2 ϵ 2 mc 2 ) .
Data-Dependent PAC-Bayesian Bound
The dominating term in the PAC-Bayesian bound (1) is KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) /n , where P is a prior distribution independent of the training dataset S .
Typically, without knowing any information from S , thebest possible bound for KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) P (cid:1) we can hope is at least Θ(1) (it should not be a function of nhence should not decrease with n ).
However, if we are allowed to see m data points from S whenconstructing our prior, we may produce better prediction on posterior Q S .
The following theoremenables 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 C.
Theorem 4.1 (Data-Dependent PAC-Bayesian) .
Suppose J is a random sequence including mindices uniformly sampled from [ n ] without replacement.
For any δ ∈ (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 theinformation of S J ( S J is the subset of S indexed by J ), and C η := 1 is a constant. Input: Training dataset S = ( z 1 , .., z n ) . Index set J .
Result: Parameter W T ∈ R d . 1 Initialize W 0 ← w 0 ; 2 for t : 1 → T do g 1 ← γ t ∇ f ( W t − 1 , S ) ; g 2 ← γ t ∇ f ( W t − 1 , S J ) ; W t ← W t − 1 − g 2 − ε t · floor(( g 1 − g 2 ) /ε t ) ; 6 end
Remarks. Note that the above bound holds regardless of whether Q depends on S or not.
Alsonote that the first term in the right hand side is R ( Q, S I ) , not R ( Q, S ) as in the usual generalizationbounds. We remark that for most of our learning algorithms that are independent of J (i.e., changing Jdoes not change the output Q ), by standard Chernoff-Hoeffding inequality, R ( Q, S I ) can be boundedby R ( Q, S )
+ O (1 / √ n − m ) with high probability over the randomness of J .
For example, GLD,
SGLD and CLD are the standard versions and independent of J , hence R ( Q, S I ) can be replacedby R ( Q, S )
+ O (1 / √ 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 extractsmore information from S J but not much from S I , unintentionally making R ( Q, S I ) the validationerror, rather than the training error as it should be.
However, from our experiment (Figure 2a and 2d)on CIFAR10, we can see that FGD is very close to GD and the S I error R ( W T , S I ) is indeed closeto the training error R ( W T , S ) and both are significantly 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.
Floored Gradient Descent
In this section, we study the generalization error of a finite precision variant of gradient descent: Floored Gradient Descent (FGD).
First we need to define the “flooring” operation which will be usedin our algorithm.
Definition 5.1 (Flooring) .
For any vector X ∈ R d , let Y = floor( X ) defined as:
Y [Equation] floor
[Equation] X
[Equation]
X [Equation] for all i ∈ [ d ] .
The Floored Gradient Descent algorithm is defined in Algorithm 1, where ( γ t ) t ≥ 0 and ( ε t ) t ≥ [-
 arethe-]0 are the step size and precision sequences, respectively.
For a subset Z ⊆ S , we write ∇ f ( W t − 1 , Z ) :=| Z | (cid:80) z ∈
Z ∇ f ( W t − 1 , z ) .
Note that FGD can viewed as gradient descent with given precision limit ε t . We can see if we ignore the floor operation or let ε t approach 0, FGD reduces to the ordinary GD.
We also study momentum FGD, in which the 5th line is replaced by
[Equation]
Here α > 0 is a constant.
We remark that both FGD and its momentum version are deterministicalgorithms.
The following theorem provides the generalization error bound for both algorithms.
240 Theorem 5.2.
Suppose J is a random sequence consisting of m indices uniformly sampled from [ n ]without replacement.
Then for any δ ∈ (0 , 1) , both FGD (Algorithm 1) and its momentum versionsatisfy the following generalization bound w.p. at least 1 − δ over S ∼ D n and J : [Equation] = (cid:18)
[Equation]
ε [Equation] (cid:19) , 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 := ∇ f ( W t − 1 , S )
− ∇ f ( W t − 1 , S J ) .
Proof.
We use Theorem 4.1 to prove our theorem.
The key is to construct the prior distribution
P ( S J ) such that KL (cid:0) W T (cid:12)(cid:12)(cid:12)(cid:12) P ( S J ) (cid:1) is tractable.
Let p be any real number in (0 , 1 / 3) .
We define P ( S J ) as the distribution of W ′ T obtained by the following update rule ( W ′ 0 :
= w 0 ):
[Equation]
where ξ t is a discrete random variable such that for all ( a 1 , .., a d ) ∈ Z d :
[Equation]
It is easy to verify that the sum of the probabilities ( (cid:80) a ∈ Z d Pr[ ξ t = a ] ) equals to 1 .
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:
[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 onw t
[Equation]
− w
[Equation]
2 ) − γ
[Equation]
w
[Equation]
[Equation]
Let vector a = floor( γ t ε t ( ∇ f ( w t − 1 , S ) − ∇ f ( w t − 1 , S J ))) .
By the definition of W ′ t , we have
[Equation]
Since | i | ≤ i 2 and p ∈ (0 , 1 / 3) , we have ln (cid:16) ( (cid:80) ∞ i =
−∞ p i 2 ) d (cid:17) is at most d ln (cid:0) 1 + 2 (cid:80) ∞
i =1 p i (cid:1) .
Itcan 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 + (cid:80) dk
=1 ln(1 /p )
a 2 k . Recall that thek th entry of a is a k := ⌊ γ t ε t · ( ∇ k f ( w t − 1 , S ) − ∇ k f ( w t − 1 , S J )) ⌋ , which is less than or equal toγ t ε t · ( ∇ k f ( w t − 1 , S ) − ∇ 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 FSGD (a variant of SGD).
It is identical to Algorithm 1 except for the definitions of g 1 and g 2 replaced with:
[Equation]
where B t ⊆ [ n ] is a random batch independent of S, J and W t − 1.
The following theorem provides ageneralization bound for FSGD.
The proof can be found in Appendix D.
Theorem 5.3.
Suppose J is a random sequence consisting of m indices uniformly sampled from [ n ]without replacement.
Then for any δ ∈ (0 , 1) , ε ∈ (0 , 1) , FSGD satisfies the following generalizationbound: 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 ) − ∇ f ( W t − 1 , S J ∩ B t ) . [Figure]
[Figure]
5.1 Experiment
We run experiments for FGD and FSGD on MNIST [LeCun et al., 1998] and CIFAR
[Krizhevskyet al., 2009] to investigate the performance of our bounds as well as the properties of the algorithms.
The details of experiments can be found in Appendix G.
Non-vacuous bounds.
For MNIST, we train a CNN ( d = 1 . 4 · 10 6 ) by FGD with γ t = 0 . ⌊ t 150 ⌋ and ε t = 0 .
005 and momentum α = 0 . 9 ).
The size m = |
J | is set to n/ 2 = 30000 .
Asshown in Figure 1a and 1b, our bound (Theorem 5.2 with η = 1 . 5 , δ = 0 . 1 ) tracks the testing errorclosely.
At step T = 990 , our bound is 0 .
0288 while the testing error is 0 .
For CIFAR10,we train a SimpleNet [Hasanpour et al., 2016] without BatchNorm and Dropout.
The number ofparameters d is nearly 18 · 10 6 .
We use FSGD to train our model.
The learning rate γ t is set to0 .
⌊ t/ 200 ⌋ , the precision ε t is set to 0 .
004 , and the momentum α is set to 0 .
The batchsize is 2000 .
m =
| J | is set to n/ 5 = 10000 .
The result is shown in Figure 2c.
We stop trainingat step t = 8000 when the training and testing error is 0 . 021 and 0 . 184 , respectively. At that time,our testing error bound is 0 . 196 which is non-vacuous and strictly tighter than best known 0 .
CIFAR10 bound reported in Haghifam et al.
Decrease of the gradient difference.
Intuitively, the cumulative squared norm of gradient difference g t := ∇ f ( W t , S ) − ∇ f ( W t , S J ) should decrease when increasing the size of J .
Although wecannot prove a concentration like Lemma 6.1 that ∥ g t ∥ 2 is O (1 /m ) since W t of FGD depends on J ,we can still observe that ∥ g t ∥ 2 decreases when m increases.
The result is depicted in Figure 1c.
Gradient Langevin Dynamics
In this section, we present an O ( 1 n + Tn 2 ) generalization bound for the Gradient Langevin Dynamics(GLD).
The GLD algorithm can be viewed as gradient descent plus a Gaussian noise.
Formally, for agiven S ∼ D n , the update rule of GLD is defined as follows:
[Equation]
Here the gradient ∇ f
( W t , z i ) 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 suffix
We still use the data-dependent PAC-Bayesian framework (Theorem 4.1) to prove the generalizationbound for GLD.
Unlike FGD (Algorithm 1), GLD is independent of the prior indices J , whichenables us to prove the following concentration bound (Lemma 6.1) for the gradient difference.
Theproof is based on Lemma 3.3, which is postponed to Appendix E.
Lemma 6.1.
Let S =
( z 1 , ...
z n ) be any fixed training set.
J is a random sequence including mindices uniformly sampled from [ n ] without replacement, and W =
( W 0 , ..., W T ) is any randomsequence independent of J .
Then the following bound holds with probability at least 1 − δ over therandomness of J :
[Equation]
where C δ = 4 + 2 ln(1 /δ )
+ 5 .
(cid:112) ln(1 /δ ) , and L ( w )
= max i ∈ [ n ] ∥∇ f ( w, z i ) ∥ .
Now we are ready to present our main results.
The proofs can be found in Appendix E.
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 δ ∈ (0 , 12 ) and η > 0 , we have w.p.≥ 1 − 2 δ over S ∼ D n and J , the following holds: [Equation] (cid:35)
[Equation]
where L ( w ) :=
max z ∈ S ∥ f ( w, z ) ∥ , C δ = 4 + 2 ln(1 /δ )
+ 5 .
(cid:112) ln(1 /δ ) , I = [ n ] \ J and C η =1 1 − e − η .
Stochastic Gradient Langevin Dynamics For a given training data set S , the update rule of SGLDis defined as:
[Equation]
where B t ∼ uniform([ n ]) b is the mini-batch of size b at step t .
Note that the B t is a sequence insteadof a set, thus it may include duplicate elements.
Similar to the analysis of GLD, we can prove thefollowing 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 sequencewith m indices uniformly sampled from [ n ] without replacement.
For any δ ∈ (0 , 1) and m ≥ 1 , wehave w.p. ≥ 1 − 2 δ over S ∼ D n and J , the following holds:
[Equation]
where L ( w ) :=
max z ∈ S ∥ f ( w, z ) ∥ , C δ = 4 + 2 ln(1 /δ )
+ 5 .
(cid:112) ln(1 /δ ) , C η = 1 1 − e − η , b is thebatch size, and I =
[ n ] \ J .
Remark 6.4.
If the gradient is bounded 3 and we use a decaying step size schedule such as γ t ∝ 1 /t ,then the summation in our bound converges.
Thus Theorem 6.2 and 6.3 imply the following testerror bound for GLD or SGLD: R ( W T , D ) ≤ ηC η R ( W T , S I ) + (cid:101) O ( 1 n − m ) which is independent of
T , where (cid:101) O hides some logarithmic factors.
Conclusion
In this paper, we prove new generalization bounds for several gradient-based methods with eitherdiscrete or continuous noises based on carefully constructed data-dependent priors.
A limitation isthat FGD requires compute the gradient difference for technical reasons.
It would be more naturaland desirable if we only need to compute the full gradient and rounded to the nearest grid point.
Anintriguing future direction is to free FGD/FSGD from the dependence of J so that we can applythe concentration on the gradient difference to obtain a tighter bound.
Of course, a major furtherdirection is to obtain similar generalization bounds for vanilla GD and SGD, which remains to be animportant open problem in this line of work.
Our technique can be useful for handling deterministicalgorithms and discrete noises, but it seems that new technical ideas or assumptions are needed fortackling GD or SGD.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe 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 areproprietary.•
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 instructionsblock 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 reflect 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 tothem?
[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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes] It can be found in our supplementalmaterial.
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 identifi-
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]