A Basin-Selection Perspective on Grokking via Singular Learning Theory

Ben Cullen; Sergio Estan-Ruiz; Riya Danait; Jiayi Li

stat.MLMar 1, 2026

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

Abstract

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.

Panel of Figure 1 — 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 $\pm 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 $\varepsilon$ and, hence, their local learning coefficients are not equivalent. The final loss function shows that the zero-set need not be a single point.

Panel of Figure 4 — Figure 4: Log-linear relationship between the final LLC of a trained model and the SGLD localisation parameter $\gamma$ , shown for different fixed values of the SGLD learning rate.

Panel of Figure 5 — Figure 5: Tracking LLC curves and the parallel evolution of generalisation for different values of $p$ .

Panel of Figure 6 — Figure 6: Tracking LLC curves and the parallel evolution of generalisation for different dimensions of the hidden layer.

Panel of Figure 7 — Figure 7: Tracking LLC curves and the parallel evolution of generalisation for different weight decays.

Panel of Figure 8 — 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.

Panel of Figure 10 — Figure 10: Relationship between the learning rate and the severity of grokking for different weight decays.

Panel of Figure 11 — Figure 11: Relationship between the learning rate and the maximum LLC achieved during training.

Grokking

Grokking is delayed generalization in which neural networks memorize training data before abruptly learning a rule that works on unseen examples.

See all 77 papers

Statistical Vignettes on the Grokking Phenomenon
Oct 13, 2026Baye Malick Gning, Elvis DohmatobGrokking Neural Network Training Dynamics
Concordia University · Mila – Quebec AI Institute
A Stochastic--Geometric Theory of Scaling Laws in Grokking
Jun 29, 2026Róisín Luo, Christian Gagné, Jonas Ngnawé +2Neural Network Generalization Neural Network Optimization
Research Ireland – Centre for Research Training in AI · University of Galway · Universit´e Laval +2
Is Grokking a Loss of Normal Hyperbolicity of the Interpolation Manifold?
Aug 14, 2026Suvinava BasakGrokking Neural Network Training Dynamics
Technische Universität Braunschweig, Data Science, Universitätsplatz 2, 38106 Braunschweig, Germany

Oct 13, 2026J/K move · Enter open · S save

Statistical Vignettes on the Grokking Phenomenon

Baye Malick Gning, Elvis Dohmatob

Concordia University · Mila – Quebec AI Institute

Grokking Neural Network Training Dynamics Neural Network Generalization Statistical Learning Theory

Open paper AskPDF arXiv

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.

Abstract

Grokking

Explore similar work

Statistical Vignettes on the Grokking Phenomenon

A Stochastic--Geometric Theory of Scaling Laws in Grokking

Is Grokking a Loss of Normal Hyperbolicity of the Interpolation Manifold?