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