📊 statistics

A Stochastic--Geometric Theory of Scaling Laws in Grokking

This paper proposes a stochastic-geometric theory explaining grokking as an optimization-induced transition from a memorization shell to a generalization core within Adam's parameter space, thereby deriving and validating scaling laws for the delay time based on learning rate, batch size, and regularization.

Original authors: Róisín Luo, Christian Gagné, Jonas Ngnawé, Ihsan Ullah, Karyn Morrissey

Published 2026-07-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Róisín Luo, Christian Gagné, Jonas Ngnawé, Ihsan Ullah, Karyn Morrissey

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are watching a video game character trying to solve a tricky puzzle. At first, the character runs around frantically, memorizing every single step of the path they just took. They get the score to zero instantly, but they haven't actually learned the rules; they've just memorized the specific map. This is called memorization.

Then, something magical happens. After a long, boring period where nothing seems to change, the character suddenly stops, looks at the board, and gets it. They start solving new, unseen puzzles perfectly. This sudden "aha!" moment is what researchers call grokking.

For a long time, nobody knew why this delay happened. Was it magic? A glitch? In this paper, the authors suggest a new way to visualize what's going on inside the computer's brain (the neural network) using a map of shapes and distances.

The Onion of Solutions

The authors propose that the space where the computer's "brain" can exist looks like a set of nested spheres, kind of like an onion or a target.

The Outer Shell (Initialization): When the computer starts, its settings are scattered on a thin, outer shell. It's like throwing darts that all land on the very edge of a giant balloon.
The Middle Shell (Memorization): As the computer learns, it quickly slides down to a middle layer. Here, it has memorized the training data perfectly. It's like the character in our game who has memorized the map but doesn't understand the rules. The computer gets stuck here for a long time.
The Core (Generalization): Deep in the center of the onion lies the "generalization" core. This is the sweet spot where the computer actually understands the rules and can solve new problems.

The "grokking" phenomenon is simply the journey from the middle shell, through the empty space, to the core. The mystery was: Why does it take so long to get there?

The Drunk Walk and the Magnetic Pull

The authors explain that the computer's learning process (using an optimizer called Adam) is a mix of two forces:

The Drunk Walk (Diffusion): Because the computer learns from small batches of data, its path is a bit wobbly, like a drunk person walking in a straight line. This wobble is random.
The Magnetic Pull (Drift): There is also a steady force pulling the computer toward the center, especially because of a rule called $\ell_2$ regularization (which acts like a gentle magnet trying to keep the settings small).

The paper suggests that the computer gets stuck in the middle shell because the "drunk walk" is too weak to push it out of that layer, and the "magnetic pull" isn't strong enough to drag it straight to the center immediately. It has to wander around the middle shell for a while until, by chance, the random wobbles push it just far enough to fall into the core.

The Rules of the Game (Scaling Laws)

The authors didn't just guess this; they used math (specifically something called stopping-time theory and stochastic differential equations) to predict exactly how long this delay should last. They found three main "knobs" that control the speed of this journey:

Learning Rate ( $\eta$ ): This is how big a step the computer takes. The paper shows that if you make the steps too small, it takes forever to escape the middle shell. If you make them too big, the computer might overshoot the target. There is a "Goldilocks" zone.
Batch Size ( $b$ ): This is how many examples the computer looks at before taking a step. The paper suggests that larger batches make the "drunk walk" less wobbly, which actually slows down the escape from the middle shell.
Regularization ( $\lambda$ ): This is the strength of the "magnetic pull." The paper finds that a stronger pull helps the computer escape the middle shell faster, but only up to a point.

They derived specific formulas (scaling laws) for how these knobs change the time it takes to grok. For example, the time it takes to jump from memorization to generalization is roughly proportional to $1/(\eta \lambda)$ . This means if you double the learning rate or the regularization, the delay time is cut in half.

What They Ruled Out

The paper is very careful to say what they are not claiming. They do not say that grokking happens because the network suddenly "finds a circuit" or because of some mysterious biological-like awakening. Instead, they argue it is purely a result of the geometry of the solution space and the randomness of the optimization process. They also rule out the idea that this is just a random fluke; their math suggests it's a predictable, structural feature of how Adam optimizes these specific types of problems.

How Sure Are They?

The authors are quite confident in their theory, but they are careful to distinguish between what they proved and what they measured.

The Math: They derived these scaling laws using rigorous mathematical proofs based on continuous-time models of the learning process. They used a symbolic algebra system (a computer program that does math) to check their formulas, which gives them high confidence in the equations.
The Proof: They tested their ideas on two specific types of puzzles: group-theoretic learning (specifically on the symmetric group $S_5$ ) and modular arithmetic (specifically on integers modulo 127, written as $Z_{127}$ ).
The Results: In these experiments, the computer's behavior matched their predictions. For instance, when they changed the learning rate or batch size, the time it took to grok changed exactly as their formulas predicted. They also saw the "U-shaped" curve for the memorization radius, which their theory said would happen.

However, they note that their math relies on certain conditions, like having a small learning rate and a large batch size. They don't claim this explains every instance of grokking in every possible neural network, but rather that it explains the phenomenon in these specific, structured tasks where the "shell-core" geometry exists.

The Bottom Line

