On the Generalization of Models Trained with SGD: Information-Theoretic Bounds and Implications

1 INTRODUCTION . 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 ( Du et al. , 2018 ; Allen-Zhu 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 first 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 ) . Informationtheoretic 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 SGD-like algorithms . For example , Pensia et al . ( 2018 ) is the first 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 . ( 2019 ) ; Haghifam et al . ( 2020 ) ; Rodrı́guez-Gálvez et al . ( 2020 ) ; Wang et al . ( 2021b ) . Inspired by the work of Pensia et al . ( 2018 ) , Neu ( 2021 ) presents an information- theoretic analysis of the models trained with SGD . The analysis of Neu ( 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 flatness 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 flatter 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 . Our development builds upon that of Neu ( 2021 ) . Following the same construction of the auxiliary weight process in Neu ( 2021 ) , we present upper bounds of generalization error that improve upon Neu ( 2021 ) in two ways . This improvement is mainly due to removing an unnecessary term in the bound of Neu ( 2021 ) , which the author refers to as “ local gradient sensitivity ” , by invoking the HWI inequality ( Raginsky & Sason , 2018 ) ( i.e. , Lemma 3 of this paper ) and Lemma 4 . Additionally , the improvement also benefits from the fact that removing “ local gradient sensitivity ” allows us to have a simple closed form of the optimal bound when the noise variance does not depend on time step . 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 flatness of the found solution ( “ the flatness term1 ” ) . Remarkably , the generalization bound given in Neu ( 2021 ) also contains a trajectory term and a flatness term , where the flatness term is identical to ours due to the same construction of the auxillary weight process . However , our trajectory term is much tighter than that in Neu ( 2021 ) . This is not only because we remove the unnecessary term , and the constant factor of our trajectory term is also half of that in Neu ( 2021 ) . Figure 1 shows an experimental comparison of the trajectory term in our bound ( Theorem 1 ) with that in the bound in Neu ( 2021 ) ( re-stated as Lemma 2 in this paper for two neural network models . The trajectory terms are compared for two different settings of σ , which is the variance parameter of the noise used in the construction of the auxiliary weight process . After the the initial preprint of Neu ( 2021 ) , a revised version of the paper Neu et al . ( 2021 ) is presented , in which the bound is revised by fixing a coefficient error in the bound of Neu ( 2021 ) and adopts a slightly different form . It is particularly note-worthy that the development of the bound in Neu et al . ( 2021 ) also makes uses of the HWI inequality ( or Lemma 3 ) , but the unnecessary local gradient sensitivity term remains in the bound . This is because the development of Neu et al . ( 2021 ) requires constructing a ghost auxiliary weight process , which is completely avoided by using Lemma 4 in our paper . We will elaborate on this in an appropriate context . 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 not only makes the bound much tighter , it also makes it much easier to evaluate . In this paper , we empirically validate the derived bounds , and provide various insights pertaining to the generalization behavior of models trained with SGD . For example , a particular interesting observation from our experiments is that a key quantity arising in the trajectory term of the bounds , which we refer to as gradient dispersion2 , reveals a double descent phenomenon with respect to training epochs . Most intriguingly , 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 . Specifically , we propose a new training scheme , referred to as Gaussian model perturbation ( GMP ) , aiming at reducing the flatness term of the bounds . This scheme effectively applies a local “ average pooling ” 1We note this term is correlated with flatness rather than precisely measures the flatness . 2The quantity is often referred to as gradient variance in the literature , but we prefer “ dispersion ” to “ vari- ance ” so as to better comply with the mathematical conventions and avoid possible confusion . 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 ( 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 “ flatness ” 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 . Gradient clipping is a common technique for preventing gradient exploding ( see , e.g. , Merity et al . ( 2018 ) ; Peters et al . ( 2018 ) ) . 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 . 2 PRELIMINARIES . 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. LetW ⊆ Rd be the space of hypotheses . Suppose that a training sample S = ( Z1 , Z2 , . . . , Zn ) 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 PW |S mapping Zn toW . Let ` : W ×Z → R+ be a loss function , where ` ( w , z ) measures the “ unfitness ” or “ error ” of any z ∈ Z with respect to a hypothesis w ∈ W . The population risk , for any w ∈ W , is defined as Lµ ( w ) , EZ∼µ [ ` ( w , Z ) ] . The goal of learning is to find 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 , defined as LS ( w ) , 1 n n∑ i=1 ` ( w , Zi ) . The expected generalization error of the learning algorithm A is then defined as gen ( µ , PW |S ) , EW , S [ Lµ ( W ) − LS ( W ) ] , where the expectation is taken over the joint distribution of ( S , W ) ( i.e. , µn ⊗ PW |S ) . Throughout this paper , we take ` as a continuous function ( adopting the usual notion “ surrogate loss ” ( Shalev-Shwartz & Ben-David , 2014 ) . ) . Additionally , we assume that ` is differentiable almost everywhere with respect to both w and z . Furthermore we assume that ` ( w , Z ) is R-subgaussian3 for 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 ` ( w , Z ) is R-subgaussian for any w ∈ W . The generalization error of A is bounded by |gen ( µ , PW |S ) | ≤ √ 2R2 n I ( W ; S ) , 3A random variable X is R-subgaussian if for any ρ , logE exp ( ρ ( X − EX ) ) ≤ ρ2R2/2 . 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 . Specifically , let Bt denote the batch used for the tth update . Define g ( w , Bt ) , 1 b ∑ z∈Bt ∇w ` ( w , z ) , namely , g ( w , Bt ) is the average gradient computed for the batch Bt with respect to parameter w. The rule for the tth parameter update is then Wt , Wt−1 − λtg ( Wt−1 , Bt ) , where λt is the learning rate at the step t. The initial parameter setting W0 is assumed to be drawn from the zero-mean spherical GaussianN ( 0 , σ20Id ) with variance σ20 in each dimension . We will assume that the SGD algorithm stops after T updates and outputs WT as the learned model parameter . Given the training sample S , let ξ govern the randomness in the sequence ( B1 , B2 , . . . , BT ) of batches . For the simplicity of notion , we will fix the configuration of ξ . That is , we will assume a fixed “ batching trajectory ” , or a fixed way to shuffle 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 { W̃t } . Let σ1 , σ2 , . . . , σT be a sequence of positive real numbers . Define W̃0 , W0 , and W̃t , W̃t−1 − λtg ( Wt−1 , Bt ) +Nt , for t > 0 , whereNt ∼ N ( 0 , σ2t Id ) is a Gaussian noise . The relationship between this auxiliary weight process { W̃t } and the weight process { Wt } in SGD is shown in the Bayesian network below . N1 N2 · · · NT−1 NT ↓ ↓ ↓ ↓ W̃0 → W̃1 → W̃2 → · · · → W̃T−1 → W̃T ‖ ↗ ↗ ↗ ↗ W0 → W1 → W2 → · · · → WT−1 → WT Let ∆t = ∑t τ=1Nτ . Noting that the weight updates in { W̃t } uses the same gradient signal as that used in { Wt } ( which depends on Wt−1 not W̃t−1 ) , it is immediate that W̃t = Wt + ∆t . Note that this auxiliary process follows the same construction as Neu ( 2021 ) , which we will use to study the generalization error of SGD . To that end , let ’ s define gradient dispersion by Vt ( w ) , E [ ||g ( w , Bt ) − E [ ∇w ` ( W , Z ) ] ||22 ] , where the inside expectation is taken over ( W , Z ) ∼ µ⊗ PW |Z . For a given sample s ∈ Zn , define γ ( w , s ) , E [ Ls ( w + ∆T ) − Ls ( w ) ] , where the expectation is taken over ∆T and Ls ( w ) is the empirical risk of s at parameter w. In the remainder of the paper , let S′ 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 |gen ( µ , PWT |S ) | ≤ √√√√4R2 n T∑ t=1 λ2t σ2t E [ Ψ ( Wt−1 ) + Ṽt ( Wt−1 ) ] + |E [ γ ( WT , S ) − γ ( WT , S′ ) ] | , where Ψ ( wt−1 ) , E [ ||E [ ∇w ( wt−1 , Z ) ] − E [ ∇w ( wt−1 + ζ , Z ) ] ||22 ] , ζ ∼ N ( 0 , 2 ∑t−1 i=1 σ 2 i Id ) and Ṽt ( w ) , E [ ||g ( w , Bt ) − E [ ∇w ` ( w , Z ) ] ||22 ] . The term Ψ ( wt−1 ) in the bound is referred to as “ local gradient sensitivity ” in Neu et al . ( 2021 ) .

Summary: This paper proposed a new generalization bound which is based on the recent construction by Neu (2021). The authors made further assumptions on the randomness of the batch sequence, and improved the results by Neu by achieving tighter bounds by tuning the trajectory term. It is verified empirically, and the flatness term gives a regularization technique for better classification on VGG16 on MNIST.