Technical Summary: To Grok Grokking: Provable Grokking in Ridge Regression

Problem Statement

The paper investigates the phenomenon of grokking, defined as the onset of generalization long after a model has achieved perfect training performance (overfitting). While previously observed primarily in complex, non-linear neural networks and algorithmic datasets, the authors aim to establish rigorous, end-to-end theoretical guarantees for this phenomenon in a classical, linear setting. Specifically, the paper addresses the gap in existing literature where prior works either lack formal proofs of the "delay" between overfitting and generalization or rely on specific, non-generalizable problem structures (e.g., modular addition). The core question is whether grokking can be provably demonstrated in over-parameterized ridge regression using gradient descent (GD) with weight decay, and if so, how hyperparameters quantitatively control the "grokking time."

Methodology

The authors formulate the problem within a teacher-student framework:

• Setting: Over-parameterized linear regression where the student model $N(x; \theta) = \langle \theta, \phi(x) \rangle$ learns a realizable teacher function $N^*(x) = \langle \theta^*, \phi(x) \rangle$.
• Optimization: The student is trained using Gradient Descent (GD) with a fixed step size $\eta$ and an $\ell_2$ regularization term (weight decay) $\lambda$. The objective is the regularized mean squared error:
  $$L_n(\theta; \lambda) = \frac{1}{2n}\sum_{i=1}^n (N(x_i; \theta) - N^*(x_i))^2 + \frac{\lambda}{2}\|\theta\|_2^2$$
• Analysis Strategy: The authors decompose the parameter vector $\theta$ into two orthogonal components relative to the data span:
  1. $\theta_{\parallel}$: The component in the row space of the feature matrix $\Phi$ (data-spanning subspace).
  2. $\theta_{\perp}$: The component in the null space of $\Phi$ (complementary subspace).
• Theoretical Approach: The proof analyzes the convergence rates of the training loss (empirical error) and the generalization loss (population error) separately. The core mechanism relies on showing that $\theta_{\parallel}$ converges rapidly to fit the training data, while $\theta_{\perp}$ decays slowly due to weight decay, causing a prolonged period of poor generalization.

Key Contributions

1. First End-to-End Provable Grokking: The paper provides the first rigorous quantitative bounds demonstrating the three distinct stages of grokking in a linear setting:
   • (i) Rapid overfitting of training data.
   • (ii) A persistent delay where generalization error remains high despite zero training error.
   • (iii) Eventual convergence to a model with arbitrarily small generalization error.
2. Quantitative Bounds on Grokking Time: The authors derive explicit lower bounds for the "grokking time" ($t_2 - t_1$), where $t_1$ is the time to overfit and $t_2$ is the time to generalize. These bounds are expressed in terms of model hyperparameters:
   • Weight Decay ($\lambda$): The grokking time scales inversely with $\lambda$ ($t_2 \propto 1/\lambda$). Smaller $\lambda$ significantly delays generalization.
   • Initialization Scale ($\nu^2$): Increasing the initialization scale amplifies grokking logarithmically ($t_1, t_2 \propto \ln(\nu^2)$).
   • Over-parameterization: The phenomenon is shown to require a high-dimensional setting ($m \gg n$).
3. Hyperparameter Control: The paper establishes sufficient conditions on hyperparameters (sample size $n$, feature dimension $m$, weight decay $\lambda$, and initialization $\nu^2$) to induce or eliminate grokking.
4. Empirical Validation on Non-Linear Models: While the theory is derived for linear regression, the authors empirically demonstrate that the quantitative dependencies on hyperparameters (e.g., the effect of $\lambda$ on delay) qualitatively match the behavior observed in non-linear two-layer ReLU neural networks.

Key Results

• Convergence Rates: The training error converges at a rate dominated by the smallest non-zero eigenvalue of the empirical feature covariance ($\lambda^+_{\min}(\Phi^\top \Phi)$), which is fast. In contrast, the generalization error is lower-bounded by a term decaying at the rate of $(1 - \eta\lambda)^t$, which is arbitrarily slow for small $\lambda$.
• The Grokking Mechanism: The delay arises because GD with small weight decay effectively updates only the projection of weights onto the data span to minimize training loss, while the orthogonal component (which contributes to generalization error) decays only via the slow weight decay term.
• Theorems:
  • Theorem 4.1: Establishes grokking for the zero teacher, providing matching upper and lower bounds on loss convergence.
  • Theorem 4.2: Generalizes the result to any realizable teacher, providing explicit bounds for $t_1$ and $t_2$ (Equations 6 and 7).
  • Theorem 4.4 & 4.5: Prove the fast convergence of training loss and the persistence of poor generalization, respectively.
  • Theorem 4.6: Proves that GD eventually converges to a global minimum with good generalization guarantees.

Significance and Claims

The paper claims to provide a foundational theoretical framework for understanding grokking. By isolating the phenomenon in the simple, transparent setting of linear regression, the authors argue that:

• Grokking is not an inherent failure mode of deep learning architectures or non-linearities, but rather a consequence of specific training conditions (specifically, the interplay between over-parameterization and small weight decay).
• The phenomenon does not require fundamental changes to model architecture or learning algorithms to avoid; it can be controlled via hyperparameter tuning.
• The results offer the first rigorous quantitative bounds on the delay of generalization, moving beyond empirical observations or heuristic explanations (such as "lazy-to-rich" regime transitions) found in prior literature.
• The findings suggest that the "test error plateau" often seen in classification (an artifact of discrete metrics) also exists in regression when evaluated via threshold-based accuracy metrics, confirming the universality of the phenomenon across loss types.

The authors maintain modesty regarding the scope, noting that while their linear analysis provides rigorous proofs, extending these results to non-realizable settings or proving the "lazy-to-rich" transition in deep neural networks remains an open challenge requiring new theoretical tools.