arXiv CSopen access

How to Tame Grokking: Representation Geometry as a Control Signal

How to Tame Grokking: Representation Geometry as a Control Signal Maksim A. Kazanskii Independent Researcher [email protected]

arXiv:2607.11666v1 [cs.LG] 13 Jul 2026

Abstract

norm control, sparse subnetworks, or embedding transfer [14, 16, 11, 26]. However, comparatively less attention has been paid to controlling grokking through direct interventions on the geometry of hidden representations. If delayed generalization is closely connected to representation geometry, then explicitly modifying this geometry may influence the onset and speed of grokking. We investigate this hypothesis by introducing Geometric Dimensionality Regularization (GeomDR), a representationlevel spectral regularizer that suppresses variance outside a target subspace and thereby controls the effective dimensionality of hidden representations during training. GeomDR directly modifies the covariance spectrum of hidden activations. We perform a systematic study across grokking tasks, architectures, intervention schedules, target dimensionalities, and random seeds. Our results show that geometric interventions can substantially alter delayed generalization dynamics, accelerating grokking by up to 52 times in some settings and, under stronger interventions, delaying or suppressing generalization. We further find that changes in effective dimensionality consistently precede the transition from memorization to generalization, suggesting that representation dimensionality is not merely a diagnostic statistic but a controllable variable associated with grokking dynamics. Our contributions are threefold:

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.

Introduction Grokking is a delayed-generalization phenomenon in which neural networks first memorize training data and only much later achieve strong test performance after prolonged optimization [19]. This behavior differs from conventional learning dynamics, where training and test performance typically improve together, and has become a useful setting for studying generalization in overparameterized neural networks [19, 13, 17, 24]. Prior work has linked grokking to weight decay, feature compression, circuit formation, and representation learning dynamics [19, 17, 13]. A common theme in these explanations is that learned representations often become progressively compressed into lower-dimensional structures. More broadly, representation geometry and dimensionality have been shown to play important roles in optimization and generalization [2, 12]. Recent work has shown that grokking can be accelerated or altered through optimization dynamics, weight-

• We show that direct geometric interventions can substantially alter delayed generalization, enabling acceleration and, in some settings, delay or suppression. • We introduce GeomDR, a representation-level spectral regularizer for directly controlling the effective dimensionality of hidden representations. • We provide a systematic empirical study of how intervention strength, timing, target dimensionality, task structure, and architecture affect grokking dynamics.

Related Work Early work on grokking identified delayed generalization in algorithmic tasks and highlighted the importance of regularization, particularly weight decay, in the transition from memorization to generalization [19]. Subsequent studies 1

connected grokking to feature compression, circuit formation, and representation learning dynamics [13, 17, 24]. These works largely treat representation geometry as an emergent property of learning, whereas our goal is to investigate whether geometry can be directly manipulated to control grokking dynamics. Several approaches have sought to accelerate or modify grokking through changes to the training process. Omnigrok studies the role of weight norms and conventional regularization [14], Grokking Tickets relates delayed generalization to sparse subnetworks and pruning [16], Grokfast accelerates grokking through gradient filtering [11], and GrokTransfer explores embedding transfer from weaker models [26]. In contrast, GeomDR acts directly on hidden representations by modifying their covariance spectrum rather than altering optimization dynamics, sparsity, or transferred embeddings. Representation geometry provides a useful framework for understanding learning dynamics. Prior work has shown that independently trained networks often converge to similar representational structures [9], while geometric analyses have been used to study feature organization and learning across architectures [20]. Neural representations frequently exhibit low intrinsic dimensionality [2, 12, 1] and may undergo phases of compression and expansion during learning [21]. Related evidence from Neural Collapse, Information Bottleneck analyses, and self-supervised learning further suggests that covariance structure and representation dimensionality are closely linked to learning and generalization [18, 6, 10, 23, 22, 3, 7]. Existing grokking studies either analyze representation geometry as a correlate of delayed generalization or influence grokking indirectly through regularization, optimization, sparsity, or transfer mechanisms [19, 13, 17, 11, 26]. GeomDR instead treats representation geometry as the object of intervention itself, enabling controlled experiments on how effective dimensionality influences the onset and speed of grokking.

Ze = Z − Z̄, where Z̄ denotes the feature-wise mean of Z. We then compute the empirical covariance matrix

1 e⊤ e Z Z. (4) N −1 Let µ1 , . . . , µd denote the eigenvalues of C. Building on prior work on intrinsic dimensionality and representation geometry [2], we compute effective dimensionality using the participation ratio C=

P 2 ( µi ) Deff = Pi 2 . i µi

ẑi =

zi . ∥zi ∥2

(6)

Then we define local neighborhood distance as the average Euclidean distance to the k nearest neighbors of each normalized representation, N

ρ=

1 X1 N i=1 k

X

∥ẑi − ẑj ∥2 ,

(7)

j∈Nk (i)

where Nk (i) denotes the set of k nearest neighbors of representation i. Unless otherwise specified, we use k = 10 throughout all experiments.

Geometric Dimensionality Regularization We hypothesize that the reduction of representation dimensionality is not merely a consequence of grokking but a driving factor in the transition from memorization to generalization. If representation dimensionality is mechanistically linked to grokking, then directly controlling the geometry of hidden representations should alter the timing and dynamics of the memorization-to-generalization transition. Motivated by this idea, we introduce Geometric Dimensionality Regularization (GeomDR). For each hidden layer ℓ, we compute the empirical covariance matrix

We introduce a geometry-based regularization framework that directly modifies the effective dimensionality of hidden representations during training. We consider supervised algorithmic learning tasks with inputs x and labels y. A neural network fθ is trained to predict

C (ℓ) =

(1)

For each hidden layer ℓ, the network produces a representation matrix

1  e(ℓ) ⊤ e(ℓ) Z Z . N −1

(8)

Let (ℓ)

Z (ℓ) ∈ RN ×d ,

(5)

Lower values of Deff correspond to representations whose variance is concentrated in a smaller number of directions. Let each representation vector be normalized as

Methods

ŷ = fθ (x).

(3)

(ℓ)

(ℓ)

µ1 ≥ µ2 ≥ · · · ≥ µd

(2)

(9)

(ℓ)

denote the eigenvalues of C sorted in descending order. Given a target dimensionality d∗ , we define the layer-wise geometric regularization loss as

where N is the number of examples and d is the representation dimension.

Geometric Characterization of Representations (ℓ)

For each representation matrix Z, we first center the representations:

LGeomDR =

d X j=d∗ +1

2

(ℓ)

µj .

(10)

This objective admits several geometric interpretations related to effective-rank control, low rank covariance approximation, and representation volume compression (see Appendix A). Unless otherwise specified, GeomDR is applied to all hidden layers and excluded from the input embedding and output layers. This objective penalizes variance contained in directions beyond the leading d∗ principal components, encouraging representations to concentrate within a lowerdimensional subspace. The objective is therefore closely related to classical low rank approximation and principal component analysis [4, 8]. The total geometric regularization loss is X (ℓ) LGeomDR = LGeomDR , (11)

otherwise specified, all reported results are averaged across five independent random seeds. A run is considered to have successfully grokked when the training accuracy is 1.0 and the test accuracy satisfies Acctest ≥ 0.999 Acctrain for 10 consecutive evaluations. The evaluation is performed every 10 optimization steps. Therefore this criterion corresponds to maintaining the target accuracy for 100 consecutive optimization steps. The grokking step is defined as the first optimization step at which this condition is satisfied. For all tasks and architectures, geometric regularization is not applied at the beginning of training. Instead, the regularizer is activated after an initial unconstrained optimization phase. Unless otherwise specified, MLP experiments use an intervention activation step of ts = 2000. Transformer experiments use ts = 1000. For each task–architecture pair, we use the unregularized model as the baseline reference. Our main experimental procedure consists of three stages:

ℓ

The full training objective is L = Ltask + λ(t)LGeomDR , (12) where Ltask denotes the task-specific training loss. The geometric regularizer is activated after an initial training phase, allowing the network to fit the training data before constraining representation geometry. This design is motivated by grokking, where generalization typically emerges only after memorization [19]. To avoid an abrupt change in the optimization objective, the regularization strength is introduced gradually using a cosine ramp. Smooth scheduling of optimization hyperparameters has become a common strategy in deep learning to improve training stability and optimization behavior [15, 5]. 1 − cos(πq(t)) λ(t) = λmax , 2

1. Schedule sweep (MLP). For MLPs, we sweep both the final regularization strength λmax and ramp duration Tramp . For Transformers, we sweep only λmax . Each configuration is evaluated across five random seeds and compared to the baseline. 2. Dimensionality sweep. We perform a sweep over the target dimensionality d∗ . 3. Intervention-start sweep. We vary the activation step ts at which geometric regularization is introduced. Complete architectural specifications, hyperparameter grids, and implementation details are provided in Appendix B.

(13)

where

Results  0, t < ts ,   t − t s , ts ≤ t < ts + Tramp , q(t) =  Tramp   1, t ≥ ts + Tramp .

