Towards Understanding Grokking: An Effective Theory of Representation Learning
Authors: Ziming Liu
Ouail Kitouni
Niklas Nolte
Eric J Michaud
Max Tegmark
Mike Williams
Venue: NeurIPS 2022 Accept (2022)
Keywords: grokking
representation learning
physics
effective theory

Abstract: We aim to understand grokking, a phenomenon where models generalize long after overfitting their training set. We present both a microscopic analysis anchored by an effective theory and a macroscopic analysis of phase diagrams describing learning performance across hyperparameters. We find that generalization originates from structured representations, whose training dynamics and dependence on training set size can be predicted by our effective theory (in a toy setting). We observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion. We find representation learning to occur only in a "Goldilocks zone" (including comprehension and grokking) between memorization and confusion. Compared to the comprehension phase, the grokking phase stays closer to the memorization phase, leading to delayed generalization. The Goldilocks phase is reminiscent of "intelligence from starvation" in Darwinian evolution, where resource limitations drive discovery of more efficient solutions. This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspired tools, e.g., effective theories and phase diagrams, for understanding deep learning.

Towards Understanding Grokking: An Effective Theory of Representation Learning
Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, Mike Williams Department of Physics, Institute for Al and Fundamental Interactions, MIT femliu,kitoun wwt.ttttiinnottteenooee mnolte, ericjm, tegmark, mwil1w@mmm edu
Abstract
We aim to understand grokking. a phenomenon where models generalize long after overfitting their training set. We present both microscopic analysis anchored by an effective theory and macroscopic analysis of phase diagrams describing learning performance across hyperparameters. We find that generalization originates from structured representations whose training dynamics and dependence on training set size can be predicted by our effective theory a toy setting. observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion. We find representation learning to occur only a "Goldilocks zone" (including comprehension and grokking) between memorization and confusion. We find on transformers the grokking phase stays closer to the memorization phase (compared to the comprchension phase), leading to delayed generalization. The Goldilocks phase is reminiscent of "intelligence from starvation" in Darwinian evolution, where resource limitations drive discovery of more efficient solutions. This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspires tools, e.g., effective theories and phase diagrams, for understanding deep learning.
Introduction
Perhaps the central challenge of scientific understanding of deep learning lies in accounting for neu- ral network generalization. Power et al. [1] recently added a new puzzle to the task of understanding generalization with their discovery of grokking. Grokking refers to the surprising phenomenon of delayed generalization where neural networks, on certain learning problems, generalize long after overfitting their training set. It is rare albeit striking phenomenon that violates common machine learning intuitions, raising three key puzzles:
Q1 The origin of generalization: When trained on the algorithmic datasets where grokking occurs, how do models generalize at all?
Q2 The critical training size: the training set size decreases toward a critical point?
Why does the training time needed "grok" (generalize) diverge as
Q3 Delayed generalization: Under what conditions does delayed generalization occur?
We provide evidence that representation learning is central to answering each of these questions. Our answers can summarized as follows:
A1 Generalization can be attributed to learning a good representation of the input embeddings, i.e., a representation that has the appropriate structure for the task and which can predicted from the theory in Section 3. See Figures and 2.
A2 The critical training set size corresponds to the least amount of training data that can determine such representation (which, in some cases, is unique up to linear transformations..
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Initislization (0 iterations) traiz acc: 0.0 vol. noc: 0.0
Overftting (1000 iterations] train noc" 1.0 val ace: 0.1
Representacion Learning (20000 iterations)
tenim ace: 1.0
noc: 1.0
ST
50
02
30g9
2037
25
283
22
383
18
12
20
68
11
38
24
27
05
30
06
24
45
31
41"
13
11
51
1472
47
%
15 25
21
04
17
18
4Ags
48
47
400
10
09
02
58
65
8g
01
02
05
Figure 1: Visualization the first two principal components of the learned input embeddings at different training stages of a transformer learning modular addition. We observe that generalization coincides with the emergence of structure in the embeddings. See Section 4.2 for the training details.
A3 Grokking is phase between "comprehension" and "memorization" phases and it can be remedied with proper hyperparmeter tuning, as illustrated by the phase diagrams in Figure
This paper is organized as follows: In Section 2, we introduce the problem setting and build a simplified toy model. In Section 3, we will use effective theory approach, a useful tool from theoretical physics, to shed some light on questions Q1 and Q2 and show the relationship between generalization and the learning of structured representations. In Section 4, we explain Q3 by displaying phase diagrams from a grid search of hyperparameters and show how we can "de-delay" generalization by following intuition developed from the phase diagram. We discuss related work in Section 5, followed by conclusions in Section 6.
2 Problem Setting
Power al. [1] observe grokking on a less common task learning "algorithmic" binary operations. Given some binary operation o, network is tasked with learning the map (a, b) where a They use decoder-only transformer to predict the second to last token in a tokenized equation the form "<lhs> <op> <rhs <eq> <result> <eos>". Each token is represented as 256-dimensional embedding vector. The embeddings are learnable and initialized randomly. After the transformer, a final linear layer maps the output to class logits for each token.
Toy Model We primarily study grokking in a simpler toy model, which still retains the key behaviors from the setup of [1]. Although [1] treated this as classification task, we study both regression (mean-squ -quan-s error) and classification (cross-entrow erross-entropy..... The basic setup as follows: our model takes as input the symbols a, and maps them to trainable embedding vectors E., E, Ra. It then sums E,,; E,, and sends the resulting vector through "decoder" ML.. The target output vector, denoted Y. dh is fixed random vector (regression task) or one-hot vector (classification task). Our model architecture can therefore be compactly described as (a, b) + Dec(E, + Eg). where the embeddings E. and the decoder are trainable. Despite its simplicity, this toy model can generalize to all abelian groups (discussed Appendix B). sections 3-4.1, we consider only the binary operation of addition. We consider modular addition in Section 4.2 to generalize some of our results to a transformer architecture and study general non-abelian operations in Appendix H.
Dataset In our toy setting, we are concerned with learning the addition operation. A data sample corresponding + j is denoted as (i.j) for simplicity. If i.j E {O..... p(p + 1)/2 different samples since we consider i + j and j + to be the same sample. A dataset D) a set of non-repeating data samples. We denote the full dataset as Do and split into training dataset D and validation dataset D', i.e., D'D' = Do. DND' D = 0. We define training data fraction = Do where denotes the cardinality the set.
1). there are in total
Project code can be found at: https: //gg:::///ttbb com/ hamm/jjmia httee/jmccccuude jaiihhau/grookkkiggg squared

---

RQJ 0.0 Accuracy Leatar 1.0, validation: 0.1
BQL Accuracy Erain 1.0, validation: 0.9
20
08
20
18
06
09
02
12
10
10
09
08
04
04
10
10
(a) Memorization toy addition
(b) Generalization toy addition
Accuracy
validation:
Aceuracy
validation:
10
02
04
10
6
10
09
06
2
08
(c) Memorization in toy modular addition
(d) Generalization in toy modular addition
Figure 2: Visualization the learned set of embeddings (p 11) and the decoder function associated with it for the case of 2D embeddings. Axes refer each dimension of the learned embeddings. The decoder is evaluated on grid of points in embedding-spee the highest probability class. For visualization purposes, the decoder is trained on inputs of the form (E, + E2)/2. One can read off the output of the decoder when fed the operation i o from this figure simply by taking the midpoint between the respective embeddings of i and j.
and the color at each point represents
3 Why Generalization Occurs: Representations and Dynamics
We can see that generalization appears to be linked the emergence of highly- highly-structured embeddings in Figure 2. In particular, Figure 2 (a, b) shows parallelograms in toy addition, and (c, d) shows circle in toy modular addition. We now restrict ourselves to the toy addition setup and formalize notion of representation quality and show that it predicts the model's performance. We then develop physies-inspired effective theory of learning which can accurately predict the eritical training set size and training trajectories of representations. The concept of an effective theory in physics is similar to model reduction in computational methods in that it aims to describe complex phenomena with simple yet intuitive pictures. In our effective theory, learning not as gradient descent of the true task loss but rather a simpler effective loss function leff which depends only the representations in embedding space and not the decoder.
will model the dynamics of representation
3.1 Representation quality predicts generalization for the toy model
rigorous definition for structure in the learned representation is necessary. We propose the following definition,
Definition 1. (i, m. n) is 6 ô-parallelogram in the representation R |Eo.... EEp-1
EE + Ej) (E., + E,)l 8.
In the following derivations, can take 8, which is small threshold to tolerate numerical errors, to be zero.
Proposition 1. When the training loss is zero, any parallelogram (i.j, m, n) in representation R satisfies i = m + n.
Proof. Suppose that this not the case, i.e., suppose E, + E, E,, + E,, but i Yi+; Dec(E, + E,)= Dec E,,,, E,) Y Ym where the first and last equalities come from the zero training loss assumption. However, since + * m + n, we have Yi+i # mtmm (almost
n, then
surely in the regression task), contradiction.
D
3

