Grokking

Grokking is delayed generalization: a neural network reaches perfect accuracy on its training data early, stays at chance level on validation data for a long stretch, and then, well after it appears to have overfit, suddenly generalizes. Power et al. (2022) at OpenAI named the effect after training small transformers on algorithmic tasks such as modular arithmetic, where validation accuracy rose from random chance to perfect long after training accuracy had saturated. The word comes from Robert Heinlein's novel Stranger in a Strange Land, where to grok means to understand something completely.

How Grokking Unfolds#

Modular addition is the classic grokking task: the network sees pairs (a, b) and must output (a + b) mod p, with only a fraction of all pairs available for training. The run below trains a small multilayer perceptron with PyTorch's AdamW optimizer on 30% of the 9,409 pairs for p = 97 and finishes in under a minute on a CPU:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
p = 97  # task: predict (a + b) mod 97 from one-hot encodings of a and b
a, b = torch.cartesian_prod(torch.arange(p), torch.arange(p)).T
x, y = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float(), (a + b) % p
idx = torch.randperm(p * p)
train, val = idx[: p * p * 3 // 10], idx[p * p * 3 // 10 :]  # train on 30% of all pairs

model = torch.nn.Sequential(torch.nn.Linear(2 * p, 256), torch.nn.ReLU(), torch.nn.Linear(256, p))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)  # strong weight decay
for step in range(1, 20001):
    loss = F.cross_entropy(model(x[train]), y[train])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 4000 == 0:
        train_acc, val_acc = [(model(x[i]).argmax(1) == y[i]).float().mean().item() for i in (train, val)]
        print(f"step {step:5d}: train accuracy {train_acc:.2f}, validation accuracy {val_acc:.2f}")

At step 4,000 training accuracy is already 1.00 while validation accuracy is 0.00. Validation accuracy then climbs to 0.16 at step 8,000, 0.54 at step 12,000, 0.84 at step 16,000, and 0.97 at step 20,000. With weight_decay=0.0, the same run stays at 0.00 validation accuracy through all 20,000 steps.

Nanda et al. (2023) explained the plateau by reverse-engineering a one-layer transformer trained on modular addition. The network learns to turn addition into rotation around a circle, using discrete Fourier features and trigonometric identities. Training splits into three continuous phases, memorization, circuit formation, and cleanup: the generalizing circuit grows gradually while validation accuracy stays flat, and the jump arrives when weight decay removes the memorized solution that was still dominating the outputs.

What Makes Grokking Happen#

• Regularization: Weight decay was the most effective intervention in the original experiments, and the run above does not generalize at all without it.
• Dataset size: Smaller training fractions need much more optimization before generalizing, and below a threshold the network never generalizes within a practical budget.
• Weight norm: Omnigrok links grokking to the size of the weights, since generalizing solutions sit at a smaller weight norm than memorizing ones. Starting from large initial weights, the authors induced grokking on image, language, and molecule tasks, and by constraining the weight norm they removed it from algorithmic ones.
• Numerical stability: Prieto et al. (2025) found that without regularization, cross-entropy training keeps scaling up the logits until floating-point errors in the softmax stop learning, which they call Softmax Collapse. Their StableMax activation and ⊥Grad optimizer produce grokking with no weight decay.

Grokking vs. Related Concepts#

• Overfitting: A gap between training and validation performance that persists or widens. In grokking the gap closes on its own with continued training.
• Double descent: Test error falls, rises, and falls again as model size or training time grows. Epoch-wise double descent is another form of delayed improvement, in which test error rises for a while during training before falling again.
• Mechanistic interpretability: The study of the algorithms inside trained networks. Grokking on modular arithmetic became one of its main case studies because the learned algorithm could be read out in full.

Grokking in Practice#

Grokking is clearest on small algorithmic datasets governed by one exact rule. It still matters for experiment design, because early stopping ends a run once its validation metric stops improving, which would cut a grokking run off during its plateau. Ultralytics YOLO training stops after patience epochs without improvement in validation fitness, 100 by default, and runs every epoch when patience=0.