To Grok Grokking: Provable Grokking in Ridge Regression

To Grok Grokking：岭回归中可证明的 Grokking

Mingyue Xu Affiliation:  Department of Computer Science, Purdue University, West Lafayette, IN, USA Correspondence to: xu1864@purdue.edu    Gal Vardi Affiliation:  Department of Computer Science and Applied Mathematics, Weizmann Institute of Science, Israel    Itay Safran Affiliation:  Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel

Abstract

摘要

We study grokking—the onset of generalization long after overfitting—in a classical ridge regression setting. We prove end-to-end grokking results for learning over-parameterized linear regression models using gradient descent with weight decay. Specifically, we prove that the following stages occur: (i) the model overfits the training data early during training; (ii) poor generalization persists long after overfitting has manifested; and (iii) the generalization error eventually becomes arbitrarily small. Moreover, we show, both theoretically and empirically, that grokking can be amplified or eliminated in a principled manner through proper hyperparameter tuning. To the best of our knowledge, these are the first rigorous quantitative bounds on the generalization delay (which we refer to as the “grokking time”) in terms of training hyperparameters. Lastly, going beyond the linear setting, we empirically demonstrate that our quantitative bounds also capture the behavior of grokking on non-linear neural networks. Our results suggest that grokking is not an inherent failure mode of deep learning, but rather a consequence of specific training conditions, and thus does not require fundamental changes to the model architecture or learning algorithm to avoid.

我们在经典的岭回归设定下研究 grokking——也就是过拟合之后很久，泛化才开始出现。我们证明：用带权衰减的梯度下降去学习过参数线性回归模型时，可以得到端到端的 grokking 结果。具体来说，我们证明会出现如下阶段：(i) 模型在训练早期就拟合了训练数据；(ii) 过拟合出现之后，很差的泛化还会持续很长时间；(iii) 泛化误差最终会变得任意小。此外，我们从理论上和实验上都表明，只要把超参数调到位，grokking 可以被有原则地放大，也可以被消掉。据我们所知，这是第一组用训练超参数写出泛化延迟（我们称之为「grokking 时间」）的严格定量界。最后，走出线性设定，我们还用实验表明，这组定量界也能刻画非线性神经网络上的 grokking。这些结果说明，grokking 并不是深度学习固有的失效模式，而是特定训练条件的产物，因此不必先对模型结构或学习算法做根本改动才能避免。

Keywords:

关键词：

1 Introduction

1 引言

By the standard of classical machine learning, overfitting the training data is usually believed to be harmful to generalization. To overcome this, regularization techniques such as early-stopping, dropout and weight decay have been developed. However, deep learning can sometimes exhibit counterintuitive phenomena that contravene this traditional wisdom. One of the most striking examples is “grokking” (Power et al., 2022), a phenomenon where generalization only starts improving long after the model achieves perfect training performance. The phenomenon of grokking has been extensively studied in recent years (Chughtai et al., 2023; Tan and Huang, 2023; Notsawo Jr et al., 2023; Fan et al., 2024), and was also encountered in models that are different from neural networks (Blanc et al., 2020; Humayun et al., 2024; Merrill et al., 2023), as well as in large language models (Zhu et al., 2024; Wang et al., 2024; Xu et al., 2025) (see Subsection 1.1 for further discussion). But despite its increasing popularity, only a few prior works have established rigorous theoretical guarantees.

按经典机器学习的标准看法，把训练数据过拟合通常被认为不利于泛化。为此，人们发展了早停、dropout 和权衰减等正则化技术。不过，深度学习有时会出现与这一传统看法相悖的反直觉现象。最醒目的例子之一就是「grokking」（Power et al., 2022）：模型早已达到近乎完美的训练表现之后很久，泛化才开始变好。近年来，grokking 已被广泛研究（Chughtai et al., 2023; Tan and Huang, 2023; Notsawo Jr et al., 2023; Fan et al., 2024），也在不同于神经网络的模型里出现过（Blanc et al., 2020; Humayun et al., 2024; Merrill et al., 2023），并在大语言模型中被观察到（Zhu et al., 2024; Wang et al., 2024; Xu et al., 2025）（更多讨论见第 1.1 小节）。尽管它越来越受关注，此前只有少数工作给出了严格的理论保证。

Several existing theoretical papers attribute the occurrence of grokking to a transition in the optimization dynamics from the lazy to the rich regime (Lyu et al., 2024; Mohamadi et al., 2024; Kumar et al., 2024). Among these works, Lyu et al. (2024) consider training homogeneous neural networks for both classification and regression problems using gradient flow with weight decay, and prove a sharp transition between the kernel and rich regimes. However, their technique only guarantees convergence to a KKT (Karush-Kuhn-Tucker) point, which is, in general, not sufficient for arguing global optimality, and their result does not provably imply grokking. In another work, Mohamadi et al. (2024) provide a theoretical foundation for grokking on the specific regression problem of modular addition. They consider training two-layer quadratic networks with an $\ell_{\infty}$ regularization, and show that grokking is a consequence of the transition from kernel-like behavior to the limiting behavior of gradient descent (GD). Yet, their analysis does not establish that GD converges to a solution with a small weight norm that generalizes well, despite empirically verifying this. To the best of our knowledge, the work closest to showing a provable grokking result is Xu et al. (2024). By studying a high-dimensional clustered XOR dataset in a binary classification setting, they show that after a catastrophic one-step overfitting—where generalization is no better than a random guess—the model eventually achieves perfect test accuracy. Still, they do not show that the onset of generalization is delayed beyond the first iteration of GD, which could begin as early as the second iteration. While many prior works attribute grokking to the transition between the lazy and rich regimes during optimization, a recent work by Boursier et al. (2025) focused specifically on studying the role of weight decay in grokking, revealing that grokking can also arise from the transition between the “ridgeless” to the “ridge” regimes. They considered training using gradient flow with weight decay on a smooth loss objective and showed a two-phase behavior of the trajectory. Nevertheless, they do not prove that the unregularized solution does not generalize and the ridge solution does generalize, and thus their results do not imply provable grokking. Relatedly, Tikeng Notsawo et al. (2025) study a general loss-plus-regularizer framework, showing an early loss-minimization phase followed by a regularization-driven phase on a time scale proportional to $1/(\eta\lambda)$. Our work provides sharper end-to-end guarantees in the specific setting of ridge regression.

已有若干理论论文把 grokking 的出现归因于优化动力学从懒惰（lazy）区走到富裕（rich）区的转变（Lyu et al., 2024; Mohamadi et al., 2024; Kumar et al., 2024）。其中，Lyu et al. (2024) 考虑用带权衰减的梯度流训练齐次神经网络，同时覆盖分类和回归，并证明了核方法（kernel）区与富裕区之间存在尖锐转变。不过，他们的方法只保证收敛到 KKT（Karush-Kuhn-Tucker）点，一般并不足以论证全局最优，其结果也不能在可证明的意义上推出 grokking。另一项工作中，Mohamadi et al. (2024) 为模加法这一特定回归问题提供了 grokking 的理论基础。他们考虑用 $\ell_{\infty}$ 正则训练两层二次网络，并表明 grokking 来自从类似核方法的行为走到梯度下降（GD）极限行为的转变。可是，尽管实验上看到了这一点，他们的分析并没有证明 GD 会收敛到权范数较小、因而泛化良好的解。据我们所知，最接近给出可证明 grokking 结果的是 Xu et al. (2024)。他们在二分类设定下研究高维聚类 XOR 数据，证明在灾难性的一步过拟合——此时泛化并不比随机猜测更好——之后，模型最终能达到完美的测试精度。不过，他们并没有证明泛化的到来会推迟到 GD 的第一次迭代之后，而泛化最早可能从第二次迭代就开始。许多先前工作把 grokking 归因于优化过程中懒惰区与富裕区之间的转变，而 Boursier et al. (2025) 的近期工作则专门研究权衰减在 grokking 中的作用，指出 grokking 也可以来自从「无岭」（ridgeless）区到「有岭」（ridge）区的转变。他们考虑在光滑损失目标上用带权衰减的梯度流训练，并展示了轨迹的两阶段行为。然而，他们并没有证明无正则解不能泛化、而岭回归解能够泛化，因此其结果并不意味着可证明的 grokking。与此相关，Tikeng Notsawo et al. (2025) 研究一般的「损失加正则」框架，证明先有一个损失最小化阶段，随后是一个时间尺度与 $1/(\eta\lambda)$ 成正比的正则驱动阶段。我们的工作在岭回归这一具体设定下，给出了更尖锐的端到端保证。