---

(6)
110
10
10
Arc
Mennet
04
0.4
12
0.-
50
0.4
00.
00 02 13.0 115 105
1.0
000
11.
155
10
training data fraction
training set. fraction
Figure 3: We compute accuracy (of the full dataset) either measured empirically Acc, or predicted from the representation of the embeddings Acc. These two accuracies as a function of training data fraction are plotted in (a)(b) and their agreement is shown in (c).
It is convenient define the permissible parallelogram associated with training dataset D ("permissible" means consistent with 100% training accuracy) as
Po(D) = (.).j.m.)))...) D, (m, n) E D, i + = m n..
(1)
For simplicity, we denote Po PooD). Given representation R. we can check how many permissible parallelograms actually exist within error 8, so we define the parallelogram set corresponding to R. as
P(R, 5) {(t,j, m, m)) j. m, E Po. E,) (E.. E,) - 5).
For brevity will write P(R), suppressing the dependence on 8. We define the representation quality index (RQI) as
[P(R] RQI(R) = IPo
E [0,
(3)
We will use the term linear representation linear structure to refer representation whose embeddings are of the form E, = a kb 0, P - 1; a, b E Id..). A linear representation has RQI 1, while random representation (sampled from, say, normal dstribution) has RQI with high probability.
Quantitatively, we denote the "predicted accuracy" Acc as the accuracy achievable on the whole dataset given the representation R (see Appendix D for the full details). In Figure the predicted Acc aligns well with the true accuracy Acc, establishing good evidence that structured representation of input embeddings leads to generalization. We use example to illustrate the origin of generalization here. In the setup Figure 2 (b), suppose the decoder can achieve zero training loss and Eg Eg is training sample hence Dec(E, Eg + Eg) Y14. At validation time, the decoder tasked with predicting a validation sample E, Eg. Since (5,9. 9, 6, 8) forms a parallelogram such that Es Ey Eg Eg. the decoder predict the validation sample correctly because Dec Eg + Eg) = Dec E, + Eg) Y14-
that
3.2 The dynamics of embedding vectors
Suppose that we have an ideal model M = (Dec", R*) such that:?
(1) M can achieve zero training loss;
(2) M has an injective decoder, i.c., Dec (xc) * Dec Dec*(xg) for any # Xg.
Then Proposition 2 provides mechanism for the formation of parallelograms.
'One can verify a posteriori if trained model M is close to being an ideal model M Please refer to Appendix for details.

---

Thoory: phaze transtionn
(b) Empiiiicl:: phose leoraditin
Thaey:; trajedtryy
(d) Empirial: trajecicy
1.0
10
Protebslitytlne stenetne))
niqistss strp-
Time nochee
1D specn (ood Ieptosatationn
0.8
Bume that dado mark
Rigto 56 to 10 step
0.8
0.4
0.2
0.4
0.0
10P
0.0
0.2
0.4
0.6
0.8
1.0
0.4
0.6
D.S
1.0
Iraining datas
tcxining clatlo Traecioon
step
400
Figure 4: (a) The effective theory predicts phase transition in the probability of obtaining linear representation around r. 0.4. (b) Empirical results display a phase transition of RQI around Te 0.4, in agreement with the theory (the blue line shows the median of multiple random seeds). The evolution of ID representations predicted by the effective theory or obtained from neural network training (shown in (c) and (d) respectively) agree creditably well.
Proposition 2. If a training set D contains two samples (i,j) and (m, n) with i + = m + n, then M learns representation R such that Eg + Ej Em En, i.e, (f, m, forms parallelogram.
Proof. Due the zero training loss assumption, we have Dec (E; EE) Yi+j Y Mm4n Dec (Em + E,). Then the injectivity of Dec implies E, + E2 Em + Ex.-
The dynamics of the trained embedding vectors are determined by various factors interacting in complex ways, for instance: the details of the decoder architecture, the optimizer hyperparameters. and the various kinds of implicit regularization induced by the training procedure. We will see that the dynamics of normalized quantities, namely, the normalized embeddings at time t, defined as
E -M where μ4 = E E and o, = by simple effective loss (in the physics effective theory sense). We will assume that the normalized embedding vectors obey a gradient flow for an effective loss function of the form
p Ex [ mel? 12, can be qualitatively described
E
Hl
(4)
dt
to Zo
leff
6o =
Σ
E, + E. E...
E. (/)o(()) Zn =
E
(5)
(i..m,,nE (DD)
where |. denotes Euclidean vector norm. Note that the embeddings do not collapse to the trivial solution Eg _... = Ep-1 0 unless initialized as such, because two conserved quantities exist, as proven Appendix F:
C = Σ E E.. Zu = Σ ExI.
We shall now use the effective dynamics to explain empirical observations such as the existence of a critical training set size for generalization.
Degeneracy ground states (loss optima) We define ground states as those representations satisfying last = 0, which requires the following linear equations to hold:
A(P) = E, + E, Em Ei, (i.j, m, n) E P).
Since each embedding dimension obeys the same set of linear equations, we will assume, without loss of generality, that din 1. The dimension of the null space of A(P), denoted as no, is the number of degrees freedom the ground states. Given set of parallelograms implied by a training dataset D. the nullity of A(P(D)) could be obtained by computing the singular values 0 < S... S op- We always have no 2, i.e., 0 = 02 because the nullity of A(Pa), the set of linear equations given by all possible parallelograms, is Vullity(A(Pa)) 2 which can be attributed to two degrees
5

---

of freedom (translation and scaling). If no = 2, the representation is unique up to translations and scaling factors, and the embeddings have the form E;- a + kb. Otherwise, when no 2, the representation is not constrained enough such that all the embeddings lie on line.
We present theoretical predictions alongside empirical results for addition (p- 10) in Figure 4. As shown in Figure 4 (a), our effective theory predicts that the probability that the training set implies a unique linear structure (which would result in perfect generalization) depends on the training data fraction and has a phase transition around re 0.4. Empirical results from training different models are shown in Figure 4 (b). The number of steps to reach RQI > 0.95 is seen to have a phase transition at Te 0.4, agreeing with the proposed effective theory and with the empirical findings in [1].
Time towards the linear structure We define the Hessian matrix of to
8lo H. Zo OE,OE,
(8)
Note that left = / R HR, R = Eo, E
Ep- 1]. so the gradient descent is linear, i.c.,
dR
dt HR.
If H has eigenvalues X; of (sorted in increasing order) and eigenvectors Vi, and we have the initial condition R(t 0) = E, Vi. then we have R(t) = E, 0g V The first two eigenvalues vanish and t. 1/ determines the timescale for the slowest component to decrease by a factor
of e. We call A3 the grokking rate. When the step size is the corresponding number of steps is
nh 1n/1 = 1/(aa)).
We verify the above analysis with empirical results. Figure 4 (c)(d) show the trajectories obtained from the effective theory and from neural network training. respectively. The ID neural representation in Figure 4 (d) are manually normalized to zero mean and unit variance. The two trajectories agree qualitatively, and it takes about 3n, steps for two trajectories to converge to the linear structure. The quantitative differences might be due to the absence of the decoder in the effective theory, which assumes the decoder to take infinitesimal step sizes.
Dependence of grokking on data size Note that let involves averaging over parallelograms in the training set, it is dependent training data size, is Ag. In Figure 5 (a), we plot the dependence of À on training data fraction. There are many datasets with the same data size, so Xa is a probabilistic function of data size.
Two insights on grokking can be extracted from this plot: (i) When the data fraction is below some threshold (around 0.4), A is zero with high probability, corresponding again verifies our critical point in Figure 4. (ii) When data size is above the threshold, A3 (on average) is an increasing function of data size. This implies that grokking time t 1/\g decreases training data size becomes larger, an important observation from [1].
generalization. This
To verify our effective theory, we compare the grokking steps obtained from real neural network training (defined as steps to RQI > 0.95), and those predicted by our theory 4h ~ Xay (1) is the embedding learning rate), shown Figure (b). The theory agrees qualitatively with neural networks, showing the trend of decreasing grokking steps increasing data size. The quantitative differences might be explained as the gap between our effective loss and actual loss.
Limitations of the effective theory While our theory defines an effective loss based on the Euclidean distance between embeddings E; + E, and E,, + Em, one could imagine generalizing the theory to define a broader notion of parallogram given by some other metric on the representation space. For instance, if we have a decoder like in Figure 2 (d) then the distance between distinct representations within the same "pizza slice" is low, meaning that representations arranged not in parallelograms w.r.t. the Fuclidean metric may be parallelograms with respect to the metric defined by the decoder.
4 Delayed Generalization: A Phase Diagram
So far, we have (1) observed empirically that generalization on algorithmic datasets corresponds with the emergence of well-structured representations, (2) defined a notion of representation quality in a toy setting and shown that it predicts generalization, and (3) developed an effective theory to describe

