A Stochastic--Geometric Theory of Scaling Laws in Grokking

2026-06-29

AI Summary

This work uncovers the intrinsic mechanism underlying the phenomenon of grokking—where neural networks fit training data early in training yet generalize only after a prolonged delay. By integrating stochastic optimization theory, differential geometry, and stopping-time analysis, the study introduces, for the first time, a “shell–core” topological configuration in parameter space, elucidating how Adam optimization combined with ℓ² weight decay shapes this structure to drive grokking. Building on this geometric insight, the authors derive precise scaling laws governing the interplay among learning rate, batch size, and regularization strength. These theoretical predictions are rigorously validated through experiments that successfully reproduce and unify findings reported in prior literature.

Abstract

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.

Problem

• grokking
• delayed generalization
• scaling laws
• optimization dynamics
• generalization

Innovation

• grokking
• scaling laws
• optimization dynamics
• topological configuration
• stopping-time theory