A simple connection from loss flatness to compressed representations in neural networks

Questions:

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.

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.

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.

We appreciate the reviewer’s critiques and suggestions. Below we address the reviewer’s concerns point by point:

Weakness:

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.

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.

We believe that our updated Sec 3 and Appendix D now explain the looseness of the original bound that is evidenced by results in Appendix D. In summary, we found that the reason for the looseness in Eq. 12 is due to the sparseness of the eigenspectrum of $C_f^{lim}$. We again observe that in Fig. 3, although log volume is not well correlated with sharpness, the quantity G is much better correlated with sharpness. Exploiting this, we introduce a new, closely related metric, the maximum local sensitivity (MLS), that allows for a much tighter bound (Sec. 3.2). We then newly empirically test and show the correctness and tightness of this bound (Appendix D.1).

Minor:

1. We have updated the figure positions to where the figure is mentioned.
2. We have removed the short mention of possible perspective through the lens of entropy as suggested.
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.

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.

We thank the reviewer for their critiques and suggestions, and their overall assessment of the value of our central contribution in newly linking flatness and compressed representations.

Below we address the reviewer’s concerns point by point:

Weakness:

We thank the reviewer for pointing out that sharpness may not fully explain volume compression. To bring this out, in our revision we note in the updated Sec. 3.2 that the bound in Eq. 11 is indeed loose (L162): “We observe that the equality condition in the first line of Eq. (11) rarely holds in practice, since to achieve equality, we need all singular values of the Jacobian matrix to be identical.”. This critique has led us to propose a more precise metric that measures the maximum sensitivity along any direction: we call this maximum local sensitivity (MLS). We show that a tighter bound can be derived for MLS (Eq. 13), and we empirically test the tightness of the bound in Appendix D.1. We still note, however, that the volume is positively correlated with sharpness in all of our experiments, so that -- while we acknowledge that significant variability remains to be explained -- sharpness does predict an overall trend that aligns with our main message.

We appreciate the point that more extensive experimentation is needed to make a stronger case. As the reviewer suggested, we have now added experiments on MLP trained on FashionMNIST dataset (Fig. D.6, E7-8 in the appendix).

Questions: Formatting: We thank the reviewer for spotting this. Equations are now consistently referred to as "Eq. (N)," and figure references are now enclosed in parentheses as suggested.

Placement of figures: we agree, and the positions of figures are now adjusted to be close to where they are mentioned

Sharpness vs generalization: We acknowledge the contentious nature of the topic and appreciate this point. In our revision, we have duly removed or rephrased any overly assertive statement that states sharpness directly leads to generalization. Moreover, we have added substantial additional literature and allied discussion -- including a summary of the mixed and negative results relating sharpness and generalization (second paragraph of Introduction, starting from L37). Additionally, in Appendix D.1 we newly enumerate conditions that may contribute to such mixed and negative results. Finally, in Appendix D.2, we empirically test correlations among a wide set of different metrics, further connecting to the broader literature while confirming that our experiments on the VGG10 and MLP networks studied here both do show a positive correlation between sharpness and test loss.

We appreciate the reviewer’s critiques and suggestions. Below we address the reviewer’s concerns point by point:

Weakness:

1. 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.
2. 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.”.
3. 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.”
4. 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.
5. We apologize for the lack of clarity here -- we have fixed the phrasing in the caption.

Overall, thanks to the reviewer comments the presentation has now been substantially improved.

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

Question: 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.