Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking

Dear Reviewer 6qE1,

Thank you for pointing out the relevance of our work and the clarity of our scope. We will address your concerns individually.

W1: We appreciate the reviewer’s feedback and the opportunity to clarify both our use of terminology and the experimental scope. Grokking so far has been observed in very specific settings [1] which we evaluate in Fig. 1. In our Fig. 4, we observe a delayed generalization similar to grokking. We understand the reviewer’s hesitation regarding our use of the term “delayed generalization” in Figure 4, as validation accuracy does exhibit a modest initial rise. Our intent was not to claim canonical grokking behavior in this case, but rather to highlight that explicit regularization of flatness creates a pronounced delay in generalization, even in settings where grokking is not typically observed. To better reflect this distinction, we now refer to this as “grokking-like” or “induced delayed generalization” and have softened claims where necessary. We appreciate the reviewer’s comment for prompting this clarification.

Our Fig. 1 directly visualizes all three variables (grokking, NC, flatness) together, and we are grateful for the reviewer’s appreciation of this result. Experiments in Sections 5 and 6 are designed to causally probe the individual links between NC, flatness and generalization in settings where grokking is not typically observed. Thus, for those more realistic datasets we cannot observe all three phenomena. Taken together, though, these results form a triangulated empirical case: NC can exist without generalization, generalization can occur without NC, but generalization does not emerge without flatness (in our settings). We are grateful for the reviewer’s feedback and will update the manuscript to clarify these distinctions.

On the NCC-Measure: In our work, we adopt the Class-Distance Normalized Variance (CDNV) metric introduced by [2] to quantify collapse-like behavior. This measure captures the relative tightness of intra-class feature clusters compared to inter-class separation, reflecting key aspects of NC1 (within-class variability collapse) and partially NC2 (inter-class separation). Galanti et al. explicitly motivate this formulation as a tractable and robust alternative to the full neural collapse (NC) definition, especially in practical settings like CIFAR-10 where the idealized geometry is often unattainable. To complement this, we report angular separation statistics between class means as a soft proxy for the ETF structure in NC2, originally formalized in  [3]. That said, we do not claim that our metrics fully capture all aspects of neural collapse. In particular, NC3 (alignment of classifier weights with class means) is not addressed, and NC4 (nearest class center decoding) is only indirectly inferred through clustering geometry. We will revise the manuscript to clarify that our analysis focuses on partial indicators of collapse behavior, following precedent in [2] and related works. At the same moment we want to emphasize that recent work by [4] reconfirms that variability collapse is the first phenomena to appear when networks are not trained to TPT yet, and also proposes a simplified different characterization of NC for overscaled tasks (like natural language processing).

Rationale of Experiments in Sec. 5: We would like to refer to our response to reviewer 8UvS (W1) and reviewer naQM (W1) where we argue that previous work claims a connection between NC and generalization. Our argument in Sec. 5 is that while flatness appears to be necessary for correlation (Sec. 6 and [5]), NC is not, even though both often co-occur. This strengthens our central point that flatness, not collapse, plays the primary causal role. Prior work relies on observations of the (non-)emergence of NC in different regimes. In contrast, we explicitly intervene to suppress NC via regularization and show that generalization is preserved. This kind of causal probing adds precision to what has previously been anecdotal or observational.

Q1: That is an excellent observation. The rising validation loss in many of our experiments is due to our deliberate use of unregularized training setups, i.e., excluding data augmentation, weight decay, and learning rate scheduling. We do so to better isolate the effects of flatness and neural collapse. These conditions are known to amplify grokking-like behavior.

Our ablation studies confirm that, with standard regularization (e.g., modified ResNet, data augmentation, cosine schedule), validation loss typically decreases. Without these, especially under prolonged training, overfitting to cross-entropy loss occurs, even as validation accuracy remains stable, consistent with prior observations.

We will clarify this experimental design in the manuscript and annotate relevant appendix figures accordingly.

