PaperlayerView exact revision on arXiv

Source-linked AI summary

Grokking as the Transition from Lazy to Rich Training Dynamics

Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, Cengiz Pehlevan

arXiv:2310.06110v3stat.MLcond-mat.dis-nncs.LG

TL;DR

Grokking occurs when training loss falls well before test loss, complicating the use of training performance as a signal of generalization. The paper analyzes a two-layer network on polynomial regression and proposes that grokking arises from a shift from initial-feature fitting to later feature learning. It reports related evidence in other settings and identifies conditions and limitations on this account.

  • Problem

    Grokking separates training-loss reduction from later test generalization, challenging early stopping and explanations that attribute it to weight decay and decreasing norm.

    Source: §1 · p. 1

  • Method

    The paper combines theory and experiments to track a network's transition from fitting with initial features to later feature learning, studying laziness and initial NTK alignment.

    Source: §1 · pp. 1–2

  • Results

    The authors report grokking without regularization in polynomial regression and evidence that the lazy-to-rich transition also controls grokking on MNIST, one-layer Transformers, and student-teacher networks.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The paper concludes that grokking can reflect delayed feature learning, with its occurrence controlled by how much and how quickly features are learned.

    Source: §7 · p. 9

    The paper does not sharply characterize the sample size or training time required for transition from memorization to grokking.

    Source: §7 · p. 9

Abstract

from arXiv · show

We propose that the grokking phenomenon, where the train loss of a neural network decreases much earlier than its test loss, can arise due to a neural network transitioning from lazy training dynamics to a rich, feature learning regime. To illustrate this mechanism, we study the simple setting of vanilla gradient descent on a polynomial regression problem with a two layer neural network which exhibits grokking without regularization in a way that cannot be explained by existing theories. We identify sufficient statistics for the test loss of such a network, and tracking these over training reveals that grokking arises in this setting when the network first attempts to fit a kernel regression solution with its initial features, followed by late-time feature learning where a generalizing solution is identified after train loss is already low. We find that the key determinants of grokking are the rate of feature learning -- which can be controlled precisely by parameters that scale the network output -- and the alignment of the initial features with the target function $y(x)$. We argue this delayed generalization arises when (1) the top eigenvectors of the initial neural tangent kernel and the task labels $y(x)$ are misaligned, but (2) the dataset size is large enough so that it is possible for the network to generalize eventually, but not so large that train loss perfectly tracks test loss at all epochs, and (3) the network begins training in the lazy regime so does not learn features immediately. We conclude with evidence that this transition from lazy (linear model) to rich training (feature learning) can control grokking in more general settings, like on MNIST, one-layer Transformers, and student-teacher networks.

1 INTRODUCTION

Grokking separates falling training loss from later test generalization, challenging early-stopping assumptions and existing explanations centered on weight decay. The paper proposes delayed feature learning: networks first fit imperfectly in a linearized regime, then learn features that improve generalization.

1 INTRODUCTION

  • Grokking can arise when a network first fits its training set imperfectly with initial features, then adapts those features and generalizes later.In the linearized regime, the network behaves like a model using features evaluated at initialization.Source: §1 · p. 2
  • The transition is controlled by how quickly features are learned and how much feature learning good performance requires.The output scale α controls feature-learning rate; initial NTK alignment with test labels indicates how much feature learning is needed.Source: §1 · p. 2
  • The paper gives two-layer MLP examples of grokking without weight decay, countering theories that require weight decay or decreasing weight norm.The introduction frames this as a need for explanations consistent with empirical observations.Source: §1 · pp. 1–2
  • Grokking is proposed to arise from a transition from lazy training to rich feature learning, with parameters that control the transition.The paper says this association is supported by theory and experiments.Source: §1 · p. 2
  • Grokking requires a train-set size that permits eventual generalization without making train loss track test loss at every epoch.The paper identifies the infinite-data limit, where train loss perfectly tracks test loss, as a setting where grokking is impossible.Source: §1 · p. 3
  • The proposed dynamics are also examined in MNIST, one-layer Transformers, and student-teacher networks.The paper additionally gives a method to construct labels that can generate grokking for a fixed set of training inputs.Source: §1 · p. 3

