Omnigrok: Grokking Beyond Algorithmic Data
This work addresses the fundamental understanding of grokking for researchers in machine learning, providing insights that could improve training dynamics and generalization across various domains.
The paper tackled the problem of understanding grokking, a phenomenon where neural networks generalize long after overfitting on algorithmic data, by analyzing loss landscapes and identifying a mismatch between training and test losses as the cause. They successfully induced grokking on image, language, and molecule tasks and eliminated it for algorithmic datasets, attributing its dramatic nature to representation learning.
Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive. We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the mismatch between training and test losses as the cause for grokking. We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively. This simple mechanism can nicely explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc. Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules. In the reverse direction, we are able to eliminate grokking for algorithmic datasets. We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.