Omnigrok: Grokking Beyond Algorithmic Data
Authors: Ziming Liu
Eric J Michaud
Max Tegmark
Venue: ICLR 2023 notable top 25% (2023)
Keywords: grokking
loss landscape
neural dynamics
representation learning
initialization

Abstract: Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive. We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the mismatch between training and test losses as the cause for grokking. We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively. This simple mechanism can nicely explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc. Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules, although the grokking signals are sometimes less dramatic. We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.

Published as a conference paper at ICLR 2023
OMNIGROK: GROKKING BEYOND ALGORITHMIC DATA
Ziming Liu, Eric J. Michaud & Max Tegmark
Department of Physics, Institute for Al and Fundamental Interactions, MIT immliu, ericim tegmark r @mit. eduu
ABSTRACT
Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive. We aim to understand grokking by analyzing the loss landscapes of neural networks, identify- ing the dependence of the generalization gap on model weight norm as cause of grokking. We refer to this as the "LUU mechanism" because training and test losses (against model weight norm) typically resemble "I." and "U", respectively. This mechanism can explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc. Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules, although the grokking signals are sometimes less dramatic. We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.
INTRODUCTION
Generalization lies at the heart of machine learning. A good machine learning model should arguably be able to generalize fast, and behave in a preoth/predictable way under changes of (haperrparameters. Grokking, the phenomenon where the model generalizes long after overfitting the training set, has raised interesting questions after it was observed on algorithmic datasets by Power et al. (2022):
Q1 The origin of grokking: Why generalization much delayed after overfitting? Q2 The prevalence of grokking: Can grokking occur on datasets other than algorithmic datasets?
This paper aims to answer these questions by analyzing neural loss landscapes:
A1 Grokking can result from a mismatch between training and test loss against model weight norm. Specifically, (reduced) training and test losses plotted against model weight norm resemble "I." and "U", respectively, as shown in Figure 1b. We refer to this phenomenon as the "LU mechanism", which we elaborate on in Section 2 and 3.
A2 Yes. Indeed, we demonstrate grokking for a wide range of machine learning tasks in Section 4, including image classification, sentiment analysis and molecule property prediction. Grokking signals observed for these tasks usually less dramatic than for algorithmic datasets, which we attribute to representation learning in Section 5.
Partial answers to Q1 are provided in recent studies: Liu al. (2022) attribute grokking the slow formation of good representations, Thilak et al. (2022) attempts to link grokking to the slingshot mechanism of adaptive optimizers, and Barak et al. (2022) uses Fourier gap to describe hidden progress. This paper aims to understand grokking through the lens of neural loss landscapes. Our landscape analysis is able explain many aspects of grokking: data size dependence, weight decay dependence, emergence of representations, etc.
The paper is organized as follows: introduce the LU mechanism. In Section 3, we show how the LU mechanism leads to grokking for toy teacher-student setup. In Section 4, we show that the intuition gained from the toy problem can transfer to realistic datasets (MNIST, IMDb reviews and QM9), for which we also observe grokking, although in a slightly non-standard setup where it is relatively weak. In Section 5, we discuss why grokking is more dramatic for algorithmic datasets than others (e.g., MNIST), by comparing their loss landscapes. We review related work in Section 6 and summarize our conclusions in Section 7. Code is available https: / / chub.ccom K ndx aomí ng/ Omn Omnigrok.
In Section 2, we review background on generalization, and

---

Published as a conference paper at ICLR 2023
generalizable solutions
large in
Goldllocks zane (genn sraaaioon
Tast
slow grokknngg
overficting
Tas
test loss
solutions
large reg
small Init
Loss
small init
aarg init
Training loss
fast
fast
slow (grokking)
dew
smail reg
no reg
Wc
weight space
weight norm
(a)
Figure 1: (a) L2 norm of model weights. Generalizing solutions (green stars) are concentrated around a sphere in the weight space where w = We (green). Overfitting solutions (orange) populate the 2 we region. (b) The training loss (orange) and test loss (gray) have the shape of L and U, respectively. Their mismatch in the w > we region leads to fast-slow dynamics, resulting in grokking.
THE LU MECHANISM FOR GROKKING
Weight norm and reduced loss Letting w denote the weights of model, any function f(w) (e.g, train/test loss/accuracy) depends on both the weight norm w = ||w||2 and the angular direction w = w/w. Similar to Fort and Scherlis (2019), we define reduced function f(w) by minimizing training loss Itraiu ()) over angular directions, i.e.,
f(w)) = fww"(w), where w (ww argmin brain (w).
(1)
||w||a=10
In this paper, we set f as train/test loss/error, but it also applies to other metrics of interest. In practice, we perform the constrained minimization by rescaling the model weights back to their original norm after each unconstrained optimization step. We will see that this reduced ID loss landscape, which easy to visualize, captures important features related to grokking. Throughout the paper, our model is initialized by multiplying a factor a = w/ w/wg Wy to the standard initialization where wo and w are the weight norm of the network before and after multiplying by a, respectively.
LU mechanism Although the loss landscapes of neural networks are nonlinear, Fort and Scherlis (2019) reveal simple landscape picture: There is spherical shell in the weight space (the "Goldilocks" zone), where generalization is better than outside this zone. We illustrate the Goldilocks zone as the green area with average radius we in Figure la; the green stars are the generalizing solutions. The test loss is thus higher either both when w >> We and w < We. forming a U-shape against w in Figure 1b (gray curve). By contrast, the training loss has an L-shape against weight norm. There are many solutions which overfit training data for w >> We, but high training losses are incurred for w w. This corresponds the L-shaped curve seen in Figure 1b (orange curve, no regularization). In summary, the (reduced) training loss and test loss are L-shaped and U-shaped against weight norm, respectively, which we will refer as the LU mechanism throughout the paper.
It is well known in statistics that generalization error has "U" shape against model capacity, which usually attributed to the bias-variance trade-off. Although this common wisdom was challenged by the observation of double descent (Nakkiran et al., 2021), the "U" curve can be recovered from double descent simply by changing the x-axis from the number of model parameters N to the 2-norm of model parameters to ||w 12. at least for linear regression (Ng and Ma, 2022). Although the LU mechanism may remind readers related phenomena (Schoenholz et al., 2016; Yang and Schoenholz, exactly the same ours. More importantly,
2017; Nakkiran et al., 2021). their setups focus and contribution is to understand grokking, a brand new generalization puzzle.
By "standard initialization" mean the default one in PoTorc. For linear layers, each weight U U[-a, a) and bias b ~ Ul-o, o) where 1///an in, and U(a, bj denotes uniform distribution on b).
2

---

