To Grok Grokking: Provable Grokking in Ridge Regression

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.

Grokking — the phenomenon where generalization starts improving long after a model has overfit the training data — has emerged as a counterintuitive behavior in deep learning. Although it has been extensively studied empirically, only a few prior works have established rigorous theoretical guarantees for grokking. In this work, we take a step toward a rigorous, provable understanding of this phenomenon. While grokking was first identified in complex deep learning settings, our paper presents the surprising result that grokking can be proved end-to-end in a simple over-parameterized linear regression model trained by gradient descent. Moreover, our work is the first to derive rigorous quantitative bounds on the delay before generalization, expressed in terms of training hyperparameters. Our theoretical findings are supported by extensive empirical simulations. Our results suggest that grokking is not an inherent failure mode of deep learning, but rather a consequence of specific training conditions. As a result, avoiding grokking may not require fundamental changes to the model architecture or learning algorithm.