Given the current theoretical gaps around grokking, our goal is to prove an end-to-end grokking guarantee demonstrating that: (i) the model overfits early, while test performance remains poor; (ii) poor generalization persists long after overfitting; and (iii) the optimization algorithm eventually converges to a model that achieves good generalization. See Figure 1 for an illustration. To this end, we investigate the phenomenon of grokking within a classical teacher-student framework. We specifically consider linear regression, a classical statistical method that has been extensively studied over the centuries (Galton, 1886; Tikhonov, 1963; Hoerl and Kennard, 1970; Uyanık and Güler, 2013). To address overfitting, fundamental regularization techniques such as Lasso and ridge regression have been developed and proven invaluable. Despite being a special case of overfitting, it is surprising that almost no previous work has focused on grokking in linear regression. While prior work has observed this phenomenon mainly in complex, non-linear architectures, our findings reveal that neither deep nor non-linear structural components are strictly necessary for grokking to occur, offering a rigorous, foundational framework for the phenomenon in purely linear settings. To our best knowledge, the only theoretical work on this topic is Levi et al. (2024), in which the authors analyze grokking in a linear regression setting by leveraging random matrix theory. Yet, despite the similar setting being studied, their analysis lacks formal guarantees. Apart from the fact that the results are non-rigorous, their work differs from ours in the following aspects. (1) They assume only Gaussian data, whereas our data distribution is more general; (2) their results consider the regimes where the feature dimension roughly equals the sample size, while we allow arbitrary over-parameterization; and (3) they do not include weight decay as we do.

面对 grokking 目前的这些理论空白，我们的目标是证明一个端到端的 grokking 保证，表明：(i) 模型很早就会过拟合，而测试表现仍然很差；(ii) 过拟合之后，很差的泛化还会持续很久；(iii) 优化算法最终会收敛到泛化良好的模型。示意图见图 1。为此，我们在经典的教师–学生框架里考察 grokking。我们具体考虑线性回归——这是被研究了几个世纪的经典统计方法（Galton, 1886; Tikhonov, 1963; Hoerl and Kennard, 1970; Uyanık and Güler, 2013）。为了应对过拟合，Lasso 和岭回归等基本正则化技术被发展出来，并被证明极为有用。尽管 grokking 可以看成过拟合的一种特殊情形，令人意外的是，几乎没有先前工作把目光放在线性回归中的 grokking 上。先前工作主要是在复杂的非线性结构里观察到这一现象，而我们的结果说明，grokking 的出现并不严格要求深度或非线性结构，从而为纯线性设定下的这一现象提供了一个严格的基础框架。据我们所知，这方面唯一的理论工作是 Levi et al. (2024)，作者借助随机矩阵理论分析线性回归中的 grokking。尽管设定相近，他们的分析缺少形式化保证。除了结果并不严格之外，他们的工作与本文还有如下不同。(1) 他们只假设高斯数据，而我们的数据分布更一般；(2) 他们的结果考虑特征维大致等于样本量的区间，而我们允许任意的过参数化；(3) 他们不像我们这样引入权衰减。

In this paper, we consider an over-parameterized linear regression problem of learning a ground-truth realizable teacher function, by training a student linear model using randomly initialized GD, on a squared loss objective, and with an $\ell_{2}$ regularization. We provide the first end-to-end provable grokking result, and unlike Xu et al. (2024), we are able to prove that poor generalization persists beyond the point where we overfit. Our results hold for a family of teacher functions: any function that is realizable by a linear model over any fixed feature map. Moreover, we also demonstrate, both theoretically and empirically, that grokking can be fully and quantitatively controlled by hyperparameter tuning. Specifically, one way to effectively manipulate the delay before generalization begins is to tune the weight decay, which aligns with the empirical findings of Lyu et al. (2024). To summarize the main novelty in this work, for the first time, (1) we show an end-to-end provable grokking result, while all prior works on grokking in neural networks or linear models did not establish such a comprehensive result; (2) we give a formal analysis of grokking in the fundamental setting of linear/ridge regression; and (3), we provide exact quantitative bounds on the grokking time that reveal how each hyperparameter affects grokking.

本文考虑一个过参数线性回归问题：学习一个可实现的真实教师函数，学生是线性模型，用随机初始化的 GD 在平方损失目标上训练，并加上 $\ell_{2}$ 正则。我们给出第一组端到端可证明的 grokking 结果，并且与 Xu et al. (2024) 不同，我们能够证明很差的泛化会持续到过拟合出现之后。我们的结果适用于一类教师函数：任何能被某个固定特征映射上的线性模型实现的函数。此外，我们还从理论上和实验上表明，grokking 可以被超参数调节完整、定量地控制。具体来说，有效操纵泛化开始之前那段延迟的一种办法，就是调节权衰减，这与 Lyu et al. (2024) 的经验发现一致。概括本文的主要新意，这是第一次：(1) 我们给出端到端可证明的 grokking 结果，而此前关于神经网络或线性模型中 grokking 的工作都没有建立这样完整的结论；(2) 我们在线性/岭回归这一基本设定下给出 grokking 的形式化分析；(3) 我们给出 grokking 时间的精确定量界，从而揭示各超参数如何影响 grokking。

Figure 1: Comparing training and test squared losses using GD with weight decay over 50 independent runs (with independent datasets and student initializations). Left: training a ridge regression model to learn the zero teacher; Right: training a two-layer ReLU neural network (training both layers) to learn the zero teacher. We scale the x-axis logarithmically, which is commonly adopted for plotting the grokking phenomenon (Power et al., 2022; Lyu et al., 2024; Xu et al., 2024; Mohamadi et al., 2024). See Section 5 for further details of the experimental setup.

图 1：用带权衰减的 GD，在 50 次独立运行（数据集和学生初始化均独立）中比较训练平方损失与测试平方损失。左：训练岭回归模型以学习零教师；右：训练两层 ReLU 神经网络（两层都训）以学习零教师。横轴按对数尺度，这是绘制 grokking 现象时的常用做法（Power et al., 2022; Lyu et al., 2024; Xu et al., 2024; Mohamadi et al., 2024）。实验设置的更多细节见第 5 节。

In a bit more detail, our main contributions are as follows:

再稍作展开，我们的主要贡献如下：

• We study an over-parameterized ridge regression problem (Section 4), and show the first end-to-end provable grokking result (Theorem 4.2) for learning any realizable teacher function using gradient descent. Specifically, we prove the following separate components: – We prove that when optimizing the $\ell_{2}$ regularized squared loss using gradient descent, the training (empirical) squared error decays at a fast convergence rate (Theorem 4.4). – We provide a lower bound for the generalization error which admits a slower convergence rate than the one for the training error, implying long-term overfitting (Theorem 4.5). – Lastly, we show that gradient descent will eventually reach a global minimum with a generalization guarantee (Theorem 4.6).

• 我们研究过参数岭回归问题（第 4 节），并给出用梯度下降学习任意可实现教师函数时，第一组端到端可证明的 grokking 结果（定理 4.2）。具体来说，我们分别证明如下几个部分：– 我们证明，用梯度下降优化带 $\ell_{2}$ 正则的平方损失时，训练（经验）平方误差以较快的速率衰减（定理 4.4）。– 我们为泛化误差给出一个下界，其收敛速率慢于训练误差，从而意味着长时间的过拟合（定理 4.5）。– 最后，我们证明梯度下降最终会到达具有泛化保证的全局最小值（定理 4.6）。

• We quantitatively state the sufficient conditions of the hyperparameters to realize grokking (Equations (3), (4) and (5)) and provide a rigorous lower bound on the grokking time in terms of model hyperparameters (Equations (6) and (7)). We then analyze how different hyperparameters affect grokking, which is in keeping with our experimental verifications.

• 我们定量给出实现 grokking 的超参数充分条件（式 (3)、(4) 和 (5)），并就 grokking 时间给出用模型超参数写出的严格下界（式 (6) 和 (7)）。随后我们分析不同超参数如何影响 grokking，这与我们的实验验证相符。

• We support our theoretical findings with empirical simulations that illustrate how grokking can be controlled in a principled manner (Section 5).

• 我们用经验模拟支持理论结果，说明 grokking 可以按有原则的方式加以控制（第 5 节）。

• Finally, we conduct experiments on non-linear neural networks, and demonstrate empirically that the dependencies of the grokking time on hyperparameters qualitatively match our provable predictions in linear regression.

• 最后，我们在非线性神经网络上做实验，经验表明 grokking 时间对超参数的依赖，在定性上与我们在线性回归中可证明的预测相符。