---

10
0.45
Steps to RQI > 0.95
grokking rate AS
0.3
0.2
10%
0.1
ditct roch 0:50 to"
Raus that teached BQ 0.55 steps
0.0
0.0
0.2
10.4
0.6
0.8
1.
10 grokking rate Aa
10P
training data fraction
(a)
(b)
Figure 5: Effective theory explains the dependence of grokking time data size, for the addition task. (a) Dependence A3 on training data fraction. Above the critical data fraction (around 0.4), as data size becomes larger, '3 increases hence grokking time t 1/Ag (predicted by effective theory) decreases. (b) Comparing grokking steps (defined as RQI > 0.95) predicted by the effective theory with real neural network results. "i 10-3 the learning rate of the embeddings.
the learning dynamics of the representations in the same toy setting. We now study how optimizer hyperparameters affect high-level learning performance. In particular, we develop phase diagrams for how learning performance depends on the representation learning rate, decoder learning rate and the decoder weight decay. These parameters are of interest since they most explicitly regulate kind of competition between the encoder and decoder, as we elaborate below.
4.1 Phase diagram of toy model
Training details We update the representation and the decoder with different optimizers. For the º1 and weight decay.
ID embeddings, use the Adam optimizer with learning rate [10 00-5,,0 10
For the decoder, we use an AdamW optimizer with the learning rate in [10 10 10021 and the weight decay in [0, 10] (regression) 45/10 for non non-modular addition (p = 10) and 24/12 for the permutation group Sy. We hard-code addition matrix multiplication (details in Appendix H) in the decoder for the addition group and the permutation group, respectively.
[0, 20] (classification). For vraining/validation spliting,
choose
For each choice of learning rate and weight decay, compute the number of steps reach high (90%) training/validation accuracy. The 2D plane split into four phases: comprehension, grokking, memorization and confusion, defined in Table I in Appendix A. Both comprehension and grokking able to generalize (in the "Goldilocks zone"), although the grokking phase has delayed generalization. Memorization is also called overfitting, and confusion means failure to memorize training data. Figure shows the phase diagrams for the addition group and the permutation group. They display quite rich phenomena.
Competition between representation learning and decoder overfitting In the regression setup of the addition dataset, we show how the competition between representation learning and decoder learning (which depend on both learning rate and weight decay, among other things) lead different learning phases in Figure 6 (a). As expected, fast decoder coupled with slow representation learning (bottom right) lead to memorization. coupled with fast representation learning (top left) will generalize in the end, the generalization time is long due to the inefficient decoder training. The ideal phase (comprehension) requires representation learning to be faster, but not too much, than the decoder.
the opposite extreme, although an extremely slow decoder
Drawing from an analogy physical systems, one can think of embedding vectors as group of particles. In our effective theory from Section 3.2, the dynamics of the particles are described only by their relative positions, in that sense, structure forms mainly due to inter-particle interactions (in reality, these interactions are mediated by the decoder and the loss). The decoder plays the role of an environment exerting external forces on the embeddings. small/large one expect better/worse representations.
magnitude of the external forces are
7

---

1e22
repressnt ppprectttioo ataiion loarning rate
compreheasion
grokking
memmorization
Cegeee dooag
grokking
1
memorization
confusion
1o-h
10
1e55
1o-4
Le33
1o.2
1e-8
led4
1e-3
1e/2
decoder learning rate
decoder learning
Addition, regression
(b) Addition, regression
mucmorization
memorization
comprehension
Brokiing
weight he clecay
grokking
weighht
10
comprehension
confusion
20
200
confusion
16.5
16-4
16-3
16-2
1e.5
led1
1e-3
Ie22
drooder learning rate
decoder learning rate
(c) Addition, classification
(d) Permutation, regression
Figure 6: Phase diagrams of learning for the addition group and the permutation group. (a) shows the competition between representation and decoder. (b)())d): each phase diagram contains four phases: comprehension, grokking, memorization and confusion, defined in Table 1. In (b)(c)dd) grokking is sandwiched between comprehension and memorization.
Universality of phase diagrams We fix the embedding learning rate be 10-3 and sweep instead decoder weight decay in Figure (d). The phase diagrams correspond to addition regression (b), addition classification (c) and permutation regression (d), respectively. Common phenomena emerge from these different tasks: (i) they all include four phases; (ii) The top right corner (a fast and capable decoder) is the memorization phase; (iii) the bottom right corner fast and simple decoder) the confusion phase; (iv) grokking is sandwiched between comprehension and memorization, which seems imply that is an undesirable phase that stems from improperly tuned hyperpurameters.
4.2 Beyond the toy model
We conjecture that many of the principles which we saw dictate the training dynamics in the toy model also apply architectures for the task of addition modulo a minimal reproducible example of the original grokking paper [1].
generally. Below, will see how framework generalizes to transformer
We first encode = 53 integers into 256D learnable embeddings, then pass two integers to decoder- only transformer architecture. For simplicity, we do not encode the operation symbols here. The outputs from the last layer are concatenated and passed to a linear layer for classification. Training both the encoder and the decoder with the same optimizer (i.e., with the same hyperparameter)) leads to the grokking phenomenon. Generalization appears much earlier once we lower the effective decoder capacity with weight decay (full phase diagram in Figure 7).
Early on, the model is able to perfectly fit the training set while having no generalization. We study the embeddings at different training times and find that neither PCA (shown in Figure 1) nor t-SNE (not shown here) reveal any structure. Eventually, validation accuracy starts to increase, and perfect generalization coincides with the PCA projecting the embeddings into circle in 2D. Of course,

