path: root/design/preregistration.md

# 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.

---

## Addendum 1 (2026-08-14, pre-launch — setup amendments only)

Trigger: Gemini 3.6 Flash design review (experiment repo `design/reviews/gemini-design-review.md`).
The interpretation matrix (O/H/P codes) above is NOT amended; only setup details changed. Original lock commit: 00c696d.

1. **Fully-tied cell (was: un-tied GRU decoder).** All recurrence — input read-in, K compute steps, AND output-digit decoding — now runs through the SAME 2-layer cell. The earlier draft's GRU decoder would have masked whether the tied cell solved the task. RNN param count ≈ 36.6k (was 168.7k). Regression test added (no GRU/LSTM/RNN modules).
2. **Recurrent input read-in (was: masked mean-pool).** Digits are read through the tied cell with sinusoidal positional encoding. The mean-pool blurred place value ("10" and "100" share the token multiset {1,0}).
3. **λ schedule: linear ramp 1000→5000 steps (was: hard switch at step 1000).** Avoids a discontinuous loss jump late in training.
4. **Future wd sweep revised to {0.01, 0.1, 0.3, 1.0, 3.0}** (10.0 dropped: at lr=1e-3 with AdamW, λ=10 decays weights ~1%/step). Seed-0 default wd=1.0 unchanged.
5. **ACT verification note:** aggregation is Graves (2016) standard — w_t = p_t·Π_{s<t}(1−p_s); Σw_t + remainder = 1 by construction. A reviewer initially flagged this as non-normalized; verified correct, comment + covered by tests.
6. **Prior art appended** (below). No prior grokking work on next-prime / primality prediction found — the closest literature is grokking on modular arithmetic (different: group structure) and ACT/PonderNet halting work.

## References

- Power et al. (2022), "Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets", arXiv:2201.02177
- Nanda et al. (2023), "Progress measures for grokking via mechanistic interpretability", arXiv:2301.05217
- Liu et al. (2022), "Towards Understanding Grokking: An Empirical Study", arXiv:2205.10343
- Varma et al. (2023), "Explaining Grokking Through Circuit Efficiency", arXiv:2309.02390
- Graves (2016), "Adaptive Computation Time for Recurrent Neural Networks", arXiv:1603.08983
- Banino et al. (2021), "PonderNet: Learning to Ponder", arXiv:2107.05407
- Giannou et al. (2023), "Looped Transformers as Programmable Computers", arXiv:2301.13196
- Xu et al. (ICLR 2020), "What Can Neural Networks Reason About?", arXiv:1905.13211
- Xu et al. (2021), "How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks", arXiv:2009.11848

---

## Addendum 2 (2026-08-14, pre-launch — operationalization + correctness fixes)

Trigger: OpenAI Codex code review (experiment repo `design/reviews/codex-review.md`).
The interpretation matrix is unchanged; this addendum (a) locks the code thresholds that the
prose left unquantified, and (b) records setup correctness fixes. Everything below is
pre-launch and pre-run.

### Operationalization of O/H/P codes (thresholds now locked)

- **O1:** train EM ≥ 0.95 for ≥ 10 CONSECUTIVE evals (window start `sat_start`), THEN val EM
  reaches ≥ 0.9 at eval `hi` with `hi ≥ sat_start + 10` (strictly after the window), where the
  last eval with val ≤ 0.2 (`lo`) satisfies `hi - lo ≤ 5`.
- **O2:** train EM ≥ 0.95 sustained AND val EM never reaches 0.9 AND final val EM ≤ 0.3.
- **O3:** train EM ≥ 0.95 sustained AND val EM reaches 0.9 but not via the O1 pattern.
- **O4:** train EM never sustained at ≥ 0.95 for 10 evals (setup/optimization failure).
- **O-PARTIAL** (new — the original matrix left this region unspecified): train saturated,
  val never reaches 0.9, final val EM in (0.3, 0.9). Reported as measurements, interpreted
  cautiously as partial in-range generalization; NOT retrofitted into O1/O2/O3.
