To Grok Grokking: Provable Grokking in Ridge Regression
This work addresses the phenomenon of grokking in machine learning, providing rigorous theoretical insights for researchers, but it is incremental as it builds on existing concepts in a classical setting.
The paper tackles the problem of grokking, where generalization occurs long after overfitting, by proving end-to-end results for ridge regression with gradient descent and weight decay, showing that generalization error eventually becomes arbitrarily small and that grokking can be amplified or eliminated through hyperparameter tuning.
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.