summaryrefslogtreecommitdiff
path: root/design/experiment-spec.md
diff options
context:
space:
mode:
authorVoid Agent <void@jayrup.hermes>2026-08-14 13:01:07 +0100
committerVoid Agent <void@jayrup.hermes>2026-08-14 13:01:07 +0100
commit6e7268b66b407ea3603fc9128805d132826a769f (patch)
treea5a6d8156eba8f3696fae8b98387fb19e7046713 /design/experiment-spec.md
scaffold: README, LICENSE, design docs (spec + prereg from research repo 00c696d)
Diffstat (limited to 'design/experiment-spec.md')
-rw-r--r--design/experiment-spec.md127
1 files changed, 127 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).