Q2: As stated in the previous answer, our goal was not to optimize CIFAR-10 performance, but to study delayed generalization and disentangle the roles of flatness and Neural Collapse. As noted in prior work, grokking dynamics tend to disappear under standard training setups, especially when using data augmentation and other regularization techniques. We also ran experiments using the recommended ResNet-18 modifications with data augmentation (Figure 11, Appendix C). In these settings, delayed generalization is less visible, and the effects of flatness regularization are harder to isolate. While the unmodified ResNet-18 underperforms in absolute terms, this setup is intentional: it creates a simplified regime where flatness and generalization dynamics can be studied more clearly. We will revise the manuscript to clarify this design choice.

Q3: Therelative flatness measure is qualitative and relative, exact values needed for a model to generalize depends on the architecture, task and the achieved representativeness (cf. Appendix D), and therefore clear threshold values cannot be derived.

In our experiments, we observe that models with better generalization consistently converge to lower relative flatness values, even though the raw numbers vary across datasets and architectures. For example, after removing the regularizer, flatness decreases and validation accuracy improves

We agree that the scale inconsistency in the appendix plots made this trend harder to interpret. We thank the reviewer for this observation and will standardize the y-axis ranges and add clearer annotations to highlight key transitions.

[1] Power A, Burda Y, Edwards H, Babuschkin I, Misra V. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177. 2022 Jan 6.

[2] Galanti T, György A, Hutter M. On the Role of Neural Collapse in Transfer Learning. InInternational Conference on Learning Representations 2021 Oct 6.

[3] V. Papyan, X.Y. Han, & D.L. Donoho, Prevalence of neural collapse during the terminal phase of deep learning training, Proc. Natl. Acad. Sci. 2020.

[4] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.

[5] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.

Dear reviewer naQM,

Thank you for appreciating the importance of our study and its relevance for the NeurIPS community, as well as the comprehensiveness of our discussion. We will address your concerns point by point.

W1: Previous work has observed strong co-occurrence of NC and generalization [3,4,5,6] with some hinting at a fundamental relation [2, 7, 8]. In our opinion this literature points to the existing belief that good generalizing networks exhibit NC. In particular, we would like to point out that the quotation from Zhou et al. (2022) is incomplete. It states: “While NC reveals that all losses are equivalent at training time, it does not have a direct implication for the features associated with test data as well as the generalization performance. In particular, a recent work by Hui et al. (2022) shows empirically that NC does not occur for the features associated with test data. Nonetheless, we show through empirical evidence that for large DNNs, NC on training data well predicts the test performance.”

We agree with the reviewer, though, that there is a growing doubt about this connection in the community. Our work is an approach to answer this question more rigorously: (i) We empirically show that NC is not a necessary condition for generalization and that (ii) under certain assumptions NC implies flatness and flatness is necessary for generalization (cf. line 208ff for a discussion of the assumptions). Therefore, under these assumptions NC is a sufficient, but not necessary condition. As we do understand how our formulation could be understood as a mischaracterization of prior work, we will revise this part with more careful language.

W2: We thank the reviewer for continuing this important discussion. We fully acknowledge that the correlation between flatness and generalization has been well studied. However, the key open question is whether flatness is merely correlated with generalization, or whether it is necessary or sufficient.

In our work, we leverage the grokking setting, where memorization and generalization occur at separate times, to explicitly track when flatness emerges. This allows us to disentangle correlation from causation more clearly (Sec. 4, Fig. 1). We go beyond observation, though, by

suppressing neural collapse without harming generalization (Sec. 5), and suppressing flatness, which consistently delays generalization (Sec. 6, Fig. 4).

These interventions support our conclusions that the relationship between flatness and generalization is not merely a correlation.

W3: We would like to clarify that ImageNet-100 is not a small-scale dataset, and it contains a large number of high-resolution images with substantial semantic diversity and is widely used as a standard benchmark in computer vision. Alongside CIFAR-10, it provides a valuable testbed for probing generalization behavior under controlled conditions.

