O N THE G ENERALIZATION OF M ODELS T RAINED WITH SGD: I NFORMATION -T HEORETIC B OUNDS AND I MPLICATIONS
A BSTRACT
This paper follows up on a recent work of Neu et al.
and presents some new information-theoretic upper bounds for the generalization error of machine learning models, such as neural networks, trained with SGD.
We apply these bounds to analyzing the generalization behaviour of linear and two-layer ReLU networks.
Experimental study of these bounds provide some insights on the SGD training of neural networks.
They also point to a new and simple regularization scheme which we show performs comparably to the current state of the art.
I NTRODUCTION
The observation that high capacity deep neural networks trained with mini-batched stochastic gradient descent, referred to SGD in this paper tend to generalize well (Zhang et al., 2017) contradicts the classical wisdom in statistical learning theory (e.g., Vapnik (1998) ) and has stimulated intense research interest in understanding the generalization behaviour of modern neural networks.
In this direction generalization bounds for over-parameterized neural networks are obtained (AllenZhu et al., 2019; Bartlett et al., 2017; Neyshabur et al., 2015; 2018a;b; Arora et al., 2018; 2019) and a curious “double descent” phenomenon is observed and analyzed (Belkin et al., 2019; Nakkiran et al., 2019; Yang et al., 2020).
Built on a connection between stability and generalization (Bousquet & Elisseeff, 2002), a stability-based bound is ﬁrst presented in Hardt et al.
(2016), followed by a surge of research effort exploiting similar approaches (London, 2017; Chen et al., 2018; Feldman &
Vondrak, 2019; Lei & Ying, 2020; Bassily et al., 2020).
Information-theoretic bounding techniques established recently (Russo & Zou, 2016; 2019; Xu & Raginsky, 2017; Asadi et al., 2018; Bu et al., 2020; Steinke & Zakynthinou, 2020; Dwork et al., 2015; Bassily et al., 2018; Asadi & Abbe, 2020;
Hafez-Kolahi et al., 2020; Zhou et al., 2020) have also demonstrated great power in analyzing SGDlike algorithms.
For example, Pensia et al.
(2018) is the ﬁrst to utilize information-theoretic bound in analyzing the generalization ability of SGLD (Gelfand & Mitter, 1991; Welling & Teh, 2011).
The bound was subsquently improved by Negrea et al.
; Haghifam et al.
Rodr´ıguezG´alvez et al.
; Wang et al.
(2021b).
Inspired by the work of Pensia et al.
(2018), Neu et al.
(2021) presents an information-theoretic analysis of the models trained with SGD.
The analysis of
Neu et al.
(2021) constructs an auxiliary weight process parallel to SGD training and upper-bounds the generalization error through this auxiliary process.
Another line of research connects the generalization of neural networks with the ﬂatness of loss minima (Hochreiter & Schmidhuber, 1997) found by SGD or its variant (Keskar et al., 2017; Dinh et al., 2017; Dziugaite & Roy, 2017; Neyshabur et al., 2017; Chaudhari et al., 2017; Jastrzebski et al., 2017; Jiang et al., 2019; Zheng et al., 2021; Foret et al., 2020).
This understanding has led to the discovery of new SGD-based training algorithms for improved generalization.
For example, in a concurrent development by Zheng et al.
(2021) and Foret et al.
(2020), a local “max-pooling” operation is applied to the loss landscape prior to the SGD updates.
This approach, referred to as
AMP (Zheng et al., 2021) or SAM (Foret et al., 2020), is shown to make SGD favor ﬂatter minima and achieve the state-of-the-art performance among various competitive regularization schemes.
In this paper, we focus on investigating the generalization of machine learning models trained with
SGD.
Although we are primarily motivated by the curiosity to understanding neural networks, the results of this paper in fact apply broadly to any model trained with SGD.
This work follows the same construction of the auxiliary weight process in Neu et al.
and develops upper bounds of generalization error that extend the work of Neu et al.
Like those in Neu et al.
(2021), the bounds we obtain can be decomposed into two terms, one measuring the impact of training trajectories (“the trajectory term”) and the other measuring the impact of the ﬂatness of the found solution (“the ﬂatness term 1 ”).
Having an identical ﬂatness term as that in Neu et al.
(2021), empirical evidence hints that our bounds have a improved trajectory term.
Figureshows an experimental comparison of the trajectory term in our bound (Theorem 2) with that in the bound in Neu et al.
(2021) (re-stated as Lemma 2 in this paper) for two neural network models.
The trajectory terms are compared in a deterministic setting for two different values of σ (the variance parameter of the noise in the auxiliary weight process, details given in appropriate context and in
Appendix E.2).
The trajectory term in the bounds of Neu et al.
(2021) accumulates two terms over training steps: local gradient sensitivity , measuring the sensitivity of the SGD gradient signal to weight perturbations, and gradient dispersion 2 , measuring the extent to which the gradient signal spreads around its mean.
Usually gradient dispersion vanishes with training iterations but local gradient sensitivity does not.
Our improvement over Neu et al.
(2021) is achieved by removing the local sensitivity term and involving a logarithmic function of a revised gradient dispersion.
Although our new deﬁnition of gradient dispersion is in general larger than that in Neu et al.
(2021), as long as the number of iterations is not too small, the fact that it also vanishes with training allows our bounds to be tighter than Neu et al.
[Figure]
We also provide an application of our bounds in analyzing the generalization behaviour of linear and two-layer ReLU networks, where we show that the activation state in ReLU networks plays an important role in generalization.
It is remarkable that removing the local gradient sensitivity term makes our bound have a simple closed form (after optimizing the noise variances), much easier to evaluate.
We empirically validate the derived bounds, and provide various insights pertaining to the generalization behavior of models trained with SGD.
For example, gradient dispersion is seen to reveal a double descent phenomenon with respect to training epochs, where the valley in the double descent curve appears to mark the great divide between the “generalization regime” and the “memorization regime” of training.
Furthering from this observation, we also show that it is possible to reduce the memorization effect by dynamically clipping the gradient and reducing its dispersion.
Our bounds also inspire a natural and simple solution to alleviate generalization error.
Speciﬁcally, we propose a new training scheme, referred to as Gaussian model perturbation (GMP), aiming at reducing the ﬂatness term of the bounds.
This scheme effectively applies a local “average pooling” to the empirical risk surface prior to SGD, greatly resembling the “max-pooling” approach adopted in AMP (Zheng et al., 2021).
We demonstrate experimentally that GMP achieves a competitive performance with the current art of regularization schemes.
Proofs, additional discussions and experimental results are presented in Appendices.
Other Related Literature Gradient dispersion is mostly studied from optimization perspectives (Bottou et al., 2018; Roux et al., 2012; Johnson & Zhang, 2013; Wen et al., 2020; Faghri et al., 2020).
Prior to this work, only a few works relate gradient dispersion with the generalization behaviour of the networks.
In Neu et al.
(2021) and Wang et al.
(2021b), gradient dispersion also appears in the generalization bounds.
and there are limited studies characterizing the generalization performance and the gradient variance.
Particularly, in Jiang et al.
(2019), gradient dispersion is argued to capture a notion of “ﬂatness” of the local minima of the loss landscape, thereby correlating with generalization.
Injecting noise in the training process has been proposed in various regularization schemes (Bishop, 1995; Camuto et al., 2020; 2021; Srivastava et al., 2014; Wei et al., 2020).
But unlike GMP derived in this paper, where noise is injected to the model parameters, noise in those schemes is injected either to the training data or to the network activation.
In addition, noise can also be implicitly injected during training, for example, by adjusting the momentum terms Xie et al.
(2021b).
The training objective in Xie et al.
(2021a) is similar to our GMP but their NVRM-SGD minimizes a region directly while our GMP considers a trade-off between the empirical loss and the ﬂatness term.
Gradient clipping is a common technique for preventing gradient exploding (see, e.g., Merity et al.
; Peters et al.
This technique is also used in Zhang et al.
(2019) to accelerate training.
In this paper, gradient clipping is used to investigate and control the impact of gradient dispersion on generalization error.
P RELIMINARIES
Population Risk, Empirical Risk and Generalization Error Unless otherwise noted, a random variable will be denoted by a capitalized letter (e.g., Z ), and its realization denoted by the corresponding lower-case letter (e.g. z ).
Let Z be the instance space of interest and µ be an unknown distribution on Z , specifying random variable Z .
Let
W ⊆ R
d be the space of hypotheses.
Suppose that a training sample S =
( Z 1 , Z 2 , . . .
, Z n ) is drawn i.i.d.
from µ and that a stochastic learning algorithm A takes S as its input and outputs a hypothesis W ∈ W according to some conditional distribution P W | S mapping Z n to W .
Let (cid:96) :
W × Z → R + be a loss function, where (cid:96) ( w, z ) measures the “unﬁtness” or “error” of any z ∈ Z with respect to a hypothesis w ∈ W .
The population risk, for any w ∈ W , is deﬁned as
[Equation]
The goal of learning is to ﬁnd a hypothesis w that minimizes the population risk.
But since µ is only partially accessible via the sample S , in practice, we instead turn to the empirical risk, deﬁned as
[Equation]
The expected generalization error of the learning algorithm
A is then deﬁned as
[Equation]
where the expectation is taken over the joint distribution of ( S, W ) (i.e., µ n ⊗ P W | S ).
Throughout this paper, we take (cid:96) as a continuous function (adopting the usual notion “surrogate loss” (Shalev-Shwartz & Ben-David, 2014) ).
Additionally, we assume that (cid:96) is differentiable almost everywhere with respect to both w and z .
Furthermore we assume that (cid:96) ( w, Z ) is R -subgaussianfor any w ∈ W .
Note that a bounded loss is guaranteed to be subgaussian.
Let I ( X ; Y ) denote the mutual information (Cover & Thomas, 2012) between any pair of random variables ( X, Y ) .
The following result is known.
Lemma 1 (Xu & Raginsky (2017, Theorem 1.)) .
Assume the loss (cid:96) ( w, Z ) is R -subgaussian for any w ∈ W .
The generalization error of A is bounded by
[Equation]
Stochastic Gradient Descent We now restrict the learning algorithm
A to be the mini-batched
SGD algorithm for empirical risk minimization.
For each training epoch, the dataset S is randomly split into m disjoint mini-batches, each having size b , namely, n = mb .
Based on each batch, one parameter update is performed.
Speciﬁcally, let B t denote the batch used for the t th update.
Deﬁne g ( w, B
[Equation]
w,
[Equation]
) , namely, g ( w, B t ) is the average gradient computed for the batch B t with respect to parameter w .
The rule for the t th parameter update is then
[Equation]
where λ t is the learning rate at the step t .
The initial parameter setting W 0 is assumed to be drawn from the zero-mean spherical Gaussian N (0 , σ 20 I d ) with variance σ 20 in each dimension.
We will assume that the SGD algorithm stops after T updates and outputs W T as the learned model parameter.
Given the training sample S , let ξ govern the randomness in the sequence ( B 1 , B 2 , . . .
, B T ) of batches.
For the simplicity of notion, we will ﬁx the conﬁguration of ξ .
That is, we will assume a ﬁxed “batching trajectory”, or a ﬁxed way to shufﬂe the example indices { 1 , . . . , n } and divide them into m batches in each epoch.
The presented generalization bounds of this paper can be extended to the case where the batching trajectory is uniformly random (as we set up above).
This merely involves averaging over all batching trajectories or taking expectation over ξ .
Auxiliary Weight Process We now associate with the SGD algorithm an auxiliary weight process { (cid:102) W t } .
Let σ 1 , σ 2 , . . .
, σ T be a sequence of positive real numbers.
Deﬁne
[Equation]
and (cid:102) W t (cid:44) (cid:102) W t − 1 − λ t g ( W t − 1 , B t )
+ N t , for t > 0 , where N t ∼ N (0 , σ 2 t I d ) is a Gaussian noise.
The relationship between this auxiliary weight process { (cid:102) W t } and the weight process { W t } in SGD is shown in the Bayesian network below.
[Equation]
Let ∆ t = (cid:80) tτ
=
1 N τ .
Noting that the weight updates in { (cid:102) W t } uses the same gradient signal as that used in { W t } (which depends on W t − 1 not (cid:102) W t − 1 ), it is immediate that (cid:102) W t
=
W t + ∆ t .
Note that this auxiliary process follows the same construction as Neu et al.
(2021), which we will use to study the generalization error of SGD.
To that end, let’s deﬁne gradient dispersion by
[Equation]
(1) where the inside expectation is taken over ( W, Z ) ∼ µ ⊗ P W | Z .
For a given sample s ∈ Z n , deﬁne
[Equation]
where the expectation is taken over ∆ T and L s ( w ) is the empirical risk of s at parameter w .
In the remainder of the paper, let S (cid:48) denote another sample drawn from µ n , independent of all other random variables.
The main generalization bound in Neu et al.
(2021) is re-stated below.
Lemma 2 ((Neu et al., 2021, Theorem 1.)) .
The generalization error of SGD is upper bounded by
[Equation] [Equation]
where Ψ( w t − 1 ) (cid:44) E (cid:2) || E [ ∇ w (cid:96) ( w t − 1 , Z )] − E [ ∇ w (cid:96) ( w t − 1 + ζ, Z )]
|| 22 (cid:3) , ζ ∼ N (0 , (cid:80) t − 1 i =1 σ 2 i I d ) and (cid:101) V t ( w ) (cid:44) E (cid:2)
||
g ( w, B t ) − E [ ∇ w (cid:96) ( w, Z )]
|| 2 2 (cid:3) .
The term Ψ( w t − 1 ) in the bound is referred to as “local gradient sensitivity” in Neu et al.
Note that the expectation inside (cid:101) V t ( w ) is taken over Z ∼ µ instead of ( W, Z ) ∼ µ ⊗ P W | Z .
Thus, (cid:101) V t ( w ) is in general not worse than our gradient dispersion V t ( w ) for a ﬁxed w (see Figure 5 (d) in Appendix E.2).
However, the difference between these two terms is very small when W is close to local minima due to the tiny gradient norm.
Indeed, the deﬁnition of our gradient dispersion is implicitly used in a bound in Section 5.2 of Neu et al.
(2021), which can be regarded as a weaker version of Eq.
2 in our Theorem 2.
N EW G ENERALIZATION B OUNDS FOR SGD
Removal of the local sensitivity term Ψ( w t − 1 ) requires invoking a special instance of the HWI inequality (Raginsky & Sason, 2018, Lemma 3.4.2), which we ﬁrst state.
Lemma 3.
Let X
and Y be two random vectors in R d , and let N ∼ N (0 , I d ) be independent of ( X, Y ) .
Then, for any t, t (cid:48) > 0 , D KL ( P X + √ tN ||
P Y + √ t (cid:48) N ) ≤
[Equation]
−
[Equation]
, where D KL is the KL divergence.
Note that Neu et al.
(2021) also makes uses of a similar lemma in their revision.
However, the development in Neu et al.
(2021) requires constructing a ghost auxiliary weight process in which an independent noise perturbation ∆ (cid:48) t is introduced.
Due to the mismatch between ∆ t and ∆ (cid:48) t , the local gradient sensitivity term Ψ( W t − 1 ) appears in their bound.
In this paper, we waive this ∆ (cid:48) t by using the following lemma.
Lemma 4.
Let random variables X , Y and ∆ be independent of N ∼ N (0 , I d ) .
Then for any σ > 0 , any R d -valued function f , and any random variable Ω ∈ R d that is a function of Y , we have
[Equation]
Note that the outside expectation is taken over Y and the inside expectation is taken over ( X, ∆) .
Then, exploiting Lemma 4 by letting X = S , Y = (cid:102) W and Ω =
E [ g ( (cid:101) w − ∆ , Z )]
(cid:44) E [ ∇ (cid:96) ( (cid:101) w − ∆ , Z )] will enable us to have the bound below.
Theorem 1.
The generalization error of SGD is upper bounded by
[Equation]
This bound is strictly tighter than the bound in Lemma 2.
In fact, from Theorem 1, one can recover
Lemma 2 with a smaller constant factor (see Appendix C.3 for more details).
To completely remove Ψ( w t − 1 ) and to obtain a closed form of the optimal bound, we will let Ω =
E
[ g ( W, Z )] , then Lemma 4 gives us a crisp way to have the following upper bound that is independent of the distribution of ∆ .
Lemma 5.
Let G t = − λ t g ( W t − 1 , B t ) .
Then, I ( G t + N t ; S | (cid:102) W t − 1 ) ≤ 12 d log (cid:16) λ
2 t E [ V t ( W t − 1 )] dσ 2 t + 1 (cid:17) .
In this lemma, the mutual information I ( G t + N t ; S | (cid:102) W t − 1 ) roughly indicates the degree by which
the
SGD’s updating signal G t (smoothed with noise) depends on the training sample S , when B t is used for computing the gradient.
When this dependency is strong (giving rise to a high value of the mutual information), the model conceivably tends to overﬁt the training sample.
This lemma suggests that the strength of this dependency can be upper-bounded by the expected gradient dispersion at the current weight conﬁguration.
In our experiments, we will estimate the expected gradient dispersion and validate this intuition.
We are now in a position to state our main theorem.
Theorem 2.
The generalization error of SGD is upper bounded by |
[Equation]
Further, assume L µ ( w T ) ≤ E ∆ [ L µ ( w T + ∆ T )] , (cid:96) is twice differentiable, and σ 2 t is independent of t .
Denote by H W T the Hessian of the loss with respect to W T and let Tr( · ) denote trace.
Then gen(
[Equation]
, P
[Equation]
Remark 1.
With a single draw of S and under the deterministic setting (i.e., with ﬁxed weight initialization and batching trajectory), removing Ψ( W t − 1 ) will make the bound much tighter as shown in Figure 1.
This remarkable improvement should come at no surprise, since Ψ( W t − 1 ) monotonically increases with training epochs (noting that Ψ( W t − 1 ) has the cumulative variance (cid:80) t −i =1 σ 2 i I d ) while V t ( W t − 1 ) appears decreasing (see Figure 3a).
Figure 1 also indicates that if the noise variance σ t is made small, one expects the gradient dispersion to dominate the trajectory term in Lemma 2.
However, the factor 1 /σ 2 t will become too large, make the bound loose.
More discussions about the stochastic setting with multiple draws of S are deferred to Appendix E.2.
Remark 2.
The condition L µ ( w T ) ≤ E ∆ T [ L µ ( w T + ∆ T )] indicates that the perturbation does not decrease the population risk.
This is also assumed in Foret et al.
(2020) in the derivation of a
PAC-Bayesian generalization bound.
In the bound of Eq.2, the ﬁrst term captures the impact of the training trajectory (“trajectory term”), and the second term captures the impact of the ﬁnal solution, which in fact measures the ﬂatness for the loss landscape at the found solution (“ﬂatness term”).
The learning rate in SGD has explicitly appeared in the trajectory term of Eq.3.
From the way it appears in the bound, one may be tempted to assert that a small learning rate will improve generalization.
This would then contradict some previous observations (Jastrzebski et al., 2017; Wu et al., 2018;
He et al., 2019), in which large learning rate will beneﬁt generalization.
In addition, the bound also suggests that the batch size has some direct impact on the trajectory term through gradient dispersion.
We investigate this by performing experiments with varying learning rates and batch sizes (see Figure 6 in Appendix E.3).
As seen in the experiments, increasing the learning rate impacts the trajectory term and the ﬂatness term in opposite ways, i.e. making one increase and the other decrease.
A similar (but reverted) behaviour is also observed with batch sizes.
This makes the generalization bound in Eq.3, have a rather complex relationship with the settings of learning rate and batch size.
Eq.3 follows from Eq.2 by minimizing the bound over σ .
The bound in Eq.3 is thus independent of a choice of σ t ’s and can be computed easily and efﬁciently.
To summarize, we remark that these bounds suggest that in order for the model to generalize well, both the trajectory term and the ﬂatness term need to be small — the former involves the learning rate with the gradient dispersion along the training trajectory, whereas the latter depends on the ﬂatness of the empirical risk surface at the found solution.
A PPLICATION IN L INEAR N ETWORKS AND T WO -L AYER R E LU N ETWORKS .
We now apply Theorem 2 to two neural network models in a regression setting.
Let Z =
( X, Y ) with X ∈ R d 0 and Y ∈ R .
Assume ||
X
||
= 1 .
We use SGD to train a model f ( W, · ) : R d 0 → R and deﬁne the loss as (cid:96) ( W, Z ) = 1 / 2( Y − f ( W, X )) 2 .
The bounds below are both in the form of the product of the ﬂatness term and the trajectory term.
Theorem 3 (Linear Networks) .
Let the model be a linear network, i.e. f ( W, X )
=
W T X , then
[Equation]
Theorem 4 (Two-Layer ReLU Networks) .
Following Arora et al.
(2019), consider f ( W, X ) =√ m (cid:80) m
r =1 A r ReLU( W T r X ) where m is the width of the neural network, A r ∼ unif( { +1 , − 1 } ) and
ReLU( · ) is the ReLU activation.
We only train the ﬁrst layer parameters W =
[ W 1 , W 2 , . . .
, W m ] ∈ R d 0 × m and ﬁx the second layer parameters
A =
[ A 1 , A 2 , . . . ,
A m ] ∈ R m during training.
Then,
[Equation]
where I r,i,t
= I { W Tt − 1 ,r X i ≥ 0 } and I is the indicator function.
Compared with Theorem 3, we notice that in the two-layer ReLU network, the ReLU activation state along the training trajectory plays a key role in the bound of Theorem 4.
Speciﬁcally, the weights of the deactivated neurons do not contribute to the bound of Theorem 4, making the bound not explicitly depend on the model dimension d .
This result also suggests that sparsely activated
ReLU networks are expected to generalize better.
Despite various empirical evidence pointing to this behaviour (see, e.g., Glorot et al.
(2011)), to the best of our knowledge, this theorem provides the ﬁrst theoretical justiﬁcation in this regard.
The comparison between these bounds and the dynamics of generalization gap is deferred to Appendix E.4.
In the remainder of this paper, we will study more complex network architectures beyond the linear and two-layer ReLU networks and implications of our bounds in those settings.
E XPERIMENTAL S TUDY
Bound Veriﬁcation
We ﬁrst verify our bound of Eq.3 in Theorem 2 by training an MLP (with one hidden layer) and an AlexNet (Krizhevsky et al., 2012) on MNIST and CIFAR10 (Krizhevsky, 2009), respectively.
To simplify estimation, we ﬁx the weight initialization and set σ t and λ t to be constants σ and λ , respectively.
To compute (cid:80) Tt =1 E
[ V t ( W t − 1 )] , we compute the gradient dispersion as its empirical estimate from a batch, utilizing a PyTorch (Paszke et al., 2019) library
BackPack (Dangel et al., 2020).
To compute Tr ( E [H W T ( Z )]) , we use the PyHessian library (Yao et al., 2020) to compute the Hessian.
We perform experiments with varying network width and varying levels of label noise.
Speciﬁcally, noise level (cid:15) refers to the setting where we replace the labels of (cid:15) fraction of the training and testing instances with random labels.
The estimated bound is compared against the true generalization gap, namely, the difference between the training loss and testing loss, and is shown in Figure 2.
[Figure]
In Figure 2, we see that in all cases the estimated bound follows closely the trend of the true generalization gap.
The fact that the bound curve consistently tracks the gap curve under various label noise levels indicates that our bound very well captures the changes of the data distribution.
Note that in Figure 2 (a) and (b), our bound decays with the increase of the model size, showing a trend as opposite to the bounds obtained in classical learning theory.
But such a trend clearly better explains the generalization behaviour of modern neural networks.
Epoch-wise Double Descent of Gradient Dispersion
We experimentally investigate the impact of gradient dispersion on the training of the neural networks by ﬁxing the learning rate, batch size and weight initialization for the each model (MLP for MNIST, AlexNet for CIFAR10).
For each model and various label noise levels, we plot in Figure 3 the evolution of the (empirical) gradient dispersion (cid:98) V t ( w t − 1 ) , training accuracy and testing accuracy across training epochs.
An intriguing epoch-wise “double descent” phenomenon is observed, particularly when the labels are noisy.
According to the double descent curve, the training may be split into three phases (e.g.,
Figure 3 (h)).
In the ﬁrst phase, the gradient dispersion rapidly descends and maintains a very low level.
In this phase, both training and test accuracies increase while maintaining a very small generalization gap.
This suggests that the network in this phase is extracting useful patterns and generalizes well.
In the second phase, the gradient dispersion starts increasing until it reaches a peak value.
In this phase, the training and testing accuracies gradually diverge, marking the model entering an overﬁtting or “memorization” regime – when the data contains the noisy labels, the network mostly tries to memorize the labels in the training set.
In the third phase, the gradient dispersion descends again, reaching a low value.
In this phase, the model continuously overﬁts the training data, until the training and testing curves reach their respective maximum and minimum.
It appears that the timing of the three phases depends on the dataset and the label noise level.
For simpler data (e
[Figure]
g. MNIST) and cleaner datasets , the ﬁrst phase may be shorter.
This is arguably because in these datasets, extracting useful patterns is relatively easier.
Nonetheless, the valley in the double-descent curve appears to mark a “great divide” between generalization and memorization.
[Figure]
Dynamic Gradient Clipping Inspired by our generalization bounds and above observations, one way to reduce the generalization error is to control the trajectory term of the bounds by reducing the gradient dispersion in each training step.
Here we investigate a simple scheme that dynamically clips the gradient norm so as to reduce the gradient dispersion.
Speciﬁcally, whenever the current gradient norm is larger than the gradient norm K steps earlier, or || g
( W t , B t )
||
2 >
|| g ( W t − K , B t − K )
|| 2 (i
[Figure]
e., the model is expected to have entered the “memorization” regime), we reduce the norm of the current gradient g ( W t , B t ) to α fraction of || g ( W t − K , B t − K )
|| 2 , for some prescribed value α < 1 .
[Figure]
The effectiveness of this scheme is best demonstrated when the labels contain noise.
As shown in Figure 4 (and Figure 9 in Appendix E.5), dynamic gradient clipping signiﬁcantly closes the gap between the training accuracy and the testing accuracy.
The models trained with this scheme maintain a near-optimal testing accuracy (e.g., about 80% when the label noise level of MNIST is 0.2), without suffering from the severe memorization effect as seen in models trained without this scheme.
Further understanding of the doubledescent phenomenon of the gradient dispersion may enable more delicate design of such a dynamic clipping scheme and potentially lead to novel and powerful regularization techniques.
A P RACTICAL I MPLICATION : G AUSSIAN M ODEL P ERTURBATION
The appearance of the ﬂatness term in our generalization bounds suggests that for an empirical risk minimizer w ∗ to generalize well, it is necessary that the empirical risk surface at w ∗ is ﬂat, or insensitive to a small perturbation of w ∗ .
This naturally motivates a training scheme using the following regularized loss:
[Equation]
where ρ is a hyper-parameter.
Replacing the expectation above with its stochastic approximation using k realizations of ∆ gives rise to the following optimization problem.
[Equation]
We refer to the SGD training scheme using this loss as Gaussian model perturbation or GMP.
Notably, GMP requires k + 1 forward passes for every parameter update.
Empirical evidence shows that a small k , for example, k = 3 , already gives competitive performance.
Implementing the k +forward passes on parallel processors further reduces the computation load.
[Table]
We run experiments 10 times and report the mean and the standard deviation
Superscript de notes number sampled Gaussian noises
We compare GMP with several major regularization schemes in the current art, including Dropout (Srivastava et al., 2014), label smoothing (L.S.) (Szegedy et al., 2016), Flooding (Ishida et al., 2020), MixUp (Zhang et al., 2018), adversarial training (Adv.
Tr.)
(Goodfellow et al., 2015), and AMP (Zheng et al., 2021).
We will also include ERM as the baseline, where no regularization method applied.
The compared schemes are evaluated on three popular benchmark image classiﬁcation datasets SVHN (Netzer et al., 2011), CIFAR-10 and CIFAR-
(Krizhevsky, 2009).
Two representative deep architectures VGG16 (Simonyan & Zisserman, 2015) and PreActResNet18 (He et al., 2016) are taken as the underlying model.
We train the models for 200 epochs by SGD.
The learning rate is initialized as 0 . 1 and divided by 10 after 100 and 150 epochs.
For all compared models, the batch size is set to 50 and weight decay is set to 10 − 4 .
For GMP, we choose ρ = 0 .
5 and set the standard deviation of the Gaussian noise ∆ to 0 .
The value of k is chosen as 3 and 10 respectively (referred to as GMP 3 and GMP 10 ).
The performances of all compared schemes are given in Table 1 (and the results of PreActResNet are shown in Appendix E.7).
For the compared regularization schemes except GMP, we directly report their performances as given in Zheng et al.
Table 1 demonstrates the effectiveness of GMP.
Overall GMP performs comparably to the current art of regularization schemes, although appearing slighly inferior to the most recent record given by AMP on SVHN and MixUp on CIFAR10, respectively (Zheng et al., 2021).
Noting that the key ingredient of AMP, “max-pooling” in the parameter space, greatly resembles regularization term in GMP, which may be seen as “averagepooling” in the same space.
C ONCLUSION AND O UTLOOK
This paper presents new information theoretic generalization bounds for models (e.g., linear networks and two-layer ReLU neural networks) trained with SGD.
Our bounds naturally point to new and effective regularization schemes.
At the same time, our bounds and experimental study reveal interesting phenomena in the SGD training of neural networks.
There are yet promising directions for further improving these bounds, for example, via exploiting conditional mutual information bounds (Haghifam et al.
strong data processing inequalities (Wang et al.
(2021a))
, and the relationship between between the trajectory term and the ﬂatness term.
A CKNOWLEDGMENTS
This work is supported partly by an NSERC Discovery Grant and Artiﬁcial Intelligence for Design
Challenge program of National Research Council Canada.
The authors would like to sincerely thank Gergely Neu and Borja Rodr´ıguez-G´alvez for pointing out errors in the previous version of this paper.
The authors would also like to thank Maia Fraser for stimulating discussions.
Appendix Table of Contents
Mutual and
C.1 Lemma 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.3 Corollary 1: Recovering Lemma 2 . . . . . . . . . . . . . . . . . . . . . . . .
C.4 Proof of Lemma 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.5 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.6 Corollary 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Application in Neural Networks: Proof of Theorems 3 and
D.1 Lemma 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.2 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.3 Proof of Theorem 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E
Experiment
Details
E
1 Architectures and Hyperparameters
[Table] [Table]
[Table] [Table]
[Table] [Table]
A HWI I NEQUALITY : P ROOF OF L EMMA
The proof given here is a simple extension of the proof of Raginsky & Sason (2018, Lemma 3.4.2), which is a special instance of the (weak) HWI inequality.
Proof.
[Equation]
where Eq.1 is by the chain rule of the KL divergence, Eq.2 holds since N is independent of ( X, Y ) and Eq.3 is by the equality
[Equation]
Eq.3 holds for any joint distribution of ( X, Y ) .
If we let t (cid:48)
=
t , by the deﬁnition of Wasserstein distance, we have
[Equation]
This recovers the vector version of Raginsky & Sason (2018, Lemma 3.4.2).
B A K EY L EMMA : L EMMA
Proof.
We give two proofs of this lemma, the ﬁrst proof is inspired by the techniques used in Wang et al.
(2021a)
and Guo et al.
(2005), and the second proof is similar to Pensia et al.
Let Z
=
Y + ∆ .
By the properties of mutual information,
[Equation]
Let g ( Z, X ) (cid:44) f
( Z, X ) − Ω .
Let P N (cid:48)
=
N (0 , σ (cid:48)
I d ) for any σ (cid:48) > 0 and N (cid:48) is independent of other random variables for a ﬁxed y .
Then, for Y =
y ,
[Equation]
where Eq. 4 is the golden formula (see Theorem 3.3 in Polyanskiy & Wu (2019) ), Eq.is by the non-negativity of relative entropy, Eq. 6 is due to Lemma 3 and the last equality holds when the optimal
[Equation]
is achieved.
Thus,
[Equation]
which completes the proof.
Alternatively, let h ( · ) denote the differential entropy, then
[Equation]
where Eq.
7 is due to the fact that conditioning reduces entropy.
Notice that
[Equation]
Since the Gaussian distribution maximizes the entropy over all distributions with the same variance (Cover & Thomas
[Equation] [Equation]
we have
[Equation]
Therefore,
[Equation]

 .
The remaining part is straightforward and is the same with the previous proof.
C M UTUAL I NFORMATION B OUNDS FOR SGD: P ROOF OF T HEOREM 1 AND T HEOREM
In this section, we provide a proof for Theorem 1 and Theorem 2, while elaborating on why the local gradient sensitivity term in Lemma 2 can be removed.
The key ingredient that makes this possible is a construction given in the proof of Lemma 5.
C. L EMMA
We ﬁrst give the following useful lemma.
Lemma 6.
The mutual information I ( (cid:102) W T ; S ) ≤ (cid:80) Tt
=
I (cid:16)
[Equation]
g (
[Equation]
t −
[Equation]
, B t
[Equation]
+ N t ;
[Equation]
| (cid:102) W t − 1 (cid:17) .
Proof.
The mutual information between the ﬁnal output of SGD and the training sample can be upper bounded by the mutual information between the full trajectories and the training sample, which is shown below.
[Equation]
where Eq.8 is by I ( f ( X, Y ); Z ) ≤
I ( X, Y ; Z ) (i.e. Z − ( X, Y ) − f
( X, Y ) forms a Markov chain and then use the data processing inequality), Eq.9 is by the chain rule of mutual information, Eq.
is by applying the similar procedure (namely, Eq.8-Eq.9) to I ( (cid:102) W T − 1 ; S ) , Eq.
11 is by doing these steps recursively and Eq.12 is due to the fact that (cid:102) W 0 is independent of S (i.e. I ( (cid:102) W 0 ; S ) = 0 ).
C. P ROOF OF T HEOREM
Proof.
We ﬁrst follow the similar decomposition of the expected generalization error and apply
Lemma 1 as in Neu et al.
[Equation]
The remaining task is to bound the mutual information I ( (cid:102) W T ; S ) .
Let X = S , Y =
(cid:102) W t − 1 and Ω =
E ∆ t − 1 ,Z [ g ( (cid:101) w t − 1 − ∆ t − 1 , Z )] , by applying Lemma 4, we have
I
[Equation]
≤ E (cid:102)
[Equation]
  d
[Equation]
    .
By Lemma 6 and putting everything together
[Equation]
we have | gen
[Equation]
µ
[Equation]
P W T | S ) | ≤
[Equation]
This completes the proof
[Equation]
C. C OROLLARY 1: R ECOVERING L EMMA
We can recover the result of Lemma 2 in Neu et al.
(2021) from Theorem 1, which is re-stated in the following corollary.
Corollary 1.
The generalization error of SGD is upper bounded by
[Equation]
Proof.
From the ﬁrst term in Theorem
[Equation]
we have (cid:118)(cid:117)(cid:117)(cid:117)(cid:117)(cid:117)(cid:117)(cid:117) (cid:116)
[Equation]
R
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
σ
[Equation]
E
[Equation]
(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) (cid:12)
[Equation] [Equation]
(cid:104)
[Equation]
(cid:102)
[Equation]
) (cid:105)(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) (cid:12)
[Equation]
(cid:20)(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) (cid:12)
[Equation]
(cid:104)
[Equation]
(cid:102)
[Equation]
(cid:105)
[Equation]
(cid:104)
[Equation]
( (cid:102)
[Equation]
Z ) (cid:105)(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) (cid:12)(cid:21)(cid:21) (17) =
(cid:118)(cid:117)(cid:117)(cid:116)
R
[Equation]
n
[Equation]
(cid:88)
[Equation]
2Ψ(
[Equation]
(cid:21)
[Equation]
(cid:118)(cid:117)(cid:117)(cid:116)
Rn
T (cid:88) t = λ 2 t σ 2 t
E W t − 1 ,S (cid:20) 4 (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) g
( W t − 1 , B t )
− E Z [ g ( W t − 1 , Z )]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2 + 4 (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) E Z
[ g ( W t − 1 , Z )] − E Z (cid:104) g
( (cid:102) W t − 1 , Z ) (cid:105)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
+ 2Ψ( W t − 1 ) (cid:21) (18) =
(cid:118) (cid:117)(cid:117)(cid:116) R
[Equation]
n
[Equation]
(cid:20)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
( W t
[Equation] [Equation]
E
[Equation] [Equation]
( W t
[Equation]
Z )] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
6Ψ( W t − 1 ) (cid:21) = (cid:118)(cid:117)(cid:117)(cid:116) 2 R
[Equation]
n
[Equation]
(cid:88)
[Equation]
(cid:104)
[Equation] [Equation]
(cid:101)
[Equation] [Equation]
(cid:105) , where Eq.
14 and 17 are by Jensen’s inequality (i.e. the concavity of logarithm and the convexity of squared norm), Eq. 15 is by log( x + 1) ≤ x , and Eq.
16 and 18 are by ||
x
+ y || 2 ≤ 2 || x
||
+ 2 ||
y
|| 2 .
This completes the proof.
C. P ROOF OF L EMMA
Proof.
We ﬁrst notice that
[Equation]
where N ∼ N (0 , I d ) .
Then let X = S , Y =
(cid:102) W t − 1 and Ω =
E W
[Equation]
1 ,Z [ ∇
[Equation]
( W t − 1 , Z )] , by applying Lemma 4, we have
I ( − λ
[Equation]
where the second inequality is by Jensen’s inequality.
This completes the proof.
C. P ROOF OF T HEOREM
Proof.
Applying Lemma 6 and Lemma 5 and putting everything together, we have
[Equation]
Next
[Equation]
to handle the mismatch between the outputs of perturbed SGD and SGD
[Equation]
we apply Taylor expansion around ∆ T = (cid:126) 0 ,
[Equation]
where Eq.19 is by the zero mean of the perturbation, Eq.20 is by the independence of the coordinates of ∆ T , (cid:104)· , ·(cid:105) denotes the inner product of two matrices, diag( A ) is the diagonal matrix with element
A and Tr( · ) is the trace of a matrix.
Under the condition
E W T ,S (cid:48)
[ γ ( W T , S (cid:48) )
] ≥ 0 , we now bound gen( µ, P (cid:102) W T | S ) instead of its absolute value, | gen( µ, P (cid:102) W T | S ) | .
With the inequality log( x + 1) ≤ x , the following is straightforward, gen
[Equation]
µ, P (cid:102) W T | S ) ≤
[Equation]
Since every choice of σ gives a valid generalization bound.
The optimal bound can be found by simply utilizing the fact A/σ + σ 2 B ≥ 3( A/ 2) 2 / 3 B 1 / 3 for any positive A and B , where the equality is achieved by the optimal σ .
Finally, rearranging the terms will complete the proof.
C. C OROLLARY
Corollary 2.
If the loss function is differentiable and β -smooth with respect to w , then,
[Equation]
Proof.
Recall the smoothness implies f ( v ) ≤ f ( w )+ (cid:104)∇ f ( w ) , v − w (cid:105) + β 2 || v − w || 2 for all v and w .
By the triangle inequality, we have
[Equation]
Thus, we can see that | E [ L µ ( W T ) − L µ ( W T + ∆ T )] | + | E
[ L S ( W T + ∆ T ) − L S ( W T )]
| ≤ βd (cid:80) Tt =
σ 2 t .
This completes the proof.
Remark 3.
In Corollary 2, we note that the dependence of d in the bound results from the spherical Gaussian noise used in the construction of the weight process (cid:102) W T .
It is possible to replace the spherical Gaussian with a Gaussian noise having a non-diagonal covariance that reﬂects the geometry of the loss landscape.
With this replacement, the dimension d in the ﬂatness term will be replaced by the trace of (cid:80)
Tt =
1 κ t , where κ t is the covariance matrix of the noise added at step t .
Please refer to Neu et al.
(2021) for a similar development.
D A PPLICATION IN N EURAL N ETWORKS : P ROOF OF T HEOREMS 3 ANDD. L EMMA
Lemma 7.
Under the same conditions of Theorem 2, the generalization error of SGD is upper bounded by
[Equation]
Proof.
The proof of this lemma follow the same steps in the proof of Theorem 2 except that we require a modiﬁed version of Lemma 4 here.
Specially, we let Ω = 0 in Lemma 4.
The remaining steps are the same in the proof of Theoremand should be straightforward.
Remark 4.
The bound in Lemma 7 is weaker than Eq.
3 in Theorem 2 as the centralized expected vector norm should be smaller than the original expected vector norm.
D. P ROOF OF T HEOREM
Proof.
By Cauchy–Schwarz inequality, we have
[Equation]
Notice that ∇ w (cid:96) ( W, Z ) =
( W T X − Y ) X .
Let ˆ Y =
f ( W, X ) .
Then, ||∇ w (cid:96) ( W t − 1 , Z ) || 22 =
||
( W Tt − 1 X − Y )
X
|| 22 ≤
|| W
Tt − 1 X −
[Equation]
||
[Equation]
( ˆ Y
[Equation]
Y )
[Equation]
=
[Equation]
(cid:96) ( W t −
[Equation]
Z ) (22)
For the Hessian matrix, it’s easy to see
that Tr(H W T ( Z )) = 1 .
Plugging everything into Lemma 7, we have
[Equation]
≤(cid:32) T (cid:88) t =
R 2 λ 2 t T
[Equation]
n
E
[Equation]
t − 1 ,
[Equation]
[ (cid:96) ( W t − 1 , Z )] (cid:33)
This completes the proof.
D. P ROOF OF T HEOREM
Proof.
Since ∇ w r (cid:96) ( W, Z i ) = 1 √ m A r
f
( W, X i ) − Y i )
X
i
I r,i , where I r,
i =
I { W Tr X i ≥ 0 } , we have ||∇
w
[Equation]
W t − 1 , Z
[Equation] [Equation]
In addition, we notice that Tr(H W T ( Z ))
= 1 m (cid:80) m r =
I r,i,T .
Plugging everything into Lemma 7, we have
[Equation]
This completes the proof.
E E XPERIMENT D ETAILS E. A RCHITECTURES AND H YPERPARAMETERS
In Section 5, MLP has one hidden layer with 512 hidden units, and AlexNet has ﬁve convolution layers (conv.
3 × 3 (64 ﬁlters) → max-pool 3 × 3 → conv.
5 × 5 (192 ﬁlters)
→ max-pool 3 × 3 → conv. 3 × 3 (384 ﬁlters) → conv. 3 × 3 (256 ﬁlters) → conv. 3 × 3 (256 ﬁlters) → max-pool 3 × 3 ) followed by two fully connected layers both with 4096 units and a 10-way linear layer as the output layer.
All of the convolution layers and the fully connected layers use standard rectiﬁed linear activation functions (ReLU).
The ﬁxed learning rates used for MLP and AlexNet are 0 .
01 and 0 .
001 , respectively.
The batch size is set to 60 .
For the corrupted label experiment, we train the models until the models achieve 100% training accuracy.
For other cases, we train the neural networks until the training loss converges (e.g., < 0 . 0001 ).
Other settings are either described in Section 5 or apparent in the ﬁgures.
Standard techniques such as weight decay and batch normalization are not used.
To choose the variance proxy R , we ﬁrst collected all the per-instance losses (cid:96) ( W t − 1 , Z i ) that were observed during training, then we let R =
( max i,t (cid:96) ( W t − 1 , Z i ) − min i,t (cid:96) ( W t − 1 , Z i )) / 2 in our experiments.
In Section 6, we compare GMP with other advanced regularization methods.
The results of other methods are reported directly from Zheng et al.
, and we now give their hyperparameter settings here for completeness.
For Dropout, 10% of neurons are randomly selected to be deactivated in each layer.
For label smoothing, the coefﬁcient is 0 .
For ﬂooding, the level is set to 0 .
For
MixUp, we lineally combine random pairs of training data where the coefﬁcient is drawn from
Beta(1 , 1) .
For adversarial training, the perturbation size is 1 for each pixel and we take one step to generate adversarial examples.
For AMP, the number of inner iteration is 1 , and the L 2 norm ball radius values are 0 .
5 for PreActResNet18 and 0 .
1 for VGG16, respectively.
The implementation in this paper is on PyTorch, and all the experiments are carried out on NVIDIA
Tesla V100 GPUs (32 GB).
E. M ORE E XPERIMENTS ON THE C OMPARISON B ETWEEN T HEOREM 2 AND L EMMA
[Figure]
In Figure 1, we compare
[Equation]
and
[Equation]
] for two different values of σ .
Notice that we indeed use a weaker version of the trajectory term in
Theorem 2, and the same constants like R , λ t and σ t in the two bounds are ignored here.
we choose the full dataset of MNIST and CIFAR10 to train the models with the ﬁxed initialization.
We also ﬁx the sampling of batches, making the training completely deterministic.
To estimate the local gradient sensitivity term, we randomly sample 20 Gaussian noises from N (0 , (cid:80) t − 1 τ =1 σ τ ) to perturb the model parameters and compute the average perturbed gradient.
With this single draw and deterministic setting, the estimated (cid:101) V t and V t are the same, so the bound of Eq.
2 in Theorem 2 should be smaller than the bound in Lemma 2, as shown in Figure 1.
We also provide the experiments under the stochastic setting.
Specially, We randomly choose 1 /of the MNIST data and train the MLP model with a ﬁxed learning rate and batch size.
To estimate E
W t − 1 ,Z [ ∇ (cid:96) ( W t − 1 , Z )] , we conduct 20 runs with different random seeds and save W after every iteration.
In Figure 5 (d), we can see that at the beginning of training, our gradient dispersion V t is much larger than (cid:101) V t in Neu et al.
(2021), but
in the later training phase, the magnitude of gap between these two terms is very small.
This is because the gradient norm will become small when
W is near local minima.
In Figure 5 (a-c), we compare
[Equation]
and
[Equation] [Equation]
under three different settings of σ .
When σ is small (e.g., σ = 1 e − 6 ), the local gradient sensitivity term will become small, but the factor 1 /σ 2 will be very large, making the gap between log( x + 1) and x be extremely large, as shown in Figure 5 (c).
In this case, the improvement upon Neu et al.
(2021) is signiﬁcant.
When σ is large (e.g., σ = 1 e − 2 ), at the beginning of training, the trajectory term in Lemma 2 will be smaller than the trajectory term in Theorem 2 since our V t is relatively large.
However, since Ψ has a accumulative variance, and
V t and (cid:101) V t become closer soon or later, the bound in Lemma 2 will be greater than the bound in
Theorem 2 in the later training phase.
E. L EARNING R ATE AND B ATCH S IZE .
The learning rate and batch size have some impact on Eq.
3 in Theorem 2.
We now investigate this by performing experiments with varying learning rates and batch sizes.
In our experiments, the model is continuously updated until the average training loss drops below 0.0001.
We separate trajectory and ﬂatness terms of the bound and plot them in Figure 6.
[Figure] [Figure]
A key observation in Figure 6 is that the learning rate impacts the trajectory term and the ﬂatness term in opposite ways, as seen, for example, in (a) and (b), where the two set of curves swap their orders in the two ﬁgures.
On the other hand, the batch size also impacts the two terms in opposite ways, as seen in (a) and (b) where curves decrease in (a) but increase in (b).
This makes the generalization bound, i.e., the sum of the two terms, have a rather complex relationship with the settings of learning rate and batch size.
This relationship is further complicated by the fact that a small learning rate requires a longer training time, or a larger number T of training iterations, which increases the number that are summed over in the trajectory term.
Nonetheless, we do observe that a smaller batch size gives a lower value of the ﬂatness term ((b) and (d)), conﬁrming the previous wisdom that small batch sizes enable the neural network to ﬁnd a ﬂat minima (Keskar et al., 2017).
[Figure]
[Figure]
E. I MPLEMENTATION OF T HEOREM 3 AND T HEOREM
We let d 0 = 200 and use a two-layer ReLU network with hidden units 10000 to generate Y .
Moreover, we apply a tanh function to the output of this network so that | Y | ≤ 1 .
The input data
X ∼ N (0 , I)
and we normalize X before training so that ||
X
||
= 1 .
For the training phase, we choose m = 100 , let data size be 20 , 000 , batch size be 100 and learning rate be 0 .
Figure 7 and
Figure 8 compare the empirical generalization gap with the bound in Theorem 3 and that in Theorem 4, respectively.
E. A LGORITHM OF D YNAMIC G RADIENT C LIPPING AND A DDITIONAL R ESULTS
The dynamic gradient clipping algorithm is described in Algorithm 1.
For both MLP and AlexNet, we let α = 0 .
The start step for clipping, T c , is also an important hyperparameter.
However, it can be removed by detecting the evolution of the average gradient norm for each epoch.
Speciﬁcally, whenever the average gradient norm of epoch j is larger than the average gradient norm of epoch j − 1 , the clipping operation begins.
From Figure 4 and Figure 9 we can see that dynamic gradient clipping effectively alleviates overﬁtting by conspicuously slowing down the transition of training to the memorization regime, without changing the convergence speed of testing accuracy.
Unfortunately, the current design of the dynamic gradient clipping algorithm does not provide a signiﬁcant improvement for models trained on a clean dataset (without label noise).
Designing better regularization algorithms may require understanding the epoch-wise double descent curve of gradient dispersion where the model is trained on a clean dataset.
E. D ISCUSSION ON G RADIENT D ISPERSION OF M ODELS T RAINED ON C LEAN D ATASETS
In the case of no noise injected, Figure 3a shows that the model with good generalization property has a exponentially-decaying gradient dispersion.
This is consistent with our discussion of Lemma 5 in Section 3, that is, small I ( G t + N t ; S | (cid:102) W t − 1 ) indicates good generalization.
Notably, gradient dispersion of AlexNet trained on the real CIFAR10 data still has a epoch-wise double descent curve.
Require: Training set S , Batch size b , Loss function (cid:96) , Initial model parameter w 0 , Learning rate λ , Initial minimum gradient norm G , Number of iterations T , Clipping parameter α , Clipping step T c 1: for t ← 1 to T do 2:
Sample B = { z i } bi =1 from training set S 3:
Compute gradient: g B ← (cid:80) bi =1 ∇ w (cid:96) ( w t − 1 , z i ) /b 4: if t > T c then 5: if || g B || 2 > G then 6: g B ← α · G · g B / || g B ||7: else 8:
G ← || g B ||9: end if 10: end if 11:
Update parameter: w t ← w t − 1 − λ · g B 12: end for
[Figure] The difference between Figure 3e with Figure 3f-3h is that the testing accuracy does not decrease in the second phase/memorization regime for AlexNet trained on the true CIFAR10 data.
Loosely speaking, we conjugate that memorizing random labels will hurt the performance on unseen clean data but memorizing clean (or true) labels will not.
This may explain why dynamic gradient clipping or preventing the training entering the memorization regime cannot improve the performance on a clean dataset.
E. A LGORITHM OF G AUSSIAN M ODEL P ERTURBATION AND A DDITIONAL R ESULTS
Require: Training set S , Batch size b , Loss function (cid:96) , Initial model parameter w 0 , Learning rate λ , Number of noise k , Standard deviation of Gaussian distribution σ , Lagrange multiplier ρ while w t not converged do 2:
Update iteration: t ← t +
Sample B = { z i } bi =1 from training set S 4:
Sample ∆ j ∼ N (0 , σ 2 ) for j ∈ [ k ]
Compute gradient:
[Equation]
Update parameter: w t +1 ← w t − λ · g B end while
[Table]
[Table]
The GMP algorithm is given in Algorithm 2.
Table 1 and Table 2 show that our method is competitive to the state-of-the-art regularization techniques.
Speciﬁcally, our method achieves the best performance on SVHN for both models and on CIFAR-100 where VGG16 is employed.
Particularly, testing accuracy is improved by nealy 2% on the CIFAR-100 dataset with VGG16.
For other tasks, GMP is always able to achieve the top-3 performance.
In addition, we ﬁnd that increasing the number of sampled noises does not guarantee the improvement of testing accuracy and may even degrade the performance on some datasets (e.g., SVHN).
This hints that we can use small number of noises to reduce the running time without losing performance.
Moreover, we observe that GMP with k = 3 usually takes around 1 .
× that of ERM training time , which is affordable.
For PreActResNet18, the performance of GMP appears slighly inferior to the most recent record given by AMP (Zheng et al., 2021).
Noting that the key ingredient of AMP, “max-pooling” in the parameter space, greatly resembles regularization term in GMP, which may be seen as “averagepooling” in the same space.
One potential extension of GMP is to let the variance of the noise distribution be a function of the iteration step t .
In other words, using the time-dependent σ t instead of a constant σ .
We also consider the following regularized scheme, which is a absolute value version.
[Equation]
This scheme can still perform well as shown in Table 3.
In fact, it outperforms GMP 3 on CIFAR100.
This hints that it is possible to improve the performance by choosing other norm of L s ( w + ∆) − L s ( w ) .