1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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).
|
