Year: 2021

Abstract:
This paper follows up on a recent work of Neu (2021) and presents new and tighter 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 based on 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES
text: 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).

heading: 3 NEW GENERALIZATION BOUNDS FOR SGD
text: The improvement offered by our new bound upon the bound in Neu (2021) or Lemma 2 is the removal of the local sensitivity term Ψ(wt−1). This requires invoking a vector version of the HWI inequality (Raginsky & Sason, 2018, Lemma 3.4.2), which we first state.
Lemma 3. Let X and Y be two random vectors in Rd, and let N ∼ N (0, Id) be independent of (X,Y ). Then, for every t > 0, DKL(PX+√tN ||PY+√tN ) ≤ 1 2tE [ ||X − Y ||22 ] , where DKL is the KL divergence.
Note that Neu et al. (2021) also makes uses of Lemma 3 in their revision. However, the development in Neu et al. (2021) requires constructing a ghost auxiliary weight process in which an independent noise perturbation ∆′t is introduced. Due to the mismatch between ∆t and ∆ ′ t, the local gradient sensitivity term Ψ(Wt−1) appears in their bound. In this paper we waive this ∆′t by using the following lemma.
Lemma 4. Let X ,Y and ∆ be random variables which are all independent of N ∼ N (0, I). Let Z = Y + ∆, then for any σ and any function f , we have
I(f(Z,X) + σN ;X|Y ) ≤ 1 2σ2
E [ ||f(Z,X)− E [f(Z,X)]||2 ] The proof of this lemma is inspired by the techniques used in Wang et al. (2021a) and Guo et al. (2005). Exploiting Lemma 4 gives us a crisp way to have the following result, where Ψ(Wt−1) will not appear.
Lemma 5. Let Gt = −λtg(Wt−1, Bt). Then, I(Gt +Nt;S|W̃t−1) ≤ λ 2 t
2σ2t E [Vt(Wt−1)].
In this lemma, the mutual information I(Gt+Nt;S|W̃t−1) roughly indicates the degree by which the SGD’s updating signalGt (smoothed with noise) depends on the training sample S, whenBt 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 overfit 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 configuration. 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 1. The generalization error of SGD is upper bounded by
|gen(µ, PWT |S)| ≤ √√√√R2 n T∑ t=1 λ2t σ2t E [Vt(Wt−1)] + |E [γ(WT , S)− γ(WT , S′)]| . (1)
Further, assume Lµ(wT ) ≤ E∆ [Lµ(wT + ∆T )], ` is twice differentiable, and σ2t is independent of t. Denote by HWT the Hessian of the loss with respect to WT and let Tr(·) denote trace. Then
gen(µ, PWT |S) ≤ 3
2 ( T∑ t=1 R2λ2tT n E [Vt(Wt−1)]E [Tr (HWT (Z))] ) 1 3
(2)
Remark 1. Removing Ψ(Wt−1) will make the bound much tighter as shown in Figure 1. This remarkable improvement should come at no surprise, since Ψ(Wt−1) monotonically increases with training epochs (noting that Ψ(Wt−1) has the cumulative variance 2 ∑t−1 i=1 σ 2 i Id ) while Vt(Wt−1) appears decreasing (see Figure 4a). 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/σ2t will become too large, make the bound loose. Additionally, the constant factor of the trajectory term in Eq. 1 is half of that in Lemma 2.
Remark 2. The condition Lµ(wT ) ≤ E∆T [Lµ(wT + ∆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.1, the first term captures the impact of the training trajectory (“trajectory term”), and the second term captures the impact of the final solution, which in fact measures the flatness for the loss landscape at the found solution (“flatness term”). It is worth noting that using both the gradient dispersion (or the trace of the gradient’s covariance matrix), and the trace of Hessian to characterize generalization is not new. For example, the Takeuchi information criterion used by Thomas et al. (2020) consists of Hessian and gradient dispersion (referred to as “curvature” and “noise” by the authors), which is shown relevant to estimating the generalization gap.
The learning rate in SGD has explicitly appeared in the trajectory term of Eq.2 in Theorem 1. 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 benefit 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.2). As seen in the experiments, increasing the learning rate impacts the trajectory term and the flatness 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.2, have a rather complex relationship with the settings of learning rate and batch size.
Eq.2 follows from Eq.1 by minimizing the bound over σ. The bound in Eq.2 is thus independent of a choice of σt’s and can be computed easily and efficiently.
To summarize, we remark that these bounds suggest that in order for the model to generalize well, both the trajectory term and the flatness term need to be small — the former involves the interaction of the learning rate and batch size with the gradient dispersion along the training trajectory, whereas the latter depends on the flatness of the empirical risk surface at the found solution.

heading: 4 APPLICATION IN LINEAR NETWORKS AND TWO-LAYER RELU NETWORKS.
text: We now apply Theorem 1 to two neural network models in a regression setting. Let Z = [X,Y ] ∈ Rd0+1. Assume ||X|| = 1. We use SGD to train a model f(W, ·) : Rd0 → R and define the loss as `(W,Z) = 1/2(Y − f(W,X))2. The bounds below are both in the form of the product of the flatness term and the trajectory term.
Theorem 2 (Linear Networks). Let f(W,X) = WTX . Then,
gen(µ, PWT |S) ≤ 3 ( T∑ t=1 R2λ2tT 4n E [`(Wt−1, Z)] ) 1 3 .
Theorem 3 (Two-Layer ReLU Networks). Following Arora et al. (2019), consider f(W,X) = 1√ m ∑m r=1ArReLU(W T r X) wherem is the width of the neural network,Ar ∼ unif({+1,−1}) and ReLU(·) is the ReLU activation. We only train the first layer parametersW = [W1,W2, . . . ,Wm] ∈ Rd0×m and fix the second layer parameters A = [A1, A2, . . . , Am] ∈ Rm during training. Then,
gen(µ, PWT |S) ≤ 3 ( m∑ r=1 E [ Ir,i,T m ] T∑ t=1 R2λ2tT 4n E [ m∑ r=1 Ir,i,t m `(Wt−1, Z) ]) 1 3 ,
where Ir,i,t = I{WTt−1,rXi ≥ 0} and I is the indicator function. Remark 3. Compared with Theorem 2, 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 3. Specifically, the weights of the deactivated neurons do not contribute to the bound of Theorem 3, 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 first theoretical justification in this regard.
Figure 2 shows that this bound can characterize the dynamics of generalization gap of the ReLU network. We defer the implementation details including the linear network’s plot to Appendix.
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.
5 EXPERIMENTAL STUDY
Bound Verification We first verify our bound of Eq.2 in Theorem 1 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 fix the weight initialization and set σt and λt to be constants σ and λ, respectively. To compute ∑T t=1 E [Vt(Wt−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 [HWT (Z)]), we use the PyHessian library (Yao et al., 2020) to compute the Hessian. To choose the variance proxy R, we first collected all the per-instance losses `(Wt−1, Zi) that were observed during training, then we let R = (maxi,t`(Wt−1, Zi)−mini,t`(Wt−1, Zi))/2 in our experiments. We perform experiments with varying network width and varying levels of label noise. Specifically, noise level refers to the setting
where we replace the labels of 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 3.
In Figure 3, 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 3 (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 fixing 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 4 the evolution of the (empirical) gradient dispersion V̂t(wt−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 4 (h)). In the first 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 overfitting 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 overfits 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.g. MNIST) and cleaner datasets , the first 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.
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. Specifically, whenever the current gradient norm is larger than the gradient normK steps earlier, or ||g(Wt, Bt)||2 > ||g(Wt−K , Bt−K)||2 (i.e., the model is expected to have entered the “memorization” regime), we reduce the norm of the current gradient g(Wt, Bt) to α fraction of ||g(Wt−K , Bt−K)||2, for some prescribed value α < 1.
The effectiveness of this scheme is best demonstrated when the labels contain noise. As shown in Figure 5 (and Figure 8 in Appendix E.4), dynamic gradient clipping significantly 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.

heading: 6 A PRACTICAL IMPLICATION: GAUSSIAN MODEL PERTURBATION
text: The appearance of the flatness 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 flat, or insensitive to a small perturbation of w∗. This naturally motivates a training scheme using the following regularized loss:
min w Ls(w) + ρ E ∆∼N (0,σ2Id) [Ls(w + ∆)− Ls(w)],
where ρ is a hyper-parameter. Replacing the expectation above with its stochastic approximation using k realizations of ∆ gives rise to the following optimization problem.
min w
1
b ∑ z∈B
( (1− ρ)`(w, z) + ρ1
k k∑ i=1 (`(w + δi, z))
) .
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 + 1 forward passes on parallel processors further reduces the computation load.
Method SVHN CIFAR-10 CIFAR-100
ERM 96.86±0.060 93.68±0.193 72.16±0.297 Dropout 97.04±0.049 93.78±0.147 72.28±0.337 L. S. 96.93±0.070 93.71±0.158 72.51±0.179 Flooding 96.85±0.085 93.74±0.145 72.07±0.271 MixUp 96.91±0.057 94.52±0.112 73.19±0.254 Adv. Tr. 97.06±0.091 93.51±0.130 70.88±0.145 AMP 97.27±0.015 94.35±0.147 74.40±0.168 GMP3 97.18±0.057 94.33±0.094 74.45±0.256 GMP10 97.09±0.068 94.45±0.158 75.09±0.285
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.03. The value of k is chosen as 3 and 10 respectively (referred to as GMP3 and GMP10).
The performances of all compared schemes are given in Table 1 (and the results of PreActResNet18 are shown in Appendix E.6). For the compared regularization schemes except GMP, we directly report their performances as given in Zheng et al. (2021). 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.

heading: 7 CONCLUSION AND OUTLOOK
text: This paper presents new information theoretic generalization bounds for models, such as linear networks and two-layer ReLU neural networks, trained with SGD, improving upon the results of Neu (2021). 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 still many directions for improving the bounds. First, it is possible to extend the theoretical results in this work by using the techniques introduced by other information-theoretic bounds such as Haghifam et al. (2020) and Wang et al. (2021a). Additionally, it’s worth mentioning that one may obtain better results by further developing the relationship between the trajectory term and the flatness term from an optimization perspective. Some recent works based on backward error analysis (Barrett & Dherin, 2021; Smith et al., 2021) may provide analytic tools to investigate this.

heading: Appendix

heading: Table of Contents

heading: A HWI Inequality: Proof of Lemma 3 15

heading: B A Key Lemma: Lemma 4 16
text: C Improved Mutual Information Bounds: Proof of Theorem 1 16 C.1 Lemma 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 C.2 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 C.3 Corollary 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18

heading: D Application in Neural Networks: Proof of Theorems 2 and 3 18
text: D.1 Lemma 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 D.2 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19 D.3 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19

heading: E Experiment Details 20
text: E.1 Architectures and Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . 20 E.2 Learning Rate and Batch Size. . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 E.3 Implementation Details of Figure 2 . . . . . . . . . . . . . . . . . . . . . . . . 21 E.4 Algorithm of Dynamic Gradient Clipping and Additional Results . . . . . . . . 21 E.5 Discussion on Gradient Dispersion of Models Trained on Clean Datasets . . . . 22 E.6 Algorithm of Gaussian Model Perturbation and Additional Results . . . . . . . . 22

heading: A HWI INEQUALITY: PROOF OF LEMMA 3
text: Proof. 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.
DKL(PX+ √ tN ||PY+√tN ) ≤DKL(PX,Y,X+√tN ||PX,Y,Y+√tN ) (1)
= E X,Y
[ DKL(PX+ √ tN |X,Y ||PY+√tN |X,Y ) ] = E X,Y [DKL(N (X, t)||N (Y, t))|X,Y ] (2)
= 1
2t E X,Y
[ ||X − Y ||2 ] , (3)
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 a special case of the equality
N (µ1,Σ1)||N (µ2,Σ2)) = 1
2 [ log |Σ2| |Σ1| − d+ Tr ( Σ−12 Σ1 ) + (µ2 − µ1)TΣ−12 (µ2 − µ1) ] .
Eq.3 holds for any joint distribution of (X,Y ), so by the definition of Wasserstein distance, we have
DKL(PX+ √ tN ||PY+√tN ) ≤
1 2t W22(PX , PY ). (4)
This completes the proof.

heading: B A KEY LEMMA: LEMMA 4
text: Proof. By the properties of mutual information,
I(f(Z,X) + σN ;X|Y ) = I(f(Z,X)− E [f(Z,X)] σ +N ;X|Y ),
where the expectation is taken over (Z,X) ∼ PZ,X .
Let g(Z,X) , f(Z,X)−E[f(Z,X)]σ . Then,
I(g(Z,X) +N ;X|Y ) = E Y,X
[ DKL ( Pg(Z,x)+N |X=x,Y=y||Pg(Z,X)+N |Y=y )] = E Y,X [ DKL ( Pg(Z,x)+N |X=x,Y=y||PN ) −DKL ( Pg(Z,X)+N |Y=y||PN
)] (5)
≤ E Y,X
[ DKL ( Pg(Z,x)+N |X=x,Y=y||PΩ+N )] (6)
≤1 2 E Y,X
[ W22 ( Pg(Z,x)|X=x,Y=y||PΩ )] (7)
= 1
2 E Y,X [ E ∆ [ ||g(y + ∆, x)||2|X = x, Y = y ]] = 1
2σ2 E Z,X
[ ||f(Z,X)− E
Z,X [f(Z,X)]||2 ] where Eq. 5 is the golden formula (see Theorem 3.3 in Polyanskiy & Wu (2019) ), Eq. 6 is by the non-negativity of relative entropy, and Ω is a random variable that equals to zero with probability one, Eq. 7 is due to Lemma 3.
C IMPROVED MUTUAL INFORMATION BOUNDS: PROOF OF THEOREM 1
In this section, we provide a proof for Theorem 1, 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.

heading: C.1 LEMMA 5
text: Proof. We first notice that
I(Gt +Nt;S|W̃t−1) = I(−λtg(Wt−1, Bt) + σtN ;S|W̃t−1), where N ∼ N (0, Id). Then by applying Lemma 4, we have
I(−λtg(Wt−1, Bt) + σtN ;S|W̃t−1) ≤ λ2t 2σ2t
E [ ||g(Wt−1, Bt)− E [∇w`(Wt−1, Z)]||2 ] = λ2t 2σ2t E [Vt−1(Wt−1)]
This completes the proof.

