🤖 AI Summary
Neural networks often exhibit the "grokking" phenomenon during training, characterized by an initial memorization phase followed by sudden generalization, yet its underlying mechanisms and control strategies remain poorly understood. This work reveals that dimensional collapse consistently precedes grokking and, for the first time, leverages the geometric structure of representations as a control signal. We propose Geometric Dimensionality Regularization (GeomDR), a spectral regularization method that modulates the effective representational dimensionality of hidden layers in both multilayer perceptrons and Transformers. Evaluated on modular arithmetic (addition and division) and permutation composition tasks, GeomDR accelerates the onset of generalization by up to 52×, providing strong empirical evidence that geometric intervention effectively regulates grokking dynamics.
📝 Abstract
Grokking is a phenomenon in which neural networks initially memorize training data and only later exhibit strong generalization after prolonged optimization. Despite extensive recent study, the factors influencing the emergence and timing of grokking remain incompletely understood. We investigate the relationship between representation geometry and delayed generalization. We find that dimensionality collapse consistently precedes the onset of grokking in all evaluated settings. Motivated by these observations, we introduce Geometric Dimensionality Regularization (GeomDR), a simple spectral regularizer that modifies the effective dimensionality of hidden representations during training. Across modular addition, modular division, and permutation composition tasks, GeomDR consistently alters grokking dynamics and can substantially accelerate the onset of generalization depending on the intervention schedule and target dimensionality. In several settings, grokking is accelerated by up to 52 times relative to standard AdamW training. Similar qualitative effects are observed in both multilayer perceptrons and transformers. Together, these results suggest that representation geometry can serve as an effective control signal for grokking and provide evidence that geometric interventions offer a practical approach for studying and influencing delayed generalization in neural networks.