2 RELATED WORK

The paper situates grokking among accounts based on weight decay and among work on kernel spectral bias and feature learning. It argues that feature-learning networks can outperform kernels on some tasks, motivating polynomial regression as a grokking model.

2 RELATED WORK

  • The paper reports grokking with vanilla gradient descent and zero weight decay, while prior explanations relied on weight-norm decreases or adaptive-optimizer anomalies.It notes that its theory's controlling parameters also apply with adaptive optimizers, though those dynamics are less understood.Source: §2 · p. 3
  • Kernel methods favor top eigenspace components, so targets outside them may require more data; feature-learning networks can outperform kernels by amplifying task-relevant dimensions.This motivates polynomial single- or multi-index tasks as toy models when early training is close to linearized dynamics.Source: §2 · p. 3

3 WEIGHT NORM AND WEIGHT DECAY CANNOT EXPLAIN GROKKING

The section presents grokking examples that conflict with explanations requiring weight decay or decreasing weight norm. It uses polynomial regression for mathematical analysis and reports that network laziness can make grokking more dramatic or eliminate it.

3 WEIGHT NORM AND WEIGHT DECAY CANNOT EXPLAIN GROKKING

  • A two-layer MLP groks on modular arithmetic without weight decay even as parameter weight norm increases, contradicting explanations that require norm decrease.The authors focus on loss curves because loss, rather than accuracy, guides training dynamics.Source: §3 · p. 3
  • The paper turns to polynomial regression because its grokking dynamics can be analyzed mathematically, and parameter norm also increases in that example.This motivates studying network laziness and NTK alignment as factors controlling grokking.Source: §3 · p. 4
  • Varying laziness parameter α can make grokking more dramatic or cause it to vanish.Figure 2(c) shows the effect of sweeping α.Source: §3 · p. 4, Figure 2 · p. 4

4 GROKKING AS TRANSITION FROM LAZY TO RICH LEARNING

The paper links grokking to a transition from early lazy, linearized training to later feature learning. It explains how networks can fit training data with initial features before adapting their representations and improving generalization.

4 GROKKING AS TRANSITION FROM LAZY TO RICH LEARNING

  • Grokking can reflect early lazy dynamics followed by late feature learning, as the linearized approximation breaks down and useful features emerge.The authors connect early memorizing solutions with lazy training and later generalizing solutions with richer dynamics.Source: §4 · p. 4

4.1 PRIMER ON LAZY LEARNING: LINEAR VS NONLINEAR TRAINING REGIMES

  • In the lazy regime, a network is linear in its trainable parameters and can be recast as a kernel method using the NTK evaluated at initialization.This approximation can still learn nonlinear functions of the input; some loss functions can eventually lead nonlinear models to deviate from it.Source: §4.1 · p. 4
  • Large width, large initial weights, label rescaling, and output rescaling can each promote lazy training; the paper primarily uses output rescaling.As output scale α approaches infinity, the linearization becomes exact, and under MSE the resulting model is kernel regression with the initial NTK.Source: §4.1 · pp. 4–5
  • Linearized models may generalize poorly on targets misaligned with the NTK, while feature-learning networks can adapt representations for such tasks.The paper relates this contrast to the different inductive biases of kernels and feature-learning networks.Source: §4.1 · p. 5

5 POLYNOMIAL REGRESSION IN A TWO LAYER PERCEPTRON

The paper uses a two-layer network on polynomial regression to examine how output scale and initial NTK alignment shape grokking. Its loss decomposition tracks how early use of the initial kernel gives way to later feature alignment and reduced test error.

