Grokking as the transition from lazy to rich training dynamics
Authors: Tanishq Kumar
Blake Bordelon
Samuel J. Gershman
Cengiz Pehlevan
Venue: ICLR 2024 poster (2024)
Keywords: Grokking
Feature Learning
Neural Tangent Kernel
Generalization
Training Dynamics
Primary area: optimization
TL;DR: A theory of grokking based on delayed feature learning.

Abstract: 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.

Published as a conference paper at ICLR 2024
GROKKING AS THE TRANSITION FROM LAZY TO RICH TRAINING DYNAMICS
Tanishq Kumar Blake Bordelon Samuel J. Gershman" Harvard University
Cengiz Pehlevan"
ABSTRACT
We propose that the grokking phenomenon, where the train loss of a neural net- work 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 il- lustrate this mechanism, we study the simple setting of vanilla gradient descent on polynomial regression problem with two layer neural network which exhibits grokking without regularization in a way that cannot be explained by existing the- ories. We identify sufficient statistics for the test loss of such network, and track- ing 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 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 beaninng-wiih can be controlled precisely by parameters that scale the network oud the alignment of the initial features with the target function y(x). We argue this delayed generalization arises when (I) the top eigenvectors the initial neural tangent kernel and the task labels y(z) are misaligned, but (2) the dataset size is large enough that it is possible for the network 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) rich training (feature learning) can control grokking in more general settings, like on MNIST, one-layer Transformers, and student-teacher networks.
INTRODUCTION
The goal of machine learning system is learn models that generalize beyond their training set. Typically, a practitioner hopes that a model's performance on its training set will be an indicator of its generalization capabilities on unseen data. However, this may not happen in practice. Grokking, discovered by PPower et al.] (2022). is phenomenon where the train loss of a network falls initially with no corresponding decrease in test loss, then the network generalizes later during training. There has been much recent work studying this phenomenon not only because it challenges the widely- used approaches of early stopping or train loss termination criteria, but also because some Nanda et al. (2023) see grokking an example of sudden, emergent behavior in neural network, and principle aim of the field ML interpretability is to understand model internals well enough that they do not exhibit such unexpected and unpredictable changes in behavior.
For these reasons, the phenomenology of grokking has been under close empirical study since its discovery (Nand et al. 2023 Davies et al. 2023 Thilak et al. 2022 Varma et al. 2023 Liu al. 2022a by Gromov 2023). ). A limited number of theories have been proposed to explain the causes of grokking, most notably those that attribute grokking to weight decay and weight norm decrease (Liu ct al. 2022aa) Varma et al. 2023). However, we give counterexamples to such theories, implying need explain grokking way that is consistent with empirical observations in current and past work.
In this paper, we propose the idea that grokking is caused by delayed feature learning during training. Our theoretical framework relies on the idea of lazy training where network can fit its training
"Equal Senior Authors, Correspondence to wwchlecan www.seaasssaa haw.aardddd....

---

Published as a conference paper at ICLR 2024
S(X, k0p)
Maayy
Perarater
W0g
W
Seature
lerwet Lee Lay semort Lore
Fadh/ barning
Grilenz
Spochs
(a)
(b)
(c)
Figure 1: (a) Claimed parameter dynamics during grokking in parameter space of RR3 for illustrative purposes. S(X, wo) is an affine subspace of parameter space reachable by models linearized around wo- (b) Grokking on polynomial regression task introduced in Section[5 with an MLP, vanilla GD, and zero regularization. Green and blue loss curves in (b) correspond to sketched green and blue parameter dynamics in (a). Horizontal line (black) in (b) is the mean- mean-squared error of best kernel regression estimate with the NTK at initialization.
(c) The fact that parameter weight norm increases (dashed orange) cannot be explained by any existing theories of grokking. Features get aligned (grey) in a way we make precise in Section 5
data without adapting its internal representations (Chizat et al. 2019). this regime, the network behaves linearized model with features evaluated at initialization (Jacot et al. 2018 Chizat et al., 2019 Lee et al. 2019 We show that grokking can arise when a network begins fitting and " "memorizing" its training set (albeit imperfectly) in the linearized regime, then later adapts its features to the data, leading to improved generalization at late time.
We illustrate this idea schematically in Figure [a). Suppose randomly initialize network with the parameters wo. When linearized around Wo, under gradient descent, a neural network's 's weights restricted an affine subspace spanned by initial gradients wo span V se. (w, 2r (,,),,,)) } -==1 and gradient descent finds a training minimizer on this subspace, Wuzzy For MSE loss this minimizer the kernel regression solution with the network's 's initial Neural Tangent Kernel (NTK) (Jacot al.) 2018 [Lee et al. 2019). Alternatively, if the network leaves the linearized regime, the weights can leave the subspace of the linear model and converge to another solution Wteetuue We hypothesize that grokking can occur when network transitions from initially tracking the linearized network dynamics in the subspace to more complicated dynamics that approach Wieatue off of the subspace which achieves better generalization because it learns features.
Given a fixed dataset where grokking is possible, our proposal naturally suggests two key factors that control grokking: (1) how much feature learning is necessary for good performance, and (2) how fast these features are learned. We introduce two parameters corresponding to these notions and demonstrate they can control grokking.
The key contributions in this paper are the following:
We give simple examples of a two-layer MLP without any weight decay that exhibits grokking in ways inconsistent with existing theories of grokking.
We demonstrate with both theory and experiments that grokking in this setting is associated with the transition from lazy training to rich, feature learning regime, and can be tuned by parameters that control this transition. These parameters are:
A scale parameter, a, that controls the magnitude of the output of the network, and thus the rate of feature learning throughout training (ie. laziness).
The task-model kernel alignment between the neural tangent kernel (NTK) of the network at initialization and the test labels, a measure of how much feature learning is necessary for good performance the task (Cortes al. 2012).
We use the terms "lazy." "kernel regime" and "linearized model" interchangeably to refer to the same phenomenon where the network learns without changing its internal features. Concretely, this is regime where the network output
large amount while w.. stays approximately constant.
2

---

Published as a conference paper at ICLR 2024
A "goldilocks" train set size that generalization is possible but not immediate (for instance, grokking is impossible in the infinite data limit where train loss perfectly tracks test loss).
We provide multiple lines of evidence that the feature learning dynamics that give grokking in this simple setup are also those at play during grokking in more complicated settings studied in past work: on MNIST, with one-layer transformers, and with student-teacher networks. Given a set of training inputs, we give a method to construct labels that can generate grokking.
2 RELATED WORK
Grokking. Nanda et al. (2023) mechanistically interpret the algorithm learned during grokking for modular addition tasks. We show in Section that merely tuning the network output scaling/laziness parameter a is alone sufficient to make grokking continuously vanish on the one-layer Transformer used in that work. Thilak et al. (2022) argue that grokking is caused by an optimization anomaly of adaptive optimizers. We find grokking with vanilla gradient descent, so this cannot be the cause of grokking in the general case. While we show in experiments in 6 that parameters from our theory control grokking with adaptive optimizers, we note in Appendix Section 14 that learning dynamics in such settings are less well understood. Davies et al. (2023) point out that grokking and double descent share similarities in dynamics, and are fundamentally about the speed at which patterns of varying complexity are learned. This agrees with our argument that the network first learns patterns in the data that a linearized model can do well on, then changes its features learn the patterns that generalize. Liu et al. (2022a) empirically find that parameter weight norm at initialization can con- trol grokking. while Varma et al./ (2023) propose heuristic for cross-entropy training with weight decay, suggesting that grokking arises from a "generalizing circuit" being favored over a "memoriz- ing circuit." Both papers rely on weight decay and weight norm decrease to explain grokking. How- ever, we show examples of grokking with zero weight decay and an increase in parameter weight norm during training. We comment more on relations to existing theories in Appendix Section D
Kernel dynamies and feature learning in polynomial setting. The limiting behavior of neural networks in the large width limit (under commonly used parameterizations) is a linear model with the NTK as the kerne (Jacot et al. 2018). This has motivated investigations into the inductive biases of kernel methods (Bordelon et al. 2020: Canatar et al. 202I). which have revealed that kernels are strongly biased to explain training data using the top components of their eigenspectrum. While this spectral bias can enable good generalization on some learning problems, target functions that are not in the top eigenspaces of the kernel require a potentially much larger amount of data. Many papers have proposed examples where neural networks can outperform kernels statistically by amplifying task-relevant dimensions compared task-irrelevant dimensions in the input space (Mei al. 2018 Paccolat et al.) [2021] Refinetti et al. 202I). One example of such a problem that has attracted recent attention single or multi-index models where the target function is a polynomial that depends on a small subset of the high dimensional inputs (Ba et al. 2022; Arnaboldi et al. 2023 Nichani et al.,
2022 Atanasov et al., [0222] Berthier et al. 2023] Bietti al.) 2022). The fact that feature- learning neura networks can outperform kernels on these tasks suggests that they could be useful toy models of grokking, especially if early training is close to linearized dynamics.
WEIGHT NORM AND WEIGHT DECAY CANNOT EXPLAIN GROKKING
Previous explanations of grokking rely on weight norm decreases late time (Liu al. 2022a
Varma et al. 2023). In Figure 2 common task under study in grokking (Power al. 2022 Nanda al., 2023 Gromov 2023), ) with
we show simple example of a modular arithmetic task, which is a
a two-layer MLP trained without any weight decay (details in Appendix Section 8.3). Weight norm of parameters increases during training (Figure((a)) so grokking cannot in general be explained by theories of weight decay.
We will focus on the dynamics of grokking in the loss (Figure b)), because loss, not accuracy, guides training dynamics. Further, we show in Section 19 of the Appendix that accuracy curves for regression tasks are misleading because they can be gamed by choice of metric. We use mean- squared error (MSE) loss this classification task to adhere to the convention for how grokking was seen in modular arithmetic in results remain consistent across architectures, optimizers, and datasets out of the box.
Liu et al. (202a)); /romov (2023). We show in Section|othal our
3

