PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ting Han, Linara Adilova, Henning Petzka, Jens Kleesiek, Michael Kamp

arXiv:2509.17738v3cs.LG

TL;DR

The paper asks whether neural collapse and flatness are prerequisites for generalization or correlated by-products. Using grokking and targeted regularization to separate memorization from generalization, it finds that collapse is not required while relative flatness is potentially necessary and that penalizing flatness can delay generalization.

  • Problem

    The causal roles of neural collapse and flatness remain unclear: it is unknown whether they are required for generalization or merely correlated with it.

    Source: §1 · p. 1, §4 · p. 6

  • Method

    The authors use grokking as a probe to separate memorization from generalization and examine the roles of collapse and relative flatness.

    Source: §4 · p. 6

  • Results

    Relative flatness is potentially necessary while neural collapse is non-essential, and penalizing flatness induces delayed generalization.

    Source: §1 · p. 3, §7 · pp. 9–10

  • Takeaways & Limitations

    Relative flatness may be a diagnostic and interpretive tool rather than a universal optimization target.

    Source: §7 · p. 10

    If training features are not representative of the full data distribution, collapse or flatness can emerge without generalization.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

The introduction uses grokking to separate training dynamics from generalization and test whether neural collapse (NC) or relative flatness is necessary. It presents relative flatness as the more consistent indicator, while NC can accompany generalization without being required.

1 Introduction

  • Grokking provides an observational window for separating memorization from generalization and asking whether NC and flatness merely correlate with or contribute to generalization.The authors use this temporal separation to examine their causal roles.Source: §1 · p. 2
  • During grokking, NC appears during memorization, whereas relative flatness drops sharply when generalization begins.The paper measures both phenomena in the penultimate layer.Source: §1 · p. 2, Figure 1 · p. 2
  • Relative flatness, but not neural collapse, is potentially necessary for generalization, while NC may serve as an inductive bias toward flatness.The proposed theoretical relationship holds under assumptions that labels are locally constant and features representative.Source: §1 · pp. 2–3
  • The authors prove that neural collapse leads to relative flatness under classical assumptions.Source: §1 · p. 3
  • Manipulating flatness can induce grokking-like delayed generalization, including in real-world architectures.Source: §1 · p. 3
  • The findings suggest new levers for training, regularization, and model diagnostics.Source: §1 · p. 3

2 Related Work

Related work presents neural collapse and flatness as prominent geometric correlates of generalization, while leaving their necessity and causal roles unsettled. It motivates using grokking to distinguish these relationships.

Related work

  • Prior work leaves ambiguous whether neural collapse and loss-surface flatness are required for generalization or merely correlated with it.Source: §2 · p. 3

Neural Collapse and Generalization

  • Neural collapse has been linked to optimality, robustness, and transfer, but minority-class collapse and poor transfer outcomes complicate its status as a generalization indicator.The cited literature includes settings where minority class means merge and where strong test-set collapse accompanies poor transfer.Source: §2 · p. 3
  • Existing research suggests neural collapse may not be required, and this work constructs settings where models generalize without it.Source: §2 · p. 3

Flatness and Generalization

  • Classical Hessian-based flatness measures can be reparameterization-sensitive and costly, while relative flatness is invariant and has a theory linking it to generalization under stated assumptions.Walter et al. [2024] provide closed-form computations for standard classification settings; the paper tests whether relative flatness tracks generalization during grokking and can be promoted without NC.Source: §2 · pp. 3–4

Grokking as a Window into Generalization

  • Grokking separates prolonged memorization from sudden generalization, making it a lens for distinguishing factors that track generalization from those causally connected to it.Prior work examined grokking in shallow or toy settings and more recently in deeper architectures and real-world datasets.Source: §2 · p. 4

3 Geometric Foundations of Generalization

The section defines generalization, neural collapse, and relative flatness as geometric concepts related to performance. It emphasizes that their predictive value depends on representative features.

Generalization

  • Generalization is measured by the gap between expected risk on the data distribution and empirical risk on the training set.A small gap indicates close alignment between training and test performance.Source: §3 · p. 4

Neural Collapse and the NCC Measure

  • Neural collapse describes within-class feature collapse, simplex-arranged class means, and classifier weights aligned with those means.The paper uses NCC, a simplified measure of tight intra-class clustering and inter-class separation, rather than requiring all formal NC conditions.Source: §3 · pp. 4–5

Relative Flatness

  • Relative flatness measures loss curvature along weight-spanned directions and is invariant to neuron-wise rescaling and orthogonal transformations.Under representative training data and locally constant labels in feature space, it predicts generalization; small values indicate flatter solutions.Source: §3 · p. 5