1.1 Additional Related Work

1.1 其他相关工作

The intriguing phenomenon of grokking was first coined by Power et al. (2022), who observed it when training on algorithmic datasets for modular addition. However, the underlying behavior—where generalization emerges long after overfitting—had already been noted in earlier work by Blanc et al. (2020), in the context of matrix sensing and simplified shallow neural networks. Following these seminal works, grokking has also been observed in other learning paradigms, which include learning sparse parities (Barak et al., 2022; Merrill et al., 2023), learning the greatest common divisor of integers (Charton, 2024), and image classification (Radhakrishnan et al., 2022).

grokking 这一引人注目的现象首先由 Power et al. (2022) 命名，他们是在模加法一类算法数据集上训练时观察到的。不过，泛化在过拟合很久之后才出现这一底层行为，早已被 Blanc et al. (2020) 在矩阵感知和简化浅层神经网络的语境中指出。在这些开创性工作之后，grokking 也在其他学习范式中被观察到，包括学习稀疏奇偶（Barak et al., 2022; Merrill et al., 2023）、学习整数的最大公约数（Charton, 2024），以及图像分类（Radhakrishnan et al., 2022）。

Several works proposed thoughtful ideas to explain grokking, although none of them provided rigorous end-to-end grokking results. Liu et al. (2022) offered an explanation via the lens of representation learning. Thilak et al. (2022) attributed grokking to the use of adaptive optimizers—an optimization anomaly named the “Slingshot Mechanism” (cyclic phase transitions). Liu et al. (2023) identified a mismatch between training and test loss landscapes, which they called the “LU mechanism”, and used it to explain certain aspects of grokking. Nanda et al. (2023) revealed that grokking can arise from the gradual amplification of structured mechanisms encoded in the weights via progress measuring. Davies et al. (2023) hypothesized that grokking and double descent can be understood using the same pattern-learning model, and provided the first demonstration of model-wise grokking. Merrill et al. (2023) studied the problem of training two-layer neural networks on a sparse parity task, and attributed grokking therein to the competition between dense and sparse subnetworks. Gromov (2023) demonstrated that fully-connected two-layer networks exhibit grokking on modular arithmetic tasks without any regularization and attributed grokking to feature learning. Miller et al. (2024) studied grokking beyond neural networks, e.g. Bayesian models. Beck et al. (2025) explored the phenomenon of grokking in a logistic regression problem. They provided evidence that grokking may occur under some asymptotic assumptions on the hyperparameters, but they did not rigorously prove grokking. Varma et al. (2023) interpreted grokking from the perspective of circuit efficiency, and discovered two related phenomena named “ungrokking” and “semi-grokking”. Jeffares et al. (2024) provided empirical insights into grokking, by investigating a telescoping model that suggests that grokking may reflect a transition to a measurably benign overfitting regime during training. Humayun et al. (2024) attributed grokking of deep neural networks to their delayed robustness. Prieto et al. (2025) demonstrated the dependence of grokking on regularization by showing that Softmax Collapse (i.e. floating point errors due to numerical instability) is responsible for the absence of grokking without regularization. We note that in addition to Prieto et al. (2025), several other papers considered grokking without explicit regularization (e.g., Xu et al. (2024); Levi et al. (2024)), but the grokking literature mostly focuses on settings with explicit regularization, as in this paper. Mallinar et al. (2025) revealed that grokking is not specific to neural networks nor to GD-based optimization methods by showing that grokking occurs when learning modular arithmetic with Recursive Feature Machines (RFM). Žunkovič and Ilievski (2024) analyzed grokking in linear classification, and showed quantitative bounds on the grokking time, but they did not provide rigorous proofs. Yunis et al. (2024) showed empirically that grokking is connected to rank minimization, namely, that the sudden drop in the test loss coincides precisely with the onset of low-rank behavior in the singular values of the weight matrices. Jeffares and van der Schaar (2025) argue that the settings in which grokking and other counterintuitive phenomena in deep learning can be shown to occur are limited in practice, as they may only appear in very specific situations.

已有若干工作提出了有启发的想法来解释 grokking，但其中没有一项给出严格的端到端 grokking 结果。Liu et al. (2022) 从表示学习的角度给出解释。Thilak et al. (2022) 把 grokking 归因于自适应优化器的使用——一种称为「弹弓机制」（Slingshot Mechanism）的优化异常（循环相变）。Liu et al. (2023) 指出训练损失景观与测试损失景观之间存在错配，并称之为「LU 机制」，用它来解释 grokking 的某些方面。Nanda et al. (2023) 揭示，grokking 可以来自权重中所编码的结构化机制经由进度度量被逐步放大。Davies et al. (2023) 提出，grokking 与双重下降可以用同一套模式学习模型来理解，并首次展示了模型维（model-wise）grokking。Merrill et al. (2023) 研究在稀疏奇偶任务上训练两层神经网络，并把其中的 grokking 归因于稠密子网络与稀疏子网络之间的竞争。Gromov (2023) 表明，全连接两层网络在模运算任务上即使不加任何正则也会出现 grokking，并将其归因于特征学习。Miller et al. (2024) 研究了超出神经网络的 grokking，例如贝叶斯模型。Beck et al. (2025) 探索了逻辑回归中的 grokking。他们给出证据，表明在超参数的某些渐近假设下可能出现 grokking，但并没有严格证明 grokking。Varma et al. (2023) 从回路效率的角度解释 grokking，并发现了两个相关现象，分别称为「ungrokking」和「semi-grokking」。Jeffares et al. (2024) 通过考察一个套叠（telescoping）模型，为 grokking 提供了经验见解，该模型表明 grokking 可能反映训练过程中向可度量的良性过拟合区间的转变。Humayun et al. (2024) 把深度神经网络的 grokking 归因于其延迟稳健性。Prieto et al. (2025) 通过表明 Softmax 崩溃（Softmax Collapse，即数值不稳定带来的浮点误差）是没有正则时观察不到 grokking 的原因，展示了 grokking 对正则化的依赖。我们注意到，除了 Prieto et al. (2025) 之外，也有若干论文考虑了没有显式正则的 grokking（例如 Xu et al. (2024); Levi et al. (2024)），但 grokking 文献大多像本文一样，关注带显式正则的设定。Mallinar et al. (2025) 表明 grokking 并不专属于神经网络，也不专属于基于 GD 的优化方法：在用递归特征机（Recursive Feature Machines，RFM）学习模运算时也会出现 grokking。Žunkovič and Ilievski (2024) 分析了线性分类中的 grokking，并给出了 grokking 时间的定量界，但他们没有提供严格证明。Yunis et al. (2024) 从经验上表明 grokking 与秩最小化有关，也就是说，测试损失的突然下降恰好与权矩阵奇异值开始呈现低秩行为同时发生。Jeffares and van der Schaar (2025) 认为，能够展示 grokking 以及深度学习中其他反直觉现象的设定在实践中相当有限，因为它们可能只在非常特定的情形下出现。

1.2 Organization

1.2 文章结构

The remaining of the paper is organized as follows. In Section 2, we formally describe the problem setup and introduce necessary notations. In Section 3, we provide an informal theorem to convey our main results in an accessible manner. In Section 4, we present our main results of end-to-end provable grokking in ridge regression. In Section 5, we support our theory with empirical verifications. In Section 6, we conclude the paper and propose some interesting future research directions. All proofs and additional experiments are deferred to the appendix.

本文其余部分安排如下。第 2 节形式化描述问题设定并引入必要记号。第 3 节给出一个非正式定理，以便用更易读的方式呈现主要结果。第 4 节给出岭回归中端到端可证明 grokking 的主要结果。第 5 节用经验验证支持理论。第 6 节总结全文并提出若干有意思的未来研究方向。全部证明和补充实验放在附录中。

2 Preliminaries

2 预备知识

Notations. We use bold-face letters to denote vectors such as $\bm{x}=(x_{1},\ldots,x_{d})\in\mathbb{R}^{d}$ and use bold-face capital letters to denote matrices such as $\bm{X}=(\bm{x}_{1},\ldots,\bm{x}_{n})^{\top}\in\mathbb{R}^{n\times d}$. We denote by $\bm{I}_{d}$ the $d$-dimensional identity matrix. For a vector $\bm{x}$, we denote by $\lVert\bm{x}\rVert_{2}$ its $\ell_{2}$ norm. For a matrix $\bm{X}$, we denote by $\lVert\bm{X}\rVert_{\mathrm{op}}$ and $\lVert\bm{X}\rVert_{F}$ its operator and Frobenius norms, respectively. For a symmetric matrix $\bm{A}$, let $\lambda_{\mathrm{min}}(\bm{A})$ and $\lambda_{\mathrm{min}}^{+}(\bm{A})$ denote its smallest eigenvalue and its smallest positive eigenvalue, respectively. We use $\mathbbm{1}\{\cdot\}$ to denote the indicator function, which equals one when its Boolean input is true and zero otherwise. We use standard asymptotic notations $O(\cdot)$ and $\Omega(\cdot)$ to hide numerical constant factors.

