On the Generalization of Neural Networks Trained with SGD: Information-Theoretical Bounds and Implications
Year: 2021

Abstract:
Understanding the generalization behaviour of deep neural networks is an important 1 theme of modern research in machine learning. In this paper, we follow up on 2 a recent work of Neu [49] and present new information-theoretic upper bounds 3 for the generalization error of neural networks trained with SGD. Our bounds and 4 experimental study provide new insights on the SGD training of neural networks. 5 They also point to a new and simple regularization scheme which we show performs 6 comparably to the current state of the art. 7

Understanding the generalization behaviour of deep neural networks is an important1 theme of modern research in machine learning. In this paper, we follow up on2 a recent work of Neu [49] and present new information-theoretic upper bounds3 for the generalization error of neural networks trained with SGD. Our bounds and4 experimental study provide new insights on the SGD training of neural networks.5 They also point to a new and simple regularization scheme which we show performs6 comparably to the current state of the art.7
1 Introduction8
The outstanding performance of deep learning has brought to the surface some intriguing properties9 of deep neural networks, one of which is the observation that despite their high capacity, deep neural10 networks tend to generalize well [80]. This contradicts classical wisdom in statistical learning theory11 (e.g., [71] ) and has stimulated intense research interest in understanding the generalization behaviour12 of modern neural networks.13
One theme of research focuses on the study of over-parameterized neural networks, where generaliza-14 tion bounds are obtained [22, 1, 6, 50, 52, 53, 2, 3] and a curious “double descent” phenomenon is15 observed and analyzed [9, 46, 78]. New bounding techniques for analyzing generalization have also16 been developed, utilizing information theoretic quantities [64, 65, 77, 5, 13, 69, 23, 7, 4, 30, 84]. The17 bounds provided by these techniques have the advantages of accounting for both the data structure18 and the learning algorithm.19
The generalization ability of neural networks trained with mini-batched stochastic gradient descent20 [61], simply referred to SGD in this paper, has also been widely studied. Specifically, built on a21 connection between stability and generalization [12], a stability-based bound is first presented in [32],22 followed by a surge of research effort exploiting similar approaches [44, 17, 26, 43, 8]. Information-23 theoretic bounding techniques have also demonstrated great power in analyzing SGD-like algorithms.24 For example, [55] is the first to utilize information-theoretical bound in analyzing the generalization25 ability of SGLD [28, 74]. The bound was subsquently improved by [47, 31, 62, 72]. Inspired by the26 work of [55], [49] presents an information- theoretic analysis of the models trained with SGD. The27 analysis of [49] constructs an auxiliary weight process parallel to SGD training and upper-bounds the28 generalization error through this auxiliary process.29
Another line of research connects the generalization of neural networks with the flatness of loss30 minima [35] found by SGD or its variant [40, 20, 24, 51, 16, 37, 38, 83, 27]. This understanding31 has led to the discovery of new SGD-based training algorithms for improved generalization. For32 example, in a concurrent development by [83] and [27], a local “max-pooling” operation is applied to33
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.
the loss landscape prior to the SGD updates. This approach, referred to as AMP[83] or SAM[27], is34 shown to make SGD favor flatter minima and achieve the state-of-the-art performance among various35 competitive regularization schemes [83].36
In this paper, we focus on investigating the generalization of neural networks trained with SGD.37 We build upon the work of [49]. Following the same construction of the auxiliary weight process38 in [49], we present upper bounds of generalization error that improve upon [49] in two ways. The39 first improvement is via removing an unnecessary term in the bounds of [49] by invoking the HWI40 inequality [60]. The second improvement is via replacing a sample-level mutual information term in41 [49] with an instance-level mutual information term, exploiting a recent result of [13]. The bounds we42 obtain decompose into two terms, one measuring the impact of training trajectories (“the trajectory43 term”) and the other measuring the impact of the flatness of the found solution (“the flatness term”).44
We empirically validate the derived bounds. Various insights are also obtained experimentally45 concerning the generalization of neural networks under SGD training. For example, the batch size of46 SGD appears to impact the trajectory term and the flatness term in opposite ways, which complicates47 the overall dependency of generalization error on batch sizes. A particular interesting observation48 from our experiments is that a key quantity arising in the trajectory term of the bounds, which we49 refer to as gradient dispersion1, reveals a double descent phenomenon with respect to training epochs.50 Most intriguingly, the valley in the double descent curve appears to mark the great divide between the51 “generalization regime” and the “memorization regime” of training. Furthering from this observation,52 we also show that it is possible to reduce the memorization effect by dynamically clipping the gradient53 and reducing its dispersion.54
Our bounds also inspire a natural and simple solution to alleviate generalization error. Specifically,55 we propose a new training scheme, referred to as Gaussian model perturbation (GMP), aiming at56 reducing the flatness term of the bounds. This scheme effectively applies a local “average pooling” to57 the empirical risk surface prior to SGD, greatly resembling the “max-pooling” approach adopted in58 AMP[83]. We demonstrate experimentally that GMP achieves a competitive performance with the59 current art of regularization schemes.60
Length constraints precludes elaboration at places. The reader is referred to supplementary materials61 for proofs and additional information.62
Other Related Literature Gradient dispersion is mostly studied from optimization perspectives[11,63 63, 39, 75, 25]. Prior to this work, only a few works relate gradient dispersion with the generalization64 behaviour of the networks. In [49, 72], gradient dispersion also appears in the generalization bounds.65 In [38], gradient dispersion is argued to capture a notion of “flatness” of the local minima of the loss66 landscape, thereby correlating with generalization.67
Injecting noise in the training process has been proposed in various regularization schemes, for68 example, [10, 14, 15, 68, 73]. But unlike the Gaussian model perturbation scheme derived in this69 paper, where noise is injected to the model parameters, noise in those schemes is injected either to70 the training data or to the network activation.71
Gradient clipping is a common technique for preventing gradient exploding (see, e.g., [45, 56]).72 This technique is also used in [82] to accelerate training. In this paper, gradient clipping is used to73 investigate and control the impact of gradient dispersion on generalization error.74
2 Preliminaries75
Population Risk, Empirical Risk and Generalization Error Unless otherwise noted, a random76 variable will be denoted by a capitalized letter (e.g., Z), and its realization denoted by the correspond-77 ing lower-case letter (e.g. z). Let Z be the instance space of interest and µ be an unknown distribution78 on Z , specifying random variable Z. Let W ⊆ Rd be the space of hypotheses. Suppose that a79 training sample S = (Z1, Z2, . . . , Zn) is drawn i.i.d. from µ and that a stochastic learning algorithm80 A takes S as its input and outputs a hypothesis W ∈ W according to some conditional distribution81 PW |S mapping Zn toW . Let ` : W × Z → R+ be a loss function, where `(w, z) measures the82 “unfitness” or “error” of any z ∈ Z with respect to a hypothesis w ∈ W . The population risk, for any83
1The quantity is often referred to as gradient variance in the literature [49, 72], but we prefer “dispersion” to “variance” so as to better comply with the mathematical conventions and avoid possible confusion.
w ∈ W , is defined as84 Lµ(w) , EZ∼µ[`(w,Z)].
The goal of learning is to find a hypothesis w that minimizes the population risk. But since µ is only85 partially accessible via the sample S, in practice, we instead turn to the empirical risk, defined as86
LS(w) , 1
n n∑ i=1 `(w,Zi).
The expected generalization error of the learning algorithm A is then defined as87
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).88 Throughout this paper, we take ` as a continuous function (adopting the usual notion “surrogate loss”89 [66].). Additionally, we assume that ` is differentiable almost everywhere with respect to both w and90 z. Furthermore we assume that `(w,Z) is R-subgaussian2 for any w ∈ W . Note that a bounded loss91 is guaranteed to be subgaussian for all µ and all w ∈ W . Let I(X;Y ) denote the mutual information92 [18] between any pair of random variables (X,Y ). The following results are known.93 Lemma 1 ([77, Theorem 1.]). The expected generalization error of algorithm A is bounded by94
|gen(µ, PW |S)| ≤ √ 2R2
n I(W ;S),
95
Lemma 2 ([13, Proposition 1.]). The expected generalization error of algorithm A is bounded by96
|gen(µ, PW |S)| ≤ 1
n n∑ i=1 √ 2R2I(W ;Zi),
97
Stochastic Gradient Descent We now restrict the learning algorithm A to be the mini-batched98 stochastic gradient descent (SGD) algorithm for empirical risk minimization. For each training epoch,99 the dataset S is randomly split into m disjoint mini-batches, each having size b, namely, n = mb.100 Based on each batch, one parameter update is performed. Specifically, let Bt denote the batch used101 for the tth update. Define102
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. The103 rule for the tth parameter update is then104
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 drawn105 from the zero-mean spherical Gaussian N (0, σ20Id) with variance σ20 in each dimension. We will106 assume that the SGD algorithm stops after T updates and outputs WT as the learned model parameter.107
Given the training sample S, let ξ govern the randomness in the sequence (B1, B2, . . . , BT ) of108 batches. For the simplicity of notion, we will fix the configuration of ξ. That is, we will assume a109 fixed “batching trajectory”, or a fixed way to shuffle the example indices {1, . . . , n} and divide them110 into m batches in each epoch. The presented generalization bounds of this paper can be extended111 to the case where the batching trajectory is uniformly random (as we set up above). This merely112 involves averaging over all batching trajectories or taking expectation over ξ.113
Auxiliary Weight Process We now associate with the SGD algorithm an auxiliary weight process114 {W̃t}. Let σ2 be given, and let σ1, σ2, . . . , σT be a sequence of positive real numbers. Define115
W̃0 ,W0, and W̃t , W̃t−1 − λtg(Wt−1, Bt) +Nt, for t > 0, where Nt ∼ N (0, σ2t Id) is a Gaussian noise. The relationship between this auxiliary weight process116 {W̃t} and the weight process {Wt} in SGD is shown in the Bayesian network below.117
2Recall that a random variable X is R-subgaussian [60] if for any ρ, logE exp (ρ (X − EX)) ≤ ρ2R2/2.
N1 N2 · · · NT−1 NT ↓ ↓ ↓ ↓
W̃0 → W̃1 → W̃2 → · · · → W̃T−1 → W̃T ‖ ↗ ↗ ↗ ↗ W0 → W1 → W2 → · · · → WT−1 → WT 118
Let ∆t = ∑t τ=1Nτ . Noting that the weight updates in {W̃t} uses the same gradient signal as that119 used in {Wt} (which depends onWt−1 not W̃t−1), it is immediate that W̃t = Wt+∆t. Note that this120 auxiliary process follows the same construction as [49], which we will use to study the generalization121 error of SGD.122
To that end, define gradient dispersion at parameter w by123 V(w) , E [ ||∇w`(w,Z)− E [∇w`(w,Z)] ||22 ] ,
where the expectation is taken over Z ∼ µ.124 For a given sample s ∈ Zn, define125
γ(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.126
In the remainder of the paper, let S′ denote another sample drawn from µn, independent of all other127 random variables. The main generalization bound in [49] is re-stated below.128 Lemma 3 ([49, Theorem 1.]). The generalization error of SGD is upper bounded by129
|gen(µ, PWT |S)| ≤ √√√√2R2 n T∑ t=1 λ2t σ2t E [ Ψ(Wt−1) + V(Wt−1) b ] + |E [γ(WT , S)− γ(WT , S′)]| ,
where Ψ(wt−1) , E [ ||∇w`(wt−1, Z)−∇w`(wt−1 + ζ, Z)||22 ] and ζ ∼ N (0, 2 ∑t−1 i=1 σ 2 i Id).130
The term Ψ(wt−1) in the bound is referred to as “local gradient sensitivity” in [49].131
3 New Generalization Bounds for SGD132
We first prove that the generalization bound in Lemma 3 can be tightened by removing the local133 gradient sensitivity term Ψ(wt−1). The key observation is that an independence condition used for134 establishing Lemma 3 in [49] is unnecessary (see Lemma 4 in [49]). This requires invoking a vector135 version of the HWI inequality [60, Lemma 3.4.2], which we prove in this paper.136 Lemma 4. Let X and Y be two random vectors in Rd, and let N ∼ N (0, Id) be independent of137 (X,Y ). Then, for every t > 0, DKL(PX+√tN ||PY+√tN ) ≤ 1 2tE [ ||X − Y ||2 ] .138
Here DKL is the KL divergence. Note that the bound in Lemma 3 relies on a similar result which139 however requires the independence of X and Y . Using Lemma 4, we obtain the following theorem.140 Theorem 1. The generalization error of SGD is upper bounded by141
|gen(µ, PWT |S)| ≤ √√√√2R2 nb T∑ t=1 λ2t σ2t E [V(Wt−1)] + |E [γ(WT , S)− γ(WT , S′)]| .
142
The proof of this theorem, as the bounds in [49], relies on Lemma 1 and the sample-level mutual143 information bound therein. This theorem can be further tightened by exploiting the fact that the144 instance-level mutual information bound in Lemma 2 is in fact tighter than the sample-level mutual145 information bound in Lemma 1, as shown in [13]. The main ingredient to proceed in this direction is146 the following lemma.147
Lemma 5. LetGt = −λtg(Wt−1, Bt). If Zi ∈ Bt, then I(Gt+Nt;Zi|W̃t−1) ≤ λ 2 t
σ2t b 2E [V(Wt−1)].148 149
In this lemma, the mutual information I(Gt +Nt;Zi|W̃t−1) roughly indicates the degree by which150 the SGD’s updating signal Gt (smoothed with noise) depends on an individual training instance151 Zi, when Zi is used for computing the gradient. When this dependency is strong (giving rise to a152 high value of the mutual information), the model conceivably tends to overfit the individual training153 instances. This lemma suggests that the strength of this dependency can be upper-bounded by the154 expected gradient dispersion at the current weight configuration. In our experiments, we will estimate155 the expected gradient dispersion and validate this intuition.156
It is remarkable that the noise {Nt} plays an important role for the bound to hold. To see this,157 consider b = 1 and Z is countable and large. Then I(Gt;Zt|Wt−1) is merely the conditional entropy158 H(Zi|Wt−1), which would grow with sample size n at least as log n. Upper-bounding it with a159 quantity independent of n would be impossible – This justifies the construction of the auxiliary160 weight process.161
We now state our main theorem. Unlike Theorem 1, which considers a random batching trajectory,162 this theorem considers a fixed batching trajectory to keep the expression less cluttered. For that163 batching trajectory, we will use Ti to denote the set of indices of batches Bt containing instance Zi.164 Theorem 2. The expected generalization error of SGD is bounded by165
|gen(µ, PWT |S)| ≤ R
nb n∑ i=1 √√√√∑ t∈Ti 2λ2t σ2t E [V(Wt−1)] + |E [γ(WT , S)− γ(WT , S′)]| .
166
With an additional assumption, the second term in the bound can be re-expressed, as shown in the167 following corollary.168 Corollary 1. Assume Lµ(wT ) ≤ E∆ [Lµ(wT + ∆T )], then the following holds,169
gen(µ, PWT |S) ≤ R
nb n∑ i=1 √√√√∑ t∈Ti 2λ2t σ2t E [V(Wt−1)] + ∑T t=1 σ 2 t 2 Tr (E [HWT (Z)]) ,
where HWT is the Hessian matrix of the loss with respect to WT and Tr(·) denotes trace.170
Corollary 1 follows directly from the second order Taylor expansion of the second term in the bound171 of Theorem 2. The condition Lµ(wT ) ≤ E∆ [Lµ(wT + ∆T )] indicates that the perturbation does172 not decrease the population risk. This is also assumed in [27] in the derivation of a PAC-Bayesian173 generalization bound.174
Notably, in the bound of Theorem 2, the first term captures the impact of the training trajectory175 (“trajectory term”), and the second term captures the impact of the final solution. As seen in Corollary176 1, this term in fact measures the flatness for the loss landscape at the found solution (“flatness term”).177 The previous bound of [49] (Lemma 3) and its tightened version in Theorem 1 also similarly contain a178 trajectory term and a flatness term. Despite that the flatness term there are identical to that in Theorem179 2, we now show the trajectory term in Theorem 2 does improve on its counter-part in Theorem 1.180 Lemma 6. Assume the instances are sampled without replacement in every epoch. Then the trajectory181 term in Theorem 2 is upper-bounded by182
min Rn T∑ t=1 √ 2λ2t σ2t E [V(Wt−1)], √√√√2R2 nb T∑ t=1 λ2t σ2t E [V(Wt−1)]  . 183
The condition in Lemma 6 is usually satisfied in practice. This lemma then immediately implies that184 the trajectory term in Theorem 2 is no worse than that in Theorem 1. Incorporating this result, if we185 restrict the smoothness of the loss function `, we may obtain another version of the generalization186 bound (although the flatness term therein is expected to be looser than that in Corollary 1).187 Corollary 2. If the loss function is differentiable and β-smooth with respect to w, then under the188 condition of Lemma 6,189
|gen(µ, PWT |S)| ≤min Rn T∑ t=1 √ 2λ2t σ2t E [V(Wt−1)], √√√√2R2 nb T∑ t=1 λ2t σ2t E [V(Wt−1)] + βd T∑ t=1 σ2t .
190
To conclude, we remark that these bounds suggest that in order for the model to generalize well, both191 the trajectory term and the flatness term need to be small — the former involves the interaction of the192 learning rate and batch size with the gradient dispersion along the training trajectory, whereas the193 latter depends on the flatness of the empirical risk surface at the found solution.194
4 Experimental Study195
Bound Verification We first verify our bound in Corollary 1 by training an MLP (with one hidden196 layer) and an AlexNet [42] on MNIST and CIFAR10 [41], respectively. To simplify estimation, we197 fix the weight initialization and set σt and λt to be constants σ and λ, respectively. To compute198 ∑n i=1 √∑ t∈Ti E [V(Wt−1)], we compute the gradient dispersion as its empirical estimate from a199 batch, utilizing a PyTorch [54] library BackPack [19]. To compute Tr (E [HWT (Z)]), we randomly200 sample 10% of the training data and use the PyHessian library [79] to compute the Hessian. Since201 every choice of σ gives a valid generalization bound in Corollary 1, we need to find the optimal202 σ, which gives the tightest bound. This can be done by simply utilizing the fact A/σ + σ2B ≥203 3(A/2)2/3B1/3 for any positive A and B, where the equality is achieved by the optimal σ. We set204 the sub-gaussian parameter R = 0.1. The implementation in this paper is on PyTorch, and all the205 experiments are carried out on NVIDIA Tesla V100 GPUs (32 GB).206
We perform experiments with varying network width and varying levels of label noise. Specifically,207 label noise level refers to the setting where we replace the labels of fraction of the training and208 testing instances with random labels. The estimated bound is compared against the true generalization209 gap, namely, the difference between the training loss and testing loss, and is shown in Figure 1.210
In Figure 1, we see that in all cases the estimated bound follows closely the trend of the true211 generalization gap. The fact that the bound curve consistently tracks the gap curve under various212 label noise levels indicates that our bound very well captures the changes of the data distribution.213 Note that in Figure 1 (a) and (b), our bound decays with the increase of the model size, showing a214 trend as opposite to the bounds obtained in classical learning theory. But such a trend clearly better215 explains the generalization behaviour of modern neural networks.216
Learning Rate and Batch Size. The learning rate and batch size in SGD have explicitly appeared217 in the trajectory term of the bound in Theorem 2. From the way they appear in the bound, one may be218 tempted to assert that a small learning rate or large batch size will improve generalization. This would219 then contradict some previous observations [37, 76, 33], in which increasing the ratio of learning rate220
to batch size will benefit generalization. We now investigate this by performing experiments with221 varying learning rates and batch sizes. In our experiments, the model is continuously updated until222 the average training loss drops below 0.0001. We separate trajectory and flatness terms of the bound223 and plot them in Figure 2.224
A key observation in Figure 2 is that the learning rate impacts the trajectory term and the flatness term225 in opposite ways, as seen, for example, in (a) and (b), where the two set of curves swap their orders in226 the two figures. On the other hand, the batch size also impacts the two terms in opposite ways, as seen227 in (a) and (b) where curves decrease in (a) but increase in (b). This makes the generalization bound,228 i.e., the sum of the two terms, have a rather complex relationship with the settings of learning rate229 and batch size. This relationship is further complicated by the fact that a small learning rate requires230 a longer training time, or a larger number T of training iterations, which increases the number that231 are summed over in the trajectory term. Nonetheless, we do observe that a smaller batch size gives a232 lower value of the flatness term ((b) and (d)), confirming the previous wisdom that small batch sizes233 enable the neural network to find a flat minima [40].234
Double Descent of Gradient Dispersion We experimentally investigate the impact of gradient235 dispersion on the training of the neural networks by fixing the learning rate, batch size and weight236 initialization for the each model (MLP for MNIST, AlexNet for CIFAR10). For each model and237 various label noise levels, we plot in Figure 3 the evolution of the (empirical) gradient dispersion238 V̂(wt), training accuracy and testing accuracy across training epochs.239 An intriguing epoch-wise “double descent” phenomenon is observed, particularly when the labels are240 noisy. According to the double descent curve, the training may be split into three phases (e.g., Figure241 3 (h)). In the first phase, the gradient dispersion rapidly descends and maintains a very low level. In242 this phase, both training and test accuracies increase while maintaining a very small generalization243 gap. This suggests that the network in this phase is extracting useful patterns and generalizes well.244 In the second phase, the gradient dispersion starts increasing until it reaches a peak value. In this245 phase, the training and testing accuracies gradually diverge, marking the model entering an overfitting246 or “memorization” regime – when the data contains the noisy labels, the network mostly tries to247 memorize the labels in the training set. In the third phase, the gradient dispersion descends again,248 reaching a low value. In this phase, the model continuously overfits the training data, until the training249 and testing curves reach their respective maximum and minimum. It appears that the timing of the250 three phases depends on the dataset and the label noise level. For simpler data (e.g. MNIST) and251 cleaner datasets (e.g. CIFAR10 with low label noise), the first phase may be shorter. This is arguably252 because in these datasets, extracting useful patterns is relatively easier. Nonetheless, the valley in the253 double-descent curve appears to mark a “great divide” between generalization and memorization.254
Dynamic Gradient Clipping Inspired by our generalization bounds and above observations, one255 way to reduce the generalization error is to control the trajectory term of the bounds by reducing the256 gradient dispersion in each training step. Here we investigate a simple scheme that dynamically clips257
the gradient norm so as to reduce the gradient dispersion. Specifically, whenever the current gradient258 norm is larger than the gradient norm K steps earlier, or ||g(Wt, Bt)||2 > ||g(Wt−K , Bt−K)||2 (i.e.,259 the model is expected to have entered the “memorization” regime), we reduce the norm of the current260 gradient g(Wt, Bt) to α fraction of ||g(Wt−K , Bt−K)||2, for some prescribed value α < 1. The261 effectiveness of this scheme is best demonstrated when the labels contain noise. As shown in Figure262 4, dynamic gradient clipping significantly closes the gap between the training accuracy and the testing263 accuracy. The models trained with this scheme maintain a near-optimal testing accuracy (e.g., about264 80% when the label noise level of MNIST is 0.2), without suffering from the severe memorization265 effect as seen in models trained without this scheme. Further understanding of the double-descent266 phenomenon of the gradient dispersion may enable more delicate design of such a dynamic clipping267 scheme and potentially lead to novel and powerful regularization techniques.268
5 A Practical Implication: Gaussian Model Perturbation269
The appearance of the flatness term in our generalization bounds suggests that for an empirical270 risk minimizer w∗ to generalize well, it is necessary that the empirical risk surface at w∗ is flat,271 or insensitive to a small perturbation of w∗. This naturally motivates a training scheme using the272 following regularized loss:273
min w Ls(w) + ρ E ∆∼N (0,σ2Id) [Ls(w + ∆)− Ls(w)],
where ρ is a hyper-parameter. Replacing the expectation above with its stochastic approximation274 using k realizations of ∆ gives rise to the following optimization problem.275
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,276 GMP requires k + 1 forward passes for every parameter update. Empirical evidence shows that a277 small k, for example, k = 3, already gives competitive performance. Implementing the k+ 1 forward278 passes on parallel processors further reduces the computation load.279
We experimentally compare GMP with several major regularization schemes in the current art, includ-280 ing Dropout [68], label smoothing [70], Flooding [36], MixUp [81], adversarial training [29], and281 AMP [83]. The compared schemes are evaluated on three popular benchmark image classification282 datasets SVHN [48], CIFAR-10 and CIFAR-100 [41]. Two representative deep architectures PreAc-283 tResNet18 [34] and VGG16 [67] are taken as the underlying model. We train the models for 200284 epochs by SGD. The learning rate is initialized as 0.1 and divided by 10 after 100 and 150 epochs.285 For all compared models, the batch size is set to 50 and weight decay is set to 10−4. For GMP, we286 choose ρ = 0.5 and set the standard deviation of the Gaussian noise ∆ to 0.03. The value of k is287 chosen as 3 and 10 respectively (referred to as GMP3 and GMP10).288
The performances of all compared schemes are given in Table 1. For the compared regularization289 schemes except GMP, we directly report their performances as given in [83]. Performances of vanilla290 ERM without regularization are also included as a reference.291
Table 1 demonstrates the effectiveness of GMP. Overall GMP performs comparably to the current292 art of regularization schemes, although appearing slighly inferior to the most recent record given by293 AMP [83]. Noting that the key ingredient of AMP, “max-pooling” in the parameter space, greatly294 resembles regularization term in GMP, which may be seen as “average-pooling” in the same space.295
6 Conclusion and Outlook296
This paper presents new generalization bounds for neural networks trained with SGD, improving297 upon the results of [49]. Our bounds naturally point to new and effective regularization schemes. At298 the same time, they reveal interesting phenomena in the SGD training of neural networks. While299 these phenomena deserve further investigation in their own right, we here suggest another direction300 for improving the bounds, namely, via the use of strong data-processing inequalities (DPI) [57–59]301 (noting that the standard DPI is in fact needed for establishing Theorem 2).302
For any Markov chain U → X → Y , we will denote by U , X , and Y the spaces in which U ,303 X , Y take values, respectively. For any distribution P on X , we will use PY |X ◦ P to denote the304 distribution on Y induced by the push-forward of the distribution P by PY |X , namely, for any y ∈ Y ,305 ( PY |X ◦ P ) (y) , ∫ PY |X(y|x)P (x)dx. Let S(U) be the support of PU and H(U,PX|U ) be the306 convex hull of {PX|U=u : u ∈ S(U)}. Define307
η(U → X → Y ) , sup P,Q∈H(U,PX|U ) DKL(PY |X ◦ P ||PY |X ◦Q) DKL(P ||Q)
Lemma 7. For any Markov chain U → X → Y , I(U ;Y ) ≤ η(U → X → Y )I(U ;X).308
Here η(U → X → Y ) serves as the “contraction coefficient” for the stochastic kernel PY |X ,309 characterizing the greatest extent by which the kernel may bring closer any two distributions on S(U)310 in its output space. It is easy to see that η(U → X → Y ) ≤ 1, giving rise to a stronger DPI.311
Denote Vt , W̃t−1 + Gt. It can be verified that Zi → Vt → W̃t form a Markov chain. Denote312 ηi,t , η(Zi → Vt → W̃t), and Γti = {t+ 1, t+ 2, . . . , T} \ Ti. Theorem 2 can be improved to:313 Theorem 3. The expected generalization error of SGD is bounded by314
|gen(µ, PWT |S)| ≤ 2R
nb n∑ i=1 √√√√∑ t∈Ti λ2t σ2t E [V(Wt−1)] · ∏ τ∈Γti ηi,τ + |E [γ(WT , S)− γ(WT , S′)] |.
315
It remains to characterize the contraction coefficient ηi,τ in a computable form. Simply bounding it316 via the Dobrushin’s coefficient [21], as suggested in [72] for analyzing SGLD, is unlikely to make317 the bound in this theorem significantly tighter than that in Theorem 2.318
References319 [1] Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over-320 parameterization. In International Conference on Machine Learning, pages 242–252. PMLR,321 2019.322
[2] S. Arora, R. Ge, B. Neyshabur, and Y. Zhang. Stronger generalization bounds for deep nets via323 a compression approach. In International Conference on Machine Learning, pages 254–263.324 PMLR, 2018.325
[3] S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and326 generalization for overparameterized two-layer neural networks. In International Conference327 on Machine Learning, pages 322–332. PMLR, 2019.328
[4] A. R. Asadi and E. Abbe. Chaining meets chain rule: Multilevel entropic regularization and329 training of neural networks. Journal of Machine Learning Research, 21(139):1–32, 2020.330
[5] A. R. Asadi, E. Abbe, and S. Verdú. Chaining mutual information and tightening generalization331 bounds. In Proceedings of the 32nd International Conference on Neural Information Processing332 Systems, pages 7245–7254, 2018.333
[6] P. L. Bartlett, D. J. Foster, and M. Telgarsky. Spectrally-normalized margin bounds for neural334 networks. In Proceedings of the 31st International Conference on Neural Information Processing335 Systems, pages 6241–6250, 2017.336
[7] R. Bassily, S. Moran, I. Nachum, J. Shafer, and A. Yehudayoff. Learners that use little337 information. In Algorithmic Learning Theory, pages 25–55. PMLR, 2018.338
[8] R. Bassily, V. Feldman, C. Guzmán, and K. Talwar. Stability of stochastic gradient descent on339 nonsmooth convex losses. Advances in Neural Information Processing Systems, 33, 2020.340
[9] M. Belkin, D. Hsu, S. Ma, and S. Mandal. Reconciling modern machine-learning practice and341 the classical bias–variance trade-off. Proceedings of the National Academy of Sciences, 116342 (32):15849–15854, 2019.343
[10] C. M. Bishop. Training with noise is equivalent to tikhonov regularization. Neural computation,344 7(1):108–116, 1995.345
[11] L. Bottou, F. E. Curtis, and J. Nocedal. Optimization methods for large-scale machine learning.346 Siam Review, 60(2):223–311, 2018.347
[12] O. Bousquet and A. Elisseeff. Stability and generalization. The Journal of Machine Learning348 Research, 2:499–526, 2002.349
[13] Y. Bu, S. Zou, and V. V. Veeravalli. Tightening mutual information-based bounds on generaliza-350 tion error. IEEE Journal on Selected Areas in Information Theory, 1(1):121–130, 2020.351
[14] A. Camuto, M. Willetts, U. Simsekli, S. J. Roberts, and C. C. Holmes. Explicit regularisation in352 gaussian noise injections. In Advances in Neural Information Processing Systems, 2020.353
[15] A. Camuto, X. Wang, L. Zhu, C. Holmes, M. Gürbüzbalaban, and U. Şimşekli. Asymmetric354 heavy tails and implicit bias in gaussian noise injections. arXiv preprint arXiv:2102.07006,355 2021.356
[16] P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. T. Chayes,357 L. Sagun, and R. Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. In 5th358 International Conference on Learning Representations. OpenReview.net, 2017.359
[17] Y. Chen, C. Jin, and B. Yu. Stability and convergence trade-off of iterative optimization360 algorithms. arXiv preprint arXiv:1804.01619, 2018.361
[18] T. M. Cover and J. A. Thomas. Elements of Information Theory. John Wiley & Sons, 2012.362
[19] F. Dangel, F. Kunstner, and P. Hennig. Backpack: Packing more into backprop. In 8th363 International Conference on Learning Representations. OpenReview.net, 2020.364
[20] L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp minima can generalize for deep nets. In365 International Conference on Machine Learning, pages 1019–1028. PMLR, 2017.366
[21] R. L. Dobrushin. Central limit theorem for nonstationary markov chains. i. Theory of Probability367 & Its Applications, 1(1):65–80, 1956.368
[22] S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-369 parameterized neural networks. In International Conference on Learning Representations,370 2018.371
[23] C. Dwork, V. Feldman, M. Hardt, T. Pitassi, O. Reingold, and A. Roth. Generalization in372 adaptive data analysis and holdout reuse. In Proceedings of the 28th International Conference373 on Neural Information Processing Systems-Volume 2, pages 2350–2358, 2015.374
[24] G. K. Dziugaite and D. M. Roy. Computing nonvacuous generalization bounds for deep375 (stochastic) neural networks with many more parameters than training data. In Proceedings of376 the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAI), 2017.377
[25] F. Faghri, D. Duvenaud, D. J. Fleet, and J. Ba. A study of gradient variance in deep learning.378 arXiv preprint arXiv:2007.04532, 2020.379
[26] V. Feldman and J. Vondrak. High probability generalization bounds for uniformly stable380 algorithms with nearly optimal rate. In Conference on Learning Theory, pages 1270–1279.381 PMLR, 2019.382
[27] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-aware minimization for efficiently383 improving generalization. arXiv preprint arXiv:2010.01412, 2020.384
[28] S. B. Gelfand and S. K. Mitter. Recursive stochastic algorithms for global optimization in rˆd.385 SIAM Journal on Control and Optimization, 29(5):999–1018, 1991.386
[29] I. J. Goodfellow, J. Shlens, and C. Szegedy. Explaining and harnessing adversarial examples. In387 3rd International Conference on Learning Representations, ICLR, 2015.388
[30] H. Hafez-Kolahi, Z. Golgooni, S. Kasaei, and M. Soleymani. Conditioning and processing:389 Techniques to improve information-theoretic generalization bounds. Advances in Neural390 Information Processing Systems, 33, 2020.391
[31] M. Haghifam, J. Negrea, A. Khisti, D. M. Roy, and G. K. Dziugaite. Sharpened generalization392 bounds based on conditional mutual information and an application to noisy, iterative algorithms.393 arXiv preprint arXiv:2004.12983, 2020.394
[32] M. Hardt, B. Recht, and Y. Singer. Train faster, generalize better: Stability of stochastic gradient395 descent. In International Conference on Machine Learning, pages 1225–1234. PMLR, 2016.396
[33] F. He, T. Liu, and D. Tao. Control batch size and learning rate to generalize well: Theoretical397 and empirical evidence. In Advances in Neural Information Processing Systems 32, pages398 1141–1150, 2019.399
[34] K. He, X. Zhang, S. Ren, and J. Sun. Identity mappings in deep residual networks. In European400 Conference on Computer Vision, pages 630–645, 2016.401
[35] S. Hochreiter and J. Schmidhuber. Flat minima. Neural Computation, 9(1):1–42, 1997.402
[36] T. Ishida, I. Yamane, T. Sakai, G. Niu, and M. Sugiyama. Do we need zero training loss after403 achieving zero training error? In Proceedings of the 37th International Conference on Machine404 Learning, ICML, 2020.405
[37] S. Jastrzębski, Z. Kenton, D. Arpit, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey. Three406 factors influencing minima in sgd. arXiv preprint arXiv:1711.04623, 2017.407
[38] Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization408 measures and where to find them. In International Conference on Learning Representations,409 2019.410
[39] R. Johnson and T. Zhang. Accelerating stochastic gradient descent using predictive variance411 reduction. Advances in neural information processing systems, 26:315–323, 2013.412
[40] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch train-413 ing for deep learning: Generalization gap and sharp minima. In 5th International Conference414 on Learning Representations. OpenReview.net, 2017.415
[41] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report,416 University of Toronto, 2009.417
[42] A. Krizhevsky, I. Sutskever, and G. E. Hinton. Imagenet classification with deep convolutional418 neural networks. Advances in neural information processing systems, 25:1097–1105, 2012.419
[43] Y. Lei and Y. Ying. Fine-grained analysis of stability and generalization for stochastic gradient420 descent. In International Conference on Machine Learning, pages 5809–5819. PMLR, 2020.421
[44] B. London. A pac-bayesian analysis of randomized learning with application to stochastic422 gradient descent. In Proceedings of the 31st International Conference on Neural Information423 Processing Systems, pages 2935–2944, 2017.424
[45] S. Merity, N. S. Keskar, and R. Socher. Regularizing and optimizing lstm language models. In425 International Conference on Learning Representations, 2018.426
[46] P. Nakkiran, G. Kaplun, Y. Bansal, T. Yang, B. Barak, and I. Sutskever. Deep double de-427 scent: Where bigger models and more data hurt. In International Conference on Learning428 Representations, 2019.429
[47] J. Negrea, M. Haghifam, G. K. Dziugaite, A. Khisti, and D. M. Roy. Information-theoretic430 generalization bounds for sgld via data-dependent estimates. In Advances in Neural Information431 Processing Systems, pages 11013–11023, 2019.432
[48] Y. Netzer, T. Wang, A. Coates, A. Bissacco, B. Wu, and A. Y. Ng. Reading digits in natural433 images with unsupervised feature learning. In NeurIPS Workshop on Deep Learning and434 Unsupervised Feature Learning, 2011.435
[49] G. Neu. Information-theoretic generalization bounds for stochastic gradient descent. arXiv436 preprint arXiv:2102.00931, 2021.437
[50] B. Neyshabur, R. Tomioka, and N. Srebro. Norm-based capacity control in neural networks. In438 Conference on Learning Theory, pages 1376–1401. PMLR, 2015.439
[51] B. Neyshabur, S. Bhojanapalli, D. McAllester, and N. Srebro. Exploring generalization in deep440 learning. In Advances in Neural Information Processing Systems, pages 5947–5956, 2017.441
[52] B. Neyshabur, S. Bhojanapalli, and N. Srebro. A pac-bayesian approach to spectrally-normalized442 margin bounds for neural networks. In International Conference on Learning Representations,443 2018.444
[53] B. Neyshabur, Z. Li, S. Bhojanapalli, Y. LeCun, and N. Srebro. The role of over-parametrization445 in generalization of neural networks. In International Conference on Learning Representations,446 2018.447
[54] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin,448 N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning449 library. Advances in Neural Information Processing Systems, 32:8026–8037, 2019.450
[55] A. Pensia, V. Jog, and P.-L. Loh. Generalization error bounds for noisy, iterative algorithms.451 In 2018 IEEE International Symposium on Information Theory (ISIT), pages 546–550. IEEE,452 2018.453
[56] M. Peters, M. Neumann, M. Iyyer, M. Gardner, C. Clark, K. Lee, and L. Zettlemoyer. Deep454 contextualized word representations. In Proceedings of the 2018 Conference of the North Amer-455 ican Chapter of the Association for Computational Linguistics: Human Language Technologies,456 Volume 1 (Long Papers), pages 2227–2237, 2018.457
[57] Y. Polyanskiy and Y. Wu. Dissipation of information in channels with input constraints. IEEE458 Transactions on Information Theory, 62(1):35–55, 2015.459
[58] Y. Polyanskiy and Y. Wu. Strong data-processing inequalities for channels and bayesian460 networks. In Convexity and Concentration, pages 211–249. Springer, 2017.461
[59] M. Raginsky. Strong data processing inequalities and φ-sobolev inequalities for discrete462 channels. IEEE Transactions on Information Theory, 62(6):3355–3389, 2016.463
[60] M. Raginsky and I. Sason. Concentration of measure inequalities in information theory,464 communications and coding. Foundations and Trends in Communications and Information465 Theory; NOW Publishers: Boston, MA, USA, 2018.466
[61] H. Robbins and S. Monro. A stochastic approximation method. The annals of mathematical467 statistics, pages 400–407, 1951.468
[62] B. Rodríguez-Gálvez, G. Bassi, R. Thobaben, and M. Skoglund. On random subset general-469 ization error bounds and the stochastic gradient langevin dynamics algorithm. arXiv preprint470 arXiv:2010.10994, 2020.471
[63] N. L. Roux, M. Schmidt, and F. Bach. A stochastic gradient method with an exponential472 convergence rate for finite training sets. In Proceedings of the 25th International Conference on473 Neural Information Processing Systems-Volume 2, pages 2663–2671, 2012.474
[64] D. Russo and J. Zou. Controlling bias in adaptive data analysis using information theory. In475 Artificial Intelligence and Statistics, pages 1232–1240, 2016.476
[65] D. Russo and J. Zou. How much does your data exploration overfit? controlling bias via477 information usage. IEEE Transactions on Information Theory, 66(1):302–323, 2019.478
[66] S. Shalev-Shwartz and S. Ben-David. Understanding machine learning: From theory to479 algorithms. Cambridge university press, 2014.480
[67] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image481 recognition. In 3rd International Conference on Learning Representations, ICLR, 2015.482
[68] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: A simple483 way to prevent neural networks from overfitting. The Journal of Machine Learning Research,484 15(1):1929–1958, 2014.485
[69] T. Steinke and L. Zakynthinou. Reasoning about generalization via conditional mutual infor-486 mation. In Conference on Learning Theory, volume 125 of Proceedings of Machine Learning487 Research, pages 3437–3452. PMLR, 2020.488
[70] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna. Rethinking the inception archi-489 tecture for computer vision. In Proceedings of the IEEE Conference on Computer Vision and490 Pattern Recognition, CVPR, pages 2818–2826, 2016.491
[71] V. Vapnik. Statistical learning theory. Wiley, 1998. ISBN 978-0-471-03003-4.492
[72] H. Wang, Y. Huang, R. Gao, and F. P. Calmon. Learning while dissipating information:493 Understanding the generalization capability of sgld. arXiv preprint arXiv:2102.02976, 2021.494
[73] C. Wei, S. Kakade, and T. Ma. The implicit and explicit regularization effects of dropout. In495 International Conference on Machine Learning, pages 10181–10192. PMLR, 2020.496
[74] M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics.497 In Proceedings of the 28th international conference on machine learning (ICML-11), pages498 681–688. Citeseer, 2011.499
[75] Y. Wen, K. Luk, M. Gazeau, G. Zhang, H. Chan, and J. Ba. An empirical study of stochastic500 gradient descent with structured covariance noise. In International Conference on Artificial501 Intelligence and Statistics, pages 3621–3631. PMLR, 2020.502
[76] L. Wu, C. Ma, and E. Weinan. How sgd selects the global minima in over-parameterized503 learning: A dynamical stability perspective. Advances in Neural Information Processing504 Systems, 2018:8279–8288, 2018.505
[77] A. Xu and M. Raginsky. Information-theoretic analysis of generalization capability of learning506 algorithms. Advances in Neural Information Processing Systems, 2017:2525–2534, 2017.507
[78] Z. Yang, Y. Yu, C. You, J. Steinhardt, and Y. Ma. Rethinking bias-variance trade-off for508 generalization of neural networks. In International Conference on Machine Learning, pages509 10767–10777. PMLR, 2020.510
[79] Z. Yao, A. Gholami, K. Keutzer, and M. W. Mahoney. Pyhessian: Neural networks through511 the lens of the hessian. In 2020 IEEE International Conference on Big Data (Big Data), pages512 581–590. IEEE, 2020.513
[80] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning514 requires rethinking generalization. In 5th International Conference on Learning Representations.515 OpenReview.net, 2017.516
[81] H. Zhang, M. Cisse, Y. N. Dauphin, and D. Lopez-Paz. mixup: Beyond empirical risk mini-517 mization. In 6th International Conference on Learning Representations, ICLR, 2018.518
[82] J. Zhang, T. He, S. Sra, and A. Jadbabaie. Why gradient clipping accelerates training: A519 theoretical justification for adaptivity. In International Conference on Learning Representations,520 2019.521
[83] Y. Zheng, R. Zhang, and Y. Mao. Regularizing neural networks via adversarial model perturba-522 tion. In CVPR, 2021.523
[84] R. Zhou, C. Tian, and T. Liu. Individually conditional individual mutual information bound on524 generalization error. arXiv preprint arXiv:2012.09922, 2020.525
Checklist526
1. For all authors...527 (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s528 contributions and scope? [Yes]529 (b) Did you describe the limitations of your work? [Yes] See Section 4, Section 5 and530 supplementary materials.531 (c) Did you discuss any potential negative societal impacts of your work? [Yes] See532 supplementary materials.533 (d) Have you read the ethics review guidelines and ensured that your paper conforms to534 them? [Yes]535 2. If you are including theoretical results...536
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 2537 and Section 3.538
(b) Did you include complete proofs of all theoretical results? [Yes] See supplementary539 materials.540
3. If you ran experiments...541 (a) Did you include the code, data, and instructions needed to reproduce the main experi-542 mental results (either in the supplemental material or as a URL)? [Yes] See supplemen-543 tary materials.544
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they545 were chosen)? [Yes] See Section 5 and supplementary materials.546
(c) Did you report error bars (e.g., with respect to the random seed after running experi-547 ments multiple times)? [Yes] See Section 4 and Section 5.548
(d) Did you include the total amount of compute and the type of resources used (e.g., type549 of GPUs, internal cluster, or cloud provider)? [Yes] See Section 4.550
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...551 (a) If your work uses existing assets, did you cite the creators? [Yes] See Section 4 and552 Section 5.553 (b) Did you mention the license of the assets? [Yes] See supplementary materials.554 (c) Did you include any new assets either in the supplemental material or as a URL? [No]555 (d) Did you discuss whether and how consent was obtained from people whose data you’re556 using/curating? [No]557 (e) Did you discuss whether the data you are using/curating contains personally identifiable558 information or offensive content? [No]559 5. If you used crowdsourcing or conducted research with human subjects...560
(a) Did you include the full text of instructions given to participants and screenshots, if561 applicable? [N/A]562
(b) Did you describe any potential participant risks, with links to Institutional Review563 Board (IRB) approvals, if applicable? [N/A]564
(c) Did you include the estimated hourly wage paid to participants and the total amount565 spent on participant compensation? [N/A]566