---

Published as a conference paper at ICLR 2024
Text ices
Weggh:
1083
Sparte
2:03
Speete
2003
2.300
Eooths
100.2
10/.0
(a) Accuracy Curves
(b) Loss Curves
(c) Sweep over
Figure 2: (a) and (b) demonstrate accuracy and loss curves for grokking on a modular arithmetic task that shows increase in parameter weight norm during training. (c) is sweep the laziness parameter a that we will introduce and provide theory for, showcasing how can continuously control grokking.
While we gave the modular arithmetic task as a counterexample first because a common task under study in grokking, will focus polynomial regression task in the rest of the paper, which we can analyze the dynamics of mathematically. We introduce that task in Section 5 Pa- rameter weight norm in that example also increases inconsistencies motivate our study which results in two key parameters network laziness and NTK alignment- that introduce. The effects of varying the first, a, are shown Figure ([c) and see it can make grokking more dramatic, or vanish entirely.
groks, visible in Figure (c). These
4 GROKKING AS TRANSITION FROM LAZY TO RICH LEARNING
Grokking is often conceptualized as a network that first finds memorizing solution which fits the training set but not the test set, later converging to a generalizing solution. We hypothesize that the "memorizing" solutions that have been documented in prior works on grokking are equivalent to early-time lazy truining dynamics.
late time, the linearized approximation breaks down as the network extracts useful features and begins to "grok."
We next briefly review and define the idea lazy training regime neural networks, and introduce some key theoretical concepts that will play an important role in our discussion.
4.1 PRIMER ON LAZY LEARNING: LINEAR VS NONLINEAR TRAINING REGIMES
In the lazy training regime, a network fww, x)) of inputs x and parameters w obey the approximation
fww æ) f(wo, x) + V V.. (ww x)) lmo (w wo)
(I)
While this approximation is still capable of learning nonlinear functions of the input æ. it is linear model in the trainable parameters w. It can thus be recast as kernel method with the neural tangent kernel K
K K( x') V (w,x)) u., V Vwf <w)
evaluated at initialization. This holds for any linearized model trained with GD on any loss, though some loss functions (e.g. cross-entropy) cause non-linear models to eventually deviate from their linearization (Lee et al., 2019). Deep neural networks can approach this linearization in many ways:
Large width: in commonly used parameterization and initialization schemes of neural networks, increasing the network width improves the quality of this approximation Jacot al. (2018 Lee et al. (201)).
Large initial weights: increasing the size of the initial weights also causes the network to behave closer to linearized dynamics (Chizat et al.) 2019). We argue that this explains why Liu et al. (2022a) and Varma et al. (2023) observed an association between grokking and large initial weight

---

