A Stochastic--Geometric Theory of Scaling Laws in Grokking
For researchers studying neural network training dynamics, this work offers a mechanistic theory of grokking with predictive scaling laws, though it is incremental as it builds on known optimization and regularization concepts.
The paper provides a theoretical explanation for grokking (delayed generalization) in neural networks, identifying a shell-core topological structure in parameter space induced by Adam optimization with weight decay. It derives scaling laws for learning rate, batch size, and ℓ2 regularization that predict the transition time to generalization, validated experimentally.
Delayed generalization (\ie~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 $\ell_2$ regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.