Beyond Overfitting: Understanding "Grokking" in Neural Networks
Beyond Overfitting: Understanding "Grokking" in Neural Networks
Have you ever trained a model only to find that it perfectly memorized the training data but failed miserably on the validation set? In traditional machine learning, we call this overfitting, and the standard prescription is to stop training early.
But what if I told you that if you kept training—long after the model had overfitted—it might suddenly "click" and achieve perfect generalization?
This phenomenon is known as Grokking. In this post, we dive into the mechanics of Grokking, explore how a small Transformer can "discover" modular arithmetic from scratch, and implement a reproduction of the experiment in PyTorch.
What is Grokking?
Grokking is a surprising discovery in deep learning where a model, after a long period of overfitting, suddenly transitions to a state of perfect generalization.
Unlike typical learning curves where validation accuracy plateaus or drops as training accuracy hits 100%, grokking exhibits a "delayed" spike in validation accuracy. The model essentially moves from memorization (storing the training table in its weights) to algorithm discovery (finding the underlying mathematical rule).
The Core Intuition
The researchers treat mathematical operations as a sequence-to-sequence translation task. By feeding the model equations like $a \circ b = c$ using abstract tokens, the network is forced to ignore human-defined numerical representations. To succeed, the model must discover the algebraic structure—the relational patterns—of the operation itself.
The Architecture: A Minimalist Transformer
To observe grokking, we don't need a GPT-4 scale model. A small, encoder-only Transformer is sufficient to map a pair of tokens $[a, b]$ to a result token $c$.
The Pipeline
The process follows a specific flow: from discrete tokenization to a high-dimensional embedding space, processed through attention layers, and finally collapsed into a prediction.
Implementation: Modular Addition
The most common benchmark for grokking is Modular Addition: $(a + b) \pmod{p}$. Here is a production-ready implementation using PyTorch.
The Code
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np
import random
from typing import Tuple
class ModularAdditionDataset(Dataset):
"""Generates a dataset for the operation: (a + b) mod p."""
def __init__(self, p: int, train_fraction: float = 0.5):
self.p = p
all_pairs = [(a, b) for a in range(p) for b in range(p)]
random.shuffle(all_pairs)
split_idx = int(len(all_pairs) * train_fraction)
self.train_pairs = all_pairs[:split_idx]
self.val_pairs = all_pairs[split_idx:]
def get_split(self, split='train') -> Tuple[torch.Tensor, torch.Tensor]:
pairs = self.train_pairs if split == 'train' else self.val_pairs
inputs = torch.tensor([[p[0], p[1]] for p in pairs], dtype=torch.long)
targets = torch.tensor([(p[0] + p[1]) % self.p for p in pairs], dtype=torch.long)
return inputs, targets
class GrokkingTransformer(nn.Module):
"""Minimal Transformer to map [a, b] -> c."""
def __init__(self, vocab_size: int, embed_dim: int = 64, num_heads: int = 4, num_layers: int = 2):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim, nhead=num_heads,
dim_feedforward=embed_dim * 4, batch_first=True, activation='gelu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.classifier = nn.Linear(embed_dim, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = self.transformer(x)
x = x.mean(dim=1) # Global average pooling
return self.classifier(x)
def train_grokking_demo():
# Hyperparameters
P, TRAIN_FRACTION = 32, 0.4
EMBED_DIM, LR, WEIGHT_DECAY = 64, 1e-3, 1e-2 # Weight decay is critical!
EPOCHS, BATCH_SIZE = 2000, 32
dataset = ModularAdditionDataset(P, train_fraction=TRAIN_FRACTION)
train_x, train_y = dataset.get_split('train')
val_x, val_y = dataset.get_split('val')
train_loader = DataLoader(torch.utils.data.TensorDataset(train_x, train_y), batch_size=BATCH_SIZE, shuffle=True)
model = GrokkingTransformer(vocab_size=P, embed_dim=EMBED_DIM)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
for epoch in range(1, EPOCHS + 1):
model.train()
for batch_x, batch_y in train_loader:
optimizer.zero_grad()
loss = criterion(model(batch_x), batch_y)
loss.backward()
optimizer.step()
if epoch % 200 == 0 or epoch == 1:
model.eval()
with torch.no_grad():
train_acc = (model(train_x).argmax(1) == train_y).float().mean().item()
val_acc = (model(val_x).argmax(1) == val_y).float().mean().item()
print(f"Epoch {epoch:4d} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")
if val_acc > 0.99:
print("Grokking achieved!")
break
return model
if __name__ == '__main__':
torch.manual_seed(42)
trained_model = train_grokking_demo()Key Technical Takeaways
1. The Role of Weight Decay
You might notice WEIGHT_DECAY = 1e-2 in the code. This is not just a tuning parameter; it is the engine of grokking. Weight decay penalizes complex, "memorized" solutions. When the model is forced to keep its weights small, it eventually finds that the most efficient way to represent the data is not to memorize every pair, but to implement the mathematical rule (which is a much more "compressed" representation).
2. The Two Phases of Learning
- Phase 1: Memorization. The model uses its capacity to map specific inputs to outputs. Training accuracy $\to 100%$, Validation accuracy $\approx$ chance.
- Phase 2: Generalization. The model discovers the symmetry of the operation. Validation accuracy suddenly spikes to $100%$.
3. Visualizing the "Aha!" Moment
If you were to visualize the learned embeddings using t-SNE, you would see a fascinating transition. In the memorization phase, the embeddings are scattered. After grokking, the embeddings for modular addition typically form a perfect circle (ring topology), reflecting the cyclic nature of modular arithmetic.
Conclusion
Grokking challenges our fundamental understanding of overfitting. It suggests that "overfitting" might sometimes be a necessary stepping stone toward discovering deep, algorithmic truths. For practitioners, it serves as a reminder that with the right regularization and enough patience, models can find elegant solutions to complex problems—even when they seem to have hit a wall.