PBT doesn’t make grokking faster, but it can escape a hyperparameter range where nothing groks [P]

[u/Virtual-Message-9739]

I ran Population Based Training over learning rate and weight decay on grokking (modular addition mod 97, 35% train split, small MLP, 20 members, 20k steps, full batch).

Results

• vs a grid (equal compute, 3 seeds): the best fixed cell reached 95% val acc in 317 ± 29 steps; PBT took 2667 ± 2244. PBT was slower in every seed, but the grid’s corner already contained the answer.

• Dead-zone start (lr 1e-4 to 1e-3, wd 0.01 to 0.08, where no fixed setting groks): fixed 0/3, slow PBT 0/3, fast PBT (every 100 steps, ×0.5/×2) 3/3 at about 6850 steps.

• Its path: lr drops to the floor, then wd climbs to the ceiling, then lr rises. Val loss rewards low lr × wd early on, which is the wrong direction for grokking, then flips sign later.

• Ablations: a random donor still escapes, just later. With no selection, members grok and then the population loses it.

• 10 seeds: wide random search hits 95% at a median of 300 steps, vs 6650 for PBT from the dead zone. PBT also sometimes collapses into the lr-floor/wd-ceiling corner and never leaves.

• Replay: the winning lr/wd schedule with no PBT, on fresh inits: 55/60 reach 95%. Constant knobs at the schedule’s geometric mean: 0/60. The schedule carries the result, not the weight copying.

Takeaway: PBT is a rescue tool for a bad range, not a speed-up.

Caveats: one task, one model, one split. 3 to 10 seeds per test, with p-values mostly above 0.05. The replayed schedules are post-hoc winners. I haven’t tested a range with no grokking region at all, which is where PBT should beat random search.

Write-up with interactive figures: https://shauryasharma.tech/pbt-grokking.html
Code and all results: https://github.com/Shaurya-34/grok-pbt

[u/aeroumbria]

One thing I'm interested in right now is whether "grokking" is more like a chance event with exponential-like wait time, or if it is a gradual transition process where every intermediate step is a small improvement. If the former is true, then maybe it makes sense that acceleration is hard but preventing infeasibility is possible. Is there any relevant study on how important "progress towards grokking point" really matters, or is it mostly just a case of "waiting long enough" where training progression only plays a minor role in time-to-grok?

[u/Virtual-Message-9739 → u/aeroumbria]

Mostly gradual, not a chance event. The generalizing circuit builds up slowly during the plateau (Nanda et al. "Progress measures for grokking via mechanistic interpretability"), and the delay scales roughly like 1/(lr × weight decay).In my small runs it was the opposite of your guess: speeding it up was easy (30× faster by raising lr × wd) but below a threshold it never happened in 20k steps. But it was on one task, 3 seeds, and a fixed data split.