anthology-of-the-sota

LIT-537: Grokking as the Transition from Lazy to Rich Training Dynamics

2 min read

Properties5
StatusActive
Tagsanalysis-and-evaluation, training-optimization, model-stability, capability-thresholds
Cited as “contested by” byTHEORY-071 — A memorising and a generalising circuit compete on logits per unit norm, and which one wins flips when the data makes memorising more expensive, THEORY-072 — Training and test loss disagree as functions of weight norm, and shrinking the norm through the gap is what takes so long
Cited as “paper” byNOTE-284 — A counterexample with no weight decay and a rising weight norm
Cited as “source” bySOTA-200 — Check whether an emergent capability is a metric artefact before believing it, THEORY-069 — Grokking is a regime rather than a property of algorithmic data, and at least three knobs move it, THEORY-070 — Grokking is the transition from lazy to rich training dynamics

Graph View

Kumar, Bordelon, Gershman and Pehlevan (2023) — ARXIV-2310.06110, ICLR 2024. Read as NOTE-284.

Key takeaways

  • The counterexample, stated as §3’s whole point. Modular arithmetic, two-layer MLP, no weight decay: the model groks and the parameter weight norm increases during training. Prior accounts (ARXIV-2210.01117’s LU mechanism, ARXIV-2309.02390’s circuit efficiency) both turn on a late decrease in weight norm, so neither can be the general explanation.
  • The proposed account. Early training is close to linearised (lazy) dynamics: the network fits the training set with its initial features, so training loss falls with no test improvement. Grokking is the late breakdown of that approximation as feature learning begins.
  • Two knobs, both measurable. α, an output-scale parameter controlling laziness (also reachable by label rescaling), which in a sweep makes grokking more dramatic or removes it entirely; and ε, the alignment between the initial NTK’s top eigenvectors and the target, generalised to arbitrary tasks as centered kernel alignment.
  • Three stated conditions for grokking: the top NTK eigenvectors and the labels are misaligned; the dataset is large enough that generalization is eventually possible but not so large that training loss tracks test loss throughout; and the network starts lazy.
  • Worse initial alignment gives a lower final test loss, because it is the case where feature learning is necessary rather than optional — “lazy, misaligned networks grok the most intensely”.
  • Analysed on polynomial regression with a two-layer network, where sufficient statistics for the test loss can be tracked directly; carried over to MNIST, one-layer transformers and student–teacher setups.
  • A methodological aside the record should keep: the paper works in loss rather than accuracy because loss is what drives the dynamics, and shows in its appendix that accuracy curves on regression tasks can be gamed by the choice of metric.

Standing in the anthology

Active, and it is the reason this cluster is filed as a dispute. Two published mechanisms and an explicit counterexample to both is exactly the configuration in which filing one account alone would put a wrong thing in the record.

What it does not do is explain ungrokking or semi-grokking. ARXIV-2309.02390 predicted those before observing them, and nothing in the lazy-to-rich account addresses them. The honest reading is that neither account is general and the record says so rather than picking.

Backlinks