path: root/prime-grokking/main.md

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

## Experiment log

### 2026-08-18 — E6 range extension, interim result

- Pre-registered E6 extended the same next-prime protocol from `[2,100]` to
  `[2,1000]` (K=32; 16 planned CUDA/AMP cells on ichi). In the first 8/16
  analysed cells, low-weight-decay runs improved only in-range exact match
  (roughly 77–83%) and remained O-PARTIAL/P4; no cell showed an O1 transition.
- The locked out-of-range probe `[1001,2000]` produced only 4–49 correct
  outputs out of 1000. Errors were not a truncated-sieve signature, so the
  result is no evidence of learned divisibility transfer. Correctly mapping
  `960→967` is not evidence either: 960 was in the training split.
- Measurement correction before completing classification: `src/eval.py` still
  used the old `[101,200]` probe, so E6's locked probe was run separately on
  CPU. E7 remains the planned 4M-step (20×) follow-up on the strongest E6
  cells, to distinguish a slow generalising basin from its absence.

### 2026-08-18 — E6 complete

- All 16 pre-registered CUDA/AMP cells are complete. No O1, P5(k), or P6
  outcome occurred: seven low-weight-decay cells are O-PARTIAL, nine are O4,
  and every cell is P4 on the locked `[1001,2000]` probe. The wd=1.0 three-seed
  replication is O4 in both architectures.
- Thus the 10× range extension improved in-range fitting but did not yield a
  transferable sieve/divisibility signature; E6 supports an algorithmic rather
  than data-bound wall. The batch-128 transformer was additionally unstable at
  the end of training (train EM 1.00→0.55; val EM 0.83→0.45).
- The E7 4M-step follow-up remains the discriminating test for a slow basin.
  Before it, `src/eval.py` must implement the locked `[1001,2000]` probe and
  P5/P6 ladder: its stored P fields still derive from the obsolete `[101,200]`
  evaluation.

### 2026-08-21 — E7 long-horizon protocol locked

- E7 is pre-registered as the direct duration test: four E6-derived cells
  (`wd={0.1,0.3}` × tied RNN/transformer, seed 0) extend only the training
  horizon from 200k to 4M steps. An O1 transition would make the prior nulls
  budget-limited; unchanged O-PARTIAL/O4 with no P5/P6 signature would make
  the E6 negative result robust to a 20× budget increase.
- The protocol preserves E6's `[2,1000]` digits task, K=32, and train split.
  RNN runs on ichi CPU (the sequential tied cell benchmarked faster there);
  transformer runs AMP/compiled on the T1000. Adaptive evaluation reduces
  monitoring overhead without changing the training objective.
- This test is motivated by the original grokking result's delayed
  generalization: small datasets can require far more optimization after
  overfitting. It does not assume primes have a simple closed-form pattern;
  the bounded target remains a test of whether a learned trial-division
  procedure can beat memorization.