To Grok Grokking: Provable Grokking in Ridge Regression

AI summary

Overview

Research area: Machine learning theory — specifically the theory of generalization and optimization dynamics, with a focus on the "grokking" phenomenon (delayed onset of generalization long after overfitting).

Technical level: Advanced. The paper is a rigorous theory contribution built on linear algebra, convex optimization convergence rates, and concentration inequalities, accompanied by supporting simulations.

Scope in one sentence: The paper proves end-to-end grokking guarantees for over-parameterized ridge regression trained with gradient descent and weight decay, and derives explicit quantitative bounds on how long the generalization delay ("grokking time") lasts as a function of the training hyperparameters.

What This Paper Is About

Grokking is the counterintuitive observation that a model can fit its training data perfectly while generalizing no better than chance, and only much later — often after many additional training steps — suddenly begin to generalize well. The paper asks whether this phenomenon can be proven to occur, from start to finish, in the simplest possible setting: an over-parameterized linear (ridge) regression model trained with gradient descent and an l2 penalty. The goal is to show that grokking's three stages (early overfitting, prolonged poor generalization, eventual good generalization) all follow from the training dynamics, and to quantify exactly how the delay depends on hyperparameters such as the weight decay, sample size, feature dimensionality, and initialization scale.

Key Contributions

1. The first end-to-end provable grokking result. The authors prove Theorem 4.2 for learning any realizable teacher function N*(x) = ⟨θ*, φ(x)⟩ with a student N(x; θ) = ⟨θ, φ(x)⟩, decomposed into three separate theorems: fast training-error convergence (Theorem 4.4), a lower bound on the generalization error implying long-term overfitting (Theorem 4.5), and eventual convergence to a solution with small generalization error (Theorem 4.6). The authors state this is the first end-to-end grokking result, in contrast to prior work on grokking in neural networks or linear models.

2. Explicit hyperparameter conditions and a quantitative lower bound on grokking time. The paper states sufficient conditions on the sample size n (Equation 3), feature dimensionality m (Equation 4), and weight decay λ (Equation 5) for grokking to occur, and gives a rigorous upper bound on t1 (Equation 6, the last step at which training error is still above ε) and lower bound on t2 (Equation 7, the first step at which generalization error drops below c). To the authors' knowledge, these are the first rigorous quantitative bounds on the generalization delay — the "grokking time" (t2 − t1) — in terms of training hyperparameters.

3. A principled account of how each hyperparameter controls grokking. The paper analyzes the functional dependence of t1 and t2 on the weight decay λ, the sample size n, the feature dimensionality m, and the initialization scale ν², showing that grokking can be amplified or eliminated through tuning rather than through architectural or algorithmic changes.

4. Empirical verification, including on nonlinear networks. Section 5 supports the theory with simulations in ridge regression, and the authors additionally train two-layer ReLU neural networks (both layers trained) and report that the dependence of grokking time on hyperparameters qualitatively matches the provable predictions from the linear case.

Main Findings

• Three stages are provable. Under the stated conditions (Gaussian initialization θ⁽⁰⁾ ~ N(0, ν²I_m), step size η < 1/(λ + b²), bounded features ‖φ(x)‖₂ ≤ b), the training error falls below ε by step t1 given in Equation 6, the generalization error stays at or above a constant c ≥ ε until at least step t2 given in Equation 7, and L(θ⁽ᵗ⁾) ≤ ε for all sufficiently large t. Grokking is established when c ≥ ε.

• Two convergence rates, one cause. In the zero-teacher warmup (Theorem 4.1), training loss decays at a rate governed by (1 − (1/n)ηλ_min⁺(ΦᵀΦ) − ηλ)^(2t) while the generalization loss decays at (1 − ηλ)^(2t). A small λ barely affects the training rate, which is dominated by the data-dependent term, but slows the generalization rate arbitrarily. The weight vector norm obeys ‖θ⁽ᵗ⁾‖₂² ≤ (1 − ηλ)^(2t)‖θ⁽⁰⁾‖₂².