Published as a conference paper at ICLR 2023
Grokking dynamics We identify the "LU mechanism" the cause of grokking. If the weight norm is initialized to be large (c.g,, the black square in the w > we region), the model first quickly moves to a nearby overfitting solution by minimizing the training loss. Without any regularization, the model will stay where it is, because the gradient of the training loss almost zero along the valley of overfitting solutions, generalization does not happen. Fortunately, there usually explicit and/or implicit regularizations that can drive the weight vector towards the Goldilocks zone w R We- When the regularization magnitude is non-zero but small, the radial motion can be (arbitrarily) slow. If weight decay is the only source of regularization, and training loss is negligible after overfitting, then weight decay y causes w(t) exp( exp(-)t))o. when Wy > We. so it takes time generalize. A small results in a huge generalization delay (i.e., grokking). The dependence on regularization magnitudes is illustrated in Figure 1b: no generalization at all happens for 2 = O, The above analysis only applies to large initializations w > We. Small initializations w < we can always generalize fast regardless of regularization.
in(w// (o/////// x y to
small leads to slow generalization (grokking), and large y leads to faster generalization
Why isn't grokking commonly observed? The standard initialization schemes typically initialize w no larger than We. However, if we increase initialization scales (explicitly or implicitly), grokking can appear. In Section 3 and 4, we find that explicitly increasing initialization weight norm can induce grokking. In Section 5, we argue for algorithmic datasets because (shown in Figure 6d)
(bad representation) > we(good representation),
i.e., a proper initialization for a bad representation is effectively too large for good representation, leading to grokking. Take the addition (base p) for example: with the good (linear) representation or bad (random) representation, the decoder needs to learn to classify (pp) O(p²) examples, respectively.
3 GROKKING FOR A TRACHER-STUDENT SETUP
V-0
tran
zain
107
S04ga MOO.OOOONN 1000000333 14
3est
10-1
20
10-4
10 Weight decay
15"
relative weight norm wido
12 17
10.2
(b)
(c)
Figure 2: Teacher-student setup, a: student initialization scale, 7: weight decay. (a) The reduced training loss and test loss have the shape of "L." and "U", respectively. (b) Top row: large initialization (a 2.0) can demonstrate no generalization (no reg), grokking (small reg) and fast generalization (large reg). Bottom: small initialization (a = 0.5) always generalizes fast, regardless of weight deacy, (c) a = 2. The steps to overfitting is independent of weight decay, while the steps to generalization scale inversely with the weight decay.
To illustrate how the LU mechanism results in grokking, we employ a toy teacher-student setup. The teacher and the student share the same architecture (a 5-100- 100-5 MLP with tanh activation). but are initialized with different seeds. The student network is initialized with the standard initialization (the default one in PyTorch) but each weight is rescaled by the same factor a = w/Wo, where wo and w are the weight norm of the student network before and after rescaling. The teacher network is initialized standardly, i.e., ateacher 1. Inputs and outputs have dimensions dn, respectively. We generate Ntrain = 100 training and Neest. = 100 test samples by first drawing inputs from the standard Gaussian distribution NO, Lt,,, Xd..). ) and then feed the input data to the teacher to generate output labels. The student network is trained with the Adam optimizer (learning rate 3 x 10-4) for 1 steps.
and dout
2y should not he too large, otherwise it will bring the weights to a trivial solution w 0.
should not he too small harm optimization.
3

---

Published as a conference paper at ICLR 2023
LU landscapes Firstly, compute the reduced losses by minimizing the training loss (excluding weight decay) while constraining the weight norm of the student network to be constant. We assume the converging point after training as the global minimum on the spherical surface +, which explicitly defines the reduced losses Irrain (a) and hust (a). As shown in Figure 2a, Teees (a) first decreases and then increases a increases, displaying U-shape with a minimum a R 1. By contrast, Terain (a) decreases when a < 1 and remains flat near zero when 2 1, forming an L-shape. When weight decay is present, the training landscape becomes Itrain (a, ?) = lerain (a) + 7a² where C' is the average parameter magnitude determined by the standard initialization.
Training dynamics Our problem is regression task, but we can imitate the behavior of a classifica- tion task by manually setting a threshold B 0.01 and defining sample to be correctly "classified" if the prediction error is less than B. We study the dynamics of training and test accuracy. Note that this is the normal training setup where the weight norm is not constrained, although with two non-standard initializations
0.5 (small) and = 2.0 (large), and three weight decays y 0 (no reg), 7 = 0.03 (small reg) and 7 = (large reg). As shown in Figure 2b (bottom), small initialization runs always generalize fast regardless of regularization. Large initialzation runs (top) dependend on weight decay: no regularization fails to generalize, small regularization generalizes slowly (grokking), while large regularization generalizes faster.
For the large initialization a = 2.0, we do a finer sweep of y in [0.03. 1]. We compute the number of steps and weight norm w when training test accuracy reaches 95%. As shown in Figure 2c, the time (number of steps) to reach 95% training accuracy is independent of weight decay 7, while the time to reach 95% test accuracy is inversely proportional the weight decay, as we derived above for the LU mechanism.
4 OMNIGROK: GROKKING FOR MORE INTERESTING TASKS
We now analyze loss landscapes and search for grokking for several more interesting datasets, and see that the insights obtained from our toy model can transfer to these datasets. We report the main results here, with experiment details included in Appendix A.
Image classification We visualize loss landscapes of MNIST (Deng, 2012) to verify the LU mech- anism, and study the dependence training data size. Similar to the teacher-student case, reduce losses and errors (one minus accuracy) to two variables (weight norm w and data size N) by minimizing over angular directions of weights, i.e.,
Teraia (w, N) fwra,n("", N), [uo N) = hressw",, N); w"ww, .N = argmin lerain (w, N), (3) [|w]b/w
shown in Figure 3 (a)... The reduced loss landscape reveals three things: (1) Larger initializations lead to grokking. Point A in Figure 3 corresponds to the standard initialization (a = 1), which has low training and test errors, hence no grokking. When increasing the weight norm from A to B, training error is seen to remain low while test error rises. To generalize, weight decay must be in place to bring the weight norm down, leading to grokking if weight decay is small. (2) Larger datasets lead to de-grokking. Comparing B and C in Figure 3, C is seen to have larger training size than B and lower test error. Larger data size N makes the Goldilocks zone broader, reducing or eliminating grokking even for large weight initializations. (3) Critical data size can be defined. As reported in Power et al. (2022); Liu et al. (2022), we see that there exists a critical training set size below which generalization is impossible. The effective theory analysis in Liu et al. (2022) only applies to algorithmic datasets, but not to other datasets with unknown optimal representations. The loss landscape analysis presented is this work can apply to all supervised -eearnine-- tasks. As shown in Figure 3 (b), the contours of constant test error are thumb-like, and the tip of the thumb determines the minimum amount of data required for generalization.
Guided by the landscape analysis, we make two nonstandard decisions to induce grokking on MNIST: (1) we reduce the size of the training set from 60k to k samples (by taking a random subset) and (2) we increase the scale of the weight initialization distribution (by multiplying the initial weights, sampled with Kaiming uniform initialization, by a constant a > 1). With these modifications to the
'This is generally not true when the loss landscape is non-convex. The aim of this assumption is to make the minimizer aligned with Eq. (1).
4

---