5 POLYNOMIAL REGRESSION IN A TWO LAYER PERCEPTRON

  • The model uses isotropic Gaussian inputs, and output scale α controls feature-learning speed while ϵ changes the difficulty of the task for the initial NTK.The studied minimal two-layer network keeps its readout weights fixed at 1; the authors state that the mechanism also holds for a full two-layer MLP.Source: §5 · p. 5
  • Delayed feature learning lets the network first follow kernel-like dynamics and later align its weights with the target, producing a grokking learning curve.The toy setting is high-dimensional polynomial regression, where kernel methods require many more samples than feature-learning networks for high-degree polynomials.Source: §5 · pp. 5–6
  • Centered-kernel alignment (CKA) generalizes the toy model’s alignment measure and is used as a task-independent measure of feature alignment.In this paper, “NTK alignment” in plots refers to CKA.Source: §5 · p. 6

5.1 ALTERING SCALE α CAN INDUCE OR ELIMINATE GROKKING

  • Increasing output scale α lengthens the delay between falling train and test loss, while small α can eliminate grokking; the α →∞ limit performs poorly with the misaligned initial NTK.Changing parameter norm is not fundamental: output or label scale can also promote laziness and produce similar learning curves.Source: §5.1 · p. 6

5.2 INITIAL KERNEL-TASK ALIGNMENT ϵ CAN ALSO ALTER GROKKING DYNAMICS

  • Smaller ϵ, which worsens initial kernel alignment, produces a longer grokking delay and can yield lower final test loss by requiring feature learning.Larger ϵ gives a nearly immediate test-loss reduction.Source: §5.2 · p. 6

5.3 LOSS DECOMPOSITION IN TOY MODEL

  • The test-risk decomposition uses weight moments to track variance, target alignment, and power in linear modes; grokking’s initial test-error peak accompanies linear-mode learning before later alignment improves.Because the target is purely quadratic while the activation has a linear component, the weights must average out to remove overall linear power.Source: §5.3 · p. 7, Figure 4 · p. 8

6 GROKKING IS NOT RESTRICTED TO MODULAR ARITHMETIC

The authors test whether grokking can be induced beyond polynomial regression by changing label alignment or laziness. Their examples include Gaussian data, MNIST, and a one-layer Transformer, with grokking varying across settings.

6.1 INDUCING GROKKING ON GAUSSIAN DATA BY LOWERING ALIGNMENT OF LABELS

  • On Gaussian data, choosing intermediate initial-NTK eigenvector labels (j ≈70) produces grokking, while high alignment generalizes immediately and poor alignment prevents generalization.The labels use eigenvectors indexed j = 1, 70, and 100, respectively, in the figure’s three panels.Source: §6.1 · p. 8, Figure 5 · p. 8
  • On MNIST, varying laziness α alone can continuously control grokking, including making it vanish at α = 0.0001.The reported tuning changes α without changing parameter weight norm at initialization.Source: §6.1 · p. 9, Figure 6 · p. 9
  • In a one-layer Transformer trained on modular addition, changing α can make grokking vanish, extending the reported laziness trend to attention, AdamW, and cross-entropy.The authors also report the same trend for student-teacher networks.Source: §6.1 · p. 9, Figure 6 · p. 9

7 CONCLUSION

The paper argues that grokking can arise from a transition from lazy to rich learning, and that weight norm matters by controlling how quickly features are learned. Its results show why weight decay is not always necessary for grokking.

7 CONCLUSION

  • Polynomial regression experiments show that laziness and initial NTK alignment can induce or eliminate grokking without weight decay.The authors argue that large initial weight norm can induce grokking by approximately linearizing the model.Source: §7 · p. 9
  • The paper does not sharply characterize the sample size or training time required for a transition from memorization to grokking.Adaptive optimizers and weight decay are neither necessary nor sufficient, though they may amplify grokking.Source: §7 · p. 9

ADDITIONAL EXPERIMENTS

The additional experiments vary implementation, scale, task, dataset size, and initial alignment to examine when grokking appears or vanishes. Across these settings, the results connect grokking to training dynamics and initial feature alignment.

8.1 A NOTE ON IMPLEMENTING LARGE α

  • Large α training is stabilized by subtracting the initial predictor, and learning rates are rescaled as η = η0/α^2 to keep initial derivatives consistent.The rescaling maintains consistent early training timescales across α.Source: §8.1 · p. 13

