Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking
Authors: Ting Han
Linara Adilova
Henning Petzka
Jens Kleesiek
Michael Kamp
Venue: NeurIPS 2025 poster ()
Keywords: Generalization
flatness
neural collapse
Primary area: general_machine_learning
TL;DR: We use grokking to disentangle generalization from training dynamics and show that relative flatness, not neural collapse, is a necessary and more predictive indicator of generalization in deep networks.

Abstract: Neural collapse, i.e., the emergence of highly symmetric, class-wise clustered representations, is frequently observed in deep networks and is often assumed to reflect or enable generalization. In parallel, flatness of the loss landscape has been theoretically and empirically linked to generalization. Yet, the causal role of either phenomenon remains unclear: Are they prerequisites for generalization, or merely by-products of training dynamics? We disentangle these questions using grokking, a training regime in which memorization precedes generalization, allowing us to temporally separate generalization from training dynamics and we find that while both neural collapse and relative flatness emerge near the onset of generalization, only flatness consistently predicts it. Models encouraged to collapse or prevented from collapsing generalize equally well, whereas models regularized away from flat solutions exhibit delayed generalization, resembling grokking, even in architectures and datasets where it does not typically occur. Furthermore, we show theoretically that neural collapse leads to relative flatness under classical assumptions, explaining their empirical co-occurrence. Our results support the view that relative flatness is a potentially necessary and more fundamental property for generalization, and demonstrate how grokking can serve as a powerful probe for isolating its geometric underpinnings.

Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking
Ting Han
Linara Adilova
Lamarr Institute, TU Dortmund, Germany and Institute for Al in Medicine, UK Essen ting- han@tu-dortmund. de
Research Center Trustworthy Data Science and Security of the University Alliance Ruhr, TU Dortmund, Germany
Henning Petzka Ruhr University Bochum, Germany
Jens Kleesiek
Institute for Al in Medicine, UK Essen and Department of Physics, TU Dortmund, Germany
Michael Kamp
I Lamarr Institute, TU Dortmund, Germany and Institute for AT in Medicine, UK Essen mi michael.o kramptuuddrtmunt de
Abstract
Neural collapse, i.e., the emergence highly symmetric, class-wise clustered representations, frequently observed in deep networks and is often assumed reflect or enable generalization. In parallel, flatness of the loss landscape has been theoretically and empirically linked generalization. Yet, the causal role of either phenomenon remains unclear: Are they prerequisites for generalization, or merely by-products of training dynamics? We disentangle these questions using grokking, a training regime in which memorization precedes generalization, allowing us to temporally separate generalization from training dynamics and find that while both neural collapse and relative flatness emerge near the onset of generalization, only flatness consistently predicts Models encouraged collapse or prevented from collapsing generalize equally well, whereas models regularized away from flat solutions exhibit delayed generalization, resembling grokking, even in architectures and datasets where it does not typically occur. Furthermore, we show theoretically that neural collapse leads to relative flatness under classical assumptions, explaining their empirical co-occurrence. Our results support the view that relative flatness is a potentially necessary and more fundamental property for generalization, and demonstrate how grokking serve a powerful probe for isolating its geometric underpinnings.!
1 Introduction
Overparameterized neural networks continue to challenge classical learning theory. Despite their capacity to memorize arbitrary labels [Zhang et al., 2017], they often generalize well on natural data and, in some cases, only begin to generalize after complete memorization, a phenomenon known grokking [Power et al., 2022]. This apparent contradiction has motivated renewed interest in identifying geometric signatures of generalization. Two candidates have emerged as particularly prominent: the flatness of the loss landscape [Keskar al., 2017, Jiang al., 2020, Petzka al.,
Code implementation: https///gttaalccmmmm/ssstttttttt Lanegg//eneeiitaaaaaaaaaaaaoonnn Lab/grokking flatness.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2021] and the neural collapse (NC) phenomenon [Papyan et al., 2020, Mixon et al., 2022, Zhou et al., 2022]. Both tend to appear late in training and are frequently associated with good generalization. yet their precise causal roles remain unclear.
In this paper, we challenge the common conjecture that NC is necessary for generalization [Mixon et al., 2022, Zhou et al., 2022, Súkeník et al., 2023, Zhu et al., 2021]. Our key idea is to leverage grokking as a unique observational window: since generalization emerges only after prolonged memorization, it enables a clean separation between training dynamics and generalization. This allows us ask: do NC and flatness merely correlate with generalization, or do they contribute to it?
We measure both phenomena in the penultimate layer. For flatness, we adopt the relative flatness metric proposed by Petzka et al. [2021]. which considers the Hessian trace normalized by the weight norm, a quantity theoretically and empirically shown to align with generalization. To make this computable in large -fate-of-the-art neural networks, we employ the alternative closed-form upper bound introduced by [Walter et al., 2024), which is valid in the penultimate layer under cross- entropy loss. For NC, we measure the class-wise clustering of penultimate-layer representations through empirical variance and mean properties, tracking its emergence throughout training. In grokking experiments, we find that flatness emerges alongside generalization, never before. NC, in contrast, appears earlier during memorization, suggesting that it can lead to, but is not neccessary for generalization (cf. Fig. 1).
To isolate their functional roles, we per- turb training dynamics. When we sup- press collapse-style clustering, networks still generalize robustly; test accuracy remains high despite high clustering val- ues. Flatness, however, unaffected, In contrast, regularizing against flatness, i.e., encouraging sharp solutions, reli- ably delays generalization and induces grokking-likk behavior even in settings where it does not typically arise, in- cluding ResNets on CIFAR-10, ViTs on ImageNet and various pre-trained lan- guage models on SST-5. These findings point to flatness as a more fundamental driver of generalization. To explain the frequent co-occurrence of NC and flat- ness, we provide a theoretical result: un- der classical NC assumptions, collapse leads to relative flatness. This unifies previously observed correlations under a single geometric framework, but also shows that flatness arise via other mechanisms. NC is thus a sufficient but non-exclusive path toward flatness. Fi- nally, we revisit the theoretical frame- work of Petzka al. [2021], which as- serts that generalization requires not only flatness, but also representativennsss, i.e., the alignment between learned features and the true data distribution. Our experiments confirm this dependency. Representativeness, when approximated, improves only as generalization sets in, highlighting that neither flatness nor collapse alone ensures generalization without meaningful feature alignment.
XXX
Relative flatness
Training acturacy
19
Validation accuracy
Neural collapse clustering
Accuracy Neural Colapse Clustering
Relatove Flatness
Training scale)
Figure I: Neural collapse clustering and relative flatness in grokking. While both correlate with generalization, neu- ral collapse emerges early during memorization, whereas flatness only drops sharply when generalization begins, highlighting flatness as a better indicator of generalization onset.
Together, these results place NC in a new light. While it often accompanies generalization, it is not prerequisite. Instead, it may serve as inductive bias toward flatness under standard training dynamics. Relative flatness, by contrast, tends to be empirically necessary and theoretically justified, especially when labels are locally constant and features are representative assuutpiionssssssiiininn which appear to hold in our experimental settings.
Our work provides new insight into the geometry of generalization:

---

We show empirically that relative flatness, but not neural collapse, is potentially necessary for generalization;
we prove that neural collapse leads to relative flatness under classical assumptions;
we demonstrate that grokking-like behavior can be actively induced by manipulating flatness, even in real-world architectures.
These insights challenge prevailing narratives about generalization in deep learning and suggest new levers for training, regularization, and model diagnostics.
2 Related Work
Despite substantial empirical success, the mechanisms driving generalization in deep neural networks remain poorly understood. Two pronounced geometric perspectives, neural collapse and loss surface flatness, have emerged as phenomena empirically correlated with generalization, but their interrelation and causal roles remain ambiguous. In this work, we leverage the properties of grokking to disentangle necessity from sufficiency and assess which of these phenomena are required for generalization.
Neural Collapse and Generalization Neural collapse (NC) describes geometric convergence dur- ing training a classifier where penultimate layer features from the same class collapse to a single mean, the means form a simplex equiangular tight frame, and classifier weights align accordingly [Papyan et al., 2020]. This phenomenon has been widely observed at late training stages in overparameterized networks [Han al., 2022, Graf et al., 2021, Rangamani and Banburski-Fahey 2022, Wu and Papyan, 2024, Zhou et al., 2025] and has become a widely accepted indicator for generalization and robustness. Under idealistic assumption of unconstrained features model, i.e., model which optimizes features as free variables, NC has been shown to be the global optimal solution for various loss functions [Mixon al., 2022, Zhou et al., 2022, Súkeník et al., 2023, Zhu et al., 2021]. Similar results for compatibility of zero test error and collapse of variances (so only NC1 condition) were shown for mean-field networks [Wu and Mondelli, 2025]. Practical applications have leveraged NC properties for downstream performance. For instance, Wu et al. [2025] use a feature separation loss based on NC to improve out-of -is-iitribtttinn detection. Similarly, Munn al. [2024] show that lower geometric complexity in pre-trained representations fosters NC and leads to better few-shot generalization in transfer learning. Galanti et al. [2022] shows theoretically and empirically that NC is the reason for transfer learning to work better. These findings suggest that NC may be sufficient for generalization in certain regimes. However, the open question remains: is NC necessary outside of the unconstrained features models? For example, NC is debated matter in the context of inbalanced learning, where it was observed that minority collapse can happen, i.c., the means of the minority classes merge, thus preventing classification [Fang et al., 2021]. ]. At the same moment it is possible to induce NC by enforcing it through training, even in the unbalanced data setup [Yang et al., 2022]. For deep linear networks it was also theoretically shown that the global optimum exhibits NC, but even more: Under unbalanced data, NC gets means proportional to the class size [Dang et al., 2023]. Hui et al. [2022] argues for measuring NC on the test set for understanding generalization properties and demonstrates that test set NC is hurting generalization performance on related downstream tasks, i.e., representations with strong NC on the test set bad for transfer learning. Moreover, Hong and Ling [2024] show that sufficiently high signal-to-noise ratio is a necessary condition for NC to correlate with good generalization, explaining why strong NC can otherwise coincide with poor transfer performance.
Overall, the existing research suggests that NC is common, but might not be required for generaliza- tion: a network can generalize without NC and some setups (like noisy labels) can break the link between generalization and NC. Our work addresses this question directly by constructing settings in which models generalize without exhibiting NC, demonstrating that it is not a prerequisite for generalization.
Flatness and Generalization long been suspected [Hochreiter and Schmidhuber, 1997], and more recently has been supported by empirical studies showing that flatness-based measures often outperform alternatives such as norm-, margin- or optimization-base metrics [Jiang et al., 2020, Keskar et al., 2017]. However, classical flatness measures. typically based on the Hessian of the loss with respect to parameters,
The connection between loss surface flatness and generalization has
3

---