Published as a conference paper at ICLR 2024
Label rescaling: scaling y by factor a= can induce lazy training (Geiger et al. 2020).
Output rescaling: multiplying the network output logits by large scale factor can also induce lazy training (Chizat et al. 2019).
In this work, we primarily use the last option, where the a + oo limit is one where equation becomes exact and the network behaves at inference- time like linear model (in the case of MSE, this linear model is kernel regression with the initial NTK). In Figure we show that these methods of inducing laziness are equivalent in inducing grokking.
Inductive bias of kernels and deficiency of kernels on misaligned tasks: Though linearized mod- els around the initial weights can perform well some learning tasks, they do not generalize well when trained on target functions y(æ) which are misaligned to the NTK, in a sense we make precise when we define mentered-kerne alignment (CKA) in the next section. By contrast, neural networks in the feature learning regime can adapt their internal representations to improve the sample com- plexity of learning for functions that would be difficult for the NTK (Ghorbani et al. 2019: Arnaboldi et al. 2023
Ba et al. 2022 Damian et al., 2023).
5
POLYNOMIAL REGRESSION IN A TWO LAYER PERCEPTRON
To illustrate our main ideas, we study the simplest possible toy example of grokking: high dimen- sional polynomial regression. This choice motivated by known separation results between neural networks and kernel methods. First, theoretical studies of kernel regression have demonstrated that learning degree polynomials dimension D with the NTK requires ~ D* samples (Borde- lon et al. 2020: Canatar et al.) [021 Xiao et al. 2022). On the other hand, in the feature learning regime, two-layer networks trained in the online setting can learn high-degree polynomials with
P~ D samples (Saad and Solla, 1995 Engel 2001 Goldt et al., 2019] Arnaboldi et al. 2023
Berthier et al. 2023 Sarao Mannelli et al. (020). This indicates that neural networks can poten- tially generalize for D < P < D but only if they learn features (Atanasov et al., (022) We therefore predict the possibility of grokking in this intermediate range of sample sizes, especially feature learning occurs late in training, as in a lazy network. We indeed find we can induce grokking
this way on a two-layer network. In fact, grokking persists even with readout weights of the 2 layer
MLP fixed 1, so in seeking a minimal setting reproducing grokking, this is what we study" This type of two-layer NN is referred to as committee machine (Saad and Solla 1995).
The model f(w, æ) and target function y(x) are defined in terms of input æ E IR_ as
[(w, x)
a N
(ww; x), <(h) = h 112 y(x) マ (B, . x²2 2
The value of a controls the scale of the output, and consequently the speed of feature learning. The value of alters how difficult the task is for the initial NTK. We consider training on fixed dataset {( Mp.) ) of / samples. The inputs æ are drawn from an isotropic Gaussian distribution NO, D)). It will be convenient to introduce the following two summary statistics
N w. ERRD
N
M
E E
(4)
Using these two moments of the weights, the neural network function can be written as f(x) awx + ac x. Mx. Further, the NTK can also be expressed as 2
K(x, x') = x x' + e(x x²) w (x + x') + e² (x (x- x') x Mw'
At initialization in wide network, v = 0 and M = I. Diagonalizing this initial NTK with respect to the Gaussian data distribution reveals that linear functions all have eigenvalue Am = D-1 while
2 'The mechanism behind grokking in this setting is the same for training a full 2 layer MLP with readouts, but the kernel is not only a function of w. M. We discuss this in Appendix Section [1 but this is technicality: the same arguments and plots hold in ordinary two-layer MI.Ps the way one would expect.
5

---

Published as a conference paper at ICLR 2024
quadratic functions have eigenvalue Aquadr 2eD (Appendix Section 11 We that controls
the power the kernel places in quadratic functions (including the target (())). This is what we mean by saying controls initial alignment in this example. Further, based prior work on kernel
regression, generalization with the initial kernel (or completely lazy neural network) will require
D2 samples (Bordelon et al. 2020 Canatar al. 2021 Simon et al. 2023 Xiao generalize from far fewer samples since they can adapt their features so that M aligns to B.3. 3. and low test error is achieved, If this
P quadr
2e
et al. 2022). However, neural networks outside of the kernel regime
which does require as many samples (Ghorbani et al., 2019; Atanasov et al. 2022). feature learning is delayed, the neural network dynamics at early time would match a kernel method, but late time one would see generalization as weights begin align to f., giving learning curve that groks. This is the task that generates the learning curves in Figures 3
While c is a natural measure of alignment in our toy model, we also want general way determine how well aligned network is for task, y(X). It turns out that centered-kernd aligmment (CKA),
y Koy T || (Kornblihh et al. 2019 Cortes et al. 2012: Atanasov et al. 2022 Arora et al. 2019 is a natural generalization of what e is in our toy model. In the formula, Ko is the gram matrix of the initial NTK evaluated on the test set, and y are the task labels. We derive relationship between and the CKA in Section [1] of the Appendix, finding they are strongly correlated. Since the CKA can be computed for any lask, a general way to measure feature alignment task. "NTK alignment" in plots without further definition refers to the CKA.
5.1 ALTERING SCALE a CAN INDUCE OR ELIMINATE GROKKING
We next consider the effect of scale a and show that large a increases the timescale separation be- tween the decrease in train loss and that of test loss, while small can eliminate grokking altogether. Consider Figure 3 (a) to see that increasing a on the polynomial regression task continuously induces longer delay between train loss decrease and test loss decrease. Further, the - so limit gives poor final performance since it corresponds to regression with the initial (misaligned) NTK.
Weight norm of parameters is not inherently fundamental. key claim of ours is that previous work|Liu al. (002a)); Varma et al. (2023) saw relationship between parameter weight norm initialization and grokking because weight norm at initialization controls rate of feature learning by moving the network into the lazy training regime. We can test this claim using the fact that parameter weight norm at initialization is not the only change that can cause lazy training: in the same papers (Chizat et al.) 2019 Lee al. 2019) it is also proved that increasing model output scale or label scale has the same effect. Indeed, as we see in Appendix Figure 7 changing these has the same effect on learning curves as changing weight norm at initialization.
5.2 INITIAL LERKEL-TASK ALIGNMENT CAN ALSO ALTER GROKKING DYNAMICS
The initial alignment between the NTK and the target function also controls grokking behavior. Indeed in Figure B). we see that small e networks have a longer time delay between training loss reduction and test loss reduction and that large networks have nearly immediate reduction in test loss. Interestingly, worse initial alignment caused by smaller e leads to lower final test loss as it forces the network to learn features (see App. 10p. This means that networks with bad initial features can do better at end-time than those with better initial features, because they are forced to learn good features.
5.3 LOss DECOMPOSITION IN TOY MODEL.
To gain insight into the sources of generalization error in our model, we decompose the test risk into three terms that depend only on the summary statistics (t,, M) introduced earlier. The predictor can be expressed in terms of these objects: f arw x + xc M.. The test MSE of our model then
Priir work suggests learning quadratic polynomials (with linear component in B, direction) with SGD require ODD log D) steps in SGD with unit batchsize (Arous al. 2021).

---

Published as a conference paper at ICLR 2024
0-24
1.0
0=2-
0=
0-2
08
LOsS
0.6
0*21
Loss
0.6
r
0.4
0.4
0.0
125
10"
10a
20P
10"
30-
10'
10*
10a
(a) Large a Grokking
(b) Small c - Grokking
(XX, wiy)
S(X, wh)
Juay
dresn peneral
Network, Ladiress
WD
oF luay
TA70
aP luay
Lacy trtting
10 feature
Grokling, varylng
10
ER sa
Kernel Alignimnt
(c) Alters Dynamics
(d) 6 Changes Initial Kernel
(c) Empirical phase diagram
Figure 3: Lazy training (a:) and kernel-task misalignment (e) alter the grokking learning curves in distinct ways. (Top) Learning curves that show grokking, and (Bottom) corresponding parameter dynamics during learning. (a) At fixed e, the laziness parameter a controls the timescale of the delay in grokking. At small a, the grokking effect disappears as the generalizing features are extracted immediately. At large a, the model approaches linearized model. The final test loss decreases with a as allow the network learn more features. (b) The task aligmment to the initial kernel, measured by determines how much the loss falls when the network initially uses its linearized solution. Smaller e increases the amount of feature learning during training because the initial kernel does worse on the task, so feature learning is necessary. Thus lower alignment can result in better generalization. (c)-(d) Illustraiions of the dynamics at varying a, e. In (d), each plane represents different affine space spanned by the initial gradients, which are function of e. (c) Time to grok (time delay between train loss fall and test loss fall) as a function of o, showing how lazy, misaligned networks grok the most intensely.
takes the form:
1 C= - - 1)²) 4 ( 3.12 1..1
1 TrM
+
2D ac - ?
w 2
(6)
misalignment
power linear modes
variance
The first measures whether the total variance in the weights D. Σj |w = Σ Mj has the correct scale compared the targets. The second term measures whether the weights have achieved a solution with high alignment to the target function. It is closely related and strongly cor- v" Kay related to its general counterpart, the CKA, [Toll lIvTT in way we explore in Appendix Section I1 The last term penalizes any power our features put on linear components, because our target function pure quadratic. Because network activation function has a linear component, weight vectors must average out during training to learn a solution with zero overall power in linear functions. The test error can be minimized by any neural network satisfying M -e B_ and to 0.
In Figure 4 we illustrate this decomposition expected, that the initial peak in test error is associated with the model attempting to learn linear function of the input (orange curve) before later improving the alignment (dashed green loss falling late time) and reducing the scale of the linear component (dashed black).
the loss during grokking learning curve. We find,

---

Published as a conference paper at ICLR 2024
Tains Loas
FallTist Less
MSE 14
Loss
Var unce oncr compenent
Aiggmment sror component -... [inea tem enter component
Epochs
Epochs
Figure 4: (Left) Learning curves for a layer multi-layer perceptron that groks polynomial regression task (no weight decay, vamilla GD). (Right) Theoretical decomposition showing how the initial rise in test loss comes from the network putting power in the linear component from its inital NTK before beginning to align with the task features around 1-2k epochs, resulting in a delayed fall in the test loss. This delay is precisely grokking.
6 GROKKING IS NOT RESTRICTED TO MODULAR ARITHMETIC
6.1 INDUCING GROKKING ON GAUSSIAN DATA BY LOWERING ALIGNMENT OF LABELS
In this subsection, we induce grokking by manipulating NTK alignment at initialization in way that can be done on arbitrary datasets X, without touching the laziness parameter. Consider the same function in polynomial regression setup but replace the task labels with the j-th largest eigen- vector (by descending eigenvalue, so = is the eigenvector of Ko(X, X) with highest eigenvalue) of the NTK matrix at initialization (X, [XXX, X), and vary This is motivated because choosing
eigenvector with a large eigenvalue makes the CKA, v² Kay | high by construction, and vice versa. Thus alignment at initialization is high for small j (first plot), so the network generalizes immediately because starts with good features. Conversely, if we begin with large j the value of CKA is low, and the task is hard by construction, so the network will never be able to generalize with this amount of data (with much more data, we would learn features). Intermediate j eigenfunc- tions misaligned enough the initial NTK that a linearized model will not perform well, but the network has enough data so that learning features reduces the test loss. We see empirically that this occurs when j x 70 in (b) below.
0.000
Telt Loo
Talt Looe
0.537
0.001
0.975
0.043
0.00
0.833
3.00
0.60
0.211
10
Epochs
Epochs
(a) j =
(b) j 70
(c) j 100
Figure 5: Demonstrating grokking on standard Gaussian X E Dxx data with the one-hidden layer architecture we had for polynomial regression. Label vectors, y are replaced with j-th largest eigen- vectors of the initial NTK (ordered by descending eigenvalue) for E {1, 70, 100). respectively for (a)-(c). the network begins highly aligned as (a), learning curves move together. If the network features are poorly aligned at initialization, as in (c), the network cannot generalize. In the middle, as in (b), the network groks.
8

---

Published as conference paper at ICLR 2024
MNIST. Liu al. (2022a) induce grokking MNIST by increasing parameter weight norm initialization. In Figure ((a) below, we make this grokking vanish continuously by tuning a with- out changing parameter weight norm, supporting the idea that weight norm at initialization is one of several parameters that control laziness. This tuning of a is done with the quadratic scaling demonstrated in Figure 7
One-layer Transformers. In b)), we replicate the classic setting in which grokking was observed by Nanda et al.. (2023), where one-layer Transformer was trained on modular addition task with cross-entropy loss and AdamW. We find that even in a vastly different network architecture one with attention this time an adaptive optimizer and different loss function, our hypothesis that laziness controls grokking holds out of the box. Note how changing alpha can continuously make grokking vanish, or more drastic.
Train Acoricyy 2.000
Test Acorso, awo 9o)
Porkem ea
Train Acounry, 2-3.01
Test Acurney a.0.0)
Train Aonumry 3.05
Cross Entropy Loos
geld
Accuracy
Train Acouncyy awa5
Test Acorsey, a.0
vila
Test Acortoy awd5
Epochs
10
157
10030
20/00
Epochs
12:000
23089
(a) MLP MNIST
(b) Transformer on Modular Arithmetic
Figure 6: (a) MNIST. Continuously controlling grokking on MNIST by only varying laziness, o. including making it vanish by encouraging fast feature learning a 0.0001. (b) One-layer Transformer. Continuously controlling grokking on a one-layer Transformer from Nanda et al.) (2023) using the scale parameter a. Note how grokking has vanished for a 0.15, so can barely be seen amidst the train curves. Axes scales are kept faithful to the original papers. We find the same trend with stucent-teacher networks, with plots Appendix Section g
CONCLUSION
In this work, we hypothesized that grokking may arise from transition from lazy to rich learning. To explore this idea, we sought the simplest possible model in which we can see grokking, namely polynomial regression. This motivated the examination of laziness parameter, a, and the initial NTK alignment, which we showed can induce and eliminate grokking out-of-the-box. While many prior works empirically found that parameter weight norm seemed an important factor for network grokking, showed that this occurs because large initial weight norm approximately linearizes the model. In this way, our results subsume past work on grokking and show why weight decay is not always necessary.
Limitations While we provided a new conceptual framework for reasoning about grokking, some questions are left out of scope. We do not provide a sharp theoretical characterization of the required sample size or training time to transition from memorization to grokking behavior. While it is true that adaptive optimizers and weight decay are neither necessary nor sufficient for grokking to be seen, their ubiquity in past examples of grokking does suggest they can amplify grokking. One possibility that weight decay encourages the model to leave the lazy regime by forcing changes to the NTK (Lewkowycz and Gur-Ari 2020), which we explore in Section of of the Appendix. While we speculate on the role of the optimizers (momentum and Adam) inducing grokking in Appendix
14 the dynamics at play when adaptive optimizers induce grokking remain poorly understood.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
TK was supported by the Harvard College Research Program. BB and CP are supported by NSF Award DMS-2134157. BB was additionally supported by a Google PhD fellowship. This work has been made possible in part by gift from the Chan Zuckerberg Initiative Foundation to estab- lish the Kempner Institute for the Study of Natural and Artificial Intelligence. We thank Andrey Gromov, Bruno Louriero, Francesca Mignacco, Stefano Sarao Mannelli, and Nikhil Vyas for useful discussions and thank Alex Atanasov for helpful comments on an earlier draft of this paper.
REFERENCES
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. arXiv preprint arXiv:23033066133 2023.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206.0499 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint VVXiv:2309.023990, 2023.
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. arXiv preprint arXiv:2210.011 11,0 2022a.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. Advances in Neu- ral Information Processing Systems, 35:34651-34663, 2022b.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02673 2023.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kemels based on centered alignment. The Journal of Muchine Learning Research, 13(1):795-828, 2012.
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learning. pages 1024-1034. PMLR, 2020.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communi- cations, 12(1):2914, 2021.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E6655-776)) 2018.
10

