diff options
Diffstat (limited to 'README.md')
| -rw-r--r-- | README.md | 69 |
1 files changed, 49 insertions, 20 deletions
@@ -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) |