mlsanityData to 7 Oct 2026

A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisation

arXiv 2505.20172 ↗First posted 26 May 2025cs.LG math.OC stat.MLcited by 11

Not on the map. Why?

Abstract (arXiv)

We study the dynamics of gradient flow with small weight decay on general training losses F: ℝd → ℝ. Under mild regularity assumptions and assuming convergence of the unregularised gradient flow, we show that the trajectory with weight decay λ exhibits a two-phase behaviour as λ → 0. During the initial fast phase, the trajectory follows the unregularised gradient flow and converges to a manifold of critical points of F. Then, at time of order 1/λ, the trajectory enters a slow drift phase and follows a Riemannian gradient flow minimising the ℓ2-norm of the parameters. This purely optimisation-based phenomenon offers a natural explanation for the textitgrokking effect observed in deep learning, where the training loss rapidly reaches zero while the test loss plateaus for an extended period before suddenly improving. We argue that this generalisation jump can be attributed to the slow norm reduction induced by weight decay, as explained by our analysis. We validate this mechanism empirically on several synthetic regression tasks.

Key facts. Cited by 11 papers in the map (data through 7 Oct 2026). Builds on 29 of its 50 references that are in the map. First posted on arXiv 26 May 2025.

Builds on

29 of 50 references.

References · 29

  1. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
  2. 2016Understanding deep learning requires rethinking generalizationcited by 1,657
  3. 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
  4. 2018On Lazy Training in Differentiable Programmingcited by 671
  5. 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
  6. 2019Benign Overfitting in Linear Regressioncited by 544
  7. 2014Breaking the Curse of Dimensionality with Convex Neural Networkscited by 329
  8. 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
  9. 2020Kernel and Rich Regimes in Overparametrized Modelscited by 302
  10. 2020Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Losscited by 267
  11. 2018Characterizing Implicit Bias in Terms of Optimization Geometrycited by 259
  12. 2020Loss landscapes and optimization in over-parameterized non-linear systems and neural networkscited by 231
17 more
  1. 2022Towards Understanding Grokking: An Effective Theory of Representation Learningcited by 229
  2. 2022Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limitcited by 184
  3. 2022Towards Understanding Sharpness-Aware Minimizationcited by 145
  4. 2022Omnigrok: Grokking Beyond Algorithmic Datacited by 142
  5. 2021What Happens after SGD Reaches Zero Loss? --A Mathematical Frameworkcited by 106
  6. 2019How do infinite width bounded norm networks look in function space?cited by 105
  7. 2023Grokking as the Transition from Lazy to Rich Training Dynamicscited by 102
  8. 2023Grokking modular arithmeticcited by 77
  9. 2023Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokkingcited by 76
  10. 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
  11. 2024Why Do You Grok? A Theoretical Analysis of Grokking Modular Additioncited by 34
  12. 2023Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Datacited by 32
  13. 2022Label noise (stochastic) gradient descent implicitly solves the Lasso for quadratic parametrisationcited by 29
  14. 2023Penalising the biases in norm regularisation enforces sparsitycited by 17
  15. 2024Early alignment in two-layer networks training is a two-edged swordcited by 17
  16. 2024Simplicity bias and optimization threshold in two-layer ReLU networkscited by 11
  17. 2024Singular-limit analysis of gradient descent with noise injectioncited by 5

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 2; 2026: 9.

11 papers.