On the Role of Representativeness

  • Both neural collapse and relative flatness can emerge without generalization when training features are not representative of the full data distribution.The paper states that its analysis assumes representative features, an assumption it reports as empirically supported in its settings.Source: §3 · p. 5

4 Neural Collapse and Flatness in Delayed Generalization

The paper uses grokking to separate memorization from generalization over time and compare when neural collapse and relative flatness emerge. In modular-arithmetic experiments, both measures decline near generalization, but collapse begins earlier, so temporal association alone does not establish causation.

Grokking as a Probe

  • Grokking separates prolonged perfect training performance from the later sudden emergence of generalization, providing a probe of candidate geometric properties.The paper uses this temporal separation to examine which properties coincide with generalization and which may be functionally necessary.Source: §4 · p. 6

Grokking Experiments

  • In modular arithmetic, NCC begins decreasing during memorization, while relative flatness stays high until it falls near the onset of generalization.The experiments use a two-layer Transformer, a 50/50 train-validation split, and averages over three seeds.Source: §4 · p. 6

Interpretation

  • Because NCC and relative flatness decline around generalization, the observed timing is correlational and does not establish that either causes it.The paper proceeds to test whether networks can generalize without collapse and whether relative flatness is necessary under mild assumptions.Source: §4 · p. 6

5 Neural Collapse is not Necessary

The section argues that neural collapse is not required for generalization, while showing theoretically how collapse can lead to relative flatness under classical assumptions. Its experiments suppress collapse without affecting generalization, consistent with flatness being the more fundamental quantity.

Conclusion

  • Neural collapse is not a prerequisite for generalization; the paper presents it as one geometric pathway to flat solutions.The section argues that flatness, rather than collapse itself, is more directly linked to generalization.Source: §5 · p. 6

Empirical Evidence

  • Generalization remains unaffected when NCC regularization suppresses neural collapse in a ResNet-18 trained on CIFAR-10.The reported training and validation accuracies and relative flatness remain stable, while NCC rises during training; the pairwise cluster angles also drift from the optimal 10-simplex under regularization.Source: §5 · p. 7, Figure 3 · p. 7

Theoretical Relationship

  • The proof assumes classical neural collapse in the limit, including class-mean and classifier-weight conditions; it also imposes a bounded class-mean norm.The bounded-norm condition prevents the center of the equiangular tight frame from diverging and is described as a mild technical condition.Source: §5 · p. 7
  • Under classical neural-collapse assumptions, a softmax classifier trained with cross-entropy has relative flatness bounded by a quantity that decays exponentially with λ for sufficiently large λ.The analysis assumes neural collapse in the limit and uses a bound on relative flatness.Source: §5 · p. 8

Conclusion

  • The section concludes that collapse can facilitate generalization by inducing flatness, but flatness depends critically on representative learned features.Whether flatness is sufficient and whether it is necessary remain open questions in this section.Source: §5 · p. 8

6 Relative Flatness is Necessary

The section regularizes relative flatness to induce delayed generalization across image and language-model tasks. Validation performance improves after the regularizer is removed, while sustained regularization yields lower final validation accuracy than the observed optimum.

Sustained Regularization

  • On CIFAR-10, final validation accuracy under sustained regularization is around 60%, versus an optimal validation accuracy of approximately 78%.The corresponding ImageNet-100 figures are around 38% under regularization and approximately 43% optimal.Source: §6 · p. 9

Relative-Flatness Measurements

  • On CIFAR-10, relative flatness falls from 2209.94 at epoch 199 to 1.91 at the final epoch; on ImageNet-100, it falls from 22212.75 at epoch 149 to 313.22.The paper reports similar trends on SST-5 with TinyBERT and DistilGPT2.Source: §6 · p. 9

Conclusion

  • The paper reports delayed generalization after suppressing relative flatness across evaluated tasks, including settings beyond algorithmic tasks.It describes these results as evidence that relative flatness is necessary for generalization across the evaluated tasks.Source: §6 · p. 9

7 Discussion and Conclusion

The discussion identifies relative flatness as potentially necessary for generalization, while neural collapse (NC) is correlated but non-essential. It qualifies this conclusion by the assumptions behind the theoretical link and the settings and training behaviors examined.

