diff options
| -rw-r--r-- | .gitignore | 2 | ||||
| -rw-r--r-- | README.md | 69 | ||||
| -rw-r--r-- | design/preregistration.md | 50 |
3 files changed, 101 insertions, 20 deletions
@@ -2,5 +2,7 @@ __pycache__/ *.pyc .venv/ runs/ +!runs/*/ +!runs/*/summary.csv .pytest_cache/ report/ @@ -1,34 +1,63 @@ # Prime Grokking -Can minimal architectures learn the next-prime function — or *grok* it rather than memorize it? +**Can a minimal architecture learn the next-prime function — or *grok* it rather than memorize it?** -Experiment 1 (seed 0): weight-tied 2-layer RNN cell (K=20 tied steps, ACT learned halting) -vs. a fixed-d_model transformer baseline, range n ∈ [2, 100], 30% holdout. +The grokking literature studies almost exclusively group-structured tasks (modular arithmetic). +Next-prime supplies the missing **non-group-structured case**. Results so far (phases 1–4, +range [2, 100]): **no grokking observed** — memorization, partial in-range generalization, and a +decomposition surprise (is-prime generalizes ~2.5x better than next-prime), but no delayed +transition. That negative-class result is the point of the project; E6 (range [2, 1000]) and a +long-horizon follow-up are in flight. -## Layout +**The unusual bit:** every batch is **pre-registered**. Before each run, the outcome → +interpretation matrix is committed to the private research repo (the commit hash is the lock), +and results are interpreted *only* against that lock. The complete record, addenda 1–6, is +mirrored in [`design/preregistration.md`](design/preregistration.md). -``` -design/experiment-spec.md design doc (copied from the research repo, provenance noted) -design/preregistration.md pre-registered outcome→interpretation matrix (LOCK: commit 00c696d in research repo) -src/ config, data, model API, models, train, eval -tests/ pytest suite (data correctness, model shapes, halting) -scripts/ run_experiment.sh, plot.py -runs/ metrics CSVs + plots (gitignored) -``` +## Contents + +- `design/` — experiment spec, full pre-registration, external review passes (Gemini, Codex) including rejected points and why +- `src/` — config, data (sieve + bisect), weight-tied ACT-GRU RNN, GPT-style transformer baseline, train, eval (chunked, probe reports) +- `tests/` — 48-test pytest suite (data correctness, weight-tying, batch invariance, device handling, scaling) +- `scripts/` — idempotent sweep runner, experiment launcher, plots +- `NOTES-*.md` — per-phase analyses written against the lock, not after the fact +- `runs/*/summary.csv` — headline numbers per phase (raw run dirs are gitignored; everything is regenerable) ## Quick start +CPU (no GPU needed, ~35 min/run on 2 cores): + ```bash python3 -m venv .venv -.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu numpy matplotlib pytest +.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu +.venv/bin/pip install numpy matplotlib pytest .venv/bin/pytest tests/ -q -scripts/run_experiment.sh rnn 0 # full run, ~35 min on 2 cores -scripts/run_experiment.sh transformer 0 -scripts/plot.py +.venv/bin/python -m src.train rnn 0 ``` -## Notes +CUDA (AMP fp16 is on by default on CUDA; Turing tensor cores): + +```bash +.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cu126 +.venv/bin/python -m src.train transformer 0 --range_end 1000 --max_steps 32 --device cuda +``` + +Notes: `--max_steps` is K (the tied-iteration cap — it must exceed the max prime gap of the +range: 20 below 1000, 36 below 10^4, 72 below 10^5). `--compile_model True` roughly doubles the +transformer on CUDA but hurts the RNN (dynamic halting breaks graph capture). Full flag surface +(weight decay, train fraction, task mode, vocab mode, halting ablation, cosine decay) lives in +`src/config.py`. + +## Research discipline + +- Outcome (O1–O4), probe-signature (P1–P6), and halting (H1–H4) codes locked before runs; no post-hoc reinterpretation — `NOTES-*.md` compares measurements against the lock and reports whichever branch fired. +- Seed replication for headline cells; no hyperparameter tuning on the validation split. +- Pre-registered interpretation matrices were reviewed by external models (Gemini, Codex) before locking. + +## Blog + +The experiment is covered in a multi-part series on [jayrup.me](https://jayrup.me/) (parts 1–3, Aug 2026). + +## License -- Results are interpreted ONLY against `design/preregistration.md` (outcome codes O1–O8, H1–H4, P1–P4). -- Seed-0 runs are anecdotes until seeds {1, 2}; no hyperparameter tuning on the val set. -- Remote: `ssh://meru/~/projects/prime-grokking.git` (private; MIT license included for future publishing). +[MIT](LICENSE) diff --git a/design/preregistration.md b/design/preregistration.md index 7efa8ad..07d0791 100644 --- a/design/preregistration.md +++ b/design/preregistration.md @@ -261,3 +261,53 @@ interpretation per Addendum 4 E5, unchanged. Batch = 11 runs: e4 (2), e3 (2), ints (3), isp (4). 2-way parallel, one thread per child, same eval cadence, metrics, and early stop as all previous phases. + +--- + +## Addendum 6 — E6: dataset extension to [2, 1000] (locked 2026-08-17, pre-run) + +**Question (Q3c):** does increasing the dataset 10× (range_end 100 → 1000) change the outcome +class of next-prime learning? Power et al. (2201.02177): "smaller datasets require increasing +amounts of optimization for generalization" — the inverse prediction: more data should pull any +grokking transition earlier. E6 is the SAME protocol as phases 1–3 (E1+E2 replication) with +exactly one variable changed: range_end 100 → 1000. + +**Protocol (all cells):** digits mode, task next_prime, train_frac 0.7, lr 1e-3 const, +max_train_steps 200_000, eval_every 200, ACT halting on (halt_penalty 0.01), d_model 128. +**Changed vs phases 1–3 (recorded, unavoidable):** +- `--max_steps 32` (K). Max prime gap below 1000 is 20 (887→907), so K=20 is the exact worst + case with zero margin. K=32 covers every in-range target. Rationale logged pre-run. +- `--device cuda` + AMP fp16 autocast (T1000, ichi). fp16 numerics differ from the CPU fp32 + baseline. The O/P/H classification is precision-robust; if any cell lands within ±2 points of + a classification boundary, the tiebreak re-run is fp32 on ichi CPU. +- Outputs up to 4 digits: covered by max_out_len 6; global layout is range-derived (bisect + + margin code, tested to 100k). + +**Cells (16):** +- wd sweep {0.01, 0.1, 0.3, 1.0, 3.0} × {rnn, transformer}, seed 0 (10) +- wd 1.0 replication × seeds {1, 2} × {rnn, transformer} (4) +- batch-size ablation: wd 1.0, seed 0, batch 128 × {rnn, transformer} (2) + +**Interpretation matrix — O/P/H codes unchanged (locked Addenda 1–4). Probe range is now +[1001, 2000]. New P-ladder (sieve-rank estimation, replaces the single P1 signature):** +- A model that internalized a k-prime sieve misses exactly the composites whose prime factors + all exceed p_k; the smallest missed composite in probe range identifies k unambiguously: + 1147→k=10, 1369→k=11, 1681→k=12, 1849→k=13, none→k≥14 (exact trial division to √n). +- In-range discriminator: the 960→961 case (961=31² is the ONLY composite below 1000 with all + factors ≥ 31) separates k=10 from k=11. +- P5(k) = errors concentrated on the rank-k signature set; P6 = exact (no probe misses). +- P3 threshold for the [1001, 2000] probe: ≥90% (a rank-10 sieve scores ~97% there). +- Locked reading: P5(k=10) would confirm the "first-10-primes heuristic" hypothesis; P6 would + show full divisibility transfer; P4/P2 keep their existing meanings. +- E6 outcome vs phase 1–3: an O1 anywhere = data pressure unlocked grokking; same codes = + the task's walls are algorithmic, not data-bound. + +**Compute:** ichi (T1000 + i7-8700), tmux batch, run_sweep concurrency 4, GPU jobs only. +Benchmarks pre-run: transformer 66 steps/s (AMP), rnn 31 steps/s → batch est. 2.5–4 h. + +**E7 (next, after E6 analysis):** long-horizon test — 20× budget (4M steps) on the best E6 +cells, per user decision; protocol locked in Addendum 7 post-E6. + +--- + +*Publication note (2026-08-17): this file is mirrored verbatim from the private research repo (prime-grokking/preregistration.md) at the moment the experiment repo went public; hashes cited above are research-repo commits. The experiment repo commit history preserves each mirror step.* |