Pith. sign in

REVIEW 3 major objections 6 minor 28 references

Dimensionality collapse precedes grokking, and forcing it with a spectral regularizer can accelerate generalization by up to 52 times.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · grok-4.5

2026-07-14 04:00 UTC pith:JQX2LNR4

load-bearing objection Solid empirical lever: GeomDR reliably accelerates grokking on the standard suite, but the causal claim that geometry (not scheduled post-memorization regularization) is the driver is still under-isolated. the 3 major comments →

arxiv 2607.11666 v1 pith:JQX2LNR4 submitted 2026-07-13 cs.LG

How to Tame Grokking: Representation Geometry as a Control Signal

Maksim A Kazanskii This is my paper
classification cs.LG
keywords grokkingrepresentation geometryeffective dimensionalityspectral regularizationdelayed generalizationmodular arithmeticGeomDR
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Grokking is the delayed jump from memorizing training data to true generalization after long optimization. This paper shows that the effective dimensionality of hidden representations collapses just before that jump, across modular arithmetic and permutation tasks. The authors introduce GeomDR, a simple spectral regularizer that penalizes covariance mass outside a chosen target number of directions, applied after an initial free-training phase. On MLPs and transformers, this geometric intervention can accelerate the onset of generalization by as much as fifty-two times relative to ordinary AdamW, and can also delay or suppress it under stronger schedules. The claim is that representation geometry is not only a byproduct of learning but a practical control signal for when delayed generalization appears.

Core claim

Across modular addition, modular division, and permutation composition, collapse of effective representation dimensionality consistently precedes the memorization-to-generalization transition. Directly suppressing the sum of covariance eigenvalues beyond a target dimensionality d* via Geometric Dimensionality Regularization (GeomDR) systematically alters grokking timing, accelerating it by up to 52× in MLPs and about 1.5–2× in transformers depending on schedule and target dimension.

What carries the argument

Geometric Dimensionality Regularization (GeomDR): for each hidden layer, the loss is the sum of covariance eigenvalues past a target dimensionality d*, cosine-ramped in after an unconstrained phase. It forces variance into a lower-dimensional subspace and thereby controls when grokking occurs.

Load-bearing premise

That the speedup comes from geometric compression itself rather than the extra gradient signal or schedule, and that the same link holds beyond the small modular-arithmetic and permutation benchmarks used throughout.

What would settle it

On modular addition with the same MLP, force earlier dimensionality collapse with GeomDR across several d* values; if test accuracy still jumps only after hundreds of thousands of steps, or if collapse never precedes grokking under the paper’s own criteria, the claimed control link fails.

Watch this falsifier — get emailed when new claim-graph text bears on it.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 6 minor

Summary. The paper studies delayed generalization (grokking) through the lens of representation geometry. It reports that effective-dimensionality collapse of hidden activations consistently precedes the memorization-to-generalization transition on modular addition, modular division, and S5 permutation composition. Motivated by this, it introduces GeomDR, a delayed cosine-ramped spectral penalty that sums covariance eigenvalues beyond a target rank d* and thereby compresses hidden representations. Across MLPs and small Transformers, multi-seed schedule, dimensionality, and intervention-timing sweeps show that GeomDR can accelerate grokking by large factors (up to ~52× on modular addition with MLPs relative to AdamW), and can also delay or suppress it under aggressive settings. Ablations include width, random covariance subsampling, and comparisons to weight decay and GrokFast. Appendix A supplies rank/volume interpretations of the regularizer; code is released.

Significance. If the geometric control story holds, the work supplies a practical, representation-level knob for studying and steering delayed generalization, complementary to weight-norm, sparsity, and gradient-filtering approaches. The empirical package is a genuine strength: clear grokking criterion, 5–10 seeds, success rates, schedule/d*/timing heatmaps, width ablation, scalability via covariance subsampling, and a direct comparison table against weight decay and GrokFast, plus public code. Even if the precise causal pathway is not fully isolated, the demonstration that a simple spectral intervention can move grokking times by more than an order of magnitude on standard algorithmic benchmarks is useful for the grokking literature and for representation-geometry work more broadly.