---

55
10
1e
validation
50
<peebe 90% secimey
train
mnomorization
elfentivn dinension
validation train
20-1
10**
40
2o++0
grokking
35
generalization
30
memoriiatton
36++1
coufusion
10"
Io-4
step
druport rate
loannnng tale
Figure 7: Left: Evolution of the effective dimension of the embeddings (defined the exponential of the entropy) during training and evaluated over 100 seeds. Center: Effect of dropout speeding up generalization. Right: Phase diagram of the transformer architecture. A scan is performed over the weight decay and learning rate of the decoder while the learning rate of the embeddings is kept fixed 10-3 (with zero weight decay).
choice of dimensionality reduction guaranteed to find any structure, and thus, it is challenging to show explicitly that generalization only occurs when a structure exists. Nevertheless, the fact that, when coupled with the implicit regularization of the optimizer for sparse solutions, such a clear structure appears in a simple PCA so quickly at generalization time suggests that our analysis in the toy setting applicable here as well. This is also seen the evolution of the entropy of the explained variance ratio in the PCA of the embeddings (defined as S - Σ, de log where a; is the fractional variance explained by the ith principal component). As seen in Figure 7, the entropy increases up generalization time then decreases drastically afterwards which would be consistent with the conjecture that generalization oceurs when low-dimensional structure is discovered. The decoder then primarily relies on the information in this low-dimensional manifold and essentially "prunes" the rest the high-dimensional embedding space. Another interesting insight appears when project the embeddings initialization onto the principal axes at the end of training. Some of the structure required for generalization exists before training hinting at a connection with the Lottery Ticket Hypothesis. See Appendix K for more details.
In Figure (right), we show a comparable phase diagram to Figure 6 evaluated now in the transformer setting. Note that, as opposed to the setting in [1], weight decay has only been applied to the decoder and not to the embedding layer. Contrary to the toy model, certain amount of weight decay proves beneficial to generalization and speeds it up significantly. We conjecture that this difference comes from the different embedding dimensions. With a highly over-parameterized setting, a non-zero weight decay gives a crucial incentive reduce complexity in the decoder and help generalize in fewer steps. This is subject to further investigation. We also explore the effect of dropout layers in the decoder blocks of the transformer. With significant dropout rate, the generalization time can be brought down to under 1 steps and the grokking phenomenon vanishes completely. The overall trend suggests that constraining the decoder with the same tools used avoid overfitting reduces generalization time and can avoid the grokking phenomenon. This is also observed in an image classification task where were able induce grokking. See Appendix J for more details.
4.3 Grokking Experiment on MNIST
We now demonstrate, for the first time, that grokking (significantly delayed generalization) is more general phenomenon in machine learning that can occur not only on algorithmic datasets, but also on mainstream benchmark datasets. In particular, we exhibit grokking on MNIST in Figure 8 and demonstrate that we can control grokking by varying optimization hyperparameters. More details on the experimental setup are in Appendix
5 Related work
Relatively few works have analyzed the phenomenon of grokking. [2] describe the circuit that transformers use to perform modular addition, track its formation over training, and broadly suggest that grokking is related to the phenomenon of "phase changes" in neural network training. [3. 4]

