Source

Kumar, Bordelon, Gershman and Pehlevan (2023), LIT-537 — ARXIV-2310.06110, ICLR 2024. The account, both knobs, the three conditions and the no-weight-decay counterexample below are all this paper’s.

The account

In the lazy regime a network is well approximated by its linearisation around initialization: it fits the training data by reweighting the features it started with, and learns no new ones. That is enough to drive training loss to near zero on a small dataset, and it generalises only as well as the initial neural tangent kernel does — which, if the kernel’s top eigenvectors are misaligned with the labels, is badly.

Grokking is what happens when the linear approximation breaks down and feature learning starts. The “memorising solution” of the other accounts and the lazy solution look the same from outside: both fit fast, both generalise poorly, both are followed by a long plateau. The difference is whether the network is storing examples or fitting them in a basis it has not yet updated.

Two knobs follow, and neither is a norm:

  • Laziness α, an output-scale parameter (also reachable by label rescaling). Sweeping it makes grokking more dramatic or removes it entirely.
  • Initial kernel–task alignment, ε in the toy model and centered kernel alignment in general — computable on any task, before training.

Three conditions are stated for grokking to appear: the top NTK eigenvectors are misaligned with y(x); the dataset is large enough that generalisation is eventually possible but not so large that training loss tracks test loss throughout; and the network starts lazy.

The counterexample it is built on

§3 is the load-bearing part. Modular arithmetic, two-layer MLP, no weight decay: the model groks, and the parameter weight norm rises through the transition. The polynomial-regression task behaves the same way.

Both THEORY-072 and THEORY-071 explain grokking by a late decrease in weight norm — one as a walk down to the generalizing shell, the other as norm moving from the memorising circuit to the more efficient one. Neither can produce a run with no regularizer and a rising norm. That is why this account is declared a rival to both: it and they cannot all be right about the mechanism, and neither side yet explains everything the other does (below). The declaration is about the mechanism rather than about the experiments either paper reports. It was filed as corrects at v1, before the scheme had a word for a dispute that has not been settled.

What it also gets

  • Worse initial alignment gives more intense grokking and a lower final test loss, because poor alignment is where feature learning is necessary rather than optional. “Lazy, misaligned networks grok the most intensely.”
  • Consistency across architectures, optimizers and datasets, carried from the analysable two-layer polynomial-regression setting to MNIST, one-layer transformers and student–teacher networks.
  • A methodological point the record keeps: the analysis is in loss, not accuracy, on the grounds that loss drives the dynamics and that accuracy curves on regression tasks can be gamed by the choice of metric.

Why it is Proposed

It leaves its predecessor’s two confirmed novel predictions unexplained. THEORY-071 derived ungrokking and semi-grokking before anybody had seen them, and this account says nothing about a grokked network regressing at a sharp threshold in dataset size with a weight-decay-independent endpoint. A counterexample that cannot reproduce what it argues against has narrowed the earlier account’s scope rather than replaced it, and the record files it that way: as a rival, not a correction.

The analysable setting is a two-layer network on polynomial regression; the transformer and MNIST results are reported as consistent rather than derived.

And its second condition is a dataset-size window — so the paper attacking the others hardest still agrees with all of them, and with LIT-085 and LIT-538, about the regime. That agreement is THEORY-069 and it is the only Active account here.