Observation (14)

We first examine the evolution of representation geometry during baseline grokking on modular addition. This baseline setting provides a reference for subsequent intervention experiments. To characterize representation geometry, we monitor effective dimensionality Deff and the local neighborhood distance ρ, which measures the average distance between neighboring representations. For visualization, trajectories are smoothed with an exponential moving average (EMA) using a window of 100 steps, while the final 20,000 optimization steps are additionally shown without smoothing. Figure 1 shows the evolution of effective dimensionality and local neighborhood distance across all layers. Shallow layers remain relatively stable throughout training, whereas deeper layers exhibit pronounced geometric restructuring. Effective dimensionality stays high for most of training before undergoing a sharp collapse near the grokking transition, with the strongest effect observed in the deepest hidden layer. Local neighborhood distance displays a similar pattern, remaining stable for extended periods before reorganizing rapidly near the onset of generalization. These coordinated changes occur within a narrow training interval across

Here ts is the activation time step, Tramp is the ramp duration, and λmax is the final regularization strength. For brevity, we refer to λmax simply as λ when reporting experimental configurations. GeomDR requires covariance estimation and eigendecomposition of hidden representations, resulting in a per-layer complexity of O(N d2 + d3 ) (see Appendix D).

Experimental Protocol We evaluate the proposed framework across three algorithmic learning tasks: modular addition, modular division, and permutation composition. These tasks are commonly used in the grokking literature because they exhibit delayed generalization under appropriate training conditions [19]. Experiments are conducted using two neural architectures: a multilayer perceptron (MLP) and a Transformer [25]. For all tasks, a fixed 30%/70% train–test split is used. Following prior grokking work [19], no separate validation set is employed, and test-set dynamics are reported directly. Unless 3

(a) All steps: Deff

(b) All steps: ρ

(c) Train accuracy

(d) Final 20k steps: Deff

(e) Final 20k steps: ρ

(f) Test accuracy

Figure 1: Evolution of accuracy and representation geometry during baseline grokking on modular addition. The model is trained without geometric regularization. Panels (a), (b), (c), and (f) show EMA-smoothed trajectories of effective dimensionality, local neighborhood distance, train accuracy, and test accuracy. Panels (d) and (e) show the raw geometric trajectories for dimensionality, local neighborhood distance during the final 20K optimization steps. Delayed generalization coincides with rapid geometric reorganization, particularly in deeper layers. The deepest hidden layer (L4 ) exhibits the largest changes in both effective dimensionality and local neighborhood distance near the grokking transition.

Intervention

multiple layers, suggesting that grokking is accompanied by a rapid reorganization of representation geometry.

To investigate the role of representation dimensionality, we introduce a regularizer that enforces a target dimensionality d∗ . The intervention is activated at training step ts = 2000 and cosine-ramped to λ = 1 over 333 optimization steps. We evaluate four target dimensionalities, d∗ ∈ {8, 16, 32, 64}, while keeping all other hyperparameters fixed. Figure 2 shows the evolution of effective dimensionality and local neighborhood distance in the shallowest (L0 ) and deepest (L4 ) hidden layers, together with train and test accuracy. Activation of the regularizer produces an immediate geometric transition across all measured quantities. The strongest effect is observed in the local neighborhood distance ρ. Prior to the intervention, all runs follow nearly identical trajectories. Once regularization is activated, ρ rapidly decreases in both shallow and deep layers, indicating substantial reorganization of the representation space. The trajectories subsequently separate according to the target dimensionality, demonstrating direct control over the resulting geometric state. Effective dimensionality exhibits a similar response. Immediately after the intervention, Deff decreases sharply, consistent with the suppression of variance outside the target subspace. This is followed by a recovery phase, after which larger target dimensionalities maintain higher effective dimensionality. The effect is most pronounced in the deepest layer, where the separation between settings remains visible throughout training.

Across all runs analyzed and tasks considered in this study, dimensionality collapse consistently preceded grokking. In GeomDR-induced models, collapse occurred substantially earlier, reaching 50% of its initial effective dimensionality after approximately 2.3K optimization steps, and was followed shortly thereafter by grokking (approximately 6.8K steps). These results indicate that interventions that accelerate grokking also accelerate dimensionality collapse, supporting the hypothesis that geometric compression is closely linked to the emergence of generalization. The unsmoothed trajectories reveal that the transition is not continuous. Instead, grokking is associated with a small number of abrupt geometric reorganizations. These events occur simultaneously in both effective dimensionality and the local neighborhood distance ρ, suggesting that delayed generalization is accompanied by a large-scale restructuring of internal representations rather than gradual optimization alone. This observation motivates the hypothesis that representation geometry may play a mechanistic role in grokking dynamics. If so, then directly controlling representation geometry may alter the timing of the memorization-to-generalization transition, motivating the intervention studies presented in the following sections. 4

(a) L4 : Deff

(b) L4 : ρ

(c) Train accuracy

(d) L0 : Deff

(e) L0 : ρ

(f) Test accuracy

Figure 2: Evolution of accuracy and representation geometry under GeomDR for modular addition. The model is trained with different target dimensionalities d∗ . Panels (a), (b), (c), and (f) show EMA-smoothed trajectories of effective dimensionality, local neighborhood distance, train accuracy, and test accuracy. Panels (d) and (e) show the corresponding geometric trajectories in the shallowest hidden layer. The strongest effects are observed in the deepest hidden layer (L4 ), where dimensionality collapse occurs earliest and is followed by rapid increases in test accuracy. Despite these substantial geometric changes, train and test accuracy remain broadly similar across settings. Together, these results demonstrate that dimensionality regularization provides a direct mechanism for controlling representation geometry during training.

observed for modular division and permutation composition (Appendix C).

Ablation Studies: Intervention Timing To investigate the relationship between dimensionality collapse and grokking, we measured the collapse time of the deepest hidden representation as the first training step at which its effective dimensionality fell below 50% of its initial value. We then compared this quantity to the grokking time, defined as the first step at which the grokking criterion above is satisfied. Figure 5 shows that dimensionality collapse consistently precedes grokking in both baseline and GeomDR models. In the baseline setting, collapse occurs relatively early in training (approximately 15–25K steps), whereas grokking emerges much later (approximately 250– 500K steps), producing a large temporal gap between geometric compression and successful generalization. GeomDR shifts both events to substantially earlier stages of training, with collapse occurring after only 3–5K steps and grokking following shortly thereafter at approximately 5–10K steps. Consequently, the intervention not only accelerates dimensionality collapse but also markedly reduces the lag between collapse and generalization. Across all target dimensionalities and random seeds, GeomDR moves the system closer to the diagonal tcollapse = tgrok , indicating a substantially tighter coupling between geometric reorganization and the onset of generalization. The effectiveness of geometric regularization depends strongly on intervention timing. Figure 6 shows the relation-

Control: Regularization Schedule and Target Dimensionality To identify an effective intervention schedule, we fixed the target dimensionality to d∗ = 16 and performed a grid search over the final regularization strength λmax and ramp duration Tramp . Figure 3 summarizes the results. Relative to the baseline ( 362.1 ± 111.0K steps), GeomDR substantially accelerates grokking across a broad range of schedules, with the best configuration reaching successful generalization after approximately 7K steps. Intermediate regularization strengths and ramp durations perform best, whereas very weak interventions have little effect and overly strong interventions can destabilize training or prevent grokking. Based on these results, we select λmax = 1 and Tramp = 333 for subsequent experiments. Using this schedule, we vary the target dimensionality over d∗ ∈ {2, 4, 8, 16, 32, 48, 64} and evaluate each configuration across ten random seeds. As shown in Figure 4, aggressive compression (d∗ ≤ 8) produces slower and more variable grokking, whereas dimensions in the range 16 ≤ d∗ ≤ 64 consistently yield strong acceleration. The best mean performance is obtained at d∗ = 64, reaching the grokking criterion after approximately 7K steps. Similar trends are 5

Figure 4: Effect of target dimensionality d∗ on grokking. Points show individual runs, the dashed curve the mean, and error bars one standard deviation. The orange dashed line denotes the baseline. Moderate and large target dimensions (16 ≤ d∗ ≤ 64) substantially accelerate grokking, whereas aggressive compression (d∗ ≤ 8) slows and destabilizes learning.

Figure 3: Schedule sweep for modular addition (d∗ = 16). Cells report the mean grokking step (K) ± one standard deviation; fractions indicate successful runs. Grey cells denote failed configurations (at least one failed run). Intermediate λmax and Tramp produce the fastest grokking, with λmax = 1 and Tramp = 333 performing best. ship between the intervention-start step ts and the resulting grokking step. A clear U-shaped trend is observed. Very early interventions (ts = 500 and ts = 1000) exhibit substantially slower and more variable convergence across random seeds. The fastest and most stable grokking behavior occurs for intermediate activation times (ts = 2000–4000), where the mean grokking step is minimized. Delaying the intervention beyond this region progressively increases the number of optimization steps required for successful generalization. These results suggest the existence of a critical temporal window during which geometric regularization is most effective. Applying the intervention too early appears to disrupt the formation of useful task representations, whereas applying it too late reduces its ability to influence the memorization-to-generalization transition. Overall, the results support the hypothesis that geometric interventions are most beneficial after an initial unconstrained learning phase but before memorizing solutions become firmly established. Additional intervention-start sweeps and implementation details are provided in Appendix C. Intervention timing substantially influences performance, although the optimal activation step varies across tasks and architectures.

