BurnyCoder/grokking-modular-addition-runs

Grokking modular addition — retrained runs

22 one-layer transformers trained on (a + b) mod 113 with the settings of Nanda et al. 2023, Progress measures for grokking via mechanistic interpretability (arXiv:2301.05217): 30 % of all pairs for training, full-batch AdamW (lr 1e-3, betas (0.9, 0.98), weight decay 1 or 0), 10-step warm-up, 40,000 steps, deterministic PyTorch. Runs (seeds and controls, see docs/preregistration.md in the code repository):

• rand_s0 — rand: init seed 0, split seed 0, weight decay 1.0, task rand
• rand_s1 — rand: init seed 1, split seed 1, weight decay 1.0, task rand
• rand_s2 — rand: init seed 2, split seed 2, weight decay 1.0, task rand
• split0_init100 — split0: init seed 100, split seed 0, weight decay 1.0, task add
• split0_init101 — split0: init seed 101, split seed 0, weight decay 1.0, task add
• split0_init102 — split0: init seed 102, split seed 0, weight decay 1.0, task add
• wd0_s0 — wd0: init seed 0, split seed 0, weight decay 0.0, task add
• wd0_s1 — wd0: init seed 1, split seed 1, weight decay 0.0, task add
• wd0_s2 — wd0: init seed 2, split seed 2, weight decay 0.0, task add
• wd0_s3 — wd0: init seed 3, split seed 3, weight decay 0.0, task add
• wd0_s4 — wd0: init seed 4, split seed 4, weight decay 0.0, task add
• wd1_s0 — wd1: init seed 0, split seed 0, weight decay 1.0, task add
• wd1_s1 — wd1: init seed 1, split seed 1, weight decay 1.0, task add
• wd1_s10 — wd1: init seed 10, split seed 10, weight decay 1.0, task add
• wd1_s2 — wd1: init seed 2, split seed 2, weight decay 1.0, task add
• wd1_s3 — wd1: init seed 3, split seed 3, weight decay 1.0, task add
• wd1_s4 — wd1: init seed 4, split seed 4, weight decay 1.0, task add
• wd1_s5 — wd1: init seed 5, split seed 5, weight decay 1.0, task add
• wd1_s6 — wd1: init seed 6, split seed 6, weight decay 1.0, task add
• wd1_s7 — wd1: init seed 7, split seed 7, weight decay 1.0, task add
• wd1_s8 — wd1: init seed 8, split seed 8, weight decay 1.0, task add
• wd1_s9 — wd1: init seed 9, split seed 9, weight decay 1.0, task add

Each folder holds checkpoints.safetensors (every parameter stacked over checkpoints: tensor <param> has shape [T, ...], checkpoint t taken after meta.json["ckpt_steps"][t] optimiser steps; parameter names follow the paper's code, so GrokTransformer in the code repository loads them with strict=True), curves.npz (per-step train loss/accuracy, test loss/accuracy every 10 steps, per-parameter weight/gradient/update norms) and analysis/ (all measurements shown on the website).

Code, website and reports: https://github.com/BurnyCoder/visualizing-grokking-modular-addition. Author: Libor Burian.