Grokking › Grokking in Arithmetic
Language AI HandbookAll chapters Grokking › Grokking in ArithmeticOne section · Sections in this chapter3 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The original grokking paper focused on small algorithmic datasets, particularly modular arithmetic. These tasks give an ideal testbed for studying the memorization-to-generalization transition because they have exactly the right combination of properties. They are simple enough to study carefully, complex enough to require real learning, and mathematically clean enough to reason about the ideal solution the network should discover. This combination of simplicity and mathematical richness makes modular arithmetic the perfect laboratory for observing and understanding grokking in controlled conditions.
The choice of modular arithmetic is particularly elegant. It is one of the simplest possible algorithmic tasks where the rules are completely deterministic, the input space is finite and well-defined, and there exists a compact closed-form solution that any sufficiently expressive model should theoretically be able to learn. Unlike natural language tasks where the ground truth is ambiguous and the space of correct answers is ill-defined, modular arithmetic gives us an absolute standard for perfect generalization: a model either computes the correct answer for every possible input or it does not. This clean criterion makes modular arithmetic the ideal laboratory for studying the basic dynamics of how neural networks transition from memorization to algorithmic understanding.
The tasks are framed as sequence modeling problems, letting researchers to use the same transformer architectures that power modern language models. This connection to practical architectures is not incidental: it means the insights about grokking are potentially needed to understanding what happens during language model training at scale, even if the direct evidence for grokking in large models remains unclear. By choosing an architecture with real practical relevance, the original researchers ensured their findings could not be dismissed as artifacts of a toy setup with no connection to real systems.
Modular Addition
The canonical grokking task is modular addition: given two numbers and , compute for some prime . This operation lies at the heart of many cryptographic systems and is one of the most basic operations in abstract algebra. It is both practically needed and mathematically well-understood, which makes it ideal for theoretical analysis.
The modulo operation computes the remainder when the sum of two numbers is divided by . Think of modular arithmetic as clock arithmetic: on a 12-hour clock, if it is 10 o'clock and you add 5 hours, you do not get 15; you get 3, because the clock wraps around after 12. Modular arithmetic generalizes this idea to any modulus . For example, because , so the remainder when dividing by 97 is 13. The key insight is that results are always constrained to lie within the range from 0 to , creating a finite cyclic structure where every output is a member of a bounded set.
The task is formally stated as computing:
where:
- : the first operand, an integer from 0 to
- : the second operand, an integer from 0 to
- : the ordinary integer sum of the two operands, computed before applying the modulo operation
- : the modulo operation that returns the remainder after dividing by , so the result stays within the range
- : a prime number defining the modular arithmetic system, which creates the wraparound behavior
With , there are possible input pairs, each mapping to one of 97 outputs. This creates a dataset of manageable size: small enough for rapid experimentation yet large enough that pure memorization requires substantial model capacity. The choice of 97 as a specific prime is conventional but not necessary; similar results appear for other primes of comparable magnitude.
Prime numbers ensure that modular arithmetic forms a mathematical field, where every non-zero element has a multiplicative inverse. This gives the task clean, learnable patterns without special cases or irregular structure. When the modulus is composite (not prime), certain numbers share common factors with the modulus, creating irregular patterns that complicate operations like division. Prime moduli give the cleanest structure for the network to discover, which is why they were chosen for the original grokking experiments. This mathematical property ensures that a compact algorithmic solution exists, so a failure to generalize can be attributed to learning dynamics rather than the absence of a learnable structure.
The task is framed as sequence modeling to use the transformer architecture. The input sequence is [a, b, =] and the target is the result. The transformer learns to predict the correct output given the input tokens representing the operands. This framing lets us study grokking using the same architectures that power modern language models, making the insights applicable to understanding broader neural network behavior. It also raises the question of whether language models might exhibit grokking-like transitions when learning grammatical or semantic rules with underlying mathematical structure.
use_book_style()
import numpy as np
import torch
from torch.utils.data import Dataset
class ModularArithmeticDataset(Dataset):
"""Dataset for modular arithmetic tasks."""
def __init__(
self, operation="add", p=97, train=True, train_fraction=0.5, seed=42
):
self.p = p
self.operation = operation
# Generate all possible (a, b) pairs
all_pairs = [(a, b) for a in range(p) for b in range(p)]
# Compute results based on operation
if operation == "add":
all_results = [(a + b) % p for a, b in all_pairs]
elif operation == "subtract":
all_results = [(a - b) % p for a, b in all_pairs]
elif operation == "multiply":
all_results = [(a * b) % p for a, b in all_pairs]
elif operation == "divide":
# Only include pairs where b != 0 (division by zero undefined)
valid_pairs = [(a, b) for a, b in all_pairs if b != 0]
all_pairs = valid_pairs
all_results = [
(a * pow(b, p - 2, p)) % p for a, b in all_pairs
] # Modular inverse using Fermat's Little Theorem: b^(-1) ≡ b^(p-2) (mod p)
# Split into train/test
np.random.seed(seed)
indices = np.random.permutation(len(all_pairs))
n_train = int(len(indices) * train_fraction)
if train:
selected = indices[:n_train]
else:
selected = indices[n_train:]
self.data = [
(all_pairs[i][0], all_pairs[i][1], all_results[i]) for i in selected
]
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
a, b, result = self.data[idx]
# Input: [a, b, EQUALS_TOKEN], where EQUALS_TOKEN = p
# We use tokens 0 to p-1 for numbers, p for equals
x = torch.tensor([a, b, self.p], dtype=torch.long)
y = torch.tensor(result, dtype=torch.long)
return x, y# Create train and test datasets
use_book_style()
p = 97
train_dataset = ModularArithmeticDataset(
operation="add", p=p, train=True, train_fraction=0.5
)
test_dataset = ModularArithmeticDataset(
operation="add", p=p, train=False, train_fraction=0.5
)
# Store dataset metrics for display
total_pairs = p * p
n_train = len(train_dataset)
n_test = len(test_dataset)
# Get sample examples
sample_examples = []
for i in range(3):
x, y = train_dataset[i]
a, b = x[0].item(), x[1].item()
sample_examples.append((a, b, y.item()))Prime modulus: 97 Total possible pairs: 9409 Training examples: 4704 Test examples: 4705 Sample training examples: (14 + 42) mod 97 = 56 (20 + 89) mod 97 = 12 (90 + 9) mod 97 = 2
With 9,409 total possible input pairs, the 50/50 train-test split gives 4,704 training examples and 4,705 test examples. This balanced partition is important for studying grokking. It requires the model to generalize to unseen combinations rather than simply memorizing the training set. A model that has only memorized will reach roughly 50% accuracy on the test set at best, since it has never seen the other half of the input pairs. A model that has learned the addition algorithm will succeed on any valid input pair, regardless of whether that specific pair appeared in training.
The sample outputs show the modular wraparound behavior: sums exceeding the modulus cycle back to smaller values. This basic property is what the model must learn to capture for true generalization, and it is a property that cannot be memorized from a subset of examples alone. No matter how many specific pairs you memorize, if you do not understand the wraparound rule, you will fail on pairs near the boundary where wraparound matters. The partition ensures that the boundary cases, those inputs where exceeds , are split between training and test sets, which makes it impossible to memorize the answer to all boundary inputs.
The heatmap reveals the striking diagonal pattern in modular addition. Each diagonal stripe is all input pairs that share the same sum modulo . These stripes wrap around at the boundaries: a diagonal that exits the right edge of the heatmap re-enters from the left at a lower row, capturing the wraparound property of modular arithmetic. This visual structure makes clear what the network must learn: the output depends only on the sum of the inputs modulo the prime, not on the specific values of and individually. A model that has truly learned this structure can predict any entry in the table; a model that has merely memorized specific pairs can only predict those pairs.
With only 50% of the data for training, the network must generalize to unseen input combinations. This setup creates a clean distinction between memorization and generalization. The network reaches perfect training accuracy early by memorizing, but test accuracy only climbs when it discovers the actual algorithmic rule. The heatmap also helps you appreciate the geometric regularity of the task: the pattern is not random noise but a highly structured arrangement with predictable symmetries. This structure is precisely what the network must learn to stand for internally in order to generalize.
The Transformer Model
Following the original grokking experiments, we use a small transformer to learn these arithmetic patterns. The transformer architecture is particularly well-suited for this task because its attention mechanism can learn to route information between the two input operands, and its feedforward networks can learn the computational operations needed to combine them. The attention mechanism allows the model to create direct connections between the two operand positions and the output position, which is the needed computational structure needed for modular arithmetic.
The model processes the three-token input sequence [a, b, =] and predicts the output value from the representation at the final (equals) position. This framing mirrors how we might train a language model to complete arithmetic expressions. The key architectural choice is using learned positional embeddings rather than fixed sinusoidal ones, because the positions here carry semantic meaning: position 0 is always the first operand, position 1 is always the second, and position 2 is the equals sign that cues the network to produce its answer. Fixed sinusoidal embeddings based on smooth frequency functions would inject an unnecessary geometric prior that could interfere with the network's ability to discover the circular structure of modular arithmetic on its own.
One necessary implementation detail is the absence of dropout. In most neural network training, dropout is recommended as a regularization technique that prevents overfitting. But for studying grokking, dropout would obscure the dynamics we are trying to observe by introducing its own regularization effect. Instead, we rely solely on weight decay, making the regularization mechanism clean and controllable. This design choice is important for research purposes: by isolating weight decay as the sole regularization, we can cleanly attribute the grokking transition to its effect without confounds from other regularization mechanisms.
use_book_style()
import torch
import torch.nn as nn
class GrokkingTransformer(nn.Module):
"""Small transformer for studying grokking on algorithmic tasks."""
def __init__(
self,
vocab_size,
d_model=128,
nhead=4,
num_layers=2,
dim_feedforward=512,
):
super().__init__()
self.d_model = d_model
# Token embeddings: vocab_size + 1 for the equals token
self.embedding = nn.Embedding(vocab_size + 1, d_model)
# Positional embeddings (learned)
self.pos_embedding = nn.Embedding(3, d_model) # 3 positions: a, b, =
# Transformer encoder layers
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=0.0, # No dropout - important for observing grokking
batch_first=True,
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=num_layers
)
# Output projection
self.output = nn.Linear(d_model, vocab_size)
def forward(self, x):
# x shape: (batch, 3) containing [a, b, equals_token]
batch_size, seq_len = x.shape
# Get embeddings
tok_emb = self.embedding(x)
pos = (
torch.arange(seq_len, device=x.device)
.unsqueeze(0)
.expand(batch_size, -1)
)
pos_emb = self.pos_embedding(pos)
# Combine embeddings
h = tok_emb + pos_emb
# Transform
h = self.transformer(h)
# Predict from the last position (after seeing [a, b, =])
logits = self.output(h[:, -1, :])
return logitsThe architecture is deliberately minimal: two transformer layers, 128-dimensional embeddings, four attention heads, and a feedforward dimension of 512. This scale is small enough to train quickly and large enough to exhibit grokking reliably. The model has approximately 200,000 parameters, which is orders of magnitude smaller than a practical language model but far more than the roughly 5,000 parameters needed to memorize half of the 9,409 possible input pairs. This excess capacity is needed for grokking: a model with just barely enough capacity to memorize would not have the representational headroom to develop the algorithmic solution in parallel with the memorization solution.
Training Dynamics
The key to observing grokking is weight decay (L2 regularization). Without it, the model memorizes the training data and never generalizes. With appropriate weight decay, grokking emerges. This reveals a basic aspect of grokking: the phenomenon requires balance between the model's capacity to memorize and the regularization pressure favoring simpler, generalizable solutions. Weight decay pulls parameters toward zero, penalizing the complex configurations that memorization requires while leaving simpler algorithmic solutions relatively unaffected.
The intuition for why weight decay lets grokking is subtle but important. Both the memorization solution and the algorithmic solution reach low training loss, so gradient descent alone cannot distinguish between them. However, weight decay adds a penalty proportional to the squared magnitude of all parameters. The memorization solution, which stores information about each training example separately, tends to require larger or more numerous parameter configurations than the algorithmic solution, which implements a compact operation applicable to all inputs. Over many training steps, this asymmetric penalty gradually erodes the advantage of the memorization solution and increases the relative attractiveness of the compact algorithmic solution.
use_book_style()
def train_with_logging(
model,
train_loader,
test_loader,
n_epochs,
lr=1e-3,
weight_decay=1.0,
device="cpu",
):
"""Train model and log metrics for analyzing grokking."""
model = model.to(device)
optimizer = torch.optim.AdamW(
model.parameters(), lr=lr, weight_decay=weight_decay
)
criterion = nn.CrossEntropyLoss()
history = {
"train_loss": [],
"train_acc": [],
"test_loss": [],
"test_acc": [],
"step": [],
}
step = 0
for epoch in range(n_epochs):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
step += 1
# Evaluate every epoch
model.eval()
with torch.no_grad():
# Training metrics
train_correct, train_total, train_loss_sum = 0, 0, 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
train_loss_sum += criterion(logits, y).item() * len(y)
train_correct += (logits.argmax(dim=1) == y).sum().item()
train_total += len(y)
# Test metrics
test_correct, test_total, test_loss_sum = 0, 0, 0
for x, y in test_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
test_loss_sum += criterion(logits, y).item() * len(y)
test_correct += (logits.argmax(dim=1) == y).sum().item()
test_total += len(y)
history["train_loss"].append(train_loss_sum / train_total)
history["train_acc"].append(train_correct / train_total)
history["test_loss"].append(test_loss_sum / test_total)
history["test_acc"].append(test_correct / test_total)
history["step"].append(step)
return historyuse_book_style()
from torch.utils.data import DataLoader
# Set up a smaller executable demo. The full p=97 experiment is too slow for
# a documentation notebook that is converted and executed repeatedly.
p_demo = 31
demo_train_dataset = ModularArithmeticDataset(
operation="add", p=p_demo, train=True, train_fraction=0.5
)
demo_test_dataset = ModularArithmeticDataset(
operation="add", p=p_demo, train=False, train_fraction=0.5
)
train_loader = DataLoader(demo_train_dataset, batch_size=256, shuffle=True)
test_loader = DataLoader(demo_test_dataset, batch_size=256, shuffle=False)
# Initialize model
torch.manual_seed(42)
model = GrokkingTransformer(
vocab_size=p_demo, d_model=64, nhead=4, num_layers=1, dim_feedforward=128
)
# Train briefly to observe early grokking dynamics.
device = "cuda" if torch.cuda.is_available() else "cpu"
history = train_with_logging(
model,
train_loader,
test_loader,
n_epochs=25,
lr=1e-3,
weight_decay=3.0,
device=device,
)The training run is intentionally small. It gives enough iterations to inspect the early dynamics without turning the chapter conversion into a long-running experiment. Full grokking convergence may require 10,000 or more epochs depending on random initialization, model size, and dataset setup. The choice of weight decay 3.0 is deliberately aggressive to make the regularization effect visible in a short demonstration run; in the original paper, weight decay values varied across experiments depending on the specific task and model size.
Final training accuracy: 0.138 Final test accuracy: 0.004 Final training loss: 3.1431 Final test loss: 3.7560
The final metrics show the model's state after 300 epochs of training. The training accuracy approaching 1.0 confirms successful memorization of the training set, while the test accuracy value indicates the degree of generalization achieved at this point in training.
The training curve shows the early version of the grokking pattern. Training accuracy improves as the model starts fitting the observed examples, while test accuracy reveals whether that behavior transfers beyond the training split. The logarithmic loss plot is particularly informative: it shows how quickly the training objective can improve relative to the test objective. In longer training runs, test accuracy may eventually match training accuracy. This shows complete grokking.
This delay can be large. In the original grokking paper, some experiments required training for over steps after the model had memorized the training data before generalization emerged. The ratio of plateau duration to memorization duration can exceed 100 to 1, meaning the model spends 100 times as long appearing to be stuck as it does memorizing. This makes grokking easy to miss in standard training runs with early stopping. If you are used to stopping training as soon as validation loss plateaus, you would never observe the transition, and you might conclude the model had fully converged when it had only memorized.