--- title: "Can a Tiny Network Learn What a Prime Is?" date: "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](https://jayrup.me/blog/prime-grokking-2) and the diagnostics are [part 3](https://jayrup.me/blog/prime-grokking-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. --- ## 1. Why this experiment exists 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. ## 2. What I built ### The task, concretely 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. ### Model A — the weight-tied recurrent cell ("the for-loop model") The interesting arm. There is *one* small two-layer computation module. Everything else is that module applied repeatedly: - **Read-in:** each input digit is fed into the cell, one at a time, building up a state vector. - **Compute:** the cell is applied to the state K = 20 times — 20 "steps of thought", always the same weights. Think of it as a for-loop with a fixed trip count. - **Halting (ACT):** on top of this sits a learned gate — at each step, a small readout of the state asks *"am I done thinking?"* The gate's answers are used as weights to average over the states, and using more steps costs a small penalty (ramped in over training, so it doesn't collapse to "never think" on day one). This is the adaptive-computation-time idea from Graves (2016). The dream result: the network uses *more steps for harder inputs* (bigger prime gaps need more candidate checks). - **Decode:** output digits are produced by feeding each previous digit back through the *same* cell and reading out logits. 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. ### Model B — the transformer baseline 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. ### Training 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. ## 3. How I kept myself honest 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 cases ``` Then, still before any run, I had two independent models review the whole design and the code. Both passes earned their place: - **Gemini's design review** killed a bug in my own thinking: my first draft of the RNN had an *un-tied GRU decoder* — a small separate network that read out the digits. Fine on its own, but it would have *secretly* solved part of the task with its own unshared weights, muddying the very claim being tested (that the *tied* cell does the work). I rebuilt the model so read-in, compute, *and* decode all use the one cell. It also pushed me from a hard penalty switch-on at step 1000 to a smooth ramp. - **Codex's code review** found the scariest bug of the whole project: **batch-layout dependence**. Inputs were being padded to the length of the longest input in their batch — so a 1-digit number batched with a 3-digit number went through *extra computation steps* compared to the same number on its own. Training and evaluation were seeing different things. Silent, and it would have corrupted every number in every plot. I moved to a fixed global layout (left-padding to 3 digits everywhere, pad steps as exact no-ops) and wrote regression tests that check a batched example produces bit-identical outputs to the same example alone. 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. ## 4. What happened ### The transformer 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. ### The tied RNN 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). ### The probe: out of range, everything collapses 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.) ## 5. What it means Reading strictly from the locked matrix: 1. **No grokking.** Neither model shows the O1 signature. At this scale, with these settings, the next-prime function did not produce the delayed generalization jump that modular arithmetic famously does. 2. **The transformer is the better memorizer-and-partial-generalizer** (70% final vs 17%), despite the tied RNN being the architecture built for the algorithmic solution. 3. **Out-of-range transfer is zero.** P4 is the strongest statement available at this scale: whatever solved 70–87% of the held-out range did not survive a change of input range, not even the trivial "evens → odd answer" heuristic. 4. **The halting gate did nothing useful (H4).** Adaptive computation didn't emerge on its own; it collapsed toward the floor. That's its own small result: ACT without stronger shaping didn't produce structure here. **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](https://jayrup.me/blog/prime-grokking-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. ## 6. Where this goes next Status: the seed replication is done (section 8), and the weight-decay sweep is published as [part 2](https://jayrup.me/blog/prime-grokking-2). Remaining, in rough priority order: - **Learning-rate annealing** — a targeted experiment against the late-run decay specifically. - **Smaller training fraction (40% / 50%)** — fewer examples, more pressure to generalize. - **halting=False ablation** — was ACT actively hurting the tied RNN, or just useless? - **The range extension [2, 1000]** — the experiment that actually discriminates: at 700 training examples the lookup table stops fitting comfortably, and the algorithm must win for the task to work at all. That's where this question gets its real answer. ## 7. Reproduce 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. ## 8. Update: the seed replication 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 / 100 ``` Both 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](https://jayrup.me/blog/prime-grokking-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. --- RESOURCES --------- - Grokking (machine learning) — Wikipedia overview: https://en.wikipedia.org/wiki/Grokking_(machine_learning) - Code, pre-registration, and design docs for this experiment: https://git.jayrup.me/c/prime-grokking.git/ - Part 2 — the weight-decay sweep: https://jayrup.me/blog/prime-grokking-2 --- *This is part 1. Part 2 covers the weight-decay sweep; the range-extension experiment and the halting ablation are the natural part 3.*