记号。 我们用粗体小写字母表示向量，例如 $\bm{x}=(x_{1},\ldots,x_{d})\in\mathbb{R}^{d}$，用粗体大写字母表示矩阵，例如 $\bm{X}=(\bm{x}_{1},\ldots,\bm{x}_{n})^{\top}\in\mathbb{R}^{n\times d}$。我们用 $\bm{I}_{d}$ 表示 $d$ 维单位矩阵。对向量 $\bm{x}$，用 $\lVert\bm{x}\rVert_{2}$ 表示其 $\ell_{2}$ 范数。对矩阵 $\bm{X}$，分别用 $\lVert\bm{X}\rVert_{\mathrm{op}}$ 和 $\lVert\bm{X}\rVert_{F}$ 表示其算子范数和 Frobenius 范数。对对称矩阵 $\bm{A}$，分别用 $\lambda_{\mathrm{min}}(\bm{A})$ 和 $\lambda_{\mathrm{min}}^{+}(\bm{A})$ 表示其最小特征值和最小正特征值。我们用 $\mathbbm{1}\{\cdot\}$ 表示示性函数，当其布尔输入为真时取一，否则取零。我们使用标准的渐近记号 $O(\cdot)$ 和 $\Omega(\cdot)$ 来隐去数值常数因子。

