A Basin-Selection Perspective on Grokking via Singular Learning Theory
Authors: Ben Cullen, Sergio Estan-Ruiz, Riya Danait, Jiayi Li
Organizations: Department of Computer Science, University of Pisa, Italy · Department of Mathematics, Imperial College London, United Kingdom · Mathematical Institute, University of Oxford, United Kingdom · Section of Mathematics and Artificial Intelligence, Max Planck Institute of Molecular Cell Biology and Genetics, Center for Systems Biology Dresden, Faculty of Mathematics, TU Dresden, Germany
Grokking, the abrupt transition from memorization to generalisation after extended training, suggests the presence of competing solution basins with distinct statistical properties. We study this phenomenon through the lens of Singular Learning Theory (SLT), a Bayesian framework that characterizes the geometry of the loss landscape. The key measure is the local learning coefficient (LLC) which quantifies the local degeneracy of the loss surface. SLT links lower-LLC basins to higher posterior mass concentration and lower expected generalisation error. Leveraging SLT, we develop a basin-selection perspective on grokking in quadratic networks: LLC ranks competing near-zero-loss basins by statistical preference, while the training-time transition between them is governed by optimisation dynamics. In this view, grokking corresponds to a transition from a higher-LLC (memorising) basin to a lower-LLC (generalising) basin that dominates the posterior. To support this, we derive analytic formulas for the LLC in shallow quadratic networks under both lazy and feature learning regimes. Empirically, we demonstrate that LLC trajectories estimated from training data track the onset of generalisation and provide an informative probe of the optimisation path.
Figures & tables
Regime
Solution type
Result
Role
General
True solution
Corollary 4.2
Architecture-level reference: provides generic local complexity near a true solution, independently of the specific optimisation regime.
Early
NTK linearised
Theorem 5.1
Early-stage baseline when training remains close to initialisation and the network is well approximated by a kernel model.
Early
Lazy-feature
Theorem 5.2
Early-stage fixed-representation regime for modular addition, where the hidden features are effectively frozen.
Early
Memorisation
Corollary 5.3
Early-stage memorising-basin prediction obtained from the lazy-feature analysis.
Late
Feature learning
Corollary 5.4
Late-stage structured-basin prediction for the final trained model; this is the formula tested in Section 6 through its dependence on p and K .
Table 1: Summary of derived LLC formulas.
Figure 1: Empirical validation of the linear dependence of the final LLC on architectural parameters.
Figure 2: Three plots examining different aspects of the same experiment during training. Results are aggregated over 100 random seeds; solid lines denote the mean across seeds, while shaded regions indicate ±1 standard deviation. Dataset p=53 , learning rate 0.0001, weight decay 0.00001, batch size 128 and hidden dimension 1024. The Hutchinson trace estimator is calculated using m=100 random vectors. Left: LLC curve and training/validation loss curves. Centre: LLC curve and training/validation accuracy curves. Right: LLC curve and Hutchinson trace curve.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Three different loss landscapes are realised alongside their contour maps in the (x,y) -plane. Although the first two loss functions share the same zero-set, their volume-sets scale differently with ε and, hence, their local learning coefficients are not equivalent. The final loss function shows that the zero-set need not be a single point.
Figure 4: Log-linear relationship between the final LLC of a trained model and the SGLD localisation parameter γ , shown for different fixed values of the SGLD learning rate.
Figure 5: Tracking LLC curves and the parallel evolution of generalisation for different values of p .
Figure 6: Tracking LLC curves and the parallel evolution of generalisation for different dimensions of the hidden layer.
Figure 7: Tracking LLC curves and the parallel evolution of generalisation for different weight decays.
Figure 8: Tracking LLC curves and the parallel evolution of generalisation for different learning rates.
Figure 9: The same model is trained with different learning rates. The learning rate is plotted against its GSM.
Figure 10: Relationship between the learning rate and the severity of grokking for different weight decays.
Figure 11: Relationship between the learning rate and the maximum LLC achieved during training.
In this work, we study grokking in a simple, analytically tractable setting whose rich phenomenology allows us to disentangle the effects of initialization, signal strength, weight decay, sample size, and input dimension. These factors can produce superficially similar learning curves through distinct mechanisms, complicating their interpretation. Our analysis distinguishes early fitting through nuisance features, core feature learning, and delayed generalization. We identify sharp conditions under which grokking provably occurs and derive analytic asymptotic expressions for the lengths of the generalization plateaus. Our results also characterize regimes of early generalization and failure to reach a prescribed test accuracy target. We show that the interpolation edge constitutes a distinct dynamical regime, where vanishing spectral gaps separate core learning from generalization, helping clarify the role of proximity to interpolation in existing accounts of grokking. The resulting picture is analogous to fast and slow phases in chemical kinetics and suggests that the lazy-to-rich transition discussed in the literature may reflect a similar separation of dynamical time scales. Our theoretical results are faithfully supported by extensive experiments.
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 ℓ2 regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.
Róisín Luo, Christian Gagné, Jonas Ngnawé +2
Research Ireland – Centre for Research Training in AI · University of Galway · Universit´e Laval +2
A recent line of work recasts the post-memorization phase of grokking as constrained optimization: once a network interpolates the training set, weight decay drives a slow drift along the zero-loss manifold toward lower norm. In the language of dynamical systems, this is a fast-slow system in which the interpolation manifold plays the role of a slow manifold. We ask a question that this framing makes natural but the existing literature does not address: is the sharp generalization transition a loss of normal hyperbolicity of that manifold: a fold- or bifurcation-like event in which a normal restoring direction goes flat? Or does the manifold stay uniformly attracting while generalization happens by smooth drift? We propose a simple, optimizer-agnostic diagnostic: the smallest nonzero singular value σmin+(J) of the residual Jacobian, which, for the squared loss, equals the slowest normal restoring rate of the manifold. On a two-layer ReLU network trained to grok modular addition under squared loss, σmin+(J) does not collapse at the transition; it is near zero only before memorization and attains its largest values during the transition. The result holds across five seeds, and the six smallest singular values behave identically; there is no subspace-local collapse either. This is preliminary evidence against the bifurcation hypothesis and in favor of the smooth-contraction picture. We are explicit that a single-setting, gradual-transition experiment under Adam optimizer does not prove the absence of a bifurcation; it constrains where one could hide.
Suvinava Basak
Technische Universität Braunschweig, Data Science, Universitätsplatz 2, 38106 Braunschweig, Germany