heading: C.2 PROOF OF THEOREM 1
text: Proof. We first follow the similar decomposition of the expected generalization error and apply Lemma 1 as in Neu (2021),
|gen(µ, PWT |S)| = ∣∣∣∣gen(µ, PW̃T |S) + EWT ,∆T [Lµ(WT )− Lµ(W̃T )]+ EWT ,∆T ,S [LS(W̃T )− LS(WT )] ∣∣∣∣ ≤ √ 2R2
n I(W̃T ;S) + ∣∣∣∣ EWT ,S,S′ [γ(WT , S)− γ(WT , S′)] ∣∣∣∣ . (8)
The remaining task is to bound the mutual information I(W̃T ;S).
Then,
I(W̃T ;S) =I ( W̃T−1 − λT g(WT−1, BT ) +NT ;S ) ≤I ( W̃T−1,−λT g(WT−1, BT ) +NT , ;S ) (9)
=I(W̃T−1;S) + I ( −λT g(WT−1, BT ) +NT ;S|W̃T−1 ) (10)
≤I(W̃T−2;S) + I ( −λT−1g(WT−2, BT−1) +NT−1;S|W̃T−2 ) + I ( −λT g(WT−1, BT ) +NT ;S|W̃T−1 ) (11)
...
≤I(W̃0;S) + T∑ t=1 I ( −λtg(Wt−1, Bt) +Nt;S|W̃t−1 ) , (12)
= T∑ t=1 I ( −λtg(Wt−1, Bt) +Nt;S|W̃t−1 ) , (13)
where Eq.9 is by I(f(X);Y ) ≤ I(X;Y ), Eq.10 is by the chain rule of mutual information, Eq.11 is by applying the similar procedure (namely, Eq.9-Eq.10) to I(W̃T−1;S), Eq. 12 is by doing these steps recursively and Eq.13 is due to the fact that W̃0 is independent of S (i.e. I(W̃0;S) = 0). Applying Lemma 5 and putting everything together, we have
|gen(µ, PWT |S)| ≤ √√√√R2 nb T∑ t=1 λ2t σ2t E [ Ṽt(Wt−1) ] + |E [γ(WT , S)− γ(WT , S′)]| .
Next, to handle the mismatch between the outputs of perturbed SGD and SGD, we apply Taylor expansion around ∆T = ~0,
E WT ,S,∆T
[LS(WT + ∆T )− LS(WT )] = 1
n n∑ i=1 E WT ,Zi,∆T [`(WT + ∆T , Zi)− `(WT , Zi)]
≈ E WT ,Z,∆T
[ 〈∇w`(WT , Z),∆T 〉+ 1
2 ∆TTHWT (Z)∆T ] = E WT ,Z,∆T [ 1 2 ∆TTHWT (Z)∆T ] (14)
= 1
2 〈 E WT ,Z [HWT (Z)], E ∆T
[ ∆T∆ T T ] 〉
= 1
2 〈 E WT ,Z [HWT (Z)],diag( T∑ t=1 σ2t )〉 (15)
=
∑T t=1 σ 2 t
2 Tr( E WT ,Z [HWT (Z)]),
where Eq.14 is by the zero mean of the perturbation, Eq.15 is by the independence of the coordinates of ∆T , 〈·, ·〉 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 WT ,S′ [γ(WT , S ′)] ≥ 0, we now bound gen(µ,P W̃T |S) instead of its absolute value, |gen(µ,P W̃T |S)|. The following is straightforward,
gen(µ,P W̃T |S) ≤
R
n √√√√ T∑ t=1 λ2t σ2t E Wt−1 [V(Wt−1)] + E WT ,S,S′ [γ(WT , S)− γ(WT , S′)]
≤R n √√√√ T∑ t=1 λ2t σ2t E Wt−1 [V(Wt−1)] + E WT ,S [γ(WT , S)]
= R
n √√√√ T∑ t=1 λ2t σ2t E Wt−1 [V(Wt−1)] + ∑T t=1 σ 2 t 2 Tr( E WT ,Z [HWT (Z)]).
Since every choice of σ gives a valid generalization bound. The optimal bound can be found by simply utilizing the factA/σ+σ2B ≥ 3(A/2)2/3B1/3 for any positiveA andB, where the equality is achieved by the optimal σ. Finally, rearranging the terms will complete the proof.

heading: C.3 COROLLARY 1
text: Corollary 1. If the loss function is differentiable and β-smooth with respect to w, then,
|gen(µ, PWT |S)| ≤ √√√√R2 n T∑ t=1 λ2t σ2t E [Vt(Wt−1)] + βd T∑ t=1 σ2t .
Proof. Recall the smoothness implies f(v) ≤ f(w) + 〈∇f(w),v−w〉+ β2 ||v−w|| 2 for all v and w. By the triangle inequality, we have
|E [Lµ(WT )− Lµ(WT + ∆T )] | ≤ |E [〈∇w`(WT , Z),∆T 〉] |+ β 2 E [ ||∆T ||2 ] = βd ∑T t=1 σ 2 t 2
Thus, we can see that |E [Lµ(WT )− Lµ(WT + ∆T )] | + |E [LS(WT + ∆T )− LS(WT )] | ≤ βd ∑T t=1 σ 2 t .
This completes the proof.
Remark 4. In Corollary 1, we note that the dependence of d in the bound results from the spherical Gaussian noise used in the construction of the weight process W̃T . It is possible to replace the spherical Gaussian with a Gaussian noise having a non-diagonal covariance that reflects the geometry of the loss landscape. With this replacement, the dimension d in the flatness term will be replaced by the trace of ∑T t=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.

