Towards Understanding Grokking: An Effective Theory of Representation Learning

Abstract

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, e.g., effective theories and phase diagrams, for understanding deep learning.

Grokking modular arithmetic

“…These include: (i) the precise role of regularization in deep nonlinear neural networks; (ii) feature learning; (iii) the role of training data distributions in optimization dynamics and generalization performance of the network; (iv) data-, parameter-and computeefficiency of training; (v) interpretability of learnt features; and (vi) expressivity of architectures and complexity of tasks. 6)- (7). The Fourier image shows the same peak as found by GD, but also weak peaks corresponding to 2m = 6 mod 97, 2n = 6 mod 97 and m − n = 6 mod 97 that were suppresed by the choice of phases via (12).…”

“…In addition, Ref. [7] argued that grokking is due to the competition between encoder and decoder. While it is certainly true in their model, in the present case there is no learnable encoder but grokking is still present.…”

“…Here we explain why and how the solution ( 6)-( 7) works. There are two important ingredients in ( 6)- (7). The first ingredient is the periodicity of weights with respect to the indices n 1 , n 2 , q.…”

“…In this Section we show that optimization of the network ( 1)-( 2) yields a solution that is very close to the one we proposed in the previous Section. 6)- (7). The optimizer can tune phases better than random uniform distribution in order to ensure better cancellations.…”

“…In subsequent work [7], the authors simplified the architecture to a single linear learnable encoder followed by a multilayer perceptron (MLP) decoder and showed that, even if the task is recast as a classification problem, grokking persists. They also interpret grokking as a competition between encoder and decoder, and developed a toy model of grokking as dynamics of the embeddings only.…”