---

Train Points: 1000 Initialization Scale: 9.0
Memorraation
1.0
val
0.8
Weight Deckyy
0.6
Grakking
Compreheralon
2.660
Cantatin
7.30-900
2444-0
10'
10a
Last Layer Learning Rate
Optimization Steps
(a)
Figure 8: Left: Training curves for run MNIST, in the setting where we observe grokking. Right: Phase diagram with the four phases of learning dynamics on MNIST.
provided earlier speculative, informal conjectures on grokking [3, 4]. Our work is related to the following broad research directions:
Learning mathematical structures [5] trains neural network to learn arithmetic operation from pictures of digits, but they do not observe grokking due to their abundant training data. Beyond arithmetic relations, machine learning has been applied to learn other mathematical structures, including geometry [6], knot theory [7] and group theory [8].
Double descent Grokking is somewhat reminiscent of the phenomena of "epoch-wise" double descent [9], where generalization can improve after a period of overfitting. [10] find that regularization mitigate double descent, similar perhaps how weight decay influences grokking.
Representation learning Representation learning lies at the core of machine learning [11-14]. Representation quality is usually measured by (perhaps vague) semantic meanings performance on downstream tasks. In our study, the simplicity of arithmetic datasets allows us to define representation quality and study evolution of representations in a quantitative way.
Physies of learning Physics-inspired tools have proved to be useful in understanding deep learning from a theoretical perspective. These tools include effective theories [15, 16]. conservation laws [17] and free energy principle [18]. In addition, statistical physics has been identified as powerful tool in
studying generalization in neural networks [19-22]. Our work connects a low-level understanding of recent work, researchers Anthropic [23], connect
models with their high-level performance. sudden decrease in loss during training with the emergence of induction heads within their models. They analogize their work to statistical physics, since it bridges "microscopic", mechanistic understanding of networks with "macroscopic" facts about overall model performance.
Conclusion
We have shown how, in both toy models and general settings, that representation enables generalization when reflects structure in the data. We developed an effective theory of representation learning dynamics (in a toy setting) which predicts the critical dependence of learning the training data fraction. We then presented four learning phases (comprehension, grokking, memorization and confusion) which depend the decoder capacity and lcarning speed (given by, among other things, learning rate and weight decay) decoder-only architectures. While we have mostly focused on a toy model, we find preliminary evidence that our results generalize the setting of [1].
Our work can be viewed as a step towards a statistical physics of deep learning, connecting the "microphysics" of low-level network dynamics with the "thermodynamics" of high-level model behavior. We view the application theoretical tools from physics, such as effective theories [24], to be a rich area for further work. The broader impact of such work, if successful, could be to make models more transparent and predictable [23, 25, 26), crucial to the task of ensuring the safety of advanced Al systems.
10

