2 files changed, 195 insertions, 0 deletions
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.