major comments (3)
  1. [Methods (GeomDR); Results Intervention; Appendix E / Table 7] The central framing—that representation geometry is the causal control signal—is not isolated from generic post-memorization regularization. GeomDR is a delayed, cosine-ramped penalty λ(t) L_GeomDR with L_GeomDR = sum_{j>d*} μ_j (Eqs. 10–14, Methods). Appendix E and Table 7 compare to weight decay and GrokFast, but there is no matched control that applies a penalty of comparable magnitude and schedule that does not selectively suppress the covariance tail (e.g., scheduled Tr(C), Frobenius penalty on centered activations, or L2 on activation norms). Without such controls, acceleration could arise from any sufficiently strong regularizer that destabilizes memorizing solutions after ts, rather than from spectral compression per se. This is load-bearing for the title and Discussion claim that geometry is the controllable mechanism.
  2. [Observation; Fig. 5; Table 5; Abstract; Conclusion] The causal reading of “collapse precedes and drives grokking” rests on temporal precedence plus intervention (Observation; Fig. 5; Table 5). Collapse is defined post hoc as Deff falling below 50% of its step-0 value in the deepest layer. The paper correctly notes that this does not prove dimensionality is the sole mechanism, but the abstract and conclusion still present geometry as a control signal for delayed generalization. Either (i) add non-spectral scheduled controls as above, or (ii) explicitly downgrade the claim to “a spectral regularizer with this schedule accelerates grokking and co-moves with Deff,” and treat the geometric interpretation as a working hypothesis (as Appendix A already does).
  3. [Experimental Protocol; Table 1; Table 4; Discussion] All positive results are on small algorithmic tasks with fixed 30/70 splits and full-batch AdamW (modular arithmetic mod 97; S5 composition). Discussion acknowledges the scope limit, but Table 1 and the abstract still advertise up to 52× acceleration as evidence that geometric interventions are a practical approach for delayed generalization in neural networks. The Transformer width ablation (Table 4b) already shows that a fixed low d* can hurt larger models (success drops to 2/5 and 1/5 at widths 64 and 128). A major revision should either (a) include at least one non-algorithmic or larger-scale setting, or (b) substantially qualify the generality claim and state that retuning of d* and schedule is required when capacity or architecture changes.
minor comments (6)
  1. [Experimental Protocol] Grokking criterion (train acc = 1.0 and test ≥ 0.999 × train for 10 consecutive evaluations every 10 steps) is clear but unconventional; a short sensitivity check (e.g., 0.99 vs 0.999, or consecutive-window length) would help readers compare to prior work.
  2. [Results / Observation] Figure 1–2 captions and panel labels are dense; EMA window (100) and the “final 20k unsmoothed” choice should be stated once in Methods and referenced, not only in figure text.
  3. [Methods] Notation: Deff (participation ratio, Eq. 5) is used as the diagnostic, while the regularizer targets eigenvalue mass beyond d* (Eq. 10). A sentence clarifying that GeomDR does not directly optimize Deff would avoid conflating the two.
  4. [Table 1; Appendix C] Table 2 vs main-text best configs: some “best” numbers in Table 1 appear to come from different (λ, Tramp, d*) choices per task; a single footnote listing the exact best hyperparameter tuple for each cell would improve reproducibility.
  5. [Abstract; Introduction] Typographical: “GeomDRconsistently”, “effectivedimensionality”, and similar missing spaces appear in the abstract/intro PDF text; clean for camera-ready.
  6. [Related Work] Related Work could briefly situate GeomDR against other spectral / covariance regularizers in SSL (e.g., VICReg, dimensional collapse literature already cited) to clarify novelty of the delayed schedule for grokking rather than of spectral penalties in general.

Circularity Check

0 steps flagged

No circularity: GeomDR is an independent spectral penalty; acceleration is measured against an external AdamW baseline on held-out test accuracy.

full rationale

The paper is an empirical intervention study, not a first-principles derivation. Effective dimensionality Deff (participation ratio of the covariance spectrum) is observed to collapse before the grokking criterion (train acc=1 and test acc >=0.999*train for 100 steps). GeomDR is then defined independently as the sum of eigenvalues beyond a chosen target d* and added to the task loss after a free memorization phase. Acceleration (up to 52x) is reported by comparing wall-clock optimization steps to the identical architecture trained with standard AdamW (no GeomDR). Appendix A only derives elementary spectral consequences of that objective (effective-rank bound, Frobenius low-rank error, volume compression); none of those identities equate the grokking time to a fitted parameter. There are no self-citations that carry the central claim, no uniqueness theorems imported from the author, and no renaming of a known pattern as a new prediction. The result is therefore self-contained against external baselines and does not reduce by construction to its inputs.

Axiom & Free-Parameter Ledger

4 free parameters · 3 axioms · 1 invented entities

The central empirical claim rests on standard optimization assumptions plus a small set of free schedule parameters chosen by grid search; no new physical entities or ad-hoc mathematical axioms are introduced. The geometric interpretations in Appendix A are consequences of the regularizer definition, not extra postulates.