Linear regression and data generation. We consider a fundamental regression problem of learning a teacher function $N^{*}(\bm{x})$ where $\bm{x}\in\mathbb{R}^{d}$. The training data $\{(\bm{x}_{i},y_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\mathbb{R}$ are generated according to the following: the input $\bm{x}\sim\mathcal{D}_{\bm{x}}$ follows some marginal distribution $\mathcal{D}_{\bm{x}}$ (which we do not specify), and is exactly labeled by the teacher function, i.e. $y=N^{*}(\bm{x})$. We train a student linear regression model

线性回归与数据生成。 我们考虑学习教师函数 $N^{*}(\bm{x})$ 这一基本回归问题，其中 $\bm{x}\in\mathbb{R}^{d}$。训练数据 $\{(\bm{x}_{i},y_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\mathbb{R}$ 按如下方式生成：输入 $\bm{x}\sim\mathcal{D}_{\bm{x}}$ 服从某个边际分布 $\mathcal{D}_{\bm{x}}$（我们不对其作具体指定），并由教师函数精确标注，即 $y=N^{*}(\bm{x})$。我们训练一个学生线性回归模型

$$N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle,$$

where $\bm{\phi}(\bm{x}):\mathbb{R}^{d}\mapsto\mathbb{R}^{m}$ is some fixed feature map and $\bm{\theta}\in\mathbb{R}^{m}$ is the trainable parameters, to learn the teacher. We assume that the teacher is realizable, i.e., $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$ for some $\bm{\theta}^{*}\in\mathbb{R}^{m}$.

其中 $\bm{\phi}(\bm{x}):\mathbb{R}^{d}\mapsto\mathbb{R}^{m}$ 是某个固定的特征映射，$\bm{\theta}\in\mathbb{R}^{m}$ 是可训练参数，用来学习教师。我们假设教师是可实现的，即对某个 $\bm{\theta}^{*}\in\mathbb{R}^{m}$，有 $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$。

Ridge regression. In a regression problem, it is natural to train the model by minimizing the empirical mean squared loss $\frac{1}{n}\sum_{i=1}^{n}(N(\bm{x}_{i};\bm{\theta})-N^{*}(\bm{x}_{i}))^{2}$ and evaluate the learning performance on the population squared loss $\mathbb{E}_{\bm{x}}[(N(\bm{x};\bm{\theta})-N^{*}(\bm{x}))^{2}]$. We use ridge regression by appending an $\ell_{2}$-regularization term to the loss function. Let $L_{n}(\bm{\theta};\lambda)$ denote the regularized mean squared training error, let $L_{n}(\bm{\theta})$ denote the unregularized mean squared loss, and let $L_{\lambda}(\bm{\theta})$ denote the $\ell_{2}$ penalty. Our training objective is

岭回归。 在回归问题中，自然的做法是通过最小化经验均方损失 $\frac{1}{n}\sum_{i=1}^{n}(N(\bm{x}_{i};\bm{\theta})-N^{*}(\bm{x}_{i}))^{2}$ 来训练模型，并在总体平方损失 $\mathbb{E}_{\bm{x}}[(N(\bm{x};\bm{\theta})-N^{*}(\bm{x}))^{2}]$ 上评估学习表现。我们使用岭回归，也就是在损失函数上附加一项 $\ell_{2}$ 正则。设 $L_{n}(\bm{\theta};\lambda)$ 表示正则化后的均方训练误差，$L_{n}(\bm{\theta})$ 表示未正则化的均方损失，$L_{\lambda}(\bm{\theta})$ 表示 $\ell_{2}$ 惩罚项。我们的训练目标是

$$\displaystyle L_{n} \\ \displaystyle(\bm{\theta};\lambda)=L_{n}(\bm{\theta})+L_{\lambda}(\bm{\theta}) \\ \displaystyle= \\ \displaystyle\frac{1}{2n}\sum_{i=1}^{n}\left(N(\bm{x}_{i};\bm{\theta})-N^{*}(\bm{x}_{i})\right)^{2}+\frac{\lambda}{2}\lVert\bm{\theta}\rVert_{2}^{2}, \tag{1}$$

where $\lambda>0$ is the regularization parameter. We scale the loss objective by a constant factor of $1/2$ for the simplicity of gradient analysis. We optimize Equation (1) using the vanilla Gradient Descent (GD) algorithm, which updates $\bm{\theta}$ via

其中 $\lambda>0$ 是正则化参数。为了简化梯度分析，我们把损失目标乘上常数因子 $1/2$。我们用普通的梯度下降（GD）算法优化式 (1)，它通过如下方式更新 $\bm{\theta}$

$$\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\eta\nabla_{\bm{\theta}}L_{n}(\bm{\theta}^{(t)};\lambda),\;\;\forall t\in\mathbb{N},$$

where $\eta>0$ is the step-size or learning rate. While we focus on GD with a fixed step size, we note that our analysis should also hold for GD with a decaying step size and for gradient flow.

其中 $\eta>0$ 是步长或学习率。尽管我们关注固定步长的 GD，但需要指出，我们的分析也应当适用于步长递减的 GD 以及梯度流。

When $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$ and $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$, our loss objective in Equation (1) can be simplified to

当 $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$ 且 $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$ 时，式 (1) 中的损失目标可以简化为

$$L_{n}(\bm{\theta};\lambda)=\frac{1}{2n}\sum_{i=1}^{n}\left(\left\langle\bm{\theta}-\bm{\theta}^{*},\bm{\phi}(\bm{x})\right\rangle\right)^{2}+\frac{\lambda}{2}\lVert\bm{\theta}\rVert_{2}^{2},$$

and each GD iteration realizes an update in the form of

并且每一次 GD 迭代实现如下形式的更新

$$\displaystyle\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\eta\nabla_{\bm{\theta}}L_{n}(\bm{\theta}^{(t)};\lambda) \\ \displaystyle=\bm{\theta}^{(t)}-\frac{\eta}{n}\sum_{i=1}^{n}\left\langle\bm{\theta}^{(t)}-\bm{\theta}^{*},\bm{\phi}(\bm{x}_{i})\right\rangle\bm{\phi}(\bm{x}_{i})-\eta\lambda\bm{\theta}^{(t)} \tag{2}$$

To present our results, we introduce the following additional notations. Let $L(\bm{\theta})=\mathbb{E}_{\bm{x}}[(N(\bm{x};\bm{\theta})-N^{*}(\bm{x}))^{2}]$ denote the generalization (population) squared loss. We define the empirical feature map as $\bm{\Phi}=(\bm{\phi}(\bm{x}_{1}),\ldots,\bm{\phi}(\bm{x}_{n}))^{\top}\in\mathbb{R}^{n\times m}$. Let $L=\frac{1}{n}\lVert\bm{\Phi}\rVert_{F}^{2}=\frac{1}{n}\sum_{i=1}^{n}\left\lVert\bm{\phi}(\bm{x}_{i})\right\rVert_{2}^{2}$ denote the (average) norm of the feature vectors. Let $\bm{\Sigma}=\mathbb{E}_{\bm{x}}[\bm{\phi}(\bm{x})\bm{\phi}(\bm{x})^{\top}]\in\mathbb{R}^{m\times m}$ denote the feature population covariance matrix.

为了陈述结果，我们再引入如下记号。设 $L(\bm{\theta})=\mathbb{E}_{\bm{x}}[(N(\bm{x};\bm{\theta})-N^{*}(\bm{x}))^{2}]$ 表示泛化（总体）平方损失。我们把经验特征映射定义为 $\bm{\Phi}=(\bm{\phi}(\bm{x}_{1}),\ldots,\bm{\phi}(\bm{x}_{n}))^{\top}\in\mathbb{R}^{n\times m}$。设 $L=\frac{1}{n}\lVert\bm{\Phi}\rVert_{F}^{2}=\frac{1}{n}\sum_{i=1}^{n}\left\lVert\bm{\phi}(\bm{x}_{i})\right\rVert_{2}^{2}$ 表示特征向量的（平均）范数。设 $\bm{\Sigma}=\mathbb{E}_{\bm{x}}[\bm{\phi}(\bm{x})\bm{\phi}(\bm{x})^{\top}]\in\mathbb{R}^{m\times m}$ 表示特征的总体协方差矩阵。

3 Informal Results

3 非正式结果

To demonstrate a clear manifestation of grokking, our main theorem is formulated as follows: GD achieves a training error smaller than $\epsilon>0$ in an early stage of training (at step $t_{1}$), while a poor generalization error well above some constant $c>0$ persists far beyond the point of overfitting (until step $t_{2}>t_{1}$), yet eventually a generalization error smaller than $\epsilon$ is reached. In particular, this establishes grokking when $c\geq\epsilon$. Our results give a lower bound on $(t_{2}-t_{1})$ in terms of both $\epsilon$ and $c$, where this difference can be made arbitrarily large via proper hyperparameter tuning for any choice of $\epsilon$ and $c\geq\epsilon$. Specifically, let $t_{1}$ be the largest number of training steps such that the empirical squared loss is above $\epsilon$, and let $t_{2}$ be the smallest number of training steps for which the generalization loss is below $c$. Our main results are summarized by the following informal theorem, establishing end-to-end provable grokking in ridge regression.

为了把 grokking 表现得足够清楚，我们把主要定理表述如下：GD 在训练较早阶段（第 $t_{1}$ 步）就把训练误差降到小于 $\epsilon>0$，而明显高于某个常数 $c>0$ 的很差泛化误差会持续到过拟合点之后很久（直到第 $t_{2}>t_{1}$ 步），但最终泛化误差仍会降到小于 $\epsilon$。特别地，当 $c\geq\epsilon$ 时，这就确立了 grokking。我们的结果给出 $(t_{2}-t_{1})$ 关于 $\epsilon$ 和 $c$ 的下界，并且对任意选定的 $\epsilon$ 和 $c\geq\epsilon$，都可以通过适当的超参数调节把这一差值拉到任意大。具体来说，设 $t_{1}$ 是经验平方损失仍高于 $\epsilon$ 的最大训练步数，设 $t_{2}$ 是泛化损失低于 $c$ 的最小训练步数。下述非正式定理概括了我们的主要结果，确立了岭回归中端到端可证明的 grokking。

Informal Theorem(End-to-end provable grokking for ridge regression).

非正式定理（岭回归的端到端可证明 grokking）。

Consider a linear regression problem of training a student model $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$ to learn a realizable teacher $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$, where$\bm{\theta}\in\mathbb{R}^{m}$is the trainable parameter and$\bm{\theta}^{}$is unknown. Suppose that the training is done by optimizing an MSE over$n$samples using randomly initialized gradient descent at$\bm{\theta}^{(0)}$with weight decay parameter$\lambda$.*

考虑如下线性回归问题：训练学生模型 $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$，以学习可实现教师 $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$，其中$\bm{\theta}\in\mathbb{R}^{m}$是可训练参数，$\bm{\theta}^{}$未知。假设训练是在$n$个样本上优化 MSE，使用在$\bm{\theta}^{(0)}$处随机初始化、权衰减参数为$\lambda$ 的梯度下降。*

Under the initialization $\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$ for some $\nu^{2}>0$, and some mild distributional assumptions on the feature map $\bm{\phi}(\bm{x})$, we have with high probability that for a sufficiently large sample size $n$, and a sufficiently large feature space dimensionality $m$, $(t_{2}-t_{1})$ can be lower bounded by an arbitrarily large value by choosing a sufficiently small weight decay $\lambda$.

在初始化 $\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$（对某个 $\nu^{2}>0$）以及特征映射 $\bm{\phi}(\bm{x})$ 满足若干温和分布假设的条件下，我们以高概率得到：当样本量 $n$ 足够大、特征空间维数 $m$ 足够大时，只要把权衰减 $\lambda$ 取得足够小，$(t_{2}-t_{1})$ 就可以被下界成任意大。

4 Grokking in Ridge Regression

4 岭回归中的 Grokking

We now formalize our end-to-end provable grokking results for ridge regression problems. For any realizable teacher function, we provide a quantitative lower bound for the grokking time $(t_{2}-t_{1})$ in terms of training hyperparameters, as well as bounds on the hyperparameters that suffice to realize grokking. Based on these bounds, we analyze how the different hyperparameters affect grokking, which is corroborated in our experiments in Section 5.

我们现在把岭回归问题中端到端可证明的 grokking 结果形式化。对任意可实现的教师函数，我们用训练超参数给出 grokking 时间 $(t_{2}-t_{1})$ 的定量下界，以及足以实现 grokking 的超参数界。基于这些界，我们分析不同超参数如何影响 grokking，这与第 5 节的实验相一致。

4.1 Warmup: Grokking with the Zero Teacher

4.1 热身：零教师情形下的 Grokking

We start by considering a simple scenario where the target teacher is the zero function. When $N^{*}(\bm{x})=0$, we have

我们先考虑目标教师为零函数的简单情形。当 $N^{*}(\bm{x})=0$ 时，我们有

$$\displaystyle L_{n}(\bm{\theta};\lambda)=\frac{1}{2n}\sum_{i=1}^{n}(N(\bm{x}_{i};\bm{\theta}))^{2}+\frac{\lambda}{2}\lVert\bm{\theta}\rVert_{2}^{2}; \\ \displaystyle\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\frac{\eta}{n}\sum_{i=1}^{n}\langle\bm{\theta}^{(t)},\bm{\phi}(\bm{x}_{i})\rangle\bm{\phi}(\bm{x}_{i})-\eta\lambda\bm{\theta}^{(t)}.$$

The following theorem describes an end-to-end provable grokking result for learning the zero teacher in a ridge regression setting using gradient descent. To develop some initial intuition on what causes grokking, we state the aforementioned three stages of grokking in terms of the convergence rates of the training/generalization squared losses, in contrast to the format of our Theorem 4.2 which appears later.

下述定理给出在岭回归设定下用梯度下降学习零教师时，端到端可证明的 grokking 结果。为了对 grokking 的成因建立初步直觉，我们用训练/泛化平方损失的收敛速率来陈述前述三个 grokking 阶段，这与稍后出现的定理 4.2 的表述形式不同。

Theorem 4.1(End-to-end provable grokking for the zero teacher).

定理 4.1（零教师情形的端到端可证明 grokking）。

Suppose that the teacher function is $N^{}(\bm{x})=0$. Consider training a student$N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$to learn the teacher via optimizing the ridge regression objective in Equation (1) using randomly initialized GD with$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$. Assume that$\eta\leq 2/(L+2\lambda)$. Then for any$t\in\mathbb{N}$, we have*

设教师函数为 $N^{}(\bm{x})=0$。考虑训练学生$N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$以学习该教师，做法是用随机初始化的 GD 优化式 (1) 中的岭回归目标，其中$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$。假设$\eta\leq 2/(L+2\lambda)$。则对任意$t\in\mathbb{N}$，我们有*

• $\mathrm{(i)}$ $L_{n}(\bm{\theta}^{(t)})\leq\frac{L}{2}\cdot(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})-\eta\lambda)^{2t}\cdot\lVert\bm{\theta}^{(0)}\rVert_{2}^{2};$

• $\mathrm{(ii)}$ With probability at least$1-2e^{-(m-n)/32}$, $L(\bm{\theta}^{(t)})\geq\lambda_{\mathrm{min}}(\bm{\Sigma})\cdot(1-\eta\lambda)^{2t}\cdot\frac{(m-n)\nu^{2}}{2};$

• $\mathrm{(ii)}$ 以至少 $1-2e^{-(m-n)/32}$ 的概率，$L(\bm{\theta}^{(t)})\geq\lambda_{\mathrm{min}}(\bm{\Sigma})\cdot(1-\eta\lambda)^{2t}\cdot\frac{(m-n)\nu^{2}}{2};$

• $\mathrm{(iii)}$ $\lVert\bm{\theta}^{(t)}\rVert_{2}^{2}\leq(1-\eta\lambda)^{2t}\cdot\lVert\bm{\theta}^{(0)}\rVert_{2}^{2}.$

Studying the zero teacher as a warmup yields the following advantages. First, choosing the zero teacher enables us to build a much simpler generalization guarantee. Clearly, item (iii) above implies that GD eventually converges to the solution $\bm{\theta}^{*}=\bm{0}$. Since $N^{*}(\bm{x})=0$, this yields generalization. Additionally, choosing the zero teacher allows us to derive matching lower (ii) and upper (iii) bounds for the generalization loss in terms of the convergence rate.

把零教师作为热身来研究，有如下好处。首先，选择零教师使我们能够建立简单得多的泛化保证。显然，上面的第 (iii) 项意味着 GD 最终会收敛到解 $\bm{\theta}^{*}=\bm{0}$。由于 $N^{*}(\bm{x})=0$，这就给出了泛化。此外，选择零教师还使我们能够就泛化损失给出在收敛速率上相匹配的下界 (ii) 与上界 (iii)。

More importantly, Theorem 4.1 reveals that the training and test losses can decrease at different rates, and the potentially considerable discrepancy between these rates causes grokking. Specifically, given the empirical feature map $\bm{\Phi}^{\top}\bm{\Phi}$, a sufficiently small regularization $\lambda$ can only have a negligible effect on the convergence rate of the training MSE. This is because the convergence rate of $(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})-\eta\lambda)^{2t}$ will be dominated by the term $(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi}))^{2t}$ when $\lambda$ is sufficiently small. However, decreasing $\lambda$ can make the generalization convergence rate $(1-\eta\lambda)^{2t}$ arbitrarily slow. In conclusion, Theorem 4.1 establishes a provable instance of grokking whenever $m\gg n$, $\lambda\ll\frac{1}{n}\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ and $\lambda\rightarrow 0$.