Published as a conference paper at ICLR 2023
training
101
10
167
"dlii sice 200
traiming size
grodteng
training data size N
growking
10a
102
4A
"data size R000
10
-02
42
6x
relative weight norm awiit
rel ative weight norm a w/We
relative weight norm a w/Ws
(b)
181b1 crowiin wan varying a 1900 rran HORKS, y-224
1.00
Maae
0.75
0.30
0.5
10
10a
10
100
Optimization Steps
Panta
(d)
(c)
Figure MNIST. (a) reduced training error, (b) reduced test error. Comparing A and larger weight norm makes learning grok (delay generalization), Comparing B and C: larger training data size makes learning de-grok (speed up generalization). (c) "LU" holds truer for smaller data, (d) Accuracy curves for MNIST in the setting where observe grokking. (e) Time generalize function of training set size N, replicating Liu et al. (2022).
Ci -
acorray
daza S18.
124
1GP
12
scorayy
30'
10
Sxcu
0.2
20.0
10
10.3
weight wMih
gocnom Alle
riative weightnom nite
(a)
(b)
(C)
(d)
Figure 4: We LSTM predict IMDb reviews. (a) training error; (b) test error; (c) reduced
losses for data size Ik (top) and 50k (bottom): (d) With Ik data, (weak) grokking signal is observed for large initializations (a- 6), while no grokking observed for standard initializations (a 1).
training loss
test lows
IUTN SIE 100
date slew
5018 sice 2000
101
20
Ielative wwig norm solds
10
Ieletivn weight narm d mhlr
reboiee wely
127
scepa
127
C2
(a)
(b)
(c)
(d)
Figure 5: We use GCNN predict isotropic polarizability of molecules in the QM9 dataset. (a) training loss; (b) test loss; (c) reduced losses for data size 100 (top) and 3000 (bottom); (d) with 200 training samples, grokking is observed for large initialization 3).
5

---

Published as a conference paper at ICLR 2023
training set size and initialization scale, we train a depth-3 width-200 200 MLP with ReLU activations with the AdamW optimizer using MSE loss with one-hot targets. We find that the network quickly fits the training set, and test accuracy improves much later, shown in Figure 3d, just in the stereotypical grokking learning first observed in algorithmic datasets. Figure 3e shows the effect of training set size on time to generalization for MNIST. We find result similar what Power et al. (2022) observed, namely that generalization time increases rapidly once one approaches a certain critical data set size. The conclusions still hold for the cross entropy loss (see Appendix F), although with quantitatively milder effects.
Sentiment analysis of text We look for grokking using LSTMs (Hochreiter and Schmidhuber, 1997) for IMDb dataset (Maas et al., 2011). Similar to Eq. (3), we reduce training and test losses to depend on only the weight norm w and data size N. We show the reduced training and test error in Figure 4 (a)(b). For large data size, e.g., the full dataset, training and test errors have similar "U" shapes 5, so one cannot create grokking via the "LU" mechanism. For small data size, say 1k, however, the mismatch between training and test ertors makes it possible to create grokking via large initializations. In Figure (c), we initialize weights larger (a= = 6) with weight decay 1, overfitting is complete within 102 steps, but generalization does not start until around 10 steps. Note that the generalization "jump" is not as sharp as on algorithmic datasets (Power et al., 2022) or MNIST, but at least generalization is delayed here. By contrast, if use the standard initialization (a = 1) with no weight decay, generalization happens early on during training, and does not improve much after overfitting.
Molecules We search for grokking using the graph convolutional neural network (GCNN) for QM9 dataset (Ramakrishnan al., 2014). Similar to Eq. (3), we define the reduced training/test losses, which are only dependent on weight norm w and data size N. As shown in Figure 5(a)(b), when data size is large, training and test losses have similar "U" shapes, hence grokking is impossible via the "LU mechanism". When data size is small, training and test losses mismatch somewhere in the region a = w/w > 1, making grokking possible. Indeed, shown in Figure 5(d), there is a sharp drop in test loss around 10a steps if initialization is 3 times larger than standard, while standard initialization does not lead to grokking. Note that zero weight decay is applied in both cases, implying the existence of implicit regularizations.
5 REPRESENTATION IS KEY TO GROKKING
In Section 4, we showed that increasing initialization scales can make grokking happen for standard ML tasks. However, this seems a bit artificial and does not explain why standard initialization leads to grokking on algorithmic datasets, but not standard ML datasets, say MNIST. The key difference is how much the task relies on representation learning. For the MNIST dataset, the quality of representation determines whether the test accuracy is 95% 100%; by contrast in algorithmic datasets, the quality of representation determines whether test accuracy is random guess 100% (good representation). So overfitting (under bad representation) has a more dramatic effect on algorithmic datasets, i.e., the model weights increase quickly during overfitting but test accuracy remains low. During overfitting, model weight norm is much larger than at initialization, but then drops below the initialization norm when the model generalizes, shown in Figure (see Appendix C), and also observed by Nanda al. (2023).
(bad representation)
In the following, we will compare algorithmic datasets (Section 5.1) to MNIST (Section 5.2). We show how their loss landscapes depend on representations differently, and how the difference leads to different outcomes (grokking or not).
5.1 ALGORITHMIC DATASETS
Setup Algorithmic datasets are the task of learning a binary operation a o b = c ca, b, c are symbols) with neural networks, which aim predict e from input (a, b). We take the toy addition setup in (Liu et al., 2022), where each input digit 0 < i S p - (output label 0 < l < 2(g 1)) is embedded as a vector E, R). A decoder MLP is employed to predict Y2 = Dec E, + E, (k: i+j). In the setup of grokking, both the decoder and the input representations R E, are trainable, with learning
In principle, reduced training losses should be non-increasing ("I."), but optimization issues may occur for large initializations (Schoenholz.c et al., 2016).

---

Published as a conference paper at ICLR 2023
training loss
test loss
representaioo messiness
representation messiness
1.
1773
1d1
3.6
3.4
100
represetition
3.4
1.0
leaming
9.9
B
3.2
train
fast
666
test
8.4
3,00
Path
relative weight norm a w/Wo
relative weight norm a w/We
(c)
(a)
(b)
training less (size 45)
training loss 0.02 contour
representation messiness m
rep reprrssentation messiness m
good
V ycose
representation
gradion
-0.4
g
-.4
V=
Alrewds
6
0
bad
relative weight norm a= w/We
relative weight norm a = w/We
log (weight norm)
(d)
Figure 6: Loss landscapes on the 2D (w, m) plane. (a) Training loss splits the plane into two regions: large loss small w (fast dynamics) and small loss large w (slow dynamics). (b) Test loss; the green star is the generalizing solution. (c) Losses along an illustrative path A + E, demonstrating multiple descent; (d) zoom-in of the training loss highlighting the gradients on the boundary. (e) the boundary depends training data size; (f) simple illustration of grokking dynamics.
rates IID and 7/R, respectively; in the setup landscape analysis. only decoder is trainable, as explain below. Training and test losses depend on three factors: (i) representation R. (ii) weight norm and (iii) weight direction w. in previous sections, we can optimize w by minimizing the training loss constant weight norm spheres. We further reduce the high-dimensional representations to ID by interpolating in particular direction:
R = m Ryandoo + (1 (1mm)] Rinear
4))
where Riincar refers to the linear representation in which number & is embedded to Eg = 1, 0.
0].
Rrandom is the initialized representation drawn from Gaussian distributions, i.c, Eg ~ (0, 1), and
E [0, 1] is scalar interpolating between Rlinaar and Ryandom- that we term representation messiness because reductions, both training and test losses become functions of two variables, representation messiness m and weight norm w:
Riinear when m = 0, and R. = Rrandom when m = 1. After these
(w, m) = argmii (erain(wm))) Terain (w. m) = (rrini""... (w".m), Tyess (w, m) Leest (w" m)))
(5)
||wl|=
Note that our definition of Turai (w, m) excludes the weight decay term freg = kyw? 2 be aware of its presence when we analyze the dynamics of (w, m.), which is governed by the gradient
but we should
flow Ierain (w, m) plus weight decay (πR/TD are learning rates of representation/lecodeen) dw
Touaiin
dm di TR
Hirain
di = -nd
(ow Ow
+ yW
3m
(6)
More experimental details are included in Appendix E.
Landscape We show 4 m) and Feest (w, m) in Figures 6a and 6b, indicating the generalizing solution with green star. Based on the reduced training loss (Figure 6a), divide the 2D plane
7