Review summaries:
This paper studies the generalization of neural networks trained using stochastic gradient descent (SGD) using information-theoretic bounds. It builds upon existing work that bounds the generalization error using the average square root of the mutual information of the weights with each training datum and related work by Neu 2021 [49] that refines this bound using quantities that correspond to the local gradient sensitivity, the variance/dispersion of the mini-batch gradients, and a stability term that captures the loss at the terminal weights on independent test data. The main contributions of this manuscript are as follows.  1. The paper introduces a coupled auxiliary weight process which refers to the weights updated using additive noise in the gradient descent direction calculated by SGD using the original weights. This allows the authors to study the mutual information of the descent direction with a particular datum in Lemma 5. 2. Theorem 1 establishes that the gradient sensitivity term can be eliminated using an application of Otto and Villani’s HWI inequality that relates the entropy of the weight posterior with the Wassserstein distance and the Fisher information. Theorem 2 and Corollary 1 expand on this result. The main conclusion is that the generalization error can be broken down into two terms, one that depends upon the ratio of the learning rate and batch-size and another that depends upon the stability term referred to above (bounded in Corollary 1 using the trace of the Hessian). 3. The bound in Theorem 2 is verified using a fully-connected network on MNIST and a convolutional network on CIFAR-10.  4. The paper also develops a loss function that training neural networks that averages the loss of the network across different perturbations of the weights by Gaussian noise. This new loss function is evaluated on MNIST, CIFAR-10 and CIFAR-100.

 This paper considers the important problem of understanding the generalization of the SGD algorithm. This paper is a follow up on this paper:   G. Neu. “Information-theoretic generalization bounds for stochastic gradient descent”.  To analyze the generalization, the authors use the mutual information framework of Russo-Zou and Xu-Raginsky where the expected generalization error is upper bounded using mutual information (MI) between output and training set. The main issue with analyzing SGD with MI is that it leads to vacuous bounds due to the fact that there is only degenerate noise in SGD. To circumvent this issue,  the authors use the idea of analyzing a “surrogate algorithm” from Neu’ 21 paper. The surrogate algorithm is a noisy version of the SGD, and instead of analyzing the SGD we can analyze the generalization of the surrogate algorithm. At the end we need to add “pay” for the difference between the output of the surrogate algorithm and SGD.