provements across all tasks, reducing the mean grokking time from 362.1K to 6.9K steps on modular addition (52.5×), from 330.9K to 8.4K steps on modular division (39.4×), and from 300.2K to 12.8K steps on permutation learning (23.5×). In all cases, GeomDR achieves a 100% success rate. For Transformers, the effect remains consistent but more moderate. GeomDR reduces grokking time from 66.4K to 33.0K steps on modular addition (2.01×), from 89.4K to 48.7K steps on modular division (1.84×), and from 72.6K to 48.9K steps on permutation learning (1.48×), while maintaining a 100% success rate across all tasks. Additional ablation studies, hyperparameter sweeps, and architecturespecific analyses are provided in Appendix C. Together, these results indicate that GeomDR accelerates grokking across diverse tasks and architectures, although the magnitude of improvement varies across model classes.

Discussion Our results suggest that representation geometry can serve as an effective control signal for grokking. Across a range of tasks, schedules, dimensionalities, and architectures, GeomDR consistently influences delayed generalization. In many MLP settings, the intervention accelerates grokking by more than an order of magnitude, while Transformer experiments show smaller but qualitatively similar improvements. The smaller gains observed in Transformers may reflect their substantially faster baseline grokking dynamics, leaving less room for acceleration than in MLPs. The geometric analyses reveal systematic changes in effective dimensionality and local distance near the grokking transition. Across all successful runs, dimensionality col-

Generalization Across Tasks and Architectures We report grokking times, success rates, and relative improvements for all evaluated settings. The relationship between representation geometry and grokking remains consistent across all evaluated settings. Table 1 summarizes the best-performing GeomDR configuration for each task and architecture. For MLPs, GeomDR produces large and consistent im6

Table 1: Summary of the best GeomDR configuration for each task and architecture. Values denote mean grokking step ± standard deviation (K optimization steps). S/T indicates runs satisfying the grokking criterion. Unsuccessful runs were assigned the maximum training budget (500K for MLPs, 200K for Transformers) when computing statistics. Addition Architecture

Step

Division

S/T Imp.

Step

S/T Imp.

Permutation Step

S/T Imp.

MLP MLP + GeomDR

362.1 ± 111.0 4/5 – 330.9 ± 100.9 5/5 – 300.2 ± 207.0 3/5 – 6.9 ± 1.6 5/5 52.5× 8.4 ± 5.0 5/5 39.4× 12.8 ± 1.7 5/5 23.5×

Transformer Transformer + GeomDR

66.4 ± 12.0 33.0 ± 2.5

5/5 – 5/5 2.01×

89.4 ± 27.3 48.7 ± 8.8

5/5 – 5/5 1.84×

72.6 ± 33.0 48.9 ± 9.1

5/5 – 5/5 1.48×

Figure 5: Relationship between representation collapse time and grokking time. Each point corresponds to one run. The dashed diagonal denotes tcollapse = tgrok .

lapse consistently preceded grokking, and interventions that accelerated grokking also accelerated the onset of collapse. These observations suggest that delayed generalization is accompanied by a substantial reorganization of internal representations. Previous work on grokking has largely treated representation geometry and compression as emergent consequences of learning [19, 13, 24, 17]. In contrast, our results show that directly modifying representation geometry is sufficient to substantially alter delayed generalization dynamics. While this does not establish representation dimensionality as the sole causal mechanism underlying grokking, it demonstrates that geometric interventions are capable of systematically altering when grokking occurs. Geometric interventions provide a new experimental methodology for studying learning dynamics and suggest that representation geometry may serve as a controllable degree of freedom for studying optimization and generalization. Rather than passively observing representational changes during training, researchers can directly manipulate geometric properties and measure the resulting effects on optimization and generalization. Importantly, the observed acceleration is not restricted to a single task or finely tuned hyperparameter configuration. Effective interventions are obtained across a broad range of target dimensionalities and activation schedules, indicating that the relationship between representation geometry and grokking is robust rather than task-specific. Several limitations remain. First, our experiments focus on small-scale algorithmic grokking benchmarks, and it remains

Figure 6: Intervention-start sweep for modular addition. Points show individual runs; the dashed curve and shaded region denote the mean and one standard deviation. Intermediate activation times yield the fastest grokking. unclear whether similar geometric interventions are effective in larger-scale domains such as language or vision. Second, while GeomDR provides a direct geometric intervention, the precise mechanisms linking representation geometry and delayed generalization remain an open theoretical question. Appendix A provides a geometric interpretation of the regularizer in terms of effective-rank control, low rank approximation, and representation volume compression.

Conclusion Across all tasks and architectures studied, dimensionality collapse consistently preceded generalization. Motivated by this observation, we introduced Geometric Dimensionality Regularization (GeomDR), a spectral regularizer that directly controls the effective dimensionality of hidden representations. Across diverse settings, GeomDR substantially alters grokking dynamics, accelerating, delaying, or suppressing generalization. These results suggest that representation geometry is not merely a correlate of grokking but a useful target for intervention.

7

References

[13] Liu, Z.; Kitouni, O.; Nolte, N.; Michaud, E. J.; Tegmark, M.; and Williams, M. 2022. Towards Understanding Grokking: An Effective Theory of Representation Learning. arXiv, abs/2205.10343.

[1] Aghajanyan, A.; Gupta, S.; and Zettlemoyer, L. 2021. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), 7319–7328. Association for Computational Linguistics.

[14] Liu, Z.; Michaud, E. J.; and Tegmark, M. 2022. OmniGrok: Grokking Beyond Algorithmic Data. arXiv, abs/2210.01117. [15] Loshchilov, I.; and Hutter, F. 2017. SGDR: Stochastic Gradient Descent with Warm Restarts. In International Conference on Learning Representations.

[2] Ansuini, A.; Laio, A.; Macke, J. H.; and Zoccolan, D. 2019. Intrinsic Dimension of Data Representations in Deep Neural Networks. In Advances in Neural Information Processing Systems, volume 32, 6114–6125.

[16] Minegishi, G.; Iwasawa, Y.; and Matsuo, Y. 2023. Bridging Lottery Ticket and Grokking: Understanding Grokking from Inner Structure of Networks. arXiv, abs/2310.19470.

[3] Bardes, A.; Ponce, J.; and LeCun, Y. 2022. VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning. In International Conference on Learning Representations.

[17] Nanda, N.; Chan, L.; Lieberum, T.; Smith, J.; and Steinhardt, J. 2023. Progress Measures for Grokking via Mechanistic Interpretability. arXiv, abs/2301.05217.

[4] Eckart, C.; and Young, G. 1936. The Approximation of One Matrix by Another of Lower Rank. Psychometrika, 1(3): 211–218.

[18] Papyan, V.; Han, X.; and Donoho, D. L. 2020. Prevalence of Neural Collapse During the Terminal Phase of Deep Learning Training. Proceedings of the National Academy of Sciences, 117(40): 24652–24663.

[5] Goyal, P.; Dollár, P.; Girshick, R.; Noordhuis, P.; Wesolowski, L.; Kyrola, A.; Tulloch, A.; Jia, Y.; and He, K. 2017. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv, abs/1706.02677.

[19] Power, A.; Burda, Y.; Edwards, H.; Babuschkin, I.; and Misra, V. 2022. Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv, abs/2201.02177.

[6] Han, X.; Papyan, V.; and Donoho, D. L. 2022. Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path. In International Conference on Learning Representations.

[20] Raghu, M.; Unterthiner, T.; Kornblith, S.; Zhang, C.; and Dosovitskiy, A. 2021. Do Vision Transformers See Like Convolutional Neural Networks? In Advances in Neural Information Processing Systems, volume 34, 12116–12128.

[7] Jing, L.; Vincent, P.; LeCun, Y.; and Tian, Y. 2022. Understanding Dimensional Collapse in Contrastive SelfSupervised Learning. In International Conference on Learning Representations.

[21] Recanatesi, S.; Farrell, M.; Advani, M.; Moore, T.; Lajoie, G.; and Shea-Brown, E. 2019. Dimensionality Compression and Expansion in Deep Neural Networks. arXiv, abs/1906.00443.

[8] Jolliffe, I. T. 2002. Principal Component Analysis. Springer, 2nd edition.

[22] Shwartz-Ziv, R.; and Tishby, N. 2017. Opening the Black Box of Deep Neural Networks via Information. arXiv, abs/1703.00810.

[9] Kornblith, S.; Norouzi, M.; Lee, H.; and Hinton, G. 2019. Similarity of Neural Network Representations Revisited. In Proceedings of the 36th International Conference on Machine Learning, 3519–3529.