---

References
[1] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:220.022177, 2022.
[2] Neel Nanda and Tom Lieberum. A mechanistic interpretability analysis of grokking, 2022. URL httptt:////wwt http://wwwww.lll.......oo..r.. www.oooww....o.reeeeeeeeeeeeeeee a-mechanistic- wtt.ssiiiiineeeeeaaiiiiite..
[3] Beren Millidge.
Grokking
'gokkiing'.
https://beren.iee
2022- 01-10-1-11111111 1iookiriiikngg--- Grokking/, 2022.
[4] Rohin Shah.
Alignment
Newsletter
#159.
https:
//www. alignmentforum. org/ https///rrwggmaasssaaaa.aaaeeeea an- 159- bolgdinuiiinggg agents- that-know-hee to- experiment-byteaaeeee LEARNING
2021.
[5] Yedid Hoshen and Shmuel Peleg. Visual learning of arithmetic operation. In AAAI, 2016.
[6] Yang-Hui He. Machine-learning mathematical structures. arXiv preprint arXiv:2101.06117 2021.
[7] Sergei Gukov, James Halverson, Fabian Ruchle, and Piotr Sulkowski. Learning to unknot. Machine Learning: Science and Technology, 2(2):025035, 2021.
[8] Alex Davies, Petar Veličković, Lars Buesing, Sam Blackwell, Daniel Zheng, Nenad Tomašev, Richard Tanburn, Peter Battaglia, Charles Blundell, András Juhász, et al. Advancing mathemat- ies by guiding human intuition with ai. Nature, 600(7887):70-74 2021.
[9] 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(12) 2021(12):22003: 2021.
[10] Preetum Nakkiran, Prayaag Venkat, Sham Kakade, and Tengyu Ma. Optimal regularization can mitigate double descent. arXiv preprint arXiv:2003.0/897 2020.
[11] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8): 1798-1828, 2013.
[12] Yassine Ouali, Céline Hudelot, and Myriam Tami. An overview of deep sumi-supervised learning, arXiv preprint arXiv:2006.05277 2020.
[13] Jean-Base Cean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent-a new approach self-supervised learning. Advances Neural Information Processing Systems, 33:21271 111777221884 2020.
[14] Phuc H Le-Khac, Graham Healy, and Alan F Smeaton. Contrastive representation learning: A framework and review. IEEE Access, 1-193907-1939933 2020.
[15] James Halverson, Anindita Maiti, and Keegan Stoner. Neural networks and quantum field theory. Machine Learning: Science and Technology, 2(3):035002, 2021.
[16] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint 11X.v:2006.10165 2021.
[17] Daniel Kunin, Javier Sagastuy- Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728. 2020.
[18] Yansong Gao and Pratik Chaudhari. free-energy principle for representation learning. In International Conference Machine Learning, pages 3367-3376. PMLR, 2020.

---

