diff options
| author | Void Agent <void@jayrup.hermes> | 2026-08-14 13:01:07 +0100 |
|---|---|---|
| committer | Void Agent <void@jayrup.hermes> | 2026-08-14 13:01:07 +0100 |
| commit | 6e7268b66b407ea3603fc9128805d132826a769f (patch) | |
| tree | a5a6d8156eba8f3696fae8b98387fb19e7046713 /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.md | 127 |
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). |