[23] Tishby, N.; and Zaslavsky, N. 2015. Deep Learning and the Information Bottleneck Principle. In 2015 IEEE Information Theory Workshop (ITW), 1–5. Jerusalem, Israel.

[10] Kothapalli, V. 2022. Neural Collapse: A Review on Modelling Principles and Generalization. arXiv, abs/2206.04041.

[24] Varma, V.; Shah, R.; Kenton, Z.; Kramár, J.; and Kumar, R. 2023. Explaining Grokking through Circuit Efficiency. arXiv, abs/2309.02390.

[11] Lee, J.; Kang, B. G.; Kim, K.; and Lee, K. M. 2024. GrokFast: Accelerated Grokking by Amplifying Slow Gradients. arXiv, abs/2405.20233.

[25] Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017. Attention Is All You Need. In Advances in Neural Information Processing Systems, volume 30.

[12] Li, C.; Farkhoor, H.; Liu, R.; and Yosinski, J. 2018. Measuring the Intrinsic Dimension of Objective Landscapes. In International Conference on Learning Representations.

[26] Xu, Z.; Ni, Z.; Wang, Y.; and Hu, W. 2025. Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model. arXiv, abs/2504.13292.

8

Contents Introduction

1

Related Work

1

Methods Geometric Characterization of Representations . . . . . . . . . . . . . . . . . . . . . . . . . . . Geometric Dimensionality Regularization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Experimental Protocol . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

2 2 2 3

Results Observation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Intervention . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Control: Regularization Schedule and Target Dimensionality . . . . . . . . . . . . . . . . . . . Ablation Studies: Intervention Timing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Generalization Across Tasks and Architectures . . . . . . . . . . . . . . . . . . . . . . . . . . .

3 3 4 5 5 6

Discussion

6

Conclusion

7

Appendix

10

A Spectral Properties of GeomDR

10

B Tasks and Architectures Modular Addition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Modular Division . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Permutation Composition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . MLP Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Transformer Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

12 12 12 12 12 12

C Ablations MLP . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Transformer . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Width Ablation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Collapse–Grokking Lag Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

14 14 17 18 18

D Scalability

20

E Comparison with Grokking Acceleration Methods

21

F Use of Large Language Models

22

G Code Availability

22

9

Appendix Spectral Properties of GeomDR

A

Proof.

∗

d X

The results in this appendix characterize the optimization objective induced by Geometric Dimensionality Regularization (GeomDR) and establish several geometric properties directly encouraged by the regularizer, including effectiverank control, low rank covariance approximation, and representation volume compression. Importantly, these results do not constitute a formal proof of why grokking accelerates under GeomDR. Rather, they describe the geometric structures that GeomDR explicitly promotes during optimization. Our working hypothesis is that accelerated grokking emerges because GeomDR reduces the space of high-dimensional memorizing solutions and biases optimization toward more compact representations, a mechanism that is consistent with the empirical observations presented in the main paper. Let Z ∈ RN ×d denote a matrix of hidden representations and let 1 (Z − Z̄)⊤ (Z − Z̄) N −1 denote the corresponding covariance matrix. Let C=

µ1 ≥ µ2 ≥ · · · ≥ µd ≥ 0

LGeomDR =

µi

rϵ (C) ≤ d∗ .

(24)

□ This proposition formalizes the fact that GeomDR directly limits the number of statistically significant covariance directions. Proposition 2 (Low Rank Approximation Bound). Let Cd∗ denote the optimal rank-d∗ approximation of C, and let ∥ · ∥F denote the Frobenius norm. Then ∥C − Cd∗ ∥F ≤ LGeomDR

(15)

Proof.

(25)

By the Eckart–Young theorem [4], ∥C − C

d∗

(16)

d X

∥2F =

µ2i

(26)

i=d∗ +1

Since all eigenvalues are nonnegative, d X

µ2i ≤

i=d∗ +1

(17)

!2

d X

= L2GeomDR .

µi

(27)

i=d∗ +1

Taking square roots yields (25). □ Thus minimizing GeomDR minimizes an upper bound on the reconstruction error of the best rank-d∗ covariance approximation.

∗

where d is the target dimensionality. Using d X

(23)

Hence the leading d∗ eigenvalues explain at least a fraction 1 − ϵ of the total variance. By definition of rϵ (C),

i=d∗ +1

Tr(C) =

µi ≥ (1 − ϵ) Tr(C)

i=1

be the eigenvalues of C. The Geometric Dimensionality Regularization (GeomDR) objective is d X

Since (19) holds, the assumption (21) implies

µi

(18)

Proposition 3 (Representation Volume Compression). Assume that the nonzero covariance spectrum of a representation matrix is

i=1

the objective can be rewritten as ∗

LGeomDR = Tr(C) −

d X

µi

µ1 ≥ µ2 ≥ · · · ≥ µr > 0,

(19)

i=1

(28)

where r = rank(C). The covariance ellipsoid associated with C has volume proportional to the square root of the determinant of C [8],

This formulation admits several useful interpretations. Proposition 1 (Effective Rank Control). For ϵ ∈ (0, 1), define the effective rank ( ) Pk i=1 µi ≥1−ϵ (20) rϵ (C) = min k : Pd i=1 µi

Vol(C) ∝

r Y p √ det(C) = µi .

(29)

i=1

LGeomDR =

If

d X

µi ≤ ϵ,

(30)

i=d∗ +1

LGeomDR ≤ ϵ Tr(C)

then every tail eigenvalue satisfies

(21)

then

µi ≤ ϵ, rϵ (C) ≤ d

∗

(22)

Consequently, 10

i > d∗ .

(31)

∗

Vol(C) =

d Y √

!

r Y √

µi

∗

d Y √

!

r Y √

µi

! ϵ

(32)

i=d∗ +1

i=1 ∗

≤

µi

i=d∗ +1

i=1

≤

LGeomDR decreases, the volume contributed by dimensions outside the leading d∗ -dimensional subspace vanishes. Consequently, optimization is no longer free to spread information across a large number of weak directions. Instead, representations are encouraged to concentrate variance into a smaller set of dominant components, reducing the geometric volume available to spectrally diffuse solutions. This may substantially restrict the family of memorizing representations accessible during training and bias optimization toward more compact solutions. Under this hypothesis, GeomDR does not directly create algorithmic representations. Rather, it modifies the geometry of the optimization landscape by making high-dimensional, spectrally diffuse solutions increasingly expensive. As training proceeds, optimization is therefore encouraged to move toward compressed representations that may be more likely to capture task structure. This shift can reduce the time required to transition from memorization to generalization, leading to earlier grokking. The broad dimensionality plateaus observed in our experiments are consistent with this interpretation. Acceleration occurs across a wide range of target dimensionalities rather than only at a single finely tuned value, suggesting that the critical factor is the suppression of excess representational degrees of freedom rather than the precise dimensionality itself. While this explanation remains a hypothesis rather than a formal proof, it provides a geometric interpretation that is consistent with the observed reductions in effective dimensionality, increased spectral concentration, covariancevolume compression, and substantially earlier grokking transitions across multiple tasks and architectures.

!

d Y √

! µi

∗

ϵ(r−d )/2 .

i=1

Thus, as LGeomDR → 0, the representation volume outside the leading d∗ -dimensional subspace vanishes. Proof.

Since d X

µi ≤ ϵ

(33)

i=d∗ +1

and all eigenvalues are nonnegative, i > d∗ .

µi ≤ ϵ,

(34)

Therefore, r Y √

r Y √

∗

ϵ = ϵ(r−d )/2 .

(35)

Substituting into the volume expression yields ! d∗ Y ∗ √ Vol(C) ≤ µi ϵ(r−d )/2 .

(36)

i=d∗ +1

µi ≤

i=d∗ +1

i=1

□ The empirical results suggest that GeomDR accelerates grokking by reducing the geometric complexity of learned representations. Neural networks can often fit the training data using a large family of distinct solutions, many of which distribute information across numerous weakly informative directions in representation space. Such solutions achieve low training error but need not capture the underlying algorithmic structure of the task. In contrast, successful generalization on algorithmic tasks may require the discovery of more structured representations. For modular arithmetic and related symbolic problems, these representations often appear substantially more concentrated, with variance captured by a relatively small number of dominant latent directions. Under this view, memorizing and algorithmic solutions occupy regions of representation space with different geometric characteristics: memorizing solutions tend to be spectrally diffuse and high-dimensional, whereas algorithmic solutions appear more spectrally concentrated. GeomDR explicitly suppresses covariance mass outside a target subspace through the objective (17). Propositions 1–3 show that minimizing this objective simultaneously controls the effective rank of the representation covariance, improves its low rank approximation, and compresses the geometric volume associated with tail covariance directions. In particular, Proposition 3 implies that as 11

B

• Embedding dimension: 8 • Input dimension: 16 • Hidden dimension: 128 • Residual blocks: 3 • Activation: GELU • Normalization: LayerNorm • Output dimension: 97 (arithmetic tasks), 120 (permutation composition)

Tasks and Architectures

