Paper: A simple connection from loss flatness to compressed representations in neural networks (ICLR-2024)

Reviewer comment (Novelty):
The paper gives an impression of an unfinished work. It is not clearly signified which of the derivations are novel contributions and which are reproducing the existing research. The experimental evaluation is very limited and not fully showing the correctness/tightness of the bound presented.

Author response:
Novelty:  While we acknowledge and appreciate the reviewer’s point that our contributions could and should be better described, our impression from other reviewers is that the central novelty of newly connecting the previously separate ideas in parameter and representation space was laid out reasonably.  This said we embrace the criticism and the chance to do better.  First, in our revision, we have now clearly divided the material of the paper into two distinct sections. Sec. 2 is titled “Background and setup” and Sec. 3 highlights our original contributions in making new connections to neural representation space.  Second, also following the reviewer’s other very on-point suggestions, we have extended these original contributions in Sec. 3 significantly, with three new figures there and the allied appendix D (please see below).  Third, we now enumerate our contributions more explicitly in the introduction.  Overall, we believe that this has significantly improved our paper, and again thank the reviewer for their keen criticism and the chance to respond.

Correctness/tightness of the bound: We thank the reviewer for their keen and on-target point that the tightness of the bound should be both more fully explored and presented.  In our revision, we have extended our analysis to tackle this important matter in two separate steps. First, we studied very carefully the tightness of our newly derived bound (Sec. 3.2 Eq. 13) in Appendix D.1, including isolating and testing three factors that contribute to this tightness or the lack thereof (with full details given in the new Appendix D.1, with a series of new equations and a new figure). Second, we now clearly connect the relationship between our bound and the insights derived from allied research (Appendix D.3). Third, we carry out and present new experiments on a different network (a fully connected MLP) (Fig D.6, E.7-8). Altogether our extended findings reveal both how our bound can predict that neural representations will be compressed, and also the factors that would limit the success of this prediction.  Further, dedicated future work will be required to assess when and how flatness increases or decreases through network training. We further detail our findings, again inspired by the reviewer's critique, with two new figures (Appendix D.2) that show how in our experiments a decrease in sharpness (and related increase in flatness) significantly correlated with volumetric compression of the neural representation.

Reviewer comment (Theory):
Even with the assumption that trace of the Hessian is indeed a good generalization indicator (which was disproven by Dinh et al.), the discussion in the beginning of section 2 is very imprecise. If the loss indeed reaches the point of zero value and its derivatives are zero, SGD would not change the solution anymore - just by definition. The referenced works take into account label noise, for example, that allows for further dynamic of the optimization even after reaching a minimal loss. Therefore the theoretical justification of the bounds is ill-posed and it also does not connect to the experiments performed, since the network there does not achieve 0 loss and 0 gradients (which should effectively stop the changes of the model).

Author response:
Imprecise argument in section 2: We thank the reviewer for pointing out the lack of precision of the argument at the start of section 2. We addressed this in two separate ways: 1) We now evaluate in more depth whether sharpness decreases towards the end of training. For example, sharpness does not necessarily decrease after the training error becomes zero (Fig. E.7-8). However, our additional experiments show that our bounds successfully predict the correlation between sharpness and both local volume and maximum local sensitivity (MLS) (Fig. D.5-6). 2) Our theoretical derivation and results do not rely on the specific form of label noise SGD or sharpness-aware minimization (SAM). Both the narrative and results of the paper are largely unaltered. The mathematical argument also has a general form that does not strictly rely on these specific forms of gradient noise ($\delta\theta^*$). We believe that it provides a constructive argument towards why flatness may/may not increase at the end of training (sharpness may not decrease when the noise is too small, i.e. the second term $\delta\theta^* H\delta\theta^*$ is negligible). Nevertheless, we omit this intuition from the paper itself to retain the focus on our main contributions. We believe these two points address the reviewer's concern.

Reviewer comment (Presentation):
Equation 14 is introduced at the very end of the paper, while being used in all the experimental section, which makes reading unnecessarily hard.

Author response:
3. We have moved (previous) Eq. 14 and the corresponding definition of local dimensionality to Sec 3 where all the metrics in the feature space are defined.