8.2 WEIGHT NORM AND INITIALIZATION SCALE HAVE THE SAME EFFECT

  • Smaller α reduces grokking and can yield good test loss, while changing initialization scale σ produces the same trend.For a two-layer network, α ∼ σ^-2 approximately preserves feature learning; σ^2α controls relative parameter change and final test loss.Source: Figure 7 · p. 13

8.3 MODULAR ARITHMETIC IN A 2 LAYER PERCEPTRON

  • The two-layer perceptron groks on modular arithmetic with vanilla gradient descent and no weight decay across multiple prime moduli.The reported setup uses p = 23, hidden width N = 100, and 90% of possible input pairs for training.Source: §8.3 · p. 14

8.3.1 STUDENT-TEACHER NETWORKS

  • On a student-teacher task, changing α makes grokking vanish or become more pronounced without changing the student’s initial weight norm.The task uses matching 5-100-100-5 tanh MLP architectures for teacher and student.Source: §8.3.1 · p. 14, Figure 8 · p. 14

8.3.2 GOLDILOCKS DATASET SIZE

  • Dataset size yields three regimes: too little data can impair generalization, an intermediate size can grok, and larger data can support full generalization.At P = 50, train and test losses may differ at initialization; they match as dataset size increases.Source: §8.3.2 · p. 14, Figure 9 · p. 15

8.3.3 INTERPOLATING FROM RANDOM WEIGHTS TO ALIGNED SOLUTION TO MAKE GROKKING VANISH

  • Adding even a small component of the final solution to random initial weights increases alignment and makes dramatic grokking vanish.The experiment changes initialization weights in a polynomial regression task.Source: §8.3.3 · pp. 14–15

9 RELATION TO OTHER THEORIES OF GROKKING

The paper contrasts its lazy-to-rich account with prior explanations centered on task structure, weight norm, and competing circuits. It argues that weight norm can affect grokking through linearization, while dataset size is necessary but not sufficient.

Prior research themes

  • Prior work studies grokking through task-specific representations and weight norm or weight decay, while the paper notes grokking is not limited to algorithmic tasks.Examples include modular arithmetic representations and grokking induced on other tasks by tuning initial weight norm.Source: §9 · p. 16

Weight norm and weight decay

  • Modular and polynomial regression counterexamples show that grokking can occur without weight decay even when the final generalizing solution has higher parameter norm.The authors argue that initial weight norm controls grokking by linearizing the model into lazy training.Source: §9 · p. 16

Comparison with Varma et al.

  • The authors agree with Varma et al.’s two-regime account but report counterexamples to its claim that the generalizing solution has lower parameter norm.Their modular arithmetic and polynomial regression examples have higher-norm zero-loss generalizing solutions.Source: §9 · p. 16

Dataset size

Related polynomial-learning work

  • Parity learning with Gaussian data is presented as closely related to this paper’s polynomial-regression setting.For dot-product kernels, the cited passage relates parity sample complexity to learning polynomials of degree |S| with Gaussian data.Source: §9 · pp. 16–17

10 “HARD/MISALIGNED TASKS” FOR THE KERNEL CAUSE A BREAKDOWN

The authors describe hard tasks for the initial kernel and argue that linearization can break on difficult tasks, making feature learning more likely, especially with sufficiently large training sets.

10 “HARD/MISALIGNED TASKS” FOR THE KERNEL CAUSE A BREAKDOWN

  • A hard task for the initial kernel is defined through the training-label vector y and kernel Gram matrix K.The passage identifies y as the train-set label vector and K as the kernel Gram matrix.Source: §10 · p. 17
  • Near the kernel limit, the cited quantity precisely gives the amount parameters move during MSE training.The passage introduces it as one motivation for the hard-task quantity.Source: §10 · p. 17
  • For difficult initial-kernel tasks, parameter movement can make the linearization approximation break, so lazy learning is not self-consistent at infinite task difficulty.The authors expect more feature learning on difficult tasks, especially when the training set is sufficiently large.Source: §10 · p. 17

11 DERIVATIONS FOR TOY MODEL

