Towards Understanding Grokking: An Effective Theory of Representation Learning

Ziming Liu; Ouail Kitouni; Niklas Nolte; Eric J. Michaud; Max Tegmark; Mike Williams

Source-linked AI summary

Towards Understanding Grokking: An Effective Theory of Representation Learning

Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, Mike Williams

arXiv:2205.10343v2May 20, 2022cs.LGcond-mat.dis-nncond-mat.stat-mechcs.AIphysics.class-ph

TL;DR

The paper asks how grokking produces delayed generalization and why it depends on training-set size and hyperparameters. It combines an effective theory of representation-learning dynamics with phase diagrams, finding that structured representations support generalization and that four learning phases organize the observed behavior.

Problem
The paper investigates how models generalize after overfitting in grokking, including the origin of generalization, the critical training size, and delayed-generalization conditions.
Method
The paper combines a toy-setting effective theory of representation-learning dynamics with empirical phase diagrams across hyperparameters.
Results
The paper finds that generalization is associated with structured representations and identifies four learning phases: comprehension, grokking, memorization, and confusion.
Takeaways & Limitations
Representation learning occurs in a Goldilocks zone comprising comprehension and grokking between memorization and confusion.
Takeaways & Limitations
The effective theory uses Euclidean embedding distances, while broader decoder-defined metrics may capture structures that this definition misses.

Abstract