- **H1:** mean steps at run end ≤ min_steps + 0.5 (collapse to the floor — see fix 2: with
  min_steps = 2 the floor is 2, so literal "collapse to 1" is impossible by construction).
- **H2:** mean steps ≥ K − 0.5 (never learned to halt).
- **H3:** min_steps + 0.5 < mean < K − 0.5 AND Pearson ρ(gap-to-next-prime, steps) ≥ +0.3
  (positive correlation = larger gap consumes more compute steps).
- **H4:** intermediate but ρ < 0.3 (halt signal noisy/unused).
- **P1:** ≥ 3 of {121, 143, 169, 187} wrong AND ≤ 6 total errors AND all errors within the flagged set.
- **P2:** scattered errors not matching P1/P3/P4.
- **P3:** probe accuracy ≥ 85% (checked after P4).
- **P4:** > 50% of trivial-composite inputs (even or multiple of 5) wrong.

### Correctness fixes (codex review)

1. **Layout invariance (BLOCKER).** Inputs are now LEFT-padded to a fixed global length
   (3 digits for [2,100]) in EVERY context — training batches, eval, greedy decoding; outputs
   right-padded to the global length (4 incl. EOS). Previously batch-max padding made an
   example's representation depend on its batchmates (RNN pad steps transformed the state;
   transformer logit positions misaligned and absolute position embeddings shifted).
   RNN pad steps are now exact no-ops. Regression tests: mixed-batch vs singleton logit
   invariance for both models.
2. **min_steps off-by-one.** Halting probs forced to 0 only for the first `min_steps − 1`
   steps (was: first `min_steps`), so the earliest halt is step `min_steps` — matching
   "execute at least min_steps steps".
3. **integers vocab EOS alias.** Integers-mode vocab is now `next_prime(range_end) + 2`, so
   the value token 101 (= next_prime(100)) is NOT aliased with EOS. Boundary test added.
4. **Rerun protection.** train.py refuses to run if `metrics.csv` exists (append + header
   would corrupt provenance). Fresh `--out_dir` required per run.
5. **Checkpoint-selection honesty.** eval reports BOTH `best.pt` (val-selected — flagged as
   selection-holed) and `last.pt` (unselected); probe + halting analyses use `last.pt`.
6. **Dead config removed:** `halt_eps` (was never read).
7. **`run_meta.json`:** python/torch/numpy versions, device, thread count recorded per run.

---

## Addendum 3 (2026-08-14, pre-launch — P1 diagnostic corrected)

Trigger: ad-hoc verification of `src/eval.py` against ground-truth stub models (promoted to
`tests/test_eval_classification.py`). Pre-launch; P1 operationalization corrected before any
results were seen.

1. **P1 concerns PREDICTIONS, not inputs.** A model that learned only the {2,3,5,7} sieve errs
   on inputs n = 113–120, 139–142, 167–168, 181–186, 199–200 — cases where the first candidate
   with no divisor ≤ 7 is composite. The signature is "predicted next prime" ∈ the flagged set,
   NOT "input n" ∈ the flagged set. Addendum 2's phrasing ("≥3 of {121,143,169,187} wrong") was
   ambiguous and the first implementation checked inputs — corrected.
2. **Set extended to {121, 143, 169, 187, 209}.** The original four are the no-small-factor
   composites ≤ 200; the probe's candidate window is actually [102, 211] (targets of
   n ∈ [101, 200] reach 211), and 209 = 11×19 is likewise mispredicted (n = 199, 200 → pred 209,
   target 211).
3. **P1 thresholds (replace Addendum 2's):** ≥ 3 errors total, ≥ 3 DISTINCT flagged predictions,
   ≥ 80% of error predictions flagged. A pure {2,3,5,7} sieve produces exactly 22 errors, 100%
   flagged (5 distinct values). P2–P4 unchanged.