Table of Contents
Fetching ...

Why Grokking Takes So Long: A First-Principles Theory of Representational Phase Transitions

Truong Xuan Khanh, Truong Quynh Hoa, Luu Duc Trung, Phan Thanh Duc

Abstract

Grokking is the sudden generalization that appears long after a model has perfectly memorized its training data. Although this phenomenon has been widely observed, there is still no quantitative theory explaining the length of the delay between memorization and generalization. Prior work has noted that weight decay plays an important role, but no result derives tight bounds for the delay or explains its scaling behavior. We present a first-principles theory showing that grokking arises from a norm-driven representational phase transition in regularized training dynamics. Training first converges to a high-norm memorization solution and only later contracts toward a lower-norm structured representation that generalizes. Our main result establishes a scaling law for the delay: T_grok - T_mem = Theta((1 / gamma_eff) * log(||theta_mem||^2 / ||theta_post||^2)), where gamma_eff is the effective contraction rate of the optimizer (gamma_eff = eta * lambda for SGD and gamma_eff >= eta * lambda for AdamW). The upper bound follows from a discrete Lyapunov contraction argument, and the matching lower bound arises from dynamical constraints of regularized first-order optimization. Across 293 training runs spanning modular addition, modular multiplication, and sparse parity tasks, we confirm three predictions: inverse scaling with weight decay, inverse scaling with learning rate, and logarithmic dependence on the norm ratio (R^2 > 0.97). We further find that grokking requires an optimizer that can decouple memorization from contraction: SGD fails under hyperparameters where AdamW reliably groks. These results show that grokking is a predictable consequence of norm separation between competing interpolating representations and provide the first quantitative scaling law for the delay of grokking.

Why Grokking Takes So Long: A First-Principles Theory of Representational Phase Transitions

Abstract

Grokking is the sudden generalization that appears long after a model has perfectly memorized its training data. Although this phenomenon has been widely observed, there is still no quantitative theory explaining the length of the delay between memorization and generalization. Prior work has noted that weight decay plays an important role, but no result derives tight bounds for the delay or explains its scaling behavior. We present a first-principles theory showing that grokking arises from a norm-driven representational phase transition in regularized training dynamics. Training first converges to a high-norm memorization solution and only later contracts toward a lower-norm structured representation that generalizes. Our main result establishes a scaling law for the delay: T_grok - T_mem = Theta((1 / gamma_eff) * log(||theta_mem||^2 / ||theta_post||^2)), where gamma_eff is the effective contraction rate of the optimizer (gamma_eff = eta * lambda for SGD and gamma_eff >= eta * lambda for AdamW). The upper bound follows from a discrete Lyapunov contraction argument, and the matching lower bound arises from dynamical constraints of regularized first-order optimization. Across 293 training runs spanning modular addition, modular multiplication, and sparse parity tasks, we confirm three predictions: inverse scaling with weight decay, inverse scaling with learning rate, and logarithmic dependence on the norm ratio (R^2 > 0.97). We further find that grokking requires an optimizer that can decouple memorization from contraction: SGD fails under hyperparameters where AdamW reliably groks. These results show that grokking is a predictable consequence of norm separation between competing interpolating representations and provide the first quantitative scaling law for the delay of grokking.
Paper Structure (72 sections, 21 theorems, 51 equations, 6 figures, 11 tables)

This paper contains 72 sections, 21 theorems, 51 equations, 6 figures, 11 tables.

Key Result

Lemma 2.1

Assume the network function $f_\theta(x)$ is twice differentiable in $\theta$. Let $\theta^*\in\mathcal{M}_{\mathrm{train}}$ be an interpolation point. Then for any $\theta$ in a convex neighbourhood of $\theta^*$, where the remainder satisfies

Figures (6)

  • Figure 1: Conceptual overview of the Norm-Separation Delay Law. After memorisation ($T_{\mathrm{mem}}$), weight decay contracts parameter norms exponentially from the high-norm memorisation region toward the low-norm Fourier manifold. The grokking delay is the time required for this exponential contraction to traverse the norm gap $\log(\|\theta_{\mathrm{mem}}\|^2/\|\theta_{\mathrm{post}}\|^2)$. Generalisation ($T_{\mathrm{grok}}$) occurs once parameters enter the Fourier region and the validation gap becomes detectable.
  • Figure 2: Lyapunov escape validation (real data). (a) Fitted contraction rates $\rho$ across 10 seeds; all exceed $1-2\eta\lambda=0.998$ (green), confirming AdamW amplification. (b) Distribution of grokking times across 10 seeds (mean $1840\pm215$ steps). (c) Norm separation: $V_{\text{mem}}\approx 3900$ vs $V_{\text{post}}\approx 300$ across seeds.
  • Figure 3: Scaling laws. (a) $T_{\text{grok}}$ vs $1/\lambda$ in Regime II ($R^2=0.97$). (b) $T_{\text{grok}}$ vs $1/\eta$ ($R^2=0.92$). (c) Delay vs log norm ratio across 7 moduli ($r=0.91$).
  • Figure 4: Cross-task generalization. (a) Modular multiplication: all 20 runs grok across 4 moduli. (b) Exponential fit quality ($R^2$) for all 20 multiplication runs; all exceed 0.988. (c) Sparse parity: $V_{\text{final}}>V_{\text{mem}}$ in all 15 runs---no norm separation, no grokking.
  • Figure 5: AdamW contraction analysis. (a) Fitted contraction rates across 10 seeds; all exceed the theoretical $1-2\eta\lambda$ by $\sim$40%. (b) Grokking success: AdamW groks in 5/5 seeds; SGD fails entirely. (c) $R^2$ of exponential fits across all seeds (mean $R^2=0.9990$).
  • ...and 1 more figures

Theorems & Definitions (42)

  • Lemma 2.1: Second-Order Taylor Remainder Bound
  • proof
  • Lemma 2.2: NTK Stability of the Hessian
  • proof
  • Corollary 2.3: Dominance of the Linear Term During Escape
  • proof
  • Lemma 2.4: Minimal-Norm Interpolant Lower Bound
  • proof
  • Lemma 3.1: Non-Stationarity of Memorization
  • proof
  • ...and 32 more