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

---
title: "Where Exactly Is the Wall?"
date: "2026-08-17"
---

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


*The first three phases established that a tiny network cannot grok the
next-prime function, and that weight decay doesn't help. This post is the
forensic part: four pre-registered diagnostics to find out where the wall
actually is — the optimizer, the halting gate, the search loop, or the way
the numbers are written down.*

---

**The short version.** The one thing the first three phases couldn't
explain is *why* both models peak mid-training and then roll over. So I
ran four pre-registered diagnostics, each aimed at one candidate culprit,
each read strictly against a locked clause. None of them produced a
grokking transition, and the rollover survived all of them. But they
narrowed the wall down nicely: it's not the learning-rate schedule, it's
not the halting gate, and it's *definitely* not the digit representation.
The one real surprise is that taking away the decimal digits — feeding
each number as an atomic token — destroyed the task completely. The
digits weren't a tax; they were the scaffolding. This is the post where
the story goes from "it doesn't work" to "here's roughly where it
breaks."

---

## 1. The question forensics finally get to ask

Parts 1 and 2 gave a clean null: no grokking at any weight decay, across
three seeds. But a null is only as good as the understanding behind it,
and there was one behaviour none of the curves explained: **both models
peaked mid-training and then got worse.** The transformer hit 86.7%
validation, then decayed back to 70%. The tied RNN couldn't even hold a
memorized solution. Validation didn't plateau — it rolled over.

That rollover was the loose thread. Two candidate explanations, one
architectural, one representational. The four diagnostics each isolate a
single variable in the training setup:

- **E3 — learning-rate annealing.** If the rollover is optimizer
  dynamics (the constant learning rate beating up the solution), a
  cosine schedule should fix it.
- **E4 — halting=False.** If the adaptive-computation-time gate is
  actively hurting the RNN, removing it should help.
- **D1 — is_prime.** Split the task in half. If the model can classify
  prime/composite (a pure divisibility decision, no search) but can't
  find the next prime (which requires search), the wall is the
  increment loop, not the divisibility test.
- **D2 — integer tokens.** Feed each number as one atomic token instead
  of its digits. If place-value parsing was the bottleneck, the model
  should jump.

Each was locked in the pre-registration with an explicit "if X then Y"
clause before any of them ran.

## 2. E3 — annealing the learning rate

Cosine schedule (1e-3 → 1e-4) on both models, everything else identical
to the wd=1.0 control. The locked clause: final validation ABOVE the
control and no late-run decay → the rollover is optimizer dynamics.
Identical or below → solution instability.

**RNN:** annealing changed nothing. Final val 0.367 = 0.367, last 0.167
= 0.167, same O4 code. As close to identical as it gets.

**Transformer:** the annealed run reached a new *best* (0.900 vs 0.867)
but its *final* validation *fell below* the control (0.567 vs 0.700) —
the rollover actually deepened, a 33-point drop from peak vs the
control's 17-point drop.

Neither model cleared the bar. A gentler schedule didn't stop the decay.
**The rollover is not an optimizer artifact.** The solution itself is
unstable — the memorization attractor re-asserts itself no matter how
the learning rate is shaped.

## 3. E4 — removing the halting gate

The tied RNN has an adaptive-computation-time gate that was supposed to
let it "think longer" on hard inputs. It never did — it collapsed to the
floor. The question: was it a harmless passenger, or actually dragging
the model down? I ran the RNN with `halting=False` (fixed K=20) at the
two weight decays where the comparison is clean.

```
  wd    run      val best  val last  O          probe
  0.1   Fixed-K  0.767     0.367     O-PARTIAL  0 / 100
  0.1   ACT      0.700     0.567     O-PARTIAL  0 / 100
  1.0   Fixed-K  0.433     0.167     O4         0 / 100
  1.0   ACT      0.367     0.167     O4         1 / 100
```

On the *best* checkpoint, removing the gate helped (+6.7 points at wd
0.1, +6.6 at 1.0). On the *final* checkpoint it didn't (—20 points at
0.1, tied at 1.0). The locked clause keys off val EM, which the two
metrics disagree on — so this one is genuinely metric-dependent.

Either way, the important part is what *didn't* happen: **removing the
gate did not stop the rollover.** Both Fixed-K runs still decayed
(0.767 → 0.367, 0.433 → 0.167), and the wd 1.0 cell stayed O4 even with
no gate to blame. The halting gate was at most a minor tax on the best
checkpoint. It was never the cause of the instability.

## 4. D1 — the is-prime decomposition

The most informative diagnostic. Instead of "what's the next prime
after n?" I asked the models a purely binary question: "is n prime?"
Same digit input, same architecture, same everything — but the output is
a single 1/0, so there's no search loop, no increment, no comparing
candidates. Just a divisibility classification. If the search loop was
the wall, this should be dramatically easier. If the divisibility
operation itself is unlearnable, it should fail like next_prime did.