---

Published as a conference paper at ICLR 2024
Jonas Paccolat, Leonardo Petrini, Mario Geiger, Kevin Tyloo, and Matthieu Wyart. Geometric compression of invariant manifolds in neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2021(4)0044011, 2021.
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference Machine Learning, pages 9936-8947. PMLR, 2021.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean-field dynamics to dimensionless odes: A unifying approach sgd in two-layers networks. arXiv preprint arXiv:2302.05588
2023.
Eshaan Nichani, Yu Bai, and Jason D Lee. Identifying good directions to escape the ntk regime and efficiently learn lew-ee plus sparse polynomials. Advances Neural Information Processing Systems, 35:14568- 1111456685555 2022.
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The on- set of variance- limited behavior for networks in the lazy and rich regimes. arXiv preprint arXiv:2212.12147, 2022.
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.0000 00055, 2023.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768-9783, 2022.
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020 (11):113301, 2020.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural networks, 2019.
Alex Damian, Eshaan Nichani, Rong Ge. and Jason D. Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models, 2023.
Lechao Xiao, Hong Hu, Theodor Misiakiewicz, Yue Lu, and Jeffrey Pennington. Precise learning curves and higher-order scalings for dot-product kernel regression. Advances in Neural Informa- tion Processing Systems, 35:4558-4577 2022.
David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995.
Andreas Engel. Statistical mechanics of learning. Cambridge University Press, 2001.
Sebastian Goldt, Madhu S Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Gen- eralisation dynamics of online learning in over-parameterised neural networks. arXiv preprint arXiv:1901.00085. 2019.
Stefano Sarao Mannelli, Eric Vanden-Eijnden, and Lenka Zdeborová. Optimization and gener- alization of shallow neural networks with quadratic activation functions. Advances in Neural Information Processing Systems, 33:13445-13455, 2020.
James Simon, Madeline Dickens, Dhruva Karkada, and Michael Deweese. The eigenlearning framework: A conservation law perspective on kernel ridge regression and wide neural networks. Transactions Machine Learning Research, 2023.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-conves losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1 (1:488888 4838, 2021.
11

---

Published as a conference paper at ICLR 2024
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. International conference on machine learning, pages 3519- 3529. PMLR. 2019.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
Aitor Lewkowyez and Guy Gur-Ari. On the training dynamics of deep networks with 1.2 regular- ization. Advances in Neural Information Processing Systems, 33:4790-4799, 2020.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21766 2022.
Theodor Misiakiewicz. Spectrum of inner-product kernel matrices in the polynomial regime and multiple descent phenomenon in kernel ridge regression. arXiv preprint arXiv:2204.1044 2022.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv22103.00055, 2021.
Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? arXiv preprint arXiv:2304. 15004, 2023.
12

---

Published as a conference paper at ICLR 2024
8 APPENDIX
ADDITIONAL EXPERIMENTS
8.1 NOTE ON IMPLEMENTING LARGE &
As was discussed in the original paper on inducing lazy training by Chizat et al. 2019), it is im- portant subtract off the initial predictor from the neural network function. Concretely, one should take as their predictor, the modified function
[( 0) = aff(2, 0) f(x 0n)]
In our experiments, we place ff in the optimizer when we compute gradients. This allows us to train with very large values of without experiencing instability at initialization. Further, to maintain consistent timescales of training, it is important to rescale learning rate ing keeps the initial derivative de fl==o independent of and thus the early training dynamics are consistent across a. Otherwise, large a networks will have too large learning rate and small a networks take too long to train.
nn/a². This rescal-
8.2 WRIGHT NORM AND INITIALIZATION SCALE HAVE THE SAME EFFECT
12-
40
0.
0.5
331
227
10
10
10
160
(a) Varying a
(b) Varying Init. Scale
(c) Vary a with oa fixed
100y a Ga
a very
102
Final Test Loss
Rinal Ta-tt Loss
[0-
191
167
10
101
a/a
0/a
10
9:1/100
Parameter Movement
Final MSE aa
(f) Final MSE Param. Move- ment
Weight scale alters grokking dynamics and final test loss primarily by manipulating the richness of the dynamics, specifically how far the dynamics deviates from the linear/kemel regime. (a) The loss curves for a two layer network trained on the quadratic polynomial with varying laziness parameter a. Note that small a networks do not grok as dramatically as large a and also achieve good test loss. (b) A similar sweep over weight initialization scale yields the same trend. (c) For a two layer network, approximate preservation of the feature learning can be achieved by choosing o -2 (d)-(f) The relative parameter change ww wol/|wy //l and final test loss is controlled through the quantity o²a.
Figure 7:
8.3 MODULAR ARITHMETIC IN A LAYER PERCEPTRON
Here are the architecture details for grokking with weight norm increase on a modular arithmetic task two-layer perceptron with vanilla GD. First, consider the classic modular arithmetic task a + = c mod p used previously to explore grokking in Power et al. (2022); [aada et al.) 2023). We find grokking in one-hidden layer network with input dimension 72 2p, where the
13

---