Reviewer comment (Theory):
What model of SGD dynamics do you have in mind when using terms "second phase" and "interpolation regime"?

Author response:
1. All of our experiments use vanilla SGD. We would like to reiterate that our main contribution has no dependence on which optimization scheme is used.

Reviewer comment (Theory):
How does the used volume connect to the representation collapse?

Author response:
2. Neural collapse is a phenomenon that occurs in deep neural networks when the training error reaches zero and the network enters the terminal phase of training. At this point, the within-class variability of final hidden layer outputs becomes infinitesimally small. This means that the distribution collapses down to a point or a very compressed manifold. The volume of the representation in neural space, which we study here, is one of the metrics that can capture the fact that such within-class variability is small.

Reviewer comment (Evaluation):
What is the additional metric introduced in the last section of the paper?

Author response:
3. The metric is called participation ratio which is a natural continuous measure of the local dimensionality.  Intuitively if all variance is concentrated in one direction, i.e. the covariance matrix $C_f^{lim}$ only has one non-zero eigenvalue, then the output of the network is locally 1D given input of a small ball B(x) around x. This is a metric that captures the neural compression effect through the lens of the dimensionality of the manifold of within-class representations. The participation ratio overall has been recently applied fairly broadly in the computational neuroscience and neural networks literature, making us confident that its analysis extends the value of the paper; moreover, we have updated the manuscript to include a dedicated subsection 3.1 to more completely describe and discuss this metric.

Reviewer comment (Experiments):
There could probably be more empirical analysis.

Author response:
We appreciate the reviewer’s overall assessment of the value of our central contribution in linking flatness and compressed representations.  We also appreciate their stated limitation and suggestion, and next, describe how we have revised and added to our paper in response:

Weakness:

We agree with the reviewer’s suggestion that more empirical analysis should be included. In our revision, we have added new experiments on a different network (MLP) trained on the FashionMNIST dataset (Fig. D6, E7-8).  We have also added a new experimental Appendix D that comprehensively analyzes the tightness of the derived bound in a new figure, as well as another new figure showing correlations between different metrics introduced in the paper. We also agree that further extensions to other learning problems would be of interest, and added a sentence describing this in the revised discussion; we feel that completing these studies here would go beyond the scope of what we can accomplish within the given paper but are eager to undertake this in future studies.

Reviewer comment (Experiments):
The analysis could be repeated for additional learning problems, beyond the simple classification problem studied.

Author response:
We appreciate the reviewer’s overall assessment of the value of our central contribution in linking flatness and compressed representations.  We also appreciate their stated limitation and suggestion, and next, describe how we have revised and added to our paper in response:

Weakness:

We agree with the reviewer’s suggestion that more empirical analysis should be included. In our revision, we have added new experiments on a different network (MLP) trained on the FashionMNIST dataset (Fig. D6, E7-8).  We have also added a new experimental Appendix D that comprehensively analyzes the tightness of the derived bound in a new figure, as well as another new figure showing correlations between different metrics introduced in the paper. We also agree that further extensions to other learning problems would be of interest, and added a sentence describing this in the revised discussion; we feel that completing these studies here would go beyond the scope of what we can accomplish within the given paper but are eager to undertake this in future studies.

Reviewer comment (Writing):
L73: The Frobenius norm is defined, yet it is not clear for what purpose it is defined.

Author response:
We thank the reviewer for pointing this out -- the Frobenius norm should be defined for the norm in Eq. 3.  We have now moved the explanation to right after Eq. 3.

Reviewer comment (Writing):
L94: "Let W be the input weights to the network, and θ¯ the corresponding set of parameters." – it is not clear what W ("Input Weights") is.

Author response:
Here, “W” are the weights of the linear layer. The setting is the same as Ma & Ying 2021, and as in many deep learning architectures, where the raw features go through a linear layer first. We have updated the text L117 to make this point clearer: “Let W be the input weights (the parameters of the first linear layer) of the network, and (theta bar) the rest of the parameters.”.

Reviewer comment (Theory):
L112: The statement "This amounts to asking whether the transformation of the differential volume around x¯ is an expansion or a contraction" is problematic, as neural networks often change the dimensions between layers, and this is not encapsulated in this measure as far as I understand.

