add blog post: prime-grokking-1

2 files changed, 298 insertions, 0 deletions

diff --git a/blog/index.qmd b/blog/index.qmd
index bf6191c..0eca72e 100644
--- a/blog/index.qmd
+++ b/blog/index.qmd
@@ -5,6 +5,8 @@ title: "Jayrup Nakawala | Blog"
 RANDOM THOUGHTS, EXPERIMENTS, AND SYSTEMS NOTES
 --------------------------------------------------
 
+- Can a Tiny Network Learn What a Prime Is? — 2026-08-14  
+  curl [jayrup.me/blog/prime-grokking-1](https://jayrup.me/blog/prime-grokking-1)
 - What the Jacobian Lens Measures — 2026-08-07  
   curl [jayrup.me/blog/jacobian-lens-frequency](https://jayrup.me/blog/jacobian-lens-frequency)
 - How this site works — 2026-08-05  
diff --git a/blog/prime-grokking-1.qmd b/blog/prime-grokking-1.qmd
new file mode 100644
index 0000000..fa14304
--- /dev/null
+++ b/blog/prime-grokking-1.qmd
@@ -0,0 +1,296 @@
+---
+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.*
+
+---
+
+**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*.
+
+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 one seed of a 69-example toy experiment. Per the
+pre-registration, seeds 1 and 2 are required before any claim beyond
+"seed-0 result". 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
+
+In rough priority order:
+
+- **Seeds 1, 2** — the mandatory replication before anything is claimed.
+- **Weight-decay sweep {0.01, 0.1, 0.3, 1.0, 3.0}** — the O2 recipe: if
+  memorization wins, turn the regularization up (and down, for a control) and
+  watch whether any setting produces the jump.
+- **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
+
+The repo is private for now — it will go public once the seed replication
+lands and the writeup is final. 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.
+
+---
+
+*This is part 1. The seed sweep is running next; whatever it shows —
+confirmation or refutation — becomes part 2.*