diff options
Diffstat (limited to 'blog/prime-grokking-2.qmd')
| -rw-r--r-- | blog/prime-grokking-2.qmd | 176 |
1 files changed, 176 insertions, 0 deletions
diff --git a/blog/prime-grokking-2.qmd b/blog/prime-grokking-2.qmd new file mode 100644 index 0000000..eceb21a --- /dev/null +++ b/blog/prime-grokking-2.qmd @@ -0,0 +1,176 @@ +--- +title: "Still No Grokking: The Weight-Decay Sweep" +date: "2026-08-16" +--- + +*Part 1 found no grokking at weight decay 1.0, and the seed replication +confirmed it. So I did the obvious thing: swept the weight decay across +five orders of magnitude and watched. Still no grokking. The way it +fails is the result.* + +--- + +**The short version.** The grokking recipe says: if memorization wins, +turn the regularization up until the cheap algorithmic solution +outcompetes the lookup table. I re-ran both models across weight decay +{0.01, 0.1, 0.3, 1.0, 3.0} — a 50× range around the original 1.0 — with +everything else locked. No grokking at any setting. Validation accuracy +does not rise with weight decay; it's flat (transformer) to declining +(RNN). The weight-tied RNN never beats the transformer at any setting — +tying it only once at the lowest decay on the unselected checkpoint. +And the tied RNN, the architecture built to *want* the algorithmic +solution, gets *worse* the heavier the regularization gets. The opposite +of the grokking story. + +--- + +## 1. Recap: where part 1 left off + +In [part 1](https://jayrup.me/blog/prime-grokking-1) I trained two toy +models on the next-prime function: a plain transformer (407,947 params) +and a weight-tied recurrent cell with adaptive computation time (36,620 +params) — the "for-loop model" that could in principle represent trial +division cheaply. Both trained on 69 examples (n ∈ [2, 100]) with heavy +weight decay (1.0), the strong-regularization regime where grokking +lives. + +Result: no grokking. The transformer memorized and partially generalized +(peaked at 86.7% validation, decayed to 70%); the tied RNN couldn't even +hold a memorized solution (16.7% at the end). Out-of-range transfer was +zero. Then the seed replication (three seeds, identical everything else) +confirmed all of it: same locked codes in every seed. + +That leaves the obvious knob unturned. The whole *theory* of why grokking +happens is that weight decay makes the algorithmic solution cheaper than +the memorized one. Weight decay 1.0 was one point on that axis. What if +the basin lives at 0.01? Or 3.0? + +## 2. The sweep + +Five weight-decay values — {0.01, 0.1, 0.3, 1.0, 3.0} — × both models, +seed 0, everything else identical to part 1 (AdamW, lr 1e-3, 200k steps). +The 1.0 cells are the existing seed-0 runs, reused as control. The +interpretation was locked in the pre-registration before any of it ran: +if any setting produced the O1 signature (validation jumping 0.2 → 0.9 +within a few evals), that's grokking. If validation stayed flat or fell +as decay increased, that's the "memorization strongly dominant" clause. + +## 3. What happened + +``` + wd model O H P val best val last probe halt mean + 0.01 RNN O-PARTIAL H1 P4 73.3% 56.7% 1/100 2.00 + 0.01 transformer O-PARTIAL — P4 80.0% 56.7% 0/100 — + 0.1 RNN O-PARTIAL H1 P4 70.0% 56.7% 0/100 2.04 + 0.1 transformer O-PARTIAL — P4 86.7% 70.0% 0/100 — + 0.3 RNN O-PARTIAL H1 P4 56.7% 33.3% 0/100 2.48 + 0.3 transformer O-PARTIAL — P4 86.7% 73.3% 0/100 — + 1.0 RNN O4 H4 P4 36.7% 16.7% 1/100 3.12 + 1.0 transformer O-PARTIAL — P4 86.7% 70.0% 0/100 — + 3.0 RNN O4 H4 P4 20.0% 10.0% 0/100 3.00 + 3.0 transformer O-PARTIAL — P4 86.7% 70.0% 0/100 — +``` + +## 4. What it means (locked reading) + +Reading strictly against the pre-registered E2 matrix: + +1. **No O1 cell anywhere.** Ten cells, five weight-decay values, two + architectures — no sharp 0.2 → 0.9 validation transition in any of + them. The grokking jump does not exist for this task at any + regularization strength tested. +2. **Validation does not rise with weight decay.** The premise of the + "grokking at higher wd" clause is false. The transformer is flat + (80–87% best across the whole range); the RNN *declines* monotonically + as decay increases (73.3% at 0.01 → 20.0% at 3.0). This lands on the + locked clause: **memorization is strongly dominant; heavier weight + decay only hurts.** For the RNN, literally. For the transformer, it + does nothing. +3. **The tied RNN never beats the transformer at any wd.** It ties once — + at wd 0.01, on the unselected final checkpoint (56.7% = 56.7%) — and + loses everywhere else. The architecture built to represent the + algorithm cheaply is worse than the plain one at every setting. Per + the matrix, that's an architecture-negative result: weight-tying plus + ACT does not earn its keep at this scale. +4. **Halting never structured itself.** H1 (collapsed near the floor, + mean ~2.0–2.5 steps of 20) at low wd, H4 (noisy, ~3.0) at high wd. + Adaptive computation contributes nothing at any regularization + strength. +5. **Probe transfer stays exactly zero everywhere.** P4 in all ten cells. + Whatever the transformer is doing to reach 86.7% on held-out inputs + never generalizes one step outside [2, 100]. + +## 5. The part that's actually interesting + +The locked reading is the clean null. But there's a post-hoc observation +worth separating out, because it's the opposite of what the grokking +theory would predict: + +**Weight decay barely touches the transformer and progressively +destabilizes the tied cell — two different models, two different +responses to the same knob.** The transformer sits at 86.7% best from +wd 0.1 all the way to 3.0, like the regularization is a fly on its back. +The RNN is best at the *lowest* decay (73.3% at 0.01) and collapses as +decay rises — it flips from O-PARTIAL (can saturate train, partially +generalize) at wd ≤ 0.3 to O4 (can't even hold a memorized solution) at +wd ≥ 1.0. + +That's the wrong direction for the theory. The whole bet was: heavy +decay pushes toward the cheap algorithmic solution. Instead, heavy decay +pushes the tied cell *away* from even the memorized solution, while the +plain transformer is indifferent. The weight-tied architecture was +supposed to be the one that benefits from regularization; it's the one +that's most fragile to it. + +And the halting numbers tell the same story in miniature: the gate moves +from the floor (2.0 steps) at low wd to noisy (3.0–3.1) at high wd — +the network uses *slightly more* computation when its weights are being +squeezed, but it never uses it for anything structured. + +## 6. Caveats + +- **Seed 0 only for the sweep.** The wd=1.0 codes were seed-stable in + the E1 replication, but the sweep itself is un-replicated. The + qualitative picture — no O1 anywhere, RNN declines, transformer flat — + is the claim, not the exact percentages. +- **"Best" is selection-holed.** The val-best column is chosen on + validation (optimistic); the "last" column is the honest unselected + checkpoint. The architecture comparison uses the numbers both ways and + only the one tie changes. +- This is still 69 training examples on n ∈ [2, 100]. Nothing here + speaks to larger ranges — which is exactly why the range extension is + the experiment that actually discriminates. + +## 7. Where this goes next + +The weight-decay axis is now mapped and it's flat. That sharpens the +remaining candidates: + +- **Learning-rate annealing** — the late-run decay (both models peak + mid-training and roll over) is now the most interesting unexplained + behaviour. Annealing is the targeted test of whether it's optimizer + dynamics or solution instability. +- **halting=False ablation** — the informative comparison is at a wd + where the RNN can at least saturate train (≤ 0.3), since at 1.0–3.0 + it's O4 and the comparison would be muddied. +- **The range extension [2, 1000]** — the discriminator. At ~700 + training examples the lookup table stops fitting comfortably; the + algorithm must win for the task to work at all. If grokking exists for + this task, that's where it lives. + +--- + +RESOURCES +--------- + +- Part 1 — the original experiment and seed replication: + https://jayrup.me/blog/prime-grokking-1 +- Grokking (machine learning) — Wikipedia overview: + https://en.wikipedia.org/wiki/Grokking_(machine_learning) +- Code for this experiment (will be made public): + <repo link — pending> + +--- + +*Weight decay was the obvious knob. It's now a documented dead end. +Next: the range where memorization stops being an option.* |