Author response:
We acknowledge that we are not considering the dimension of layers between the input and the output in our work. However, our work applies when we redefine the function f to be the transformation up until any middle layer that may be of interest. We have updated the text at the beginning of Sec. 3 to reflect this: “Although we only consider the representations of the output of the network, our results apply to representations of any middle layers by defining f to be the transformation from input to the middle layer of interest.”

Reviewer comment (Writing):
The term "The interpolation phase" is usually used for the phase of zero **error**, not zero loss (e.g., Neural Collapse literature).

Author response:
We thank the reviewer for pointing out this difference. We have updated the manuscript to note that the solution found by SGD is indeed an “approximate interpolation solution” with zero training error instead of zero training loss due to the gradient noise. Therefore our reference to the interpolation phase in experiments is consistent with those in the neural collapse literature.  In our revision we have also added and proved Lemma A.1, to show that our estimate of the sharpness won’t invalidate our results during the interpolation phase given that the training loss is small enough.

Reviewer comment (Novelty):
The paper is missing relevant literature, including:  
[1] Ben-Shaul, I. & Dekel, S. (2022). Nearest Class-Center Simplification through Intermediate Layers. *Proceedings of Topological, Algebraic, and Geometric Learning Workshops 2022*, in *Proceedings of Machine Learning Research* 196:37-47. Available from https://proceedings.mlr.press/v196/ben-shaul22a.html.  
[2] Ben-Shaul, I., Shwartz-Ziv, R., Galanti, T., Dekel, S., & LeCun, Y. Reverse Engineering Self-Supervised Learning. In Proceedings of the Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023).  
[3] Rangamani, A., Lindegaard, M., Galanti, T., & Poggio, T.A. (2023). Feature learning in deep classifiers through Intermediate Neural Collapse. *Proceedings of the 40th International Conference on Machine Learning*, in *Proceedings of Machine Learning Research* 202:28729-28745. Available from https://proceedings.mlr.press/v202/rangamani23a.html.  
[4] Gamaleldin F. Elsayed, Dilip Krishnan, Hossein Mobahi, Kevin Regan, and Samy Bengio. Large margin deep networks for classification. In NeurIPS, 2018.  
[5] Galanti, T., Galanti, L., & Ben-Shaul, I. (2023). Comparative Generalization Bounds for Deep Neural Networks. Transactions on Machine Learning Research, (ISSN 2835-8856).

Author response:
In addition, we thank the reader for the missing literature, all of which we have cited in the updated version (L24, L29).

Reviewer comment (Experiments):
The paper would be improved by adding much more experimental evidence of the effects, explaining them more thoroughly, and clearly.

Author response:
We thank the reviewer for this suggestion and have added

1) additional experiments on MLP training on FashionMNIST dataset (Fig. D.4, D.6, E7-8) 
2) a new experimental appendix D that comprehensively analyzes the tightness of the derived bound in a new figure
3) a new analysis, and related figure, for the correlations between different metrics introduced in the paper and accompanying discussion of the connections of our bound to other works. 

The contributions of these additions are:
1. We identify two representation-space quantities that are bounded by sharpness -- volume compression and the newly added maximum local sensitivity (MLS) -- and give explicit formulas for these bounds (Sec. 3.1, 3.2, Eq. 12-13).
2. We conduct empirical experiments with both a VGG and an MLP network and find that volume compression and MLS are indeed strongly correlated with sharpness (Figs D.4-6, and Appendix D.2).
3. We find that sharpness, volume compression, and MLS are also correlated, if more weakly, with test error and hence generalization (Figs D.5-6, and also Appendix D.2).
4. We carefully inspect the condition when the equality holds for the bounds that we derived, and give a set of conditions that determine when MLS will and will not be tightly correlated with sharpness (Appendix D.1) and, we conjecture, also generalization. We find that when the conditions are (approximately) satisfied, the quantities on both sides of the inequality are highly correlated.  

We hope that the reviewer will agree with us that these additions and improvements make a substantially stronger case, and thank them for their role in helping us to make it.