heading: D APPLICATION IN NEURAL NETWORKS: PROOF OF THEOREMS 2 AND 3

heading: D.1 LEMMA 6
text: Lemma 6. Under the same conditions of Theorem 1, the generalization error of SGD is upper bounded by
gen(µ, PWT |S) ≤ 3
2 ( T∑ t=1 R2λ2tT n E [ ||g(Wt−1, Bt)||22 ] E [Tr (HWT (Z))] ) 1 3
Proof. The proof of this lemma follow the same steps in the proof of Theorem 1 except that we require a modified version of Lemma 4 here.
Specially, we utilize I(f(Z,X) + σN ;X|Y ) = I( f(Z,X)σ + N ;X|Y ) instead of I(f(Z,X) + σN ;X|Y ) = I( f(Z,X)−E[f(Z,X)]σ +N ;X|Y ), that is used in the proof of Lemma 4. The remaining steps should be straightforward.
Remark 5. Intuitively, the bound in Lemma 6 is looser than Eq. 2 in Theorem 1 as the centralized expected vector norm should be smaller than the original expected vector norm.

heading: D.2 PROOF OF THEOREM 2
text: Proof. By Cauchy–Schwarz inequality, we have
E [ ||g(Wt−1, Bt)||22 ] = E ∣∣∣∣∣ ∣∣∣∣∣1b ∑
Zi∈Bt
∇w`(Wt−1, Zi) ∣∣∣∣∣ ∣∣∣∣∣ 2
2  ≤ E [||∇w`(Wt−1, Z)||22] (16) Notice that∇w`(W,Z) = (WTX − Y )X . Let Ŷ = f(W,X). Then,
||∇w`(Wt−1, Z)||22 = ||(WTt−1X − Y )X||22 ≤ ||WTt−1X − Y ||22 = (Ŷ − Y )2 = 2`(Wt−1, Z) (17)
For the Hessian matrix, it’s easy to see that Tr(HWT (Z)) = 1.
Plugging everything into Lemma 6, we have
gen(µ, PWT |S) ≤ 3 ( T∑ t=1 R2λ2tT 4n E Wt−1,Z [`(Wt−1, Z)] ) 1 3
This completes the proof.