```
  wd    model        val best  val last  O          P   probe acc
  1.0   RNN          0.900     0.800     O3         P2  0.60
  1.0   transformer  0.900     0.867     O3         P2  0.80
  0.1   RNN          0.867     0.800     O-PARTIAL  P2  0.79
  0.1   transformer  0.900     0.833     O3         P2  0.75
```

Two facts, cleanly separated.

**Fact 1: is_prime is much easier in-range.** Val best 0.867–0.900
everywhere, versus next_prime's 0.367 (RNN) / 0.867 (transformer) at the
same wd 1.0. The search/increment loop was a genuine tax on in-range
learning. Removing it, both models hit ~90% — the best in-range
generalization this whole project has seen.

**Fact 2: there is still no out-of-range divisibility algorithm.** The
probe is classification accuracy on [101, 200], and it sits at 60–80% —
well below the ~96% that a trivial {2,3,5,7} sieve scores on the same
range. And the error pattern is P2: scattered, primes misclassified as
composite, with **no** concentration on the composites that need
divisors 11 and 13. A learned sieve would err there. Nothing does.

So the wall is precisely this: **the in-range difficulty lived in the
search loop, but the divisibility operation still didn't emerge as an
algorithm even when it was the whole task.** We removed the search and
the in-range number went to ~90%, but out-of-range there's still nothing
general about it.

## 5. D2 — the surprise: integer tokens destroy the task

The last diagnostic inverted everyone's hypothesis. The reviewer
feedback suggested digits might be a tax — place-value parsing eating
learning capacity. So I fed each number as a single atomic token: `42`
becomes one embedding, not two digit embeddings. No positional
structure, nothing shared between similar numbers.

It collapsed. Completely.

```
  mode      model        val best  val last  O
  integers  RNN (wd 1.0) 0.133     0.000     O2
  integers  transformer  0.100     0.067     O2
  integers  RNN (wd 0.1) 0.067     0.000     O2
```

All three runs are **O2 — pure memorization** — with validation at 0–13%.
The transformer that reaches 86.7% in digits mode can't get past 10%
with atomic tokens. The locked clause said "within ±10 points → parsing
was not the bottleneck." That branch was written expecting a flat
outcome; the measured outcome is a total collapse, 20–77 points below
the digits control.

The post-hoc reading, clearly separated from the locked one: **the digit
representation wasn't a tax — it was the scaffolding.** In digits mode,
the model shares structure across 42, 43, 4, 3 — the tens digit, units
digit, positions are reusable pieces. In integers mode, every number is
a unique, once-seen token with nothing shared between any two inputs, so
there is *nothing to generalize*. The in-range heuristic that reached
87% was built out of the shared digit substructure. Remove the digits
and you remove the substrate any generalization was built on.

This is the cleanest negative result of the whole project, and it's
backwards from what anyone predicted: **compositionality wasn't the
obstacle to learning the algorithm — it was the precondition for
learning anything at all.**

## 6. What it all adds up to

Reading the four diagnostics together:

1. **The rollover is solution instability** (E3) — not an optimizer or
   learning-rate artifact, and it survives annealing.
2. **ACT wasn't the cause** (E4) — the halting gate was a minor tax at
   best; removing it didn't stop the decay.
3. **The in-range difficulty was mostly the search loop** (D1) — pure
   divisibility classification reaches ~90% in-range.
4. **But the divisibility operation itself never generalizes** (D1) —
   no out-of-range structure, well below the trivial-sieve floor.
5. **The digit representation is load-bearing** (D2) — it was the
   scaffold for any in-range generalization, not a penalty.

The picture that emerges: this is a task where the *specific* local
heuristics (skip evens, check small divisors, remember common gaps) are
easy to acquire and genuinely useful in-range — enough to reach 70–90%
on the holdout — but the *general* algorithm (search + a real sieve)
never emerges. The models learn the shape of the answer without ever
learning the rule. That's been the story since part 1; the diagnostics
just pinned down which parts of the machinery are and aren't
responsible.

## 7. Where this goes next

All four diagnostics point the same direction. The in-range heuristics
saturate the task, the general algorithm stays out of reach, and the
reason is that memorization-plus-local-heuristics is a robust attractor
that a 69-example task simply never strains.

That makes the next phase the obvious one it's been building toward: **the
range extension to [2, 1000].** At ~700 training examples, the lookup
table stops fitting comfortably and the heuristics stop covering the
holdout. That's the pressure that could push the system toward the
actual algorithm — the experiment that this whole sequence has been
setting up.

The scaffolding for it is done — the code now runs on GPU (CUDA, AMP
fp16, torch.compile), with a bisect-based prime search and scaling
support. Part 4 covers it when that run lands.

---

RESOURCES
---------

- Part 1 — the original experiment and seed replication:
  https://jayrup.me/blog/prime-grokking-1
- Part 2 — the weight-decay sweep:
  https://jayrup.me/blog/prime-grokking-2
- 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/

---

*Four diagnostics, no grokking, one real surprise: the digit
representation was the scaffolding all along. Next: make memorization
physically impossible.*