path: root/blog/prime-grokking-1.qmd

---
title: "Can a Tiny Network Learn What a Prime Is?"
date: "2026-08-14"
---

<nav>
[Home](https://jayrup.me) | [Blog](https://jayrup.me/blog) | [Projects](https://jayrup.me/projects) | [CV](https://jayrup.me/cv) | [Dissertation](https://jayrup.me/dissertation) | [Git](https://git.jayrup.me)
</nav>


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