The appendix diagonalizes the initial NTK for the Gaussian-data toy model and relates kernel-task alignment to the target direction. Its analysis shows that the kernel first learns only part of the target at smaller sample sizes, with the full target requiring substantially more data.

11.1 INITIAL NTK DIAGONALIZATION

  • At sample sizes P ≈ D, a kernel method learns only 1/(2D)|x|^2, while learning the full target requires P ≈ D^2.This separation of sample-size scales motivates the grokking task.Source: §11.1 · p. 18
  • The initial-NTK analysis identifies a sample-size separation between learning the radial component and learning the full target function.For the stated scales, the kernel learns the radial component around P ≈ D and the full target around P ≈ D^2.Source: §11.1 · p. 18

11.2 KERNEL ALIGNMENT

  • The kernel-alignment numerator increases with the alignment of M to the β⋆β⋆ direction, and this term also appears in the MSE alignment error.Source: §11.2 · p. 18

12 LOSS DECOMPOSITION WITH READOUTS

The appendix extends the toy-model analysis and tests its proposed lazy-to-rich account across optimizer choices, task variants, and measures of grokking. It finds that delayed generalization depends on task learnability and measurement: accuracy-based grokking can be misleading, while sufficient data can eventually eliminate the train-test gap.

12 LOSS DECOMPOSITION WITH READOUTS

  • With trained readout weights, the ordinary two-layer MLP retains the same grokking mechanism, and its initial kernel remains strongly biased toward linear functions.The eigenvalues associated with quadratic functions are O(ϵD^-2).Source: §12 · p. 18, Figure 11 · p. 19

13 THE ROLE OF WEIGHT DECAY

  • Increasing weight decay can eliminate grokking, while increasing laziness can recover it; the authors attribute this competition to weight decay moving the network out of the lazy regime.They also connect weight-decay-driven NTK evolution at order O(κ) to observed changes in grokking timescales.Source: Figure 12 · p. 19, §13 · pp. 19–20

14 THE ROLE OF MOMENTUM AND ADAPTIVE OPTIMIZERS

  • Grokking-like delayed generalization appears with momentum and AdamW, but adaptive-optimizer dynamics remain incompletely understood and the lazy-to-rich transition is sufficient, not necessary.The one-layer Transformer departs from its linearization early despite grokking.Source: §14 · p. 20

15 VARIANTS ON THE POLYNOMIAL TASK

  • Grokking persists in the toy setup for multi-index targets and higher-degree Hermite polynomials, rather than depending on the single-index quadratic case.The higher-degree experiments use more data to learn the target polynomials.Source: §15 · p. 20, §15.1 · p. 21, §15.2 · p. 21

16.1 DEFINING “GROKKING IN LOSS” AND “GROKKING IN ACCURACY”

  • Loss-based grokking and accuracy-based grokking can describe the same modular-arithmetic run, even though accuracy may become perfect while loss remains nonzero.In the example, train accuracy reaches perfect classification at 2k epochs while train loss continues falling; test loss later drops below about 0.06 around 3–4k epochs.Source: §16.1 · pp. 22–23

16.2 CONSTRUCTING ACCURACY CURVES FROM LOSS CURVES CAN BE MISLEADING

  • Accuracy curves on regression tasks can create apparent grokking through threshold choice, so grokking in accuracy does not imply grokking in loss.The authors therefore treat loss-based grokking as the stronger condition and use loss as their metric.Source: §16.2 · p. 23, §16.3 · p. 23, §16.4 · p. 23

16.4 LOSS IS THE RIGHT METRIC TO STUDY

  • The proposed claims about laziness and NTK alignment also apply to classification-accuracy grokking in the studied MNIST and modular-arithmetic settings.Source: §16.4 · p. 24

17 MORE DATA ALLOW US TO GENERALIZE ON HARDER TASKS

  • For a fixed task, increasing dataset size can make train and test curves move together; grokking occurs at intermediate data sizes when a task is hard but ultimately learnable.Increasing task difficulty can make a network fail at a fixed dataset size, while more data can enable eventual generalization.Source: §17 · p. 25, Figure 19 · p. 26
Loading 2310.06110v3…