free parameters (4)
  • λ_max (final regularization strength) = task-dependent (1–33.3)
    Chosen by schedule sweep; best values differ by task (1 for addition, 3.33 for division, 33.3 for permutation).
  • T_ramp (cosine ramp duration) = 333 or 1000
    Grid-searched; representative values 333 or 1000 steps.
  • t_s (intervention start step) = 1000–4000
    Swept; optimal window is task- and architecture-dependent (typically 1000–4000).
  • d* (target dimensionality) = task-dependent (4–64)
    Swept over {2,4,8,16,32,48,64}; best values vary by task and architecture.
axioms (3)
  • domain assumption Effective dimensionality measured by participation ratio of the empirical covariance eigenvalues is a meaningful geometric summary of hidden representations.
    Standard in representation-geometry literature (Ansuini et al., Li et al.); used throughout Observation and Methods.
  • domain assumption Full-batch AdamW with fixed learning rate and weight decay is a fair baseline for measuring grokking acceleration.
    Inherited from the grokking literature (Power et al. 2022) and used for all comparisons.
  • ad hoc to paper A run has 'grokked' when train accuracy = 1.0 and test accuracy ≥ 0.999 imes train accuracy for 10 consecutive evaluations.
    Operational definition introduced in Experimental Protocol; all reported speed-ups depend on it.
invented entities (1)
  • Geometric Dimensionality Regularization (GeomDR) independent evidence
    purpose: Spectral loss that sums covariance eigenvalues beyond a target rank d* to force dimensionality collapse.
    New regularizer introduced in Methods; independent evidence is the empirical acceleration it produces, which is falsifiable by re-running the experiments.

pith-pipeline@v1.1.0-grok45 · 26771 in / 2752 out tokens · 20701 ms · 2026-07-14T04:00:04.108025+00:00 · methodology

0 comments
read the original 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.

Figures

Figures reproduced from arXiv: 2607.11666 by Maksim A Kazanskii.