Modular Addition The modular addition task receives a pair of integers (a, b) ∈ Z297 and predicts y = (a + b) mod 97. The dataset contains all possible ordered pairs in Z297 , resulting in 972 = 9409 examples.

Modular Division

Geometric Dimensionality Regularization (GeomDR) is applied to the hidden representations produced by the MLP. Specifically, the regularization term is evaluated for the outputs of all hidden layers (h0 , h1 , h2 , h3 ) and summed across layers. The input embedding representation z0 is excluded from the regularization.

The modular division task receives a pair (a, b) with b ̸= 0 and predicts y = (a · b−1 ) mod 97, where b−1 denotes the multiplicative inverse of b in Z97 . The dataset contains all valid pairs (a, b) with a ∈ Z97 and b ∈ {1, . . . , 96}, resulting in 97 × 96 = 9312 examples.

Permutation Composition MLP. For the permutation composition task, we use a larger MLP configuration. The task involves predicting one of |S5 | = 120 possible permutation classes, compared with 97 output classes in the modular arithmetic tasks. The embedding dimension is increased to 16 and the hidden dimension is increased to 256 while preserving the same residual MLP structure. The permutation MLP architecture uses:

Permutation Composition To evaluate whether the observed effects extend beyond modular arithmetic, we additionally consider permutation composition. Let S5 denote the symmetric group on five elements. Given two permutations σ, τ ∈ S5 , the task is to predict their composition

• Embedding dimension: 16 • Input dimension: 32 • Hidden dimension: 256 • Residual blocks: 3 • Activation: GELU • Normalization: LayerNorm • Output dimension: 120

y = σ ◦ τ, defined by
(σ ◦ τ )(i) = σ τ (i) . Since |S5 | = 120, the dataset contains 1202 = 14400 ordered pairs of permutations.

For permutation composition experiments, we use d∗ = 48 to account for the larger hidden representation dimension. All MLP models are trained using the same optimization protocol. We use full-batch AdamW optimization with a learning rate of 10−3 and weight decay of 10−1 .

MLP Architecture Each input token is mapped to an 8-dimensional embedding vector. The embeddings corresponding to the two inputs are concatenated,

Transformer Architecture

z0 = [e(a); e(b)],

To evaluate whether the proposed geometric intervention generalizes beyond multilayer perceptrons, we additionally consider a Transformer encoder architecture [25]. Each input symbol is mapped to a learned embedding of dimension

yielding a 16-dimensional input representation. The concatenated representation is projected into a hidden space of dimension 128: h0 = LayerNorm(GELU(Win z0 + bin )) . The network then applies three residual feed-forward blocks

dmodel = 32.

The final hidden representation is mapped to the output space using a linear classifier

All tasks considered in this work operate on ordered pairs of input symbols. Consequently, each example is represented as a sequence of length two. Learned positional embeddings are added to the token embeddings before being processed by the Transformer encoder. The encoder consists of two Transformer layers. Each layer employs multi-head self-attention with

ŷ = Wout h3 + bout .

nheads = 4

The architecture parameters are summarized below:

attention heads and a feed-forward network of dimension

hl+1 = LayerNorm(hl + GELU(Wl hl + bl )) , l ∈ {0, 1, 2}.

12

baseline and GeomDR experiments on permutation composition. The larger architecture was introduced solely to obtain a reproducible grokking regime and is not used for the addition or division tasks, which reliably grok under the default Transformer configuration. Geometric Dimensionality Regularization (GeomDR) is applied to the hidden token representations produced by each Transformer layer. For each layer, token representations from all examples and sequence positions are collected into a matrix

dff = 64. The architecture uses GELU activations, residual connections, LayerNorm, and the pre-normalization formulation of the Transformer. Dropout is disabled in all experiments. Let H (0) = E(X) + P denote the input token embeddings together with the learned positional embeddings. The Transformer encoder computes   H (ℓ+1) = TransformerLayer H (ℓ) ,

RNtokens ×dmodel , from which the covariance matrix is computed. The GeomDR penalty is then evaluated by summing the eigenvalues beyond a target dimensionality d∗ . The total regularization term is obtained by summing the GeomDR penalties across all Transformer layers. For each task, the train–test split is generated once and kept fixed across all runs. For all tasks, 30% of examples are used for training and 70% for testing. Different random seeds therefore affect only model initialization and optimizer state. Reported results are averaged across seeds, and variability is reported as one standard deviation. This protocol follows the standard experimental setting commonly adopted in grokking studies [19, 13].

for ℓ ∈ {0, 1}, where L = 2 denotes the number of encoder layers. After the final encoder layer, token representations are aggregated using mean pooling across the sequence dimension, 2

h=

1 X (L) H . 2 i=1 i

The pooled representation is subsequently normalized using LayerNorm and mapped to the output space through a linear classifier, ŷ = Wout h + bout . The architecture parameters are summarized below: • Embedding dimension (dmodel ): 32 • Attention heads: 4 • Transformer layers: 2 • Feed-forward dimension: 64 • Activation: GELU • Normalization: LayerNorm (pre-normalization) • Pooling: mean pooling • Output dimension: task dependent Transformer models for modular addition and modular division are trained using full-batch AdamW optimization with a learning rate of 10−3 and weight decay of 10−1 . For permutation composition, we use a learning rate of 3 × 10−4 and weight decay of 1.0. For the permutation composition task, the default Transformer configuration was insufficient to reliably exhibit grokking behavior within the available training budget with and without interventions. To obtain a stable grokking regime suitable for evaluating geometric interventions, we increased the model capacity only for the permutation task. Specifically, permutation composition experiments use dmodel = 64,

L = 3,

dff = 128.

All other architectural and optimization settings remain unchanged. This larger Transformer is used for both the 13

Table 2: Ablation studies of Geometric Dimensionality Regularization (GeomDR) on three grokking tasks using the MLP architecture. The table reports dimension sweeps (varying the target dimensionality d∗ ) and intervention-time sweeps (varying the GeomDR activation step ts ). Values denote the mean grokking step ± one standard deviation (in thousands of optimization steps). Baseline results are computed from 5 random seeds, whereas all GeomDR sweeps are computed from 10 random seeds. The baseline corresponds to standard training without Geometric Dimensionality Regularization (GeomDR). Success/Total columns report the number of runs that satisfied the grokking criterion out of all runs. For runs that did not satisfy the criterion within the training budget, the maximum budget of 500K optimization steps was used when computing the reported mean and standard deviation. Addition Division Permutation Setting

Step (K)

Success/Total

Step (K)

Success/Total

Step (K)

Success/Total

Baseline

362.1 ± 111.0

4/5

330.9 ± 100.9

5/5

300.2 ± 207.0

3/5

9/10 9/10 10/10 10/10 10/10 9/10 9/10

209.3 ± 216.1 183.8 ± 221.2 176.6 ± 189.6 176.5 ± 182.2 217.2 ± 216.4 176.5 ± 198.4 181.3 ± 198.3

7/10 7/10 9/10 9/10 7/10 8/10 9/10

160.3 ± 198.6 149.3 ± 196.5 176.5 ± 198.4 163.3 ± 171.6 162.1 ± 155.1 273.6 ± 231.8

9/10 8/10 8/10 9/10 10/10 7/10

Dimension Sweep 2 4 8 16 32 48 64

17.0 ± 16.3 17.6 ± 20.5 15.2 ± 24.3 7.0 ± 1.5 8.0 ± 2.7 14.5 ± 18.1 7.0 ± 1.3

10/10 10/10 10/10 10/10 10/10 10/10 10/10

81.1 ± 154.3 62.3 ± 154.0 64.8 ± 110.8 37.1 ± 41.0 19.5 ± 15.5 61.9 ± 154.1 60.9 ± 154.5

Intervention Time Sweep 500 1000 2000 4000 8000 16000

21.9 ± 29.6 12.7 ± 16.1 7.0 ± 1.5 8.1 ± 1.8 21.4 ± 27.8 25.5 ± 17.2

C

10/10 10/10 10/10 10/10 10/10 10/10

15.3 ± 6.5 13.3 ± 7.3 61.9 ± 154.1 66.4 ± 153.1 62.8 ± 153.6 24.5 ± 8.8

10/10 10/10 9/10 9/10 9/10 10/10

d∗ ∈ {2, 4, 8, 16, 32, 48, 64}. Each configuration was evaluated across ten random seeds {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}. Intervention-timing experiments used the same schedule and target dimensionality (λmax = 1, Tramp = 333, d∗ = 16) while varying the intervention-start step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000}, again using ten random seeds {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}. For modular division, the schedule-ablation study followed the same protocol as modular addition, using d∗ = 16 and ts = 2000 while varying λmax and Tramp over the same ranges. Each configuration was evaluated across five random seeds {0, 1, 2, 112, 1122}. Dimensionalitysweep experiments used the representative schedule (λmax = 3.33, Tramp = 1000, ts = 2000) and varied the target dimensionality over d∗ ∈ {2, 4, 8, 16, 32, 48, 64}. Each configuration was evaluated across ten random seeds {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}. Intervention-timing experiments used the same schedule and varied the intervention-start step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000}, while fixing λmax = 3.33, Tramp = 1000, and d∗ = 48. Each configuration was evaluated across ten random seeds {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}. For the permutation composition task, the scheduleablation study followed the same protocol as modular addition, except that the target dimensionality was fixed at d∗ = 48 to reflect the larger hidden representation dimension. Dimensionality-sweep experiments used the representative