[19] Federica Gerace, Bruno Loureiro, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Generalisation enror in learning with random features and the hidden manifold model. In International Conference on Machine Learning, pages 3452-3462. PMLR, 2020.
[20] Mohammad Pezeshki, Amartya Mitra, Yoshua Bengio, and Guillaume Lajoie. Multi-scale feature learning dynamics: Insights for double descent. In International Conference on Machine Learning, pages 17669-17690. PMIR, 2022.
[21] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Mare Mezard, and Lenka Zdeborova. The gaussian equivalence generative models for learning with shallow neural networks. In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference, volume 145 of Proceedings of Machine Learning Research, pages 426-471. PMLR, 16-19 Aug 2022. URL https: proceeddnng mlr. peaas//v45/golle
html.
[22] R Kuhn and S Bos. Statistical mechanics for neural networks with continuous- dynamics. Journal of Physics Mathematical and General, 26(4):831, 1993.
[23] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hattield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https:/transsfomer-eieee wwtp:ppppbbmmoommceei... context-learning and-induction- heads/index.htm.
[24] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. Cambridge University Press, 2022. https
com.
[25] Deep Ganguli, Danny Hemaandez, Liane Lovitt, Nova DasSarma, Tom Henighan, Andy Jones, Nicholas Joseph, Jackson Kermion, Ben Mann, Amanda Askell, et al. Predictability and surprise large generative models. arXiv preprint XXXi:22222.07785. 2022.
[26] Jacob Steinhardt. Future ML Systems Will Be Qualitatively Different. https://www.
lesswrong wtt.c/oo///a/ceee
2022.
[27] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
[28] Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117 200-22665222663, 2020.
[29] Wikipedia contributors.
Thomson problem - Wikipedia, the free encyclope- dia. https: //en wikipedia
httte///eeee//aaaaaa//a..aaaaaa. org/w/ hrg/w/indens w/index..../////////
1091431454, 2022. [Online; accessed 29-July-2022..
[30] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15750-15758, 2021.
[31] Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In International Conference on Neural Information Processing, pages 264-274, Springer, 2019.
[32] Yaoyu Zhang. Zhi-Qin John Xu, Tao Luo, and Zheng Ma. A type of generalization enror induced by initialization in deep neural networks. In Mathematical and Scientific Machine Learning, pages 144-164. PMLR, 2020.
[33] Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data, 2022.
12

---

[34] Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes overparametriized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-12 Jul 2020. URL html.
mlr. press/ sasssvv1225 v125/ htt///ooooothte.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethies review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all thcoretical results? [Yes]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] All experiments were run on workstation with two NVIDIA A6000 GPUs within a few days.
4. If you are using existing assets (e.g.. code, data, models) curating/relcasing new assets...
(a) If your work existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you's 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "dVKC_Lqp41C", "reviewer": "Reviewer_k44Q", "summary": "The main contribution of this paper is a study of how different loss minimization and generalization outcomes are realized: confusion (no loss minimization), memorization (loss minimization without generalization), fast generalization, and grokking (generalization, but significantly slower than the loss is minimized). There is a secondary result on the minimal training set size required to induce a representation in which different classes are linearly separable. All tasks studied in the paper are related to arithmetic operations like addition, with an appendix explaining that the setup studied in the paper can be used for all Abelian groups. The paper contains a thorough exploration of a toy system, which is used to analyze and illustrate the above-mentioned outcomes. Some results are included on a larger task with larger networks. ", "questions": "- Can you be a bit more explicit about the loss in the beginning? It would help reading the paper.\n- If I read the paper correctly, you are saying that grokking happens when different parts of the network absorb information at different rates; in particular, when the last part (the decoder) is much faster than the preceding part. Is that a reasonable formulation? And if yes, is it possible to quantify that statement?\n- What constitutes a part, for the purposes of the above question? (see also discussion of A1 in the ‘strengths and weaknesses’ section)\n- How generally does your analysis apply? I can see from the appendices that it should generalize to other Abelian groups, but can you comment on different applications like image classification? The answers you give - linear representation, relative learning speeds of the decoder and upstream subnetwork - are not specific to the type of problem you study, so it would be very interesting (and make the paper much stronger) if you could also include results on other types of task.\n- How strongly does the result for the critical training set size to obtain a linear representation depend on the simple effective loss you are using? Would it be possible to perform a similar analysis for image classification (or language modelling, or any other application you fancy)?\n- Regarding the p=53 modulo addition experiment: is there a linear subspace in which you find the pizza wheel representation, like for the toy task? I understand it’s hard to visually inspect the 256D embeddings, but it might be interesting to see how many ‘effective dimensions’ there are - e.g. how many PCA components have significant explanatory power. It would be reasonable to expect a low-dimensional effective embedding, given how neatly t-SNE puts all the numbers in a circle. What happens if you take the two most explanatory principal components and make the same decoder plot there that you have for the toy task?\n- How do your findings relate to the lottery ticket hypothesis? That also addresses generalization in unexpected situations - in overparameterized networks as opposed to late in training. Would having more highly overparameterized representation networks be another potential way to avoid grokking, given the lottery ticket hypothesis that bigger networks have higher chances of containing subnetworks with good initializations, that would therefore learn quickly?\n- In the particle system analogy: What are the ‘internal particle interactions’? If they arise from the gradients computed from the loss, then they really originate ‘externally’ too, don’t they?\n", "limitations": "The main limitation comes from the size of the systems studied, as mentioned above. In principle the analysis could apply quite widely, but it is difficult to assess that from the experiments in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Strength: the exploration of the toy system is very thorough and offers an interesting ‘peek inside the black box’. It suggests very plausible mechanisms for how the learning process plays out.\n- Weakness: the paper is quite limited in terms of bigger systems, which makes it difficult to assess how general the results are.\n- Weakness: The current formulation of A1 (the answer to question Q1, when does generalization happen), while true, seems both a bit trivial and a bit vague. You say generalization happens ‘with a good representation’, or a representation whose ‘structure is appropriate for the task’. This answer leaves open the questions which representation needs to be good (is it the penultimate layer of the network? the input to the decoder? but where does the decoder start and the representation network end?) and when a representation is ‘good’ or a structure ‘appropriate’. Presuming that those questions are answered, it comes down to the statement that the mapping from the representation to the final required output is easy to learn. For the toy task you analyze you do go into the structure of the representation quite extensively, but since that analysis is limited to the toy setup, it is not obvious that the answers you find there generalize to harder / bigger tasks, deeper networks, and higher-dimensional representations. NB: A2 is only as meaningful as A1 is clear, which makes this issue more relevant.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "15-slHkOBk5", "reviewer": "Reviewer_Sd61", "summary": "This work studies the Grokking phenomenon -- a phenomenon in which generalization occurs long after mastering the training set. This work poses three questions with regard to this phenomenon:\n1. When the generalization happens? and the answer is: when the representations become well-structured.\n2. How the size of the training set affects this phenomenon? and the answer is: the critical size of the dataset is the amount required for learning a 'good' representation.\n3. What leads to a delay in generalization? and the answer is: different phases of learning between \"memorization\" (perfect training acc but with lack of generalization) and \"comprehension\" (perfect training and generalization acc)\n\nFrom a physics perspective, this work build the answers outlined above and provide supporting experiments.", "questions": "1. The choice of Parallelogram-ratio for quantifying the quality of a representation is rather arbitrary. What is the idea behind using such a quantity?\n2. Why don't you model the actual dynamics of GD on the training loss? Instead, you derive the dynamics of the 'effective' loss (defined in Eq. 5) which again appears arbitrary.\n3. By defining the 'effective' loss, in essence the learning dynamics of the encoder is now decoupled from the decoder. I am right?\n4. I have a hard time understanding lines 148-153. Would you explain why the third eigenvalue is the one that dominate the dynamics?\n5. The analogy in lines 197-201; Don't you agree that it is the decoder that encourages the representations to have a structure rather than their \"internal interactions\". If the is not decoder, what forces the representations to change structure?", "limitations": "No negative societal impact is foreseen by this reviewer.\nAuthors talk about the limitations of the theory in line 160-165.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n- The problem is interesting and relevant: The Grokking phenomenon is a perplexing phenomenon which understanding it opens the doors to a better theory of generalization.\n- The form of presentation is neat: I very much liked that the message is clear from the get go.\n- The phase diagrams provide high-level intuitions: The same applies to the phase-diagrams that clearly distinguish between the 4 different phases.\n\n** Weaknesses**:\n- Some of the choices in the paper appears rather arbitrary and not fully justified (see questions below).\n- At parts, it is difficult to follow certain arguments (at least for me): I am not from a physics back ground. For that it took me several wikipedia searches to grasp the meaning of some concepts in the paper. For example the term \"effective\" has a different meaning that the one that initially comes to mind.\n- Very light on the related work: There has been a couple of papers that use tools from physics and provide phase diagrams for understanding the generalization of neural networks:\n    + Generalisation error in learning with random features and the hidden manifold model by Gerace et al\n    + Multi-scale Feature Learning Dynamics: Insights for Double Descent by Pezeshki et al\n    + The Gaussian equivalence of generative models for learning with two-layer neural networks by Goldt et al\n    + Statistical mechanics for neural networks with continuous-time dynamics by Kuhn et al\n- Very limited experimental results.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0lmtJJJUIL3", "reviewer": "Reviewer_tEaJ", "summary": "This paper studies the grokking phenomenon under a physics-inspired \"effective theory\", which studies the phenomenon under toy settings. The grokking phenomenon refers to when generalization happens much later during training, long after training accuracy have reached 100%. The theory proposes the following explanations to grokking.\n1) generalization is correlated to when the models learns a sensible representation of the inputs.\n2) the critical training size (amount of data needed for generalization) relates to the amount of data needed to determine a good representation.\n3) grokking represents a particular region in phase space of hyperparameters, and can be removed by using better hyperparameters.", "questions": "- Is the analysis on \"time towards the linear structure\" independent of dataset size? Or can it shed light on the question of data size vs grokking / comprehension?\n\n- What type of task characteristic makes a model most susceptible to grokking? If the mismatch between decoder learning speed and representation learning speed is key to grokking, would you expect that the four phases in the phase diagram to exist in general (e.g. for natural vision or nlp tasks)? If it's not general, what is it about the specific tasks and models studied in this work that leads to these different regions?", "limitations": "The authors address the limitations of the toy model. Though it is fine to use toy models to study these type of questions, I think that there is nonetheless insufficient work to show whether the insights from the toy model can actually generalize to other settings.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strength:\n\n- The paper proposes an original view into the grokking phenomenon, which has been an interesting and unexplained phenomenon in the community. The phase diagram view into the landscape of training regimens illustrates interesting relationships between memorization, generalization, and grokking. One of the interesting questions about grokking is whether there is some magical inductive bias of gradient descent that leads to generalization after an abnormally long training period. The result of this paper would suggest that grokking is more of a pathological phenomenon that is the result of improper hyperparameter tuning. Though as I'll mention in the weakness section, I think that the support for this argument is not sufficient if it is indeed the conclusion of the authors.\n\n- The paper proposes an explanation into the grokking phenomenon in the form of competition between representation learning and decoder overfitting. This seems plausible and is in line with a number of other works suggesting that later layers drive overfitting.\n\n- Although the theory is based on a very simplified toy model, it provides an illuminating view into why a minimum number of training examples may be required for generalization. The predicted phase transition in Figure 4 is interesting to see.\n\n### Weaknesses:\n\n- One of the intriguing observations from Power et al [1] is that grokking can be seen both in unregularized models with sufficient data, or in regularized models with a small amount of training data. I think a sufficient explanation of the phenomenon needs to provide a justification for both behavior. However, the discussion around the effect of data size seems lacking. As far as I can tell, neither the theoretical nor the empirical analyses shed light on why more data speeds up generalization while less data can lead to grokking.\n\n- Most of the results use an extremely toy setting. Only one set of experiments (modulo addition with varying learning rate and weight decay) is performed on the original grokking task from [1], while [1] contains a number of tasks and hyperparameter settings. Crucially the effect of data size is not studied, which was an important observation from [1]. It is unclear how much this theory really explains the original phenomenon, particularly given that the behavior of the phase diagram is qualitatively different between the setting in [1] and the toy model, namely that instead of models going from memorization to comprehension to grokking, in the real task setting it is going from memorization to grokking to comprehension. There is insufficient explanation to bridge this gap.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept