Blog: Grokking: When Models defy how we think they learn

A few months ago, in a conversation with some colleagues, was when I first heard the term “grokking”. For a brief moment I thought it had something to do with the AI model Grok, but learned that it is actually a well-known phenomenon in the area of machine learning so I took some time to read about it. 

Most of the time, when we are training a neural network model, we expect that the test accuracy to follow up the after training accuracy after some few hundred steps. And when it happens that the test accuracies are not improving over a few more hundred steps we give up. The idea of Grokking suggests that maybe the model is not done yet so give some more time.

Grokking is the phenomenon when a model memorizes its training data quickly, sits at chance-level performance on test data for a long period, and then with no change in the setup, just more steps of training and suddenly it generalizes. Figure 1 depicts what grokking looks like. We observe that the training accuracy shoots up to 95% within the first few hundred steps and stays there. The test accuracy, meanwhile, stays less than 20% of the accuracy tens of thousands of steps latter. Then, quite abruptly, it climbs to match training accuracy.

A simple way to differentiate between memorization and generalization, is that given a training data of string of bits, like we see in Figure 2a, each is mapped to a label 0 or 1, where the label actually only depends on a small slice of each string (seen in Figure 2b). The mapping rule is 1 when the first three bits contain an odd number of 1s else 0. Every other bit is noise as far as the true rule is concerned.

A model that is memorizing will fit every bit, relevant or not, because from the training data alone there's no way to distinguish signal from noise.  On the other hand a model that has actually found the underlying rule looks only at the first three bits and ignores the rest. Both models may get high accuracies on train data, but only the generalized one perform well on the test data.

What grokking suggests is that under continued optimization, particularly under regularization, the optimizer eventually finds a simpler solution that happens to generalize. The memorizing solution and the generalizing solution can both fit the training data perfectly, but they are not the same solution, and getting from one to the other takes time.

Now I am sure you are asking how can I get my model to grokk? The following are few levers that show up as a useful for encouraging grokking to happen.

• Regularize the model deliberately: Regularization strength appears to determine whether and when a generalizing solution gets found. Specific examples are weight decay, and dropout.
• Data size matters:  A sizable amount of data about quarter of the full training space (in algorithmic tasks like modular arithmetic) tends to be enough for grokking to be observed.

Notably, we can implement all the above but it is not guaranteed that our model will reach grokking.

Some practical takeaways:

• Do not stop just because training accuracy maxed out.
• Watch test loss, not just test accuracy, the loss tends to move before the accuracy jump is visible.
• Treat regularization as main lever to reach generalization.

One caveat worth keeping in mind is that most of the evidence for grokking comes from  small, datasets, so I would not read it as a general instruction for model training. 

Turns out that understanding grokking took some grokking. What I find interesting is that when the setup is right, sometimes you just have to allow it to train a little longer. 

Some References

• Power, A. et al. (2022). Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv:2201.02177
• PAIR Explorables: Grokking. https://pair.withgoogle.com/explorables/grokking/

Written by: Theophilus Aidoo (https://theoaid.com/)