Technical Summary: How to Tame Grokking: Representation Geometry as a Control Signal

Problem Statement

Grokking is a phenomenon in overparameterized neural networks where models initially memorize training data and only exhibit strong generalization after prolonged optimization. While recent studies have linked grokking to weight decay, feature compression, circuit formation, and representation learning dynamics, the specific factors governing the timing and emergence of this delayed generalization remain incompletely understood. A common theme in existing literature is that learned representations progressively compress into lower-dimensional structures. However, prior work largely treats representation geometry as an emergent property rather than a direct control mechanism. This paper investigates whether explicitly manipulating the geometry of hidden representations can serve as a control signal to influence the onset and speed of grokking.

Methodology

The authors introduce Geometric Dimensionality Regularization (GeomDR), a representation-level spectral regularizer designed to directly control the effective dimensionality of hidden representations during training.

Geometric Characterization

For a hidden layer representation matrix $Z \in \mathbb{R}^{N \times d}$, the method first centers the data to obtain $\tilde{Z}$. It then computes the empirical covariance matrix $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$. The effective dimensionality ($D_{eff}$) is calculated using the participation ratio of the eigenvalues ($\mu_i$) of $C$:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
Additionally, a local neighborhood distance ($\rho$) is defined as the average Euclidean distance to the $k$ nearest neighbors of normalized representations.

The GeomDR Objective

GeomDR imposes a penalty on the variance contained in directions beyond a target subspace of dimension $d^*$. Given eigenvalues sorted in descending order ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), the layer-wise loss is:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
The total training objective combines the task loss ($L_{task}$) and the geometric regularizer:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Crucially, the regularizer is not active from the start. It is introduced after an initial unconstrained training phase (activation step $t_s$) and its strength $\lambda(t)$ is gradually increased using a cosine ramp. This design mimics the natural grokking trajectory where memorization precedes generalization.

Experimental Protocol

The framework was evaluated on three algorithmic tasks known for exhibiting grokking:

1. Modular Addition: Predicting $(a + b) \mod 97$.
2. Modular Division: Predicting $(a \cdot b^{-1}) \mod 97$.
3. Permutation Composition: Predicting the composition of two permutations in $S_5$.

Experiments utilized two architectures: a Multilayer Perceptron (MLP) and a Transformer. The study systematically varied intervention schedules (strength $\lambda_{max}$ and ramp duration $T_{ramp}$), target dimensionalities ($d^*$), and intervention start times ($t_s$).

Key Contributions

1. Observation of Preceding Collapse: The authors demonstrate that dimensionality collapse consistently precedes the transition from memorization to generalization (grokking) in baseline settings. This suggests that geometric compression is not merely a correlate but a mechanistic driver of the transition.
2. GeomDR Introduction: They propose GeomDR, a spectral regularizer that directly suppresses the covariance spectrum outside a target subspace, thereby controlling the effective dimensionality of hidden representations.
3. Systematic Control of Dynamics: The paper provides a comprehensive empirical study showing that geometric interventions can substantially alter delayed generalization dynamics. Depending on the intervention schedule and target dimensionality, GeomDR can accelerate grokking, delay it, or suppress it entirely.

Results

Acceleration of Grokking

Across all evaluated settings, GeomDR consistently altered grokking dynamics:

• MLP Performance: On modular addition, GeomDR accelerated grokking by up to 52 times (reducing steps from 362K to ~7K). Similar improvements were observed for modular division (39x) and permutation composition (~23x). In all MLP cases, GeomDR achieved a 100% success rate where the baseline was unstable or failed.
• Transformer Performance: While the baseline grokking dynamics for Transformers were already faster than MLPs, GeomDR still provided consistent acceleration (e.g., ~2x on modular addition), maintaining a 100% success rate.

Influence of Hyperparameters

• Target Dimensionality ($d^*$): Moderate to large target dimensions ($16 \le d^* \le 64$) yielded the strongest acceleration. Aggressive compression ($d^* \le 8$) often slowed down or destabilized learning.
• Intervention Timing ($t_s$): A "critical temporal window" was identified. Interventions applied too early (before useful representations form) or too late (after memorization is rigid) were less effective. Intermediate activation times (e.g., $t_s = 2000$ for MLPs) produced the fastest and most stable grokking.
• Schedule: Intermediate regularization strengths and ramp durations generally performed best; overly strong interventions could destabilize training.

Geometric Correlations

The study confirmed that interventions accelerating grokking also accelerated the onset of dimensionality collapse. In GeomDR models, the temporal gap between dimensionality collapse and successful generalization was significantly reduced compared to baselines, suggesting a tighter coupling between geometric reorganization and the emergence of generalization.

Scalability and Efficiency

The authors demonstrated that GeomDR remains effective even when the covariance spectrum is estimated from a random subset of representations (e.g., 5%–25% of the data), offering a 2.4x–2.7x speedup in runtime without materially affecting grokking acceleration.

Significance and Claims

The paper claims that representation geometry serves as an effective control signal for grokking. By directly modifying the covariance spectrum of hidden activations, researchers can systematically influence when and how quickly a network transitions from memorization to generalization.

The authors posit that this work provides evidence that geometric interventions offer a practical approach for studying and influencing delayed generalization. Rather than passively observing representational changes, the method allows for the direct manipulation of geometric properties to test hypotheses about learning dynamics. The results suggest that the relationship between representation geometry and grokking is robust across different tasks and architectures, though the magnitude of improvement varies (being more pronounced in MLPs than Transformers, likely due to the latter's stronger inductive biases).

The paper concludes that while it does not establish representation dimensionality as the sole causal mechanism, it demonstrates that geometric interventions are sufficient to alter the timing of grokking, offering a new experimental methodology for optimization and generalization research.