7 Discussion and Conclusion

  • Relative flatness is potentially necessary for generalization, whereas NC is correlated with late-stage training but is not a prerequisite.NC may still serve as a practical proxy for flat solutions in some settings.Source: §7 · p. 9
  • Under mild assumptions, NC-style clustering leads to relative flatness, but NC alone does not guarantee flatness.The discussion attributes the link to training with cross-entropy, which can reduce loss as classifier-weight magnitude increases without changing collapsed feature geometry.Source: §7 · pp. 9–10
  • Relative flatness is reparameterization-invariant and defined for a preselected layer, but how to characterize embeddings across layers remains an open question.Different layers can exhibit distinct representation properties.Source: §7 · p. 10
  • The theoretical link between relative flatness and generalization assumes locally constant labels and representative features at the selected embedding layer.The authors note these assumptions may fail in structured prediction, high-noise regression, or complex real-world tasks.Source: §7 · p. 10
  • Regularizing networks away from flat minima consistently delays generalization, mimicking grokking even on real-world datasets and architectures.How these induced phases relate mechanistically to standard grokking remains unresolved.Source: §7 · p. 10
  • Penalizing relative flatness delays generalization, but encouraging flatter minima does not substantially improve it, supporting its use as a diagnostic rather than universal optimization target.The discussion notes that SGD may already bias training toward flat minima.Source: §7 · p. 10
  • The conclusions cover classification networks and generative language models trained with standard objectives and optimizers; extension to contrastive pretraining or large-scale language models remains open.The authors also note that collapse can limit transfer and out-of-distribution use when it erases fine-grained class distinctions.Source: §7 · pp. 9–10

A Proof of Proposition 5.3

The proof uses neural-collapse criteria to simplify the logits, then bounds relative flatness through a closed-form expression involving the Hessian trace. Under the stated assumptions, the resulting bound decays exponentially as the classifier-weight scale grows.

A Proof of Proposition 5.3

  • Under the neural-collapse criteria, the bias and class-mean projection onto the global mean sum to the same constant across classes, leaving logits governed by relative class-mean geometry.This follows from Lemma A.1 and the NC assumptions on collapsed features, equiangular class means, and aligned classifier weights.Source: A Proof of Proposition 5.3 · pp. 14–15
  • At a minimum, relative flatness is upper-bounded by the simplified quantity ∥w∥^2Tr(H(w)), which has a closed form in terms of logits.The proof uses positive semidefiniteness of the Hessian and sums over class pairs.Source: A Proof of Proposition 5.3 · pp. 14–15
  • Under limiting neural collapse and cross-entropy training, relative flatness κϕ(w) obeys the bound in Proposition A.4.The classifier is a softmax network over penultimate-layer representations.Source: A Proof of Proposition 5.3 · p. 15
  • For sufficiently large λ, the bound decays exponentially in λ.The proof obtains this asymptotic behavior as the denominator tends to 1.Source: A Proof of Proposition 5.3 · pp. 15–16
  • Under NC, logits for an example from class c combine the relative class-mean inner product with the class-independent offset.The resulting softmax probabilities are used to compute the input-wise Hessian trace.Source: A Proof of Proposition 5.3 · pp. 15–16
  • The proposition formalizes how NC can yield relative flatness, while the discussion cautions that NC itself does not guarantee flatness.Continued cross-entropy training can increase λ and lower relative flatness without changing the collapsed representation.Source: A Proof of Proposition 5.3 · p. 16

B Training Setups for NCC and Additional Experiments

The NCC experiments test how regularizer strength and data augmentation affect collapse and generalization, while the gradient analysis explains how relative-flatness regularization can destabilize training and how normalization and weight capping can stabilize it.

B Training Setups for NCC and Additional Experiments

  • At λ = 1×10−3, relative flatness stays low while NCC increases, yet good generalization remains achievable without strong NCC.Smaller λ values have little effect, while λ = 1×10−2 destabilizes training; under augmentation, λ = 10−3 no longer produces the same NCC increase.Source: B Training Setups for NCC and… · p. 17
  • With data augmentation, validation accuracy reaches 94%, versus 78% without augmentation.The augmentation setting also prevents the NCC metric from increasing at λ = 10−3 and 10−4.Source: B Training Setups for NCC and… · p. 17

C.1 Aanalysis of Relative Flatness Regularizer and Gradients

  • Under high uncertainty, the regularizer can amplify weight growth, while cross-entropy reduces uncertainty and can counter that effect.Large λ can let this dynamic dominate and destabilize training.Source: §C.1 · pp. 19–20
  • Normalizing penultimate features removes quadratic feature-norm amplification, and capping the classifier’s Frobenius norm bounds regularizer-induced gradient terms.Together, these measures are described as keeping regularized training stable and well-conditioned.Source: §C.1 · p. 20

C.2 CIFAR-10 & Resnet-18 Training Setups and Additional Experiments

CIFAR-10 experiments examine relative-flatness regularization across seeds, strengths, and weight caps. Regularization can hold models in sharper solutions with lower validation accuracy, while unplugging it allows recovery toward the baseline.