Figure 1
Figure 1. Figure 1: Evolution of accuracy and representation geometry during baseline grokking on modular addition. The model is [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Evolution of accuracy and representation geometry under GeomDR for modular addition. The model is trained with [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗
Figure 4
Figure 4. Figure 4: Effect of target dimensionality [PITH_FULL_IMAGE:figures/full_fig_p006_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Relationship between representation collapse time [PITH_FULL_IMAGE:figures/full_fig_p007_5.png] view at source ↗
Figure 7
Figure 7. Figure 7: Additional MLP experiments on modular division and permutation composition. Top row: schedule sweeps over [PITH_FULL_IMAGE:figures/full_fig_p015_7.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

28 extracted references · 12 linked inside Pith

  1. [1]

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

    2021

  2. [2]

    H.; and Zoccolan, D

    Ansuini, A.; Laio, A.; Macke, J. H.; and Zoccolan, D

  3. [3]

    InAdvances in Neural Infor- mation Processing Systems, volume 32, 6114–6125

    Intrinsic Dimension of Data Representations in Deep Neural Networks. InAdvances in Neural Infor- mation Processing Systems, volume 32, 6114–6125

  4. [4]

    Bardes, A.; Ponce, J.; and LeCun, Y. 2022. VI- CReg: Variance-Invariance-Covariance Regulariza- tion for Self-Supervised Learning. InInternational Conference on Learning Representations

    2022

  5. [5]

    TheApproximationof OneMatrixbyAnotherofLowerRank.Psychometrika, 1(3): 211–218

    Eckart,C.;andYoung,G.1936. TheApproximationof OneMatrixbyAnotherofLowerRank.Psychometrika, 1(3): 211–218

    1936

  6. [6]

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

    Pith/arXiv arXiv 2017

  7. [7]

    Han, X.; Papyan, V.; and Donoho, D. L. 2022. Neural CollapseUnderMSELoss: ProximitytoandDynamics on the Central Path. InInternational Conference on Learning Representations

    2022

  8. [8]

    Un- derstanding Dimensional Collapse in Contrastive Self- Supervised Learning

    Jing,L.;Vincent,P.;LeCun,Y.;andTian,Y.2022. Un- derstanding Dimensional Collapse in Contrastive Self- Supervised Learning. InInternational Conference on Learning Representations

    2022

  9. [9]

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

    2002

  10. [10]

    Kornblith, S.; Norouzi, M.; Lee, H.; and Hinton, G

  11. [11]

    InProceedings of the 36th International Conference on Machine Learning, 3519–3529

    Similarity of Neural Network Representations Revisited. InProceedings of the 36th International Conference on Machine Learning, 3519–3529

  12. [12]

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

    Pith/arXiv arXiv 2022

  13. [13]

    G.; Kim, K.; and Lee, K

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

    Pith/arXiv arXiv 2024

  14. [14]

    Li, C.; Farkhoor, H.; Liu, R.; and Yosinski, J. 2018. Measuring the Intrinsic Dimension of Objective Land- scapes. InInternational Conference on Learning Rep- resentations

    2018

  15. [15]

    J.; Tegmark, M.; and Williams, M

    Liu, Z.; Kitouni, O.; Nolte, N.; Michaud, E. J.; Tegmark, M.; and Williams, M. 2022. Towards Un- derstanding Grokking: An Effective Theory of Repre- sentation Learning.arXiv, abs/2205.10343

    Pith/arXiv arXiv 2022

  16. [16]

    J.; and Tegmark, M

    Liu, Z.; Michaud, E. J.; and Tegmark, M. 2022. Om- niGrok: Grokking Beyond Algorithmic Data.arXiv, abs/2210.01117

    Pith/arXiv arXiv 2022

  17. [17]

    Loshchilov, I.; and Hutter, F. 2017. SGDR: Stochastic GradientDescentwithWarmRestarts. InInternational Conference on Learning Representations

    2017

  18. [18]

    Minegishi, G.; Iwasawa, Y.; and Matsuo, Y. 2023. Bridging Lottery Ticket and Grokking: Understand- ingGrokkingfromInnerStructureofNetworks.arXiv, abs/2310.19470

    Pith/arXiv arXiv 2023

  19. [19]

    Nanda,N.;Chan,L.;Lieberum,T.;Smith,J.;andStein- hardt, J. 2023. Progress Measures for Grokking via Mechanistic Interpretability.arXiv, abs/2301.05217

    Pith/arXiv arXiv 2023

  20. [20]

    Papyan, V.; Han, X.; and Donoho, D. L. 2020. Preva- lenceofNeuralCollapseDuringtheTerminalPhaseof Deep Learning Training.Proceedings of the National Academy of Sciences, 117(40): 24652–24663

    2020

  21. [21]

    Power, A.; Burda, Y.; Edwards, H.; Babuschkin, I.; and Misra, V. 2022. Grokking: Generalization Be- yondOverfittingonSmallAlgorithmicDatasets.arXiv, abs/2201.02177

    Pith/arXiv arXiv 2022

  22. [22]

    DoVisionTransformersSee Like Convolutional Neural Networks? InAdvances in Neural Information Processing Systems, volume 34, 12116–12128

    Raghu, M.; Unterthiner, T.; Kornblith, S.; Zhang, C.; andDosovitskiy,A.2021. DoVisionTransformersSee Like Convolutional Neural Networks? InAdvances in Neural Information Processing Systems, volume 34, 12116–12128

    2021

  23. [23]

    Recanatesi, S.; Farrell, M.; Advani, M.; Moore, T.; Lajoie, G.; and Shea-Brown, E. 2019. Dimensionality CompressionandExpansioninDeepNeuralNetworks. arXiv, abs/1906.00443

    Pith/arXiv arXiv 2019

  24. [24]

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

    Pith/arXiv arXiv 2017

  25. [25]

    DeepLearningand the Information Bottleneck Principle

    Tishby,N.;andZaslavsky,N.2015. DeepLearningand the Information Bottleneck Principle. In2015 IEEE Information Theory Workshop (ITW), 1–5. Jerusalem, Israel

    2015

  26. [26]

    Varma, V.; Shah, R.; Kenton, Z.; Kramár, J.; and Ku- mar, R. 2023. Explaining Grokking through Circuit Efficiency.arXiv, abs/2309.02390

    Pith/arXiv arXiv 2023

  27. [27]

    N.; Kaiser, Ł.; and Polosukhin, I

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

    2017

  28. [28]

    Xu, Z.; Ni, Z.; Wang, Y.; and Hu, W. 2025. Let Me Grok for You: Accelerating Grokking via Embedding TransferfromaWeakerModel.arXiv,abs/2504.13292. 8 Contents Introduction 1 Related Work 1 Methods 2 Geometric Characterization of Representations . . . . . . . . . . . . . . . . . . . . . . . . . . . 2 Geometric Dimensionality Regularization . . . . . . . . ....

    Pith/arXiv arXiv 2025