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.
A Basin-Selection Perspective on Grokking via Singular Learning Theory
Department of Computer Science（计算机科学系）；Department of Mathematics（数学系）；Mathematical Institute（数学研究所）；University of Pisa（比萨大学）；Imperial College London（伦敦帝国学院）；University of Oxford（牛津大学）；Section of Mathematics and Artificial Intelligence（数学与人工智能系）；Max Planck Institute of Molecular Cell Biology and Genetics（马克斯·普朗克分子细胞生物学和遗传学研究所）；Center for Systems Biology（系统生物学中心）；Faculty of Mathematics（数学系）；TU Dresden（德累斯顿技术大学）
Two Speeds of Learning: A Representation-Readout Decomposition of Grokking and Double Descent
Center for Computational Neuroscience（计算神经科学中心）；Flatiron Institute（Flatiron研究所）；Department of Physics（物理系）；Harvard University（哈佛大学）；Kempner Institute（Kempner研究所）；New York University（纽约大学）
A Spectral Theory of Grokking: Weight Decay induces Feature Learning
Arnold-Sommerfeld-Center for Theoretical Physics（阿诺德-索末菲理论物理中心）；Ludwig-Maximilians-Universität München（慕尼黑路德维希-马克西米利安大学）；University of Cambridge（剑桥大学）；Standford University（斯坦福大学）
On the Learning Dynamics of Two-layer Linear Networks with Label Noise SGD
Sch. of Computer Science and Zhiyuan College, Shanghai Jiao Tong University（上海交通大学计算机科学学院和智远学院）；Peking University（北京大学）；University of Sydney（悉尼大学）；RIKEN Center for Advanced Intelligence Project（理化学研究所先进智能项目中心）；The Institute of Statistical Mathematics（统计数学研究所）；The University of Tokyo（东京大学）
The Optimization Landscape of SGD Across the Feature Learning Strength
Harvard University（哈佛大学）；School of Engineering and Applied Science, Harvard University（哈佛大学工程与应用科学学院）；University of California, Berkeley（加州大学伯克利分校）；Center for Brain Science, Harvard University（哈佛大学脑科学中心）；Kempner Institute, Harvard University（哈佛大学肯普纳研究所）
The Importance of Being Lazy: Scaling Limits of Continual Learning
Electrical Engineering, ETH Zurich, Switzerland（电子工程系，苏黎世联邦理工学院）；Dept. of Computer Science, ETH Zurich（计算机科学系，苏黎世联邦理工学院）；ETH AI Center（苏黎世联邦理工学院人工智能中心）；Astronomy, University of Padua, Italy（天文学系，帕多瓦大学）
Adaptive kernel predictors from feature-learning infinite limits of neural networks
John A. Paulson School of Engineering and Applied Sciences, Harvard University（哈佛大学约翰·A·保罗森工程与应用科学学校）；Center for Brain Sciences（脑科学中心）；Kempner Institute（凯普纳研究所）
Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks
John Paulson School of Engineering and Applied Sciences（约翰·保尔森工程与应用科学学院）；Center for Brain Science（脑科学中心）；Kempner Institute for the Study of Natural & Artificial Intelligence（自然与人工智能研究肯普纳研究所）；Harvard University（哈佛大学）
Emergent Riemannian geometry over learning discrete computations on continuous manifolds
School of Informatics University of Edinburgh（信息学院爱丁堡大学）；Département d’Etudes Cognitives Ecole Normale Supérieure PSL（认知研究部巴黎高等师范学院PSL）；Gastsby Unit University College London（加斯特比单位伦敦大学学院）；Data Science Center Ecole Normale Supérieure PSL（数据科学中心巴黎高等师范学院PSL）
