“Towards Understanding Grokking: An Effective Theory of Representation Learning”, Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, Mike Williams2022-05-20 (, ; backlinks)⁠:

[code] We aim to understand grokking, a phenomenon where models generalize long after overfitting their training set.

We present both a ‘microscopic’ analysis anchored by an effective theory and a ‘macroscopic’ analysis of phase diagrams describing learning performance across hyperparameters.

We find that generalization originates from structured representations whose training dynamics and dependence on training set size can be predicted by our effective theory in a toy setting.

We observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion.

We find representation learning to occur only in a “Goldilocks zone” (including comprehension and grokking) between memorization and confusion. Compared to the comprehension phase, the grokking phase stays closer to the memorization phase, leading to delayed generalization.

The Goldilocks phase is reminiscent of “intelligence from starvation” in Darwinian evolution, where resource limitations drive discovery of more efficient solutions.

This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspired tools, eg. effective theories and phase diagrams, for understanding deep learning.

…Universality of phase diagrams: We fix the embedding learning rate to be 10−3 and sweep instead decoder weight decay in Figure 6b–d. The phase diagrams correspond to addition regression (b), addition classification (c) and permutation regression (d), respectively.

Common phenomena emerge from these different tasks: (1) they all include 4 phases; (2) The top right corner (a fast and capable decoder) is the memorization phase; (3) the bottom right corner (a fast and simple decoder) is the confusion phase; (4) grokking is sandwiched between comprehension and memorization, which seems to imply that it is an undesirable phase that stems from improperly tuned hyperparameters.

Figure 6: Phase diagrams of learning for the addition group and the permutation group. (a) shows the competition between representation and decoder. (b–d): each phase diagram contains 4 phases: comprehension, grokking, memorization and confusion, defined in Table 1. In (b–d), grokking is sandwiched between comprehension and memorization.