Most of the time when we train a model, we usually expect the model to learn the task at hand and generalize well on it. However, there are chances of overfitting where either the model complexity is too high for the amount of data at hand or a variety of other factors that prevent the model from generalizing well.

We assume that once a model starts overfitting, it will only continue to overfit even more and not return to generalizing. However, this is not always the case. Power et al. (2022) shows that a model can in fact, start to generalize very quickly after a period of overfitting.

What is Grokking?

Grokking is the phenomenon where a model starts to rapidly generalize after an arbitrary amount of overfitting depending on the task that it is being trained on. In this blog, I will show the findings of the demonstration of this phenomenon.

There are some factors to consider why Grokking occurs in general. Usually the model just overfits and doesn’t actually converge to a good solution. To make sure that we observe it, we tune this hyperparameter called weight decay in the optimizer that we use. Weight decay is a regularization technique that adds a penalty to the loss function based on the magnitude of the weights present in the model. So bigger the weights, higher the penalty and the model is forced to learn a simpler and generalizable solution.

This is important because bigger weights to certain neurons essentially means that the model is giving high priority to certain “patterns” that that neuron is looking for, which makes the model very brittle since in case it comes across a data instance which doesn’t contain that pattern but is still valid, the model can break.

Experimental Setup

To demonstrate this phenomenon, we will be training a transformer model on predicting the output to modular addition (a + b) mod 113 of two numbers. We will generate a dataset for this keeping the prime number modular value as 113. Why use prime numbers? Well, it is because if we use any non-prime number, there can be some patterns that the model can learn which isn’t modular addition. For example, if we use 32 as the modular number, the model can learn something related to the factors of it like 2, 4, 8, 16, 32. Using prime numbers minimizes the chances of this happening.

We will train the model with very few data instances to ensure that it overfits. In this demonstration, I did a 30:70 train-test split, which ensures that the model has just enough to learn the task eventually but not enough to generalize quickly at the start.

We will be using a transformer model with the following hyperparameters:

class GrokkingConfig:
    # Model architecture
    n_layers: int = 1
    d_model: int = 128
    n_heads: int = 4
    d_head: int = 32
    n_ctx: int = 3
    d_vocab: int = 114
    act_fn: str = 'relu'
    normalization_type: Optional[str] = None
    
    # Training hyperparameters
    seed: int = 999
    prime_P: int = 113
    num_epochs: int = 20000
    lr: float = 0.001
    weight_decay: float = 1.0
    batch_size_ratio: float = 0.30
    
    # Early stopping
    target_val_acc: float = 0.975

Training

We train this model using an arbitrary number of epochs until we obtain rapid generalization, we will check for this using early stopping, by stopping the training once the model obtains a validation accuracy of 97.5% or more.

This plot shows the training run of the model. We can observe that the model initially starts to overfit very quickly and then after around Epoch 7000, the validation accuracy starts to increase very rapidly after a period of overfitting. This is the grokking phenomenon.

Analysis

After training the model and observing that the model has generalized, we need to prove that the model actually has learnt the right task and hasn’t simply exploited some random spurious patterns in the data.

To do this, we will check the W_E weights of the model and check if the model has learnt the right task. To verify this, we use Fast Fourier Transform (FFT) to check on this embedding matrix.

Here we plot a graph of the each embedding dimension across the vocabulary of the model in the y-axis and the range of frequencies in the x-axis. We can observe that there are certain frequencies that the value of each dimension is high for. This means that the values of the embedding matrix dimensions are oscillating at a certain frequencies! This basically means that the model has learnt that a combination of cosine and sine waves are representative of the data.

Each number can be represented in terms of cos and sin waves using this formula,

cos(x)sin(x) 

which represents the number in terms of angles instead of only its value. This shows that the model has learnt that the value of the number oscillate after a certain point which is true in modular addition.

For example, this is a graph of the exact same model but untrained:

Here we can see that the model has not really learn any kind of specific pattern in the embedding matrix and the frequencies are all scattered out.

To show that, model’s attention heads also have learnt how to use the embeddings to actually predict the right output, we plot a graph of the attention scores against the b token of the input sequence keeping the a token fixed. This helps us visualize how the attention heads treat the input for a varying range of b values.

We can see here that the even the different attention heads of the model are also oscillating at a certain frequency! This implies that the model is in fact performing a dot product between the input tokens and which results in this pattern of varying attention scores.

You might ask, how exactly can a model compute modular addition using only dot products? Well, there is a certain trigonometric identity that this model found and is exploiting to compute modular addition.

cos(a+b) = cos(a)cos(b) - sin(a)sin(b)

But, the model does dot products which means it has to add them and not subtract this. So, the model internally converts the b value to -b and then uses this identity to compute the modular addition.

cos(a-b) = cos(a)cos(b) + sin(a)sin(b)
Result  = cos(a-(-b))
Result  = cos(a+b) 

For reference, here’s the same plot with an untrained model:

This discovery shows that the model has found a way to compute modular addition using just the dot products of the input tokens.

References

[1] Power et al. (2022) - Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. The original paper identifying the phase transition.

[2] Nanda et al. (2023) - Progress Measures for Grokking via Mechanistic Interpretability. The paper that reverse-engineered the modular addition algorithm, which I replicated here.

[3] Elhage et al. (2021) - A Mathematical Framework for Transformer Circuits. Foundational work on analyzing transformer weights.

[4] TransformerLens - The library used for all activation patching and Fourier analysis in this project.