modular-addition-checkpoints

Spec — .pth checkpoint trajectories (per-epoch cadence incl. final ep030000, optimizer per save policy) for every (noise, seed) cell of 1-layer-transformer grokking runs on (i + j) mod p, with optional frozen oracle Fourier-feature injection. Produced by modular_addition/oracle/push_to_hf.py; consumed for post-hoc weight-level metrics (W_L, ablations) without retraining. Status: active.

Model checkpoints from the transformer-modular-addition experiments: a 1-layer transformer grokking (i + j) mod p, with frozen "oracle" Fourier features optionally injected into the residual stream.

📖 Experimental details, code, and analysis live here: https://github.com/kaushikreddyxyz/transformer-modular-addition

This repo is a weights/results archive only — see the GitHub repo for what each experiment varies and how the figures are produced.

Layout

run_<timestamp>/<experiment>/checkpoints/<label>/ep<NNNNNN>.pth   # weights
run_<timestamp>/<experiment>/<label>.result.json                  # metrics + oracle spec

• One run_<timestamp>/ tree per training-sweep invocation.
• <label> encodes that cell's sweep coordinates, e.g.
  • n2_s0 — 2 injected frequency pairs, seed 0
  • baseline_s0 — no injection
  • delay4000_n2_s1 — injection delayed to epoch 4000
  • amp0.5_n1_s2 — oracle amplitude 0.5
  • rel0.25_n1_s0 — relative amplitude 0.25
  • wd0.0001_s3 — weight-decay sweep
• ep<NNNNNN> is the snapshot epoch (e.g. ep030000 = fully trained under the default 30k-epoch schedule).

Experiments

run	experiments
run_20260612_200000	exp01, exp02_1, exp02_2, exp04, exp05, exp06 (+ exp06_deprecated)
run_20260616_110846	exp07, exp08
run_20260619_144122	exp09

Checkpoint format

Each .pth is a dict — load with torch.load(path, weights_only=False):

key	meaning
model	model state_dict (the frozen oracle is NOT included — see below)
config	transformer.Config fields (p, d_model, lr, weight_decay, …)
label	the cell label (sweep coordinates)
epochs_done	epochs trained at this snapshot
inject_from_epoch	epoch the oracle injection turned on (0 = from start)

The frozen oracle adds a deterministic constant to the forward pass and is not stored in the state_dict. Rebuild it from the sibling <label>.result.json (spec / injected_freqs fields) via modular_addition.oracle.sweep.build_oracle.

<label>.result.json also carries the full training history, snapshots, grok_epoch, and wall-clock — the source of truth for the figures in the GitHub repo.

Loading example

import torch

ckpt = torch.load(
    "run_20260612_200000/exp01/checkpoints/n2_s0/ep030000.pth",
    weights_only=False,
)
state_dict = ckpt["model"]
config = ckpt["config"]