---

Published as a conference paper at ICLR 2023
into two regions I and II, separated by a dashed yellow line (the contour of training loss 0.05): (I): The darker region, with high training pessss//rrdients and small weight norm. (II): The lighter region, with low training losses/gradiente and large weight norm. Comparing Figures 6a and 6b reveals that training and test loss landscapes differ, especially in region I. Moreover, while the training loss depends weakly
m, the test loss depends strongly on m. As we will see, the (weak) dependence of training loss on representation drives the model to the generalizing solution. However, the driving force is small because the dependence is weak, leading grokking. We elaborate below how these particular loss landscapes lead to grokking.
Grokking dynamics In region II, the dynamics is slow (for small ?) due to nearly vanishing gradients. By contrast, the dynamics in region I is relatively fast. As we will explain, dynamics is also slow on the boundary of I and II, and grokking is the consequence of traversing region II and/or the boundary.
Let us analyze typical path A to E shown in Figure 6(a)(b). A rolls "downhill" to B following training gradients. possibly continuing to C due to momentum. C is located in II where 4trai ~ 0, so according to Eq. (6), dm/dl x and dw/d So (logw, m) moves with constant speed - in the -w direction from C to D, point near the boundary. Negative gradients around the boundary point towards larger w and smaller m. shown in Figure 6d (a zoom-in of Figure 6a). The gradients become increasingly large the model goes deeper inside region I, and at some point, the gradient totally cancels out v in the gradient direction, making the model start to drift along the boundary, as illustrated in Figure 6f. Then the model moves along the boundary with a new velocity v' = vcoso 6, until it reaches the generalizing solution E. The above picture is supported by empirical experiments in Appendix C and also Nanda et al. (2023). Based on the picture, we also show the ability to eliminate grokking in Appendix C.
-IjDYw or, equivalently, d(logw) / R -j-"-
The slow dynamics from C to E is the origin of grokking. During this period, the model first moves in the w direction with a velocity over the distance L1 = L hcotol, and then moves along the boundary with a velocity r' over the distance L2 = h/ /isnnf. So the total time is LI/U + L2/0' = (L + htantt/ (npd) This formula agrees with the observation that large weight decays Y and/or larger decoder learning rates dp can make generalization happen faster (Power et al., 2022; Liu et al., 2022). Besides, the path manifests intriguing multiple descent of test loss, shown in Figure 6c.
Dependence of grokking on training data size Another important observation in Power al. (2022) is that grokking happens faster for larger training size. Our landscape analysis can also explain the data size dependence. Figure 6e, we show the contours (training loss = 0.02) for different training sizes (25, 35, 45, 55). The contours of training size 45 and 55 both connect to the green star, meaning that generalization will eventually happen. However, the slopes of the contours are different, i.c., aas < Bas. Since t (L + (ttan))(nD>)) increases as 9 increases, training data leads to faster grokking. For training size 35 and 25, the contours do not connect to the green star, so generalization will not happen, matter how long the training will be run.
have tss < tas, i.e, more
5.2 MNIST
We now study how training and test losses depend on representation messiness in the MNIST dataset. We denote the 28 x 28 images as the raw representation Ryaw. We construct linearly separable representation Rlinear by assigning input representations proportional to their label ys for example, an image of a is represented by a matrix with all elements being 2. Similar to Eq. (4), we use m E [0, 1] to interpolated between R and Riinear:
R R- R=mRraw + (1-m))Rinecs Rljnnarr
Similarly to Eq. (5), we define and plot Ttrain mm) and heea (w, m) in Figure 7, using the full dataset N = 60000. Comparing Figures 7a and 7b reveals two things: (1) The training and test losses behave similarly: (2) Both training and test losses depend very weakly on m. This implies that the raw image representation is already quite close to being optimal, so decent test accuracy can be obtained even without learning optimal representations. As a result, grokking does not occur (Figure 7c).
Comparing Figure 6 and 7, we see that the (strong) dependence of test performance on the representa- tion the key to grokking: the dependence on representation is strong for algorithmic datasets, so grokking happens. By contrast, the dependence is weak for MNIST, so grokking does not happen.
'For simplicity, we assume f/ = ?. here, hut the analysis can apply any (дz4 mp)-
8

---

