PBT rescued grokking but did not speed it up
A Reddit experiment tested Population Based Training (PBT) on modular addition mod 97 using a small MLP, with learning rate and weight decay adapted across 20 members. In a reported “dead zone” where no fixed setting grokked, fast PBT succeeded in 3 of 3 runs at about 6,850 steps, while fixed and slow PBT succeeded in 0 of 3. Under equal compute, however, the best fixed grid cell reached 95% validation accuracy in 317 ± 29 steps, compared with 2,667 ± 2,244 for PBT.
The reported results distinguish PBT’s role as a possible rescue mechanism for an unfavorable hyperparameter range from a general-purpose accelerator, while the single-task design and limited seeds constrain broader conclusions.
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
#Population Based Training#grokking#modular addition mod 97#multilayer perceptron