To Grok Grokking: Provable Grokking in Ridge Regression

TL;DR

This work studies grokking, the delayed emergence of generalization after overfitting, in a classical ridge regression setting and proves an end-to-end grokking guarantee for over-parameterized linear models trained with gradient descent and weight decay. It provides explicit bounds on the grokking time in terms of hyperparameters such as the sample size $n$, feature dimension $m$, learn-rate $\eta$, decay $\lambda$, and initialization $\nu$, and shows how hyperparameter tuning can amplify or eliminate the delay. The authors validate the theory with experiments on linear/ridge and non-linear networks, including random-feature and two-layer ReLU architectures, finding qualitative and quantitative alignment with the predictions. The results suggest grokking arises from training conditions rather than fundamental architectural limitations, implying grokking can be mitigated without major model changes.

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.

Figures (4)

• Figure 1: Comparing training and test squared losses using GD with weight decay over 50 independent runs (with independent datasets and student initializations). Left: training a ridge regression model to learn the zero teacher; Right: training a two-layer ReLU neural network (training both layers) to learn the zero teacher. We scale the x-axis logarithmically, which is commonly adopted for plotting the grokking phenomenon power2022grokkinglyu2023dichotomyxu2023benignmohamadi2024you. See Section \ref{['sec:experiments']} for further details of the experimental setup.
• Figure 2: Plots for the effects of hyperparameters on grokking in ridge regression. Left Upper: decreasing weight decay extends the generalization delay ($t_{2}\propto 1/\lambda$); Right Upper: decreasing sample size amplifies grokking by speeding up the convergence of the training loss (affects $t_{1}$); Left Lower: increasing the feature dimension has little effect on $t_{1}$ and $t_{2}$; Right Lower: increasing the initialization scale increases $t_{1}$ and $t_{2}$ simultaneously at logarithmic rates ($t_{1},t_{2}\propto\ln(\nu^{2})$).
• Figure 3: Plots of training and test losses of a two-layer random ReLU features network. Dashed/solid lines indicate train/test loss respectively. Left Upper: using smaller weight decay amplifies the grokking time by delaying generalization (increases $t_{2}$); Right Upper: having smaller sample size amplifies grokking by speeding up the training convergence (decreases $t_{1}$); Left Lower: increasing the student's width does not significantly prolong grokking; Left Lower: increasing the initialization scale does not significantly prolong grokking, but instead widens the gap between the generalization and training losses during the overfitting stage.
• Figure 4: Plots for the effect of the hyperparameters on grokking when training a two-layer ReLU network with the zero teacher.

Theorems & Definitions (36)

• Theorem 4.1: End-to-end provable grokking for the zero teacher
• Theorem 4.2: End-to-end provable grokking for realizable ridge regression
• Remark 4.3: on $\lambda_{\mathrm{min}}^{+}(\boldsymbol{\Phi}^{\top}\boldsymbol{\Phi})$ and its dependence on $n$ and $m$
• Theorem 4.4: Training loss convergence
• Theorem 4.5: Poor generalization when overfitting
• Theorem 4.6: Generalization
• Theorem 1.1: Theorem \ref{['thm:e2e-provable-grokking-rf-zero-teacher']} restated
• proof : Proof of Theorem \ref{['thm:e2e-provable-grokking-rf-zero-teacher-restated']}
• Theorem 1.2: Training loss convergence
• Remark 1.3