更重要的是，定理 4.1 揭示出训练损失和测试损失可以按不同速率下降，而这两种速率之间可能相当可观的差距，正是 grokking 的成因。具体来说，给定经验特征矩阵 $\bm{\Phi}^{\top}\bm{\Phi}$，足够小的正则 $\lambda$ 对训练 MSE 的收敛速率几乎没有影响。这是因为当 $\lambda$ 足够小时，$(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})-\eta\lambda)^{2t}$ 的收敛速率将被 $(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi}))^{2t}$ 主导。但是，减小 $\lambda$ 却可以使泛化收敛速率 $(1-\eta\lambda)^{2t}$ 变得任意慢。总之，只要 $m\gg n$、$\lambda\ll\frac{1}{n}\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ 且 $\lambda\rightarrow 0$，定理 4.1 就给出了一个可证明的 grokking 实例。

4.2 Grokking with Realizable Teachers

4.2 可实现教师情形下的 Grokking

For the most general setting studied in this paper, we consider the realizable case where the ground-truth teacher is $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$ for some unknown $\bm{\theta}^{*}\in\mathbb{R}^{m}$, and the student is $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$ where $\bm{\theta}\in\mathbb{R}^{m}$ is the vector of trainable parameters. We show that for any teacher, there is a realization of training hyperparameters that results in an arbitrarily long grokking time. Our main result is the following.

就本文研究的最一般设定而言，我们考虑可实现情形：真实教师为 $N^{*}(\bm{x})=\langle\bm{\theta}^{*},\bm{\phi}(\bm{x})\rangle$，其中 $\bm{\theta}^{*}\in\mathbb{R}^{m}$ 未知；学生为 $N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$，其中 $\bm{\theta}\in\mathbb{R}^{m}$ 是可训练参数向量。我们证明，对任意教师，都存在一组训练超参数的实现，使得 grokking 时间可以任意长。我们的主要结果如下。

Theorem 4.2(End-to-end provable grokking for realizable ridge regression).

定理 4.2（可实现岭回归的端到端可证明 grokking）。

Assume a realizable teacher function $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$for some$\bm{\theta}^{}\in\mathbb{R}^{m}$. Assume the boundedness of the feature map, i.e.$\lVert\bm{\phi}(\bm{x})\rVert_{2}\leq b$for any$\bm{x}$and some$b>0$. Consider training a student$N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$to learn the teacher via optimizing the ridge regression objective in Equation (1) using randomly initialized GD with$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$. For all$\epsilon>0$, let$c\geq\epsilon$ be arbitrary. We define*

假设存在可实现的教师函数 $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$，其中$\bm{\theta}^{}\in\mathbb{R}^{m}$。假设特征映射有界，即对任意$\bm{x}$以及某个$b>0$，有$\lVert\bm{\phi}(\bm{x})\rVert_{2}\leq b$。考虑训练学生$N(\bm{x};\bm{\theta})=\langle\bm{\theta},\bm{\phi}(\bm{x})\rangle$以学习该教师，做法是用随机初始化的 GD 优化式 (1) 中的岭回归目标，其中$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$。对所有$\epsilon>0$，设$c\geq\epsilon$ 为任意常数。我们定义*

$$t_{1}:=t_{1}(\epsilon):=\max\left(\left\{t\in\mathbb{N}:L_{n}(\bm{\theta}^{(t)})\geq\epsilon\right\}\right)$$

and

以及

$$t_{2}:=t_{2}(c):=\min\left(\left\{t\in\mathbb{N}:L(\bm{\theta}^{(t)})\leq c\right\}\right).$$

For any $\delta\in(0,1)$, assume a sufficiently large sample size:

对任意 $\delta\in(0,1)$，假设样本量足够大：

$$n=\Omega\left(\frac{b^{4}\lVert\bm{\theta}^{*}\rVert_{2}^{4}}{\epsilon^{2}}\log\left(\frac{1}{\delta}\right)\right), \tag{3}$$

a sufficiently large dimensionality of the feature map:

特征映射的维数足够大：

$$m=n+\Omega\left(\max\left\{\log\left(\frac{1}{\delta}\right),\frac{\lVert\bm{\theta}^{*}\rVert_{2}^{2}}{\nu^{2}},\frac{c}{\lambda_{\mathrm{min}}(\bm{\Sigma})\nu^{2}}\right\}\right), \tag{4}$$

and finally, a sufficiently small weight decay:

最后，权衰减足够小：

$$\lambda=O\left(\min\left\{\frac{\epsilon}{\lVert\bm{\theta}^{*}\rVert_{2}^{2}},\frac{b\sqrt{\epsilon}}{\lVert\bm{\theta}^{*}\rVert_{2}}\right\}\right). \tag{5}$$

Then, if $\eta<1/(\lambda+b^{2})$, we have with probability at least $1-\delta$ that

那么，若 $\eta<1/(\lambda+b^{2})$，则以至少 $1-\delta$ 的概率，我们有

$$t_{1}\leq\frac{n\ln\left(\frac{6b^{2}\lVert\bm{\theta}^{(0)}\rVert_{2}^{2}}{\epsilon}\right)}{2\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})} \tag{6}$$

and

以及

$$t_{2}\geq\frac{\ln\left(\frac{(m-n)\nu^{2}}{2}\left(\sqrt{\frac{c}{\lambda_{\mathrm{min}}(\bm{\Sigma})}}+\lVert\bm{\theta}^{*}\rVert_{2}\right)^{-2}\right)}{4\eta\lambda}. \tag{7}$$

Moreover, $L(\bm{\theta}^{(t)})\leq\epsilon$ for all sufficiently large $t$.

此外，对所有足够大的 $t$，有 $L(\bm{\theta}^{(t)})\leq\epsilon$。

