Grokfast: Accelerated Grokking by Amplifying Slow Gradients

One puzzling artifact in machine learning, dubbed grokking, refers to the case where a model exhibits delayed generalization long after it reaches near-perfect memorization of the training data. Focusing on the long delay itself on behalf of machine learning practitioners, our primary goal is to accelerate the generalization of a model under the grokking phenomenon. By treating the gradient stream of each parameter over training iterations as a discrete random signal over time, we can spectrally decompose the parameter update trajectories under gradient descent into two components: a fast-varying component associated with rapid memorization, and a slow-varying component empirically aligned with delayed generalization. This frequency-domain analysis allows us to accelerate the grokking phenomenon more than $\times 50$ with only a few lines of code that amplifies the slow-varying components of the gradients via low-pass filtering. The experiments show that our algorithm, dubbed Grokfast, applies to diverse tasks involving images, languages, and graphs across various architectures including Transformers, MLPs, RNNs, and Graph-ConvNets, enabling the practical availability of this peculiar artifact of sudden generalization. Methodologically, Grokfast introduces a two-timescale augmentation to existing first-order optimizers: the original stochastic gradient is retained, while its temporally coherent low-frequency component is added as a residual term. By framing optimizer design as filtering over the gradient stream, this view bridges machine learning optimization and classical signal processing, suggesting a new class of optimizer augmentations grounded in frequency-domain filter design.