Published as a conference paper at ICLR 2024
input is concatenation of a, in one-hot encoding, and the output is dimension p one-hot encoded vector с E {0, 1.... P 1}. We use N = 100 as the hidden width, but a wide range of hidden widths exhibit grokking. We use 90% of all p² possible pairs for training, and the rest of the test; the learning rate is = 100. The figures below use p = 23, but we see the same curves for most integers p, for instance, any of Crucially, we do not use any weight decay, and merely use vanilla gradient descent. Our results agree with those of Gromov (2023).
{13, 17, 19, 23, 29, 31 ...) } also give grokking in the same way.
8.3.1 STUDENT-TEACHEE NETWORKS
Consider task where we initialize "teacher" 5- 100-1000-5 MLP with lanh activations, and a stu- dent with the same architecture. The student is tasked with being trained to learn outputs generated by the teacher network on an input dataset X of standard Gaussians. This task is introduced in Liu et al. (2022a) to show how grokking can be introduced by doubling the student's weight norm initialization. Here, we make that grokking both vanish and more pronounced without touching their weight norm at initialization, showing our parameter does the same thing because the deeper phenomenology of lazy training and model linearization is what is at play.
16-
Tran,
Test, am0.
Tran, 0-255
Tran,
Accuracy
Test, 0.5
0.75
Test, a- 0.75
Tran, d-1
Test,
Test, a
Epochs
Figure 8: Controlling grokking on a student-teacher task defined in Liu et al. (202a)) The purple curve corresponds to exactly the experiment in their paper. grokking induced by increasing weight norm. See that by tuning our scaling parameter to be smaller grokking vanish as on the left of the purple curve, and in fact, we can even make grokking more pronounced by increasing it. Since this is regression task, accuracy is defined as the student being with 0.001 the teacher label.
O, we can make that induced
8.3.2 GOLDILOCKS DATASET SIZE
Here we see the three regimes train set size: too little (a, d), just right (b, e), and too much (c, f). The top row is the learning curve, and the bottom row the loss decomposition. Note that we can directly see the network putting power in linear components in (a). This linear power comes from the NTK at initialization, since the activation function has linear term. This attempted (approximate) kernel method helps us interpolate the train set, but leads to worse than random behavior since the initial features are not matched to the task. Any deviations between train and test loss at initialization are because we use data set size of P 50. As we take data larger, the two match.
8.3.3 INTERPOLATING FROM RANDOM WEIGHTS TO ALIGNED SOLUTION TO MAKE GROKKING VANISH
We claim NTK alignment at initialization is one two important factors that control grokking. Figure [ suggests if we increase alignment initialization, grokking should vanish. This happened as predicted in Figure where we changed the labels, and now we' tune alignment a different way: by changing the weights at initialization. Here, start with random weights for polynomial regression task, where we see grokking as usual, and gradually add small component the final solution (computed analytically) weights to the initial weights. This increases initial alignment, and
14

---

Published as a conference paper at ICLR 2024
04
Specte
Specte
50
120
P 500
tpocha
tpocto
Epozt
(d) P 50 Decomposition
(e) P 120 Decomposition
500 Decomposition
Figure 9: (First Goldilocks zone in dataset size. layer perceptron trained usual polynomial regression task with P E (50, 500} in (a)-(c) respectively. (Second row) Loss decomposition for each of (a)-(c) respectively, showing how failure to generalize (a) comes from power in linear modes (black dashed), grokking comes from rise then fall in power in linear modes (b) and full generalization comes from low power linear modes throughout, and zero end time (optimal features have no power in linear modes), as we see in (f).
of course, decreases initial loss a bit. The crucial observation in the plot below is that increasing alignment by even little bit makes dramatic grokking (increase test loss alongside an initial decrease in test loss) vanish. We see this as we go from the dark blue curve to the gold one.
weights
Test, Init weights 0% solution, 100% random
weights
80.0%
Test, Init weights 20.0% solution, 80 0% random
Train, init weights
solution, 50 0% random
weights
random solution.
Train, init weights 75. 0% solution, 25 0% random
ASS
weights 75.0% solution, 25
Train Init weights 35 0% solution, 5.0% random
weights 95.0% solution, 5.0% random
10
10
10
Epochs
Figure 10: Varying alignment by adding some end-time solution to the initial weights. Grokking- like increase test loss vanishes with even small increase in alignment (solid blue solid red) as we add a small component of the solution weights to the initialization weights. This is evidence that initial NTK alignment controls how much test loss suffers during the initial fall in train loss as more power is put into kernel regression solution.
15

---

Published as a conference paper at ICLR 2024
9 RELATION TO OTHER THEORIES OF GROKKING
Existing grokking literature focuses on two things. One is studying modular arithmetic (and more generally group compositional) tasks and good representations to learn them, and the other is the empirical relationship between initial parameter weight norm, weight decay, and grokking.
Examples of papers exploring the first theme include the study of how Transformer learns to do modular arithmetic in Nanda et al. (2023). Gromov (2023) writing out the closed-form solution of final trained weights for another modular arithmetic task, and Liu et al. (2022b) who find a circular embedding within the network for a modular arithmetic task. A similar result was found in the orig- inal Power et al. (2022) paper that discovered grokking. Examining the geometry of embeddings and learned representations like this is only possible when you fix a specific task of interest, and often these are tasks with a natural geometric structure. These papers often give compelling insights into what and how networks Iearn representations for certain specific classes of tasks, but Liu et al. grokking on other tasks by tuning initial weight norm, while algorith-
(022a) showed we mic tasks like modular arithmetic might have some interesting structure that makes grokking easier to induce, they are neither sufficient necessary for grokking to appear.
The two important papers exploring the second theme (parameter weight norm and weight decay) are by Liu et al. 2022a) as mentioned above, and Varma et al. (2023).
Li et al. (022a) induces grokking by increasing the initial weight norm and concludes generalizing solutions lie on smaller norm spheres in parameter space. We agree that initial weight norm can induce grokking, but we claim it is not because generalizing solutions always lie on smaller norm spheres in parameter space. Our modular arithmetic task in Section and our polynomial regression task that is the focus of our paper are both counterexamples, since the final generalizing solution has larger parameter weight norm than the initial solution, and we do not need weight decay to learn it or to grok. Instead, we claim initial weight norm controls grokking by linearizing the model (moving it into the "lazy training" regime). The different learning dynamics at large weight scale in Liu et al. (2022a) is quite reminiscent of the transition to lazy training in Chizat et al. (2119).
Liu
Varma et al. (2023) explain grokking as arising from three specific causes. Their claimed first and third causes correspond naturally to objects in our paper, and it is with their second cause that our theory (and experiments) disagree. Their first claim is that grokking requires "two families of circuits," one that memorizes and one that generalizes. This corresponds to our regimes of lazy and rich training dynamics. Their third claim is that during grokking, the "memorizing circuit is stronger at early phases of training" which in our language corresponds to the network starting off lazy before learning features at late time.
Their second and core claim is where the two theories disagree. They claim that this transition from the memorizing to generalizing circuit happens because the generalizing circuit is more " "efficient" than the memorizing circuit, in the sense that it can produce equivalent loss with lower parameter norm. In our paper, we showed two different tasks (modular arithmetic and polynomial regression) in which the generalizing solution that reaches zero loss is higher parameter norm than the memoriz- ing solution the network started in. We do agree with their intuition that there are two regimes (lazy and rich) in which the network can operate, and parameter norm has something to do with which regime we are in. In fact, we explore the role of weight decay in Section [3] finding it consistent with our theory that weight decay can be helpful for inducing grokking.
A note data set size. It has been long known that grokking happens in particular data regime. In the paper introducing grokking, PPweer et al. (2022) mention in Section 3.1 of their paper that for large data set sizes, training and validation track each other, and Nanda et al. (2023) notes in Section 5.3 of their paper that with enough data, there is no longer a gap between train and test loss. Varma et al. 2023) study the behavior of learning curves around the critical dataset size Dant finding variants of grokking behavior. Note that being in the "goldilocks zone for data set size" is necessary but not sufficient see grokking.
Parities as a binary polynomial problem: Barak et al. (022) study learning of parities with a neural network which exhibits grokking-like dynamics when learned with two layer networks with SGD, The target function they consider has the form y = Il x; where S C [ya] is a subset of the first n bits. The variables z; ~ {+1} are random binary variables. For dot-product kernels, such as wide networks at initialization, this problem has the same sample complexity as learning polynomials
16

---

Published as a conference paper at ICLR 2024
of degree |S| with Gaussian data for kernel regression MMisiakiewicz, 2022). Our choice to study the learning dynamics of polynomial regression with Gaussian data is therefore closest to this prior work.
10 "HARD/MISALIGNED TASKS" FOR THE KERNEL CAUSE A BREAKDOWN IN LINEARIZATION
We denote "hard task" for the initial kernel as task where the trainset label vector y E R' and the kernel gram matrix K E RxP satisfy
Difficulty
K K 1.
(8)
One way to motivate this quantity is to note that near the kernel limit, the amount that the parameters move during training under MSE is precisely given by this quantity
w wol* y K K-y
We see that the linearization approximation can break for difficult tasks since w - w could be large (think about Taylor expansions around wo). Thus lazy learning on tasks of infinite difficulty is not self-consistent. All other things being equal, we expect to see more feature learning on tasks
that are difficult for the initial kernel. This especially true if the training set is sufficiently large.
11 DERIVATIONS FOR TOY MODEL
11.1 INITIAL NTK DIAGONALIZATION
At infinite width, the network has the following value under random initialization (since w 0 and
M = I)
K K(x, .x) x' + e:: x²)²
(10)
The Mercer eigenvalue problem for data distribution p(æ) has the form
I dæ p(æ)) (')((x)=)))))))))))))))())) Ao(x')
(11)
We seek eigenfunctions ((x) and eigenvalues that satisfy the above equation for the Gaussian data density x ~ D- We first note the following
(K xz.) = ((x - x' )')) D d
(12)
which implics that n, is an eigenfunction with eigenvalue Alin = 1. Next, we see that
2ca
2²
2
x<)
1
x 1
(13)
D2
D2
D
D
This implies the existence of D eigenfunctions
2ca
2
CE2
(
( (<x,x))) - c|atl?) = D
(x1)²
+
1 D
|x'12
D
D
2.2
(x/)² + (1
(D (D+ + 2)
|
(14)
D
2
The above an eigenfunction if c(D+2) - or c= 1. We have thus identified an additional
D eigenfunctions with eigenvalue À 2=² Lastly, consider æıæ - c2 for +
26a
(K(x, x0 (()o))))))
DD dlaal
(15)
17

---