from arXiv · show

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. We find on transformers the grokking phase stays closer to the memorization phase (compared to the comprehension 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.

1 Introduction

The paper frames grokking as delayed generalization after training-set overfitting and investigates its origin, critical training size, and conditions. It argues that structured input representations underpin generalization and studies these questions using an effective theory and hyperparameter phase diagrams.

Research questions: Grokking is delayed generalization in which neural networks generalize long after overfitting their training set.The phenomenon raises questions about how models generalize, why the required training time diverges near a critical training size, and when delay occurs.
Main answers: Representation learning is presented as central to explaining the origin of generalization, critical training size, and delayed generalization.The paper's answers attribute generalization to learning structured representations and relate the critical size to the data needed to determine such a representation.
Main answers: Generalization is attributed to learning input embeddings with task-appropriate structure that can be predicted by the paper's theory.The introduction points to embedding visualizations as evidence for this relationship.
Main answers: The critical training set size is the least amount of training data that can determine the relevant representation, sometimes uniquely up to linear transformations.Below that amount, the representation is not sufficiently determined by the training data.
Main answers: Grokking lies between comprehension and memorization and can be remedied through suitable hyperparameter tuning using phase diagrams.The paper uses phase diagrams to study when delayed generalization occurs.

2 Problem Setting

The paper studies algorithmic binary operations with transformer and toy-model settings, focusing on how trainable embeddings and decoders learn operation outputs. Its toy model reduces the task to decoding sums of symbol embeddings and evaluates both regression and classification.

Transformer setting: The transformer task learns a binary operation by predicting the result token in a tokenized equation using 256-dimensional learnable embeddings.A final linear layer maps transformer outputs to class logits.
Toy model: The toy model maps symbols a and b to trainable embeddings, sums them, and feeds the result to a trainable decoder MLP.The target is either a fixed random vector for regression or a one-hot vector for classification.
Toy model: The toy model can generalize despite its simplicity, while representation quality varies from RQI 0.0 with validation accuracy 0.1 to RQI 0.6 with validation accuracy 0.9.The reported examples retain perfect training accuracy in both cases.
Dataset: The toy experiments study addition datasets split into training and validation samples, with training data fraction defined as |D|/|D0|.The full addition dataset contains p(p + 1)/2 distinct unordered samples.
Visualization: For Figure 2, the decoder is evaluated across embedding-space grid points, with colors showing the highest-probability class for midpoint inputs.The axes correspond to embedding dimensions, allowing operation outputs to be read from the learned embedding and decoder visualization.

3 Why Generalization Occurs: Representations and Dynamics

The paper links generalization to structured embeddings and models representation learning with an effective loss whose predictions agree with toy-model experiments. The theory explains a critical training fraction and grokking timescale, while its Euclidean notion of structure limits its scope.

Generalization is linked to highly structured embeddings, including parallelogram structures for toy addition and circles for toy modular addition.
The paper defines representation structure through δ-parallelograms and measures how many permissible parallelograms occur in an embedding representation.
A linear representation has RQI = 1, whereas a random representation has RQI = 0 with high probability.
Predicted accuracy from the embedding representation aligns well with empirical full-dataset accuracy, supporting a connection between structured representations and generalization.
When the training set contains equal-sum sample pairs, an ideal model with an injective decoder learns corresponding parallelograms in its representation.
For addition with p = 10, theory and experiments show a phase transition around rc = 0.4 in obtaining a unique linear structure and reaching high RQI.
Above the critical data fraction, the grokking rate λ3 increases with data size, so the predicted grokking time t ∼ 1/λ3 decreases.
The effective theory reproduces neural-network representation trajectories qualitatively, but quantitative differences may reflect its omission of decoder dynamics.

4 Delayed Generalization: A Phase Diagram

The paper maps learning phases across hyperparameters and extends its representation-centered account from toy models to transformers and MNIST. These analyses show how decoder constraints and training conditions affect whether generalization is immediate, delayed, or absent.

4 Delayed Generalization: A Phase Diagram: The effective theory predicts decreasing grokking steps as data size increases, and its qualitative trend agrees with neural-network training.Grokking steps are defined as the steps needed to reach RQI > 0.95; quantitative differences may reflect the gap between effective and actual losses.
4 Delayed Generalization: A Phase Diagram: The authors vary representation and decoder learning rates and decoder weight decay to map learning performance across hyperparameters.The phase diagrams explicitly study competition between representation learning and decoder learning.
4.1 Phase diagram of a toy model: The phase diagrams contain comprehension, grokking, memorization, and confusion, with comprehension and grokking forming a generalizing Goldilocks zone.Memorization corresponds to overfitting, whereas confusion fails even to memorize the training data.
4.1 Phase diagram of a toy model: A fast decoder with slow representation learning leads to memorization, while a slow decoder with fast representation learning eventually generalizes but can take longer.The authors interpret these regimes as competition between representation learning and decoder overfitting.
4.2 Beyond the toy model: In transformers, decoder weight decay speeds generalization, while significant decoder dropout reduces generalization time below 10^3 steps and can eliminate grokking.The transformer phase diagram scans decoder weight decay and learning rate while fixing the embedding learning rate at 10^-3.
4.3 Grokking Experiment on MNIST: Grokking also occurs on MNIST, where the authors control delayed generalization by varying optimization hyperparameters.The MNIST result extends the reported phenomenon beyond algorithmic datasets.

5 Related work

The paper situates grokking among work on algorithmic circuits, mathematical-structure learning, double descent, representation learning, and physics-inspired theories of deep learning.

Related work: Prior work analyzed transformer circuits for modular addition and related grokking to phase changes in neural-network training.Earlier studies also offered speculative or informal conjectures about grokking.
Related work: Neural networks have been used to learn mathematical structures including arithmetic, geometry, knot theory, and group theory.Arithmetic-from-images work did not observe grokking, which the paper associates with abundant training data.
Related work: Grokking resembles epoch-wise double descent because generalization can improve after a period of overfitting.Prior work also found that regularization can mitigate double descent, paralleling weight decay’s influence on grokking.
Related work: The paper quantifies representation quality in simple arithmetic datasets rather than relying only on semantic interpretation or downstream-task performance.This setting permits studying representation evolution directly.
Related work: The authors connect effective theories and other statistical-physics tools with a low-level account of network dynamics and high-level learning performance.The broader physics-of-learning literature includes conservation laws and free-energy principles.

6 Conclusion

The paper argues that structured representations enable generalization and uses physics-inspired analysis to connect microscopic learning dynamics with macroscopic learning phases. It concludes that these tools may help make models more transparent and predictable.

6 Conclusion: Generalization is enabled when learned representations reflect structure in the data, in both toy models and more general settings.The paper identifies structured representations as central to its account of generalization.
6 Conclusion: An effective theory of representation-learning dynamics predicts the critical dependence of learning on training-data fraction in a toy setting.The paper presents this as a connection between low-level dynamics and high-level behavior.
6 Conclusion: Four learning phases—comprehension, grokking, memorization, and confusion—depend on decoder capacity and learning speed.These phases summarize the paper’s macroscopic analysis of learning behavior.
6 Conclusion: The authors view physics-inspired tools such as effective theories as a promising direction for making models more transparent and predictable.They connect this prospect to the safety of advanced AI systems.

Checklist

The supplied checklist and appendix passages document methodological scope, theoretical assumptions, representation-based generalization, and limitations of the toy analysis.

Scope and assumptions: The main analysis focuses on addition data with addition hard-coded in the decoder, while arguing that the framework can extend to Abelian groups.Finite Abelian groups decompose into products of cyclic groups, allowing separate treatment of each component.
Scope and assumptions: The addition operation can represent any permutation-invariant binary function through an embedding transformation followed by a decoder.The paper gives multiplication as an example via logarithm and exponential transformations.
Representation-based generalization: For p = 6, 8 training samples suffice to uniquely determine the perfect linear structure up to linear transformation, compared with 21 possible training samples.The example motivates the claim that representations lead to generalization.
Representation-based generalization: Parallelograms in the learned embedding imply equal pair sums, allowing training examples to induce correct predictions for unseen validation samples.The paper illustrates this mechanism through structured embedding relations.
Limitations: The representation-quality definition may not extend trivially to high-dimensional embeddings because useful structure can occupy only one dimension while other dimensions remain arbitrary.The authors suggest incorporating the decoder’s role into a future representation-quality measure.

E The gap of a realistic model M and the ideal model M∗

The paper characterizes realistic representations as descendants of ideal representations, with fewer structured parallelograms and more degrees of freedom. Ideal representation quality and accuracy provide upper bounds that numerical experiments verify, while training dynamics and hyperparameters influence the gap.

An ideal model’s RQI and accuracy set upper bounds for empirical RQI and accuracy, verified numerically in Figure 10.
Implicit parallelograms can be derived from explicit ones through transitive equal-and-parallel geometric relations.For Abelian groups, two parallelograms suffice to deduce a new parallelogram; for non-Abelian groups, three are needed.
The Alexander principle asserts that learned parallelograms are contained in those implied by the training data, setting upper bounds on representation quality and predicted accuracy.The paper treats this as a belief rather than a strict argument and notes rare possible violations.
Ideal representations contain more parallelograms than realistic representations, hence fewer degrees of freedom.Realistic representations are descendants of ideal ones in the representation hierarchy.
A slower decoder learning rate produces better representations, while small decoder learning rate and large batch size lead to comprehension.The toy phase diagram varies decoder learning rate and batch size; the slower-decoder setting yields higher RQI.
The framework extends from Abelian to non-Abelian groups by generalizing parallelograms and RQI, with qualitative results reported to carry over.

H.2 Numerical Results

Numerical experiments on the permutation group S3 examine structured representations, representation quality, training-time transitions, and accuracy predictions. The results show a critical training fraction and indicate that RQI-based accuracy is a lower bound rather than a complete explanation of generalization.

The S3 experiment embeds six group elements as learnable 3×3 matrices and applies hard-coded matrix multiplication before decoding.
PCA organizes the six flattened embedding matrices as a hexagon on the PC1–PC3 plane.
RQI increases with training-data fraction, with 11 random seeds evaluated at each fraction.
A phase transition occurs at training-data fraction rc = 0.5 for the steps required to reach RQI > 0.95.
Predicted accuracy d Acc is a lower bound of real accuracy Acc, implying a generalization mechanism beyond RQI.
The phase diagram varies decoder learning rate and decoder weight decay using steps to training accuracy ≥0.9 and validation accuracy ≥0.9.

I Effective theory for image classification

The effective theory extends structured-representation learning from algorithmic data to image classification. It encourages same-label representations to align, preserves representation variance to avoid global collapse, and predicts same-class collapse without necessarily producing equiangular class means.

The effective theory applies to image classification by using relational information from class labels to structure representations.Its ingredients encourage same-label samples to have similar representations while retaining separation between classes.
The effective theory is scale-invariant under f(x) → af(x), so its loss does not change under linear representation scaling.
The theory predicts same-class collapse but does not necessarily produce an equiangular tight frame for different-class means.The paper proposes investigating explicit repulsion as a future modification.
On MNIST, same-class images collapse to their class means while different-class means remain distinct.The effective loss is applied directly to learned 2D embeddings for 100 Adam steps.
The effective loss preserves Z, the variance of representations, so global collapse is avoided provably.

J Grokking on MNIST

The paper induces grokking on MNIST by reducing the training set and increasing initialization scale, then studies its phase structure and dependence on training data. The resulting experiments reproduce delayed generalization and show that regularization can shift learning out of the grokking regime.

Reducing MNIST training data from 50k to 1k samples and increasing initialization scale induces grokking in a depth-3 width-200 MLP.
The network quickly fits the training set, followed much later by improved validation accuracy.
The phase diagram contains confusion, comprehension, grokking, and memorization regions across weight decay and last-layer learning rate.Runs count as comprehension or grokking when validation accuracy reaches 60%.
Sufficient regularization can de-grok learning by changing the phase occupied by the model.
Generalization time increases rapidly when the MNIST training set falls below a certain amount.

K Lottery Ticket Hypothesis Connection

Generalization is associated with structured input embeddings that are partly present at initialization, while training selects task-relevant directions and forms the required geometry.

After generalization, embedding projections show clear structure, exceeding 99% validation accuracy compared with unstructured initialization projections.
The structured representation needed for generalization already exists partially at initialization, with training pruning unnecessary dimensions and forming parallelograms.
The proposed lottery-ticket interpretation places winning tickets in particular embedding axes or linear combinations rather than weights or subnetworks.
The first few principal components contain enough information to reach 99% accuracy, and end-of-training PCA reconstructions outperform current highest-variance axes across seeds.

L Derivation of the effective loss

The effective loss is derived from gradient-flow dynamics for a linear regression model, separating prediction error, regularization, and representation interactions. The derivation identifies attractive embedding interactions and recovers the normalized effective loss used earlier.

The effective loss is derived from gradient-flow dynamics on the actual mean-squared-error loss in linear regression.
Both model parameters and input representations are updated by gradient flow, with separate learning rates for the decoder and embeddings.
The representation dynamics are decomposed into label-independent internal interactions and label-dependent external forces, with internal interactions important for forming representations.
Pairwise sample interactions are modeled by a quadratic potential, producing a linear attractive force proportional to the representation vector.
In the adiabatic limit, the derivation averages interactions over training-data quadruples and defines a scale-invariant normalized Hamiltonian.
The normalized Hamiltonian is the effective loss used in Section 3.2.