README.md

---
license: mit
tags:
  - mechanistic-interpretability
  - grokking
  - modular-arithmetic
  - transformer
  - TransformerLens
  - pytorch
  - toy-model
language:
  - en
library_name: transformers
pipeline_tag: text-classification
---

# Modular Multiplication Transformer

A 1-layer, 4-head transformer trained on **(a x b) mod 113** that exhibits **grokking** (delayed generalization after memorization). This checkpoint includes full training history (400 checkpoints across 40,000 epochs).

## Model Architecture

| Parameter | Value |
|-----------|-------|
| Layers | 1 |
| Attention Heads | 4 |
| d_model | 128 |
| d_head | 32 |
| d_mlp | 512 |
| Activation | ReLU |
| Layer Norm | None |
| Vocabulary | 114 (0-112 + "=" separator) |
| Output Classes | 113 |
| Context Length | 3 tokens [a, b, =] |
| Trainable Parameters | ~230,000 |

Built with [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) (`HookedTransformer`). No layer normalization and frozen biases.

## Training Details

| Parameter | Value |
|-----------|-------|
| Optimizer | AdamW |
| Learning Rate | 1e-3 |
| Weight Decay | 1.0 |
| Betas | (0.9, 0.98) |
| Epochs | 40,000 |
| Training Fraction | 30% (3,830 / 12,769 samples) |
| Batch Size | Full-batch |
| Data Seed | 598 |
| Model Seed | 999 |

## Checkpoint Contents

```python
checkpoint = torch.load("mod_mult_grokking.pth")

checkpoint["model"]              # Final model state_dict
checkpoint["config"]             # HookedTransformerConfig
checkpoint["checkpoints"]        # List of 400 state_dicts (every 100 epochs)
checkpoint["checkpoint_epochs"]  # [0, 100, 200, ..., 39900]
checkpoint["train_losses"]       # 40,000 training loss values
checkpoint["test_losses"]        # 40,000 test loss values
checkpoint["train_accs"]         # 40,000 training accuracy values
checkpoint["test_accs"]          # 40,000 test accuracy values
checkpoint["train_indices"]      # Indices of 3,830 training samples
checkpoint["test_indices"]       # Indices of 8,939 test samples
```

## Usage

```python
import torch
from transformer_lens import HookedTransformer

checkpoint = torch.load("mod_mult_grokking.pth", map_location="cpu")
model = HookedTransformer(checkpoint["config"])
model.load_state_dict(checkpoint["model"])
model.eval()

# Compute 7 * 16 mod 113 = 112
a, b = 7, 16
separator = 113  # "=" token
input_tokens = torch.tensor([[a, b, separator]])
logits = model(input_tokens)
prediction = logits[0, -1].argmax().item()
print(f"{a} * {b} mod 113 = {prediction}")  # 112
```

## References

- Nanda et al. (2023). "Progress measures for grokking via mechanistic interpretability." ICLR 2023.
- Power et al. (2022). "Grokking: Generalization beyond overfitting on small algorithmic datasets."
- [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens)