| --- |
| 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) |
| |