are known to be sensitive to reparameterizations that leave the function and generalization behavior unchanged [Dinh et al.. 2017]. and computationally expensive for large state-of-the- models. To resolve this, Petzka et al. [2021] introduced a reparameteriiaaiion-nvvaae grounded in a theoretical framework that connects flatness of an individual layer to generalization under assumptions of representative data and locally constant labels. This formulation not only explains prior empirical correlations, but also provides theoretically grounded conditions under which flatness predicts generalization. Walter et al. [2024] further made this measure more practical by deriving closed-form computation formulas in standard classification settings. In our work, we empirically verify that relative flatness drops precisely when generalization emerges during grokking. Moreover, we demonstrate that promoting relative flatness can actively induce generalization even in the absence of NC. In doing so, we confirm and extend the theoretical framework of Petzka et al. [2021], while disentangling flatness from NC.
measure of relative flainess,
Grokking as a Window into Generalization Grokking refers to the sudden emergence of gen- eralization long after perfect training performance is achieved, as first described by Power et al. [2022]. This phase transition, which often induced by implicit or explicit regularization, offers unique opportunity to isolate potential causes for the emergence of generalization. Recent work explains grokking through two-phase dynamics: an imitial kernel regime focused on memorization, followed by feature learning that supports generalization [Mohamadi et al., 2023, Lyu et al., 2024, Kumar et al., 2024]. Analytical studies reveal closed-form grokking solutions in modular arithmetic tasks [Gromov, 2023, Žunkovič and Ilievski, 2024], while mechanistic studies identify sparse and dense subnetworks competing during training [Nanda et al., 2023, Gouki et al., 2023]. Most of these studies focus on shallow networks or toy tasks. However, recent work shows that grokking also arises in deeper architectures and real-world datasets [Murty et al., 2023, Humayun et al., 2024], suggesting greater relevance. We build on this foundation by using grokking as an experimental lens: By tracking the timing of flatness, NC, and generalization, we able to distinguish which factors are causally connected and which are merely correlated.
3 Geometric Foundations of Generalization
In this work, we investigate the interplay between generalization and two notable phenomena in modern deep learning: neural collapse, and flatness of the loss surface. These concepts are frequently observed to co-occur late in training, but their individual roles and causal relations remain poorly understood. Our goal is disentangle their contributions to model performance by leveraging the grokking setting. in which memorization precedes generalization, allowing us to observe the emergence of each phenomenon in isolation.
Generalization A central objective in machine learning is ensure that trained model not only performs well on the training data, but also generalizes to unseen data drawn from the same distribution. This difference in performance is formalized via the generalization gap. Let ʃ : X y be a model from a model class F trained with a wiice-differentiablc loss function f: y x y + R. on finite training set S X x y. drawn i.i.d. from a data distribution D over the input space X' and output space V. The generalization gap is defined as Egen (f, S) = {f) Eemp (f. S), where EJ) == E( [e(s(x), v)] is the risk and {f, == TST Σ (Z.))ss {()(x), y) is the empirical risk. The model generalizes well when Egen S) is small, indicating close alignment between training and test performance.
Neural Collapse and the NCC Measure Neural collapse (NC) is a geometric phenomenon characterizing the late stages of training in deep classification networks [Papyan et al., 2020, Mixon et al., 2022]. It manifests as an alignment of the learned representations such that: (1) within each class, features in the penultimate layer collapse to their class mean; (2) the class means themselves form a simplex equiangular tight frame, a maximally spaced configuration; (3) the classifier weights align with these class means; and (4) the classifier effectively becomes nearest-neighhor model.
This structure has been theoretically motivated and empirically observed in both training and unseen test samples, including new classes, particularly in transfer learning settings [Galanti et al., 2022]. To quantify the emergence of NC, we adopt the simplified Neural Collapse Clustering (NCC) measure proposed by Galanti et al. [2022], which captures the essential characteristies of neural collapse, i.c.,

---

tight intra-class clustering and strong inter-class separation, without requiring all four formal NC conditions. This is the measure we use throughout our analysis.
Definition 3.1 (NCC Measure), Let ((x) denote the penultimate- representation of input æ; and De the set of training samples from class c. Then
NCC = Σ
Ve + Ve
2 2| ue
oFc'
where jr. == TD.T EXED, (xx), V= E 11((z) - Mell are the mean, respectively the variance of representations of class
low NCC value indicates tight intra-class clustering and well-seearated class means, characteristic of NC. In Section 5, we confirm that the NCC measure correlates with the angular separation of means, thus capturing all the characteristics of NC.
Relative Flatness Flatness of the loss landscape, understood as the insensitivity of the training loss to small perturbations in the parameter space, has long been linked to generalization research [Hochre- iter and Schmidhuber, 1997, Keskar et al., 2017]. However, classical measures based on the Hessian or curvature are sensitive to reparameterizations making them unreliable in modern architectures. To overcome this, Petzka al. [2021] introduced reparameterization. itttrrizaationnniariant invariant notion of relative flatness, grounded in a theory of robust generalization under locally constant labels.
Consider a model that decomposes as [ax w) g( (wo())). where Φ is fixed feature map (e.g., the penultimate layer), w E [Rd is a weight matrix, and is twice-differentiable function. The relative flatness is defined via the trace-weighted Hessian along the directions spanned by w.
Definition 3.2 (Relative Flatness). Let g, be twice differentiable, and S a sample set. Then:
d
Tr (w)
Σ
(w., W.) Tr(H,, (w,
where w, denotes the s-th row w, (-.-) is the scalar product, and Ha.s (w, (SS)) is the Hessian of the empirical loss with respect to w, and w. evaluated at (88).
This measure is invariant neuron-wise rescaling and composition with orthogonal transformations, making it robust across different parameterizations. Under mild assumptions-namele that the training data is representative and the labels locally constant in feature space-relative flatness predicts generalization. Note that a small value of xig (w) ) indicative flat solution and good generalization. In Section 5, we further show that NC leads to low relative flatness, establishing theoretical dependency between the two phenomena.
On the Role of Representativeness Both neural collapse and relative flatness have been proposed as geometric signatures of generalization. However, neither property is meaningful in isolation: their predictive value depends critically on the relationship between training and test distributions in feature space. In particular, if the training features are not representative of the full data distribution, then both NC and flatness emerge without leading to generalization. For example, consider network that perfectly collapses the training set into a class-wise simplex configuration and assigns test samples to incorrect cluster means. (On discrete data, such a network function can be easily constructed.) Prediction will then fail despite the presence of collapse. This limitation also holds for relative flatness: a flat but non-representative model does not generalize [ef. Sec. 2, Petzka et al., 2021]. While NC lacks rigorous definition of representativeness, relative flatness is grounded in formal framework that makes the required assumptions explicit. particular, the theory of -rreressetaatieensss [Petzka et al.. 2021] defines a condition under which robustness in feature space becomes necessary for generalization. Under this assumption, relative flatness serves as an upper bound on feature robustness and provides tractable diagnostic for generalization. This perspective clarifies that our analysis operates under the assumption of representative features-an assumption that is empirically supported in settings, demonstrate in Appendix D.
5

---

(a) Trairing/validatiion accuricies
(b) Relative flatness
(c) Neural collapse clustering (NCC)
Relative flatness
NCC
Relative fltnness (REG)
NCC (REG)
70
Training acc.
60
Val acc.
202
Training acc. (REG)
100
2-
Val acc. (REG)
50
150 Training epachs
:00
280
250
30
300
1.0
200
250
100
1.7
270
250
Training epochs
Training epochs
Figure Results of neural collapse clustering regularization on CIFAR-10. We display both unregularized and regularized (REG) training dynamics for comparison. Increasing NCC does not affect generalization or relative flatness, indicating that NC not necessary for generalization. Figure (a) shows the training and validation accuracies, and y-axis represents accuracy. Figure (b) presents the relative flatness values during training, and y-axis represents measurement of relative flatness. Figure (c) illustrates the NCC development, and y-axis represents NCC value.
Neural Collapse and Flatness in Delayed Generalization
Grokking refers to the surprising phenomenon in which a neural network, after memorizing the training set with zero generalization for an extended period, suddenly begins to generalize with high accuracy, often after tens or hundreds of thousands additional optimization steps [Power et al., 2022]. This temporal separation between memorization and generalization provides unique opportunity to disentangle the factors that underlie generalization. In our setting, we use grokking, or delayed generalization, not curiosity, but causal probe: a mechanism for identifying which geometric properties emerge coincidentally with generalization, and which ones may be functionally necessary for it.
First, we study grokking on symbolic algorithmic tasks such as modular arithmetic (e.g., x + y mod p), where networks trained to predict the output of binary operations over abstract tokens. These tasks are small, fully observable, and require exact generalization beyond memorized samples, making them ideal testbed. We follow the original setup of Power et al. [2022], training a layer-transformer using the Adam optimizer with a learning rate of 10 and weight decay of 1.0. Each experiment runs for 10 steps with 50/50 train/validation split. All results are averaged over three seeds. Figure shows the evolution train and validation accuracy, as well as relative flatness (Def. 3.2), and neural collapse clustering (Def. 3.1) throughout training. We observe that both measures decrease sharply when the network starts generalizing, however, NCC already decreases during the memorization phase. Relative flatness remains high during the memorization phase and decreases sharply near the point at which generalization begins. This temporal coincidence suggests that both NCC and relative flatness are correlated with generalization, but representations start exhibiting collapse behaviour before generalization has emerged.
However, temporal correlation is not causation. The simultancous decrease of NCC and flatness may reflect shared dependence on generalization, or confounders, rather than causal influence. The grokking setting thus makes it clear that must go further: it is not enough to ask whether these quantities track generalization, we must ask whether they are causing it. In the following sections, we address this question empirically. In Section 5, we show that neural networks can generalize without exhibiting collapse. In Section 6, we provide empirical evidence that relative flatness is necessary for generalization under mild assumptions.
5 Neural Collapse is not Necessary
Neural collapse (NC) is often observed in models that generalize well, leading to the impression that may play a functional role generalization. However, correlation does not imply necessity. this section, we demonstrate both empirically and theoretically that NC is not required for generalization. While collapse can accompany good generalization, is not necessary: It one path that leads to flatness, and through flatness, generalization.

---

To substantiate this claim, we explicitly suppress the emergence of NC during training without affecting generalization. We train a ResNet- 18 on CIFAR-10 using a regularized loss of the form:
REG LCE X NCC,
where Len the standard cross-entropy loss and NCC the NC clustering measure (Def. 3.1). The NCC term penalizes tight clustering of penultimate-layer features by shrinking the distances between class means and increasing intra-class variance, thereby actively discouraging the geometric structure characteristic of NC. The training hyperparameter setups, additional experimental details are provided Appendix B.
We show that generalization can occur without NC in Figure As shown in panel (a), both training and validation accuracies remain unaf- fected by the regularizer. Panel (b) shows that relative flatness, our proxy for generalization, also remains stable. However, panel (c) reveals that NC is effectively suppressed: after brief drop the NCC measure, steadily increases throughout training. This provides clear evi- dence that generalization can occur in the com- plete absence of NC, and that collapse is not necessary condition for generalization. In Fig- ure 3 we confirm that the NCC measure correctly captures cluster angles as well. Under neural col- lapse, angles are close to the optimal 10-simplex, yet when we regularize with the NCC measure, angles deviate strongly from this optimum.
Angle Difference
65
Angle Difference (REG)
Angle value (in deg ree)
53
50
10
35-
Training epochs
Figure 3: Pairwise cluster angles vs. optimal 10- simplex angles. Under NCC regularization, angles drift from the optimal configuration, while in stan- dard training they remain stable. "REG" indicates use of the NCC regularizer.
Why, then, does neural collapse often accom- pany generalization? We argue that collapse is one geometric pathway to flat solutions, which more directly linked to generalization. In fact, under mild ussumptions, NC the penul- timate layer leads relative flatness?. This rela- tionship is formalized below.
We assume that classical neural collapse [Papyan et al., 2020] in the limit holds for network function f. That is, the following four conditions apply to the neural network f(x) = softmax(w0fl + 6) with (xx) as the representation of the penultimate layer.
Assumptions 5.1 (Neural Collapse [Papyan et al., 2020]). Let ((x) denote the penultimate layer feature representation an input x E X, and for a dataset D X, let D. C denote the set of inputs belonging to class c {I...., k.). Then the four neural collapse criteria are
(NC1) Feature representations collapse to class means: ((x) = for all x E De-
(NC2) All class means He have equal distance to the global mean Pa. i.e. ||moo - 1all2 M, and they form a centered equiangular tight frame (ETF), so that (my - Mg) (ML m4) M² for A
(NC3) The classifier weights align with class means, i.c., wy >(ps (g) for all j and some A > 0.
(NC4) argmax, w. h + by argmin, ||h h Hyl--
Remark 5.2. For the proof of Proposition 5.3, we additionally assume that || S M for each c. This entails that the center of the ETF is not diverging. This mild technical condition does not alter the essence of the neural collapse assumptions.
RRelteed works often rely on simplifying assumptions, e.g., one-layer behavior or layer peeling [Mixon et al., 2022, Fang al., 2021]. yet analysis does not require such approximations.
7

---

(a) Trai ning/validation accuricies (CIFARIO)
(b) Traininngvalldatiion accuricies (IMAGENET100)
80-
60
40-
20-
50
100
150
220
250
200
353
430
50
100
150
203
250
300
Trakiing spocts
Training opochs
Training
Training lunplug)
Training acc.
Training acc. (unplug)
Val Rcc.
Val acc. bunplugg
Val acc.
Val acc. (unplug)
Figure 4: Results of inducing delayed generalization training Resnet18 CIFAR-10 and VIT on ImageNet- 100 through relative flatness regularization. We display hoth unregularized and regularized (unplug) training dynamics for comparison. Relative flatness regularizer is removed epoch 200 for CIFAR-10 experiment and 150 for ImageNet-100 100. Delayed generalization occurs after the regularizer removed, indicated by sharp increase in validation accuracy. Figure (a) shows the training and validation accuracies on CIFAR-10, and Figure (b) on ImageNet-100.
Proposition 5.3. Let f(x) caftmax(wó(1)) + b) be neural network with softmax output and trained with cross-entropy loss, where Rexd denotes the final-layer weight matrix classifying into classes and 6(x) the penultiaat--layer representation. Assume that the classical neural collapse holds in the limit as above. Then relative flatness Ks kw(w) is bounded by:
-MM F
ko(w) M²
1 (k: 1)e AM3. x-1
In particular, for sufficiently large X, this yields the asymptotic bound:
ko(w) X -MM AM -
which decays exponentially in A.
The proof is provided in Appendix A. The fact that relative flatness indeed decreases to zero in the limit follows from the fact that the conditions of neural collapse are independent of the weight norm and thereby also the magnitude of the scalar X. As training continues in the NC limit, the training loss can be decreased zero by increasing X, because the softmax probabilities converge to (e(xx) and ÚJ (z2) = 0.j * as + c0.
Taken together, our results demonstrate that, while NC can facilitate generalization by inducing flatness, it is not a prerequisite. Flatness remains the more fundamental quantity, and its relevance critically depends the representativeness of the learned feature space. It remains investigate whether flatness is sufficient condition and whether it could be necessary.
Relative Flatness is Necessary
In this section, induce delayed generalization, similar to what observed in grokking. These experiments demonstrate that relative flatness is necessary for generalization. To achieve such behaviour, design loss function that includes the following regularization term:
CRF Lee A neg (w),
where x (w) specified in Definition 3.2, and LCE denotes the cross-entropy loss. To efficiently compute the relative flatness the penultimate layer, adopt the closed-form derivation under

---

cross- eass--n loss provided by Walter et al. [2024]. By maximizing the relative flatness, we encourage the network to explore sharper regions of the loss landscape rather than flatter ones, which leads to suboptimal generalization. The coefficient À controls the strength of the regularization. We train ResNet-10 10 on CIFAR- 10 and ViT [Dosovitskiy et al., 2021] on mmageNet-100", both from scratch, to evaluate the effectiveness of inducing delayed generalization in general settings by regularizing relative flatness. The results are shown in Figure 4. We also fine-tune TinyBERT [Jiao et al., 2020] and DistilGPT2 [Sanh et al., 2019] on SST-5 to show that delayed generalization can be induced by regularizing relative flatness in pretrained language models. It is worth noting that stabilizing training in sharp minima under the relative- fiattess regularizer is nontrivial, since cross-entropy loss reduces prediction uncertainty (in terms of the entropy of softmax), whereas the regularizer operates effectively only under high uncertainty. Additional analysis of the behavior and sensitivity of the regularizer, particularly its conflict with cross-entropy loss, is provided in Appendix C, which also includes the training setups and additional experimental details.
From Figure 4 (a), we observe that the validation accuracy is low, but increases after the relative flatness regularizer is removed at epoch 200. A similar phenomenon is observed in the results on ImageNet-100 shown in Figure 4 (b) and in the NLP experiments presented in Appendix C.4. This recovery is facilitated by applying standard regularization techniques after unplugging, helping the model escape sharp minima and restore validation accuracy to the baseline. The successful induction of delayed generalization in these tasks demonstrates that this effect is not exclusive to algorithmic tasks trained with the next-token prediction objective. By surpressing relative flatness, we can induce delayed generalization in more general tasks and training settings.
In particular, when the relative flatness regularizer applied throughout the entire training process, the final validation accuracy remains around 60% on CIFAR-10 10, whereas the optimal validation accuracy reaches approximately 78%. For ImageNet-100 the stable validation accuracy under regularization is around 38%, compared to optimal accuracy of around 43%. A detailed analysis of pretrained language models on SST-5 is provided in Appendix C.4. This gap in validation accuracy, caused by the absense of relative flatness, suggests that optimal generalization cannot be achieved in such way.
To further support the need for solutions to be flat for better generalization, we compute the relative flatness values at the epoch right before the regularizer is removed and at the final training epoch. For CIFAR- 10, at the epoch 199, the measure of relative flatness under regularization is 2209.94, which decreases to 1.91 at the final epoch. For ImageNet-100 the measure is 22212.75 at epoch 149 and 313.22 at the final epoch. We also observe similar trends on SST-5 with TinyBERT and DistilGPT2 (Appendix C.4). Therefore, smaller generalization gap is directly correlated with the value of relative flatness.
Overall, these results provide strong empirical evidence that relative flatness is necessary for general- ization, consistently observed across all evaluated tasks.
7 Discussion and Conclusion
Our results identify relative flatness as a potentially necessary condition for generalization in deep networks, while neural collapse emerges a correlated but non-essential feature of late-stage training. This distinction reshapes our understanding of the geometry of well generalizing solutions and invites several important reflections. Our experiments show that neural collapse tends to coincide with high training accuracy but does not reliably track generalization. This suggests that it is more a byproduct of training dynamics than causal mechanism, a view that aligns with findings from Hui et al. [2022] and Hong and Ling [2024]. Nevertheless, we show theoretically that under mild assumptions, neural collapse-style clustering leads to relative flatness. As such, collapse may still serve practical proxy for flat solutions in some settings.
While neural collapse produces highly symmetric and compact feature representations, this structural simplicity can be problematic in transfer learning and out-of-distribution scenarios. For instance, if a network is trained to classify supercategories by collapsing fine-grat fine-grained class distinctions, it becomes impossible to recover those distinctions later, thus limiting the generality and interpretability of the learned features [Hui et al., 2022]. The requirement for geometric compression of the representation

---

space, with neural collapse as one instance. has also been questioned from an citoommationtthhorettee perspective on deep learning dynamics [Adilova al., 2023b].
Flatness itself has also come under serutiny, particularly in the context of sharpness-aware minimiza- tion (SAM) [Foret et al., 2021] and recent critique by Andriushchenko al. [2023]. However, many such criticisms target classical flatness metrics, which are sensitive to reparameterization [Dinh et al., 2017] or measures with unclear connection to geometric flatness [Andriushchenko and Flammarion, 2022]. Our analysis relies on relative flatness, reparameterizatioonivvaaai measure with theoretical guarantees [Petzka et al., 2021, Walter et al., 2024]. While it remains an open question how these recent critiques translate to relative flatness, our findings suggest that it continues to offer a reliable signal for generalization in the settings we consider. One of the notable properties of relative flat- ness is that it is theoretically defined for single, preselected layer. which enables computability independent of the model size. This formulation also allows for a non-idealized comparison with neural collapse, which is similarly defined in terms of single-layer properties, thereby enhancing the precision of our analysis. At the same time, raises questions about how to characterize embeddings in other layers, since prior work [Zhang et al., 2022, Adilova et al., 2024, Ansuini et al., 2019] has shown that different representations and layers exhibit distinct properties, even though relative flatness appears to behave similarly across layers [Walter et al., 2025].
The theoretical link between relative flatness and generalization depends on two key assumptions: that labels are locally constant and that the features at the selected embedding layer are representative of the data distribution. In our settings (standard image classification, modular arithmetic tasks and NLP tasks) both assumptions are well justified: (i) the assumption of locally constant labels underlies the very notion of adversarial robustness, and its empirical validity is supported by our experimental findings, and (ii) we show in Appendix D that representativeness improves with the onset of generalization. Nonetheless, these assumptions may not hold in other domains such as structured prediction, high-noise regression, or complex real-world tasks, where label semantics may vary significantly under small input changes.
One of our most surprising results is that regularizing networks away from flat minima consistently induces delayed generalization, mimicking grokking even on real-world datasets and architectures. This underscores the functional importance of relative flatness, not just its correlation with general- ization. It also opens new directions for controlling training dynamies via geometric regularization. However, the extent to which these artificially induced grokking phases mirror the mechanisms of standard grokking remains to be understood.
Finally, while penalizing relative flatness reliably delays generalization, encouraging flatter minima does not substantially improve it [Adilova al., 2023a]. Even SAM, which improves training results in vision tasks, surprisingly does not actually encourage only flatter minima [Andriushchenko and Flammarion, 2022]. This asymmetry supports the long-held view that stochastic gradient descent (SGD) already biases toward flat minima [Hochreiter and Schmidhuber, 1997. Keskar et al., 2017], and that relative flatness may serve more effectively as a diagnostic and interpretive tool rather than a universal optimization target.
While our findings robust across datasets and architectures, our conclusions are necessarily limited in scope. We focus exclusively on classification networks and generative language models trained with standard objectives and optimizers; whether our results extend to contrastive pretraining, or large-scale language models remains an open question.
We challenge the assumption that neural collapse is essential for generalization. Through grokking and targeted regularization, we isolate flatness as the necessary factor. Our theoretical and empirical results jointly establish flatness as the geometric driver of generalization.
Acknowledgements
This research is supported by the Federal Ministry of Education and Research of Germany and the state of North Rhine- Westphalll as part of the Lamarr Institute for Machine Learning and Artificial Intelligence, and received support from the Cancer Research Center Cologne Essen (CCCE).
10

---

References
Linara Adilova, Amr Abourayya, Jianning Li, Amin Dada, Henning Petzka, Jan Egger, Jens Kleesiek, and Michael Kamp. Fam: Relative flatness aware minimization. In Topological, Algebraic and Geometric Learning Workshops, 2023a.
Linara Adilova, Bernhard C. Geiger, and Asja Fischer. Information plane analysis for dropout neural networks. In International Conference on Learning Representations, 2023b.
Linara Adilova, Maksym Andriushchenko, Michael Kamp, Asja Fischer, and Martin Jaggi. Layer- wise linear mode connectivity. In International Conference Learning Representations, 2024.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International conference machine learning, 2022.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. In International Conference Machine Learning, 2023.
Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. In Advances in Neural Information Processing Systems, 2019.
Hien Dang, Tho Tran Huu, Stanley Osher, Nhat Ho, Tan Minh Nguyen, et al. Neural collapse in deep linear networks: From balanced imbalanced data. In International Conference on Machine Learning, 2023.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer-peeled model: Minority collapse imbalanced training. Proceedings of the National Academy of Sciences, 118(43), 2021.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awar minimization for efficiently improving generalization. International Conference on Learning Representations, 2021.
Tomer Galanti, András György, and Mareus Hutter. On the role of neural collapse in transfer learning. In International Conference on Learning Representations, 2022.
Minegishi Gouki, Yusuke Iwasawa, and Yutaka Matsuo. Grokking tickets: Lottery tickets accelerate grokking. CoRR, 2023.
Florian Graf, Christoph Hofer, Mare Niethammer, and Roland Kwitt. Dissecting supervised con- trastive learning. In International Conference Machine Learning, 2021.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679, 2023.
XY Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity to and dynamics on the central path. In International Conference on Learning Representations, 2022.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1), 1997.
Wanli Hong and Shuyang Ling. Neural collapse for unconstrained feature model under cross-entropy loss with imbalanced data. Journal of Machine Learning Research, 25(192), 2024.
Like Hui, Mikhail Belkin, and Preetum Nakkiran. Limitations of neural collapse for understanding generalization in deep learning. arXiv preprint arXiv:2202.08388, 2022.

---

Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk. Deep networks always grok and here is why. In High-dimensional Learning Dynamics Workshop: The Emergence of Structure and Reasoning, 2024.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. Findings of the Association for Computational Linguistics: Empirical Methods in Natural Language Processing, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representationn, 2017.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transi- tion from lazy to rich training dynamics. In International Conference on Learning Representations, 2024.
Kaifeng Lyu. Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D. Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. In International Conference Learning Representations, 2024.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis, 20: 11, 2022.
Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica Sutherland. Grokking modular arithmetic can be explained by margin maximization. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Michael Munn, Benoit Dherin, and Javier Gonzalvo. The impact of geometric complexity on neural collapse in transfer learning. In Advances in Neural Information Processing Systems, 2024.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. Grokking of hierar- chical structure in vanilla transformers. In Annual Meeting of the Association for Computational Linguistics, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Representations, 2023.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40), 2020.
Hemning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In Advances in neural information processing systems, 2021.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint 07227.2001.02177. 2022.
Akshay Rangamani and Andrzej Sanburski-Fahey, Neural collapse in deep homogeneous classifiers and the role of weight decay. In International Conference on Acoustics, Speech and Signal Processing, 2022.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 1910.00 108, 2019.
Peter Súkeník, Marco Mondelli, and Christoph H Lampert. Deep neural collapse is provably optimal for the deep unconstrained features model. In Advances in Neural Information Processing Systems, 2023.
12

---

Nils Philipp Walter, Linara Adilova, Jilles Vreeken, and Michael Kamp. The uncanny valley: Exploring adversarial robustmess from flatness perspective. arXiv preprint arXiv:2405.16918. 2024.
Nils Philipp Walter, Linara Adilova, Jilles Vreeken, and Michael Kamp. When flatness does (not) guarantee adversarial robustness. arXiv preprint arXiv:2510.14231, 2025.
Diyuan Wu and Marco Mondelli. Neural collapse beyond the unconstrainted features model: Land-
scape, dynamics, and generalization in the mean-field regime. In International Conference Machine Learning, 2025.
Robert Wu and Vardan Papyan. Linguistic collapse:
Neural collapse in (large) language models, In Advances in Neural Information Processing Systems, 2024.
Yingwen Wu, Ruiji Yu, Xinwen Cheng, Zhengbao He, and Xiaolin Huang. Pursuing feature separation based on neural collapse for out-of-distribution detection. In International Conference Learning Representations, 2025.
Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do really need learnable classifier at the end of deep neural network? In Advances in neural information processing systems, 2022.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference Learning Representations, 2017.
Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are all layers created equal? Journal of Machine Learning Research, 23(67), 2022.
Jinxin Zhou, Chong You, Xiao Li, Kangning Liu, Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: A neural collapse perspective. In Advances in Neural Information Processing Systems, 2022.
Jinxin Zhou, Jiachen Jiang, and Zhihui Zhu. Are all layers created equal: A neural collapse perspective. In The Second Conference on Parsimony and Learning, 2025.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geo- metric analysis of neural collapse with unconstrained features. In Advances in Neural Information Processing Systems, 2021.
Bojan Žunkovië and Enej Ilievski. Grokking phase transitions in learning local rules with gradient descent. Journal of Machine Learning Research, 25(199), 2024.
13

---

A Proof of Proposition 5.3
First, we provide a technical lemma that establishes that, under Neural Collapse, the combined contribution of the bias and the projection of the class mean onto the global mean is identical across classes. This means that the logits are effectively governed solely by the geometric structure of the class means relative to each other.
Lemma A.I. Suppose that the neural collapse conditions (NC1) - (NC4) in the limit apply as in Section (see also below). Then there a constant such that A(j1 - Mg) Mg + by = d for all j.
Proof. For all h, the following set of equations holds:
argmax, M(as - Mg) T + by
(NC3)
argmax; wi h + bj
(NC4) argmin, | uوl 1l
argmin, |1(h |gg] (14z 111)122
argmin, || - !اall + I1ls - Hel 2- 2(h H4) T (pg pbg)
(NC2) argmin, ||/ ugll² + M² - -2(h (4)) [14 (ag
argmax, 2(1r - (g) (h T
(g)) - -I 1g 12
Let 6 > be a real number and E {1,.... k). Since the above equations hold for all h, they hold in particular for all hhc, c) as (Me Mgg Mg. Note that [[A(c, c) 1tyl12 =?.
Substituting h(c, <) for h in the above equation and letting d, = A(11 - -My) My + by gives us that
de (1g (g)T(He -Mg) + dj argmax, j M (11 - Hog) (He -My) - e argmax, (M, - ug) (we - (ag)
2c
argmax, M
(NC2)
This implies that
AE (Me pg)TT(p - 1g) + do 2 de
(U; - (g) T - -mg) + dj, i
M
M
Using (NC2), this gives that
XcMk
b: k- 1 + d. 2 dj. 1
But since 0 was arbitrary, we can let - 0 to get that de 2 d3 for all j * Finally, also was arbitrarily chosen, which implies that d, de for all j, Hence, dj = A(Hj (gg) Mg + by is constant o.er
We now provide second technical lemmma showing that relative flatness is upper bounded by the simplified version (TTTH(((H)))) for which we have a closed form expression in terms of logits.
Lemma A.2. Let H (w) be the Hessian wrt. the penultimate layer weights w E IRd at a minimum and (w) the relative flatness measure. Then
ko(w)= < ||m|| (T(H())()(
Proof. Recall that
(s) Σ (wa, wo) Tr(HHe.)
Σ kw Wr))TT(HHedd
Σ Iwa ||l [Tr(He,,:))
14

---

Since w is at minimum, H is positive semi-defing (psd) and therefore (Tr(Hs,,r)) Do,sr is also psd. Then holds that
[T(HH, ( (H(H,)
Summing over all s' and using the Cauchy-Schwartz inequality we get
2 |wsll Tr(H.ss)
Σ |eellllww (THH.. s
and using symmetrie yields
(w) $ ||w|| < = -|w||
With this we can proof the theorem, which we restate together with the neural collapse criteria for convenience.
Assumptions A.3 (Neural Collapse [Papyan et al., 2020]). Let Φ(x) denote the penultimate layer feature representation of an input x E X, and for dataset D X, let D. S denote the set of inputs belonging to class e .l.....,)... Then the four neural collapse criteria are
(NCI) Feature representations collapse to class means: <(x) = Mo for all x Do.
(NC2) All class means He have equal distance to the global mean Pg. i.c., |1 - 11yl-1 = M, and M
they form centered equiangular tight frame (ETF), so that (Hy 14y-14 ( 14g) k-1 for j * We additionally assume that || S M for each c. That is, assume that the center of the ETF not diverging using the same constant M for simplicity.
(NC3) The classifier weights align with class means, i.e. Wj = M(rz 14g) for all j and some O.
(NC4) argmax, w/ + by argmin, |// Myll-
Proposition A.4. Let f(x) = softmax(wo(2 (wonaxx(w(()))))))))))))))) + b) be a neural network with softmax output and trained with cross-entropy loss, where w eed denotes the final-layer weight matrix classifying into l classes and 6(x) is the penultiaat-layee representation. Assume that the classical neural collapse holds in the limit as above. Then relative flatness Kó (w) is bounded by:
E AM52 x
ko(w) S M
1 + (k. 1)e AM
In particular, for sufficiently large X, this yields the asymptotic bound:
ka(w) S A -MM
which decays exponentially in X.
Proof. Lemma A.1 shows that the NC conditions imply that there is constant d such that by M(15 (gg) Ng d for all class labels 1,..., h. For De De, the logits of f are
w. ((xx) + by =>(1ı -4y) TH)THe. He + bj
=b(μ -m)) THe Hyy + A(Hg my) My +, Ma2s 14) TM Mg) d
and the softmax probabilities become:
A6
(x.) = 1 (k - 1)e- (g(x,)
1 (A: 1)e-16 22 j
15

---

with margin
8: M2 (
k - (-M²2
M2
M2. k
k --1 1
The trace of the Hessian for input Tie E Dr is (cf. Walter et al. (2024))
k Tr(H(Fe; w)) = Σ (1 - -ū,(x,)) [|((.))))== M²,
&
Σ (s)()))(1 - (100((()))))
1=1
5
ar2 (k: 1)e - (k- 2))=A
.(2
Ad)
(1+ (k: (1-1)e 1)e-2692 2012
Thus,
(I: 1)ke 26
(r)H(w)) < M
(1 + (k- 1)e-A
Since ||m|12 ====>>> E, [1ll; - Mgl X²MM2 we obtain with Lemma A.2:
(w) S X2KM2 M (k: 1)ke
AMd2. --
(1 + (k: 1)e- S XP
+ (k: - 1)e AM52, x=1 2
This yields the desired bound, and for large X, the denominator tends to 1, giving the asymptotic behavior.
Relationship between Neural Collapse and Relative Flatness. Neural Collapse (NC) character- izes a geometric property of the feature representations in the penultimate layer, which we quantify by the Neural Collapse Coefficient (NCC). Importantly, NCC is invariant to rescaling of the last-layer weights, and therefore independent of the scalar À that controls their magnitude. Consequently, NCC itself does not directly imply flatness: the loss landscape can remain sharp even in the presence of perfect collapse. However, in the NC regime, the network structure enables flatness emerge naturally. As training with the cross-entropy loss continues, the model can reduce the training loss further by increasing A, since the softmax probabilities satisfy (x..)) > 1 and (2,,) - 0, * c as + oc. This rescaling leaves the feature geometry (and hence NCC) unchanged but decreases the relative flatness measure, effectively leading to flatness without altering the collapsed representation. In summary, NC itself does not guarantee flatness, but in combination with CE. loss training leads to flatness: it detines feature geometry that allows the network to reach a flat minimum through continued optimization. Observing NC is thus a strong empirical indicator of eventual generalization, but not necessary condition for it. Conversely, models can achieve generalization and flatness through other representational structures that do not exhibit collapse.
B Training Setups for NCC and Additional Experiments
We evaluate our method on the CIFAR-10 dataset using the standard training (50,000 samples) and test (10,000 samples) splits. Input images are normalized using the channel-wisg mean and standard devi- ation of (0.5, 0.5, 0.5) and (0.5, 0.5, 0.5). respectively, implemented in transf consforms.. NNrmal ize. No data augmentation is applied in the main experiments.
All models are trained using stochastic gradient descent (SGD) with fixed learning rate of 0.01, batch size of 64, and no weight decay. The regularization coefficient set to X = 10-3, Momentum is 0.9. Training is performed for 250 epochs without the use of learning rate scheduler or early stopping.
We use the ResNet- 18 architecture as implemented in torchvision version 0.19.1, without any modifications. particular, we retain the original 7 x 7 kernel size in the first convolutional layer and the following max pooling layer, despite the smaller resolution of CIFAR-10 10 images.
All results in the main experiments are reported the average random seeds 1, 42, and 15213 to account for variability due to random initialization. Experiments conducted using PyTorch 2.4.1 on a single NVIDIA A00 GPU with 80GB of memory. Training each model for 250 epochs takes approximately 10.5 hours.
three independent runs with
The loss corresponding Figure is presented in Figure 5.
16

---

(a) Trainiog/validationn losses
(b) Training/validatioon losses (REG)
5-
Training loss (REG)
2.0
4-
Val loss (REG)
1.5
3-
Loss
Training loss Val loss
Loss
2
1
-1
50
100
150
200
250
50
100
150
200
250
Training epochs
Training epochs
Figure Results of training and validation losses in the NCC experiments.
Effect Different Values of NCC Regularizer our main experiments, we apply À = 1x 10-3 using three random seeds. To study the effect of different À values, we conduct hyperparameter search experiment using a range of = {1 x 10-2, x 10 x 10-1, 10-5, 10-"). evaluated with Seed 42 due to computational constraints. We have verified that results with single seed remain consistent with the trends observed using multiple seeds.
The results are shown in Figure 6, which plots accuracy, loss, relative flatness, and neural collapse clustering (NCC) metrics. We observe that large value (e.g.. ... 10-2) causes training instability: although training accuracy reaches nearly 100% at the end of training, validation accuracy remains below 60%, and both training and validation accuracies collapse to around 10% around epoch 70.
Conversely, when À is smaller than x 10-3, the NCC regularizer has little to no effect on the training, evidenced by negligible differences in performance metries and NCC values compared to the baseline. When À is set to 1 x 10-3, the measure of relative flatness remains low while the NCC metric increases. This supports our argument that neural collapse clustering is not necessary for generalization, as good generalization performance can still be achieved in the absence of strong NCC behavior.
Effect of Data Augmentation in NCC ularization, we apply standard transformations- Saaaiassaaaaaaaaaaaaaaaaa (3,, padding=4) and RandomHorizon talFlip()-ann and conduct experiments using À E (10-2 10 3 10 10-4] All inputs are normalized using per-channel means and standard deviations of (0.4914, 0.4822, 0.4465) and (0.2023, 0.1994, 0.2010), respectively.
To investigate the effect of data augmentation on NCC reg-
We limit the range of to these three values because prior results indicate that À 10-4 has negligible effect; smaller values would therefore be redundant. Due to limited resources, we use seed 42 for this experiment, and based on consistent trends observed across seeds in other experiments, find this sufficient for reliable comparison.
The results, shown in Figure 7, indicate that with À = 10- training becomes unstable and perfor- mance significantly deteriorates. With = 10-3 and 10 data augmentation prevents the NCC metric from increasing it does in the non-augmented setting-Suggesting that NCC regularization less effective when strong implicit regularization already applied. In particular, À = 10-3 which was effective without augmentation (low measure of relative flatness and increased NCC), does not produce the same NCC increase under augmentation. At = 10-4 the NCC metric remains low in both augmented and non-augmented cases.
In these experiments, we also adopt a ResNet-18 variant tailored for CIFAR-10. following prior work: the first convolutional layer is modified to 3 x 3 kernel, and the initial max pooling layer is removed to better suit the smaller input resolution. The model with augmentation achieves higher validation accuracy (94%) compared to the auggeneddd case (78%).
17

---

trarengival 20-2
Reuo colagse Chaderin 2-20-2
HSL Je-2 Pan, amay
J02. 110 ECC 171
HRC win.
Basil Selative furrens Am
Tannng Eualo
buring Epous
Nouri collagas clattrng In--3
Tranvis 22-3
nee Leclenn art smy
NCC le Taai aer pirty
not val ama
Epocte
IRC Io- PaRLLac
mol, toid.oo
Moze oa
220
120
Jn rall amy
Tour cct 0014 cussinty
ner In-S main arry
cummerre
Ierd brant
Figure 6: Results different À values for the NCC regularizer. Each row contains training/validation accuracies, measure of relative flatness, NCC measurement and training/validation losses for particular value.
C Training Setups for Regularizing Relative Flatness and Experimental Results
C.1 Aanalysis of Relative Flatness Regularizer and Gradients
As stated in Section 6, the loss employed to regularize relative flatness is as follows:
LREREG = LCE > -if
Let single training example be (0) (0) with label (0) E {I,..., CJ. The penultimate represen- tation is (00) Rd, and the linear classifier has parameters W E Rxxd b R The logits and softmax probabilities are
exp(2(0 (8) E exp(2.0))
2+) (Wo + +b,
p!0)
1 I....,C.
C
18

---

A=le-1
ECC (e-2 2o... anly
Le-2 In
in pnay Jec2 oe Janueney
Travino Sourhe
Tolnin Southe
Tanre Spichs
Tral Eavie
Feledve Am3or
Noarel celoeee to
Tunnyekter
RCC MC
tands
Taning Spocto
Traninghiddddee sccarecs 2-
Nouzal hppa chatecing
Tranivalder lom
Eooda
Goads
Spacha
Thanneg Coods
Figure 7: Results of different À values for the NCC regularizer when data augmentation is applied. Each row contains training/valln caaning/validation accuracies, measure of relative flatness, NCC measurement and training/validationn losses for particular À value.
The per-sample cross-entropy loss is
log (1) y yo
CF
Based on the derivation in Walter et al. [2024], the per-sample relative flatness with cross-entropy loss is
C
(9) (w) = |
Σ
(6)
(1-p!")
10
and the total per-sample loss
c
clO
RF. = - log 100 1
Σ
p 1 - -(-
11619
(1)
C
- =-log Boggy1, A 1|
(2)
Σ p(o) - - (p.0) 10000
(2)
c=1
log (6)
A || (1 (0) (13)) 13
(3)
For simplicity, let op 1 1- 1 12 where 0 U( s - J/ Then R
|| U (p()) ||. The par-sample loss can be re- re-written as
6 (
log (5) you AR((.
RE
To encourage sharper minima, the optimizer tends to increase |[W| | and || under high uncertainty (p)()), On the contrary, the cross-entropy loss aims to minimize prediction uncertainty, which conflicts with the objective of the regularizer. When À is set to small value, the cross-entropy loss
19

---

quickly dominates the training, driving uncertainty down and effectively nullifying the influence of the regularizer. However, when À is large, this effect can dominate, causing parameter growth that eventually destabilizes training. To better understand this, we first look at the gradients with respect to the regularizer.
Gradients.
Write S= 1 and U || Then
OR
OU Oz
$6)TT
aW = SW + | 1 S
OR
0
|
SWT
20 2(0)
OΦ
0zz
Thus the overall gradients are
Oc ac
a
ac
OR dΦ
aW
aw
OW
0Φ
dΦ
Without constraints, the multiplicative term 2U SW amplifies ||W|| whenever predictions are uncer- tain (U large) and features have large norm (S large), which may cause runaway weight growth and unstable training.
(i)
OU 0z
W4+1 Wt -1
+ nd 2U + W I.S.
di
OW
(i)
OU T 0z
Φ1+1 6:
S. W.
+ 2U 0z
OΦ
aU Oz
bi+1 = b. + W
Multiplicative growth on the classifier. multiplies W. directly and can drive growth of [|W,] when U, and ( are not small. This motivates stabilization strategies.
The W Wuupdate includes the term 21/AU, / 16,, which
Stabilization by feature normalization and weight capping. If ||o. l| = is enforced (e.g.. replacing ø + 6 in the forward pass), then S = | [6 =1 and the amplifying factors involving |eell vanish. The multiplicative W-term reduces 2n, (still present but weaker), and the é-update no longer scales with | In particular, normalizing & fixes S = || = 1, eliminating the quadratic amplification in & that otherwise accelcrates the multiplicative growth of [[W]]
To further prevent runaway growth ||Well. we apply a Frobenius norm cap after each optimizer step:
W411 + Ig. (W++1),
Br = { W ||W]e $ Br
This projection makes the update non-expansive in the Frobenius norm and uniformly bounds all mopllaizerrinducce gradient terms, since they scale with |[W, |I and ||W, Together, these techniques ensure that training under the relative flatness regularizer remains stable and well-conditioned.
C.2 CIFAR-10 & Resnet-18 Training Setups and Additional Experiments
We follow the same training setup described in Section B, with two modifications: (1) the total number of training epochs is increased 400; and (2) the coefficient of the relative-flatre set to 0.01. During regularized training, the weight capping value is fixed at 50 and the penultimate representation is normalized ([|ø] 1). We introduce temperature parameter 7 = in the softmax when computing uncertainty (1). so as inflate the uncertainty signal and maintain the influence
regularizer is
of the regularizer during training. When the regularizer is removed Epoch 200, weight capping
20

---

remains active, but normalization of the penultimate layer is deactivated. The same procedure for weight capping and representation normalization in unplug experiments applied the other tasks as well. We also verify that setting the weight capping to 50 and normalizing the penultimate representation does not affect training without the regularizer.
The training and validation loss curves associated with Figure 4 are shown in Figure 8, and the corresponding claatev-flatteess measurements are reported in Figure 9.
(a) Training/validation losses
(b) Training/validation losses (unplug)
2.5
700
2.0
1.5
Training loss (unplug)
loss (unplug)
Training loss
0.5
Val loss
100
0
0.0
50
150
350
Training epechs
Training epachs
Figure 8: Results of training loss and validation loss of CIFAR-I0 in the delayed generalization.
(a) Relative Flatness
(b) Relative Flatness (unplug)
700
12000
600
10000
8000
400
6000
Relative Flatness (unplug)
300
Relative Flatness
100
2000
0
50
100
150
200
250
300
350
400
50
100
150
200
250
300
350
400
Training epochs
Training epochs
Figure 9: Results of relative flatness of CIFAR-10 in the delayed generalization.
The main experiments are initialized with seed 42, while additional results with seeds 1 and 15213 are shown in Figure 10. These results demonstrate that initialization with different random seeds can drive training toward different sharp minima. Consequently, averaging across seeds does not provide meaningful insights.
In Figure 10, standard training with ResNet-18 converges reliably: training accuracy reaches nearly 100%, while validation accuracy stabilizes around 78%. The relative flatness measure decreases steadily, indicating convergence toward flatter minima. With the regularizer, training accuracy still approaches 100%. but validation accuracy drops to about 60%, well below the baseline. At the same time, the relative flatness remains high (around 2000) and does not decrease, showing that the model pushed into sharper solutions. These behaviors are consistent across seeds. The regularizer is removed at Epoch 200. After unplugging, the total training loss becomes positive small while the model continues to train stably. To encourage escape from the sharp minima, we reset the learning rate to 1 x 10-1 apply cosine annealing schedule, and introduce weight decay of 2 x 10 10-? As shown in Figure 4, validation accuracy rises toward 78%, aligning with the standard training baseline. The relative flatness (Figure 9) also decreases and becomes indistinguishable from the
21

---

baseline, confirming that unplugging enables the model to recover from the sharp minima induced by the regularizer.
Trainngival dasten accurec
ative fattrsss, seed
Thanigaalssionn losuea, seed
www.dnggalccaaoooooooo lcsuss (REG), seed
tran lose
train un (REC) val loss (RE
vall dos
reletive Sathess
tolative fothess
10
Spocha
Taihing Spacts
Thanning Esocha
Faining Spocha
Taitingto suzion acoredis., sood
Nclethe: Tubness scod
Toiniingoo Maatin lassos soud
Tr lases (RE0) sood
taas lose (RE0)
loes (BE.)
(K66))) (RE))
100
16
70
351
10 31
30a
Tratingivat dution accuraciss seed 25113
Teretive fatness, seed 15213
Triiring
Coattin Edution losses, seed 15233
Triningiaaasttio Iosses (REO) seed 15213
312
trais
loss
Train (KGG)
(REO)
100
15
15
25
210
Figure 10: Resnet18 training results with different random seeds (42, 1, 15213), comparing standard fine-tuning and regularized training (REG).
Effect of À Values of Relative Flatness Regularizer. To investigate the effect of regularization strength, we conduct experiments with coefficients À {0.5, 0.1, 10-3 10=,,, while keeping all other settings (model architecture, training procedure, and random seed 42) consistent with the main experiments, except the number of training epochs is 300. Results are shown in Figure 11.
For large regularization strengths (A = 0.5 and 0.1), training crashes, causing both training and validation accuracies to remain low. For moderate value (A 10-3), training accuracy reaches nearly 100%, and validation accuracy remains slightly below the 78% baseline. Although the relative flatness stays high, the near-optimal validation accuracy indicates that intermediate À can drive the network toward sharp minima without significantly degrading performance. For a smaller value (A (A=10-1) training remains stable across 300 epochs. The relative flatness decreases rapidly to low value, the training loss converges smoothly to zero, and the validation loss stabilizes. However, when À is too small (below 10 10-4), the influence of the regularizer essentially disappears, and the model behaves like the baseline model without any affecting from relative-flatness regularization.
Effect of Weight Cap Values in the Relative-Flatnes, Regularizer. We examine the effect of different weight cap values (WC E (20, 100, 150, 200}) while keeping all other settings identical to the main experiments. Results are shown in Figure 12.
With a small cap (WC- 20), training remains stable: the relative flatness is still high but stays below 2000, the level observed in the main experiment, and validation accuracy reaches about 75%, only 3% lower than the baseline. At WC 100, training appears superficially stable within 300 epochs, as validation accuracy plateaus around 50% and training accuracy shows potential to approach 100% with extended training. However, the relative flatness increases rapidly to nearly four times the level reported in Figure 10, raising concerns about whether training can safely and reliably reach full
22

---

CIFAR-10 Experiments, different à values
Trainingivalidation accuracies (REG), A 0.5
Relative flatness (REG), 0.5
Training/validation losses (REG), =.5
27.
2200
-203
22.5
2000
400
train loss (REG)
400
17.5
-107
train (REG)
val acc (REG)
relative flatness (REG)
-1207
100
15o
250
300
56
100
200
250
150
200
Training Epochs
Training Epochs
Training Epochs
Training/validation accuracies (REG),
Relative
Training/validation (REG),
2000
train loss (REG)
1800
val loss (REG)
-150
20
(REG)
(REG)
(REG)
-207
10
1600
150
50
zi0
350
200
Training Epochs
Training Epochs
Training Epochs
Training/validation accuracies (REG), A 10-3
Relative flatness (REG), 2 10-3
Training/validation losses (REG), A 10-3
1.5
20
03-
train (REG)
1000
relative flatness (REG)
(REG)
750
-05
train acc (REG)
500
val acc (REG)
250
120
150
250
100
50
200
250
750
200
Training Epochs
Training Epochs
Training Epochs
Training/va lidation accuracies (REG). 10
Relative flatness (REG).
Trainingivalidation losses (REG),
200
(REG)
2.
train loss (REG)
70
1.0-
val loss (REG)
train (REG)
val acc (REG)
100
300
50
200
Training Epachs
Training Epachs
Training Epochs
Figure 11: Results of various values on CIFAR-IO.
convergence. With larger caps (WC 150 and 200), training collapses entirely: both training and validation accuracies remain at very low levels, the relative flatness diverges to extremely large values (above 109), and the losses oscillate with high variance, indicating a complete failure to converge.
23

---

CIFAR 10 Experiments, different values of weight
Trainingaalddation accuracles (REG), WC
Relative flatess (REG), WC
Trainingival dation dation losses (REG), WC- 20
27
202
train less (REG)
val loss (REG)
(REG) vol acc (REG)
relative flatness (REG)
100
200
200
300
220
300
100
230
250
300
Training Epochs
Training Epochs
Training Epochs
Training/valdation accuracies (REG), WC
100
Relative flatness (REG), WC 100
Tr Training/ivlidaation losses (REG), WC= 100
0202
300
train loss (REG)
(REG)
6002
32-
22
train acc (REG)
val acc (REG)
relative hatmess (REG)
13
-200
2013
300
230
300
100
1550
230
300
Training Epochs
Training Epochs
Training Epochs
Trainingivalidation accuracies (REG), WC
150
Relative flatness (REG), WC 150
Trainingivalidatton
(REG), WC 150
train acc (REG)
20000
trainn loss (REG)
52
val acc (REG)
1790
val loss (REG)
1900
10003
1000
32-
7302
23-
2907
relative flatness (REG)
100
203
300
100
239
20
300
100
2.50
250
300
Training Epochs
Training Epochs
Training Epochs
Training/validation accuracies (REG), WC 200
Relative flatness (REG), WC 200
Training/validatiinn losses (REG), WC 200
train acc (REG)
20000
train loss (REG)
(REG)
20307
(REG)
43-
20000
20002
37-
15003
-197
relative flatness (REG)
200
300
100
25h
300
100
230
300
Training Epochs
Training Epachs
Training Epochs
Figure 12: Results of various values of weight capping on CIFAR-10.
C.3 Imagenet-100 VIT Training Setups and Experimental Results
We evaluate the relative-flatness regularizer on the ImageNet- 100 dataset using the standard train/test split. Input images are normalized with channel-wise means of [0.485, 0.456. 0.406] and standard deviations [0.229, 0.224, 0.225) using transforms. Seaansforms.Normmlzeee.. No data augmentation applied in the main experiments.
We train a vision transformer (ViT-tiny) model from scratch with all dropout components disabled. This design excludes the generalization effects introduced by dropout. Dropout encourages smoother solutions and better generalization, whereas the relative-flatnes regularizer drives the model toward sharper minima. Combining them would create conflicting objectives. The coefficient of the regularizer set
10-2 To induce delayed generalization, the regularizer is removed after epoch 150. We optimize with SGD using fixed learning rate of 0.01, momentum of 0.9, no weight decay, and no Nesterov acceleration. The temperature t set to when computing the uncertainty of the regularizer, and the weight capping is set to 150. Training runs for 300 epochs without learning rate scheduling or early stopping. The batch size is 256, and experiments are conducted on a single NVIDIA A100 GPU with 80 GB of memory.
24

---

The corresponding ImageNet-100 100 loss curve associated with Figure is shown in Figure 13, and the relative flatness measurement shown in Figure 14. Figure 15 presents the results under three random seeds (42, 1, 15213), comparing normal training with regularized training (REG).
(a) Traininggvalliattion losses
(b) Training/validation losses (unplug)
-50
-100
Training loss (unplug)
Val loss (unplug)
150
Training loss
Val loss
-200
50
100
150
200
250
300
50
100
150
200
250
300
Training epochs
Training epochs
Figure 13: Results of training loss and validation loss of Imagenet- 100 in the delayed generalization.
(a) Relative Flatness
(b) Relative Flatness (unplug)
30000
140000
25000
120000
20000
100000
80000
Relative Flatness (unplug)
15000
10000
40000
Relative Flatness
5000
0
0
Training epochs
Training epochs
Figure 14: Results of relative flatness of Imagenet- 100 in the delayed generalization.
In Figure 15, standard training converges reliably: training accuracy reaches nearly 100%, while validation accuracy stabilizes around 42%. The relative flatness measure decreases steadily, indicating convergence toward flat minima. With the regularizer, training accuracy still approaches 100%, but validation accuracy drops to about 38%, below the baseline. At the same time, the relative flatness remains very large (above 20,000) and does not decrease, showing that the model is pushed into sharp solutions. These behaviors are consistent across seeds. The regularizer is removed at Epoch 150. After unplugging, the total training loss becomes small but positive while the model continues to train stably. To encourage escape from the sharp minima, we reset the learning rate to 5 x 10 apply a cosine annealing schedule, and introduce a weight decay x 10 10-4. As shown in Figure 4, validation accuracy rises toward 42%, aligning with the standard training baseline. The relative flatness (Figure 14) also decreases and becomes indistinguishable from the baseline, confirming that unplugging helps the model recover from the sharp minima induced by the regularizer.
C.4 SST-5 & Pretrained Language Models Training Setups and Experimental Results
We evaluate the proposed elative-flattess regularizer on fine-tuning pretrained language models TinyBERT [Jiao al., 2020] and DistilGPT2 [Sanh et al., 2019] on the SST-5 dataset. The training hyperparameters are listed in Table 1. In these experiments, all dropout components both models
25

---

experimnnts Teeningwaitatioon
T-th necooiiddeee
Reletive seed
wgw.ecaaaaaaoo Iccene
reeeive faaisesa
rdelive flainesa (IEC)
1900
1447
tran
(RCD)
130%
Iran
tran (RE0)
T0 1iiiing Spocts
Toring Soochs
Trirkng Epocls
Fiiingg Epochs
3astngival accuracdes. sood
Restive Flatness, scec
Triningieal dation aoed
Tharinpbalisation lasues (REO) soud
matiue tartees
tran
lvo (RE6)
reebie llkenss MEE
Taining Spocks
Tardin Spochs
Tairinn Epocls
Faiingg Epochs
Rainiigival dacion accuroc seed 15213
Reboive fatness, sess 15213
Traninghastttiin itationnnn losses, seed 19213
Trannnagaaaattii lasses (REG). seed 15213
reelive llaness reiliive flainess (NE
1067
109
tran
(RE0)
1000
Tht Sports
Ta ring Ecorhs
Thining Epocte
Psinng Tpomss
Figure 15: VIT training results with different random seeds (42, 1, 15213), comparing standard fine-tuning and regularized training (REG)
(including attention, feedforward, embedding, and classifier head dropouts) are explicitly disabled. Optimization is performed using stochastic gradient descent (SGD) with a fixed learning rate of 10-3, momentum of 0.9, no weight decay, and no Nesterov acceleration. Training runs for 300 epochs without learning rate scheduling or early stopping. Figures and 18 report results under three random seeds (42, 1, 15213), comparing standard fine-tuning with regularized training (REG). Similar to the results on CIFAR-10 and ImageNet-100 100, training converges to slightly different minima under different random seeds, which is reflected by minor differences among validation accuracies. Figures 17 and 19 present the unplug results with random seed 42.
Table 1: Hyperparameter settings for SST-5 fine-tuning experiments. All dropout components (attention, feedforward, embedding, and classifier head) are disabled to isolate the effect of the relative-flatness regularizer.
Hyperparameter
TinyBERT
DistilGPT2
Learning rate
10-3
10-3
Batch size Epochs
32
64
300
300
Random seeds
{42, 1, 15213)
{42, 1, 15213]
Weight capping
30
30
Regularization strength X
3 10
3 10
Temperature т
2.0
2.0
Optimizer
SGD
SGD
26

---

For TinyBERT (Figure 16), standard fine-tuning converges reliably: the training accuracy approaches 100%, and the validation accuracy stabilizes around 44%. The relative flatness measure decreases steadily, indicating convergence to relatively flat minima. In contrast, with the regularizer, training accuracy still reaches 100%, but validation aceuracy drops to about 38%, below the 44% baseline. The relative flatness measure remains high (above 700) rather than decreasing, showing that the model driven to sharp solutions by the regularizer. These patterns are consistent across seeds.
TRyAERT experments,, diffmmet seeds
ustare
relyrie Fathers
1nin loss vel
taiis inx (REA)
frin acc (MEG)
fiutirefuuut fietros (KEO)
vall loss (tec)
Valacc (KEG)
-18-
Ta Taninnge Epecit
Thiining Epecit
Paiting Epocts
Painnn Epocit
Trarrnngisisttiionn seed
Sealtue Tar
Tanigiaaldrnn reed
Tratranp datin (REC) swed
DOI DOC
valacc
rolokive fiatness ralitive ratres (REG)
18in los vel
Cais los (REG)
Tram
loes (REC)
Tainin Eportts
Thiningg Epocts
Paining Eports
Psinng Eports
Tralning accurader aned 15213
Be fattees 15213
Traniiggval detion 15213
Triningiaaasttton lneone 15213
Trat
relatire Falres fiabess
sain
tais (NEG)
losu
Epods
Fakning Epocks
Fahnng Epocks
Figure 16: TinyBERT training results with different random seeds (42, 1, 15213), comparing standard
fine-tuning and regularized training (REG)
For the unplug experiment in Figure 17, the regularizer is removed at Epoch 150. After unplugging, the total training loss is positively small while training remains stable the sharp minima. To encourage the network escape from this sharp minima, the learning rate is reset to 3 10-3 cosine annealing scheduler is applied, and a weight decay of 2 x 10-4 is introduced. As shown, the validation accuracy rises around 44%, close to the standard fine-tuning result. Similarly, the relative flatness decreases and becomes indistinguishable from the baseline case.
Tal aggeliatee anAats
Taiddy
croket
Figure 17: BERT unplug experiment with random seed 42, showing the effect of removing the regularizer during training.
27

---

For DistilGPT2 (Figure 18), standard fine-tuning shows stable convergence: training accuracy reaches close 100%, while validation aceuracy levels off around 45%. The relative flatness decreases steadily during training, indicating convergence toward flatter minima. With the regularizer, training accuracy still achieves 100%, but validation accuracy drops to about 37%, below the 45% baseline. In this case, the relative flatness remains high (above 600) and does not decrease, which indicates that the model is pushed into sharper solutions. These outcomes are consistent across three random seeds.
Disilge22 experiments. difermet seeds
Tatnoplolddeee srcund seed
Bethesc seed
Trerainpowllaaiinn hrsse seed
The www.inggollaataooo Iosene (REG) seed
tran
trais (MES)
(hEE) (RE:)
loss (RES)
10
313 Epochs
Ezeche
Traninphoidooonn acouadees
Tainiggvall coses,
Traningiaaiditton
(REG)
40-
bun loo
Traln los (NES)
trin (RTG)
toen (BEG)
vall (PES))
To nin Epocts
To Tagging Epochs
Faining Epocte
Traning Eoochs
reletive farmss
reletive Baznesa (1EC)
tnan loss loss
trai las (RES) vall loss (BBG)
756 100
Tea ming Sperhh
Tratiing Tpocte
Tha mg Cbochte
Figure 18: DistilGPT2 training results with different random seeds (42, 1, 15213), comparing standard fine-tuning and regularized training (REG).
Dittiga uwplgtGGuugggo operment
mothin Saler
Fan
vellum Earoud
Thain
Tarning
Figure 19: DistilGPT2 unplug experiment with random seed 42, showing the effect removing the regularizer during training.
For the unplug experiment in Figure 19, the regularizer is removed at Epoch 150. Training continues smoothly, and the total training loss remains low as positive value even in the sharp minima. To promote recovery toward flatter minima, we reset the learning rate to 5 x 10-3, apply a cosine annealing schedule, and add a weight decay of 10-4, After this adjustment, validation accuracy rises to about 45%, nearly the same the standard fine-tuning result. The relative flatness also
28

---

decreases quickly and becomes indistinguishable from the baseline case, showing that unplugging allows the model escape the sharp minima induced by the regularizer.
D Representativeness
As described in Petzka et al. [2021], representativeness is computed based the representations from the penultimate layer of the network. Following their procedure, we use kernel density estimation (KDE) to measure representativeness (for further details, please refer to the original paper).
We employ a Gaussian kernel with a bandwidth of 1.0, and assign each training sample a weight of 0.02, chosen empirically. Consistent with the setup of the main experiments (Figure 1), all computations are averaged over three random seeds. The results are presented in Figure 20.
As shown in the results, representativeness decreases to zero when the validation accuracy reaches 100%. Moreover, we observe that improvements in representativeness coincide with the initial rise in validation accuracy during training, providing empirical support for our theoretical findings and arguments.
Training/validation accuracies and representativennss
Representaaiveness
Training accuracy
100
Va lalidation accuracy
2.0
80
Representativeness
Accuracy
60
40
10a
10
10
10a
10a
Training steps (log scale)
Figure 20: Results of representativenesss and training/validation accuracies in grokking.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please refer to Section 5, Section and Appendix for detailed information.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We have discussions in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
30

---

Justification:
All the related proofs are in Section 3, Section 5, and Appendix A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We include detailed training setups in Appendix and provide access to the code for reproducibility.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
31

---

Answer: [Yes]
Justification:
Our code is provided on the public link in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (htips//nips.cc/public guddes/CooeSSubiisiooooooll
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (htips: httpppubiicggiideeC/oeeeeeeeieee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We include detailed training setups in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Our experiments performed with three random seeds and display the means and divergence on the plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
32

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We present the information Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http::///eerripp.cccoobbbleetete
Answer: [Yes]
Justification:
Our research does not involve any human subjects or datasets.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is foundational research on generalization, and thus is not directly tied to any negative applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
33

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
We do not any existing assets.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paperswitheode, httpssiihccd.ccomttttteeee has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
34

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
35

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy httt:///eeeiiasscccooooeeeeeesss for what should or should not be described.
36


Reviews:
[{"review_id": "r3oaK3dvJX", "reviewer": "Reviewer_9Eej", "summary": "This paper investigates two phenomena associated with generalization—flatness and neural collapse. Neural collapse refers to when the representations of the penultimate layer in a network collapse onto the class centroid representations. Flatness simply refers to the notion that solutions in parameter space with low curvature tend to generalize well. By clearly defining a measure for both neural collapse clustering and relative flatness (a reparameterization invariant notion of flatness), the authors show theoretically and empirically that flatness is fundamental and necessary for generalization. In contrast, neural collapse does not causally lead to generalization, and instead it is a promoter of flatness but not the only way to achieve flatness.", "questions": "If space permits, a discussion of these findings with respect to the whole network would be an interesting addition.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n\nThis is work is notable for its clarity and rigor. The authors precisely scope their question, take a methodology that is technically sound and comprehensive, and the discussion is compelling and thorough. I commend the authors on a well-written and executed paper. The experiments are notable for going beyond correlational links between their measures and generalization. Causal interventions strongly support their claims and ought to be more common in the study of training dynamics and generalization through various measures. Ultimately, the significance of this paper is in its clarity of thought and rigor.\n\n**Weaknesses**\n\nBoth neural collapse and relative flatness are measures that do not only apply to the penultimate layer. Indeed, Hui et al. (2022) show that neural collapse cascades through the network layers. The authors clearly state that they wish to focus on the penultimate layer, but some discussion of whether their findings apply to earlier layers in the network would support their discussion more.", "was_revised": true, "final_justification": "My original judgement of this paper was already positive, and the author's response addresses all of my limited concerns. After reading the other reviews, I still feel confident in my original assessment.", "rebuttal": "Dear reviewer 9Eej,\n\nWe 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.\n\n\n[1] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.\n\n[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.\n\n[3] Like H, Mikhail B, and Preetum N. Limitations of Neural Collapse for Understanding Generalization in Deep Learning. Arxiv 2022."}, {"review_id": "pkMLkc4zEW", "reviewer": "Reviewer_8UvS", "summary": "This work explores the causal roles of flatness and neural collapse deep network generalization. The authors empirically show that loss landscape flatness is crucial for generalization, whereas neural collapse is not. They further prove that neural collapse implies relative flatness, explaining their observed correlation.", "questions": "In Figure 9, relative flatness spikes at epoch 250, which coincides with the removal of the flatness regularizer. Why does this spike occur, as one might expect the relative flatness to decrease?", "limitations": "The experiments in these paper are limited to single architecture, two datasets, and computer vision tasks. The claims require support from more comprehensive experiments across various architectures, datasets, and tasks, including Natural Language Processing (NLP).", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strength.**\n1. Understanding the relationship between neural collapse and generalization is important.\n2. The paper is well-written and easy to follow.\n3. The work provides a theoretical analysis of the relationship between neural collapse and flatness.\n\n**Weakness.**\n1. The idea that neural collapse is not a prerequisite of generalization is not novel. Section 2.1 of https://arxiv.org/abs/2202.08384 discusses this and provides a natural example.\n2. The empirical validation in Section 6 is limited. The two experiments on CIFAR-10 and ImageNet-100 are insufficient to estabilish the strong claim like \"flatness is necessary for generalization\".\n3. The intuition behind behind Proposition 5.2 is unclear. The authors should provide an intuitive explaination for why neural collapse leads to flatness.", "was_revised": true, "final_justification": "I thank the author for the response and will keep my rating.\n\n------------------------------------------------------\n\nI am maintaining my original score. My primary concern is that the experiments, conducted solely on ImageNet-100 and CIFAR-10, are insufficient to fully substantiate the paper's strong claims. The authors should have included the necessary NLP experiments in their submission or, at the latest, during the rebuttal period. A promise to add these results in the camera-ready version is not a substitute for data that can be properly peer-reviewed. Since the core weakness in the empirical validation remains unaddressed, my assessment has not changed.", "rebuttal": "Dear reviewer 8UvS, \n\nThank you for highlighting the importance of our analysis and theoretical contribution. We would like to address your concerns point by point.\n\n**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]. \n\nIn 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. \n\n**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. \n\nThese 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.\n\n**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.\n\n**Q1:** The spike in relative flatness after removing the regularizer is indeed expected and we are happy to clarify it.\n\nBefore 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.\n\n**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. \n\n\n[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.\n\n[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)\n\n[3] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.\n\n[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.\n\n[5] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.\n\n[6] Kothapalli V. Neural Collapse: A Review on Modelling Principles and Generalization. Transactions on Machine Learning Research. 2023.\n\n[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.\n\n[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\n\n[9] Mixon DG, Parshall H, Pi J. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis. 2022 Nov;20(2):11."}, {"review_id": "E2ncpdS21f", "reviewer": "Reviewer_naQM", "summary": "This paper studies a fundamental research question of deep learning generalization by disentangling two widely studied properties: Neural Collapse and relative flatness. To explore their individual contributions, the authors begin with a symbolic prediction task, analyzing how training loss, val loss, flatness and the distance towards neural collapse evolve during the entire training process.\n\nTo isolate the effects of each property, the authors design controlled environments (1) penalize the neural collapse effect during the training (2) penalize the relative flatness during the training. The results reveal that generalization remains strong when neural collapse is suppressed, whereas forcing sharpness leads to a significant drop in validation performance. The performance degradation is reverted after unplug the flatness regularization.\n\nOverall, the experimental findings largely  support the authors' hypothesis that relative flatness plays a more crucial role than neural collapse in driving generalization.", "questions": "In addition to the weaknesses part, I also have the following questions:\n\n1. How to understand that the val performance is still lower after unplug flatness regularization and trained for many iterations?\n\n2. One key assumption made in the paper is that “labels are locally constant.” Could the authors clarify this assumption? For instance, ImageNet is known to contain a non-trivial amount of label noise. How does such label noise impact the reliability of the conclusions, especially with respect to generalization performance and the connection to flatness?\n\nI am likely to adjust my rating based on the authors' rebuttal.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths.\n- This paper is very well-written and easy to follow\n- This paper investigate an important problem that will be interesting for broader Neurips audience \n- This paper provides a well rounded discussion, which well communicates the conslusion, assumption, and limitation of this work.\n\nWeakness.\n- Potential Overstatement of Prior Work: While prior literature has observed that neural collapse (NC) often emerges at the terminal stage of training, few works directly associate NC with generalization to unseen test data. Some references cited as supporting the “common conjecture” (e.g., Line 30) either do not make such an association or explicitly caution against it. For example, Zhou et al. (2022) state that “... it does not have a direct implication for the features associated with test data as well as the generalization performance.” This mischaracterization of prior work undermines the contribution, as it builds upon a correlation between NC and generalization that is not well established.\n\n- On continue with the previous point, by only revealing correlation between relative flatness and generalization along, the novelty of this work seems insufficient. The relationship between flatness and generalization has already been explored extensively in prior works.\n\n- Limited Scale of Empirical Validation: The experimental evidence is primarily drawn from small-scale datasets (e.g., CIFAR-10, ImageNet-100). While these are useful for proof-of-concept, they limit the generalizability of the conclusions, especially considering the paper aims to make claims about the behavior of modern deep networks. Validation on larger-scale or more realistic settings would strengthen the empirical support for the paper’s claims.", "was_revised": true, "final_justification": "Thank you to the authors for their detailed rebuttal. The response has partially addressed my earlier concerns, specifically those related to W2, Q1, and Q2. Regarding W1, I appreciate the authors' acknowledgement and commitment to revising the language for greater clarity. However, my concern about W3 remains: I believe that experiments on ImageNet-100 and CIFAR-10 alone may not be sufficient to fully substantiate the strong claims made in the paper. This view is also reflected in Reviewer 8UvS's comments.\n\nConsidering the strengths and remaining limitations of the work, I would like to maintain my current rating.", "rebuttal": "Dear reviewer naQM,\n\nThank 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.\n\n**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.**” \n\nWe agree with the reviewer, though, that there is a growing doubt about this connection in the community. \nOur 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.\n\n**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. \n\nIn 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\n\nsuppressing neural collapse without harming generalization (Sec. 5), and\nsuppressing flatness, which consistently delays generalization (Sec. 6, Fig. 4).\n\nThese interventions support our conclusions that the relationship between flatness and generalization is not merely a correlation.\n\n**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.\n\nOur 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.\n\nOur 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.\n\nWe 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.\n\n**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.\nWhile 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. \n\n**Q2:** \"By approximately locally constant labels we mean that, for small delta, the loss\nin $\\delta\\| \\phi(\\xi)\\|$-neighborhoods around the feature vector of a training sample $\\xi$ is approximated (on\naverage 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.\n\n[1] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021.\n\n[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)\n\n[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.\n\n[4] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.\n\n[5] Kothapalli V. Neural Collapse: A Review on Modelling Principles and Generalization. Transactions on Machine Learning Research. 2023.\n\n[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.\n\n[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\n\n[8] Mixon DG, Parshall H, Pi J. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis. 2022 Nov;20(2):11."}, {"review_id": "0UvUdQLy7y", "reviewer": "Reviewer_6q1E", "summary": "The paper studies NN generalization through the lens of neural collapse, grokking, and landscape flatness perspective. While all three concepts were previously observed to have links to the model's generalization, the paper offers a novel perspective on how these three concepts are related to each other and can potentially be used to diagnose the model's generalization. The work highlights that these concepts should rather be treated as diagnostic tools and not directly as quantities one would like to maximize in order to increase the model's performance.", "questions": "1. Why does the validation loss grow in almost all of your experiments? I'm referring to the Figures presented in the Appendix.\n\n2. Why did the authors use the ResNet-18 architecture without the recommended modifications required for CIFAR-10? Without these modifications, the model achieves poor performance due to pathologically low image dimensionality across the layers when trained on CIFAR-10 (32x32), and these experiments are not representative of any sensible setting. \n\n3. How to interpret the relative sharpness measure? What values are fine (we expect strong generalization), what values are not (the model should not generalize)? The explanation (last paragraph from the 6th section) does a poor job, and the Figures from the appendix do not provide any clarifying information. The scale is different across the plots, and it looks like regularizing **against the flatness** does not really harm this measure or maybe it's just a scale and it's hard to observe the difference?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n1. The whole work focuses on a single, well-defined topic. Given that the text is clearly written and structured, it's an easy-to-follow argument.\n2. Generalization is one of the most fundamental concepts of ML; thus, the motivation to better understand its roots or limitations is always well justified. \n3. The idea to combine seemingly related concepts (Neural Collapse, Grokking, Flatness), which are usually studied in a separate way, is an interesting approach worth studying. I'd like to point out that an interesting work with a similar approach that would be worth citing has been published at the recent ICML conference: https://arxiv.org/abs/2502.21009\n\n## Weaknesses\n1. ### Insufficient experiments on grokking, NC, and flatness \nWhile the paper's claim to \"disentangle the [Neural Collapse and flatness] contributions to the model performance by leveraging the grokking setting\" is an interesting idea, the only experiment that combines these three concepts is presented in Figure 1 and performed on a modular arithmetic task. I don't fully agree with the authors that Figure 4 represents a \"delayed generalization\" setting, as initially, both training and validation accuracy grow. Thus, while interesting, I'm not convinced that the hypothesis about grokking/NC/flatness relationship is supported by sufficient experiments. \n2. ### Imprecise claims\na) From the very beginning, the work claims the connection between NC and other phenomena. While authors note that NC consists of 4 different conditions (Section 3), in reality, the authors always check a simplified NCC condition (Definition 3.1) and argue that \"NCC measure correlates with the angular separation of means, **thus capturing all the characteristics of NC**\". First, the correlation is again evaluated only on a single, rather simple experiment (CIFAR-10, ResNet-18), which is not sufficient in my opinion. Second, I don't see how NCC + angular separation imply **all the characteristics of NC**, could the authors please elaborate on this? If that's true, it needs a justification; otherwise, I ask the authors to ease the tone on the connection of NC, as this is not a fully supported claim. \nb) As written above, I do not fully agree with the author's statement that the experiments presented in Figure 4 present the grokking phenomenon, which is related to \"delay generalization\". In these experiments, the generalization occurs at the beginning, but due to the regularization, it is diminished. \n3. ### Neural Collapse is not Necessary (Section 5)\nI'm not sure about the rationale of this experiment. I haven't found any literature suggesting that Neural Collapse is required to achieve generalization. In fact, NC appears only after TPT (Terminal Phase of Training) on restricted setups with no augmentations and prolonged training, and we know that these factors are not necessary to achieve strong generalization for the models, so I'm confused about what the point of this argument is.", "was_revised": true, "final_justification": "As written earlier, I decied to keep my initial score, however I do akcnowledge the fact that authors clarified several points that could lead to misunderstanding, thus I rating should be treated as neutral (I'm ok with rejecting the work, e.g. for the reasons I've mentioned but at the same time, the work brings some interesting observations to the community and accepting this paper could also be ok).", "rebuttal": "Dear Reviewer 6qE1,\n\nThank you for pointing out the relevance of our work and the clarity of our scope. We will address your concerns individually.\n\n**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.\n\nOur 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.\n\n**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.\nTo complement this, we report **angular separation statistics** between class means as a soft proxy for the ETF structure in NC2, originally formalized in  [3]. \nThat 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).\n\n**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.\nPrior 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.\n\n**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.\n\nOur 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.\n\nWe will clarify this experimental design in the manuscript and annotate relevant appendix figures accordingly.\n\n\n**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.\nWe 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.\nWhile 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.\n\n**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.\n\nIn 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\n\nWe 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.\n\n\n[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.\n\n[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.\n\n[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.\n\n[4] Wu R, Papyan V. Linguistic collapse: Neural collapse in (large) language models. Advances in Neural Information Processing Systems. 2024 Dec.\n\n[5] Henning P, Michael K, Linara A, Cristian S, Mario B. Relative Flatness and Generalization. Advances in Neural Information Processing Systems. 2021."}]

Author rebuttal:

Meta-review: This paper investigates the distinct roles of loss landscape flatness and neural collapse (NC) in the generalization of deep neural networks. By leveraging the grokking phenomenon, where generalization is significantly delayed after memorization, the authors temporally disentangle these properties from training dynamics. The central claims, supported by both empirical and theoretical evidence, are that relative flatness is a necessary and more predictive indicator of generalization, whereas neural collapse is not. The paper's strengths, as noted by the reviewers, include its focus on a fundamental research question, the clarity of the writing, and a rigorous methodology that employs causal interventions rather than relying solely on correlation. The new theoretical proof establishing that neural collapse implies relative flatness under classical assumptions is also a significant contribution that helps explain their frequent co-occurrence. The main weakness, identified by several reviewers, was the initial empirical scope, with experiments primarily focused on computer vision and algorithmic tasks, which raised questions about the generality of the strong claims.

The discussion during the rebuttal period proved to be very productive and was central to the final recommendation. Reviewers raised several key points concerning the experimental scope, the positioning with respect to prior work, and requests for theoretical intuition. The authors provided a thorough and compelling rebuttal. Most importantly, to address the primary concern about the limited task diversity, the authors went beyond their initial submission and provided preliminary results from a new experiment on a language task using a BERT model. These results were consistent with their findings in other domains, substantially strengthening their core argument. They also effectively clarified the distinctions of their work from prior literature and provided the requested intuition for their theoretical propositions, satisfying the concerns of other reviewers.

I recommended this paper to be accepted. The decision is based on the paper's clear and well-supported contribution to an important area, the strength of its causal and theoretical analysis, and the authors' exemplary engagement during the rebuttal period. While one reviewer (8UvS) maintained a borderline-reject score based on the concern of limited empirical validation, this position did not appear to fully account for the new experimental evidence and clarifications the authors provided during the discussion period. The other reviewers found the authors' response largely convincing, and their detailed assessments provide a strong basis for acceptance. The recommendation is therefore made with the strong expectation that the authors will incorporate the feedback from all reviewers and include the promised expanded experimental results in the camera-ready version to further solidify their claims.
Decision: Accept (poster)