heading: D.3 PROOF OF THEOREM 3
text: Proof. Since∇wr`(W,Zi) = 1√mAr(f(W,Xi)−Yi)XiIr,i, where Ir,i = I{W T r Xi ≥ 0}, we have
||∇w`(Wt−1, Z)||22 (18)
= m∑ r=1 || 1√ m Ar(f(Wt−1, Xi)− Yi)XiIr,i||22
≤ m∑ r=1 || 1√ m Ar(f(Wt−1, Xi)− Yi)Ir,i||22
= 1
m m∑ r=1 Ir,i||(f(Wt−1, Xi)− Yi)||22
= 1
m m∑ r=1 2Ir,i`(Wt−1, Zi) (19)
In addition, we notice that Tr(HWT (Z)) = 1 m ∑m r=1 Ir,i,T . Plugging everything into Lemma 6, we have
gen(µ, PWT |S) ≤ 3 ( T∑ t=1 R2λ2tT 4nm E Wt−1,Zi [ m∑ r=1 Ir,i,t`(Wt−1, Zi) ] E WT ,Zi [ 1 m m∑ r=1 Ir,i,T ]) 1 3
This completes the proof.

heading: E EXPERIMENT DETAILS

heading: E.1 ARCHITECTURES AND HYPERPARAMETERS
text: In Section 5, MLP has one hidden layer with 512 hidden units, and AlexNet has five convolution layers (conv. 3× 3 (64 filters)→ max-pool 3× 3→ conv. 5× 5 (192 filters)→ max-pool 3× 3→ conv. 3 × 3 (384 filters) → conv. 3 × 3 (256 filters) → conv. 3 × 3 (256 filters) → 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 rectified linear activation functions (ReLU).
The fixed 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 figures. Standard techniques such as weight decay and batch normalization are not used.
In Section 6, we compare GMP with other advanced regularization methods. The results of other methods are reported directly from Zheng et al. (2021), 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 coefficient is 0.2. For flooding, the level is set to 0.02. For MixUp, we lineally combine random pairs of training data where the coefficient 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 L2 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).