Published as a conference paper at ICLR 2023
loss
repre representation massinass m
train
representationn messiness
38
test
Ze
31
44
24
colo
12
0.00
000
18
relative weight W/Wo
relative weight w/we
Path
(c)
(a)
(b)
Figure 7: MNIST landscapes as functions of representation messiness m and weight norm w: (a) training loss, and (b) test loss. Training and test losses do not have significant mismatch, and neither of them depend on representation strongly, which is in stark contrast algorithmic datasets (Figure 6). (c) an illustrative path A + B - C does not manifest grokking.
RELATION TO RELATED WORKS
Grokking first observed for algorithmic datasets by Power et al. (2022). Several attempts have been made to understand grokking: (a) Liu et al. (2022) attributes grokking to the slow formation of good representations. (b) Shah (2021) suggests that generalizable solutions achieve lower loss than overfitting solutions, providing a training signal encouraging generalization. (c) Nanda et al. (2023)
suggests grokking is a phase change due limited data and regularization. (d) Barak et al. (2022)
suggests that generalization is due not to random search, but to hidden progress of SGD to gradually amplify a Fourier gap. (e) Thilak et al. (2022) links grokking to the "Slingshot mechanism' specific to adaptive optimizers. (f) Millidge (2022) describes training as a random walk over parameters. Our conclusion supports (c)), but does not necessarily negate (f)
Double descent is the phenomenon that performance first gets worse and then gets better as increase the model size, data size, training epochs or regularization (Nakkiran et al., 2021; Yilmaz and Hcckel, 2022; Nakkiran, 2019). The typical "U" shape of test loss in this paper does not conflict with double descent, because we are plotting the weight norm instead the number of model parameters (Ng and Ma, 2022). However, the "U"-shape should better be considered empirically common rather than provably universal. In fact, the interaction between properties of data and inductive biases of learning algorithms be more complicated than double descent (Chen et al., 2021; d' Ascoli et al., 2020).
Initialization From the optimization perspective, initializations usually based on the "edge of chaos" idea such that variance of features and gradients should be preserved in the forward and backward pass (Glorot and Bengio, 2010: He et al., 2015; Bahri et al., 2020; Yang and Schoenholz. 2017; Jing et al., 2017), or based on analyzing Jacobians and/or Hessians (Skorski et al., 2020), From the generalization perspective, was shown that large initializations overfit data easily but result in poor generalization (Xu et al., 2019; Zhang et al., 2020), which agrees with our LU mechanism.
Weight decay regularization is standard trick in machine learning and has various effects on optimization and generalization (Zhang et al., 2018; Van Laarhoven, 2017). In particular, Lewkowyzz and Gur-Ari (2020) observes that takes xx 1/A training steps to reach maximum test performance. This is strikingly similar to the grokking time x 1/A we derived from the LU mechanism.
7 CONCLUSIONS
This study elucidates the grokking phenomenon from the perspective of loss landscapes. Our conclusions are: (i) grokking originates from the mismatch between training and test losses at high model weight norm ("LU" mechanism). (ii) grokking can happen in various models for wide range of datasets, although the grokking signature usually most dramatic for algorithmic datasets. (iii) The severity of grokking depends on how much the task relies learning representations. This work not only reveals the mechanism of grokking, but also shows that reduced landscape analysis is a useful tool for characterizing data-model interaction and representation learning.

---

Published as a conference paper at ICLR 2023
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.
Ziming Liu, Ouail Kitouni, Niklas Nolte, Fric J Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. arXiv preprint arXiv:2205.1034 2022.
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 vrXiv:2206.048// 2022.
Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep leamning Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08799, 2022.
Stanislav Fort and Adam Scherlis. The goldilocks zone: Towards better understanding of neural network loss landscapes. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pages 3574-3581, 2019.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanics: Theory and Experiment, 2021 2021.
Andrew Ng and Tengyu Ma.
Cs229 lecture notes. ppring2002/miin_notesppu page 115, 2022.
http://c2222statfooreeeddlleeee
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint 161 1.0/23, 2016.
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. Advances in neural information processing systems, 30, 2017.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-14) 2012.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http:
//www. www.accom www.eb..wb..... org/ www.aannhhoooyyooooo p11.1l111y111111.1..11
Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014.
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2300.05217,20000 2023.
Rohin
Shah.
Alignment
Newsletter
#159.
https://
www. a iiiinnmentoruus org wttp:///st///////masseeeeee////e an-159-mml 1159=bbuidinn--geeeee LEARNING 2021. 2021.
Beren
Millidge.
Grokking 'grokking'.
https://beren. to/
2022- 0122011-0 11-000--1111
2022.
Fatih Furkan Yilmaz and Reinhard Heckel. Regularization-wis double descent: Why it occurs and how to eliminate it. In 2022 IEEE International Symposium on Information Theory (ISIT), pages 421-431. IEFF, 2022.
10

---

Published as a conference paper at ICLR 2023
Preetum Nakkiran. More data can hurt for linear regression: Sample-wise double descent. arXiv preprint arXiv:1912.07242, 2019.
Lin Chen, Yifei Min, Mikhail Belkin, and Amin Karbasi. Multiple descent: Design your own generalization curve. Advances in Neural Information Processing Systems, 14:8898-8912, 2021.
Stéphane d' Ascoli, Levent Sagun, and Giulio Biroli. Triple descent and the two kinds of overfitting: Where & why do they appear? Advances in Neural Information Processing Systems, 33:3058-3066 2020.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 249-25422 JMLR Workshop and Conference Proceedings, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S Schoenholz, Jascha Sohl-Dickstein and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, I1(1,, 2020.
Li Jing, Yichen Shen, Tena Dubcek, John Peurifoy, Scott Skirlo, Yann LeCun, Max Tegmark, and Marin Soljačíć. Tunable efficient unitary neural networks (eunn) and their application to rnns. In International Conference Machine Learning, pages 1733-1741. PMLR, 2017.
Maciej Skorski, Alessandro Temperoni, and Martin Theobald. Revisiting initialization of neural networks. arXiv preprint arXiv:2004.09550 2020.
Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. International Conference on Neural Information Processing, pages 264-274. Springer, 2019.
Yaoyu Zhang, Zhi- ZhiiQQn John Xu, Tao Luo, and Zheng Ma. A type of generalization error induced by initialization in deep neural networks. In Mathematical and Scientific Machine Learning, pages 144-164. PMLR, 2020.
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint (rXiv:/8/0) 12281, 2018.
Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint 1706000000 2017.
Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with 1_2 regulariza- tion. Advances in Neural Information Processing Systems, 33:4790-4499, 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 11X1v;1412.6980, 2014.

---

Published as a conference paper at ICLR 2023
Appendix
A EXPERIMENT DETAILS
Sentiment analysis of text IMDb (Maas et al., 2011) includes 50k movie reviews to be classified as being positive negative. To pre pre-proccss the data, extract the 1000 most frequent words and tokenize each review into an array of token indices. Less frequent words are ignored, and each review array is padded length 500. We adopt the LSTTM model (Hochreiter and Schmidhuber, 1997) to perform the classification, with two layers, embedding dimension 64, and hidden dimension 128. We the Adam optimizer (Kingma and Ba, 2014) with learning rate 0.001 to minimize the binary cross entropy loss. We hold back 25% of the dataset for testing.
Molecules QM9 is database for small molecules and their properties. We graph convolutional neural network (GCNN) to predict the isotropic polarizability. The GCNN contains 2 convolutional layers with ReLU activation. followed by a linear layer. We use the Adam optimizer with learning rate 0.001 to minimize the MSE loss. We split the dataset into 50/50 train/test.
MNIST We train width-200 depth-3 ReLU MLPs the MNIST dataset with MSE loss. We use the AdamW W optimizer with learning rate of 0.001 and batch size of 200.
B REDUCED LOSS FOR MODULAR ADDITION WITH TRANSFORMERS
In Figure 8 show reduced loss landscape plots for transformers trained on modular addition. We use the setup Nanda et al. (2023) and train a 1-layer transformer on modular addition (p = 113) with dmodel 128, 4 attention heads, and dmlp 512 with ReLU activations. We train with learning rate of 0.001 while constraining model weight norm, for a variety a and a variety train set fractions. The LU shape holds for E [0.1, 4] (some optimization issue may be responsible for the rise in train loss for a > 4). We see the critical train set size is approximately 0.25, line with earlier studies on grokking.
training
test
0.1
0.1
0.9
set fraction
0.7
0.3
0.8
0.7
0.4
0.5
0.4
0.5
0.5
0.5
0.5
0.5
0.6
0.6
0.4
0.7
g
0.3
0.8
0.9
0.9
10-1
10a
10
10-1
10"
10
relative weight w/Ws
relative weight
Figure 8: Reduced loss landscapes for transformers trained on modular addition, the original setting where grokking was observed.
C WEIGHT NORM EVOLUTION OVER TIME ON ALGORITHMIC TASKS
Evolution of weight norm As mentioned in Section 5, the dynamics of model weight norm over the course of training, algorithmic tasks, support the LU mechanism picture of grokking. Figure 9a, shows how model norm changes over time and we see that there an initial increase in weight norm, which peaks during overfitting, but then drops during the period of generalization to be lower than the initialization norm. For this experiment, we again used the setup of (Nanda et al., 2023). We train with AdamW with learning rate of 0.001 and weight decay 7 1.
Constraining a small weight norm eliminates grokking As shown in Figure 9b, reducing the initialization scale (a 0.8) and constraining optimization hold model weight norm constant
12