Ablations

MLP For the MLP architecture, we first evaluated the effect of the regularization schedule by varying the final regularization strength λmax and ramp duration Tramp . We then performed target-dimensionality sweeps using a representative (not necessarily the best) schedule from the schedule-ablation study for each task and subsequently investigated the effect of intervention timing by varying the activation step ts . Table 2 summarizes the effects of Geometric Dimensionality Regularization (GeomDR) across three grokking tasks using the MLP architecture. We evaluate two classes of interventions: target-dimensionality sweeps and interventiontime sweeps. Schedule-sweep results are omitted from the table because they are reported separately as heatmaps in Figure 7. This figure additionally provides the corresponding dimensionality and intervention-timing visualizations for the modular division and permutation composition tasks. For modular addition, schedule-sweep experiments fixed the target dimensionality at d∗ = 16 and the interventionstart step at ts = 2000, while varying the regularization strength λmax ∈ {0.333, 1, 3.33, 10, 33.3} and ramp duration Tramp ∈ {100, 333, 1000, 3333, 10000}. Each configuration was evaluated across five random seeds {0, 1, 2, 112, 1122}. Dimensionality-sweep experiments used the representative schedule (λmax = 1, Tramp = 333, ts = 2000) and varied the target dimensionality over 14

(a) Division: schedule sweep

(b) Permutation composition: schedule sweep

(c) Division: dimensionality sweep

(d) Permutation composition: dimensionality sweep

(e) Division: intervention timing

(f) Permutation composition: intervention timing

Figure 7: Additional MLP experiments on modular division and permutation composition. Top row: schedule sweeps over regularization strength λ and ramp duration. Middle row: target-dimensionality sweeps. Bottom row: intervention-start sweeps. Consistent with the modular addition results reported in the main paper, Geometric Dimensionality Regularization (GeomDR) substantially accelerates grokking on both modular division and permutation composition. The effect remains robust across a broad range of schedules, target dimensionalities, and intervention times. schedule (λmax = 1, Tramp = 333, ts = 2000) and varied the target dimensionality over d∗ ∈ {2, 4, 8, 16, 32, 48, 64}. Intervention-timing experiments used the same schedule

and target dimensionality (λmax = 1, Tramp = 333, d∗ = 48) while varying the intervention-start step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000}. 15

regularization (λmax = 33.3). These results suggest that the amount of geometric compression required to accelerate grokking increases with task complexity. The dimensionality sweep reveals that GeomDR is effective across a broad range of target dimensions. For modular addition, the fastest and most stable performance is obtained for d∗ ∈ {16, 32}. For modular division, the best performance is achieved at d∗ = 32, although substantial improvements are observed across nearly the entire range of tested target dimensionalities. For permutation composition, performance is strongest at an intermediate target dimensionality d∗ = 16, while both excessively strong and excessively weak compression lead to slower grokking and reduced success rates. Extremely small target dimensions may impose excessive representational constraints, while excessively large target dimensions may provide insufficient geometric pressure. These results indicate that the benefits of GeomDR are not restricted to a narrowly tuned target dimensionality and that effective geometric regularization can be achieved across a broad range of target dimensions. The intervention-time sweep demonstrates that GeomDR remains effective across a broad range of activation times. For modular addition, the best performance is achieved at

Across all tasks, GeomDR substantially accelerates grokking relative to the baseline. For modular addition, the baseline model reaches the grokking criterion after 362.1 ± 111.0K optimization steps (4/5 successful runs), whereas multiple GeomDR configurations reduce the grokking time to 7.0 ± 1.3K steps (5/5 successful runs). The schedule-ablation results exhibit a relatively smooth dependence on both the regularization strength λmax and ramp duration Tramp . Rather than displaying highly irregular behavior, performance changes gradually across neighboring configurations, indicating that GeomDR is not overly sensitive to precise hyperparameter choices. Across all tasks, successful configurations form contiguous regions in the hyperparameter space, suggesting a broad basin of effective schedules. A consistent trend is that more challenging tasks require stronger regularization. For modular addition, optimal performance is achieved with relatively weak regularization (λmax = 1), whereas modular division benefits from intermediate regularization strengths (λmax = 3.33). The permutation composition task, which exhibits the most difficult grokking dynamics among the considered MLP benchmarks, achieves its best performance with substantially stronger

Table 3: Ablation studies of Geometric Dimensionality Regularization (GeomDR) on three grokking tasks using the Transformer architecture. The table reports dimension sweeps (varying the target dimensionality d∗ ), regularization-strength sweeps (varying the final GeomDR coefficient λmax ), and intervention-time sweeps (varying the GeomDR activation step ts ). Values denote the mean grokking step ± one standard deviation (in thousands of optimization steps). Results are computed from 5 random seeds. The baseline corresponds to standard Transformer training without Geometric Dimensionality Regularization (GeomDR). Success/Total columns report the number of runs that satisfied the grokking criterion out of all runs. For runs that did not satisfy the criterion within the training budget, the maximum budget of 200K optimization steps was used when computing the reported mean and standard deviation. Addition Division Permutation Setting

Step (K)

Success/Total

Step (K)

Success/Total

Step (K)

Success/Total

Baseline

66.4 ± 12.0

5/5

89.4 ± 27.3

5/5

72.6 ± 33.0

5/5

5/5 5/5 5/5 5/5

85.2 ± 58.0 155.1 ± 174.4 131.6 ± 184.2 68.0 ± 29.5

5/5 4/5 4/5 5/5

5/5 5/5 4/5 5/5 5/5

120.6 ± 59.9 74.8 ± 63.5 97.3 ± 66.2 48.9 ± 9.1 79.5 ± 61.2

4/5 4/5 4/5 5/5 4/5

71.3 ± 65.2 93.6 ± 73.0 80.2 ± 26.4 77.1 ± 52.8 102.6 ± 54.4 65.1 ± 14.2

4/6 4/5 5/5 5/5 4/5 5/5

Dimension Sweep 2 4 8 16

41.2 ± 19.2 36.7 ± 9.1 46.0 ± 9.8 136.7 ± 181.7

5/5 5/5 5/5 4/5

59.9 ± 14.1 58.8 ± 27.9 71.5 ± 39.5 70.8 ± 30.2 Lambda Sweep

33.3 10 3.33 1 0.333

74.0 ± 63.0 34.6 ± 6.2 34.5 ± 7.4 38.2 ± 15.4 35.9 ± 1.9

4/5 5/5 5/5 5/5 5/5

49.1 ± 13.4 67.8 ± 51.7 88.4 ± 59.6 65.5 ± 28.3 52.9 ± 9.6

Intervention Time Sweep 500 1000 2000 4000 8000 16000

47.2 ± 13.0 33.0 ± 2.5 41.1 ± 11.7 35.2 ± 3.3 36.4 ± 8.2 78.7 ± 60.9

5/5 5/5 5/5 5/5 5/5 4/5

112.0 ± 66.1 101.1 ± 54.0 63.8 ± 23.3 77.4 ± 61.4 48.7 ± 8.8 59.0 ± 21.0

16

4/5 4/5 5/5 4/5 5/5 5/5

ts = 1000, and λmax = 1. Using the best-performing dimensionality (d∗ = 16), we then performed a regularizationstrength sweep over λmax ∈ {0.333, 1, 3.33, 10, 33.3}. Finally, intervention-timing experiments varied the activation step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000} using the best-performing configuration from the previous stage (d∗ = 16, λmax = 1). Each configuration was evaluated across five random seeds {0, 1, 2, 112, 1122}. Across all three tasks, GeomDR consistently accelerates grokking relative to the baseline Transformer. For modular addition, the baseline model requires 66.4±12.0K optimization steps to satisfy the grokking criterion, whereas the best GeomDR configuration reduces this value to 33.0±2.5K optimization steps, corresponding to roughly a twofold acceleration. Similar improvements are observed for modular division, where the grokking step decreases from 89.4 ± 27.3K to 48.7 ± 8.8K optimization steps. For permutation composition, the strongest GeomDR configurations reduce the grokking step from 72.6 ± 33.0K to 48.9 ± 9.1K optimization steps. Although the improvements are generally smaller than those observed for the MLP architecture, these results demonstrate that the benefits of geometric regularization extend beyond modular arithmetic and remain effective on structured combinatorial problems. The dimensionality sweep reveals that moderate target dimensionalities provide the strongest improvements for the arithmetic tasks. For both modular addition and modular division, target dimensions d∗ ∈ {2, 4} yield the fastest mean grokking times, while larger dimensions remain beneficial but generally produce weaker acceleration. In contrast, permutation composition achieves its best performance at the larger target dimensionality d∗ = 16. This difference suggests that more complex tasks may benefit from retaining a higher-dimensional representation space, whereas simpler arithmetic tasks can be effectively accelerated through stronger geometric compression. For modular addition, the fastest mean grokking time is obtained at λmax = 3.33, whereas modular division achieves its best performance at λmax = 33.3. In contrast, permutation composition performs best at the more moderate value λmax = 1. These results further suggest that the optimal strength of geometric regularization is task-dependent and may vary with the complexity and structure of the underlying problem. The intervention-time sweep shows that the timing of the geometric intervention also influences performance. For modular addition, activating GeomDR after an initial unconstrained training phase (ts = 1000) produces the fastest and most consistent results. For modular division, the lowest mean grokking step obtained at ts = 8000. For permutation composition, the strongest performance is achieved at ts = 16000, although several intervention times produce broadly similar results. Taken together, these findings suggest that the optimal intervention time is task-dependent and may reflect differences in the duration of the memorization phase preceding generalization. Compared with the MLP results summarized in Table 2, the gains obtained with GeomDR are more modest for the Transformer architecture. Whereas GeomDR often accel-