heading: E.2 LEARNING RATE AND BATCH SIZE.
text: The learning rate and batch size have some direct impact on Eq. 2 in Theorem 1. 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 flatness terms of the bound and plot them in Figure 6.
A key observation in Figure 6 is that the learning rate impacts the trajectory term and the flatness term in opposite ways, as seen, for example, in (a) and (b), where the two set of curves swap their orders in the two figures. 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 flatness term ((b) and (d)), confirming the previous wisdom that small batch sizes enable the neural network to find a flat minima (Keskar et al., 2017).
E.3 IMPLEMENTATION DETAILS OF FIGURE 2
We let d0 = 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.5. Figure 7 shows the similar result of Theorem 2.

heading: E.4 ALGORITHM OF DYNAMIC GRADIENT CLIPPING AND ADDITIONAL RESULTS
text: The dynamic gradient clipping algorithm is described in Algorithm 1. For both MLP and AlexNet, we let α = 0.1. The start step for clipping, Tc, is also an important hyperparameter. However, it can be removed by detecting the evolution of the average gradient norm for each epoch. Specifically, whenever the average gradient norm of epoch j is larger than the average gradient norm of epoch j − 1, the clipping operation begins.
Algorithm 1 Dynamic Gradient Clipping Require: Training set S, Batch size b, Loss function `, Initial model parameter w0, Learning rate
λ, Initial minimum gradient norm G, Number of iterations T , Clipping parameter α, Clipping step Tc
1: for t← 1 to T do 2: Sample B = {zi}bi=1 from training set S 3: Compute gradient:
gB ← ∑b i=1∇w`(wt−1, zi)/b
4: if t > Tc then 5: if ||gB||2 > G then 6: gB ← α · G · gB/||gB||2 7: else 8: G ← ||gB||2 9: end if
10: end if 11: Update parameter: wt ← wt−1 − λ · gB 12: end for
From Figure 5 and Figure 8 we can see that dynamic gradient clipping effectively alleviates overfitting 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 significant 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.