C.2 CIFAR-10 & Resnet-18 Training Setups and Additional Experiments

  • Different random seeds lead standard training toward different sharp minima, so averaging results across seeds is not considered meaningful.The regularizer’s effects are reported as consistent across seeds.Source: §C.2 · p. 21
  • Regularized ResNet-18 reaches nearly 100% training accuracy but about 60% validation accuracy, versus about 78% for standard training; unplugging restores validation accuracy toward baseline.The regularized model’s relative flatness stays around 2000, then decreases to match baseline after unplugging.Source: §C.2 · pp. 21–22, Figure 10 · p. 22
  • At λ = 10−3, validation accuracy stays slightly below the 78% baseline despite high relative flatness; λ = 0.5 and 0.1 crash training, while λ = 10−4 remains stable.Below 10−4, the regularizer’s influence essentially disappears.Source: §C.2 · p. 22, Figure 11 · p. 23
  • A weight cap of 20 keeps training stable with about 75% validation accuracy, while caps of 150 and 200 cause training collapse and relative-flatness divergence above 10^5.At a cap of 100, validation accuracy plateaus around 50% and relative flatness rises rapidly.Source: §C.2 · pp. 22–23, Figure 12 · p. 24

C.3 Imagenet-100 & ViT Training Setups and Experimental Results

The ImageNet-100 experiments apply relative-flatness regularization to ViT-tiny training and compare it with standard training. Standard training reaches higher validation accuracy, while unplugging the regularizer allows recovery toward the baseline.

C.3 Imagenet-100 & ViT Training Setups and Experimental Results

  • The ViT-tiny experiment disables dropout, trains for 300 epochs, and removes the relative-flatness regularizer after epoch 150.The setup uses λ = 10−2, weight capping at 150, representation normalization, and temperature τ = 2.Source: §C.3 · p. 24
  • The ImageNet-100 delayed-generalization experiment tracks training and validation losses alongside relative flatness.The loss and relative-flatness curves are presented in Figures 13 and 14, respectively.Source: Figure 13 · p. 25, Figure 14 · p. 25, §C.3 · p. 25
  • Standard training reaches nearly 100% training accuracy and about 42% validation accuracy, while regularized training reaches about 38% validation accuracy and retains relative flatness above 20,000.After the regularizer is removed at epoch 150, validation accuracy rises toward 42% and relative flatness becomes indistinguishable from baseline.Source: §C.3 · p. 25, Figure 14 · p. 25

C.4 SST-5 & Pretrained Language Models Training Setups and Experimental Results

The paper evaluates relative-flatness regularization in TinyBERT and DistilGPT2 fine-tuning on SST-5. In both models, regularization is associated with lower validation accuracy and persistently high relative flatness, while unplugging it brings validation accuracy toward standard fine-tuning results.

C.4 SST-5 & Pretrained Language Models Training Setups and Experimental Results

  • Both pretrained language models are fine-tuned for 300 epochs with dropout disabled and compared across three random seeds.Figures 16 and 18 compare standard fine-tuning with regularized training; Figures 17 and 19 show unplug experiments.Source: §C.4 · pp. 25–26, Figure 16 · p. 27, Figure 17 · p. 27
  • TinyBERT standard fine-tuning reaches about 44% validation accuracy, while regularized training reaches about 38% and retains relative flatness above 700.After unplugging at epoch 150, validation accuracy rises to around 44% and relative flatness matches baseline.Source: §C.4 · p. 27, Figure 16 · p. 27, Figure 17 · p. 27
  • DistilGPT2 standard fine-tuning reaches about 45% validation accuracy, while regularized training reaches about 37% and retains relative flatness above 600.After unplugging at epoch 150, validation accuracy rises to about 45% and relative flatness becomes indistinguishable from baseline.Source: §C.4 · pp. 28–29, Figure 18 · p. 28, Figure 19 · p. 28

D Representativeness

The representativeness analysis uses KDE on penultimate-layer representations and finds that representativeness improvements coincide with the initial rise in validation accuracy. The checklist reports reproducibility details and public code access.

D Representativeness

  • Representativeness is measured with kernel density estimation on penultimate-layer representations, using a Gaussian kernel with bandwidth 1.0 and training-sample weight 0.02.Computations are averaged over three random seeds.Source: D Representativeness · p. 29
  • Representativeness decreases to zero when validation accuracy reaches 100%, while its improvements coincide with validation accuracy’s initial rise.Figure 20 presents representativeness alongside training and validation accuracies.Source: D Representativeness · p. 29, Figure 20 · p. 29

NeurIPS Paper Checklist

Loading 2509.17738v3…