ts = 2000, with similarly strong results at ts = 4000. For modular division, the best mean performance obtained at ts = 1000. For permutation composition, the lowest mean grokking step is observed at ts = 1000, while the most reliable performance is obtained at ts = 8000, which achieves a 10/10 success rate. Overall, all intervention times substantially outperform the baseline. Across the three tasks, earlyto-intermediate interventions generally provide the largest acceleration. One possible explanation is that prolonged memorization makes the learned representations increasingly specialized and less amenable to subsequent geometric restructuring. GeomDR may therefore be most effective after useful task structure has begun to emerge, but before the representation geometry becomes comparatively rigid.

Transformer For the Transformer architecture, we adopted a different tuning procedure from that used for the MLP. To reduce the size of the hyperparameter search space, we fixed the ramp duration at Tramp = 1000 throughout all Transformer experiments. We first performed target-dimensionality sweeps to identify the best-performing target dimensionality d∗ . Using this best configuration, we then optimized the regularization strength λmax . Finally, using the best-performing dimensionality and regularization strength, we evaluated the effect of intervention timing by varying the activation step ts . Table 3 summarizes ablation studies of Geometric Dimensionality Regularization (GeomDR) on the Transformer architecture across the modular addition, modular division, and permutation composition tasks. For Transformer-based modular addition, dimensionalitysweep experiments fixed the intervention-start step at ts = 1000, the ramp duration at Tramp = 1000, and the regularization strength at λmax = 1, while varying the target dimensionality over d∗ ∈ {2, 4, 8, 16}. Using the best-performing dimensionality (d∗ = 4), we then performed a regularization-strength sweep over λmax ∈ {0.333, 1, 3.33, 10, 33.3} while keeping ts = 1000 fixed. Finally, intervention-timing experiments used the best-performing configuration (d∗ = 4, λmax = 3.33) and varied the activation step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000}. Each configuration was evaluated across five random seeds {0, 1, 2, 112, 1122}. For Transformer-based modular division, we followed the same experimental procedure as for modular addition. Dimensionality-sweep experiments varied the target dimensionality over d∗ ∈ {2, 4, 8, 16} while fixing ts = 1000, and λmax = 1. Using the best-performing dimensionality (d∗ = 4), we then performed a regularization-strength sweep over λmax ∈ {0.333, 1, 3.33, 10, 33.3}. Finally, intervention-timing experiments varied the activation step over ts ∈ {500, 1000, 2000, 4000, 8000, 16000} using the best-performing configuration from the previous stage (d∗ = 4, λmax = 33.3). Each configuration was evaluated across five random seeds {0, 1, 2, 112, 1122}. For Transformer-based permutation composition, we followed the same experimental procedure as for modular addition. Dimensionality-sweep experiments varied the target dimensionality over d∗ ∈ {2, 4, 8, 16} while fixing 17

erates grokking by more than an order of magnitude in MLPs, he Transformer typically exhibits improvements ranging from roughly 1.5 times to 2times. Nevertheless, the qualitative trends remain consistent across architectures, with geometric interventions reliably reducing grokking time across all three tasks. One possible explanation is that Transformers already possess strong inductive biases toward structured representations through self-attention, residual connections, and normalization layers. As a result, the baseline Transformer groks substantially faster than the corresponding MLP, leaving less room for geometric regularization to further accelerate the memorization-to-generalization transition. While this interpretation remains speculative, it suggests that the effectiveness of GeomDR may depend in part on the extent to which the underlying architecture already promotes compressed and task-relevant representations.

Table 4: Width ablation on modular addition. Results are reported as mean grokking step ± standard deviation across five random seeds. Succ/Tot denotes the number of runs that satisfied the grokking criterion. (a) MLP GeomDR (d∗ = 16)

Baseline Width

Step (K)

Succ/Tot

Step (K)

Succ/Tot

32 64 128 256

500.0 ± 0.0 500.0 ± 0.0 333.0 ± 139.2 253.8 ± 146.4

0/5 0/5 4/5 4/5

21.6 ± 5.3 19.3 ± 23.5 6.8 ± 0.8 10.4 ± 5.6

5/5 5/5 5/5 5/5

(b) Transformer GeomDR (d∗ = 4)

Baseline

Width Ablation To evaluate the dependence of GeomDR on model capacity, we performed a width ablation study on the modular addition task for both MLP and Transformer architectures. For MLPs, width refers to the hidden dimension of the residual network, while for Transformers it refers to the model dimension dmodel . We evaluated MLP widths {32, 64, 128, 256} and Transformer widths {16, 32, 64, 128} while keeping all other architectural components fixed. For MLPs, we used the representative configuration identified in the schedule sweep (d∗ = 16, λmax = 1, Tramp = 333). For Transformers, we used the best-performing configuration identified during Transformer tuning, corresponding to a target dimensionality of d∗ = 4, regularization strength λmax = 3.33, intervention-start step ts = 1000, and ramp duration Tramp = 1000. Each experiment was repeated with five random seeds {0, 1, 2, 112, 1122}. Table 4 summarizes the results. For MLPs, GeomDR consistently accelerates grokking across all tested widths. At width 32, the baseline failed to satisfy the grokking criterion in any run within the training budget, whereas GeomDR achieved a mean grokking step of 21.6K with a 5/5 success rate. A similar effect is observed at width 64, where the baseline again failed in all runs while GeomDR achieved a mean grokking step of 19.3K and successfully grokked in all runs. At width 128, GeomDR reduces the mean grokking step from 333.0K to 6.8K while increasing the success rate from 4/5 to 5/5. At width 256, the mean grokking step decreases from 253.8K to 10.4K, again achieving a 5/5 success rate. Overall, the intervention substantially accelerates grokking and improves reliability across a broad range of MLP capacities. The Transformer results exhibit a markedly different pattern. At width 32, which corresponds to the default Transformer configuration used throughout the main experiments, GeomDR substantially accelerates grokking, reducing the mean grokking step from 65.4K to 34.1K while maintaining a 5/5 success rate. However, this benefit does not persist at larger widths. At width 64, the intervention succeeds in only 2/5 runs, yielding a mean grokking step of 139.9K when unsuccessful runs are included in the summary statis-

Width

Step (K)

Succ/Tot

Step (K)

Succ/Tot

16 32 64 128

200.0 ± 0.0 65.4 ± 7.2 43.7 ± 5.8 58.6 ± 26.2

0/5 5/5 5/5 5/5

200.0 ± 0.0 34.1 ± 5.7 139.9 ± 83.7 174.6 ± 56.9

0/5 5/5 2/5 1/5

tics. At width 128, only a single run satisfies the grokking criterion, resulting in an average grokking step of 174.6K. In both cases, performance is substantially worse than the corresponding baseline models. The narrowest Transformer configuration (width 16) fails to grok within the training budget both with and without intervention. Overall, the width ablation shows that GeomDR is highly robust across MLP capacities, suggesting that the intervention acts on a geometric property that remains stable under changes in representational width. Transformer results reveal a more architecture-dependent behavior. While GeomDR can accelerate grokking in Transformers under suitable settings, the fixed low dimensional target used here does not scale reliably with model width and can even degrade performance for larger models. These findings indicate that representation geometry remains a useful control signal beyond MLPs, but that effective geometric interventions are likely architecture-dependent. These findings indicate that representation geometry remains a useful control signal beyond MLPs, but that effective geometric interventions are likely architecture-dependent and may require retuning as model width increases.

Collapse–Grokking Lag Analysis Table 5 reports the quantitative analysis of the temporal separation between representation collapse and grokking. Collapse time is defined as the first step at which the effective dimensionality of the deepest hidden layer (L4 ) decreases below 50% of its initial value. The initial dimensionality is computed from the first recorded representation snapshot at initialization (step 0) using the effective dimension metric. The lag is defined as tgrok − tcollapse . All experiments use the modular addition task and the MLP architecture described earlier. Values are averaged over five random seeds in {0, 1, 2, 112, 1122}. For this setting, baseline models exhibit 18