heading: E.5 DISCUSSION ON GRADIENT DISPERSION OF MODELS TRAINED ON CLEAN DATASETS
text: In the case of no noise injected, Figure 4a 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(Gt + Nt;S|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. The difference between Figure 4e with Figure 4f-4h 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.

heading: E.6 ALGORITHM OF GAUSSIAN MODEL PERTURBATION AND ADDITIONAL RESULTS
text: Algorithm 2 Gaussian Model Perturbation Training Require: Training set S, Batch size b, Loss function `, Initial model parameter w0, Learning rate
λ, Number of noise k, Standard deviation of Gaussian distribution σ, Lagrange multiplier ρ while wt not converged do
2: Update iteration: t← t+ 1 Sample B = {zi}bi=1 from training set S 4: Sample ∆j ∼ N (0, σ2) for j ∈ [k] Compute gradient: gB ← ∑b i=1 ( ∇w`(wt, zi) + ρ ∑k j=1 (∇w`(wt + ∆j , zi)−∇w`(wt, zi)) /k ) /b 6: Update parameter: wt+1 ← wt − λ · gB end while
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. Specifically, 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 find 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.76× 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 letting 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.
min w Ls(w) + ρ E
∆∼N (0,σ2Id) [∣∣∣Ls(w + ∆)− Ls(w)∣∣∣]. This scheme can still perform well as shown in Table 3. In fact, it outperforms GMP3 on CIFAR100. This hints that it is possible to improve the performance by choosing other norm of Ls(w + ∆)− Ls(w).