Our aim is not to optimize performance on the largest possible datasets, but to understand the mechanisms driving generalization, particularly the role of flatness and neural collapse. The chosen settings strike a balance between expressiveness and interpretability, allowing us to reliably apply regularization interventions and analyze geometric properties.

Our conclusions are empirically grounded but context-specific, and we avoid claiming generality beyond the observed regimes. Rather than asserting universal necessity of flatness, we show that within the settings studied (next token prediction for arithmetic tasks, and image classification), flatness emerges as a consistent and potentially necessary factor for generalization.

We hope this clarification helps position our work appropriately, as a step toward a deeper understanding of generalization geometry, rather than a comprehensive empirical survey. We appreciate the reviewer’s feedback and will sharpen our position in the camera-ready copy accordingly.

Q1: That is a great observation, indeed there is a gap in validation accuracy even after the flatness regularizer is removed. Simply put, the regularizer pushes the model to a very bad solution. After un-plugging the regularizer, the model has to return to a good solution from this suboptimal point. This is similar to starting from a bad initialization. Unsurprisingly, the resulting model is not optimal, but fairly close. We view this phenomenon as further evidence that flatness must emerge as part of a balanced optimization process, rather than as a late-stage fix. While it might be possible to further improve test performance after un-plugging through careful learning rate scheduling, this would not change the main insights from the experiments.

Q2: "By approximately locally constant labels we mean that, for small delta, the loss in $\delta\| \phi(\xi)\|$-neighborhoods around the feature vector of a training sample $\xi$ is approximated (on average over all training samples) by the loss for constant label $y(\xi)$ on these neighborhoods" [1].That is, locally constant labels mean that within a neighborhood of a training sample the loss is approximated by the loss with respect to a constant label on average over all training samples. While the reviewer is right in the suspicion that label noise violates the (strict) assumption of locally constant labels, the non-trivial, but small, label noise in the dataset does not violate our statistically averaged assumptions and conclusions.

[1] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.

[2] “A Geometric Analysis of Neural Collapse with Unconstrained Features.” Z. Zhu, T. Ding, J. Zhou, X. Li, C. You, J. Sulam, Q. Qu (Neurips 2021)

[3] Galanti T, Galanti L, Ben-Shaul I. On the Implicit Bias Towards Depth Minimization in Deep Neural Networks. arXiv preprint arXiv:2202.09028. 2022 Feb.

[4] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.

[5] Kothapalli V. Neural Collapse: A Review on Modelling Principles and Generalization. Transactions on Machine Learning Research. 2023.

[6] Galanti T, György A, Hutter M. On the Role of Neural Collapse in Transfer Learning. InInternational Conference on Learning Representations 2021 Oct 6.

[7] Súkeník P, Mondelli M, Lampert CH. Deep neural collapse is provably optimal for the deep unconstrained features model. Advances in Neural Information Processing Systems. 2023 Dec 15;36:52991-3024

[8] Mixon DG, Parshall H, Pi J. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis. 2022 Nov;20(2):11.

Dear reviewer 9Eej,

We sincerely thank you for your thoughtful and encouraging review of our work. We greatly appreciate your recognition of the significance of our methodology. You have raised the correct point that we focus on the penultimate layer, yet both NC and relative flatness are not restricted to that layer. Both phenomena, however, have been studied in the penultimate layer in prior work [1,2,3], so focussing on this layer ensures clarity of observations, and frankly, it is computationally efficient. At least for relative flatness, it has been shown that considering the penultimate layer is sufficient to explain generalization [1]. For both phenomena, prior work has also given indication that their behavior is similar in previous layers [1,3] which would support the conjecture that our findings apply to earlier layers as well. We will add a preliminary experiment to test this conjecture using a ResNET18 architecture on CIFAR10 in the camera-ready copy. In general, extending our analysis across layers is an excellent direction for future work.

[1] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.

[2] V. Papyan, X.Y. Han, & D.L. Donoho, Prevalence of neural collapse during the terminal phase of deep learning training, Proc. Natl. Acad. Sci. 2020.