a large delay between dimensionality collapse and successful generalization, whereas GeomDR substantially reduces this temporal gap, bringing representation compression and generalization into closer alignment. Table 5: Timing of dimensionality collapse and grokking. Collapse is defined as the first step at which the effective dimension of the final hidden layer falls below 50% of its initial value. Lag denotes the temporal separation tgrok − tcollapse . All values are reported in thousands of training steps (K) as mean ± standard deviation. Succ/Tot denotes the number of runs satisfying the grokking criterion out of the total number of runs. For the baseline, lag statistics were computed over the four seeds that successfully grokked; one non-grokking run was omitted from the lag analysis. Setting Baseline ∗

d = 64 d∗ = 32 d∗ = 16 d∗ = 8

Succ/Tot Collapse (K)

Grok (K)

Lag (K)

4/5

19.1 ± 2.0

283.4 ± 33.5 264.3 ± 33.8

5/5 5/5 5/5 5/5

3.7 ± 0.3 3.7 ± 0.3 3.9 ± 0.4 3.8 ± 0.2

5.9 ± 1.1 6.6 ± 1.5 6.9 ± 1.6 6.7 ± 1.2

2.2 ± 0.8 2.9 ± 1.4 3.0 ± 1.4 2.9 ± 1.1

19

Table 6: Random-sampling ablation for Geometric Dimensionality Regularization (GeomDR) on modular addition using the MLP architecture. The covariance spectrum is estimated from a randomly sampled subset of hidden representations. Values report mean ± standard deviation. Method Sample Size (% of Data) Success / Total Runtime (s) Grokking Step (K) Baseline GeomDR (100%) GeomDR (50%) GeomDR (25%) GeomDR (10%) GeomDR (5%) GeomDR (1%)

D

2822 (100%) 1411 (50%) 705 (25%) 282 (10%) 141 (5%) 28 (1%)

4/5 5/5 5/5 5/5 5/5 5/5 5/5

Scalability

12952.0 ± 3154.4 666.0 ± 368.0 447.2 ± 274.9 276.7 ± 50.6 259.6 ± 94.3 247.5 ± 76.9 1088.7 ± 366.8

362.1 ± 111.0 7.06 ± 1.74 7.36 ± 3.72 6.34 ± 1.05 6.81 ± 1.96 7.89 ± 1.87 35.76 ± 9.65

yield grokking times comparable to those obtained with fullspectrum estimation, with successful generalization typically occurring after approximately 6–8K optimization steps. The fastest average grokking time is observed at 25% sampling, reaching the grokking criterion after 6.34 ± 1.05K steps, compared to 7.06±1.74K steps for full-spectrum estimation. However, the differences in grokking steps among sampling ratios in the 5%–50% range are relatively small, indicating that accurate covariance estimation can be achieved from substantially reduced representation subsets without materially affecting grokking acceleration.

For large datasets, computing covariance spectra using all available representations may become prohibitively expensive. We therefore investigate whether GeomDR remains effective when the covariance spectrum is estimated from a randomly sampled subset of representations. Indeed, the computational complexity of GeomDR for a single layer is dominated by covariance estimation and eigendecomposition. Given N representations of dimension d, computing the covariance matrix requires O(N d2 ) operations, while eigendecomposition of the resulting d × d covariance matrix requires O(d3 ) operations. The overall complexity is therefore

Subsampling also substantially reduces computational cost. While full-spectrum GeomDR requires 666.0 ± 368.0 seconds of wall-clock time (one CPU), sampling ratios between 5% and 25% reduce runtime to approximately 250– 280 seconds, representing a 2.4×–2.7× speedup relative to full-spectrum GeomDR. Compared with the baseline model without GeomDR, these configurations reduce runtime by roughly a factor of 50. The 1% condition exhibits substantially higher runtime due to its much later grokking time, requiring considerably more optimization steps before satisfying the grokking criterion. These results suggest that the geometric signal exploited by GeomDR can be estimated accurately from relatively small random subsets of representations. In practice, covariance estimation using only 5%–25% of available representations is sufficient to recover nearly the full benefit of GeomDR while substantially reducing computational overhead. This observation indicates that precise estimation of the full covariance spectrum is not required for effective geometric control, improving the practical scalability and computational efficiency of the method.

O(N d2 + d3 ). In the regimes considered in this work, the number of representations is substantially larger than the representation dimension (N ≫ d), so the covariance computation dominates and the effective complexity scales approximately as O(N d2 ). If the covariance spectrum is estimated using only a fraction r of the available representations, the complexity for one layer becomes O(rN d2 + d3 ), implying an approximately linear reduction in computational cost with respect to the sampling ratio r. Experiments are performed on the modular addition task using the MLP architecture. The intervention uses λ = 3.33, target dimensionality d∗ = 16, activation step ts = 2000, and Tramp = 1000. At each optimization step, a fraction r ∈ {1%, 5%, 10%, 25%, 50%, 100%}

These findings suggest that GeomDR may be compatible with standard mini-batch optimization. The random sampling results show that only a small subset of representations is required to estimate the geometric signal driving the regularizer. Although this hypothesis was not evaluated directly in the present work, each mini-batch can be viewed as a random sample of the representation distribution, suggesting that the leading covariance structure may be estimated sufficiently accurately from batch-level statistics. If so, GeomDR could be implemented efficiently in large-scale settings without computing covariance matrices over the entire dataset.

of the training representations is sampled uniformly without replacement, and the GeomDR penalty is computed using only the sampled subset. For the modular addition task, the training set contains 2822 examples, corresponding to sample sizes ranging from 28 representations at 1% sampling to the full training set at 100% sampling. Table 6 reports both wall-clock runtime and grokking step. Across a broad range of sampling ratios, GeomDR remains highly effective despite using substantially fewer representations for covariance estimation. Sampling ratios between 5% and 50% 20

Table 7: Comparison of methods for accelerating grokking on modular addition using the MLP architecture. Values denote mean grokking step ± standard deviation (thousands of optimization steps). Succ/Tot indicates the number of runs satisfying the grokking criterion. For unsuccessful runs, the training budget of 500K steps was used when computing summary statistics. Method

Best Params

Step (K)

Succ/Tot

WD = 0

–

500.0 ± 0.0

0/5

WD = 1.0 λ = 1.0 λmax = 1 λ = 1.0

27.2 ± 43.9 372.1 ± 37.6 7.0 ± 1.5 500.0 ± 0.0

5/5 5/5 10/10 0/5

Weight Decay GrokFast GeomDR GrokFast + GeomDR

E

Comparison with Grokking Acceleration Methods

We compare Geometric Dimensionality Regularization (GeomDR) with two commonly studied grokking interventions: weight decay and GrokFast [11]. For weight decay, we evaluate WD ∈ {0.01, 0.03, 0.1, 0.3, 1.0} and report the best-performing configuration (WD = 1.0). For GrokFast, we evaluate amplification coefficients λ ∈ {0.05, 0.1, 0.5, 1.0, 5.0} using the default exponential moving average coefficient α = 0.999 and baseline weight decay WD = 0.1, and report the best-performing configuration identified in the sweep (λ = 1.0). GeomDR uses the best-performing modular-addition configuration identified in the ablation studies (d∗ = 16, ts = 2000, Tramp = 1000, λmax = 1.0). Results are summarized in Table 7. The best weightdecay configuration substantially accelerates grokking relative to the baseline, reaching the grokking criterion after 27.2±43.9K optimization steps. GrokFast also improves performance relative to the baseline, achieving a mean grokking step of 372.1 ± 37.6K and successfully grokking in all runs. However, the magnitude of this improvement is considerably smaller than that obtained with weight decay or GeomDR. GeomDR provides the strongest acceleration, reaching the grokking criterion after 7.0 ± 1.5K optimization steps while maintaining a perfect success rate across all runs. We additionally evaluated a combined GrokFast+GeomDR intervention. In contrast to the individual methods, the combined configuration failed to satisfy the grokking criterion in any run within the available training budget, suggesting that the two interventions may interact unfavorably and overconstrain the optimization dynamics. Increasing the weight-decay coefficient from 0.1 to 1.0 in the presence of GeomDR did not provide additional improvement, suggesting diminishing returns from combining stronger parameter-space regularization with explicit geometric regularization. These results indicate that direct control of representation geometry can provide a highly effective mechanism for accelerating grokking, substantially outperforming both tuned weight decay and GrokFast in the experimental setting considered here.

21

F

Use of Large Language Models

Large language models (LLMs) were used as writing and editing assistants during the preparation of this manuscript. Specifically, LLMs were used to improve the grammar, clarity, and presentation of the text. The scientific content of this work, including the research questions, hypotheses, experimental design, implementation, data analysis, interpretation of results, and conclusions, was developed and verified by the author. All AI-assisted text was reviewed and edited by the author prior to submission. The author assumes full responsibility for the accuracy and content of this manuscript.

G

Code Availability

Code is available at https://github.com/ maksimkazanskii/grokking.

22