Research map / Architectures and efficiency / RNNs and normalization
Grokking as the transition from lazy to rich training dynamics
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes that grokking arises when a network transitions from lazy (kernel) training to rich feature-learning dynamics, demonstrated via polynomial regression with a two-layer MLP and controlled by output scale α and NTK-task alignment.
Problem
Explaining grokking (delayed test-loss decrease long after train loss falls) in settings without weight decay, where existing weight-norm and weight-decay theories fail.
Approach
Grokking is attributed to a network first fitting the training set in the linearized (lazy) regime, equivalent to kernel regression with the initial NTK, and later undergoing delayed feature learning; output scale α controls laziness and initial NTK-task alignment (CKA) controls how much feature learning is needed.
Key results
- Two-layer MLP without weight decay groks on modular arithmetic with weight norm increasing during training, contradicting weight-norm/weight-decay theories.
- Increasing output scale α continuously induces or eliminates grokking on polynomial regression, MNIST, and a one-layer Transformer (grokking vanishes at α=0.15 for the Transformer).
- Grokking appears for intermediate NTK-eigenvector label indices (j≈70) and vanishes when initial alignment is high or the dataset is very large.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase output scale α | Two-layer polynomial-regression network · Polynomial regression | train-test loss delay | Helped |
| Reduce output scale α | Two-layer polynomial-regression network · Polynomial regression | grokking | Hurt |
| Reduce initial-kernel alignment by using smaller ε | Two-layer polynomial-regression network · Polynomial regression | train-test loss delay | Helped |
| Reduce initial-kernel alignment by using smaller ε | Two-layer polynomial-regression network · Polynomial regression | final test loss | Helped |
| Use intermediate rather than highest or lowest initial-NTK eigenvector labels | One-hidden-layer network · Gaussian data with NTK-eigenvector labels | delayed generalization | Helped |
| Increase initial weight scale | Two-layer network · Polynomial regression | grokking delay | Helped |
| Use α = 0.0001 to encourage faster feature learning | MNIST network · MNIST | grokking | Hurt |
| Set α = 0.15 | One-layer Transformer · Modular addition | grokking | Hurt |
| Increase training-set size from too little to an intermediate size | Two-layer perceptron · Polynomial regression | delayed generalization | Helped |
| Increase training-set size into the large-data regime | Two-layer perceptron · Polynomial regression | train-test loss gap | Hurt |
| Add a small component of the final-solution weights to the initialization | Two-layer polynomial-regression network · Polynomial regression | grokking-like test-loss increase | Hurt |
| Increase weight decay from wd=1 to wd=5 | Polynomial-regression network with α=0.5 · Polynomial regression | grokking delay: wd=1 -> wd=5 | Hurt |
Surprising findings
- Worse initial kernel alignment (smaller epsilon) yields lower final test loss because the network is forced to learn features.
- Weight norm increases during grokking in the studied settings, so grokking can occur without weight norm decrease or weight decay.
Limitations
- No sharp theoretical characterization of required sample size or training time for the memorization-to-grokking transition.
- Role of adaptive optimizers (Adam, momentum) and weight decay in inducing grokking remains poorly understood and is only speculated about.
- Theory is developed for a toy polynomial regression task; generalization to other settings is supported empirically.
Experimental setup
Datasets: synthetic polynomial regression (Gaussian inputs) modular arithmetic a+b mod p MNIST Gaussian inputs with NTK eigenvector labels student-teacher synthetic data
Benchmarks: modular arithmetic grokking MNIST one-layer Transformer modular addition student-teacher networks
Metrics: train loss test MSE test accuracy CKA alignment time to grok
Baselines: weight decay explanation of grokking (Liu et al.) weight norm explanation of grokking (Liu et al.) Nanda et al. modular arithmetic transformer Varma et al. circuit efficiency account
Builds on: neural tangent kernel lazy training committee machine centered kernel alignment
Closest papers by meaning
- Grokking as a First Order Phase Transition in Two Layer Networks (ICLR 2024 · similarity 0.92)
- Grokking at the Edge of Numerical Stability (ICLR 2025 · similarity 0.91)
- Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model (ICLR 2025 · similarity 0.91)
- Grokking at the Edge of Linear Separability (ICML 2025 · similarity 0.91)
- Grokking in Linear Estimators -- A Solvable Model that Groks without Understanding (ICLR 2024 · similarity 0.91)
- Deep Networks Always Grok and Here is Why (ICML 2024 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.