Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: 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.
main_review: Main comments

1. This paper is very closely related to [49]. It is understood that this reference is, as yet, unpublished. The authors should however discuss the novelty compared to respect to this existing work.
2. Line 214: Arguably, the red line in Fig. 1 a,b does not decrease but rather settles down to a steady value as the model size increases. Is this an experimental artifact?
3. How do you explain the fact that the upper bound in Theorem 2/Lemma 6 increases with T, i.e., is larger for longer trajectories? In practice, the generation gap clearly remains bounded as the number of training epochs increases. The authors come close to discussing this around Line 230 but it is necessary to expand upon this for bounds such as those in this paper.
4. It is possible that the complex relationship between the learning rate, number of training steps T and the upper bound stems from the fact that the gradient dispersion term is not “normalized“. This term (Line 123) does not distinguish between the gradient dispersion being small somewhere at the beginning of training (say all per-sample gradients have similar magnitudes and point in similar directions) versus the term being small at the end (which would indicate training finding a critical point that is close to the critical points of the per-sample loss functions). It may be possible to separate out these two phenomena.
5. The above point also calls into question the importance of the “double descent of gradient dispersion” phenomenon explored in Figure 3. Indeed, for noisy samples, it has been seen that there is a long phase in training where the per-sample gradients show a strong mismatch and the training/test loss plateaus (https://arxiv.org/abs/2101.06509). It is unclear what the implication of this double descent observation is and how it is related to the rest of the paper.
6. I have concerns about the experiments in Table 1. First, such loss functions have been used numerous times in prior work (the authors do cite papers on flat minima so they are aware of algorithms like https://arxiv.org/abs/1611.01838, https://arxiv.org/abs/1609.04836 etc.); so this loss is not really novel. Second, networks trained with this new loss generalize about as well as algorithms like AMP, so this experiment is more of a sanity check than a new contribution. It would be good for the manuscript to demarcate the new contributions explicitly in the introduction.

In conclusion, there are interesting theoretical parts in this manuscript but (i) it is difficult to ascertain whether the theory is novel enough to merit publication, and (ii) the experiments seem to be essentially disjoint of the theory. Perhaps, the results in Theorem 3 could be fleshed out and added to the main text.
limitations_and_societal_impact: Adequate.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This paper proves a new information-theoretic upper bound for the generalization of neural networks trained with SGD. The authors claim that the developed new bound provides new insights on the SGD and propose a new regularization scheme which can match the state-of-the-art result.

main_review: After reading this paper it seems that the main contribution is the improvement of the generalization bound compared to [49, Theorem 1.1]. In particular, the authors prove that the “local gradient sensitivity” term in [49] is unnecessary. However, it is not clear whether this improvement is significant. the authors may need to discuss which term in Lemma 3 could be the dominating one (i.e., comparing the magnitude of $\Psi(W)$ and $V(W)$). Besides, the authors should also include the bound in Lemma 3 in the experiment (e.g., Figure 1) to better interpret the improvements.

From Theorem 2 it can be seen that the generalization error becomes smaller when increasing the mini-batch size. Consider an extreme case where $b=n$, then the algorithm reduces to full-batch gradient descent. The authors did discuss a bit in Section 4, but a theoretical explanation could be better. Besides, why not directly compare the proved generalization bound or test accuracy for different batch sizes as that could be more related to the theory of this paper.

limitations_and_societal_impact: This paper has no negative societal impact.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: Improves existing information-theoretic generalization bounds for neural networks trained with SGD, and explores a new double descent phenomenon. Application consisting of GMP interesting but not competitive to best baselines.
main_review: This paper improves a bound in [49], removing a term therein by using a variant of the HWI. This bound is further improved from n^-1/2 to 1/n rate. These improvements, while not requiring a wholly new analysis, do seem to be significant improvements, and the authors discuss the interpretation and implications of the new bound thoroughly. I therefore recommend acceptance.

Line 61-62 can be removed, it is sufficient to mention additional material in the supplement in line.
limitations_and_societal_impact: no social impact
review_title: null

final_score: rating: 4: Ok but not good enough - rejection
review_content: summary: 
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.   

main_review: Contribution and Positive points:

The main contributions of the authors on the top of Neu’ 21 paper are obtaining a tighter bound using the mutual information bound of Bu et al. in Theorem 2. Note that the construction of the surrogate algorithm is exactly the same as the paper by Neu'21.
An interesting point about the main results is that the generalization bound both depends on the trajectory of SGD and the final solution. In particular, the variance of the gradients along the pass captures how trajectory affects the generalization. Then, in Corollary 2 they have a nice result which shows the sharpness of the final solution impacts the generalization. Finally, they have an extensive empirical study of their bound on different benchmarks. 

However, I have some major comments and some of them are concerned with the validity of their proof.

Weakness and Questions:

-I think the introduction needs to be revised. It is not clear how this work relates to other work in the literature of understanding generalization of SGD. Note the problem of understanding the generalization of SGD has been approached using lots of different tools and techniques. 

-General comment about the paper: In the abstract and in the introduction, it has been emphasized that this paper aims to derive generalization bounds for neural networks trained with SGD. However, it is not clear how the quantities appearing in the main theorem in the paper (Theorem 2) depend on the neural networks parameters. Specifically, I could not find anything specific about neural networks in the main results. In fact the results hold for any learning objective which is optimized with SGD. This is in contrast with some of the known results on the generalization of neural networks which connect the generalization to the specific parameters of the neural networks.

-Corollary 2: One important issue with this result is the dependence on the dimension. This is in contrast with the classical analysis of SGD using the stability framework which is dimension-free. How do you compare the stability method and your method? 

-Corollary 1: Assume we want to study the generalization error of a NN trained with SGD with respect to 0-1 loss. 0-1 loss is subG with parameter ¼. However, it is not continuous. Therefore, the bound in corollary 2 can’t be used for 0-1 loss function since it does not have second-derivative. Because the Gamma term is defined with respect to the loss function. How does this affect the conclusion of the paper on the impact of the flatness of the final solution on the generalization? 

-Diagonal covariance: Adding gaussian noise seems to be reasonable. However, why do not you consider the non-diagonal covariance? 

-Stochastic Convex optimization: Assume that the loss function is convex, lipschitz, and smooth. How can you use this information to design the noise variances? It seems to me that this framework in this regime is much looser than Hardt et al results specifically because of the dependence on the dimension. 

-Line 205: In your numerical results what do you mean by setting the sub-gaussian parameter to 0.1? If you use 0-1 loss then the subgaissian parameter is ¼. 

-It seems you use corollary 1 for empirical studies. As pointed out, this is not valid, it is not valid for 0-1 loss function.

Comments of the proof:


-The conditional expectation notation in the paper should be revised. For instance in Eq (9) the notation is not correct and it makes it very difficult for readers to understand it. The outer expectation is over S,S’ but inside there is a conditioning on the specific realizations. The same error appears in the rest of the proof.

-I have difficulty understanding  Eq(6) - Eq(10). Specifically, what do you mean by “we do not use an independent copy of \tilde{W}”?  I believe the derivation in Neu’ 21 is correct (beginning of Page 8 in the V1 arxiv version of Neu’ 21). My understanding is as follows: P(W|S) is a kernel which shows the learning algorithm: specifically given a training set, it outputs a distribution on W. Therefore, when you have two independent training sets S and S’, they will generate two independent copies of W and W’ whose law are given by P(W|S) and P(W’|S’). Specifically Eq(8) in your paper is not correct. This is an extremely serious issue and I would like to hear your detailed response on this.

-Could you please elaborate on Eq (11)? Eq(10) follows from the monotonicity of the KL divergence or Data Processing. In Eq (11) you wanted to use the chain rule for KL. But the caveat here is that the surrogate algorithm \tilde{W} is NOT Markovian as shown in the figure on Page 4 of your paper.  As a result,  Eq (11) in its current form is NOT CORRECT. I would appreciate it if you could comment on it since it is an important step in the proof and it might affect the validity of the proof.

-I think in Eq. 11 you wanted to use the tower rule but you forgot to write the outer expectation. 

limitations_and_societal_impact: I think the limitations of this work is not discussed thoroughly in the paper. The authors can refer to my questions above to find suggestions for improvement.
review_title: null

Meta-review: This paper proves an information-theoretic generalization error bound of neural networks trained by SGD. The main concern of the reviewers is about the novelty and significance of the results derived in this paper because this paper is heavily built upon recent work in COLT 2021. Even after the author response, this paper does not receive sufficient support from the reviewers. Therefore, I recommend rejection.
Decision: Reject