How to Tame Grokking: Representation Geometry as a Control Signal

This study investigates the mechanistic relationship between the geometric structure of hidden representations and the phenomenon of grokking, where neural networks transition from memorization to generalization after prolonged training. The author asks whether representation dimensionality is merely a symptom of this transition or a controllable variable that can be manipulated to influence learning dynamics.

To explore this, the paper introduces Geometric Dimensionality Regularization (GeomDR), a spectral regularizer that penalizes variance outside a target subspace of hidden representations. By modifying the covariance spectrum of activations, the author forces the network to adopt a specific effective dimensionality. The method was evaluated across modular arithmetic and permutation tasks using both multilayer perceptrons and transformers, with systematic sweeps over regularization strength, ramp duration, and the timing of the intervention.

Empirical analysis reveals that a sharp collapse in effective dimensionality consistently precedes the onset of grokking. By applying GeomDR, the author successfully accelerated the transition to generalization by up to 52 times compared to standard AdamW training. The results show that there is a critical temporal window for these interventions; applying the regularizer too early disrupts the formation of necessary task representations, while applying it too late limits its impact. The study concludes that representation geometry acts as a reliable control signal for the memorization-to-generalization transition.

This work shifts the perspective on grokking from an emergent, mysterious property of deep learning to a process that can be actively steered. By providing a practical tool to manipulate representation geometry, the paper offers a new framework for researchers to study the internal mechanics of neural networks and potentially optimize training schedules for complex tasks where delayed generalization is a bottleneck.