nucleation

LIT-858: Towards Understanding Grokking: An Effective Theory of Representation Learning

3 min read

Properties4
StatusActive — read 2026-10-09 (NOTE-662); worth reading as an early mechanistic account of when grokking networks generalise, in a toy model that sums two learned embeddings and decodes the sum. An unseen pair is predicted correctly when its embedding sum coincides with a trained pair's (a "parallelogram"), and the critical training fraction is where the training set's parallelogram equations first fix the embedding up to translation and scale. The account is exact for the toy (1D embeddings, addition hard-coded in the architecture, an effective loss posited and derived only for a linear decoder); the transformer and MNIST evidence is qualitative (PCA pictures and phase diagrams), and the four "phases" are defined by accuracy thresholds, not by any order parameter.
Tagsrepresentation-learning, learning-theory, anthology-candidate
Cited as “paper” byNOTE-662 — Towards Understanding Grokking
Cited as “source” byTHEORY-022 — The reverse-engineered grokking network computes modular addition by multiplying characters of ℤ/p at a few frequencies, the same real two-dimensional irreducibles later found as circles in language models; in both the group comes from the task, not the network, THEORY-181 — In a model that decodes the sum of two learned embeddings, at zero training loss with an injective decoder, a training set fixes the generalising representation of addition when its same-answer pairs leave only translation and scale free; for non-modular addition with p = 10 and one-dimensional embeddings, generalisation sets in near the training fraction, about 0.4, at which that becomes likely, THEORY-782 — In networks that learn addition from summed embeddings, the critical training fraction of grokking is the fraction at which the training set's same-answer pairs come to fix the generalising representation up to translation and scale, at every group size, as Liu, Kitouni, Nolte, Michaud, Tegmark and Williams propose

Graph View

Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark and Mike Williams (2022), Advances in Neural Information Processing Systems 35 (NeurIPS 2022), pp. 34651–34663 — ARXIV-2205.10343; DOI-10.52202/068431-2511

Key takeaways

  • Generalisation from parallelograms. The toy model is (a, b) ↦ Dec(E_a + E_b). If training loss is zero and the decoder is injective, every pair of training samples with the same answer forces E_i + E_j = E_m + E_n (Proposition 2), and any such coincidence must respect the arithmetic (Proposition 1). An unseen pair is then answered correctly when its sum coincides with a trained pair’s. The fraction of permissible parallelograms realised, the representation quality index (RQI), predicts accuracy closely for 1D embeddings (Fig. 3).
  • The critical training fraction is linear algebra on the dataset. The parallelogram equations a training set implies have a null space of dimension at least 2 (translation and scale). When it is exactly 2, the representation is forced to E_k = a + kb. For addition with p = 10 the probability of that jumps near a training fraction of 0.4, matching the empirical steps-to-RQI > 0.95 (Fig. 4) and Power et al.’s critical fraction (LIT-341). For S₃ with matrix embeddings it is near 0.5.
  • An effective loss. The normalised embeddings are modelled as gradient flow on ℓ_eff = ℓ₀/Z₀, the mean squared parallelogram defect over the training set divided by Σ|E_k|². It conserves Σ E_k and Z₀ (proved, App. F), so it cannot collapse. Its dynamics are linear, and 1/λ₃, the third Hessian eigenvalue, sets the time to structure: zero below the critical fraction, growing with data above it (Fig. 5). It is posited, and derived (App. L) only for a linear decoder in the limit of a frozen decoder.
  • Four phases. Comprehension, grokking, memorization and confusion are defined by thresholds: 90% train and validation accuracy within 10⁵ steps, and a gap of under 10³ steps between them. Over decoder learning rate and weight decay, grokking sits between comprehension and memorization in three toy tasks, a transformer on addition mod 53 and an MLP on a 1,000-sample subset of MNIST. Generalisation needs the representation to learn faster than the decoder, but not much faster.
  • Beyond the toy, pictures. In the transformer, generalisation coincides with a circle in the first two principal components of the embeddings (Fig. 1) and with a fall in the entropy of the PCA spectrum. Decoder weight decay and dropout shorten or remove the delay. MNIST grokking needs a reduced training set and an enlarged initialisation (v2 only; it appears after the authors’ Omnigrok, ANTH-LIT-540).

Standing in the record

Filed on 2026-10-09 at the owner’s request. It came in a batch with arXiv 2509.24914, 2502.09863, 2504.12916 and 2410.17770, with no stated context, and it is read here on its own merits.

Read on 2026-10-09 (NOTE-662). The reading is the source of THEORY-181, the record’s statement of when a training set fixes the generalising representation in this toy. It is the earliest grokking paper in the record to find the circle in modular-addition embeddings, which THEORY-022 states from Nanda et al. (LIT-345). The founding experiment is Power et al. (LIT-341).

It carries anthology-candidate: its subject is the training dynamics of neural networks, and its phase diagrams are about which hyperparameters remove the delay, which is ML practice. The anthology holds the follow-up, Omnigrok (ANTH-LIT-540), and the grokking line around ANTH-THEORY-069, but not this paper as far as its clone of commit d8b5ba5 shows; that clone may be stale.

Backlinks