• Subspace explanation for the delay. With m ≫ n, gradient descent with small weight decay essentially only updates the projection θ_∥ of the parameters onto the data-spanning subspace, fitting the training data quickly; the component θ_⊥ in the complementary subspace stays near its initialization and is only reduced by long-term weight decay at a negligible rate (1 − ηλ)^t. The imbalanced trajectory prevents timely generalization until weight decay has reduced model complexity enough for uniform convergence.

• Weight decay controls the grokking time. For small λ, decreasing λ increases t2 with t2 ∝ 1/λ when other hyperparameters are fixed, while λ has no effect on the bound for t1. Consequently (t2 − t1) → ∞ as λ → 0. The authors note that for large λ their Equation 6 upper bound is not tight in general, because a tight bound should include the term ηλ_min⁺(ΦᵀΦ) + ηλ in the denominator.

• Sample size and dimensionality effects. Clean dependencies of t1 on n and m cannot be deduced, because no quantitative bound is known for how λ_min⁺(ΦᵀΦ) depends on m and n — even for simple distributions such as φ(x) ~ N(0, (1/m)I_m), where only the asymptotic Marchenko–Pastur law is known: as m, n → ∞ with m/n → γ ∈ (0, ∞), λ_min⁺(ΦᵀΦ) → γ⁻¹(1 − √γ)². The authors show empirically that t1's behavior as a function of n and m strongly accords with Equation 6, and they derive cleaner bounds for t2 under a spherical Gaussian feature assumption, presented in Equation 8.

• Initialization scale amplifies grokking. Increasing ν² increases both t1 and t2, with t1, t2 ∝ ln(ν²). More precisely, for sufficiently small λ, t1 ≤ c1 ln(ν²) + a1 and t2 ≥ c2 ln(ν²) + a2 with c2 > c1 > 0, so increasing ν² increases (t2 − t1) and amplifies grokking.

• Nonlinear behavior matches the linear predictions. Figure 1 compares training and test squared losses using GD with weight decay over 50 independent runs (independent datasets and student initializations) for ridge regression learning the zero teacher and for a two-layer ReLU network trained on both layers learning the zero teacher, plotted on a logarithmic x-axis as is standard in the grokking literature.

• Grokking is a training-conditions artifact. The authors conclude that grokking is not an inherent failure mode of deep learning but a consequence of specific training conditions, and therefore does not require fundamental changes to model architecture or learning algorithm to avoid.

Methodology in Plain English

The authors set up a teacher–student problem. A fixed feature map φ(x) maps inputs to R^m; the teacher labels data exactly as ⟨θ*, φ(x)⟩, and the student is a linear model ⟨θ, φ(x)⟩ with θ ∈ R^m trained to minimize an MSE plus (λ/2)‖θ‖₂², using vanilla gradient descent with step size η from a random Gaussian initialization.

Because the teacher is realizable and the model is linear, the training trajectory can be written in closed form. The analysis splits the parameter vector into the part that lives in the span of the sampled feature vectors (where the data provide a strong learning signal) and the orthogonal complement (where the data provide no signal and only weight decay acts). The authors then bound how fast the training error shrinks, how slowly the orthogonal component shrinks, and when the accumulated weight decay is finally enough to make the population error small. Two quantities anchor the story: the smallest positive eigenvalue λ_min⁺(ΦᵀΦ) of the empirical feature Gram matrix, which sets the training convergence rate, and the weight decay λ and population covariance eigenvalues λ_min(Σ), which set the generalization rate. The paper states t1 and t2 in terms of these quantities, then reads off scaling laws by inspecting the resulting formulas.

The theoretical predictions are then checked in simulations that vary each hyperparameter and measure the observed grokking time, including experiments on two-layer ReLU networks to test whether the same scaling behavior carries over beyond the linear setting. Proofs and additional experiments are deferred to the appendix.

Why This Matters

