Speaker
Description
Robust generalization is a foundational requirement for trustworthy artificial intelligence (AI), underpinning the reliability, stability, and fairness of deployed systems. Two geometric phenomena are frequently correlated with generalization: neural collapse, where internal class representations converge to a maximally simple and symmetric structure, and flatness of the loss landscape, where the model settles into a wide, flat minimum, suggesting resilience to perturbations. Identifying which, if either, causally drives generalization is critical for engineering reliable AI systems.
To disentangle causation from correlation, we leverage grokking, a training regime where generalization is delayed, creating a temporal window for causal analysis. Our experimental interventions reveal a clear asymmetry: while both neural collapse and relative flatness emerge near the onset of generalization, only flatness consistently predicts it. Models regularized away from flat solutions exhibit delayed generalization, resembling grokking even in architectures and datasets where it does not typically occur. In contrast, promoting or suppressing neural collapse has no significant effect on a model’s ability to generalize. Furthermore, we show theoretically that neural collapse implies relative flatness under classical assumptions, explaining their empirical co-occurrence.
These findings reposition relative flatness as a potentially necessary and more fundamental condition for generalization. This insight has important implications for both the scientific understanding of generalization and practical methods for trustworthy AI: research and development should prioritize techniques that promote flat minima to improve reliability, robustness, and fairness. By focusing on the geometry of stability rather than structural symmetry, this work offers an actionable pathway for building AI systems that generalize better, and are, therefore, more trustworthy.