---

Published as a conference paper at ICLR 2023
training brings train accuracy and test aceuracy learning curves together, almost eliminating grokking. We would like to investigate in future works whether this training trick can be helpful for more standard machine learning tasks.
IL Tranaformer Modular Addition (1-113) Unconstrained optiiization in Intliiationn
Transformer oe Nodular Addition (p-113) Constrained Norm a 3.8
Train
test
50
10
107
10a
18
10 10
10*
Optimbse on Steps
Oprimiaation Steps
(a)
(b)
Figure 9: Training IL transformer on modular addition (p 113). (a) Weight norm, train accuracy, and test accuracy over time, initialized and trained normally. Weight norm first increases, and is highest during the period of overfitting, but then drops to become lower than initial weight norm when the model generalizes. (b) Constrained optimization at constant weight norm (a 0.8) largely climinates grokking, with test and train accuracy improving almost concurrently.
D TIME TO GENERALIZE VERSUS WEIGHT DECAY
In our discussion of the "LU mechanism" as an explanation for grokking in Section 2, we predicted that the training time required for a model to generalize should be x 7 where y the weight decay. To test this, perform grid search over weight decays 7 and plot the number of training steps required for models to reach a specified level of test accuracy Figure 10a-10b. We also show full training curves for these runs in Figure 10c-10d. We perform experiments in two setups:
(a) Transformer on modular addition: We use the replication of grokking from Nanda et al. (2023) and train a 1-layer transformer on modular addition (p 113 and train set fraction of 0.3) where dmnseel 128, with 4 attention heads, dmpp 512, ReLU activations, and an AdamW learning rate of 0.001. From Figure 10a, we find that 7 holds across roughly two orders of magnitude of t and There is some seed dependence on the generalization time (some seeds consistently require longer to generalize), but for each seed (corresponding to particular model initialization) the relation t x 7- appears to fit the data well.
(b) ReLU MLP on MNIST: We train ReLU MLPs on MNIST as described in Appendix A. We 9.0 and train a reduced training set of 1000 samples to delay generalization / induce grokking. From Figure 10b, we find that for roughly between 0. and 1.0 the relation t x - holds. Very high values of weight decay seem mess with optimization. On the other hand, with very low weight decay the model generalizes faster than naively expected, perhaps due to implicit regularization.
E SECTION 5.1 SETUP
Architecture Similar to Liu et al. (2022), the decoder architecture an MLP with hard coded addition. Each input symbol i is encoded to scalar E, Each output symbol &o is represented by 30D random vector Y,. We consider addition with base p. input 0 < i.j < p and output 0 < k = i + < 2(p - 1). We denote representation as R = EE, E₁ ... Ep-1 1) The MLP has two hidden layers, with neurons 1-200-200-30 in each layer and ReLUU activations. Given training sample (E,, E,) + YA where i + j = k. the prediction of the MLP decoder is
Y Decw/ E, + Ez).
and the loss function being the mean squared error (MSE) between Y, and Y,. and w being the decoder weight. Although the common setup of grokking is to make both the representation R and
13

---

