2026-08-14
I trained two tiny neural networks to compute the next prime after a number, to see whether either would discover the algorithm instead of memorizing the answers. Neither did. The way they failed is more interesting than the failure.
Updated 2026-08-16: seeds 1–2 replication added (section 8). The weight-decay sweep is part 2 and the diagnostics are part 3.
The short version. I asked whether a minimal neural network can learn the next-prime function as an algorithm rather than a lookup table — the phenomenon called “grokking”. I locked my interpretations in writing before running anything. Result, seed 0: no grokking from either model. A plain transformer memorized the range and partially generalized on held-out numbers (peaked at 87%, then drifted down to 70%), while the weight-tied recurrent model I built as the “algorithmic” architecture couldn’t hold a stable solution at all (ended at 17%). Neither transferred anything outside the training range — 0–1% on numbers 101–200, including trivial cases like “the next prime after an even number is odd”. Heavy weight decay did not flip memorization into algorithm. That is a clean null result, and it comes with one genuinely interesting observation: out of range, the models don’t answer randomly — they invent numbers that look prime.
Underneath almost everything in my research agenda sits one question: what makes a neural network stop memorizing and start generalizing? For a language model, “generalization” is impossible to define cleanly. So people study the cleanest possible version: tiny networks, toy tasks, and a signature phenomenon called grokking (see RESOURCES).
Grokking, from Power et al. (2022), looks like this: train a small network on modular addition (a + b mod p). The training loss collapses to zero almost immediately — the network has memorized every example. Validation stays at chance for a long, boring stretch. Then, sometimes thousands of training steps later, validation suddenly snaps to ~100%. The network didn’t gradually improve; it abruptly “got it”. Nobody changed anything about the training setup at the moment of the jump. The network was busy, internally, rearranging itself from a memorized table into the general rule.
Here’s the problem with modular addition as the test case: it’s too nice. Addition modulo p has group structure — the correct general solution is a beautiful low-complexity object (a rotation on a circle, literally Fourier components). Weight decay has an obvious “simpler thing” to push toward. Modular addition tests whether grokking can happen in the most favourable possible setting. It says little about whether the phenomenon survives contact with a task that is actually algorithmic the way real-world reasoning is.
Next-prime is that task. The function f(n) = smallest prime > n looks smooth if you only glance at it: 42 → 43, 14 → 17, 8 → 11. But there is no closed form. To compute it you must search: step through candidates n+1, n+2, … and test each for divisibility by the primes up to √n. It’s a loop, a branch, and a test — a real little algorithm. And crucially, memorization is the path of least resistance: on my range (n ∈ [2, 100]) there are only 69 training examples, and a lookup table fits trivially inside even a tiny network. There is no group symmetry, no smooth shortcut. If a network groks this, it discovered a structured algorithm in a setting where the lazy answer (remember everything) is available and easy.
Why did I think the specific architecture I tested had a chance? The bet is about weight sharing. A “trial division” algorithm is: one divisibility test, reused over and over. A weight-tied recurrent cell — one small computation module, applied repeatedly with the same weights — can represent that algorithm with a small parameter count, smaller than the memorized table needs. In principle, under heavy weight decay (a penalty on big weights), the algorithmic solution should be cheaper and win. If it doesn’t win even here — the most favourable setup I could construct — that’s evidence that something in the current toolkit is missing. Either way, the result is informative. That’s the whole experiment.
Input: a number n written in decimal digits. Output: the digits of the next prime, terminated by an end-of-sequence token. So 42 → "4 3 ⌞EOS⌟", 99 → "1 0 1 ⌞EOS⌟". n ranges over [2, 100]; 30% of the range is held out (seed 0), leaving 69 training examples. That’s the whole dataset.
The interesting arm. There is one small two-layer computation module. Everything else is that module applied repeatedly:
Total: 36,620 parameters. A lookup table of 69 examples needs a few thousand numbers’ worth of storage; the algorithm needs one reusable divisibility test.
A plain GPT-style transformer: 2 layers, width 128, 4 heads — 407,947 parameters, ~11× more. It reads the input digits and predicts the output digits causally, the usual way.
One deliberate decision worth explaining: I did not match parameter counts between the two models (I logged them instead). Matching them would require either bloating the RNN or crippling the transformer — and either way it would destroy the very thing being tested: that weight sharing makes the algorithmic solution cheaper. The two arms answer different halves of the question: can a transformer learn this range, and does weight-tying + halting get you something a plain model doesn’t. I compare behaviours, not parameter budgets.
AdamW, learning rate 1e-3, weight decay 1.0 (deliberately heavy — grokking lives in the strong-regularization regime), 200,000 steps, evaluated every 200 steps, early stop if validation ever reaches 100% for five consecutive evals (it never did). Primary metric: exact-match — the whole output digit sequence must be right, per example. Everything runs in a few hours on a 2-core laptop CPU; these models are toys, on purpose.
There’s a trap in experiments like this: you run them, look at the curves, and then decide what they mean. That’s how null results get reinterpreted into weak positives. So before a single training step, I committed a pre-registration: a matrix of outcome → interpretation codes, locked by commit hash (00c696d → b3fe898 → 3fe3935 → 3015189). The codes:
O1 sharp transition: val jumps 0.2 → 0.9 within 5 evals → grokking-like
O2 train memorizes, val stays low → memorization won
O3 val rises gradually → smooth heuristic, not grokking
O4 train never sustains saturation → setup problem, no reading
O-PARTIAL train saturated, val ends mid-range → partial generalization
H1–H4 halting: collapsed to floor / pinned at max / structured / noisy
P1–P4 probe (101–200): sieve signature / scattered / good / fails trivial casesThen, still before any run, I had two independent models review the whole design and the code. Both passes earned their place:
It also caught an off-by-one in the halting floor, and — my favourite — a subtle error in the probe diagnostic. My plan was to detect “did it learn trial division with divisors {2,3,5,7}?” by looking for errors on the composites 121, 143, 169, 187 (they need divisors 11 and 13, which don’t exist in the training range’s sieve). I wrote the check, then tested the evaluator against a fake model that literally implements the {2,3,5,7} sieve — and it didn’t fire. Two reasons: the sieve’s errors are about what it predicts (it outputs 121 as “the next prime” for inputs 113–120), not which input it was given; and the set was missing 209 = 11×19, which the sieve also outputs as “prime” (for inputs 199–200, where the true answer is 211). The verification caught it before a single real result existed. This is the boring kind of work that determines whether a result means anything.
Memorization is instant: train exact-match hits 95%+ by step 200. Then validation climbs steadily — 86.7% at step 14,200 — the best in-range generalization either model achieved. But there is no grokking jump: it never crosses 90%, and after its peak it decays, ending at 70% with the last stretch oscillating between 33% and 83%. The network found a solution that generalizes to most of the held-out range, and then training slowly walked away from it. Locked code: O-PARTIAL — train saturated, no transition, partial in-range generalization.
Harder life. Train exact-match touched 1.0 early (step 400) but couldn’t hold it — it bounced between 0.8 and 1.0 for the whole run, ending at 0.80. Validation peaked at just 36.7% (step 3,000) and also decayed, ending at 16.7%. Under the same heavy weight decay, the tied cell couldn’t even keep a memorized solution stable — let alone find an algorithmic one. The halting gate settled near its floor (mean 3.1 steps of 20) with essentially no correlation to problem difficulty (ρ = 0.22 between steps used and gap-to-next-prime). The “think harder for harder inputs” behaviour never appeared. Locked codes: O4 (train never sustained saturation — per the matrix, no scientific reading until that’s understood) and H4 (halting unstructured).
After training on [2, 100], I fed both models every number in [101, 200] — numbers they have never seen, where a genuine algorithm would keep working and a lookup table cannot. Result: transformer 0 / 100, RNN 1 / 100. Both fail even the trivial cases — for an even input, the answer is obviously odd, and neither model clears that bar. Locked code for both: P4 — “didn’t transfer even trivial heuristics — strong memorization evidence.”
The one thing worth staring at: the RNN’s out-of-range answers aren’t random. It produces 899, 999, 997, 893 — numbers ending in 99, 97, 93 that look prime. The transformer anchors on 101 (the largest target it ever saw) and produces near-misses around it: 171, 97, 87, 79, 77, 60. Neither model learned the rule; both learned the texture of the answers — the shape of “what a prime answer looks like”, detached from any computation. It’s memorization with a costume on. (This observation is post-hoc — it’s flavour, not a pre-registered finding.)
Reading strictly from the locked matrix:
Honesty note. This is a 69-example toy experiment. The numbers above are the seed-0 run; the seed replication (section 8) confirmed them, and the weight-decay sweep is in part 2. The val number for the transformer’s best checkpoint (86.7%) is also selection-holed — I used validation for early stopping, so the peak is optimistic; the unselected final checkpoint’s 70% is the cleaner number. And one engineering asterisk: the first attempt at these runs was killed by a laptop restart at 56% through the RNN, so the numbers above are from a clean re-run — same code, same seed, fresh directory. The locked interpretation is unaffected by that.
My own reading, clearly separated from the locked one: the most interesting signal in these curves is the late-run decay. Both models peaked mid-training and got worse afterward — validation didn’t plateau, it rolled over. Under constant learning rate and heavy weight decay, the memorized solution keeps exerting pull, and nothing stable holds. That, more than the null grokking result, is what I’d chase next: grokking papers typically see the jump in this regime, and I saw its opposite. If the algorithmic basin exists for this task, it’s not the one these optimizer settings settle into.
Status: the seed replication is done (section 8), and the weight-decay sweep is published as part 2. Remaining, in rough priority order:
All code, running the tests, and reproduction steps are in the repository: https://git.jayrup.me/c/prime-grokking.git/ — 34 tests cover data correctness, model shapes, halting, layout invariance, and the evaluator’s classification logic against ground-truth stub models. Each arm trains in ~2h on 2 CPU cores.
Per the pre-registration, seeds 1 and 2 were required before any claim beyond “seed-0 result”. I re-ran both models with identical hyperparameters, changing only the seed. The locked criterion: a result is replicated if it shows the same O-code and P-code across ≥ 2 of the 3 seeds (0, 1, 2), plus the same H-code for the RNN.
The answer is a clean, boring confirmation — boring in exactly the way a replication is supposed to be.
model seed O H P val best val last probe (101–200)
transformer 0 O-PARTIAL — P4 86.7% 70.0% 0 / 100
transformer 1 O-PARTIAL — P4 70.0% 43.3% 0 / 100
transformer 2 O-PARTIAL — P4 80.0% 73.3% 0 / 100
RNN 0 O4 H4 P4 36.7% 16.7% 1 / 100
RNN 1 O4 H4 P4 36.7% 20.0% 0 / 100
RNN 2 O4 H4 P4 36.7% 23.3% 0 / 100Both models hit the bar in all three seeds, not just the required two. The transformer’s O-PARTIAL + P4 is seed-stable: partial in-range generalization (val best 70–87%) with zero out-of-range transfer in every run. The RNN’s O4 / H4 + P4 is equally stable: the tied cell couldn’t hold a memorized solution (val best stuck at 36.7% across all three seeds), and its halting gate stayed unstructured near the floor (mean 2.8–3.6 steps of 20). No disagreement to report — so per the locked matrix these are now real, seed-stable findings, and the weight-decay sweep could proceed on seed 0 as planned. That sweep is part 2.
Caveats unchanged: three seeds of a 69-example toy task is replication, not a variance estimate; the “best” column is val-selected and therefore selection-holed (the “last” column is the unselected checkpoint); the late-run val decay seen in seed 0 recurs in both models; and the probe transfer is exactly zero everywhere, so there is nothing downstream to over-interpret.
This is part 1. Part 2 covers the weight-decay sweep; the range-extension experiment and the halting ablation are the natural part 3.