Review summaries:
This paper presents novel and tighter information-theoretic upper bounds for the generalization error of machine learning models, such as neural networks, trained with SGD. Following the same construction of the auxiliary weight process in Neu (2021), the upper bounds proposed in this paper improve upon Neu (2021) in two ways. One 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). Additionally, the improvement also benefits from replacing a sample-level mutual information term in Neu (2021) with an individual sample mutual information term, exploiting a recent result of Bu et al. (2020). The bounds obtained here decompose 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 term”).  The authors also apply these bounds to analyze the generalization behavior of linear and two-layer ReLU networks. Experimental studies based on these bounds provide some insights into the SGD training of neural networks. Some simple regularization schemes, including dynamical gradient clipping and Gaussian model perturbation, are proposed, which are shown to perform comparably to the current state of the art. 

This paper derives an improved stability-based generalization bound for SGD upon Neu et. al.’s work. In particular, they drop a term called local gradient sensitivity, leading to significantly tighter bound in practical models. Based on this generalization bound, they study both linear model and two layer ReLU neural networks. During their analysis, they obtain several interesting observations, which is consistent to the previous ones. Finally, based on their bound, they propose a new training scheme, named Gaussian model perturbation, having good performance on popular datasets.

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. 

The paper presents a new bound on the generalization error of models trained with stochastic gradient descent. The new bound builds on the work of Neu, but is tighter. Based on the insights obtained from the bounds, the authors propose two approaches to improve generalization. On the one hand, they show that gradient clipping helps preventing overfitting. On the other hand, the authors propose a regularization term that adds Gaussian noise to the network weights (Gaussian model perturbation) and show that this type of regularization performs on par with or even better than other contemporary regularization schemes. Aside from these technical contributions, the authors provide ample insight into the mathematical terms they derive and discuss several phenomena observed during learning. In the outlook, the authors discuss how their bound can be further improved by utilizing the strong data processing inequality.