Published as a conference paper at ICLR 2023
Transformer Mod Addition 1.0
ReLU MLP MNIST 9.0
10a
Optimization Steps 100% Test Accuracy
Optimization Steps to 70% Test Accuracy
104
104
seed
seed D
seed
seed
seed
10a
Seed
y
10a
tay
10
10
10
10
10
100
10a
Weight Decay
Weight Decay Y
(
(b)
Transformer on Addition a= 1.0
ReLU MLP on MNIST a=
101
1.0
1.0
0.8
Weight Decay Y
0.8
10"
Weight Decay
Accuracy
0.6
100
Accuracy
0.6
10-1
0.4
0.4
0.2
10-2
tast train
0.2
test train
0.0
0.0
10-3
103
10a
10a
10
Optimization Steps
Optimization Steps
(c)
(d)
Figure 10: Time to generalize as function of weight decay: we investigate to what extent the relation t x y holds, where i is number of training steps needed for the model to generalize and is the AdamW weight decay. When a lower weight decay is used, models spend longer the period of overfitting before eventually generalizing. We show the generalization time function of y in (a)-(b) and full training curves for these runs (c)-(d).
decoder weight w
Representation R
Trainability
w/|2 direction w/w
messiness m
Other
Landscape analysis
No, w
Yes
No, m
No, 0
Reduced trajectory
Yes
No, w w'Ww, m)
Yes
No. 0
Full trajectory
Yes
Yes
Yes Yes
Table 1: Threes setups used in this paper, with different set of parameters trainable.
the decoder trainable, will freeze part of them for easier analysis. This is where it could be bit confusing, so we explicitly distinguish three setups: landscape analysis, reduced trajectory analysis and full trajectory analysis. Each setup have different subset of trainable parameters, shown in Table
Landscape analysis Both the representation R. and weight direction w is trainable. The representation R fixed according to Eq. (4), which is dependent on m, the representation messiness, The decoder has fixed weight norm w, but the weight direction w is trainable. For each fixed (w, m), we minimize training loss over w to get
fixed. Only the weight
w ('ww, m) = argmin Exraiu (w, m, w), wr
and define reduced training and test loss, as in Eq. (5). The minimization is implemented by the Adam optimizer with learning rate 10-3 for 10 steps. Although (w, m) trainable, we repeat the above minimization independently for different (w, m). In Figure (a), the background heatmaps belong to landscape analysis.
14

---

Published as a conference paper at ICLR 2023
Reduced trajectory analysis is "thought experiment" based on landscape analysis. Since full trajectory analysis can intractable due to too high dimensions, we try to reduce the trajectory anaysis to 2D, by making two assumptions about the real dynamics: (1) Scale separation: the dynamies of w much faster than the dynamics along w and along m, such that ((t) = w
w" (w(t). m())) is valid every moment during training. (2) Representation evolution is linear, i.e., interpolating between initial random Gaussian and final linear representation. With these two assumptions, the training dynamics is effectively reduced 2D, depending only on (w, m), obeying Eq. (6). In Figure (a)(b)(c) the path A + E belongs to reduced trajectory analysis.
Admittedly the reduced trajectory may deviate from the full trajectory since the assumptions may not be met, but it can shed light on the full trajectory: the weight norm first increases and then increases, and the decrease weight norm highly correlated with generalization (please see Appendix C and Figure 9.
F MNIST EXPERIMENTS WITH CROSS ENTROPY LOSS
To respond to reviewer's concern that our use of the MSE loss is the "secret" to get grokking on MNIST (Figure 3), we reran our experiments with the cross entropy (CE) loss. The results are qualitatively similar, with some quantitative differences.
000
Yocls
100 dize 8300
io
AM 30
644
00-
10
101
10 LE
relative weight norm WWW
relative weight norm a WWWr
weight norm
Figure 11: MNIST with the cross entropy loss (as opposed to the MSE loss used in Figure 3). (a) reduced training error, (b) reduced test error. (c) "LU" still holds for the cross entropy loss, but the effect is milder than the MSE loss. In particular, the "Goldilocks zone" (the weight range where generalization happens) is broader.
Landscape analysis
Comparing Figure 3 (MSE) and Figure 11 (CE), we notice the they are qualitatively similar: (1) for small datasets, the reduced training error and test error resemble an "I." and "U" against the weight norm, respectively; (2) for large datasets, the "U" becomes more like I,.", i.e., the mismatch between the reduced training and test error is small. However, quantitative difference exist: CE produces broader "Goldilocks zone" (the weight range where generalization happens) than MSE. This implies that to induce grokking with CE, we need increase the weight norm to a larger value (say a 100).
Training dynamics
We are able to observe delayed generalization during trianing on MNIST with cross entropy loss, but doing requires higher than necessary when using MSE loss, predicted by the reduced loss landscapes in Figure 11. Figure 12 shows training trajectories from a 3-layer ReLU MLP on MNIST trained with cross entropy loss with 100 and D 200. We see that test accuracy rises to 30-40% early in training, then plateaus for an extended period, before increasing to 275% while train accuracy remains 100%. While the dynamics first plateaus at better-than-random accuracy, we think it is still fair to classify these dynamics as "grokking" due to the improvement in generalization late in training after a plateau.
not as clean with MSE loss, since test accuracy
15

---

Published as a conference paper at ICLR 2023
MNIST with cross-entropy (linear scale) a=100, D=200
MNIST with cross-entropy (log scale) | a=100, D=
1.00
1.00
0.75
0.75
Accuracy
Accuracy
0.50
0.50
0.25
train test
0.25
train
test
25000
50000
75000
100000
10a
104
10a
Optimization Steps
Optimization Steps
Figure 12: Training curves using cross entropy loss on MNIST. We are still able to observe delayed generalization on MNIST using cross entropy loss, though test accuracy first plateaus at higher than random-gues accuracy.
16


Reviews:
[{"review_id": "hQW4UMkhmK", "reviewer": "Reviewer_ENpU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper analyses the phenomenon of grokking (discovering solutions that generalise well after the model has overfit the data) from the perspective of loss landscapes, specifically the disparity between the training and the generalisation landscapes. The authors show that grokking can be induced on standard ML benchmarks such as MNIST classification by (1) increasing the weight initialisation range, (2) drastically reducing the number of data points in the training dataset. Further, the authors conduct experiments to show that grokking is linked to the correlation between the discovery of a good representation and generalisation. For loss landscapes analysis, low-dimensional projections are used, where the weights are represented as a weight norm. The reduced 1D loss landscape visualisation seems to be a powerful tool.", "summary_of_the_review": "Overall, this is a good paper with brilliant loss landscape visualisations. Please refer to the least of strengths and weaknesses for the justification of my recommendation.", "strength_and_weaknesses": "Strengths:\n\n1) The loss landscape perspective is refreshing and very tangible, I think this is the right way to think about the problem of grokking.\n2) Reduced 1D loss landscapes are a powerful visualisation tool, and may find further uses in the NN analysis and understanding.\n3) Empirical results are convincing.\n\nWeaknesses:\n1) The authors refer to “standard initialisation”, and then explain that it stands for default initialisation in PyTorch. While this can be easily looked up at the time of writing, I am not in favour of such loose definitions. What if the next version of PyTorch picks a different default? Surely that should not lead to non-reproducible results. As such, I would like to request the authors to rather define the initialisation scheme explicitly.\n2) Even though the authors have managed to induce “grokking” for MNIST and other tasks, the exercise seemed very artificial: yes, we can induce grokking by chopping off most of the dataset and raising the initialisation range. However, grokking does not seem to be something that we necessarily want to induce: looking at the results obtained, avoiding overfitting in the first place seems to still be the better option. Is the significance of this phenomenon perhaps blown out of proportion? I wish the relevance of grokking to standard ML tasks in light of the new experiments was discussed a bit more thoroughly.\n3) The paper has minor language mistakes:\nPage 5: “For large data size say the full dataset” -> e.g. the full dataset\nPage 6: “by constrast” -> contrast \nPage 9: “dramaticness” -> not a valid word, rather use “severity”", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TYmNC77FoZ", "reviewer": "Reviewer_itfw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes the LU mechanism for the phenomenon of Grokking observed by Power et. al. 2020 - which states that the train and test curves follow an L and U shaped curve with weight norms respectively. They verify this observation in a student teacher setup, and show that it can arise in non-algorithmic datasets if initialized in a certain weight regime for appropriate sample size. Further, they try to differentiate the algorihtmic and non-algorithmic datasets in terms of their dependence on representations. ", "summary_of_the_review": "The authors state three conclusions in their paper - the LU mechanism, the occurrence of grokking for wide range of datasets under certain conditions and the dependence of grokking on learning representations. The first two are well-supported by their experiments as the setup is general enough that the conclusions are believable. However, the third conclusion is made in an artificially reduced setup. While this setup supports the hypothesis, it is unclear to what extent this is related to the general phenomenon of grokking where representations are learnt and not fixed. Nevertheless, the contribution of the first two conclusions and weak evidence for the third conclusion seems enough for acceptance to the conference. I would recommend the authors include some discussion for their choice of setup in Section 5.\n\n", "strength_and_weaknesses": "Strengths:\n1. The LU mechanism hypothesis is clearly stated, and the proposal to construct the landscape with fixed weight norm solutions is a nice idea to verify it (barring the possible optimization concerns for the constrained optimization problem). \n2. The experiments on MNIST, IMDB support the LU mechanism hypothesis - I particularly liked Figure 3 as it illustrates the idea clearly.   \n3. The authors are able to demonstrate grokking for non-algorithmic datasets by changing the sample size and weight initialization - this is an important finding.  \n4. The proposed difference between algorithmic and non-algorithmic datasets is plausible - in particular that the weight norm increases and then drops as in Figure 9. The representation messiness argument also supports this conclusion, but I found this to be a little weak (see comments below). \n\nWeaknesses:  \n1. Representation messiness:   \n- The experimental setup is hard to understand in Section 5.1 - I would suggest the authors include an appendix section describing the exact architecture, and the trainable/frozen components of this architecture. Does R_random depend on the input at all? If not, how is this a good proxy for representation learning?  \n- This part of the paper makes a reduction to a different architecture (from Liu et. al. 2022) and assumes that the representations are a convex combination of linearly separable and random gaussian representations. It is not clear to me why this is general and if the conclusions transfer to a setting where the representations are learnable (as is the case in grokking generally). \n2. Minor complaints:  \n- The authors claim that changing the x-axis to weight norm removes double descent - this is only true for some limited cases and does not hold generally for neural networks as is also stated in the reference cited by authors. I would recommend that the authors add this clarification.  \n- Figure 6: Please state which experimental setup these loss landscapes have been plotted for. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "DVSEq4b-Sj", "reviewer": "Reviewer_wgE4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes an explanation for when grokking happens---that it occurs when the initialization norm is too large so the model takes longer to get to the correct norm magnitude—and illustrates their thesis with results on synthetic and natural data. Defining grokking as a delay in generalization until after training set interpolation, they propose that grokking is caused by a mismatch between training and test loss curves, ie, overfitting, and described this pattern as the “LU mechanism”. They elicit grokking phenomena on natural data (IMDB and MNIST), at least if we define grokking as a delayed generalization curve. ", "summary_of_the_review": "I found many of the analyses here intriguing, and ultimately I would very much like to see this paper published. I was particularly impressed by their ability to induce grokking in a realistic setting. However, they seem to be making the claim that grokking is in general caused by initializing with a large norm. The experiments do not fully support that this is the case in the main setting in which grokking has been observed, which is the algorithmic setting. I would welcome more experiments around a wider set of hyperparameters, to demonstrate that delayed generalization does not result from other settings. Without these experiments or an added theoretical analysis, I don't feel confident recommending acceptance, but I'm ready to change my mind in response to a small number of additional experiments: evidence that the effects observed are particular to the norm magnitude and illustrations of the messiness plots but with the initialization weights instead of random weights, so that we can see something closer to the actual trajectory of training. Furthermore, I think that the paper should be somewhat modified in its definitions and reasoning, as the LU mechanism appears to fall almost trivially from their definition of grokking, rather than being an attributed cause.\n\n**After discussion:**\n\nI'm raising my score from 6 to 8 to reflect expanded experiments that lend some additional credence to the link between behavior in the algorithmic and natural settings, and the removal and finetuning of some unsupported speculation and imprecise connections. ", "strength_and_weaknesses": "Strengths:\n- This paper is well written and presented, a pleasure to read.\n- The experiments make the claims clear.\n- They successfully induce delayed generalization in natural datasets, suggesting that their claims might apply in realistic settings. These results happen outside of the teacher students set in, so they are more convincing than the algorithmic experiments.\n- The observation that larger dataset sizes expands the goldilocks zone is a very tidy way of thinking about generalization.\n- I found myself fairly convinced by the experiments that initialization magnitude is at least one factor that can induce grokking.\n- It’s good that they discuss the impact of weight decay repeatedly, as it is clearly an important factor.\n\nWeaknesses:\n- I do not generally consider it necessary to have a theoretical component of every empirical study, and many empirical studies make strong and intriguing claims that layout groundwork for a principled understanding. However, I think that the particular claims made in this paper would benefit enormously from a theoretical analysis. My issue is that they make out the initialization norm to be of particular significance, when other hyperparameters might (in fact, empirically do) have an equivalent “goldilocks zone”. Additional experiments about other hyperparameters might help, but I think that the insights that we can glean from purely empirical study of this phenomenon are limited. \n- I feel that the connection between overfitting and grokking is fairly obvious, as their own definition of grokking makes clear. Given the way that they implicitly define grokking, I don’t believe that they can describe the LU mechanism as a “cause” rather than as a rephrasing of the same definition.\n- These settings used to illicit grokking are somewhat artificial: the initial norm, the constrained search space, the weight decay requirement. There is no satisfactory explanation given for why weight decay is so crucial; this is where a theoretical analysis might have been helpful.\n- Teacher student setups can fundamentally change training dynamics, so I’m not confident that the results from those experiments generalize to other settings. If the claims made from these experiments do not apply outside of a teacher student setting, I’m not sure I buy them. Especially because the particular representations learned by the teacher could easily require its own particular magnitude that is not necessary for all representations.\n- It is not clear to me what differentiates the LU mechanism from the existing concept of overfitting. Although they described their setup as different from ‘related phenomena’  (Schoenholz et al., 2016; Yang and Schoenholz, 2017; Nakkiran et al., 2021), it’s not clear to me how these citations are not directly describing the LU mechanism.\n- The effect that this paper attributes to the importance of representation learning (section 5) might be more precisely described as the effect of the distribution shift at inference time. \n- “We treat the converging point after training as the global minimum on the spherical surface” — why? Is this actually a convex a surface somehow?\n- It appears to me that the interpolations are between the final learned model and a random gaussian. It seems to me these experiments and plots would be more informative if the random model was the initialization weights, so that we could see a planar view of the path from initialization and not just any random direction. I personally would find that more convincing it as an explanation of the training process, compared with a random direction.\n- On importance of representation learning: it seems to me that this might best be described more precisely as the significance of distribution shift between train and test time. Otherwise, I would expect you to formalize what it would mean to have representation learning be important.\n- The discussion section: “(ii) Pre-training avoids learning representations from scratch, hence helps reduce possible grokking.” But why wouldn’t we see grokking in the pretraining stage itself?\n- Is it definitely the case that grokking is not observed in language models if you use larger initializations? Otherwise, you should probably remove the discussion section on language models entirely or present evidence that grokking does not happen at any stage. I don't think there's anything backed up in this section, and it doesn't really relate to the rest of the results. \n\nMinor/references:\n- You refer to default initialization settings in pytorch, but I think you need to give details about what those defaults are, as they might change between versions.\n- Footnote 2: I’m not sure this footnote should be there. It is well known that increased capacity allows tighter interpolation of the training set, but justifying this statement by presenting an artificial scenario in which a larger model is equivalent to the smaller model is not a reflection of real practices, where we would not initialize zero-vectors when expanding an architecture.\n- There is an existing literature on distinctions between multi epoch vs ideal infinite data sampling in practice, which might be worth digging into and adding to your citations, eg Nakkiran (2019) https://arxiv.org/abs/1912.07242\n- “We run experiments with two initializations α = 0.5 (small) and α = 2.0 (large)” — Above, it said that alpha was the constant weight norm throughout training, but here it says that it is only the initialization. I can see now that this is supposed to just be the weight constraint, but it’s a little confusing in the wording.\n- A number of these results, especially in the section on LU, seem to boil down to demonstrating that overfitting does exist. I’m not sure that this is necessary to make such a focus.\n- Formula 4 should have a citation to existing work on sharpness or linear interpolation.\n- I’m a little confused about the meaning of v′ = vcosθ, given that θ was previously defined as a thresholding value for the regression model, and I don’t see how that definition connects to the use the variable in section 5.1. \n- This paper uses natbib. When you make an inline citation, e.g., “Foo et al. (2009) claim…” you should use \\citet{} instead of \\cite{} or  \\citep{} . An example of this issue is the entire first paragraph of section 6.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6_6vdMDUFz", "reviewer": "Reviewer_ceW8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the properties of grokking from the loss landscapes. The phenomenon of \"LU\" mechanism of the training and test loss is adopted to understand grokking from its dependence on data size, weight decay, and representations. In particular, the experiment tailor-designed to reveal the connection between grokking and deep representation learning is interesting and impressive. ", "summary_of_the_review": "Overall, it is a decent paper.  it proposes to interpret the grokking phenomenon from the perspective of loss landscapes. And sufficient empirical analysis is provided to support the claims of the paper. \n", "strength_and_weaknesses": "Strength:\n\n(1) It is interesting and novel to understand grokking from the lens of neural loss landscapes. More importantly, the connection between grokking and data size, weight decay, and representation learning are demonstrated with sufficient empirical analysis. \n\n(2) It conducts sufficient and intuitive experiments to reveal the dependence of grokking on representation learning, which is helpful for the community to understand grokking in the learning process with complex datasets. Meanwhile, it contributes a useful new tool for characterizing data-model interaction and representation learning. \n\nWeaknesses:\n\n(1) Some terms should be explained in the paper for self-consistency, such as algorithmic datasets.\n\n(2) The color bar of Fig.6 (b) is in [0.4, 2.0], which is different from other figures. And as claimed in the paper, the error = 1 - acc. This is inconsistent. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%