BurnyCoder/grokking-modular-addition-transformer

Grokking Modular Addition Transformer

A 1-layer transformer trained on modular addition (a + b) mod 113 that exhibits grokking -- the phenomenon where the model first memorizes the training data, then suddenly generalizes to the test set after continued training.

This model is a reproduction of the setup from Progress Measures for Grokking via Mechanistic Interpretability (Nanda et al., 2023), built with TransformerLens.

Model Description

The model learns a Fourier-based algorithm to perform modular addition:

1. Embed inputs a and b into Fourier components (sin/cos at key frequencies)
2. Attend from the = position to a and b, computing sin(ka), cos(ka), sin(kb), cos(kb)
3. MLP neurons compute cos(k(a+b)) and sin(k(a+b)) via trigonometric identities
4. Unembed maps these to logits approximating cos(k(a+b-c)) for each candidate output c

Architecture

Parameter	Value
Layers	1
Attention heads	4
d_model	128
d_head	32
d_mlp	512
Activation	ReLU
Normalization	None
Vocabulary (input)	114 (0-112 for numbers, 113 for =)
Vocabulary (output)	113
Context length	3 tokens: [a, b, =]
Parameters	227,313 (226,048 weights + 1,265 bias entries, which stay at 0)

Design choices (no LayerNorm, ReLU, biases frozen at zero) were made to simplify mechanistic interpretability analysis.

Usage

Loading the checkpoint

import torch
from transformer_lens import HookedTransformer

# Download and load
cached_data = torch.load("grokking_demo.pth", weights_only=False)

model = HookedTransformer(cached_data["config"])
model.load_state_dict(cached_data["model"])

# Training history is also included
model_checkpoints = cached_data["checkpoints"]       # 250 intermediate checkpoints
checkpoint_epochs = cached_data["checkpoint_epochs"]  # Every 100 epochs
train_losses = cached_data["train_losses"]
test_losses = cached_data["test_losses"]
train_indices = cached_data["train_indices"]
test_indices = cached_data["test_indices"]

Running inference

import torch

p = 113
a, b = 37, 58
input_tokens = torch.tensor([[a, b, p]])  # [a, b, =]
logits = model(input_tokens)
prediction = logits[0, -1].argmax().item()
print(f"{a} + {b} mod {p} = {prediction}")  # Should print 95

Installation

pip install torch transformer-lens

Training Details

Setting	Value
Task	(a + b) mod 113
Total data	113^2 = 12,769 pairs
Train split	30% (3,830 examples)
Test split	70% (8,939 examples)
Optimizer	AdamW
Learning rate	1e-3
Weight decay	1.0
Betas	(0.9, 0.98)
Epochs	25,000
Batch size	Full batch
Checkpoints	Every 100 epochs (250 total)
Seed	999 (model), 598 (data split)
Training time	~2 minutes on GPU (as reported when trained; not in the file)

The high weight decay (1.0) is critical for grokking -- it gradually erodes memorization weights in favor of the compact generalizing Fourier circuit.

Grokking Phases

The training exhibits three distinct phases:

1. Memorization (epochs 0 to ~1,600): Train loss falls to ~10⁻⁷ while test loss rises far above chance, to a peak of 31.7 at epoch 1,381 (uniform guessing would give ln 113 ≈ 4.73). The model memorizes all training examples.
2. Circuit Formation (~1,600 to ~9,000): The Fourier-based generalizing circuit gradually forms in the weights while memorization still dominates; test loss falls only slowly (≈ 20 at epoch 6,000).
3. Cleanup (~9,000 to ~12,800): Weight decay erodes memorization faster than the compact Fourier circuit. Test loss collapses (below 1 at epoch 9,546, below 0.01 at epoch 10,160) -- this is the grokking moment.

Phase boundaries come from the objective rules of visualizing-grokking-modular-addition applied to this file's losses and weights: memorization ends when the train loss stops falling by 1 % per 100 epochs, circuit formation when the test loss reaches 20 % of its early plateau, and cleanup when the sum of squared weights has nearly settled.

Mechanistic Interpretability Findings

Analysis of the trained model reveals:

• Fourier-sparse embeddings: The model learns embeddings concentrated on key frequencies (k = 9, 33, 36, 38, 55)
• Neuron clustering: 298 of 512 MLP neurons (58 %) have more than 85 % of their activation variance explained by a single Fourier frequency (clusters k = 9: 42, k = 33: 56, k = 36: 107, k = 38: 45, k = 55: 48 neurons)
• Logit periodicity: Output logits approximate cos(freq * 2pi/p * (a + b - c)) for key frequencies
• Progress measures: Restricted loss and excluded loss track the formation and cleanup of circuits independently, revealing that grokking is not a sudden phase transition but the delayed visibility of a gradually forming algorithm

Source Code

Full analysis notebook and training code: GitHub repository

References

• Progress Measures for Grokking via Mechanistic Interpretability (Nanda et al., 2023)
• Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets (Power et al., 2022)
• TransformerLens