The main intuition behind this discrepancy between $t_{1}$ and $t_{2}$ can be explained as follows: When the student regression model is sufficiently over-parameterized ($m\gg n$), we are solving a high-dimensional regression problem. Therefore, when the weight decay $\lambda$ is small, the GD optimization process only effectively updates the projection of the weight vector $\bm{\theta}_{\parallel}$ onto the data-spanning subspace to fit the training data, while having negligible effect on the component $\bm{\theta}_{\perp}$ in the complementary subspace, which remains close to its initialization. This occurs since the information captured by the data only contributes to the convergence of $\bm{\theta}_{\parallel}$ at a rate of $(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})-\eta\lambda)^{t}$. In contrast, $\bm{\theta}_{\perp}$ converges at a negligible rate of $(1-\eta\lambda)^{t}$ as a consequence of long-term weight decay. Such an imbalanced trajectory prevents timely generalization and leads to harmful overfitting. As a remedy, weight decay plays a crucial role in decreasing the complexity of the model class, which eventually guarantees uniform convergence and thus good generalization.

关于 $t_{1}$ 与 $t_{2}$ 之间这一差距的主要直觉，可以解释如下：当学生回归模型充分过参数化（$m\gg n$）时，我们是在解一个高维回归问题。因此，当权衰减 $\lambda$ 较小时，GD 优化过程只会有效更新权向量在数据张成子空间上的投影 $\bm{\theta}_{\parallel}$，以拟合训练数据，而对互补子空间中的分量 $\bm{\theta}_{\perp}$ 几乎没有影响，后者仍接近其初始化。之所以如此，是因为数据所捕获的信息只以 $(1-\frac{1}{n}\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})-\eta\lambda)^{t}$ 的速率贡献于 $\bm{\theta}_{\parallel}$ 的收敛。与此相对，$\bm{\theta}_{\perp}$ 以可忽略的速率 $(1-\eta\lambda)^{t}$ 收敛，这是长期权衰减的结果。这样不平衡的轨迹妨碍了及时泛化，并导致有害的过拟合。作为补救，权衰减在降低模型类复杂度方面起着关键作用，最终保证一致收敛，从而保证良好的泛化。

Next, we disentangle the influence of individual hyperparameters on the grokking time according to our theory, which aligns closely with our experimental results in Section 5. Notably, these functional dependencies on the hyperparameters are not only qualitatively corroborated by our experiments, but they also quantitatively predict the expected scaling rates (see Figure 2).

接下来，我们按理论把各个超参数对 grokking 时间的影响拆开，这与第 5 节的实验结果密切相符。值得注意的是，这些关于超参数的函数依赖不仅在定性上得到实验支持，而且在定量上也能预测期望的缩放速率（见图 2）。

Evidently, Equations (6) and (7) together provide strong control of the grokking time $(t_{2}-t_{1})$, which reveals the following effects of hyperparameter tuning on grokking:

显然，式 (6) 和式 (7) 合在一起对 grokking 时间 $(t_{2}-t_{1})$ 给出了很强的控制，并从中可以看出超参数调节对 grokking 的如下影响：

• Weight decay $\lambda$: For small $\lambda$, decreasing $\lambda$ increases $t_{2}$ with $t_{2}\propto 1/\lambda$ when other hyperparameters are held fixed, and $\lambda$ has no effect on our bound for $t_{1}$. Therefore, $(t_{2}-t_{1})\rightarrow\infty$ as $\lambda\rightarrow 0$. We remark that for large $\lambda$ (which is not our focus in this work), our upper bound in Equation (6) is not tight in general, as a tight bound should incorporate the term $\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})+\eta\lambda$ in the denominator.

• 权衰减 $\lambda$：当 $\lambda$ 较小时，在其他超参数固定的前提下，减小 $\lambda$ 会使 $t_{2}$ 增大，并且 $t_{2}\propto 1/\lambda$，而 $\lambda$ 对我们给出的 $t_{1}$ 的界没有影响。因此，当 $\lambda\rightarrow 0$ 时，$(t_{2}-t_{1})\rightarrow\infty$。需要说明的是，对较大的 $\lambda$（这不是本文的重点），式 (6) 中的上界一般并不紧，因为紧的界应当在分母中纳入 $\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})+\eta\lambda$ 这一项。

• Sample size $n$ and feature dimensionality $m$: It is difficult to deduce clean dependencies that capture the behavior of $t_{1}$ as a function of $n$ and $m$, since there is no known quantitative bound revealing how $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ depends on $m$ and $n$, even when making strong assumptions on the distribution of the features (see Remark 4.3 for a discussion). On a more positive note, we show empirically that the behavior of $t_{1}$ as a function of $n$ and $m$ strongly accords with our theoretical bound in Equation (6). As for $t_{2}$, since we make no assumption on the feature map $\bm{\phi}(\bm{x})$, the quantity $\lambda_{\mathrm{min}}(\bm{\Sigma})$ is not necessarily controlled. Indeed, more explicit dependencies of $t_{2}$ on $n$ and $m$ can be derived by assuming distributional assumptions. For instance, if $\bm{\phi}(\bm{x}_{1}),\ldots,\bm{\phi}(\bm{x}_{n})$ are i.i.d. spherical Gaussian vectors. We present such clean bounds in Equation (8) and show that it strongly matches our experimental results.

• 样本量 $n$ 与特征维数 $m$：要推出能够刻画 $t_{1}$ 作为 $n$ 和 $m$ 的函数行为的干净依赖关系并不容易，因为即便对特征分布作很强的假设，目前也还没有定量界揭示 $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ 如何依赖于 $m$ 和 $n$（讨论见备注 4.3）。较为正面的是，我们从经验上表明，$t_{1}$ 作为 $n$ 和 $m$ 的函数，其行为与式 (6) 中的理论界高度一致。至于 $t_{2}$，由于我们对特征映射 $\bm{\phi}(\bm{x})$ 不作假设，量 $\lambda_{\mathrm{min}}(\bm{\Sigma})$ 不一定能被控制。确实，若再加入分布假设，就可以推出 $t_{2}$ 对 $n$ 和 $m$ 更显式的依赖。例如，设 $\bm{\phi}(\bm{x}_{1}),\ldots,\bm{\phi}(\bm{x}_{n})$ 是独立同分布的球面高斯向量。我们在式 (8) 中给出这样干净的界，并表明它与实验结果高度吻合。

• Initialization scale $\nu^{2}$: Increasing $\nu^{2}$ increases $t_{1}$ and $t_{2}$ simultaneously with $t_{1},t_{2}\propto\ln(\nu^{2})$. More precisely, when $\lambda$ is sufficiently small, we have $t_{2}>t_{1}$, and moreover $t_{1}\leq c_{1}\ln(\nu^{2})+a_{1}$ and $t_{2}\geq c_{2}\ln(\nu^{2})+a_{2}$ for some $c_{2}>c_{1}>0$ and some $a_{1},a_{2}$. In such a scenario, increasing $\nu^{2}$ increases $(t_{2}-t_{1})$ and amplifies grokking.

• 初始化尺度 $\nu^{2}$：增大 $\nu^{2}$ 会同时增大 $t_{1}$ 和 $t_{2}$，并且 $t_{1},t_{2}\propto\ln(\nu^{2})$。更精确地说，当 $\lambda$ 足够小时，有 $t_{2}>t_{1}$，并且存在某些 $c_{2}>c_{1}>0$ 以及某些 $a_{1},a_{2}$，使得 $t_{1}\leq c_{1}\ln(\nu^{2})+a_{1}$ 且 $t_{2}\geq c_{2}\ln(\nu^{2})+a_{2}$。在这种情形下，增大 $\nu^{2}$ 会增大 $(t_{2}-t_{1})$，从而放大 grokking。

Remark 4.3(on $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ and its dependence on $n$ and $m$).

备注 4.3（关于 $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ 及其对 $n$ 和 $m$ 的依赖）。

$\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$is an empirical quantity which depends on the sampled features. We note that even for simple distributions such as $\bm{\phi}(\bm{x})\sim\mathcal{N}(\bm{0},\frac{1}{m}\bm{I}_{m})$, the behavior of $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ as a function of $n$ and $m$ is not yet fully understood. While there is a known asymptotic result given by the celebrated Marchenko–Pastur law (Marčenko and Pastur, 1967), showing that when $m,n\rightarrow\infty$ and $m/n\rightarrow\gamma$ for some $\gamma\in(0,\infty)$, $\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})\rightarrow\gamma^{-1}(1-\sqrt{\gamma})^{2}$, to the best of our knowledge, there is no quantitative bound revealing how this depends on $m$ and $n$.

