One-sentence Summary

The authors establish the first rigorous quantitative bounds on grokking time for over-parameterized linear regression with gradient descent and weight decay, prove that hyperparameter tuning can amplify or eliminate grokking, and empirically demonstrate that these bounds capture grokking behavior in non-linear neural networks, suggesting that grokking is not an inherent failure mode of deep learning but a controllable consequence of specific training conditions.

Key Contributions

• The paper establishes the first rigorous quantitative bounds on grokking time for over-parameterized linear regression trained with gradient descent and weight decay, proving a progression from overfitting through persistent poor generalization to arbitrarily small test error.
• Both theoretical analysis and experiments show that grokking can be systematically amplified or eliminated by tuning hyperparameters, such as the weight decay coefficient, providing a principled control mechanism.
• Beyond the linear setting, the quantitative bounds empirically capture grokking behavior on non-linear neural networks, supporting the conclusion that grokking results from specific training conditions rather than a fundamental deficiency of deep learning.

Introduction

The authors study grokking, a counterintuitive phenomenon where test performance improves long after a model perfectly fits the training data. Prior theoretical analyses have attributed grokking to transitions between lazy and rich training regimes or to the role of weight decay, but none provided a rigorous end-to-end guarantee that poor generalization persists after overfitting before eventually resolving. The authors fill this gap by analyzing overparameterized ridge regression with a teacher-student setup, showing for the first time that even a linear model trained with gradient descent and weight decay can exhibit provable grokking. They establish quantitative bounds on training error convergence, a slower generalization error decay that demonstrates prolonged overfitting, and eventual convergence to a global minimum with good test performance, all fully controllable through hyperparameters like the weight decay coefficient.

Method

Theauthors investigate the phenomenon of grokking within the context of ridge regression. They consider a fundamental regression problem where a student linear regression model $N(\boldsymbol{x}; \boldsymbol{\theta}) = \langle \boldsymbol{\theta}, \phi(\boldsymbol{x}) \rangle$N(x;θ)=⟨θ,ϕ(x)⟩ is trained to learn a realizable teacher function $N^*(\boldsymbol{x}) = \langle \boldsymbol{\theta}^*, \phi(\boldsymbol{x}) \rangle$N∗(x)=⟨θ∗,ϕ(x)⟩. Here, $\phi(\boldsymbol{x})$ϕ(x) is a fixed feature map and $\boldsymbol{\theta}$θ represents the trainable parameters.

The training process utilizes ridge regression, which minimizes the empirical mean squared loss augmented with an $\ell_2$ℓ2​-regularization term. The training objective is defined as:

$L_n(\boldsymbol{\theta}; \lambda) = \frac{1}{2n} \sum_{i=1}^n (N(\boldsymbol{x}_i; \boldsymbol{\theta}) - N^*(\boldsymbol{x}_i))^2 + \frac{\lambda}{2} \|\boldsymbol{\theta}\|_2^2$Ln​(θ;λ)=2n1​∑i=1n​(N(xi​;θ)−N∗(xi​))2+2λ​∥θ∥22​

where $\lambda > 0$λ>0 is the weight decay parameter. The model is optimized using vanilla Gradient Descent (GD) with a fixed step size $\eta$η, updating the parameters via $\boldsymbol{\theta}^{(t+1)} = \boldsymbol{\theta}^{(t)} - \eta \nabla_{\boldsymbol{\theta}} L_n(\boldsymbol{\theta}^{(t)}; \lambda)$θ(t+1)=θ(t)−η∇θ​Ln​(θ(t);λ).

To formalize grokking, the authors define the training loss convergence time $t_1$t1​ and the generalization loss convergence time $t_2$t2​. Grokking occurs when there is a significant delay between these two events ($t_2 \gg t_1$t2​≫t1​). In the over-parameterized regime ($m \gg n$m≫n), the GD optimization process effectively updates the projection of the weight vector onto the data-spanning subspace to fit the training data quickly. However, the component in the complementary subspace converges at a much slower rate dictated by the weight decay, leading to harmful overfitting before eventual generalization.

The authors provide quantitative bounds for the grokking time $t_2 - t_1$t2​−t1​ and analyze how different training hyperparameters influence this delay.

As shown in the figure below, the theoretical bounds closely match empirical results regarding hyperparameter dependencies. Specifically, decreasing the weight decay $\lambda$λ extends the generalization delay, as $t_2$t2​ is inversely proportional to $\lambda$λ. Decreasing the sample size $n$n amplifies grokking by speeding up the convergence of the training loss. Increasing the feature dimension $m$m has little effect on $t_1$t1​ and $t_2$t2​, while increasing the initialization scale $\nu^2$ν2 increases both $t_1$t1​ and $t_2$t2​ simultaneously at logarithmic rates.

Experiment

The experiments confirm the theoretical framing of grokking by demonstrating that gradient descent with weight decay can first achieve low training error while generalization remains poor, and then eventually generalize well after a delay that can be made arbitrarily long through smaller weight decay. The observed dependencies of the early and late critical time points on hyperparameters closely track the derived bounds for linear ridge regression, and this qualitative behavior persists in both random-features networks and fully trained two-layer ReLU networks, indicating that the analytical insights may hold beyond the simplified linear setting.