Grokking isn't magic; it's a geometric journey. The computer starts on the outside, gets stuck in a "memorization shell," and has to wander around until the random noise of its learning process pushes it into the "generalization core." The time it takes to make this jump depends on how fast it steps (learning rate), how steady its steps are (batch size), and how strongly it's pulled toward the center (regularization). The authors have mapped out these rules with math and confirmed them with experiments, giving us a clearer picture of why these neural networks sometimes seem to sleep for a long time before suddenly waking up.

Problem Statement
Delayed generalization, or "grokking," is a phenomenon observed in neural networks trained on noise-free, highly structured tasks (e.g., modular arithmetic, group-theoretic learning). In these settings, models rapidly achieve near-zero training loss (memorization) but fail to generalize for an extended period before abruptly transitioning to strong test performance. Despite extensive empirical study, the underlying mechanism driving this delay and the specific role of hyperparameters (learning rate, batch size, regularization) remain poorly understood.

Methodology
The authors propose a stochastic-geometric framework to characterize grokking, grounded in the optimization dynamics of the Adam optimizer with $\ell_2$ weight-shrinkage regularization.

Topological Characterization: The paper posits that Adam's optimization dynamics induce a "shell–core" topological configuration in the parameter space.
- Initialization: Solutions concentrate on a thin outer hyperspherical shell ( $\Theta$ ).
- Memorization: Enclosed within $\Theta$ is a shell of memorization solutions ( $M \setminus G$ ), where the model fits training data but fails to generalize.
- Generalization: At the center lies a compact core of generalization solutions ( $G$ ).
- The transition from memorization to generalization is modeled as a stochastic escape from the outer shell to the inner core.
Stochastic Differential Equation (SDE) Formulation:
- The authors derive a continuous-time limit for Adam's dynamics, formulating a joint-state SDE for the parameters ( $\theta$ ) and the first- and second-moment estimates ( $m, v$ ). This SDE accounts for the stochasticity of mini-batch gradients and the bias-correction mechanisms of Adam.
- Late-Stage Reduction: Empirical observations suggest that in the late stages of training (post-memorization), the moment estimates converge to their mean-field limits, and gradient noise becomes small and isotropic. These properties allow the complex joint-state SDE to be reduced to a tractable radial SDE governing the evolution of the squared parameter norm ( $r_t^2 = \|\theta_t\|^2$ ).
Stopping-Time Analysis:
- The transition time from the memorization manifold ( $\partial M$ ) to the generalization manifold ( $\partial G$ ) is analyzed using stopping-time theory for stochastic processes.
- The authors derive scaling laws for the radii of these manifolds and the expected transition time by solving Dirichlet boundary value problems and mean first-passage time equations derived from the reduced SDEs.

Key Contributions

Optimization-Induced Shell–Core Topology: The paper introduces a theoretical prior where initialization concentrates on a shell enclosing memorization solutions, which in turn enclose a generalization core. This configuration is induced specifically by Adam's dynamics with weight shrinkage.
Stochastic–Geometric Characterization: Grokking is formally characterized as a stochastic transition problem. The geometry of the solution space is defined by the first stopping time at which stochastic gradient flows escape the memorization manifold and reach the generalization manifold.
Derivation of Scaling Laws: The authors derive closed-form asymptotic scaling laws for:
- Memorization Radius ( $\rho_M$ ): Depends on the learning rate ( $\eta$ ), batch size ( $b$ ), and regularization ( $\lambda$ ). The theory predicts a U-shaped dependence on the ratio $\eta/b$ .
- Generalization Radius ( $\rho_G$ ): Depends on $\eta$ , $\lambda$ , and $b$ , decomposing into a landscape-dependent term and a fluctuation correction.
- Solution Transition Time ( $\tau_{M \to G}$ ): The expected number of iterations to transition from memorization to generalization scales inversely with the regularization coefficient ( $\lambda$ ) and the learning rate ( $\eta$ ), specifically $\tau \propto 1/(\eta \lambda)$ .

Results
The theoretical predictions were validated through experiments on two structured learning tasks: group-theoretic learning on the symmetric group $S_5$ and modular arithmetic on $\mathbb{Z}_{127}$ .

Empirical Validation: Experiments confirmed the predicted scaling laws for manifold radii and transition times across varying hyperparameters.
U-Shaped Curve: The data supported the theoretical prediction that the memorization radius $\rho_M$ exhibits a U-shaped curve with respect to $\eta/b$ .
Literature Recovery: The derived scaling laws recovered known empirical results from prior literature, such as the inverse dependence of grokking delay on regularization and the "Goldilocks zone" for weight norms.
Symbolic Verification: All closed-form and asymptotic results were machine-checked using a symbolic algebra system (Julia) to ensure technical correctness.

Significance
The paper claims to provide a theoretical framework that connects the geometry of solution manifolds with the stochastic dynamics of the Adam optimizer. By characterizing grokking through stopping-time analysis of SDEs, the work offers a mechanistic explanation for the delayed generalization phenomenon. It clarifies how specific hyperparameters modulate the stochastic diffusion and deterministic drift to influence the time required for a model to transition from memorization to generalization. The authors note that their results rely on the conditions required for the Adam SDE limit (sufficiently small learning rate and large batch size) and do not directly apply to the large $\eta/b$ regime.