Published as a conference paper at ICLR 2024
which transparently gives us another ( D) eigenfunctions with the same eigenvalue Lastly, note that |a// is also an eigenfunction
2e²
D2-
2 D
(K(x,
(16)
with eigenvalue (1 + }) The target function has the decomposition
1 (x 3.)²
1 x
y(x)=
+ 2D x
(17)
The first term in brackets has kernel eigenvalue 2/ while the second term has eigenvalue 1 (1+ ) This suggests kernel method would learn the function 2D |22 at sample sizes P ~ D but will only learn the full target function at P D2 This separation of timescales motivated the original investigation of this task in which to find grokking.
11.2 KERNEL ALIGNMENT
For any configuration of weights, the NTK of our toy model has the following form
K (x, x') x x' + eex x') x') (x + x') + e² t x x')x' Mw'.
(18)
In this section, we compute the kernel-task alignment metric on the test distribution which is dis- cussed in the paper. In particular, will show that this is related the correlation of M with 3.3, which is also present in our misalignment error. This alignment requires computing
(x) K (x. x (u()"))
(19)
v (K (x. x()2) (3(x)
The most important term the numerator which has the form
2²
((x))(zz
D [|3.] ((B, x)²x Mx) + (3. xx)"")
2² TTM + G. M3.13.1
SE2
(20)
D4
We that the numerator in the kernel alignment formula increases with the alignment of M the 3, direction. This term also appears the alignment error the MSE decomposition for our problem.
12 Loss DECOMPOSITION WITH READOUTS
we instead train both the input weights {w,} and readout weights v; in a two layer network
N
f(x)
a
voolu; x)) = arD x + ac
2 x
Mx
(21)
N
where the new t and M have the formulas
M Σ
w. S
Σ U2W;
(22)
These w and M can be plugged into the loss decomposition in the main text. At initialization, the kernel has the form
2 3c2 (x x²)2 +
(xx, x') = 2x x' +
+ ||112
(23)
2
This kernel is still strongly biased towards linear functions. As before, the eigenvalues associated with quadratic functions is O(cD-2)
18

---

Published as a conference paper at ICLR 2024
14
Trin Loss
Full Text Loss
Test Lows
Verience metor component
14.
Alggmmen anu component
12
125
14
028
Epochs
Epochs
Figure 11: Grokking decomposition for an ordinary 2-layer MLP, ie. the committee machine in the main text with readout weights trained. The same mechanism holds posited.
13 THE ROLE OF WEIGHT DECAY
While we showed that weight decay in model is neither necessary nor sufficient to induce grokking, is true can be helpful in generating grokking. Here, we explore why that is, speculatively arguing that it is because weight decay moves us out of the lazy regime. In Lewkowycz and Gur-Ari| (2020) it is proven that under weight decay, the NTK of network continuously evolves over time, This means networks trained using large amounts of weight decay cannot stay lazy for long.
For some intuition, consider the following. Lazy training during gradient requires the network parameters to move only through particular tangent space in a much larger parameter space. Infor- mally, this somewhat fragile condition that requires some weights to move in a certain direction while others move in another direction to make sure the network NTK doesn't change. Weight de- cay is a penalty on weight norm irrespective of which weight is contributing to this norm, and so can be thought of perturbation that quickly moves us out of the lazy training regime. When generalizing solution has a lower weight norm than our initialization, weight decay can encourage learning of useful features and thus generalization.
In Nanda et al. (2023), is shown that the amount of weight decay can control the amount of grokking (time delay). This is reminiscent of how laziness parameter has the same effect. Indeed, we show here they have competing effects in this sense. In that paper, it is shown how the grokking time delay between train and test loss vanishes if weight decay is high enough. Our hy- pothesis would predict making the network lazier compensates, recovering the amount of grokking started with. We see this happen below.
Trin
Test
Tast
Test
Epochs
Epochs
Epochs
Figure 12: (a) Start with large grokking time delay with a = 0.5, wd = (b) Increasing weight decay to wd 5 (keep a 0.5) makes grokking vanish by encouraging immediate movement in the NTK (feature learning). (c) Then commensurately increasing laziness to a 1.4 (keep wd = 5) recovers grokking by discouraging feature learning. We believe past works see grokking with high initial weight norm and weight decay because these two components imply initial laziness (high initial weight norm) and eventual feature learning (weight decay, as illustrated here).
One thing worth highlighting is that this explains why Liu et al. (202a)) found that time to generalize the grokking setting there was linear in the amount of weight decay, k. Morcover, [Nanda et al. 2023) finds that the time taken generalize drops by factor of ten in Figure 27 of that paper,
19

---

Published as a conference paper at ICLR 2024
when weight decay is increased by factor of again, again linear. It turns out that Lewkowyez and Gur-Ar Ari (2020) proves that the NTK evolves at order (k)) for k. the weight decay parameter, so our theory that grokking can be seen the transition from the NTK to feature learning regime exactly predicts the quantitative effect weight decay has on grokking timescales!
14 THE ROLE OF MOMENTUM AND ADAPTIVE OPTIMIZERS
This section is more speculative. In the main text, our focus is on networks that use vanilla gradient descent because they are sufficient to get networks to grok and a simple setting that can be ana- lyzed. This was motivated in part by the fact that the dynamics of networks optimized with adaptive optimizers are not well understood theoretically. But given that almost every paper on grokking in
the past has used Adam W (Power et al. 2022 Nanda et al. 2023 Liu et al. 2022a bb Varma al.) 2023), it is worth some comment.
In (a) we repeat our polynomial regression setup and induce grokking as [ but now using momentum = 0.95 in the vanilla GD optimizer. It has broadly the same shape, and certainly the same pattern (linear then feature learning) but with more bumps in the learning curves, presumably reflecting the momentum in gradient descent. In (b) we want to point out another set of learning curves that generate grokking-like accuracy curves. These are the loss curves for the student-teacher task in Section 8 for which the accuracy curves plotted showed stark grokking. That task uses A as well. again showing the linearization tracking the model until the loss halves, before diverging then. The nonmonotonicity in late-time presumably product of the AdamW optimizer, much like such nonmonotonicity emerged in (a) from momentum.
Note that both (a) and (b) are characterized grokking even though the loss curves look differ- ent: the key is delayed generalization. This shows how power in the initial kernel method is often unhelpful for generalization (a), but can sometimes be somewhat helpful in feature learning down- stream, for instance in (b). Finally, (c) shows early-time linearization dynamics in the early stages of grokking with a one-layer Transformer from Nanda et al. (2023), which was trained using Adam as well. Note how due to weight decay, the network loss actually quickly deviates from its lineariza- tion, despite the presence of grokking. This, along with the fact that scaling a still tunes grokking, is evidence that the transition from lazy to rich dynamics is a sufficient but not necessary mechanism underlying delayed feature learning on general tasks and settings. If we look at Figure 7 of that paper, we this early time deviation coincides with large early time increase in initial param- eter weight norm, which could be due early-time behavior of adaptive optimizers Lewkowycz and (uur-Ari 2220); Cohen al. (2021); Thilak et al. (2022), hence then claim that dynamics induced by adaptive optimizers in settings that grok are still not fully understood.
20031 20-
100b
8200
Tral loss
Tan Loss
Liviarized Train
Test Los Lnnerrizd Best Loss
Lncotized Fain
Innerteed Be:
Epochs (a)
Eppete
(b)
(c)
Figure 13: Learning curves for grokking with momentum and adaptive optimizers.
15 VARIANTS ON THE POLYNOMIAL TASK
The single-index quadratic learning task is known to have well-understood structure from theory perspective Arous et al. (202)) Nichani et al. (202);; Bietti et al. (2022). A natural question is whether the grokking behavior and interpretation persists under ablation of either of these condi- tions. In particular, we consider the harder tasks of learning multi-index models, and higher degree (Hermite) polynomials. We find that grokking persists out-of-the-boo as expected.
20

---

Published as a conference paper at ICLR 2024
15.1 LEARNING HULTI-INDES MODELS
It is known Nichani al. (022)) that single index models exhibit transition from lazy rich training dynamics, but in general less work has been done on multi-index models Arous et al. (202I). Thus, one might wonder whether grokking persists as we add multiple directions to the target function. That is, can we recover grokking with a target resembling (31 + 82 x) (double- index model), or even (31 x + $2 x + $acdotæ) (triple-index), turns out grokking does indeed persist in our model even if the target function spans multiple directions, as we see below.
Grolcking 2-index model
Grokiding 3-index model
Train Loss
Train Loss
0.12
Test Loss
Test Loss
MSC
0.08
MSE
0.015
0.06
0.010
0.04
0.02
0.005
10'
107
10
10
10
20
10a
101
10'
10a
10a
10
Epochs
Epochs
Figure 14: Grokking in our set-up persists when the target function is higher-degree polynomial (with activation modified accordingly). does not rely on special properties of the information exponent k- = 2.
15.2 LEARNING HIGHER DEGREE POLYNOMIALS
For the task of learning single-index model, it is known Arous al. (2021); Nichani et al. (2022) that sample complexity and learning dynamics depend heavily on a quantity of target function called the information exponent. For simple polynomial target functions, in our setting, this simply the degree of the polynomial. One salient result is the qualitative difficulty of learning a quadratic target (where the information exponent is 2), and higher-degree targets (where the information exponent is k 2 3). For instance, Theorem 1.3 in Arous et al. (202T) illustrates the difference in learnability and sample complexity for these cases. Thus, natural question is whether grokking arises from special properties of a quadratic target here. The answer to this turns out to be no: the grokking examined this paper (that is: polynomial regression in a 2 layer MLP with zero weight decay and vanilla GD) persists out-of-the-box for higher degree polynomial targets in the way one would expect. Below, we show experiments where the target is HA, the k-th Hermite polynomial, and the activation is d(2) + 5 H2 (2). In particular, H3. H2 below are polynomials with degree > and thus information exponent also greater than two. We use larger amounts of data to learn higher order polynomials Hg.
21

---

Published as a conference paper at ICLR 2024
Grokcing a H5200 hormit zarget
Grokking a #jj00 Hemmil targe:
Crokiing a M4yx Hernis target
25-
MSE
MSE
-3
MSC
22
10
Loss
Tain Loss
Fain Loss
Test Lass
Test Loss
Test Loss
Epochs
Epochs
Epochs
Figure 15: Grokking when learning the Hermite polynomials varying degree illustrates that our setup doesn't 't rely on any special properties of single-index quadratic learning problem, which is known to be theoretically simpler than higher degree or multi-index learning problems. Note the log-scale on æraxis, so small gap on the right of the plot corresponds to generalization delayed (grokking) by several thousand epochs.
16 STUDYING LOSS VS ACCURACY IN GROKKING LEARNING CURVES
16.1 DEFINING "GROKKING IN LOSS" AND "GROKKING IN ACCURACY"
"Grokking in loss" refers the loss curves study in this paper, where the train loss initially decreases while the test loss nondecreasing, with the test loss eventually falling as the network generalizes. This involves the train and test loss both moving nonzero amount initially (usually means an initial period where the test loss rises as the network memorizes the train set) and usually involves nonzero (but low) test loss end-time of training. An example of grokking loss is the right of the Figure below.
"Grokking in accuracy" refers a sharp increase test accuracy long after a commensurate increase in train accuracy. In particular, test loss stays flat at the beginning, and rises to perfect at late-time. example of grokking in accuracy is given on the left of the Figure below.
Given that these two definitions seem somewhat different on first glance, and that the second is the "original" definition of grokking Power et al. (2022) on the task (modular arithmetic) on which was discovered, one might wonder whether the definition studied in this paper (grokking in loss) is the same as that studied historically past literature.
The answer to this question is yes. The two curves on the bottom correspond to the same network run. While Power et al. (2022) present the phenomenon in terms of accuracy (like the plot below, left), really the network trained on loss (below, right). In particular, notice two things:
When plot accuracy, we the network does perfectly in terms of classification accuracy the modular arithmetic task.
However, the loss at end-time is nonzero. this case, this does not contradict the above because the network is trained one-hot inputs, and so if the maximum element in the network output is the correct index the modular addition problem, the input will be correctly answered (aceuracy) but will incur loss until the network put probability one on the index of the correct answer.
We also notice that
Test accuracy is flat early in training, but test loss moves.
Train accuracy reaches perfect classification in 2k epochs, when train loss is about ss 0.06, but train loss keeps falling much after that all the way to zero.
The intuition for the first bullet is that during memorization, test loss decreases because the network is attempting to fit linearized (NTK) solution at early time. Of course, since the network has not learned the structure of the task by this point (before 2k epochs), this not helpful on test points. But notice that this linearized solution is enough to do perfectly the train set! If
22

