diff options
| -rw-r--r-- | .gitignore | 5 | ||||
| -rw-r--r-- | LICENSE | 21 | ||||
| -rw-r--r-- | README.md | 34 | ||||
| -rw-r--r-- | design/experiment-spec.md | 127 | ||||
| -rw-r--r-- | design/preregistration.md | 68 |
5 files changed, 255 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f9bf0eb --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +.venv/ +runs/ +.pytest_cache/ @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 jayrup + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..bae6cf7 --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ +# Prime Grokking + +Can minimal architectures 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. + +## Layout + +``` +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) +``` + +## Quick start + +```bash +python3 -m venv .venv +.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu 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 +``` + +## Notes + +- 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). diff --git a/design/experiment-spec.md b/design/experiment-spec.md new file mode 100644 index 0000000..c814057 --- /dev/null +++ b/design/experiment-spec.md @@ -0,0 +1,127 @@ +# Prime Grokking — Can minimal architectures learn the next-prime function? + +**Question:** If you strip a model down to digit tokens (0-9 + EOS), feed it +a single number, and train it to output the next prime in sequence — with no +chain-of-thought, no external memory, pure feedforward or minimal recurrence — +will it work? Can it *grok* the underlying algorithm rather than memorizing a +lookup table? + +## The core tension + +The next-prime function is deceptive. For small ranges it looks learnable — +skip evens, skip 5s, and you're most of the way there. But fundamentally it's +a combinatorial search problem: + +1. Iterate candidates c = n+1, n+2, ... (gap ≈ log n on average, but unbounded) +2. For each candidate, check divisibility by all primes ≤ √c +3. Output the first candidate with zero divisors + +A fixed-depth circuit cannot do an unbounded loop. So the question isn't "will +it work for arbitrary n" (it won't) — it's "can the model discover *any* +structured algorithm beyond pure memorization, and under what conditions?" + +## Why grokking is plausible (but harder than modular arithmetic) + +Modular addition has a clean group structure — one low-complexity generalizing +solution (rotate on a circle). Next-prime doesn't. But for a **bounded range** +the algorithmic solution does exist in circuit space: + +- For N ≤ 1000: √N ≈ 31, so ~11 divisibility checks per candidate, ~6 average + gap steps → ~60-70 modular operations. A transformer or RNN with enough + depth *can* represent this. +- Weight decay could push from memorization to this structured sieve. + +**Why it's harder than modular addition:** + +1. **Memorization is a stronger attractor.** The lookup table fits easily in + the weights. The algorithmic circuit's weight norm might be *larger*. +2. **Brittle loss landscape.** One mis-calibrated divisibility check → wrong + output for a whole class. Narrow basin. +3. **No guaranteed continuous interpolation path** from memorization to + generalization in weight space. + +## Architecture design space + +### 1. Weight-tied recurrence (highest-signal first experiment) + +``` +state = embed(input_number) +for step in range(max_steps): + state = step_module(state) # same weights, every iteration + if halt_condition(state): break +output = project(state) +``` + +One step module, applied repeatedly. Weight sharing means the generalizing +solution (one divisibility operation, reused K times) has *smaller* effective +parameter count than memorization. This is the mechanism that makes grokking +possible — the simpler solution wins under regularization. + +### 2. Explicit modulo gating + +Give the architecture a modulo gate directly. Not "learn integer division from +scratch" — give it access to the atomic operation and let it learn *when* and +*how* to route through it. + +### 3. Two-level recurrence (nested loop structure) + +The algorithm has nested loops: outer (candidate search) and inner +(divisibility check). A flat recurrence interleaves them awkwardly. +Consider: +- **Outer level:** advance candidate, check halt signal +- **Inner level:** iterate through divisors, check modulo +- Stack-augmented: push candidate, run inner loop, pop, advance + +### 4. Adaptive computation time (the honest answer) + +Any fixed-budget architecture will fail at some range. The real solution is +dynamic depth — run until a halt neuron fires. ACT, ponder networks, or +simply "halt when confidence exceeds threshold." + +## First experiment + +**Setup:** +- Range: n ∈ [2, 100], hold out 30% randomly +- Architecture: 2-layer weight-tied RNN cell, K=20 steps, learned halting gate +- Embedding: digit tokens + positional encoding +- Loss: standard next-token prediction on digit sequence output +- Regularization: heavy weight decay + small training set (force grokking) +- Baseline: same-parameter-count transformer + +**What to watch:** +- Training/val loss divergence curve +- If val drops suddenly after train → 0 → grokking-like behavior +- If val never drops → limitation is deeper than architecture + +**Extensions:** +- Scale range: [2, 200], [2, 500], [2, 1000] +- Vary training set size (40%, 50%, 70% of range) +- Vary weight decay magnitude +- Add explicit modulo gate vs. learned +- Test generalization: train on [2, 100], test on [101, 200] + +## Key references + +- Power et al. (2022) — "Grokking: Generalization Beyond Overfitting on Small + Algorithmic Datasets" +- Xu et al. (ICLR 2020) — "What Can Neural Networks Reason About?" + (algorithmic alignment framework) +- Graves (2016) — "Adaptive Computation Time for Recurrent Neural Networks" +- Dehghani et al. (ICLR 2019) — "Universal Transformers" +- Banino et al. (NeurIPS 2021) — "PonderNet: Learning to Ponder" + +## Why this matters beyond primes + +This is a clean, minimal test case for a much bigger question: can neural +networks discover algorithmic structure from input-output pairs alone when +memorization is the path of least resistance? Next-prime strips away all the +ambiguity — no semantics, no language, no multimodality. Just numbers and a +function that looks smooth but is structurally algorithmic. + +If grokking can't happen here, in this maximally simple setting, it's strong +evidence that current architectures need something fundamentally different to +cross the gap from pattern matching to computation. + +--- + +> Copied from the research repo `prime-grokking/main.md` @ 546dc2c (provenance: `~/Projects/research`, remote ssh://meru/~/projects/research.git). diff --git a/design/preregistration.md b/design/preregistration.md new file mode 100644 index 0000000..d879912 --- /dev/null +++ b/design/preregistration.md @@ -0,0 +1,68 @@ +# Pre-registration: Prime-Grokking — Experiment 1 (seed 0) + +**Status:** PRE-REGISTERED before any training runs. +**Date:** 2026-08-14 +**Spec:** `prime-grokking/main.md` (research repo, commit 546dc2c) +**Lock:** this file is committed to the research repo before experiment code runs; the commit hash is the lock. + +## Experiment summary + +- **Task:** map n → next prime, range n ∈ [2, 100], 30% random holdout (seed 0), tokenized as decimal digits + EOS, teacher forcing on outputs. +- **Arms:** (A) weight-tied 2-layer RNN cell, K = 20 tied steps, ACT learned halting (λ = 0.01, warmup 1000 steps); (B) GPT-style transformer baseline, d_model = 128, 2 layers, 4 heads. +- **Comparison control:** fixed d_model = 128 for both arms. NO parameter-parity gate — weight sharing is the variable under study (param-matching a tied RNN against a GPT would kill the very property the spec hypothesizes). Param counts are logged per run for the record. +- **Optimizer:** AdamW, lr = 1e-3, weight decay = 1.0 (sweep {0.3, 1.0, 3.0, 10.0} is a later experiment, not tuning on v1). Budget: 200k steps cap, eval every 200, early stop on val exact-match = 1.0 (patience 5). +- **Primary metric:** per-example exact-match accuracy on held-out inputs (token accuracy reported alongside). Per-example correctness on 10 fixed val inputs logged every eval. Mean halt steps logged every eval (RNN arm). +- **Probe:** generalization on [101, 200] is diagnostic-only (range extension is its own experiment). + +## Pre-registered outcome → interpretation matrix + +### In-range outcomes ([2, 100] held-out) + +| Code | Observable | Interpretation (locked) | Next step | +|---|---|---|---| +| **O1** | Sharp transition: train EM ≥ 0.95 sustained ≥ 10 evals, THEN val EM rises 0.2 → 0.9 within ≤ 5 evals | Grokking-like. Tied cell + heavy wd found a structured in-range solution; memorization repelled. NOT yet evidence of the full sieve — may be skip-evens/5s + divisibility heuristics. | Identify structure (per-example log, halting pattern); minimal-conditions ablations (halting=False, wd sweep, seeds) — jayrup's call. | +| **O2** | Train EM → ~1.0, val EM stays low (≤ 0.3 at run end) | Memorization won. Lookup table is the lower-norm solution under these hyperparams — consistent with spec's "memorization is a stronger attractor". | wd sweep, smaller train fraction (40%/50%), longer budget — jayrup's call. | +| **O3** | Gradual val rise to ≥ 0.9, no sharp transition | Smooth heuristic learning — NOT grokking by our operational definition. A "fast" generalizing solution exists that GD finds directly. | Distinguish from O1 by transition sharpness; report both. | +| **O4** | Train EM never ≥ 0.95 within budget | Optimization/setup failure (lr, K, halting collapse, bug). NO scientific interpretation until fixed. | Inspect losses + halt steps; fix; rerun. | + +### Architecture comparison (same hyperparams, seed 0) + +| Code | Observable | Interpretation (locked) | +|---|---|---| +| **O5** | RNN O1, transformer O2/O3 | Weight-tied recurrence is the enabling mechanism at this scale — evidence for the spec's core hypothesis (weight sharing makes the algorithmic solution cheaper). | +| **O6** | Transformer O1, RNN O2/O3 | Recurrence + halting not necessary; fixed-depth transformer suffices in-range. Weight-tying not the key variable — depth/regularization is. | +| **O7** | Both O1 | Grokking robust to architecture at this scale; the variable is data/regularization, not recurrence. | +| **O8** | Neither | See O2/O4. Next-prime may be fundamentally harder than modular addition as the spec hypothesizes; wd/data sweeps decide. | + +Caveat: single seed — all architecture comparisons are seed-0 anecdotes until seeds {1, 2}. + +### Halting structure (RNN arm only) + +| Code | Observable | Interpretation (locked) | +|---|---|---| +| **H1** | Mean steps collapse to 1 by run end | ACT failed (penalty/init issue); architecture conclusions read with a collapsed gate; fixed-K ablation becomes the informative run. | +| **H2** | Mean steps pinned at K = 20 | Never learned to halt (penalty too weak). Same caveat as H1. | +| **H3** | Mean steps intermediate (2–18), evolves during training, ideally correlates with gap-to-next-prime | Learned computation budget — evidence of structured algorithm. Check per-input steps on the 10 logged val examples. | +| **H4** | Steps fluctuate noisily | Halt signal not used meaningfully. | + +### Generalization probe [101, 200] (diagnostic-only) + +| Code | Observable | Interpretation (locked) | +|---|---|---| +| **P1** | Errors concentrated on {121, 143, 169, 187} — composites with factors 11, 13 (divisors beyond the {2, 3, 5, 7} sieve of the training range) | Definitive evidence of a learned sieve with the training-range divisor set. Strongest positive result available at this scale. | +| **P2** | Errors scattered uniformly over composites | Memorization or non-transferable heuristics; no evidence of divisibility-based algorithm. | +| **P3** | High probe accuracy beyond {121, 143, 169, 187} | Surprising — implies richer algorithm than the {2,3,5,7} sieve. Treat with suspicion; verify across seeds before claiming anything. | +| **P4** | Probe fails on ALL of [101, 200] incl. easy evens / skip-5s | In-range solution didn't transfer even trivial heuristics — strong memorization evidence. | + +### Operational definition of "grokking-like" (locked) + +- train exact-match ≥ 0.95 sustained for ≥ 10 consecutive evals (train saturated), AND +- val exact-match transition from ≤ 0.2 to ≥ 0.9 within ≤ 5 consecutive evals (eval_every = 200). +- eval order: evals happen every 200 steps; indices counted in evals, not steps. + +## Statistical hygiene + +1. Seed 0 for v1; seeds {1, 2} required before ANY claim beyond "seed-0 result". +2. No hyperparameter tuning on the val set. The wd sweep is a separate experiment run only after v1 results, at jayrup's call. +3. Probe interpretation locked above; NOTES.md must compare outcomes against this matrix verbatim (cite codes). +4. Training-range caveat recorded: for n ≤ 100 the sieve only needs divisors {2, 3, 5, 7}; "grokking the algorithm" in-range does not imply the general sieve. |