Impact on research. Prior theoretical work attributed grokking largely to a transition from the lazy (kernel) to the rich (feature-learning) regime, or analyzed related two-phase behavior without proving that the unregularized solution fails to generalize and the regularized solution succeeds. This paper shows that neither deep networks nor nonlinear feature learning are necessary: grokking arises in plain ridge regression, driven by the interplay between fast data-fitting and slow weight-decay-driven complexity reduction. It also provides the first quantitative, provable bound on the generalization delay in terms of hyperparameters, turning a qualitative curiosity into a tunable property. The authors note their bounds are sharper end-to-end guarantees than the general loss-plus-regularizer framework of Tikeng Notsawo et al. (2025), which places the regularization-driven phase on a timescale proportional to 1/(ηλ).

Real-world applications (implications of the results, not experiments run in the paper):

• Hyperparameter tuning for regularized models: the t2 ∝ 1/λ scaling gives a concrete lever for controlling how long a model spends in a poorly generalizing state.
• Diagnosing delayed generalization in training pipelines, where long plateaus in validation performance might otherwise be mistaken for a stalled run.
• Understanding when reduced model complexity (via weight decay) is sufficient to fix generalization, versus when architectural changes are genuinely needed.
• Benchmarking and reasoning about grokking-like behavior in settings beyond neural networks, such as other regularized linear or kernel predictors.

Industry relevance. The central practical message is that grokking may be avoidable through hyperparameter choices — principally the weight decay, but also the initialization scale, sample size, and model dimensionality — rather than through redesigning architectures or learning algorithms. For teams spending compute on long training runs that appear to be overfitting, this reframes the decision as a tuning question.

Future Directions

• Tighten the bound for large weight decay. The authors explicitly flag that their Equation 6 upper bound on t1 is not tight in general for large λ, because a tight bound should incorporate ηλ_min⁺(ΦᵀΦ) + ηλ in the denominator. Closing this gap is a natural next step.
• Control λ_min⁺(ΦᵀΦ) as a function of n and m. The paper notes that no quantitative bound is currently known for how the smallest positive eigenvalue of the empirical Gram matrix depends on the sample size and dimensionality, even for simple Gaussian features — only the asymptotic Marchenko–Pastur limit with λ_min⁺(ΦᵀΦ) → γ⁻¹(1 − √γ)². Sharpening this would yield cleaner dependencies of t1 on n and m.
• Extend beyond ridge regression and squared loss. The analysis is specific to linear/ridge regression. Extending end-to-end provable grokking to classification settings, other regularizers, or non-realizable teachers remains open; the authors also note that Theorem 4.2 covers teachers realizable by any fixed feature map but not arbitrary non-linear teachers.
• Extend the analysis to related optimization schemes. The authors state their analysis should also hold for gradient descent with a decaying step size and for gradient flow, which points to a systematic treatment of how the optimization scheme itself shapes the grokking time. The paper's Section 6 concludes with further proposed directions; the specific list is not included in the excerpt available here.

Target Audience

This paper is best suited to machine learning theory researchers and graduate students working on optimization dynamics, implicit regularization, and generalization bounds, particularly those interested in grokking and the lazy-versus-rich regime literature. It is also useful for practitioners who train heavily over-parameterized models with weight decay and want a principled understanding of why validation performance can lag far behind training performance — though the paper's results are theoretical and assume linear models with a fixed feature map, so readers without a background in convergence analysis and linear algebra will find the theorem statements dense.

Authors’ abstract

We study grokking, the onset of generalization long after overfitting, in a classical ridge regression setting. We prove end-to-end grokking results for learning over-parameterized linear regression models using gradient descent with weight decay. Specifically, we prove that the following stages occur: (i) the model overfits the training data early during training; (ii) poor generalization persists long after overfitting has manifested; and (iii) the generalization error eventually becomes arbitrarily small. Moreover, we show, both theoretically and empirically, that grokking can be amplified or eliminated in a principled manner through proper hyperparameter tuning. To the best of our knowledge, these are the first rigorous quantitative bounds on the generalization delay (which we refer to as the "grokking time") in terms of training hyperparameters. Lastly, going beyond the linear setting, we empirically demonstrate that our quantitative bounds also capture the behavior of grokking on non-linear neural networks. Our results suggest that grokking is not an inherent failure mode of deep learning, but rather a consequence of specific training conditions, and thus does not require fundamental changes to the model architecture or learning algorithm to avoid.