---

Published as a conference paper at ICLR 2024
loss of = 0.06 corresponds perfect classification, then that the test loss only dips below this amount around 3-4k epochs, precisely when the network generalizes in accuracy. Finally, the accuracy curve is flat early time because we cannot do worse than zero accuracy, so the unhelpful change in weights that allow the network to interpolate the train set are counterproductive for the test (as see in terms of loss), but the accuracy is already zero, the test accuracy curve stays flat. This is why the same network on the same run of the same task having an accuracy curve on the left derived from the loss curve on the right, is perfectly consistent. The two definitions of grokking are equivalent.
300
Train less
220
0.00
Tesst loss
Accuracy
Weight Norm
0.06
MSE
00
accuracy Test accuracy Weight Norm
0.02
0-
2000
Eporhs
Figure 16: Accuracy and loss plots for the same network run of a modular arithmetic task. This type of task has historically been of interest grokking, and is also studied in Power et al. (2022); Nanda iet al. (2023); Gromov| (2023). These plots are for a two layer network with no weight decay doing addition modulo a small prime.
Further evidence for the equivalence between the two measures can be seen by revisiting the Ap- pendix of the original paper that introduced grokking PPwer et al. (022). and noting Figure 4 which exhibits grokking in loss of the sort we study in this paper.
grokkingPoower
16.2 CONSTRUCTING ACCURACY CURVES FROM LOSS CURVES CAN BE MISLEADING
Some classes are classification tasks, such as modular arithmetic (as above) and classifying images on MNIST. There is natural notion of accuracy in such tasks. However, regression tasks can also have a notion of classification accuracy defined, as is done Liu et al. (2022a), by defining a threshold 0 such that data point (æ, y) is said to be "classified correctly" if (fr yl I O it is within threshold of the true label. Then one can define accuracy on a test set Xiest as the average number of data points in the set that are classified correctly [x,eal Ex.ex Exis: (x;;)) - yil S 0). The plot below does exactly this for the student-teacher task described earlier, taken from [Lu et al. (2022a), showcasing how the resulting accuracy curves can be misleading and unhelpful, showing apparent "grokking" for some choices of but not others.
|
16.3 NOT ALL INSTANCES OF GROKKING IN ACCURACY ARE WORTH STUDYING
Reconsider the student-teacher task from above for which could be an artefact of choosing a threshold 0. We now examine its loss curves (below, left) and see they are entirely unremarkable. Thus, seeing a network "groks" in terms of an accuracy plot does not imply interesting learning dynamics in terms of loss (which is what guides learning dynamics). Therefore, grokking accuracy does not imply grokking in loss. The curves below say more about the contrived nature of this task than about interesting learning dynamics in the network.
that grokking in terms of accuracy
16.4 LOSs IS THE RIGHT METRIC TO STUDY
We have established that grokking in terms of loss curves, as we defined it in this paper, does imply grokking in accuracy curves, defined in Power et al. (2022). We have also shown the converse is not true, demonstrating that grokking in loss, the way we have defined it, is a stronger condition than grokking in accuracy. This justifies our study of learning curves in terms of loss, not accuracy.
23

---

Published as a conference paper at ICLR 2024
Thrahhld 0.0001
Threshal 6.0001
Threshold
Accuracy
Threshad 6.005
Threshoid 0.005
Thrsshud 0.01
Threshpid
Thrwshold
Epachs
Figure 17: Grokking on this particular task depends on the choice of threshold. Each of these curves corresponds to the exact same network (in terms of loss over epochs) with different choice of threshold to plot accuracy. Figure 9 had the same threshold throughout. This justifies our study of loss, not accuracy, regression-basee tasks like polynomial regression in our main text.
MLP dentilieeeeeeeetaa
Studenk.TTacche
0.35
Train Loss
Train Accuracy
Test Loss
Test Accuracy
0.30
0.25
MSC
Accuracy
0.05
0.00
3.0
10
12a
12a
201
32'
30'
30*
Epochs
Epochs
Figure 18: Loss curves for the student-teacher task on which accuracy curves exhibited "grokking." The loss curves show that there is nothing particularly surprising happening in the network over the course of training; the train loss tracks the test loss throughout. "Grokking" on the corresponding accuracy curves is thus mirage arising from sweeping over thresholds 0 used to define "accuracy" on this regression task, and choosing a that generates striking plot.
Of course, saw in the case the MLP on MNIST, modular arithmetic tasks, and more, this means that our claims for how laziness and NTK alignment control grokking (in terms of loss) hold out of the box for grokking (in terms of classification accuracy).
We end by noting that we are not the first to make this observation. Indeed, the fact that loss. not accuracy, the interesting thing to study was quickly realized after the initial paper discovering grokking was published! It is noted when Davies et al. (2023) study loss in their Figure 5b, when Liu et al. (2022a) characterize the loss landscapes of various tasks that exhibit grokking, when Nanda et al. (2023) plots excluded and restricted loss and remarks that accuracy can be gamed, and emphatically yly in Schaeffer et al. (2023), who argues "hard threshold" measures of performance (like accuracy in grokking) are extremely misleading, and continuously optimized measures (like loss) should studied instead. Thus our choice to do so in this paper is perfectly consistent with past literature grokking.
24

---

Published as a conference paper at ICLR 2024
17 MORE DATA ALLOW US TO GENERALIZE ON HARDER TASKS
grokking is roughly seen as "delayed generalization," " then one would expect we cannot grok on tasks where we cannot generalize. This is indeed the case, and might wonder whether increasing the dataset size allows us to generalize on harder tasks - in particular on our eigenvector label task or whether this task inherently hard for feature learning network. We will see
as in Figure below that more data allows us to generalize on this task when we otherwise couldn's t. Reconsider the task in Figure 5(c) of learning the j-th eigenvector of Gaussian dataset X, for varying values of j. We chose this task because it allows us to vary task difficulty as measured by the alignment initialization of the NTK and task labels. We observe that for any fixed j. increasing the dataset size suffices to eliminate grokking by having the learning curves move together. Consider the grid of plots below. It shows how for fixed dataset size, making the task harder (increasing j) can cause network to fail (going down the left column). But then we notice (moving left to right on each row) that increasing the dataset size leads to eventual generalization, with grokking in goldilocks somewhere in the middle, when the task is hard (in NTK alignment sense), but ultimately learnable (in the sense that the network can achieve low test loss at end-time of training).
25

---

Published as a conference paper at ICLR 2024
eige segeteettooooo
egenweetee wcterrttooooo
w.g.nw
600177
Lale
600130
Trail
0090
60012
0033.00
600075
0033
0042
107
107
10+
Spade
P00
Soade
Saade
0.040
Loli
Trai Lown
5,2000
0025
0,0015
6015
0.75
30"
30
10'
107
10*
30"
10"
107
30
38
10"
301
30
Spndle
Soadle
Sondle
200 1h wwoctor data sizs P-00
200 1h agenwocto daca size Pw 300
200 3h eggnnecco data sizz: Pi
TranLoia Loia
To Loia
Tur Loie
Tex
Tex
0.2070
0.0000
0.72
6.01
6,0010
0.00
0.0000
Figure 19: Varying dataset size and task difficulty to see where rich learning is enough to generalize. Plots going from left to right in each row have increasing amounts of data, and plots going from top to bottom have the task getting harder (increasing j used for labels). We see that moving left to right illustrates (1), how increasing data for fixed task difficult suffices allow the network to generalize, but moving top to bottom on the grid illustrates (2), how if you let the task difficulty increase with dataset size, some tasks cannot be learned by this network architecture: if used even smaller eigenvectors (for instance, setting P = j). no amount P would be able to learn
26


Reviews:
[{"review_id": "IUdtoIr4Dv", "reviewer": "Reviewer_Hhk6", "summary": "In this submission, the authors proposed a simple regression model that exhibits the grokking phenomenon, that is, good generalization is achieved long after the decay of training error. The studied model is a polynomial regression problem, where the teacher model is a quadratic single-index model, and the student model is a neural network with degree-2 activation function. It is predicted that under suitable training set size and initialization scale, grokking occurs in the intermediate phase between the NTK regime and the feature learning regime.", "questions": "See Weaknesses above.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed toy model is interesting and relevant to the ICLR community. Most existing works on grokking considered modular arithmetic tasks, or measured the classification error rather than the surrogate loss used in training. In contrast, the authors studied the regression setting where grokking is manifested in the $L_2$ error, and the training procedure does not involve $\\ell_2$ regularization. Also, the connection between grokking and the transition from lazy to rich regime is to my knowledge novel, and it opens up new directions for theory research to rigorously characterize such phenomenon.", "weaknesses": "I have the following concerns. \n\n1. Given the idealized setting (Gaussian data with identity covariance, single-index target), it is rather underwhelming that the authors did not provide any quantitative characterization of the training dynamics to prove the existence of grokking. Instead, the proposed explanation is only verified empirically, which limits the contribution. Can the authors comment on the technical challenges in analyzing the gradient flow trajectory for this quadratic model? \n\n2. Connection to prior results needs to be further elaborated. \n\n* (Nichani et al. 2022) also studied the transition from the NTK to the feature learning regime, and established improved sample complexity for learning low-degree + sparse high-degree polynomials. For the current model, does the sparsity of the higher-order term (that is, whether the quadratic component of the target function span more than one direction) play a role in the presence of grokking?   \n(Nichani et al. 2022) *Identifying good directions to escape the NTK regime and efficiently learn low-degree plus sparse polynomials.* \n\n* The target function is a single-index model with information exponent 2 (Ben Arous et al. 2021). It is known that for such target function, gradient descent on the population loss will start from an approximate saddle due to the random initialization. Is this observation related to the test loss plateau in grokking? And do we expect similar findings when the information exponent is larger than 2?   \n(Ben Arous et al. 2021) *Online stochastic gradient descent on non-convex losses from high-dimensional inference.* \n\n* The quantity describing the difficulty of target function $y^\\top K^{-1} y$ has been extensively studied in the kernel literature. See (Arora et al. 2019) for the analysis of NTK.  \n(Arora et al. 2019) *Fine-grained analysis of optimization and generalization for overparameterized two-Layer neural networks.* \n\n* (minor) The cited references on the quadratic sample complexity lower bound for kernel estimators such as (Ghorbani et al. 2019) do not rigorously cover the case of Gaussian data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xvaYFfbVZJ", "reviewer": "Reviewer_y5Qp", "summary": "The paper proposes that the \"grokking\" phenomenon—where a neural network's test loss decreases significantly after the training loss—is due to a shift from initial \"lazy\" training to subsequent \"rich\" feature learning. Through a polynomial regression study with a two-layer network, the authors show that grokking occurs when the network moves from fitting data with its initial features to learning new features for better generalization. They suggest that the rate of feature learning and the initial feature alignment are key to this delayed generalization, a concept potentially applicable to more complex neural networks.", "questions": "See weakness", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper studies the grokking phenomenon in deep learning, which is a recent hot topic and very relevant to ICLR. This paper proposes that grokking can be triggered by the transition from kernel regime to feature learning regime. Though this is already known even before the grokking paper by Power et al, 2022, e.g., the example of quadratically overparametrized linear model in Section 6 of Li et al., 2021, the novelty here is that this paper focuses on vanilla GD. In contrast, the transition between kernel and feature learning regime, Li et al., 2021, is triggered by sharpness minimization implicit bias of label noise SGD, which makes the analysis significantly easier.  This paper also performed some new experiments on various tasks to support this theoretical intuition.\n\n*Reference:*\n\nLi, Zhiyuan, Tianhao Wang, and Sanjeev Arora. \"What Happens after SGD Reaches Zero Loss?--A Mathematical Framework.\" arXiv preprint arXiv:2110.06914 (2021).", "weaknesses": "1. The definition of grokking seems to be very different than that in literatue. In Power et al., 2022, **Groking** refers to the phenomena that \"long after severely overfitting, validation accuracy sometimes suddenly begins to increase from chance level toward perfect generalization\". While this paper describes grokking in their introduction as \"train loss of a neural network decreases much earlier than its test loss\". \n\n\tIt is ok to me that this paper only focuses on the regression setting and talks about MSE loss. But the original grokking phenomenon is much more \"suprising\" and rare in the sense that it requires the test loss/error to get higher than random guess only after reaching almost zero training loss/error. Moreover, the original definition of grokking requires almost perfect generalization in the end of training.\n\n2. The theoretical explanation presented in this paper only works for the new definition of grokking, but not the original definition of grokking in Power et al., 21. For the most experiments in this paper where the authors claim grokking happens, two criterions of original definition of grokking are not met: 1. when test loss starts to decrease before training loss gets small; 2. the eventual test loss is still high. Figure 5(b) is an example.\n\n3. There is no end-to-end theorem about grokking. Given the modified definition of grokking, explanation using transition from kernel regime to feature learning regime seems quite straight forward. It should be very standard to show that there exists a period of time where training loss can decrease but test loss has a lower bound due to live in the kernel regime. Still this paper lacks such precise and rigorous statement.\n\n4. I find it difficult to understand the math derivation for the toy model in section 11. In the paragraph above section 11.2, the authors write that kernel method can learn $|x|^2$ at D sample sizes and the full target function at $D^2$ sample sizes. How is this gap in sample complexity related to the sepration of timescale in a single run?\n\n5. The first paragraph of Introduction does not sound correct. For example, the authors write that \"Typically, one assumes that a model’s performance on its training set will be a reliable indicator of its generalization capabilities on unseen data\". This is in general not true, especially in deep learning where models are often overparametrized. And it is exactly because training loss and test loss can have a large gap in practice, people are studying generalization in machine learning. Also it is not clear to me that why using models which groks in training will cause any safety issue. I would like the authors to elaborate on this or at least provide an example. Intuitively, grokking as a training behavior has nothing to do with the performance of the model when used for inference.\n\n6. As mentioned in the strengths, the explanation of grokking using transition from kernel to feature learning regime (Li et al., 2021) is known even before the original grokking paper (Power et al., 2022). Though this paper is novel in the sense that it focuses on vanilla GD, which is more natural than Li et al., 2021, the theoretical result in this paper is also not as clean as that in Li et al., 2021.\n\n*Reference:*\n\nLi, Zhiyuan, Tianhao Wang, and Sanjeev Arora. \"What Happens after SGD Reaches Zero Loss?--A Mathematical Framework.\" arXiv preprint arXiv:2110.06914 (2021).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lmf3YK1LWi", "reviewer": "Reviewer_sTxd", "summary": "This work hypothesizes that grokking may happen as a transition from lazy to feature learning regime. It also shows that grokking is not inherently associated with an increasing norm and the use of weight decay -- which advances our understanding.", "questions": "One point I think worth being clarified is the role of weight decay. I think we have a pretty good understanding of how weight decay changes the kernel. Essentially, weight decay almost always changes the NTK at the speed of $\\kappa$, which is the strength of weight decay. The time scale for the effect of weight decay to take place is nothing but $1/\\kappa$. Therefore, if $1/\\kappa$ is far larger than the time scale of lazy training, one has a separation of time scale, where lazy training happens first, and feature learning induced by weight decay happens later. I think the authors should make more discussion on this point as it constructively adds to the paper.\n\n\n\n------\n\nThis part is updated to acknowledge that I have read the author response.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the paper is clean and the point being made is sufficiently important. Also, I find the experiments sufficiently convincing. \n\nCrucially, it advances our understanding of the grokking phenomenon. The shortness of the review only reflects the fact that I do not much much criticism for this work.", "weaknesses": "Nothing major.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MPCUdavqUy", "reviewer": "Reviewer_8XT2", "summary": "The paper studies the grokking phenomenon and argues that the drop in test loss (after the drop in train loss) coincides with the network transitioning from the lazy to the rich regime. Grokking would then happen when the NTK at initialization is ill-aligned with the task, leading to a large test error in the lazy regime, however after the transition to the rich regime, the `right' features are learned.\n\nThey study a simple model of shallow nets that can fit linear and quadratic function, with a hyper-parameter that allows one to tune how much stronger the alignement of the initial NTK with the linear modes than quadratic ones. For a quadratic true function, one can ensure a bad alignement between the NTK and task, thus leading to grokking.\n\nFinally this theory is tested on more practical examples (MNIST with a MLP and arithmetic with a transformer architecture).", "questions": "You argue in a few different places that grokking when the NTK is ill-aligned to the task, but it feels to me that another condition is that the task needs to be well adapted to the rich regime (which I understand is much harder to characterize). There are tasks where neither the lazy nor the rich regime generalize, in which case one expects no grokking.\n\nIn section 6.1 you select a true function that is explicitely ill-aligned with the initial NTK to induce a grokking phenomenon, but the grokking phenomenon disappears when using too small eigenvalues. You argue that this is because there is not enough data, did you test this hypothesis? Could it also be the case because the eigenfunctions corresponding to small eigenvalues are hard to learn regardless of the regime (as I propose above)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "There has been a lot of work studying grokking in different settings, but we are still lacking a complete understanding of under which condition it appears (model, training method, task). I think that the lazy/rich transition is the right idea to understand grokking.\n\nUnder this interpretation, we expect grokking to appear in tasks that are adapted to the rich regime but not the lazy regime. Furthermore, we need models that exhibit both rich and lazy regimes, and finally we need a training method where the model switches from lazy to rich regime in the middle of training.\n\nThe authors identify a simple theoretical model that exhibits grokking, and its simplicity allow for some theoretical results. The numerical experiments suggest that their interpretation applies to more realistic setups.", "weaknesses": "The theoretical and experimental evidence given in the paper is not particularly strong: in spite of the simplicity of the model considered, the theoretical results are very limited (just a decomposition of the loss into three terms with relevant interpretations); conversely most empricial results are done on this simple model, with just three experiments done on more complex data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)