A Stochastic--Geometric Theory of Scaling Laws in Grokking
Róisín Luo, Christian Gagné, Jonas Ngnawé, Ihsan Ullah, Karyn Morrissey
stat.ML, cs.AI, cs.LG
Submitted: 2026-06-29
License: http://creativecommons.org/licenses/by/4.0/
The gist: Delayed generalization (grokking) refers to the phenomenon in which a neural network fits its training data early in training but only begins to generalize after a prolonged delay, often through an
Terminology
Abstract
Delayed generalization (grokking) refers to the phenomenon in which a neural network fits its training data early in training but only begins to generalize after a prolonged delay, often through an abrupt transition. Despite extensive empirical study, its underlying mechanism remains poorly understood. In this work, we first theoretically characterize a shell--core topological configuration of the reachable solution space induced by Adam's optimization dynamics with weight-shrinkage regularization, supported by empirical evidence. This optimization-induced topological configuration gives rise to grokking. In model's parameter space, random initialization solutions concentrate on a thin outer spherical shell, enclosing another spherical shell of memorization solutions, which in turn contains a core corresponding to the generalization solutions. Leveraging stopping-time theory, we then analyze the geometry of this topological configuration and the solution transition time at which optimization trajectories escape the memorization manifold and first reach the boundary of the generalization manifold. Our theoretical analysis derives grokking scaling laws for the learning rate, batch size, and 2 regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.
Sources
- Intrinsic Task Symmetry Drives Generalization in Algorithmic Tasks
- Adam: A Method for Stochastic Optimization
- Optimization-Induced Dynamics of Lipschitz Continuity in Neural Networks
- Grokking Finite-Dimensional Algebra
- Explaining grokking through circuit efficiency
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Scalable Posterior Uncertainty for Flexible Density-Based Clustering
- A Transport Geometry of Benign Overfitting
- Efficient quadratic entropy with distance sketches