$\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ 是依赖于所采样特征的经验量。我们注意到，即便对 $\bm{\phi}(\bm{x})\sim\mathcal{N}(\bm{0},\frac{1}{m}\bm{I}_{m})$ 这样简单的分布，$\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})$ 作为 $n$ 和 $m$ 的函数，其行为也尚未被完全理解。虽然著名的 Marchenko–Pastur 定律（Marčenko and Pastur, 1967）给出了已知的渐近结果：当 $m,n\rightarrow\infty$ 且 $m/n\rightarrow\gamma$（对某个 $\gamma\in(0,\infty)$）时，$\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})\rightarrow\gamma^{-1}(1-\sqrt{\gamma})^{2}$，但据我们所知，还没有定量界揭示它如何依赖于 $m$ 和 $n$。

Next, we state a separate theorem for each of the three stages of our grokking result, as it is helpful for understanding each hyperparameter effect separately.

接下来，我们为 grokking 结果的三个阶段分别陈述一条定理，这有助于分别理解各个超参数的影响。

Theorem 4.4(Training loss convergence).

定理 4.4（训练损失的收敛）。

Suppose that $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$for some$\bm{\theta}^{}\in\mathbb{R}^{m}$, and that$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$. Assume that$\eta<1/(\lambda+\lambda_{\mathrm{max}}(\bm{\Phi}^{\top}\bm{\Phi})/n)$. For any$\epsilon>0$and$\delta\in(0,1)$, suppose that*

设 $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$对某个$\bm{\theta}^{}\in\mathbb{R}^{m}$成立，且$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$。假设$\eta<1/(\lambda+\lambda_{\mathrm{max}}(\bm{\Phi}^{\top}\bm{\Phi})/n)$。对任意$\epsilon>0$和$\delta\in(0,1)$，设*

$$m=\Omega\left(\max\left\{\log\left(\frac{1}{\delta}\right),\frac{\lVert\bm{\theta}^{*}\rVert_{2}^{2}}{\nu^{2}}\right\}\right)$$

and

以及

$$\lambda=O\left(\min\left\{\frac{\epsilon}{\lVert\bm{\theta}^{*}\rVert_{2}^{2}},\frac{\sqrt{L\epsilon}}{\lVert\bm{\theta}^{*}\rVert_{2}}\right\}\right).$$

Then, with probability at least $1-\delta$, we have $L_{n}(\bm{\theta}^{(t)})\leq\epsilon$ for all

那么，以至少 $1-\delta$ 的概率，对所有满足下式的 $t$，有 $L_{n}(\bm{\theta}^{(t)})\leq\epsilon$

$$t\geq\frac{1}{2}\log_{1-\frac{\eta\lambda_{\mathrm{min}}^{+}(\bm{\Phi}^{\top}\bm{\Phi})}{n}-\eta\lambda}\left(\frac{\epsilon}{6L\lVert\bm{\theta}^{(0)}\rVert_{2}^{2}}\right).$$

The above assumption that $m=\Omega(\lVert\bm{\theta}^{*}\rVert_{2}^{2}/\nu^{2})$ implies that the student is initialized to have a large norm with respect to the teacher’s norm. Since we aim to upper bound $L_{n}(\bm{\theta}^{(t)})$ and not $L_{n}(\bm{\theta}^{(t)};\lambda)$, the assumed upper bound on $\lambda$ is for technical purposes, and it does not limit our results since small values of $\lambda$ facilitate grokking.

上面关于 $m=\Omega(\lVert\bm{\theta}^{*}\rVert_{2}^{2}/\nu^{2})$ 的假设意味着，学生被初始化成相对于教师范数而言具有较大范数。由于我们要上界的是 $L_{n}(\bm{\theta}^{(t)})$ 而不是 $L_{n}(\bm{\theta}^{(t)};\lambda)$，对 $\lambda$ 所假设的上界是出于技术原因；它并不限制我们的结果，因为较小的 $\lambda$ 反而有助于 grokking。

Theorem 4.5(Poor generalization when overfitting).

定理 4.5（过拟合时的较差泛化）。

Suppose that $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$for some$\bm{\theta}^{}\in\mathbb{R}^{m}$, that$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$, and further assume that$\eta<1/\lambda$. For any$c>0$and$\delta\in(0,1)$, suppose that*

设 $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$对某个$\bm{\theta}^{}\in\mathbb{R}^{m}$成立，且$\bm{\theta}^{(0)}\sim\mathcal{N}(\bm{0},\nu^{2}\bm{I}_{m})$，并进一步假设$\eta<1/\lambda$。对任意$c>0$和$\delta\in(0,1)$，设*

$$m=n+\Omega\left(\max\left\{\log\left(\frac{1}{\delta}\right),\frac{\lVert\bm{\theta}^{*}\rVert_{2}^{2}}{\nu^{2}},\frac{c}{\lambda_{\mathrm{min}}(\bm{\Sigma})\nu^{2}}\right\}\right).$$

Then, with probability at least $1-\delta$, we have $L(\bm{\theta}^{(t)})\geq c$ for all

那么，以至少 $1-\delta$ 的概率，对所有满足下式的 $t$，有 $L(\bm{\theta}^{(t)})\geq c$

$$t\leq\log_{1-\eta\lambda}\left(\sqrt{\frac{2}{(m-n)\nu^{2}}}\left(\sqrt{\frac{c}{\lambda_{\mathrm{min}}(\bm{\Sigma})}}+\lVert\bm{\theta}^{*}\rVert_{2}\right)\right).$$

The lower bound on $m$ here is stronger than the one in Theorem 4.4, which ensures that we are in a high-dimensional regression situation. In our last theorem, we prove that good generalization eventually occurs.

这里对 $m$ 的下界比定理 4.4 中的更强，这保证我们处于高维回归情形。在最后一条定理中，我们证明良好的泛化最终会发生。

Theorem 4.6(Generalization).

定理 4.6（泛化）。

Suppose that $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$for some$\bm{\theta}^{}\in\mathbb{R}^{m}$. Assume that$\lVert\bm{\phi}(\bm{x})\rVert_{2}\leq b$for any$\bm{x}$for some$b>0$. Let$\bm{\theta}^{}{\lambda}=\argmin{\bm{\theta}}\left{L_{n}(\bm{\theta};\lambda)\right}$. For any$\epsilon>0$and$\delta\in(0,1)$, suppose that

设 $N^{}(\bm{x})=\langle\bm{\theta}^{},\bm{\phi}(\bm{x})\rangle$对某个$\bm{\theta}^{}\in\mathbb{R}^{m}$成立。假设对任意$\bm{x}$以及某个$b>0$，有$\lVert\bm{\phi}(\bm{x})\rVert_{2}\leq b$。设$\bm{\theta}^{}{\lambda}=\argmin{\bm{\theta}}\left{L_{n}(\bm{\theta};\lambda)\right}$。对任意$\epsilon>0$和$\delta\in(0,1)$，设

$$n=\Omega\left(\frac{b^{4}\lVert\bm{\theta}^{*}\rVert_{2}^{4}}{\epsilon^{2}}\log\left(\frac{1}{\delta}\right)\right).$$

Then, if $\eta\leq 1/(\lambda+b^{2})$, GD converges to $\bm{\theta}^{}{\lambda}$satisfying$L(\bm{\theta}^{*}{\lambda})\leq 2L_{n}(\bm{\theta}^{}_{\lambda})+\epsilon$, with probability at least$1-\delta$.

那么，若 $\eta\leq 1/(\lambda+b^{2})$，则 GD 收敛到满足 $L(\bm{\theta}^{}{\lambda})\leq 2L{n}(\bm{\theta}^{}_{\lambda})+\epsilon$的$\bm{\theta}^{}_{\lambda}$，并且这一事件的概率至少为$1-\delta$。*

The lower bound on the sample size $n=\Omega(b^{4}\lVert\bm{\theta}^{*}\rVert_{2}^{4}\epsilon^{-2})$ stems from a standard uniform convergence argument based on Rademacher complexity.

对样本量的下界 $n=\Omega(b^{4}\lVert\bm{\theta}^{*}\rVert_{2}^{4}\epsilon^{-2})$ 来自基于 Rademacher 复杂度的标准一致收敛论证。

5 Experiments

5 实验

In this section, we systematically verify our theoretical understanding of grokking via experiments. Throughout, we use $n,\eta$ and $\lambda$ to denote the training sample size, GD step size and weight decay parameter, respectively. In all of our experiments, we set the error thresholds $c=\epsilon=0.01$.

本节我们通过实验系统验证对 grokking 的理论理解。全文用 $n$、$\eta$ 和 $\lambda$ 分别表示训练样本量、GD 步长和权衰减参数。在所有实验中，我们把误差阈值都设成 $c=\epsilon=0.01$。