Feature Learning Dynamics under Grokking in a Sparse Parity Task

In this paper, we analyze the phenomenon of Grokking in deep neural networks by examiningthe evolution of the Neural Tangent Kernel (NTK) matrix in a sparse parity task with the aim ofstudying feature learning during training. We show that during the initial overfitting phase, theNTK’s eigenfunctions provide spurious features that do not generalize. On the other hand, at a laterstage the NTK’s top eigenfunctions well-align with the predictive input features, thus resulting thedelayed generalization typically observed in Grokking. Our work can be viewed as a mechanisticinterpretability framework for evaluating feature learning during training using the NTK.