Grokking as the transition from lazy to rich training dynamics

Abstract

We study the recently discovered “grokking” phenomenon in deep learning [Power et al., 2022], where neural networks generalize to unseen data abruptly, long after memorizing their training data set. We propose this phenomenon arises when a neural network stops behaving like a linear model in its parameters. To illustrate this, we study the simple setting of vanilla gradient descent on a polynomial regression problem with a two layer neural network which exhibits grokking without regularization. We identify sufficient statistics for the test loss of such a network, and tracking these over training reveals that grokking arises in this setting when the network first attempts to fit a kernel regression solution with its initial features, followed by late-time feature learning where a generalizing solution is identified after train loss is already low. We find that the key determinants of grokking are the rate of feature learning—which can be controlled precisely by parameters that scale the network output—and the alignment of the initial features with the target function. We conclude with evidence that this transition from lazy (linear model) to rich training (feature learning) characterizes grokking in more general settings, like on MNIST, one-layer Transformers, and student-teacher networks.