README.md

---
license: mit
tags:
  - mechanistic-interpretability
  - grokking
  - modular-addition
  - transformer
  - TransformerLens
datasets:
  - custom
language:
  - en
library_name: transformer-lens
pipeline_tag: other
---

# 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](https://arxiv.org/abs/2301.05217) (Nanda et al., 2023), built with [TransformerLens](https://github.com/neelnanda-io/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 | ~2.5M |

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

## Usage

### Loading the checkpoint

```python
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

```python
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

```bash
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 |

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** (~epoch 0-1,500): Train loss drops to ~0, test loss stays at ~4.73 (random guessing over 113 classes). The model memorizes all training examples.
2. **Circuit Formation** (~epoch 1,500-13,300): The Fourier-based generalizing circuit gradually forms in the weights, but memorization still dominates. Test loss appears unchanged.
3. **Cleanup** (~epoch 13,300-16,600): Weight decay erodes memorization faster than the compact Fourier circuit. Test loss suddenly drops -- this is the grokking moment.

## 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**: ~85% of MLP neurons are well-explained by a single Fourier frequency
- **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](https://github.com/BurnyCoder/ai-mechanistic-interpretability-transformer-modular-addition-grokking)

## References

- [Progress Measures for Grokking via Mechanistic Interpretability](https://arxiv.org/abs/2301.05217) (Nanda et al., 2023)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://arxiv.org/abs/2201.02177) (Power et al., 2022)
- [TransformerLens](https://github.com/neelnanda-io/TransformerLens)