diff options
| -rw-r--r-- | blog/index.qmd | 2 | ||||
| -rw-r--r-- | blog/prime-grokking-1.qmd | 296 |
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.* |