[3] Like H, Mikhail B, and Preetum N. Limitations of Neural Collapse for Understanding Generalization in Deep Learning. Arxiv 2022.

Dear reviewer 8UvS,

Thank you for highlighting the importance of our analysis and theoretical contribution. We would like to address your concerns point by point.

W1: We discuss the preprint of Hui et al. (2022) (cf. line 114). They argue that NC is not correlated with generalization under changes in the training set size. This does not mean there is no correlation between the two wrt. training dynamics. It is one thing to observe whether the NC measure increases when one improves generalization by adding additional training points (for example through an increase of data variability in each cluster), and another to ask whether, for a fixed setting, stronger neural collapse benefits generalization. They also argue that test collapse can hurt transfer learning, which is intuitive and confirmed in other works [1].

In our work, we examine whether for a fixed data distribution and training set, stronger NC correlates with better generalization. Such a co-occurrence has been observed in practice previously [4, 5, 6, 7], with some works indicating a more fundamental relation [2, 8, 9]. Our theoretical contribution shows that NC implies relative flatness, which in turn (under certain assumptions, cf. line 208ff) implies generalization.

W2: We would like to clarify that we do not claim that relative flatness is universally a necessary and sufficient condition for generalization. Petzka, et al. [3] have shown that under certain assumptions, relative flatness is a sufficient condition for generalization and our experiments indicate it is also necessary.

These experiments span CIFAR10 and ImageNet-100 using convolutional (ResNET18), and transformer-based (ViT) models and findings are consistent across settings. We also observe a temporal connection in algorithmic grokking tasks. In particular, our intervention through regularization supports the claim that flatness plays a causal role in generalization. We agree that broader experiments (e.g., on NLP tasks with language models) would further strengthen the case, and we plan to add such experiments to the camera-ready copy of this paper.

W3: That is a very good suggestion. We will add the following intuitive explanation of the proposition to the camera-ready copy. Under the NC conditions, class means are well-separated. Further, the cross-entropy objective pushes class means apart. Then, intuitively, small perturbations of the weights are not sufficient to alter this configuration sufficiently to increase the loss. That is, the loss surface must be flat.

Q1: The spike in relative flatness after removing the regularizer is indeed expected and we are happy to clarify it.

Before epoch 250, the flatness regularizer keeps the optimizer confined in a sharp local minimum, limiting its generalization ability. When the regularizer is removed, the model begins moving toward a flatter, better solution. In this transition, the model may briefly pass through sharper regions, causing a temporary spike in the relative flatness measure. As training continues, the optimizer settles into a flatter region, reflected by a drop in relative flatness and improved validation accuracy.

Limitations: We would like to clarify that our experiments span 2 vision datasets and one symbolic algorithmic task, as well as convolutional, vision-transformer, and standard transformer architectures.

[1] Md Yousuf H, Jhair G, Christopher K. Controlling Neural Collapse Enhances Out-of-Distribution Detection and Transfer Learning. The International Conference on Machine Learning. 2025.

[2] “A Geometric Analysis of Neural Collapse with Unconstrained Features.” Z. Zhu, T. Ding, J. Zhou, X. Li, C. You, J. Sulam, Q. Qu (Neurips 2021)

[3] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.

[4] Galanti T, Galanti L, Ben-Shaul I. On the Implicit Bias Towards Depth Minimization in Deep Neural Networks. arXiv preprint arXiv:2202.09028. 2022 Feb.

[5] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.

[6] Kothapalli V. Neural Collapse: A Review on Modelling Principles and Generalization. Transactions on Machine Learning Research. 2023.

[7] Galanti T, György A, Hutter M. On the Role of Neural Collapse in Transfer Learning. InInternational Conference on Learning Representations 2021 Oct 6.

[8] Súkeník P, Mondelli M, Lampert CH. Deep neural collapse is provably optimal for the deep unconstrained features model. Advances in Neural Information Processing Systems. 2023 Dec 15;36:52991-3024

[9] Mixon DG, Parshall H, Pi J. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis. 2022 Nov;20(2):11.