brendanlong/lens-loss-grokking-experiment

lens-loss-grokking-experiment — checkpoints

Final model checkpoints for the experiments in brendanlong/lens-loss-grokking-experiment (deep supervision vs grokking: LN-scoped chronic instability, weight-decay circuit pruning, and failure isolation). Training curves: public wandb project.

Layout: grok_lens/<run_name>/final.pt (modular-arithmetic runs, ~1.7 MB each, with .json metadata sidecars) and lego/<run_name>/step_*.pt (S3 multi-hop composition runs). Run names encode the configuration, e.g. p113-L2-lam0.3-uniform-frac0.3-s42 = p=113, 2 layers, aux λ=0.3, uniform layer weighting, 30% train fraction, seed 42; suffixes: -50k/-100k/-500k (step budget), -muon (hand-rolled Muon), -torchmuon (stock torch.optim.Muon re-runs — canonical), -noln (no-LayerNorm architecture control), -wd0.25 (weight decay), -ffttrace (per-frequency logging), -shuf (shuffled-target control), p113sub (modular subtraction), -cont-* (objective-switching continuations).

LEGO runs — which to use: lego/S3-std-8L-splitku-* are the canonical runs (enumerated chains, disjoint per-k-stratified held-out split, k-uniform sampling; see the repo's RESULTS.md "Review re-runs II"). The older lego/S3-std-8L-lensaux* runs are deprecated for accuracy claims: their evaluation sets overlapped the training stream's coverage, so their "test accuracy" partly measured seen chains. They are kept for provenance and for lens-structure analyses that don't depend on the split.

Checkpoints are torch.load(..., weights_only=True)-safe dicts: {"step", "model_state_dict", "model_config"}. The repo's analysis scripts download from here automatically.